diff --git a/core/download/aria2_backend.py b/core/download/aria2_backend.py index a859950..5dbc8b4 100644 --- a/core/download/aria2_backend.py +++ b/core/download/aria2_backend.py @@ -150,3 +150,387 @@ def aria2_action_error_is_ok(status: str, message: str) -> bool: if status == "downloading": return "not paused" in lowered or "has not been paused" in lowered return False + + +def read_aria2_version(executable: str) -> str: + """Read the installed aria2 version without raising process errors.""" + facade = _downloader_module() + if not executable: + return "" + try: + kwargs: Dict[str, Any] = {} + if facade.os.name == "nt": + kwargs["creationflags"] = getattr( + facade.subprocess, + "CREATE_NO_WINDOW", + 0, + ) + result = facade.subprocess.run( + [executable, "--version"], + stdout=facade.subprocess.PIPE, + stderr=facade.subprocess.STDOUT, + text=True, + timeout=5, + check=False, + **kwargs, + ) + except Exception: + return "" + + first_line = "" + for line in str(result.stdout or "").splitlines(): + text = line.strip() + if text: + first_line = text + break + if not first_line: + return "" + + for token in first_line.replace(",", " ").split(): + if token and token[0].isdigit(): + return token + return first_line + + +def get_aria2_status(settings: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + """Return aria2 availability, daemon and active-transfer information.""" + facade = _downloader_module() + active_settings = ( + settings if isinstance(settings, dict) else facade.load_settings() + ) + configured_path = str(active_settings.get("aria2c_path") or "").strip() + try: + resolved_path = facade._resolve_aria2c_executable(active_settings) + available = True + version = facade._read_aria2_version(resolved_path) + error = "" + except Exception as exc: + resolved_path = "" + available = False + version = "" + error = str(exc) + + with facade.aria2_lock: + process = facade.aria2_process + running = process is not None and process.poll() is None + if not running and process is not None: + facade.aria2_process = None + facade.aria2_process_started_by_resolver = False + managed = bool(running and facade.aria2_process_started_by_resolver) + active_transfers = len(facade.aria2_transfers) + + return { + "backend": facade._download_backend_from_settings(active_settings), + "configured_path": configured_path, + "resolved_path": resolved_path, + "available": available, + "version": version, + "running": running, + "managed": managed, + "can_stop": bool(managed and active_transfers == 0), + "active_transfers": active_transfers, + "auto_stop_enabled": bool(active_settings.get("aria2_auto_stop_daemon", True)), + "idle_stop_seconds": facade.ARIA2_IDLE_STOP_SECONDS, + "error": error, + } + + +def aria2_rpc(method: str, params: Optional[list[Any]] = None) -> Any: + """Call the local aria2 JSON-RPC endpoint.""" + facade = _downloader_module() + if not facade.aria2_rpc_url: + raise Aria2Error("aria2 RPC endpoint is not initialized") + + payload = { + "jsonrpc": "2.0", + "id": facade.secrets.token_hex(8), + "method": method, + "params": [f"token:{facade.aria2_rpc_secret}", *(params or [])], + } + with facade.aria2_rpc_lock: + response = facade.requests.post( + facade.aria2_rpc_url, + json=payload, + timeout=facade.ARIA2_RPC_TIMEOUT, + ) + text = response.text + try: + body = response.json() + except ValueError as exc: + raise Aria2Error( + f"aria2 RPC returned non-JSON response ({response.status_code}): {text[:300]}" + ) from exc + + if "error" in body: + error = body.get("error") or {} + message = error.get("message") if isinstance(error, dict) else str(error) + raise Aria2Error(message or f"aria2 RPC {method} failed") + + if response.status_code != 200: + raise Aria2Error( + f"aria2 RPC {method} returned HTTP {response.status_code}: {text[:300]}" + ) + + return body.get("result") + + +def aria2_ping() -> bool: + """Check whether the local aria2 RPC endpoint responds.""" + facade = _downloader_module() + try: + result = facade._aria2_rpc("aria2.getVersion", []) + return isinstance(result, dict) + except Exception: + return False + + +def cancel_aria2_idle_timer_locked() -> None: + """Cancel the pending idle timer while holding the aria2 lock.""" + facade = _downloader_module() + if facade.aria2_idle_timer is not None: + facade.aria2_idle_timer.cancel() + facade.aria2_idle_timer = None + + +def aria2_has_active_transfers_locked() -> bool: + """Return whether aria2 has active resolver-owned transfers.""" + return bool(_downloader_module().aria2_transfers) + + +def stop_aria2_daemon(reason: str = "manual") -> Dict[str, Any]: + """Stop the aria2 RPC process started by Model Resolver.""" + facade = _downloader_module() + with facade.aria2_lock: + facade._cancel_aria2_idle_timer_locked() + process = facade.aria2_process + running = process is not None and process.poll() is None + if not running: + facade.aria2_process = None + facade.aria2_rpc_url = "" + facade.aria2_rpc_secret = "" + facade.aria2_process_started_by_resolver = False + return { + "success": True, + "stopped": False, + "message": "aria2 daemon is not running", + } + + if not facade.aria2_process_started_by_resolver: + return { + "success": False, + "stopped": False, + "error": "This aria2 daemon was not started by Model Resolver.", + } + + if facade._aria2_has_active_transfers_locked(): + return { + "success": False, + "stopped": False, + "error": "aria2 daemon has active downloads.", + } + + try: + process.terminate() + process.wait(timeout=5) + except Exception: + try: + process.kill() + process.wait(timeout=2) + except Exception: + pass + + facade.aria2_process = None + facade.aria2_rpc_url = "" + facade.aria2_rpc_secret = "" + facade.aria2_process_started_by_resolver = False + + log.info(f"aria2 RPC daemon stopped ({reason})") + return {"success": True, "stopped": True, "message": "aria2 daemon stopped"} + + +def start_aria2_daemon(settings: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + """Start the aria2 RPC process without creating a download.""" + facade = _downloader_module() + active_settings = ( + settings if isinstance(settings, dict) else facade.load_settings() + ) + try: + facade._ensure_aria2_daemon(active_settings) + status = facade.get_aria2_status(active_settings) + return { + **status, + "success": True, + "started": bool(status.get("running")), + "message": "aria2 daemon is running", + } + except Exception as exc: + try: + status = facade.get_aria2_status(active_settings) + except Exception: + status = {} + return { + **status, + "success": False, + "started": False, + "error": str(exc), + } + + +def aria2_idle_stop_worker() -> None: + """Stop a resolver-owned daemon after its idle timeout.""" + facade = _downloader_module() + with facade.aria2_lock: + facade.aria2_idle_timer = None + + settings = facade.load_settings() + if not settings.get("aria2_auto_stop_daemon", True): + return + with facade.aria2_lock: + process = facade.aria2_process + running = process is not None and process.poll() is None + if ( + not running + or not facade.aria2_process_started_by_resolver + or facade._aria2_has_active_transfers_locked() + ): + return + facade.stop_aria2_daemon(reason="idle") + + +def schedule_aria2_idle_stop() -> None: + """Schedule daemon shutdown when no resolver transfer remains active.""" + facade = _downloader_module() + settings = facade.load_settings() + if not settings.get("aria2_auto_stop_daemon", True): + return + with facade.aria2_lock: + facade._cancel_aria2_idle_timer_locked() + process = facade.aria2_process + running = process is not None and process.poll() is None + if ( + not running + or not facade.aria2_process_started_by_resolver + or facade._aria2_has_active_transfers_locked() + ): + return + facade.aria2_idle_timer = facade.threading.Timer( + facade.ARIA2_IDLE_STOP_SECONDS, + facade._aria2_idle_stop_worker, + ) + facade.aria2_idle_timer.daemon = True + facade.aria2_idle_timer.start() + + +def ensure_aria2_daemon(settings: Optional[Dict[str, Any]] = None) -> None: + """Start or reuse the resolver-owned aria2 daemon.""" + facade = _downloader_module() + active_settings = ( + settings if isinstance(settings, dict) else facade.load_settings() + ) + with facade.aria2_lock: + facade._cancel_aria2_idle_timer_locked() + process = facade.aria2_process + if process is not None and process.poll() is None and facade._aria2_ping(): + return + + if process is not None and process.poll() is None: + try: + process.terminate() + except Exception: + pass + facade.aria2_process = None + facade.aria2_process_started_by_resolver = False + + executable = facade._resolve_aria2c_executable(active_settings) + port = facade._find_free_port() + facade.aria2_rpc_secret = facade.secrets.token_hex(16) + facade.aria2_rpc_url = f"http://127.0.0.1:{port}/jsonrpc" + + command = [ + executable, + "--enable-rpc=true", + "--rpc-listen-all=false", + f"--rpc-listen-port={port}", + f"--rpc-secret={facade.aria2_rpc_secret}", + "--check-certificate=true", + "--allow-overwrite=true", + "--auto-file-renaming=false", + "--file-allocation=none", + "--max-concurrent-downloads=5", + "--continue=true", + "--daemon=false", + "--quiet=true", + f"--stop-with-process={facade.os.getpid()}", + ] + ca_cert = facade._try_certifi_ca_path() + if ca_cert: + command.insert(5, f"--ca-certificate={ca_cert}") + + creationflags = 0 + if facade.os.name == "nt": + creationflags = getattr(facade.subprocess, "CREATE_NO_WINDOW", 0) + + log.info(f"Starting aria2 RPC daemon from {executable}") + facade.aria2_process = facade.subprocess.Popen( + command, + stdout=facade.subprocess.DEVNULL, + stderr=facade.subprocess.PIPE, + creationflags=creationflags, + ) + facade.aria2_process_started_by_resolver = True + + start_time = facade.time.time() + last_error = "" + while facade.time.time() - start_time < 10: + process = facade.aria2_process + if process.poll() is not None: + stderr = "" + try: + stderr = ( + process.stderr.read() if process.stderr else b"" + ).decode("utf-8", errors="replace") + except Exception: + stderr = "" + raise Aria2Error( + "aria2 RPC process exited early with code " + f"{process.returncode}: {stderr.strip()}" + ) + try: + if facade._aria2_ping(): + return + except Exception as exc: + last_error = str(exc) + facade.time.sleep(0.2) + + raise Aria2Error( + "Timed out waiting for aria2 RPC to become ready" + f"{': ' + last_error if last_error else ''}" + ) + + +def aria2_tell_status(gid: str) -> Dict[str, Any]: + """Read an aria2 transfer status with retries for transient resets.""" + facade = _downloader_module() + keys = [ + "gid", + "status", + "totalLength", + "completedLength", + "downloadSpeed", + "errorMessage", + "files", + ] + for attempt in range(facade.ARIA2_STATUS_RPC_RETRIES): + try: + result = facade._aria2_rpc("aria2.tellStatus", [gid, keys]) + return result if isinstance(result, dict) else {} + except (facade.requests.exceptions.ConnectionError, facade.requests.exceptions.Timeout): + if attempt + 1 >= facade.ARIA2_STATUS_RPC_RETRIES: + raise + log.debug( + f"Retrying aria2 status RPC for {gid} " + f"({attempt + 2}/{facade.ARIA2_STATUS_RPC_RETRIES})" + ) + facade.time.sleep(facade.ARIA2_STATUS_RPC_RETRY_DELAY * (attempt + 1)) + + return {} diff --git a/core/downloader.py b/core/downloader.py index e7f7cab..a41c93f 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -6,7 +6,7 @@ Handles downloading models from various sources with progress tracking. import hashlib import os -import secrets +import secrets # noqa: F401 import shutil # noqa: F401 import subprocess import threading @@ -24,10 +24,25 @@ log = create_module_logger(__name__) from .download.aria2_backend import ( Aria2Error, + get_aria2_status, # noqa: F401 + start_aria2_daemon, # noqa: F401 + stop_aria2_daemon, # noqa: F401 ) from .download.aria2_backend import ( aria2_action_error_is_ok as _aria2_action_error_is_ok, ) +from .download.aria2_backend import ( + aria2_has_active_transfers_locked as _aria2_has_active_transfers_locked, # noqa: F401 +) +from .download.aria2_backend import ( + aria2_idle_stop_worker as _aria2_idle_stop_worker, # noqa: F401 +) +from .download.aria2_backend import aria2_ping as _aria2_ping # noqa: F401 +from .download.aria2_backend import aria2_rpc as _aria2_rpc +from .download.aria2_backend import aria2_tell_status as _aria2_tell_status +from .download.aria2_backend import ( + cancel_aria2_idle_timer_locked as _cancel_aria2_idle_timer_locked, # noqa: F401 +) from .download.aria2_backend import ( delete_partial_download_files as _delete_partial_download_files, ) @@ -37,20 +52,25 @@ from .download.aria2_backend import ( from .download.aria2_backend import ( delete_xet_partial_file as _delete_xet_partial_file, ) +from .download.aria2_backend import ensure_aria2_daemon as _ensure_aria2_daemon from .download.aria2_backend import ( - find_free_port as _find_free_port, + find_free_port as _find_free_port, # noqa: F401 ) from .download.aria2_backend import ( parse_aria2_int as _parse_aria2_int, ) +from .download.aria2_backend import read_aria2_version as _read_aria2_version # noqa: F401 from .download.aria2_backend import ( resolve_aria2_completed_path as _resolve_aria2_completed_path, ) from .download.aria2_backend import ( - resolve_aria2c_executable as _resolve_aria2c_executable, + resolve_aria2c_executable as _resolve_aria2c_executable, # noqa: F401 ) from .download.aria2_backend import ( - try_certifi_ca_path as _try_certifi_ca_path, + schedule_aria2_idle_stop as _schedule_aria2_idle_stop, +) +from .download.aria2_backend import ( + try_certifi_ca_path as _try_certifi_ca_path, # noqa: F401 ) from .download.previews import ( MODEL_PREVIEW_EXTENSIONS, # noqa: F401 @@ -341,346 +361,6 @@ def _download_backend_from_settings(settings: Optional[Dict[str, Any]] = None) - return normalize_download_backend(active_settings.get("download_backend")) -def _read_aria2_version(executable: str) -> str: - if not executable: - return "" - try: - kwargs: Dict[str, Any] = {} - if os.name == "nt": - kwargs["creationflags"] = getattr(subprocess, "CREATE_NO_WINDOW", 0) - result = subprocess.run( - [executable, "--version"], - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - timeout=5, - check=False, - **kwargs, - ) - except Exception: - return "" - - first_line = "" - for line in str(result.stdout or "").splitlines(): - text = line.strip() - if text: - first_line = text - break - if not first_line: - return "" - - for token in first_line.replace(",", " ").split(): - if token and token[0].isdigit(): - return token - return first_line - - -def get_aria2_status(settings: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: - global aria2_process, aria2_process_started_by_resolver - - active_settings = settings if isinstance(settings, dict) else load_settings() - configured_path = str(active_settings.get("aria2c_path") or "").strip() - try: - resolved_path = _resolve_aria2c_executable(active_settings) - available = True - version = _read_aria2_version(resolved_path) - error = "" - except Exception as exc: - resolved_path = "" - available = False - version = "" - error = str(exc) - - with aria2_lock: - running = aria2_process is not None and aria2_process.poll() is None - if not running and aria2_process is not None: - aria2_process = None - aria2_process_started_by_resolver = False - managed = bool(running and aria2_process_started_by_resolver) - active_transfers = len(aria2_transfers) - - return { - "backend": _download_backend_from_settings(active_settings), - "configured_path": configured_path, - "resolved_path": resolved_path, - "available": available, - "version": version, - "running": running, - "managed": managed, - "can_stop": bool(managed and active_transfers == 0), - "active_transfers": active_transfers, - "auto_stop_enabled": bool(active_settings.get("aria2_auto_stop_daemon", True)), - "idle_stop_seconds": ARIA2_IDLE_STOP_SECONDS, - "error": error, - } - - -def _aria2_rpc(method: str, params: Optional[List[Any]] = None) -> Any: - if not aria2_rpc_url: - raise Aria2Error("aria2 RPC endpoint is not initialized") - - payload = { - "jsonrpc": "2.0", - "id": secrets.token_hex(8), - "method": method, - "params": [f"token:{aria2_rpc_secret}", *(params or [])], - } - # aria2's Windows RPC server can reset one connection when status polling - # overlaps a pause/resume command. Keep local RPC requests serialized. - with aria2_rpc_lock: - response = requests.post(aria2_rpc_url, json=payload, timeout=ARIA2_RPC_TIMEOUT) - text = response.text - try: - body = response.json() - except ValueError as exc: - raise Aria2Error( - f"aria2 RPC returned non-JSON response ({response.status_code}): {text[:300]}" - ) from exc - - if "error" in body: - error = body.get("error") or {} - message = error.get("message") if isinstance(error, dict) else str(error) - raise Aria2Error(message or f"aria2 RPC {method} failed") - - if response.status_code != 200: - raise Aria2Error( - f"aria2 RPC {method} returned HTTP {response.status_code}: {text[:300]}" - ) - - return body.get("result") - - -def _aria2_ping() -> bool: - try: - result = _aria2_rpc("aria2.getVersion", []) - return isinstance(result, dict) - except Exception: - return False - - -def _cancel_aria2_idle_timer_locked() -> None: - global aria2_idle_timer - if aria2_idle_timer is not None: - aria2_idle_timer.cancel() - aria2_idle_timer = None - - -def _aria2_has_active_transfers_locked() -> bool: - return bool(aria2_transfers) - - -def stop_aria2_daemon(reason: str = "manual") -> Dict[str, Any]: - """Stop the aria2 RPC process started by Model Resolver.""" - global aria2_process, aria2_rpc_url, aria2_rpc_secret, aria2_process_started_by_resolver - - with aria2_lock: - _cancel_aria2_idle_timer_locked() - running = aria2_process is not None and aria2_process.poll() is None - if not running: - aria2_process = None - aria2_rpc_url = "" - aria2_rpc_secret = "" - aria2_process_started_by_resolver = False - return {"success": True, "stopped": False, "message": "aria2 daemon is not running"} - - if not aria2_process_started_by_resolver: - return { - "success": False, - "stopped": False, - "error": "This aria2 daemon was not started by Model Resolver.", - } - - if _aria2_has_active_transfers_locked(): - return { - "success": False, - "stopped": False, - "error": "aria2 daemon has active downloads.", - } - - process = aria2_process - try: - process.terminate() - process.wait(timeout=5) - except Exception: - try: - process.kill() - process.wait(timeout=2) - except Exception: - pass - - aria2_process = None - aria2_rpc_url = "" - aria2_rpc_secret = "" - aria2_process_started_by_resolver = False - - log.info(f"aria2 RPC daemon stopped ({reason})") - return {"success": True, "stopped": True, "message": "aria2 daemon stopped"} - - -def start_aria2_daemon(settings: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: - """Start the aria2 RPC process without creating a download.""" - active_settings = settings if isinstance(settings, dict) else load_settings() - try: - _ensure_aria2_daemon(active_settings) - status = get_aria2_status(active_settings) - return { - **status, - "success": True, - "started": bool(status.get("running")), - "message": "aria2 daemon is running", - } - except Exception as exc: - try: - status = get_aria2_status(active_settings) - except Exception: - status = {} - return { - **status, - "success": False, - "started": False, - "error": str(exc), - } - - -def _aria2_idle_stop_worker() -> None: - global aria2_idle_timer - - with aria2_lock: - aria2_idle_timer = None - - settings = load_settings() - if not settings.get("aria2_auto_stop_daemon", True): - return - with aria2_lock: - running = aria2_process is not None and aria2_process.poll() is None - if not running or not aria2_process_started_by_resolver or _aria2_has_active_transfers_locked(): - return - stop_aria2_daemon(reason="idle") - - -def _schedule_aria2_idle_stop() -> None: - global aria2_idle_timer - settings = load_settings() - if not settings.get("aria2_auto_stop_daemon", True): - return - with aria2_lock: - _cancel_aria2_idle_timer_locked() - running = aria2_process is not None and aria2_process.poll() is None - if not running or not aria2_process_started_by_resolver or _aria2_has_active_transfers_locked(): - return - aria2_idle_timer = threading.Timer(ARIA2_IDLE_STOP_SECONDS, _aria2_idle_stop_worker) - aria2_idle_timer.daemon = True - aria2_idle_timer.start() - - -def _ensure_aria2_daemon(settings: Optional[Dict[str, Any]] = None) -> None: - global aria2_process, aria2_rpc_url, aria2_rpc_secret, aria2_process_started_by_resolver - - active_settings = settings if isinstance(settings, dict) else load_settings() - with aria2_lock: - _cancel_aria2_idle_timer_locked() - if aria2_process is not None and aria2_process.poll() is None and _aria2_ping(): - return - - if aria2_process is not None and aria2_process.poll() is None: - try: - aria2_process.terminate() - except Exception: - pass - aria2_process = None - aria2_process_started_by_resolver = False - - executable = _resolve_aria2c_executable(active_settings) - port = _find_free_port() - aria2_rpc_secret = secrets.token_hex(16) - aria2_rpc_url = f"http://127.0.0.1:{port}/jsonrpc" - - command = [ - executable, - "--enable-rpc=true", - "--rpc-listen-all=false", - f"--rpc-listen-port={port}", - f"--rpc-secret={aria2_rpc_secret}", - "--check-certificate=true", - "--allow-overwrite=true", - "--auto-file-renaming=false", - "--file-allocation=none", - "--max-concurrent-downloads=5", - "--continue=true", - "--daemon=false", - "--quiet=true", - f"--stop-with-process={os.getpid()}", - ] - ca_cert = _try_certifi_ca_path() - if ca_cert: - command.insert(5, f"--ca-certificate={ca_cert}") - - creationflags = 0 - if os.name == "nt": - creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0) - - log.info(f"Starting aria2 RPC daemon from {executable}") - aria2_process = subprocess.Popen( - command, - stdout=subprocess.DEVNULL, - stderr=subprocess.PIPE, - creationflags=creationflags, - ) - aria2_process_started_by_resolver = True - - start_time = time.time() - last_error = "" - while time.time() - start_time < 10: - if aria2_process.poll() is not None: - stderr = "" - try: - stderr = (aria2_process.stderr.read() if aria2_process.stderr else b"").decode( - "utf-8", - errors="replace", - ) - except Exception: - stderr = "" - raise Aria2Error( - f"aria2 RPC process exited early with code {aria2_process.returncode}: {stderr.strip()}" - ) - try: - if _aria2_ping(): - return - except Exception as exc: - last_error = str(exc) - time.sleep(0.2) - - raise Aria2Error( - f"Timed out waiting for aria2 RPC to become ready{': ' + last_error if last_error else ''}" - ) - - -def _aria2_tell_status(gid: str) -> Dict[str, Any]: - keys = [ - "gid", - "status", - "totalLength", - "completedLength", - "downloadSpeed", - "errorMessage", - "files", - ] - for attempt in range(ARIA2_STATUS_RPC_RETRIES): - try: - result = _aria2_rpc("aria2.tellStatus", [gid, keys]) - return result if isinstance(result, dict) else {} - except (requests.exceptions.ConnectionError, requests.exceptions.Timeout): - if attempt + 1 >= ARIA2_STATUS_RPC_RETRIES: - raise - log.debug( - f"Retrying aria2 status RPC for {gid} " - f"({attempt + 2}/{ARIA2_STATUS_RPC_RETRIES})" - ) - time.sleep(ARIA2_STATUS_RPC_RETRY_DELAY * (attempt + 1)) - - return {} - - class _HuggingFaceXetDownloadCancelled(Exception): """Raised by the Xet progress adapter when a download is cancelled."""