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

refactor: centralize initial download progress

Browse Source
This commit is contained in:
Dariusz L
2026-08-05 02:13:51 +02:00
parent 0b3c1785c7
commit e326a107f4
5 changed files with 111 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/download/aria2_backend.py
+9 -14
View File
@@ -4,6 +4,7 @@ import socket
from typing import Any, Dict, Optional
from ..log_system import create_module_logger
from .state import create_initial_progress
log = create_module_logger("core.downloader")
@@ -724,20 +725,14 @@ def download_file_with_aria2(
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",
}
facade.download_progress[download_id] = create_initial_progress(
url=url,
path=dest_path,
filename=filename,
directory=facade.os.path.dirname(dest_path),
download_backend="aria2",
start_time=start_time,
)
try:
facade.os.makedirs(facade.os.path.dirname(dest_path), exist_ok=True)
core/download/huggingface_xet.py
+10 -14
View File
@@ -5,6 +5,7 @@ from typing import Any, Dict, Optional
from urllib.parse import urlparse
from ..log_system import create_module_logger
from .state import create_initial_progress
log = create_module_logger("core.downloader")
@@ -334,20 +335,15 @@ def download_huggingface_xet(
)
with facade.download_lock:
facade.download_progress[download_id] = {
"status": "starting",
"progress": 0,
"total_size": expected_size,
"downloaded": 0,
"filename": filename,
"path": dest_path,
"directory": facade.os.path.dirname(dest_path),
"url": validated_url,
"error": None,
"speed": 0,
"start_time": start_time,
"download_backend": "huggingface_xet",
}
facade.download_progress[download_id] = create_initial_progress(
url=validated_url,
path=dest_path,
filename=filename,
directory=facade.os.path.dirname(dest_path),
download_backend="huggingface_xet",
start_time=start_time,
total_size=expected_size,
)
facade.cancelled_downloads.discard(download_id)
try:
core/download/orchestrator.py
+18 -28
View File
@@ -2,6 +2,8 @@
from typing import Any, Callable, Dict, Optional
from .state import create_initial_progress
def find_active_download_for_path(
download_progress: Dict[str, Dict[str, Any]],
@@ -112,20 +114,14 @@ def download_file(
last_cli_log = start_time
with download_lock:
download_progress[download_id] = {
"status": "starting",
"progress": 0,
"total_size": 0,
"downloaded": 0,
"filename": facade.get_filename_from_path(dest_path),
"path": dest_path,
"directory": os.path.dirname(dest_path),
"url": url,
"error": None,
"speed": 0,
"start_time": start_time,
"download_backend": "python",
}
download_progress[download_id] = create_initial_progress(
url=url,
path=dest_path,
filename=facade.get_filename_from_path(dest_path),
directory=os.path.dirname(dest_path),
download_backend="python",
start_time=start_time,
)
try:
destination_directory = os.path.dirname(dest_path)
@@ -596,20 +592,14 @@ def start_background_download(
f"(download_id={active_download_id})"
)
return active_download_id
download_progress[download_id] = {
"status": "starting",
"progress": 0,
"total_size": 0,
"downloaded": 0,
"filename": filename,
"path": initial_path,
"directory": initial_directory,
"url": url,
"error": None,
"speed": 0,
"start_time": time.time(),
"download_backend": facade._download_backend_from_settings(),
}
download_progress[download_id] = create_initial_progress(
url=url,
path=initial_path,
filename=filename,
directory=initial_directory,
download_backend=facade._download_backend_from_settings(),
start_time=time.time(),
)
def run_download() -> None:
try:
core/download/state.py
+27
View File
@@ -16,6 +16,33 @@ xet_transfers: Dict[str, Dict[str, Any]] = {}
xet_transfers_lock = threading.Lock()
def create_initial_progress(
*,
url: str,
path: str,
filename: str,
directory: str,
download_backend: str,
start_time: float,
total_size: int = 0,
) -> Dict[str, Any]:
"""Build the shared initial progress record for a download."""
return {
"status": "starting",
"progress": 0,
"total_size": total_size,
"downloaded": 0,
"filename": filename,
"path": path,
"directory": directory,
"url": url,
"error": None,
"speed": 0,
"start_time": start_time,
"download_backend": download_backend,
}
@dataclass(frozen=True)
class DownloadStateDependencies:
"""State and side-effect dependencies supplied by the composition layer."""
tests/test_download_state.py
+47
View File
@@ -0,0 +1,47 @@
import importlib
def test_initial_download_progress_preserves_the_shared_contract():
state = importlib.import_module("core.download.state")
create_initial_progress = getattr(state, "create_initial_progress", None)
assert callable(create_initial_progress)
progress = create_initial_progress(
url="https://example.com/model.safetensors",
path=r"C:\models\model.safetensors",
filename="model.safetensors",
directory=r"C:\models",
download_backend="aria2",
total_size=42,
start_time=12.5,
)
assert progress == {
"status": "starting",
"progress": 0,
"total_size": 42,
"downloaded": 0,
"filename": "model.safetensors",
"path": r"C:\models\model.safetensors",
"directory": r"C:\models",
"url": "https://example.com/model.safetensors",
"error": None,
"speed": 0,
"start_time": 12.5,
"download_backend": "aria2",
}
def test_initial_download_progress_defaults_total_size_to_zero():
state = importlib.import_module("core.download.state")
progress = state.create_initial_progress(
url="https://example.com/model.safetensors",
path="model.safetensors",
filename="model.safetensors",
directory="",
download_backend="python",
start_time=1.0,
)
assert progress["total_size"] == 0
assert progress["download_backend"] == "python"