210 lines
8.7 KiB
Python
210 lines
8.7 KiB
Python
#!/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()
|