#!/usr/bin/env python3 """ Standalone worker tool for the distributed download simulation. This tool is responsible for the "lock-execute-unlock" workflow for a single download task based on an info.json file. It's designed to be called by an orchestrator like `stress_policy_tool.py`. """ import argparse import json import logging import os import sys import time from copy import deepcopy try: from dotenv import load_dotenv except ImportError: load_dotenv = None # Temporarily add project root to path to allow importing from sibling packages script_dir = os.path.dirname(os.path.abspath(__file__)) project_root = os.path.abspath(os.path.join(script_dir, '..')) if project_root not in sys.path: sys.path.insert(0, project_root) from ytops_client.profile_manager_tool import ProfileManager from ytops_client.stress_policy import utils as sp_utils from ytops_client.stress_policy.state_manager import StateManager from ytops_client.stress_policy.utils import load_policy, apply_overrides from ytops_client.stress_policy.workers import _run_download_logic from ytops_client.stress_policy_tool import shutdown_event # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) def add_locking_download_emulator_parser(subparsers): """Adds the parser for the 'download-emulator' command.""" parser = subparsers.add_parser( 'download-emulator', help='(Internal) Standalone download worker.', description='Internal tool to run a single download task with profile locking. Not intended for direct user invocation.' ) # Since this is an internal tool, we expect one subcommand. download_subparsers = parser.add_subparsers(dest='download_emulator_command', help='Action to perform', required=True) run_parser = download_subparsers.add_parser( 'lock-and-run', help='Lock a profile, run a download, and unlock it.', formatter_class=argparse.RawTextHelpFormatter ) run_parser.add_argument('--policy-file', required=True, help='Path to the YAML policy file.') run_parser.add_argument('--info-json-path', required=True, help='Path to the info.json file to process.') run_parser.add_argument('--set', action='append', default=[], help="Override a policy setting using 'key.subkey=value'.") # Redis connection arguments, to be passed from the orchestrator redis_group = run_parser.add_argument_group('Redis Connection') redis_group.add_argument('--env-file', help='Path to a .env file.') redis_group.add_argument('--redis-host', help='Redis host.') redis_group.add_argument('--redis-port', type=int, help='Redis port.') redis_group.add_argument('--redis-password', help='Redis password.') redis_group.add_argument('--env', help="Environment name for Redis key prefix.") redis_group.add_argument('--key-prefix', help='Explicit key prefix for Redis.') run_parser.add_argument('--verbose', action='store_true', help='Enable verbose logging.') def main_locking_download_emulator(args): """Main logic for the 'download-emulator' tool.""" if args.verbose: logging.getLogger().setLevel(logging.DEBUG) # --- Load Policy --- policy = load_policy(args.policy_file) if not policy: return 1 policy = apply_overrides(policy, args.set) # --- Load .env if specified --- if load_dotenv and args.env_file: if load_dotenv(args.env_file): logger.info(f"Loaded environment variables from {args.env_file}") else: logger.error(f"Specified --env-file was not found: {args.env_file}") return 1 # --- Setup ProfileManager --- redis_host = args.redis_host or os.getenv('REDIS_HOST', 'localhost') redis_port = args.redis_port or int(os.getenv('REDIS_PORT', 6379)) redis_password = args.redis_password or os.getenv('REDIS_PASSWORD') if args.key_prefix: key_prefix = args.key_prefix elif args.env: key_prefix = f"{args.env}_profile_mgmt_" else: logger.error("Must provide --env or --key-prefix for Redis connection.") return 1 manager = ProfileManager( redis_host=redis_host, redis_port=redis_port, redis_password=redis_password, key_prefix=key_prefix ) download_policy = policy.get('download_policy', {}) profile_prefix = download_policy.get('profile_prefix') if not profile_prefix: logger.error("Policy file must specify 'download_policy.profile_prefix'.") return 1 # --- Main Lock-Execute-Unlock Logic --- owner_id = f"dl-emulator-{os.getpid()}" locked_profile = None lock_attempts = 0 try: # --- 1. Lock a profile (with wait & backoff) --- while not shutdown_event.is_set(): locked_profile = manager.lock_profile(owner=owner_id, profile_prefix=profile_prefix) if locked_profile: logger.info(f"Locked profile '{locked_profile['name']}' with proxy '{locked_profile['proxy']}'.") break # Simplified wait logic from stress_policy_tool backoff_seconds = [3, 5, 9, 20, 50] sleep_duration = backoff_seconds[min(lock_attempts, len(backoff_seconds) - 1)] logger.info(f"No download profiles available. Waiting {sleep_duration}s... (attempt {lock_attempts + 1})") time.sleep(sleep_duration) lock_attempts += 1 if not locked_profile: logger.warning("Could not lock a profile; shutting down.") return 1 # --- 2. Read info.json --- try: with open(args.info_json_path, 'r', encoding='utf-8') as f: info_json_content = f.read() except (IOError, FileNotFoundError) as e: logger.error(f"Could not read info.json file '{args.info_json_path}': {e}") return 1 # --- 3. Execute download logic --- # The locked profile's proxy MUST be used for the download. local_policy = deepcopy(policy) local_policy.setdefault('download_policy', {})['proxy'] = locked_profile['proxy'] # The StateManager is used by _run_download_logic for rate limiting and cooldowns, # but for this standalone worker, we don't need its persistence features. # We disable log writing to prevent creating state files. dummy_state_manager = StateManager(policy_name="locking_emulator_run", disable_log_writing=True) results = _run_download_logic( source=args.info_json_path, info_json_content=info_json_content, policy=local_policy, state_manager=dummy_state_manager, args=args, # Pass orchestrator args through profile_name=locked_profile['name'], profile_manager_instance=manager ) # --- 4. Record overall task activity --- # Note: Download-specific activity ('download'/'download_error') is recorded # inside _run_download_logic -> process_info_json_cycle. download_success = all(r.get('success') for r in results) if results else False activity_type = 'success' if download_success else 'failure' manager.record_activity(locked_profile['name'], activity_type) logger.info(f"Finished processing '{sp_utils.get_display_name(args.info_json_path)}' with profile '{locked_profile['name']}'. Overall success: {download_success}") return 0 if download_success else 1 except Exception as e: logger.error(f"An unexpected error occurred: {e}", exc_info=True) return 1 finally: # --- 5. Unlock the profile --- if locked_profile: cooldown_str = manager.get_config('unlock_cooldown_seconds') cooldown_seconds = int(cooldown_str) if cooldown_str and cooldown_str.isdigit() else None if cooldown_seconds and cooldown_seconds > 0: logger.info(f"Unlocking profile '{locked_profile['name']}' with a {cooldown_seconds}s cooldown.") manager.unlock_profile(locked_profile['name'], owner=owner_id, rest_for_seconds=cooldown_seconds) else: logger.info(f"Unlocking profile '{locked_profile['name']}'.") manager.unlock_profile(locked_profile['name'], owner=owner_id) if __name__ == '__main__': # This is a simplified parser setup for direct execution, # the real one is defined in `add_..._parser` for the main CLI. parser = argparse.ArgumentParser() subparsers = parser.add_subparsers(dest='command') add_locking_download_emulator_parser(subparsers) args = parser.parse_args() if hasattr(args, 'download_emulator_command'): sys.exit(main_locking_download_emulator(args)) else: parser.print_help()