Dariusz L
f7dfe3d8e5
- Recover completed .part files on HTTP 416 without deleting and redownloading - Add stream stall detection and automatic reconnect for interrupted socket reads - Implement persistent download state store and pause/resume lifecycle controls - Add RateLimitCoordinator and ConnectivityGuard circuit breaker for requests - Support proxy_url configuration and validate aria2 proxy settings - Add unit tests for 416 recovery, stall reconnect, state store, and network guard
441 lines
16 KiB
Python
441 lines
16 KiB
Python
"""Public download operations composed from the internal download modules."""
|
|
|
|
import hashlib
|
|
import os
|
|
import secrets
|
|
import shutil
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from collections import deque
|
|
from types import SimpleNamespace
|
|
from typing import Any, Dict, Optional
|
|
from urllib.parse import urlparse
|
|
|
|
import requests
|
|
|
|
from ..aria2_installer import ARIA2_INSTALL_ROOT
|
|
from ..log_system import create_module_logger
|
|
from ..network_guard import (
|
|
get_connectivity_guard,
|
|
get_rate_limit_coordinator,
|
|
parse_retry_after,
|
|
)
|
|
from ..network_utils import (
|
|
host_matches_domain,
|
|
request_public_url,
|
|
validate_public_http_url,
|
|
)
|
|
from ..path_utils import (
|
|
calculate_file_sha256,
|
|
get_comfy_root_path,
|
|
get_filename_from_path,
|
|
get_model_resolver_sidecar_path,
|
|
get_path_identity,
|
|
is_path_within,
|
|
read_json_safe,
|
|
write_json_atomic,
|
|
)
|
|
from ..resolver import invalidate_model_caches
|
|
from ..settings import (
|
|
load_settings,
|
|
normalize_download_backend,
|
|
normalize_relative_subfolder,
|
|
)
|
|
from ..type_utils import (
|
|
extract_response_file_size,
|
|
get_category_folder_keys,
|
|
normalize_download_category,
|
|
)
|
|
from ..type_utils import (
|
|
format_size_bytes as format_bytes,
|
|
)
|
|
from . import (
|
|
aria2_backend,
|
|
config,
|
|
directories,
|
|
huggingface_xet,
|
|
metadata,
|
|
orchestrator,
|
|
previews,
|
|
state,
|
|
state_store,
|
|
validation,
|
|
)
|
|
|
|
try:
|
|
import folder_paths
|
|
except ImportError:
|
|
folder_paths = None
|
|
|
|
|
|
log = create_module_logger("core.downloader")
|
|
|
|
SPEED_HISTORY_SIZE = 5
|
|
CHUNK_SIZE = 4 * 1024 * 1024
|
|
CLI_LOG_INTERVAL = 5
|
|
ARIA2_RPC_TIMEOUT = (2, 5)
|
|
ARIA2_STATUS_RPC_RETRIES = 4
|
|
ARIA2_STATUS_RPC_RETRY_DELAY = 0.15
|
|
ARIA2_IDLE_STOP_SECONDS = 5 * 60
|
|
MAX_ARIA2_URL_RECOVERY_ATTEMPTS = 2
|
|
PYTHON_DOWNLOAD_MAX_RETRIES = 2
|
|
PYTHON_DOWNLOAD_RETRY_DELAY = 0.5
|
|
PYTHON_DOWNLOAD_CONNECT_TIMEOUT = 20
|
|
PYTHON_DOWNLOAD_READ_TIMEOUT = 60
|
|
PYTHON_DOWNLOAD_STALL_TIMEOUT = 60
|
|
# Keep validation aligned with the install location used by the aria2 installer.
|
|
# ARIA2_INSTALL_ROOT resolves the custom node package first, so paths reached
|
|
# through a shared ComfyUI/custom_nodes symlink remain valid while external
|
|
# executables are still rejected.
|
|
MANAGED_ARIA2_ROOT = ARIA2_INSTALL_ROOT
|
|
HF_XET_ARIA2_AUTH_HOSTS = {
|
|
"cas-bridge.xethub.hf.co",
|
|
"cas-bridge-direct.xethub.hf.co",
|
|
"cas-bridge-direct.xethub-eu.hf.co",
|
|
}
|
|
|
|
|
|
context = SimpleNamespace()
|
|
context.log = log
|
|
context.os = os
|
|
context.secrets = secrets
|
|
context.shutil = shutil
|
|
context.subprocess = subprocess
|
|
context.threading = threading
|
|
context.time = time
|
|
context.hashlib = hashlib
|
|
context.deque = deque
|
|
context.requests = requests
|
|
context.urlparse = urlparse
|
|
context.folder_paths = folder_paths
|
|
|
|
context.SPEED_HISTORY_SIZE = SPEED_HISTORY_SIZE
|
|
context.CHUNK_SIZE = CHUNK_SIZE
|
|
context.CLI_LOG_INTERVAL = CLI_LOG_INTERVAL
|
|
context.ARIA2_RPC_TIMEOUT = ARIA2_RPC_TIMEOUT
|
|
context.ARIA2_STATUS_RPC_RETRIES = ARIA2_STATUS_RPC_RETRIES
|
|
context.ARIA2_STATUS_RPC_RETRY_DELAY = ARIA2_STATUS_RPC_RETRY_DELAY
|
|
context.ARIA2_IDLE_STOP_SECONDS = ARIA2_IDLE_STOP_SECONDS
|
|
context.MAX_ARIA2_URL_RECOVERY_ATTEMPTS = MAX_ARIA2_URL_RECOVERY_ATTEMPTS
|
|
context.PYTHON_DOWNLOAD_MAX_RETRIES = PYTHON_DOWNLOAD_MAX_RETRIES
|
|
context.PYTHON_DOWNLOAD_RETRY_DELAY = PYTHON_DOWNLOAD_RETRY_DELAY
|
|
context.PYTHON_DOWNLOAD_CONNECT_TIMEOUT = PYTHON_DOWNLOAD_CONNECT_TIMEOUT
|
|
context.PYTHON_DOWNLOAD_READ_TIMEOUT = PYTHON_DOWNLOAD_READ_TIMEOUT
|
|
context.PYTHON_DOWNLOAD_STALL_TIMEOUT = PYTHON_DOWNLOAD_STALL_TIMEOUT
|
|
context.MANAGED_ARIA2_ROOT = MANAGED_ARIA2_ROOT
|
|
context.HF_XET_ARIA2_AUTH_HOSTS = HF_XET_ARIA2_AUTH_HOSTS
|
|
|
|
context.aria2_process = None
|
|
context.aria2_recent_stderr = deque(maxlen=50)
|
|
context.aria2_rpc_url = ""
|
|
context.aria2_rpc_secret = ""
|
|
context.aria2_rpc_lock = threading.Lock()
|
|
context.aria2_idle_timer = None
|
|
context.aria2_process_started_by_resolver = False
|
|
|
|
context.download_progress = state.download_progress
|
|
context.download_lock = state.download_lock
|
|
context.cancelled_downloads = state.cancelled_downloads
|
|
context.aria2_lock = state.aria2_lock
|
|
context.aria2_transfers = state.aria2_transfers
|
|
context.aria2_action_locks = state.aria2_action_locks
|
|
context.aria2_desired_states = state.aria2_desired_states
|
|
context.xet_transfers = state.xet_transfers
|
|
context.xet_transfers_lock = state.xet_transfers_lock
|
|
context.python_download_controls = state.python_download_controls
|
|
context.python_download_controls_lock = state.python_download_controls_lock
|
|
context.register_python_download_control = state.register_python_download_control
|
|
context.unregister_python_download_control = state.unregister_python_download_control
|
|
context.pause_python_download = state.pause_python_download
|
|
context.resume_python_download = state.resume_python_download
|
|
context.is_python_download_active = state.is_python_download_active
|
|
|
|
context.load_settings = load_settings
|
|
context.normalize_download_backend = normalize_download_backend
|
|
context.normalize_relative_subfolder = normalize_relative_subfolder
|
|
context.generate_download_id = config.generate_download_id
|
|
|
|
context.get_connectivity_guard = get_connectivity_guard
|
|
context.get_rate_limit_coordinator = get_rate_limit_coordinator
|
|
context.parse_retry_after = parse_retry_after
|
|
context.host_matches_domain = host_matches_domain
|
|
context.request_public_url = request_public_url
|
|
context.validate_public_http_url = validate_public_http_url
|
|
context.build_download_headers = validation.build_download_headers
|
|
context.DOWNLOAD_USER_AGENT = validation.DOWNLOAD_USER_AGENT
|
|
context._get_header_value = validation._get_header_value
|
|
context._sanitize_download_error = validation._sanitize_download_error
|
|
context._strip_sensitive_url_params = validation._strip_sensitive_url_params
|
|
context.is_allowed_model_download_filename = (
|
|
validation.is_allowed_model_download_filename
|
|
)
|
|
context.sanitize_download_filename = validation.sanitize_download_filename
|
|
context.extract_response_file_size = extract_response_file_size
|
|
context.format_bytes = format_bytes
|
|
context.calculate_file_sha256 = calculate_file_sha256
|
|
context.get_comfy_root_path = get_comfy_root_path
|
|
context.get_filename_from_path = get_filename_from_path
|
|
context.get_model_resolver_sidecar_path = get_model_resolver_sidecar_path
|
|
context.get_path_identity = get_path_identity
|
|
context.is_path_within = is_path_within
|
|
context.write_json_atomic = write_json_atomic
|
|
context.read_json_safe = read_json_safe
|
|
context.get_category_folder_keys = get_category_folder_keys
|
|
context.invalidate_model_caches = invalidate_model_caches
|
|
|
|
context._extract_expected_sha256 = metadata._extract_expected_sha256
|
|
context.normalize_metadata_file_path = metadata.normalize_metadata_file_path
|
|
context._json_safe_metadata = metadata._json_safe_metadata
|
|
context._coerce_int_or_value = metadata._coerce_int_or_value
|
|
context._coerce_size = metadata._coerce_size
|
|
context._metadata_source_value = metadata._metadata_source_value
|
|
context._resolve_lora_manager_model_type = metadata._resolve_lora_manager_model_type
|
|
context.build_model_resolver_metadata = metadata.build_model_resolver_metadata
|
|
context.read_completed_metadata_sha256 = metadata.read_completed_metadata_sha256
|
|
|
|
|
|
def _bind(function):
|
|
"""Bind the composed download context to an internal operation."""
|
|
def bound(*args: Any, **kwargs: Any):
|
|
kwargs.setdefault("dependencies", context)
|
|
return function(*args, **kwargs)
|
|
|
|
return bound
|
|
|
|
|
|
context._download_backend_from_settings = _bind(
|
|
config.download_backend_from_settings
|
|
)
|
|
context.get_download_directory = _bind(directories.get_download_directory)
|
|
context.write_model_resolver_metadata = _bind(
|
|
metadata.write_model_resolver_metadata
|
|
)
|
|
context.create_model_preview = _bind(previews.create_model_preview)
|
|
context._download_preview_asset = _bind(previews._download_preview_asset)
|
|
context._download_preview_asset_with_system_trust = _bind(
|
|
previews._download_preview_asset_with_system_trust
|
|
)
|
|
context._download_preview_image = _bind(previews._download_preview_image)
|
|
context._download_preview_image_with_system_trust = _bind(
|
|
previews._download_preview_image_with_system_trust
|
|
)
|
|
context.get_existing_model_preview_path = previews.get_existing_model_preview_path
|
|
|
|
context._download_huggingface_xet = _bind(
|
|
huggingface_xet.download_huggingface_xet
|
|
)
|
|
context._run_huggingface_xet_transfer = _bind(
|
|
huggingface_xet.run_huggingface_xet_transfer
|
|
)
|
|
|
|
context.download_file_with_aria2 = _bind(aria2_backend.download_file_with_aria2)
|
|
context.start_background_sha256_verification = _bind(
|
|
aria2_backend.start_background_sha256_verification
|
|
)
|
|
context._start_background_sha256_verification = context.start_background_sha256_verification
|
|
context.get_aria2_status = _bind(aria2_backend.get_aria2_status)
|
|
context.pause_download = _bind(aria2_backend.pause_download)
|
|
context.resume_download = _bind(aria2_backend.resume_download)
|
|
context.start_aria2_daemon = _bind(aria2_backend.start_aria2_daemon)
|
|
context.stop_aria2_daemon = _bind(aria2_backend.stop_aria2_daemon)
|
|
context._aria2_ping = _bind(aria2_backend.aria2_ping)
|
|
context._aria2_rpc = _bind(aria2_backend.aria2_rpc)
|
|
context._aria2_tell_status = _bind(aria2_backend.aria2_tell_status)
|
|
context._cancel_aria2_idle_timer_locked = _bind(
|
|
aria2_backend.cancel_aria2_idle_timer_locked
|
|
)
|
|
context._delete_partial_download_files = _bind(
|
|
aria2_backend.delete_partial_download_files
|
|
)
|
|
context._delete_python_partial_download_file = _bind(
|
|
aria2_backend.delete_python_partial_download_file
|
|
)
|
|
context._delete_xet_partial_file = _bind(aria2_backend.delete_xet_partial_file)
|
|
context._ensure_aria2_daemon = _bind(aria2_backend.ensure_aria2_daemon)
|
|
context._recover_aria2_missing_control_file = _bind(
|
|
aria2_backend.recover_aria2_missing_control_file
|
|
)
|
|
context._force_remove_aria2_transfer = _bind(
|
|
aria2_backend.force_remove_aria2_transfer
|
|
)
|
|
context._get_aria2_action_lock = _bind(aria2_backend.get_aria2_action_lock)
|
|
context._read_aria2_version = _bind(aria2_backend.read_aria2_version)
|
|
context._resolve_aria2c_executable = _bind(
|
|
aria2_backend.resolve_aria2c_executable
|
|
)
|
|
context._resolve_download_url_for_aria2 = _bind(
|
|
aria2_backend.resolve_download_url_for_aria2
|
|
)
|
|
context._run_aria2_desired_state_worker = _bind(
|
|
aria2_backend.run_aria2_desired_state_worker
|
|
)
|
|
context._aria2_idle_stop_worker = _bind(aria2_backend.aria2_idle_stop_worker)
|
|
context._schedule_aria2_idle_stop = _bind(aria2_backend.schedule_aria2_idle_stop)
|
|
context._set_download_progress_status = _bind(
|
|
aria2_backend.set_download_progress_status
|
|
)
|
|
context._try_certifi_ca_path = _bind(aria2_backend.try_certifi_ca_path)
|
|
context._aria2_has_active_transfers_locked = _bind(
|
|
aria2_backend.aria2_has_active_transfers_locked
|
|
)
|
|
context._queue_aria2_desired_state = _bind(
|
|
aria2_backend.queue_aria2_desired_state
|
|
)
|
|
context._aria2_action_error_is_ok = aria2_backend.aria2_action_error_is_ok
|
|
context._is_aria2_missing_control_file_error = (
|
|
aria2_backend.is_aria2_missing_control_file_error
|
|
)
|
|
context._is_aria2_important_stderr = aria2_backend.is_aria2_important_stderr
|
|
context.is_aria2_important_stderr = aria2_backend.is_aria2_important_stderr
|
|
context._is_aria2_recoverable_url_error = (
|
|
aria2_backend.is_aria2_recoverable_url_error
|
|
)
|
|
context.is_aria2_recoverable_url_error = (
|
|
aria2_backend.is_aria2_recoverable_url_error
|
|
)
|
|
context._drain_aria2_stderr = _bind(aria2_backend.drain_aria2_stderr)
|
|
context.drain_aria2_stderr = context._drain_aria2_stderr
|
|
context.Aria2Error = aria2_backend.Aria2Error
|
|
context._find_free_port = aria2_backend.find_free_port
|
|
context._parse_aria2_int = aria2_backend.parse_aria2_int
|
|
context._resolve_aria2_completed_path = aria2_backend.resolve_aria2_completed_path
|
|
context.resolve_proxy_url = _bind(aria2_backend.resolve_proxy_url)
|
|
context._resolve_proxy_url = context.resolve_proxy_url
|
|
|
|
context.save_download_spec = state_store.save_download_spec
|
|
context.update_download_progress_state = state_store.update_download_progress_state
|
|
context.remove_persisted_download = state_store.remove_persisted_download
|
|
context.clear_completed_persisted_transfers = state_store.clear_completed_persisted_transfers
|
|
context.get_persisted_spec = state_store.get_persisted_spec
|
|
context.restore_persisted_downloads = _bind(state_store.restore_persisted_downloads)
|
|
|
|
|
|
def _state_dependencies() -> state.DownloadStateDependencies:
|
|
return state.DownloadStateDependencies(
|
|
download_progress=context.download_progress,
|
|
download_lock=context.download_lock,
|
|
cancelled_downloads=context.cancelled_downloads,
|
|
aria2_lock=context.aria2_lock,
|
|
aria2_transfers=context.aria2_transfers,
|
|
aria2_desired_states=context.aria2_desired_states,
|
|
xet_transfers=context.xet_transfers,
|
|
xet_transfers_lock=context.xet_transfers_lock,
|
|
set_download_progress_status=context._set_download_progress_status,
|
|
force_remove_aria2_transfer=context._force_remove_aria2_transfer,
|
|
thread_factory=threading.Thread,
|
|
log=log,
|
|
delete_partial_download_files=context._delete_partial_download_files,
|
|
remove_persisted_download=context.remove_persisted_download,
|
|
clear_completed_persisted_transfers=context.clear_completed_persisted_transfers,
|
|
python_download_controls=context.python_download_controls,
|
|
python_download_controls_lock=context.python_download_controls_lock,
|
|
)
|
|
|
|
|
|
def get_progress(download_id: str) -> Optional[Dict[str, Any]]:
|
|
return state.get_progress(download_id, _state_dependencies())
|
|
|
|
|
|
def get_all_progress() -> Dict[str, Dict[str, Any]]:
|
|
return state.get_all_progress(_state_dependencies())
|
|
|
|
|
|
def cancel_download(download_id: str) -> bool:
|
|
return state.cancel_download(download_id, _state_dependencies())
|
|
|
|
|
|
def clear_completed_downloads():
|
|
return state.clear_completed_downloads(_state_dependencies())
|
|
|
|
|
|
context.get_progress = get_progress
|
|
context.get_all_progress = get_all_progress
|
|
context.cancel_download = cancel_download
|
|
context.clear_completed_downloads = clear_completed_downloads
|
|
|
|
|
|
class _BoundHuggingFaceXetProgressAdapter(
|
|
huggingface_xet.HuggingFaceXetProgressAdapter
|
|
):
|
|
"""Bind the composed context for direct progress adapter consumers."""
|
|
|
|
def __init__(self, *args: Any, **kwargs: Any):
|
|
kwargs.setdefault("dependencies", context)
|
|
super().__init__(*args, **kwargs)
|
|
|
|
|
|
context._HuggingFaceXetProgressAdapter = _BoundHuggingFaceXetProgressAdapter
|
|
|
|
|
|
def _download_file(*args: Any, **kwargs: Any) -> Dict[str, Any]:
|
|
kwargs["dependencies"] = context
|
|
return orchestrator.download_file(*args, **kwargs)
|
|
|
|
|
|
def _download_model(*args: Any, **kwargs: Any) -> Dict[str, Any]:
|
|
kwargs["dependencies"] = context
|
|
return orchestrator.download_model(*args, **kwargs)
|
|
|
|
|
|
def download_file(*args: Any, **kwargs: Any) -> Dict[str, Any]:
|
|
return _download_file(*args, **kwargs)
|
|
|
|
|
|
def download_model(*args: Any, **kwargs: Any) -> Dict[str, Any]:
|
|
return _download_model(*args, **kwargs)
|
|
|
|
|
|
def start_background_download(*args: Any, **kwargs: Any) -> str:
|
|
kwargs["dependencies"] = context
|
|
return orchestrator.start_background_download(*args, **kwargs)
|
|
|
|
|
|
context.download_file = _download_file
|
|
context.download_model = _download_model
|
|
|
|
|
|
is_allowed_model_download_filename = context.is_allowed_model_download_filename
|
|
sanitize_download_filename = context.sanitize_download_filename
|
|
get_existing_model_preview_path = context.get_existing_model_preview_path
|
|
|
|
|
|
__all__ = [
|
|
"cancel_download",
|
|
"clear_completed_downloads",
|
|
"download_file",
|
|
"download_model",
|
|
"get_all_progress",
|
|
"get_aria2_status",
|
|
"get_download_directory",
|
|
"get_existing_model_preview_path",
|
|
"get_progress",
|
|
"is_allowed_model_download_filename",
|
|
"normalize_download_category",
|
|
"pause_download",
|
|
"resume_download",
|
|
"sanitize_download_filename",
|
|
"start_aria2_daemon",
|
|
"start_background_download",
|
|
"stop_aria2_daemon",
|
|
"write_model_resolver_metadata",
|
|
]
|
|
|
|
|
|
get_aria2_status = context.get_aria2_status
|
|
get_download_directory = context.get_download_directory
|
|
pause_download = context.pause_download
|
|
resume_download = context.resume_download
|
|
start_aria2_daemon = context.start_aria2_daemon
|
|
stop_aria2_daemon = context.stop_aria2_daemon
|
|
write_model_resolver_metadata = context.write_model_resolver_metadata
|
|
|
|
try:
|
|
state_store.restore_persisted_downloads(
|
|
context.download_progress,
|
|
context.download_lock,
|
|
log_fn=log,
|
|
)
|
|
except Exception as _restore_err:
|
|
log.warning(f"Failed to restore persisted downloads: {_restore_err}")
|