feat(batch-downloader): run aria2 jobs
This commit is contained in:
watarika
@@ -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)
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user