From 0437250b5aed896a3cbdec262bae62066f4b7dc2 Mon Sep 17 00:00:00 2001 From: watarika Date: Fri, 10 Jul 2026 20:25:05 +0900 Subject: [PATCH] feat(batch-downloader): run aria2 jobs --- model_batch_downloader/aria2_runner.py | 178 +++++++++++++++++++++++++ tests/test_aria2_runner.py | 144 ++++++++++++++++++++ 2 files changed, 322 insertions(+) create mode 100644 model_batch_downloader/aria2_runner.py create mode 100644 tests/test_aria2_runner.py diff --git a/model_batch_downloader/aria2_runner.py b/model_batch_downloader/aria2_runner.py new file mode 100644 index 0000000..4bb3989 --- /dev/null +++ b/model_batch_downloader/aria2_runner.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +from dataclasses import dataclass +import os +from pathlib import Path +import subprocess +import tempfile +import time +from typing import Callable, Mapping +from urllib.parse import urlsplit + +from .manifest import ResolvedItem +from .security import auth_for_url, authenticated_url, redact + + +@dataclass(frozen=True, slots=True) +class DownloadRecord: + item_id: str + model_type: str + relative_path: str + absolute_path: str + status: str + bytes: int + elapsed_seconds: float + + +@dataclass(frozen=True, slots=True) +class DownloadResult: + entries: dict[str, DownloadRecord] + + +@dataclass(frozen=True, slots=True) +class DownloadFailure: + item_id: str + source: str + exit_code: int | None + detail: str + + +class BatchDownloadError(RuntimeError): + def __init__(self, failures: tuple[DownloadFailure, ...]): + self.failures = failures + lines = ["One or more model downloads failed:"] + lines.extend( + f"- {failure.item_id} ({failure.source}, exit={failure.exit_code}): " + f"{failure.detail}" + for failure in failures + ) + super().__init__("\n".join(lines)) + + +def _safe_source(url: str) -> str: + parsed = urlsplit(url) + return f"{parsed.hostname or 'unknown'}{parsed.path}" + + +def _control_text( + item: ResolvedItem, environ: Mapping[str, str] +) -> tuple[str, tuple[str, ...]]: + auth = auth_for_url(item.url, environ) + source_url = authenticated_url(item.url, auth) + lines = [ + source_url, + f" dir={item.destination.parent}", + f" out={item.destination.name}", + " continue=true", + " auto-file-renaming=false", + " allow-overwrite=false", + f" split={item.split}", + f" max-connection-per-server={item.split}", + ] + if auth.header: + lines.append(f" header={auth.header[0]}: {auth.header[1]}") + return "\n".join(lines) + "\n", auth.secrets + + +def _existing_complete_path(item: ResolvedItem) -> Path | None: + if item.existing_path and item.existing_path.is_file(): + return item.existing_path + if item.destination.is_file(): + return item.destination + return None + + +def run_downloads( + items: tuple[ResolvedItem, ...], + executable: str, + environ: Mapping[str, str] | None = None, + run_process: Callable[..., object] = subprocess.run, +) -> DownloadResult: + environment = os.environ if environ is None else environ + entries: dict[str, DownloadRecord] = {} + failures: list[DownloadFailure] = [] + + for item in items: + sidecar = Path(str(item.destination) + ".aria2") + complete = _existing_complete_path(item) + if complete and not sidecar.exists(): + entries[item.item_id] = DownloadRecord( + item.item_id, + item.model_type, + item.relative_path.as_posix(), + str(complete), + "skipped", + complete.stat().st_size, + 0.0, + ) + continue + + item.destination.parent.mkdir(parents=True, exist_ok=True) + started_resumed = sidecar.exists() + control_text, secrets = _control_text(item, environment) + control_path: Path | None = None + started = time.monotonic() + + try: + with tempfile.NamedTemporaryFile( + "w", + encoding="utf-8", + suffix=".aria2-input", + delete=False, + ) as handle: + control_path = Path(handle.name) + handle.write(control_text) + try: + control_path.chmod(0o600) + except OSError: + pass + + completed = run_process( + [executable, "--input-file", str(control_path)], + capture_output=True, + text=True, + check=False, + ) + return_code = getattr(completed, "returncode", None) + stdout = getattr(completed, "stdout", "") or "" + stderr = getattr(completed, "stderr", "") or "" + if return_code != 0 or not item.destination.is_file() or sidecar.exists(): + detail = redact( + (stderr or stdout or "aria2 did not complete the file")[:2000], + secrets, + ) + failures.append( + DownloadFailure( + item.item_id, + _safe_source(item.url), + return_code, + detail, + ) + ) + continue + + entries[item.item_id] = DownloadRecord( + item.item_id, + item.model_type, + item.relative_path.as_posix(), + str(item.destination), + "resumed" if started_resumed else "downloaded", + item.destination.stat().st_size, + time.monotonic() - started, + ) + except Exception as exception: + failures.append( + DownloadFailure( + item.item_id, + _safe_source(item.url), + None, + redact(str(exception), secrets), + ) + ) + finally: + if control_path is not None: + control_path.unlink(missing_ok=True) + + if failures: + raise BatchDownloadError(tuple(failures)) + return DownloadResult(entries) diff --git a/tests/test_aria2_runner.py b/tests/test_aria2_runner.py new file mode 100644 index 0000000..a061950 --- /dev/null +++ b/tests/test_aria2_runner.py @@ -0,0 +1,144 @@ +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from model_batch_downloader.aria2_runner import BatchDownloadError, run_downloads +from model_batch_downloader.manifest import ResolvedItem + + +def make_item( + tmp_path, + name="model.safetensors", + item_id="model", + url=None, + existing_path=None, +): + return ResolvedItem( + url or "https://huggingface.co/a/resolve/main/" + name, + "diffusion_models", + "", + name, + item_id, + 8, + tmp_path / name, + Path(name), + existing_path, + ) + + +def test_complete_file_is_skipped_without_process(tmp_path): + target = tmp_path / "model.safetensors" + target.write_bytes(b"ready") + resolved = make_item(tmp_path, existing_path=target) + calls = [] + + result = run_downloads( + (resolved,), "aria2c", {}, lambda *args, **kwargs: calls.append(args) + ) + + assert result.entries["model"].status == "skipped" + assert result.entries["model"].bytes == 5 + assert calls == [] + + +def test_partial_sidecar_marks_success_as_resumed_and_deletes_control_file(tmp_path): + resolved = make_item(tmp_path) + sidecar = Path(str(resolved.destination) + ".aria2") + sidecar.write_text("state", encoding="utf-8") + seen = {} + + def fake_run(command, **_kwargs): + seen["command"] = command + seen["control_path"] = Path(command[-1]) + seen["control"] = seen["control_path"].read_text(encoding="utf-8") + resolved.destination.write_bytes(b"done") + sidecar.unlink() + return SimpleNamespace(returncode=0, stdout="ok", stderr="") + + result = run_downloads( + (resolved,), "aria2c", {"HF_TOKEN": "secret"}, fake_run + ) + + assert result.entries["model"].status == "resumed" + assert "secret" not in " ".join(seen["command"]) + assert "Authorization: Bearer secret" in seen["control"] + assert not seen["control_path"].exists() + + +def test_new_civitai_download_puts_token_only_in_control_file(tmp_path): + resolved = make_item( + tmp_path, + url="https://civitai.com/api/download/models/42?format=SafeTensor", + ) + seen = {} + + def fake_run(command, **_kwargs): + seen["command"] = command + seen["control"] = Path(command[-1]).read_text(encoding="utf-8") + resolved.destination.write_bytes(b"done") + return SimpleNamespace(returncode=0, stdout="ok", stderr="") + + result = run_downloads( + (resolved,), "aria2c", {"CIVITAI_API_TOKEN": "cv_secret"}, fake_run + ) + + assert result.entries["model"].status == "downloaded" + assert "cv_secret" not in " ".join(seen["command"]) + assert "format=SafeTensor&token=cv_secret" in seen["control"] + assert "allow-overwrite=false" in seen["control"] + assert "auto-file-renaming=false" in seen["control"] + + +def test_failure_continues_then_raises_redacted_summary(tmp_path): + first = make_item(tmp_path, "a.safetensors", "a") + second = make_item(tmp_path, "b.safetensors", "b") + attempts = [] + + def fake_run(command, **_kwargs): + attempts.append(command) + if len(attempts) == 2: + second.destination.write_bytes(b"ok") + return SimpleNamespace(returncode=0, stdout="", stderr="") + return SimpleNamespace( + returncode=3, stdout="", stderr="Authorization: Bearer secret" + ) + + with pytest.raises(BatchDownloadError) as captured: + run_downloads( + (first, second), "aria2c", {"HF_TOKEN": "secret"}, fake_run + ) + + assert len(attempts) == 2 + assert "secret" not in str(captured.value) + assert "[REDACTED]" in str(captured.value) + assert "a" in str(captured.value) + + +def test_process_exception_continues_to_next_item(tmp_path): + first = make_item(tmp_path, "a.safetensors", "a") + second = make_item(tmp_path, "b.safetensors", "b") + attempts = 0 + + def fake_run(_command, **_kwargs): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise OSError("launch failed") + second.destination.write_bytes(b"ok") + return SimpleNamespace(returncode=0, stdout="", stderr="") + + with pytest.raises(BatchDownloadError, match="launch failed"): + run_downloads((first, second), "aria2c", {}, fake_run) + + assert attempts == 2 + + +def test_success_requires_completed_file_and_removed_sidecar(tmp_path): + resolved = make_item(tmp_path) + + def fake_run(_command, **_kwargs): + return SimpleNamespace(returncode=0, stdout="", stderr="") + + with pytest.raises(BatchDownloadError, match="did not complete"): + run_downloads((resolved,), "aria2c", {}, fake_run)