From 59bc928f8e3d796b7848aa1fab3dfbe2581daa9d Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Mon, 26 Oct 2026 14:27:44 +0100 Subject: [PATCH] Verify aria2 download SHA256 in background to avoid blocking at 100% - 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 --- core/download/api.py | 6 ++ core/download/aria2_backend.py | 100 ++++++++++++++++++++++++++-- tests/test_downloader_boundaries.py | 97 +++++++++++++++++++++++++++ 3 files changed, 198 insertions(+), 5 deletions(-) diff --git a/core/download/api.py b/core/download/api.py index e90d375..2b390d4 100644 --- a/core/download/api.py +++ b/core/download/api.py @@ -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) diff --git a/core/download/aria2_backend.py b/core/download/aria2_backend.py index dd6e38b..4299272 100644 --- a/core/download/aria2_backend.py +++ b/core/download/aria2_backend.py @@ -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, diff --git a/tests/test_downloader_boundaries.py b/tests/test_downloader_boundaries.py index 83cb969..b229a5f 100644 --- a/tests/test_downloader_boundaries.py +++ b/tests/test_downloader_boundaries.py @@ -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__":