nodes/Azornes-Comfyui-Model-Resolver
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: extract downloader backend helpers

Browse Source
This commit is contained in:
Dariusz L
2026-08-02 16:40:54 +02:00
parent 29a051e79a
commit 6ef551877d
2 changed files with 183 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/download/aria2_backend.py
+152
View File
@@ -0,0 +1,152 @@
"""Low-level aria2 helpers used by the downloader facade."""
import importlib
import socket
from typing import Any, Dict, Optional
from ..log_system import create_module_logger
log = create_module_logger("core.downloader")
class Aria2Error(RuntimeError):
"""Raised when the aria2 backend cannot start or process a request."""
def _downloader_module():
"""Return the facade module so existing dependency patches remain effective."""
return importlib.import_module("core.downloader")
def try_certifi_ca_path() -> str:
"""Return certifi's CA bundle path when it is available."""
facade = _downloader_module()
try:
import certifi # type: ignore
path = certifi.where()
return path if path and facade.os.path.isfile(path) else ""
except Exception:
return ""
def resolve_aria2c_executable(settings: Optional[Dict[str, Any]] = None) -> str:
"""Resolve aria2c while restricting explicit paths to the managed install."""
facade = _downloader_module()
active_settings = (
settings if isinstance(settings, dict) else facade.load_settings()
)
configured = str(active_settings.get("aria2c_path") or "").strip()
candidate = facade.os.path.expandvars(
facade.os.path.expanduser(configured or "aria2c")
)
expected_names = {"aria2c", "aria2c.exe"}
candidate_name = facade.get_filename_from_path(candidate).lower()
has_path_component = bool(
facade.os.path.isabs(candidate)
or facade.os.path.dirname(candidate)
or "/" in candidate
or "\\" in candidate
)
if has_path_component:
if (
candidate_name in expected_names
and facade.os.path.isfile(candidate)
and facade.is_path_within(candidate, facade.MANAGED_ARIA2_ROOT)
):
return facade.os.path.realpath(facade.os.path.abspath(candidate))
raise Aria2Error(
"Custom aria2c paths are restricted to the managed Model Resolver install. "
"Use the built-in aria2 installer or place aria2c on PATH."
)
if candidate_name in expected_names:
resolved = facade.shutil.which(candidate)
if resolved and facade.get_filename_from_path(resolved).lower() in expected_names:
return facade.os.path.realpath(facade.os.path.abspath(resolved))
raise Aria2Error(
"aria2c executable was not found. Use the built-in installer or place aria2c on PATH."
)
def find_free_port() -> int:
"""Return an unused local TCP port for the aria2 RPC endpoint."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind(("127.0.0.1", 0))
sock.listen(1)
return int(sock.getsockname()[1])
def parse_aria2_int(value: Any) -> int:
"""Parse an aria2 numeric field without propagating malformed values."""
try:
return int(value)
except (TypeError, ValueError):
return 0
def resolve_aria2_completed_path(status: Dict[str, Any], default_path: str) -> str:
"""Use aria2's completed file path when it reports one."""
files = status.get("files")
if isinstance(files, list) and files:
first = files[0]
if isinstance(first, dict):
candidate = first.get("path")
if isinstance(candidate, str) and candidate:
return candidate
return default_path
def delete_partial_download_files(dest_path: str) -> None:
"""Delete an incomplete model and its aria2 control sidecar."""
facade = _downloader_module()
for path in (dest_path, f"{dest_path}.aria2"):
try:
if path and facade.os.path.exists(path):
facade.os.remove(path)
except Exception as exc:
log.warning(f"Could not delete incomplete download file {path}: {exc}")
def delete_python_partial_download_file(partial_path: str) -> None:
"""Remove a partial Python download without touching the final model path."""
facade = _downloader_module()
try:
if partial_path and facade.os.path.exists(partial_path):
facade.os.remove(partial_path)
except Exception as exc:
log.warning(
f"Could not delete incomplete Python download file {partial_path}: {exc}"
)
def delete_xet_partial_file(partial_path: str, attempts: int = 5) -> bool:
"""Delete a stopped Xet partial file, retrying while Windows releases it."""
facade = _downloader_module()
attempts = max(1, int(attempts or 1))
last_error: Optional[Exception] = None
for attempt in range(attempts):
try:
if not facade.os.path.exists(partial_path):
return True
facade.os.remove(partial_path)
return True
except Exception as exc:
last_error = exc
if attempt + 1 < attempts:
facade.time.sleep(0.25)
log.warning(f"Could not delete incomplete Xet file {partial_path}: {last_error}")
return False
def aria2_action_error_is_ok(status: str, message: str) -> bool:
"""Recognize idempotent aria2 pause/resume errors."""
lowered = str(message or "").lower()
if status == "paused":
return "already paused" in lowered or "is paused" in lowered
if status == "downloading":
return "not paused" in lowered or "has not been paused" in lowered
return False
core/downloader.py
+31 -124
View File
@@ -7,8 +7,7 @@ Handles downloading models from various sources with progress tracking.
import hashlib
import os
import secrets
import shutil
import socket
import shutil # noqa: F401
import subprocess
import threading
import time
@@ -23,6 +22,36 @@ from .log_system import create_module_logger
log = create_module_logger(__name__)
from .download.aria2_backend import (
Aria2Error,
)
from .download.aria2_backend import (
aria2_action_error_is_ok as _aria2_action_error_is_ok,
)
from .download.aria2_backend import (
delete_partial_download_files as _delete_partial_download_files,
)
from .download.aria2_backend import (
delete_python_partial_download_file as _delete_python_partial_download_file,
)
from .download.aria2_backend import (
delete_xet_partial_file as _delete_xet_partial_file,
)
from .download.aria2_backend import (
find_free_port as _find_free_port,
)
from .download.aria2_backend import (
parse_aria2_int as _parse_aria2_int,
)
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,
)
from .download.aria2_backend import (
try_certifi_ca_path as _try_certifi_ca_path,
)
from .download.previews import (
MODEL_PREVIEW_EXTENSIONS, # noqa: F401
MODEL_PREVIEW_MAX_DOWNLOAD_BYTES, # noqa: F401
@@ -307,60 +336,11 @@ def generate_download_id() -> str:
return str(uuid.uuid4())[:8]
class Aria2Error(RuntimeError):
"""Raised when the aria2 backend cannot start or process a request."""
def _download_backend_from_settings(settings: Optional[Dict[str, Any]] = None) -> str:
active_settings = settings if isinstance(settings, dict) else load_settings()
return normalize_download_backend(active_settings.get("download_backend"))
def _try_certifi_ca_path() -> str:
try:
import certifi # type: ignore
path = certifi.where()
return path if path and os.path.isfile(path) else ""
except Exception:
return ""
def _resolve_aria2c_executable(settings: Optional[Dict[str, Any]] = None) -> str:
active_settings = settings if isinstance(settings, dict) else load_settings()
configured = str(active_settings.get("aria2c_path") or "").strip()
candidate = os.path.expandvars(os.path.expanduser(configured or "aria2c"))
expected_names = {"aria2c", "aria2c.exe"}
candidate_name = get_filename_from_path(candidate).lower()
has_path_component = bool(
os.path.isabs(candidate)
or os.path.dirname(candidate)
or "/" in candidate
or "\\" in candidate
)
if has_path_component:
if (
candidate_name in expected_names
and os.path.isfile(candidate)
and is_path_within(candidate, MANAGED_ARIA2_ROOT)
):
return os.path.realpath(os.path.abspath(candidate))
raise Aria2Error(
"Custom aria2c paths are restricted to the managed Model Resolver install. "
"Use the built-in aria2 installer or place aria2c on PATH."
)
if candidate_name in expected_names:
resolved = shutil.which(candidate)
if resolved and get_filename_from_path(resolved).lower() in expected_names:
return os.path.realpath(os.path.abspath(resolved))
raise Aria2Error(
"aria2c executable was not found. Use the built-in installer or place aria2c on PATH."
)
def _read_aria2_version(executable: str) -> str:
if not executable:
return ""
@@ -435,13 +415,6 @@ def get_aria2_status(settings: Optional[Dict[str, Any]] = None) -> Dict[str, Any
}
def _find_free_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.bind(("127.0.0.1", 0))
sock.listen(1)
return int(sock.getsockname()[1])
def _aria2_rpc(method: str, params: Optional[List[Any]] = None) -> Any:
if not aria2_rpc_url:
raise Aria2Error("aria2 RPC endpoint is not initialized")
@@ -708,63 +681,6 @@ def _aria2_tell_status(gid: str) -> Dict[str, Any]:
return {}
def _parse_aria2_int(value: Any) -> int:
try:
return int(value)
except (TypeError, ValueError):
return 0
def _resolve_aria2_completed_path(status: Dict[str, Any], default_path: str) -> str:
files = status.get("files")
if isinstance(files, list) and files:
first = files[0]
if isinstance(first, dict):
candidate = first.get("path")
if isinstance(candidate, str) and candidate:
return candidate
return default_path
def _delete_partial_download_files(dest_path: str) -> None:
for path in (dest_path, f"{dest_path}.aria2"):
try:
if path and os.path.exists(path):
os.remove(path)
except Exception as exc:
log.warning(f"Could not delete incomplete download file {path}: {exc}")
def _delete_python_partial_download_file(partial_path: str) -> None:
"""Remove a partial Python download without touching the final model path."""
try:
if partial_path and os.path.exists(partial_path):
os.remove(partial_path)
except Exception as exc:
log.warning(
f"Could not delete incomplete Python download file {partial_path}: {exc}"
)
def _delete_xet_partial_file(partial_path: str, attempts: int = 5) -> bool:
"""Delete a stopped Xet partial file, retrying while Windows releases it."""
attempts = max(1, int(attempts or 1))
last_error: Optional[Exception] = None
for attempt in range(attempts):
try:
if not os.path.exists(partial_path):
return True
os.remove(partial_path)
return True
except Exception as exc:
last_error = exc
if attempt + 1 < attempts:
time.sleep(0.25)
log.warning(f"Could not delete incomplete Xet file {partial_path}: {last_error}")
return False
class _HuggingFaceXetDownloadCancelled(Exception):
"""Raised by the Xet progress adapter when a download is cancelled."""
@@ -2000,15 +1916,6 @@ def _set_download_progress_status(download_id: str, status: str, **updates: Any)
download_progress[download_id].update(updates)
def _aria2_action_error_is_ok(status: str, message: str) -> bool:
lowered = str(message or "").lower()
if status == "paused":
return "already paused" in lowered or "is paused" in lowered
if status == "downloading":
return "not paused" in lowered or "has not been paused" in lowered
return False
def _run_aria2_desired_state_worker(download_id: str) -> None:
while True:
with aria2_lock: