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

Verify aria2 download SHA256 in background to avoid blocking at 100%

Browse Source 
- Omit blocking checksum and check-integrity options from aria2 transfers to prevent long pauses at 100% on large models
- Run SHA-256 integrity verification asynchronously in a background thread via Python hashlib
- Update download progress and metadata sidecars upon background verification completion or mismatch
- Add unit tests for aria2 option omission and background verification
This commit is contained in:
Dariusz L
2026-10-26 14:27:44 +01:00
parent 55197621eb
commit 59bc928f8e
3 changed files with 198 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/download/api.py
+6
View File
@@ -28,6 +28,7 @@ from ..path_utils import (
get_model_resolver_sidecar_path,
get_path_identity,
is_path_within,
read_json_safe,
write_json_atomic,
)
from ..resolver import invalidate_model_caches
@@ -140,6 +141,7 @@ 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
@@ -189,6 +191,10 @@ context._run_huggingface_xet_transfer = _bind(
)
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)
core/download/aria2_backend.py
+95 -5
View File
@@ -772,11 +772,11 @@ def download_file_with_aria2(
# 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"
# Checksum and check-integrity are deliberately omitted from aria2 options.
# On Windows, aria2c's single-threaded, unaccelerated software checksum
# verification blocks download completion and freezes the process at 100%
# for many minutes on 20+ GB models. Instead, verification is handled
# asynchronously in the background via hardware-accelerated Python hashlib.
user_agent = facade._get_header_value(request_headers, "User-Agent")
referer = facade._get_header_value(request_headers, "Referer")
if user_agent:
@@ -968,6 +968,16 @@ def download_file_with_aria2(
f"Avg speed: {facade.format_bytes(int(avg_speed))}/s"
)
facade.invalidate_model_caches()
if expected_sha256:
start_background_sha256_verification(
completed_path,
expected_sha256,
download_id=download_id,
metadata_path=metadata_path,
dependencies=facade,
)
return result
if state == "error":
@@ -1030,6 +1040,86 @@ def download_file_with_aria2(
facade._schedule_aria2_idle_stop()
def start_background_sha256_verification(
completed_path: str,
expected_sha256: str,
*,
download_id: str = "",
metadata_path: Optional[str] = None,
dependencies: Any = None,
) -> Optional[Any]:
"""Verify downloaded model SHA256 in a background thread using hardware-accelerated Python hashlib."""
facade = require_download_dependencies(dependencies, "aria2 backend")
clean_expected = str(expected_sha256 or "").lower().strip()
if not clean_expected or not facade.os.path.isfile(completed_path):
return None
def _worker() -> None:
filename = facade.get_filename_from_path(completed_path)
log.debug(f"Starting background SHA256 verification for {filename}...")
try:
hasher = facade.hashlib.sha256()
with open(completed_path, "rb") as file_handle:
while True:
chunk = file_handle.read(4 * 1024 * 1024)
if not chunk:
break
hasher.update(chunk)
actual_sha256 = hasher.hexdigest().lower()
if actual_sha256 == clean_expected:
log.info(f"✓ Background SHA256 verified for: {filename}")
if download_id:
with facade.download_lock:
if download_id in facade.download_progress:
facade.download_progress[download_id]["sha256_verified"] = True
facade.download_progress[download_id]["sha256"] = actual_sha256
target_meta_path = (
metadata_path
or facade.get_model_resolver_sidecar_path(completed_path)
)
if target_meta_path and facade.os.path.isfile(target_meta_path):
meta = facade.read_json_safe(target_meta_path, {})
if isinstance(meta, dict):
meta["sha256_verified"] = True
meta["sha256"] = actual_sha256
facade.write_json_atomic(target_meta_path, meta, indent=2)
else:
log.error(
f"✗ SHA256 mismatch detected for {filename}! "
f"Expected {clean_expected}, got {actual_sha256}"
)
if download_id:
with facade.download_lock:
if download_id in facade.download_progress:
facade.download_progress[download_id]["sha256_verified"] = False
facade.download_progress[download_id]["actual_sha256"] = actual_sha256
facade.download_progress[download_id]["error"] = (
f"SHA256 mismatch: expected {clean_expected}, got {actual_sha256}"
)
target_meta_path = (
metadata_path
or facade.get_model_resolver_sidecar_path(completed_path)
)
if target_meta_path and facade.os.path.isfile(target_meta_path):
meta = facade.read_json_safe(target_meta_path, {})
if isinstance(meta, dict):
meta["sha256_verified"] = False
meta["actual_sha256"] = actual_sha256
facade.write_json_atomic(target_meta_path, meta, indent=2)
except Exception as exc:
log.warning(f"Background SHA256 verification failed for {filename}: {exc}")
thread_name = f"mr-verify-sha256-{download_id or facade.get_filename_from_path(completed_path)}"
thread = facade.threading.Thread(
target=_worker,
name=thread_name,
daemon=True,
)
thread.start()
return thread
def force_remove_aria2_transfer(
download_id: str,
gid: str,
tests/test_downloader_boundaries.py
+97
View File
@@ -7,9 +7,11 @@ from unittest.mock import MagicMock, call, patch
from core.download.api import context as downloader
from core.download.aria2_backend import (
download_file_with_aria2,
pause_download,
recover_aria2_missing_control_file,
resume_download,
start_background_sha256_verification,
)
from core.download.orchestrator import download_file, find_active_download_for_path
from core.path_utils import normalize_metadata_file_path
@@ -386,6 +388,101 @@ class DownloaderBoundaryTests(unittest.TestCase):
self.assertFalse(os.path.exists(dest_path))
self.assertFalse(os.path.exists(partial_path))
def test_aria2_options_omit_checksum_and_integrity_check(self):
with tempfile.TemporaryDirectory() as temp_dir:
dest_path = os.path.join(temp_dir, "model.safetensors")
mock_rpc = MagicMock()
mock_rpc.side_effect = ["test-gid", {}]
with patch.object(downloader, "_ensure_aria2_daemon"), \
patch.object(downloader, "_resolve_download_url_for_aria2", return_value=("https://example.com/model.safetensors", {})), \
patch.object(downloader, "_aria2_rpc", mock_rpc), \
patch.object(downloader, "_aria2_tell_status", return_value={"status": "complete", "completedLength": "100", "totalLength": "100"}), \
patch.object(downloader, "write_model_resolver_metadata", return_value=""), \
patch.object(downloader, "_schedule_aria2_idle_stop"), \
patch.object(downloader, "invalidate_model_caches"), \
patch.object(downloader, "_start_background_sha256_verification"):
result = download_file_with_aria2(
"https://example.com/model.safetensors",
dest_path,
"aria2-no-checksum-test",
metadata={"sha256": "a" * 64},
dependencies=downloader,
)
self.assertTrue(result["success"])
add_uri_options = mock_rpc.call_args_list[0].args[1][1]
self.assertNotIn("checksum", add_uri_options)
self.assertNotIn("check-integrity", add_uri_options)
def test_background_sha256_verification_success(self):
with tempfile.TemporaryDirectory() as temp_dir:
dest_path = os.path.join(temp_dir, "test_model.safetensors")
content = b"sample model weights for sha256 check"
with open(dest_path, "wb") as f:
f.write(content)
import hashlib
expected_sha = hashlib.sha256(content).hexdigest()
sidecar_path = f"{dest_path}.modelresolver.json"
downloader.write_json_atomic(sidecar_path, {"sha256": expected_sha})
download_id = "bg-verify-success"
with downloader.download_lock:
downloader.download_progress[download_id] = {"status": "completed"}
thread = start_background_sha256_verification(
dest_path,
expected_sha,
download_id=download_id,
metadata_path=sidecar_path,
dependencies=downloader,
)
self.assertIsNotNone(thread)
thread.join(timeout=5)
with downloader.download_lock:
progress = downloader.download_progress.get(download_id, {})
self.assertTrue(progress.get("sha256_verified"))
self.assertEqual(expected_sha, progress.get("sha256"))
sidecar_meta = downloader.read_json_safe(sidecar_path, {})
self.assertTrue(sidecar_meta.get("sha256_verified"))
self.assertEqual(expected_sha, sidecar_meta.get("sha256"))
def test_background_sha256_verification_mismatch(self):
with tempfile.TemporaryDirectory() as temp_dir:
dest_path = os.path.join(temp_dir, "test_corrupt.safetensors")
with open(dest_path, "wb") as f:
f.write(b"corrupted content")
expected_sha = "0" * 64
sidecar_path = f"{dest_path}.modelresolver.json"
downloader.write_json_atomic(sidecar_path, {"sha256": expected_sha})
download_id = "bg-verify-mismatch"
with downloader.download_lock:
downloader.download_progress[download_id] = {"status": "completed"}
thread = start_background_sha256_verification(
dest_path,
expected_sha,
download_id=download_id,
metadata_path=sidecar_path,
dependencies=downloader,
)
self.assertIsNotNone(thread)
thread.join(timeout=5)
with downloader.download_lock:
progress = downloader.download_progress.get(download_id, {})
self.assertFalse(progress.get("sha256_verified"))
self.assertIn("mismatch", progress.get("error", "").lower())
sidecar_meta = downloader.read_json_safe(sidecar_path, {})
self.assertFalse(sidecar_meta.get("sha256_verified"))
if __name__ == "__main__":