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

refactor: extract aria2 transfer backend

Browse Source
This commit is contained in:
Dariusz L
2026-08-02 16:55:52 +02:00
parent 9360783628
commit c35f5bb5b2
2 changed files with 257 additions and 233 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/download/aria2_backend.py
+248
View File
@@ -534,3 +534,251 @@ def aria2_tell_status(gid: str) -> Dict[str, Any]:
facade.time.sleep(facade.ARIA2_STATUS_RPC_RETRY_DELAY * (attempt + 1))
return {}
def download_file_with_aria2(
url: str,
dest_path: str,
download_id: str,
headers: Optional[Dict[str, str]] = None,
metadata: Optional[Dict[str, Any]] = None,
category: str = "",
) -> Dict[str, Any]:
"""Download a file with an aria2c JSON-RPC process."""
facade = _downloader_module()
settings = facade.load_settings()
result = {
"success": False,
"download_id": download_id,
"path": dest_path,
"error": None,
"size": 0,
}
start_time = facade.time.time()
filename = facade.get_filename_from_path(dest_path)
expected_sha256 = facade._extract_expected_sha256(metadata)
with facade.download_lock:
facade.download_progress[download_id] = {
"status": "starting",
"progress": 0,
"total_size": 0,
"downloaded": 0,
"filename": filename,
"path": dest_path,
"directory": facade.os.path.dirname(dest_path),
"url": url,
"error": None,
"speed": 0,
"start_time": start_time,
"download_backend": "aria2",
}
try:
facade.os.makedirs(facade.os.path.dirname(dest_path), exist_ok=True)
facade._ensure_aria2_daemon(settings)
aria2_url, request_headers = facade._resolve_download_url_for_aria2(
url,
headers,
)
options: Dict[str, Any] = {
"dir": facade.os.path.dirname(dest_path),
"out": filename,
"continue": "true",
"max-connection-per-server": "4",
"split": "4",
"min-split-size": "1M",
"allow-overwrite": "true",
"auto-file-renaming": "false",
"file-allocation": "none",
"no-want-digest-header": "true",
# Redirects were already resolved and validated above. Keeping them
# disabled prevents sensitive headers from reaching another host.
"max-redirect": "0",
}
if expected_sha256:
# Let aria2 verify the file as part of the transfer instead of
# requiring a second Python pass over a multi-gigabyte model.
options["checksum"] = f"sha-256={expected_sha256}"
options["check-integrity"] = "true"
user_agent = facade._get_header_value(request_headers, "User-Agent")
referer = facade._get_header_value(request_headers, "Referer")
if user_agent:
options["user-agent"] = user_agent
if referer:
options["referer"] = referer
header_values = [
f"{key}: {value}"
for key, value in request_headers.items()
if str(key).lower() not in {"user-agent", "referer"}
]
if header_values:
options["header"] = header_values
gid = facade._aria2_rpc("aria2.addUri", [[aria2_url], options])
if not isinstance(gid, str) or not gid:
raise Aria2Error("aria2 did not return a download gid")
with facade.aria2_lock:
facade.aria2_transfers[download_id] = {
"gid": gid,
"path": dest_path,
}
with facade.download_lock:
facade.download_progress[download_id]["aria2_gid"] = gid
facade.download_progress[download_id]["status"] = "downloading"
log.info(f"Starting aria2 download: {filename}")
last_cli_log = start_time
while True:
if download_id in facade.cancelled_downloads:
try:
facade._aria2_rpc("aria2.forceRemove", [gid])
except Exception:
pass
with facade.download_lock:
if download_id in facade.download_progress:
facade.download_progress[download_id]["status"] = "cancelled"
facade.download_progress[download_id]["speed"] = 0
facade._delete_partial_download_files(dest_path)
facade.cancelled_downloads.discard(download_id)
result["error"] = "Download cancelled"
return result
status = facade._aria2_tell_status(gid)
state = str(status.get("status") or "")
total_size = facade._parse_aria2_int(status.get("totalLength"))
downloaded = facade._parse_aria2_int(status.get("completedLength"))
speed = facade._parse_aria2_int(status.get("downloadSpeed"))
progress = int((downloaded / total_size) * 100) if total_size > 0 else 0
mapped_status = {
"active": "downloading",
"waiting": "downloading",
"paused": "paused",
"complete": "completed",
"error": "error",
"removed": "cancelled",
}.get(state, state or "downloading")
with facade.download_lock:
if download_id in facade.download_progress:
facade.download_progress[download_id].update(
{
"status": mapped_status,
"progress": max(0, min(progress, 100)),
"total_size": total_size,
"downloaded": downloaded,
"speed": (
0
if mapped_status in {"paused", "completed"}
else speed
),
"download_backend": "aria2",
"aria2_gid": gid,
}
)
now = facade.time.time()
if (
now - last_cli_log >= facade.CLI_LOG_INTERVAL
and mapped_status == "downloading"
):
last_cli_log = now
total_str = facade.format_bytes(total_size) if total_size else "?"
log.info(
"aria2 progress: "
f"{facade.format_bytes(downloaded)} / {total_str} "
f"({progress}%) - {facade.format_bytes(speed)}/s"
)
if state == "complete":
completed_path = facade._resolve_aria2_completed_path(
status,
dest_path,
)
size = (
facade.os.path.getsize(completed_path)
if facade.os.path.exists(completed_path)
else downloaded
)
metadata_path = facade.write_model_resolver_metadata(
completed_path,
metadata or {},
category,
url,
create_preview=True,
)
with facade.download_lock:
facade.download_progress[download_id].update(
{
"status": "completed",
"progress": 100,
"downloaded": size,
"total_size": total_size or size,
"speed": 0,
"path": completed_path,
"directory": facade.os.path.dirname(completed_path),
}
)
if metadata_path:
facade.download_progress[download_id][
"metadata_path"
] = metadata_path
result.update(
{
"success": True,
"path": completed_path,
"size": size,
"metadata_path": metadata_path,
}
)
elapsed = facade.time.time() - start_time
avg_speed = size / elapsed if elapsed > 0 else 0
log.info(f"✓ aria2 download complete: {filename}")
log.info(
f"Size: {facade.format_bytes(size)}, Time: {elapsed:.1f}s, "
f"Avg speed: {facade.format_bytes(int(avg_speed))}/s"
)
facade.invalidate_model_files_cache()
facade.invalidate_local_hash_match_cache()
return result
if state == "error":
error_msg = status.get("errorMessage") or "aria2 download failed"
with facade.download_lock:
facade.download_progress[download_id]["status"] = "error"
facade.download_progress[download_id]["error"] = error_msg
result["error"] = error_msg
return result
if state == "removed":
with facade.download_lock:
facade.download_progress[download_id]["status"] = "cancelled"
facade.download_progress[download_id]["speed"] = 0
facade._delete_partial_download_files(dest_path)
facade.cancelled_downloads.discard(download_id)
result["error"] = "Download cancelled"
return result
facade.time.sleep(0.5)
except Exception as exc:
error_msg = facade._sanitize_download_error(exc)
with facade.download_lock:
if download_id in facade.download_progress:
facade.download_progress[download_id]["status"] = "error"
facade.download_progress[download_id]["error"] = error_msg
facade.download_progress[download_id]["speed"] = 0
result["error"] = error_msg
log.error(f"✗ aria2 download failed: {filename}")
log.error(f"Error: {error_msg}")
return result
finally:
with facade.aria2_lock:
facade.aria2_transfers.pop(download_id, None)
facade.aria2_action_locks.pop(download_id, None)
facade.aria2_desired_states.pop(download_id, None)
facade._schedule_aria2_idle_stop()
core/downloader.py
+9 -233
View File
@@ -23,7 +23,8 @@ from .log_system import create_module_logger
log = create_module_logger(__name__)
from .download.aria2_backend import (
Aria2Error,
Aria2Error, # noqa: F401
download_file_with_aria2,
get_aria2_status, # noqa: F401
start_aria2_daemon, # noqa: F401
stop_aria2_daemon, # noqa: F401
@@ -39,12 +40,12 @@ from .download.aria2_backend import (
)
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 aria2_tell_status as _aria2_tell_status # noqa: F401
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,
delete_partial_download_files as _delete_partial_download_files, # noqa: F401
)
from .download.aria2_backend import (
delete_python_partial_download_file as _delete_python_partial_download_file,
@@ -52,22 +53,22 @@ from .download.aria2_backend import (
from .download.aria2_backend import (
delete_xet_partial_file as _delete_xet_partial_file, # noqa: F401
)
from .download.aria2_backend import ensure_aria2_daemon as _ensure_aria2_daemon
from .download.aria2_backend import ensure_aria2_daemon as _ensure_aria2_daemon # noqa: F401
from .download.aria2_backend import (
find_free_port as _find_free_port, # noqa: F401
)
from .download.aria2_backend import (
parse_aria2_int as _parse_aria2_int,
parse_aria2_int as _parse_aria2_int, # noqa: F401
)
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,
resolve_aria2_completed_path as _resolve_aria2_completed_path, # noqa: F401
)
from .download.aria2_backend import (
resolve_aria2c_executable as _resolve_aria2c_executable, # noqa: F401
)
from .download.aria2_backend import (
schedule_aria2_idle_stop as _schedule_aria2_idle_stop,
schedule_aria2_idle_stop as _schedule_aria2_idle_stop, # noqa: F401
)
from .download.aria2_backend import (
try_certifi_ca_path as _try_certifi_ca_path, # noqa: F401
@@ -177,7 +178,7 @@ from .download.metadata import (
)
from .download.validation import (
DOWNLOAD_USER_AGENT, # noqa: F401
_get_header_value,
_get_header_value, # noqa: F401
_is_sensitive_metadata_key, # noqa: F401
_sanitize_download_error,
_strip_sensitive_url_params,
@@ -374,231 +375,6 @@ def _download_backend_from_settings(settings: Optional[Dict[str, Any]] = None) -
return normalize_download_backend(active_settings.get("download_backend"))
def download_file_with_aria2(
url: str,
dest_path: str,
download_id: str,
headers: Optional[Dict[str, str]] = None,
metadata: Optional[Dict[str, Any]] = None,
category: str = "",
) -> Dict[str, Any]:
"""Download a file with an aria2c JSON-RPC process."""
settings = load_settings()
result = {
"success": False,
"download_id": download_id,
"path": dest_path,
"error": None,
"size": 0,
}
start_time = time.time()
filename = get_filename_from_path(dest_path)
expected_sha256 = _extract_expected_sha256(metadata)
with download_lock:
download_progress[download_id] = {
"status": "starting",
"progress": 0,
"total_size": 0,
"downloaded": 0,
"filename": filename,
"path": dest_path,
"directory": os.path.dirname(dest_path),
"url": url,
"error": None,
"speed": 0,
"start_time": start_time,
"download_backend": "aria2",
}
try:
os.makedirs(os.path.dirname(dest_path), exist_ok=True)
_ensure_aria2_daemon(settings)
aria2_url, request_headers = _resolve_download_url_for_aria2(url, headers)
options: Dict[str, Any] = {
"dir": os.path.dirname(dest_path),
"out": filename,
"continue": "true",
"max-connection-per-server": "4",
"split": "4",
"min-split-size": "1M",
"allow-overwrite": "true",
"auto-file-renaming": "false",
"file-allocation": "none",
"no-want-digest-header": "true",
# Redirects were already resolved and validated above. Keeping them
# disabled prevents sensitive headers from reaching another host.
"max-redirect": "0",
}
if expected_sha256:
# Let aria2 verify the file as part of the transfer instead of
# requiring a second Python pass over a multi-gigabyte model.
options["checksum"] = f"sha-256={expected_sha256}"
options["check-integrity"] = "true"
user_agent = _get_header_value(request_headers, "User-Agent")
referer = _get_header_value(request_headers, "Referer")
if user_agent:
options["user-agent"] = user_agent
if referer:
options["referer"] = referer
header_values = [
f"{key}: {value}"
for key, value in request_headers.items()
if str(key).lower() not in {"user-agent", "referer"}
]
if header_values:
options["header"] = header_values
gid = _aria2_rpc("aria2.addUri", [[aria2_url], options])
if not isinstance(gid, str) or not gid:
raise Aria2Error("aria2 did not return a download gid")
with aria2_lock:
aria2_transfers[download_id] = {
"gid": gid,
"path": dest_path,
}
with download_lock:
download_progress[download_id]["aria2_gid"] = gid
download_progress[download_id]["status"] = "downloading"
log.info(f"Starting aria2 download: {filename}")
last_cli_log = start_time
while True:
if download_id in cancelled_downloads:
try:
_aria2_rpc("aria2.forceRemove", [gid])
except Exception:
pass
with download_lock:
if download_id in download_progress:
download_progress[download_id]["status"] = "cancelled"
download_progress[download_id]["speed"] = 0
_delete_partial_download_files(dest_path)
cancelled_downloads.discard(download_id)
result["error"] = "Download cancelled"
return result
status = _aria2_tell_status(gid)
state = str(status.get("status") or "")
total_size = _parse_aria2_int(status.get("totalLength"))
downloaded = _parse_aria2_int(status.get("completedLength"))
speed = _parse_aria2_int(status.get("downloadSpeed"))
progress = int((downloaded / total_size) * 100) if total_size > 0 else 0
mapped_status = {
"active": "downloading",
"waiting": "downloading",
"paused": "paused",
"complete": "completed",
"error": "error",
"removed": "cancelled",
}.get(state, state or "downloading")
with download_lock:
if download_id in download_progress:
download_progress[download_id].update(
{
"status": mapped_status,
"progress": max(0, min(progress, 100)),
"total_size": total_size,
"downloaded": downloaded,
"speed": 0 if mapped_status in {"paused", "completed"} else speed,
"download_backend": "aria2",
"aria2_gid": gid,
}
)
now = time.time()
if now - last_cli_log >= CLI_LOG_INTERVAL and mapped_status == "downloading":
last_cli_log = now
total_str = format_bytes(total_size) if total_size else "?"
log.info(
f"aria2 progress: {format_bytes(downloaded)} / {total_str} ({progress}%) - {format_bytes(speed)}/s"
)
if state == "complete":
completed_path = _resolve_aria2_completed_path(status, dest_path)
size = os.path.getsize(completed_path) if os.path.exists(completed_path) else downloaded
metadata_path = write_model_resolver_metadata(
completed_path,
metadata or {},
category,
url,
create_preview=True,
)
with download_lock:
download_progress[download_id].update(
{
"status": "completed",
"progress": 100,
"downloaded": size,
"total_size": total_size or size,
"speed": 0,
"path": completed_path,
"directory": os.path.dirname(completed_path),
}
)
if metadata_path:
download_progress[download_id]["metadata_path"] = metadata_path
result.update(
{
"success": True,
"path": completed_path,
"size": size,
"metadata_path": metadata_path,
}
)
elapsed = time.time() - start_time
avg_speed = size / elapsed if elapsed > 0 else 0
log.info(f"✓ aria2 download complete: {filename}")
log.info(
f"Size: {format_bytes(size)}, Time: {elapsed:.1f}s, Avg speed: {format_bytes(int(avg_speed))}/s"
)
invalidate_model_files_cache()
invalidate_local_hash_match_cache()
return result
if state == "error":
error_msg = status.get("errorMessage") or "aria2 download failed"
with download_lock:
download_progress[download_id]["status"] = "error"
download_progress[download_id]["error"] = error_msg
result["error"] = error_msg
return result
if state == "removed":
with download_lock:
download_progress[download_id]["status"] = "cancelled"
download_progress[download_id]["speed"] = 0
_delete_partial_download_files(dest_path)
cancelled_downloads.discard(download_id)
result["error"] = "Download cancelled"
return result
time.sleep(0.5)
except Exception as exc:
error_msg = _sanitize_download_error(exc)
with download_lock:
if download_id in download_progress:
download_progress[download_id]["status"] = "error"
download_progress[download_id]["error"] = error_msg
download_progress[download_id]["speed"] = 0
result["error"] = error_msg
log.error(f"✗ aria2 download failed: {filename}")
log.error(f"Error: {error_msg}")
return result
finally:
with aria2_lock:
aria2_transfers.pop(download_id, None)
aria2_action_locks.pop(download_id, None)
aria2_desired_states.pop(download_id, None)
_schedule_aria2_idle_stop()
def download_file(
url: str,
dest_path: str,