diff --git a/changelog.d/tsk-ilhlue-download-timeout-resume.md b/changelog.d/tsk-ilhlue-download-timeout-resume.md new file mode 100644 index 000000000..8ef97b73e --- /dev/null +++ b/changelog.d/tsk-ilhlue-download-timeout-resume.md @@ -0,0 +1,27 @@ +### Fixed + +- Model downloads no longer hang forever on a stalled connection. The HTTP + transfer in `tinyagentos/download_manager.py` ran with `timeout=None`, which + disables the connect, read, write and pool timeouts together: a Wi-Fi drop, a + NAT table eviction or a CDN edge that stopped sending left the task at + `status="downloading"` with nothing ever erroring, showing a progress bar + frozen part-way with no way to tell it apart from a slow link. It now uses + finite timeouts and retries transport errors and 5xx responses with + exponential backoff, so a single transient failure from a mirror no longer + kills a multi-gigabyte transfer. Expect previously invisible stalls to start + surfacing as errors — that is the fix working. +- Interrupted model downloads resume instead of restarting. Bytes already on + disk are asked for with a `Range` header, so a 40 GB model that fails at + 39 GB continues from where it stopped; a server that ignores the header and + answers `200` restarts cleanly rather than appending a second copy. +- A failed model download no longer leaves a corrupt file at the canonical + path, where every later "is this model installed?" existence check would take + it for a real weight. Bytes are staged in a `.part` file and renamed + onto the destination only after validation passes. +- Finished download tasks are pruned after an hour instead of staying resident + for the lifetime of the process, so `/api/models/downloads` no longer grows + without bound. Pending and downloading tasks are never pruned. +- A re-download of an already-installed model no longer deletes the existing + valid file when the new attempt fails before promoting anything: the + cleanup on failure now only removes `task.dest` when this attempt actually + renamed the `.part` stage file onto it. diff --git a/docs/design/model-torrent-mesh.md b/docs/design/model-torrent-mesh.md index c3c00ed37..4240a9d40 100644 --- a/docs/design/model-torrent-mesh.md +++ b/docs/design/model-torrent-mesh.md @@ -89,6 +89,15 @@ Remaining: DownloadManager passkey fetch via the account-session proxy + clean H 3. If torrent completes: verify SHA256 against manifest, then done. 4. If HTTP completes first (web seed from inside libtorrent, or direct fallback): still verify SHA256. +- **HTTP path robustness** — the `download_url` fallback is not a bare fetch: + finite connect/read/write/pool timeouts (a half-open connection surfaces as + an error instead of a progress bar frozen at 63% forever), retry with + exponential backoff on transport errors and 5xx, and `Range`-header resume + so a 40 GB pull that dies at 39 GB does not restart from zero. Bytes are + staged in a `.part` file and renamed onto the canonical path only + after the SHA256 check passes, so a failed download can never leave a + corrupt weight where a later "is this model installed?" check would find it. + The stage file is kept on failure — it is what the next attempt resumes from. - **Seeding** — after a successful download, the torrent is kept in the libtorrent session. Seeding runs in the background with user-configurable upload limits. diff --git a/tests/test_download_manager.py b/tests/test_download_manager.py index c0b7d2fcb..448a009b5 100644 --- a/tests/test_download_manager.py +++ b/tests/test_download_manager.py @@ -1,11 +1,14 @@ import asyncio import hashlib +import time from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import pytest_asyncio +from tinyagentos import download_manager from tinyagentos.download_manager import DownloadManager, DownloadTask @@ -639,3 +642,389 @@ async def _install(): assert task.status == "complete" assert task.downloaded_bytes == 0 assert task.total_bytes == 0 + + +# --------------------------------------------------------------------------- +# _download: timeouts, resume, retry and cleanup +# --------------------------------------------------------------------------- + +class _FakeHttpServer: + """A stand-in httpx.AsyncClient that honours the transport contract the + real client documents, so a test can drive the failure modes a home + internet connection actually produces without opening a socket. + + It reproduces the two behaviours that matter here: + + * ``timeout=None`` disables the read timeout, so a server that stops + sending mid-body leaves the stream awaiting forever. A finite + ``read`` raises ``httpx.ReadTimeout`` instead. + * A request carrying ``Range: bytes=N-`` is answered with 206 and the + remaining bytes, exactly as a CDN that supports resume does. + """ + + def __init__(self, body: bytes, *, stall_after: int | None = None, + fail_attempts: int = 0, failure=None, supports_range: bool = True): + self.body = body + self.stall_after = stall_after + self.fail_attempts = fail_attempts + self.failure = failure or httpx.ReadError("connection reset") + self.supports_range = supports_range + self.requests: list[dict] = [] + self.timeouts: list[object] = [] + + def __call__(self, *args, **kwargs): + self.timeouts.append(kwargs.get("timeout")) + return _FakeClient(self, kwargs.get("timeout")) + + +class _FakeClient: + def __init__(self, server: "_FakeHttpServer", timeout): + self._server = server + self._timeout = timeout + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + def stream(self, method, url, headers=None, **kwargs): + headers = dict(headers or {}) + self._server.requests.append(headers) + return _FakeStream(self._server, self._timeout, headers) + + +class _FakeStream: + def __init__(self, server: "_FakeHttpServer", timeout, headers: dict): + self._server = server + self._timeout = timeout + self._attempt = len(server.requests) + rng = headers.get("Range") + self._offset = 0 + self.status_code = 200 + self.headers: dict[str, str] = {} + if rng and server.supports_range: + self._offset = int(rng.removeprefix("bytes=").rstrip("-")) + if self._offset >= len(server.body): + self.status_code = 416 + else: + self.status_code = 206 + self.headers["content-range"] = ( + f"bytes {self._offset}-{len(server.body) - 1}/{len(server.body)}" + ) + self._payload = server.body[self._offset:] + if self.status_code != 416: + self.headers["content-length"] = str(len(self._payload)) + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + def raise_for_status(self): + if self.status_code >= 400: + raise httpx.HTTPStatusError( + f"HTTP {self.status_code}", + request=httpx.Request("GET", "http://example.com/f.bin"), + response=httpx.Response(self.status_code), + ) + + async def aiter_bytes(self, chunk_size=65536): + failing = self._attempt <= self._server.fail_attempts + # A dropped connection cuts the body off part-way through; a failure + # that arrives only after every byte has been delivered would leave a + # complete stage file and never exercise resume. + limit = len(self._payload) // 2 if failing else None + sent = 0 + for i in range(0, len(self._payload), chunk_size): + if self._server.stall_after is not None and sent >= self._server.stall_after: + break + chunk = self._payload[i:i + chunk_size] + if limit is not None and sent + len(chunk) > limit: + chunk = chunk[:limit - sent] + if chunk: + yield chunk + break + yield chunk + sent += len(chunk) + if failing: + raise self._server.failure + if self._server.stall_after is not None: + read = getattr(self._timeout, "read", self._timeout) + if read is None: + # Half-open connection: the peer stops sending and never + # closes. With no read timeout the stream waits forever. + await asyncio.Event().wait() + # Stand-in for the real `read`-second wait; what is under test is + # that a finite read timeout turns the stall into an error at all, + # not how long that error takes to arrive. + await asyncio.sleep(0.01) + raise httpx.ReadTimeout("timed out reading response body") + + +@pytest.fixture +def fast_retries(): + """Collapse the production backoff so a retry test costs milliseconds. + + ``create=True`` so the fixture still applies against a build with no + retry at all — the red run then shows the download defect under test + rather than an AttributeError from the fixture. + """ + with patch.object(download_manager, "DOWNLOAD_RETRY_BASE_DELAY", 0.001, create=True), \ + patch.object(download_manager, "DOWNLOAD_RETRY_MAX_DELAY", 0.001, create=True): + yield + + +class TestDownloadTimeoutResumeAndCleanup: + @pytest_asyncio.fixture + def dm(self): + return DownloadManager() + + @pytest.mark.asyncio + async def test_stalled_connection_errors_instead_of_hanging(self, dm, tmp_path, fast_retries): + """A server that stops sending mid-body must surface as an error. + With timeout=None the task sits at status="downloading" forever and + the user watches a progress bar frozen at 63%.""" + dest = tmp_path / "out.bin" + task = DownloadTask(id="dl", url="http://example.com/f.bin", dest=dest) + server = _FakeHttpServer(b"x" * 200_000, stall_after=65536) + + with patch("tinyagentos.download_manager.httpx.AsyncClient", server): + await asyncio.wait_for(dm._download(task, expected_sha256=None), timeout=10) + + assert task.status == "error" + assert task.error + + @pytest.mark.asyncio + async def test_client_is_built_with_finite_timeouts(self, dm, tmp_path): + dest = tmp_path / "out.bin" + task = DownloadTask(id="dl", url="http://example.com/f.bin", dest=dest) + server = _FakeHttpServer(b"payload") + + with patch("tinyagentos.download_manager.httpx.AsyncClient", server): + await dm._download(task, expected_sha256=None) + + assert server.timeouts, "no AsyncClient was constructed" + timeout = server.timeouts[0] + assert isinstance(timeout, httpx.Timeout) + # The read timeout is the one that turns a half-open connection into + # an error; the others keep a dead peer from wedging connect/pool. + assert timeout.read is not None and timeout.read > 0 + assert timeout.connect is not None and timeout.connect > 0 + assert timeout.write is not None and timeout.write > 0 + assert timeout.pool is not None and timeout.pool > 0 + + @pytest.mark.asyncio + async def test_interrupted_download_resumes_from_byte_offset(self, dm, tmp_path): + """Bytes already on disk from an interrupted transfer are asked for + with a Range header instead of being thrown away — a 40 GB model that + died at 39 GB must not restart from zero.""" + body = bytes(i % 256 for i in range(4096)) + dest = tmp_path / "out.bin" + part = tmp_path / "out.bin.part" + part.write_bytes(body[:3000]) + task = DownloadTask(id="dl", url="http://example.com/f.bin", dest=dest) + server = _FakeHttpServer(body) + + with patch("tinyagentos.download_manager.httpx.AsyncClient", server): + await dm._download(task, expected_sha256=hashlib.sha256(body).hexdigest()) + + assert server.requests[0].get("Range") == "bytes=3000-" + assert task.status == "complete", task.error + assert dest.read_bytes() == body + assert not part.exists() + + @pytest.mark.asyncio + async def test_server_ignoring_range_restarts_cleanly(self, dm, tmp_path): + """A mirror that answers 200 to a Range request is sending the whole + file again; appending it to the stub would double the bytes.""" + body = b"abcdefghij" * 100 + dest = tmp_path / "out.bin" + part = tmp_path / "out.bin.part" + part.write_bytes(b"stale bytes") + task = DownloadTask(id="dl", url="http://example.com/f.bin", dest=dest) + server = _FakeHttpServer(body, supports_range=False) + + with patch("tinyagentos.download_manager.httpx.AsyncClient", server): + await dm._download(task, expected_sha256=hashlib.sha256(body).hexdigest()) + + assert task.status == "complete", task.error + assert dest.read_bytes() == body + + @pytest.mark.asyncio + async def test_completed_stub_answered_with_416_restarts_from_zero(self, dm, tmp_path, fast_retries): + """A .part left behind holding every byte of the file gets a 416 to + its Range request. Treating that as a hard 4xx would wedge the model + permanently: it must fall back to a clean full fetch.""" + body = b"complete already" + dest = tmp_path / "out.bin" + part = tmp_path / "out.bin.part" + part.write_bytes(body) + task = DownloadTask(id="dl", url="http://example.com/f.bin", dest=dest) + server = _FakeHttpServer(body) + + with patch("tinyagentos.download_manager.httpx.AsyncClient", server): + await dm._download(task, expected_sha256=hashlib.sha256(body).hexdigest()) + + assert task.status == "complete", task.error + assert dest.read_bytes() == body + + @pytest.mark.asyncio + async def test_failed_download_leaves_no_file_at_destination(self, dm, tmp_path, fast_retries): + """A partial file at the canonical path is read as a present, valid + model by every later "is this installed?" existence check.""" + dest = tmp_path / "out.bin" + task = DownloadTask(id="dl", url="http://example.com/f.bin", dest=dest) + server = _FakeHttpServer(b"y" * 200_000, fail_attempts=99) + + with patch("tinyagentos.download_manager.httpx.AsyncClient", server): + await dm._download(task, expected_sha256=None) + + assert task.status == "error" + assert not dest.exists() + + @pytest.mark.asyncio + async def test_transient_failure_is_retried(self, dm, tmp_path, fast_retries): + """A single dropped connection from a mirror must not kill a + multi-gigabyte pull.""" + body = b"retry me" * 500 + dest = tmp_path / "out.bin" + task = DownloadTask(id="dl", url="http://example.com/f.bin", dest=dest) + server = _FakeHttpServer(body, fail_attempts=1) + + with patch("tinyagentos.download_manager.httpx.AsyncClient", server): + await dm._download(task, expected_sha256=hashlib.sha256(body).hexdigest()) + + assert task.status == "complete", task.error + assert len(server.requests) == 2 + assert dest.read_bytes() == body + + @pytest.mark.asyncio + async def test_retry_resumes_rather_than_restarting(self, dm, tmp_path, fast_retries): + body = b"z" * 200_000 + dest = tmp_path / "out.bin" + task = DownloadTask(id="dl", url="http://example.com/f.bin", dest=dest) + server = _FakeHttpServer(body, fail_attempts=1) + + with patch("tinyagentos.download_manager.httpx.AsyncClient", server): + await dm._download(task, expected_sha256=hashlib.sha256(body).hexdigest()) + + assert task.status == "complete", task.error + assert server.requests[1].get("Range") == "bytes=100000-" + assert dest.read_bytes() == body + + @pytest.mark.asyncio + async def test_server_5xx_is_retried_but_404_is_not(self, dm, tmp_path, fast_retries): + """4xx is the server saying the URL is wrong; retrying it just delays + the error the user needs to see.""" + dest = tmp_path / "out.bin" + task = DownloadTask(id="dl", url="http://example.com/f.bin", dest=dest) + server = _FakeHttpServer( + b"never sent", + fail_attempts=99, + failure=httpx.HTTPStatusError( + "HTTP 404", + request=httpx.Request("GET", "http://example.com/f.bin"), + response=httpx.Response(404), + ), + ) + + with patch("tinyagentos.download_manager.httpx.AsyncClient", server): + await dm._download(task, expected_sha256=None) + + assert task.status == "error" + assert len(server.requests) == 1 + + @pytest.mark.asyncio + async def test_server_5xx_is_retried(self, dm, tmp_path, fast_retries): + """5xx is retried (unlike 4xx above); with_retry's own status branch + handles it, but nothing exercised that path -- a regression that + stopped retrying 5xx would leave every other test here green.""" + body = b"eventually served" + dest = tmp_path / "out.bin" + task = DownloadTask(id="dl", url="http://example.com/f.bin", dest=dest) + server = _FakeHttpServer( + body, + fail_attempts=1, + failure=httpx.HTTPStatusError( + "HTTP 503", + request=httpx.Request("GET", "http://example.com/f.bin"), + response=httpx.Response(503), + ), + ) + + with patch("tinyagentos.download_manager.httpx.AsyncClient", server): + await dm._download(task, expected_sha256=hashlib.sha256(body).hexdigest()) + + assert len(server.requests) > 1 + assert task.status == "complete", task.error + + @pytest.mark.asyncio + async def test_failed_redownload_does_not_delete_an_existing_valid_file( + self, dm, tmp_path, fast_retries + ): + """A re-download of an already-installed model must not destroy the + good copy just because the NEW attempt failed before ever promoting + anything to task.dest. _stream_to_part writes only to the .part + stage file; if with_retry exhausts its attempts, task.dest still + holds whatever a PRIOR successful download put there and that must + survive this attempt's failure.""" + dest = tmp_path / "out.bin" + existing = b"a previously installed, valid model" + dest.write_bytes(existing) + task = DownloadTask(id="dl", url="http://example.com/f.bin", dest=dest) + server = _FakeHttpServer(b"new bytes that never arrive", fail_attempts=99) + + with patch("tinyagentos.download_manager.httpx.AsyncClient", server): + await dm._download(task, expected_sha256=None) + + assert task.status == "error" + assert dest.exists(), "a failed re-download must not delete the existing file" + assert dest.read_bytes() == existing + + +class TestTaskPruning: + """self._tasks grew for the lifetime of the process: every model the user + ever downloaded stayed resident, and /api/models/downloads listed them + all.""" + + @pytest.mark.asyncio + async def test_old_finished_tasks_are_pruned(self, tmp_path): + dm = DownloadManager() + stale = DownloadTask(id="old", url="u", dest=tmp_path / "old.bin") + stale.status = "complete" + stale.completed_at = time.time() - (download_manager.TASK_RETENTION_SECONDS + 60) + dm._tasks["old"] = stale + dm._running["old"] = MagicMock() + + async def _install(): + return {"success": True} + + dm.start_installer_task("fresh", _install()) + await dm._running["fresh"] + + assert "old" not in dm._tasks + assert "old" not in dm._running + assert "fresh" in dm._tasks + + @pytest.mark.asyncio + async def test_recent_and_active_tasks_are_kept(self, tmp_path): + dm = DownloadManager() + recent = DownloadTask(id="recent", url="u", dest=tmp_path / "r.bin") + recent.status = "complete" + recent.completed_at = time.time() + running = DownloadTask(id="running", url="u", dest=tmp_path / "s.bin") + running.status = "downloading" + running.started_at = time.time() - (download_manager.TASK_RETENTION_SECONDS + 60) + dm._tasks["recent"] = recent + dm._tasks["running"] = running + + async def _install(): + return {"success": True} + + dm.start_installer_task("fresh", _install()) + await dm._running["fresh"] + + assert dm.get_progress("recent") is recent + assert dm.get_progress("running") is running diff --git a/tinyagentos/download_manager.py b/tinyagentos/download_manager.py index 2ada83e54..be56dfe1e 100644 --- a/tinyagentos/download_manager.py +++ b/tinyagentos/download_manager.py @@ -9,8 +9,77 @@ import httpx +from tinyagentos.clients.retry import with_retry + logger = logging.getLogger(__name__) +# A read timeout is what turns a half-open connection — a Wi-Fi drop, a NAT +# table eviction, a CDN edge that stops sending — into an error instead of a +# task that sits at status="downloading" forever behind a progress bar frozen +# at 63%. It bounds the gap BETWEEN chunks, not the total transfer, so a slow +# but alive multi-gigabyte pull is unaffected. +DOWNLOAD_TIMEOUT = httpx.Timeout(connect=10.0, read=60.0, write=60.0, pool=10.0) + +# Retry budget for a single transfer. Longer than the inference-client default +# because the thing being retried is a multi-gigabyte pull, not a chat call: +# waiting seconds for a mirror to recover is cheap next to restarting one. +DOWNLOAD_MAX_ATTEMPTS = 4 +DOWNLOAD_RETRY_BASE_DELAY = 1.0 +DOWNLOAD_RETRY_MAX_DELAY = 30.0 + +# How long a finished (complete/error) task stays queryable before it is +# dropped. The models UI polls /api/models/downloads for a while after a +# transfer ends; anything older than this is history nobody reads, and +# keeping it means every model ever downloaded stays resident for the +# lifetime of the process. +TASK_RETENTION_SECONDS = 3600.0 + +_TERMINAL_STATUSES = ("complete", "error") + + +class _RangeRestart(Exception): + """Raised when the server rejects our resume offset with a 416. + + The stage file holds at least as many bytes as the server is willing to + serve — a stale or over-long ``.part``. Restarting from zero recovers; + letting the 416 escape as a 4xx would wedge that model permanently. + """ + + +# Failures worth another attempt. httpx.TransportError covers timeouts, +# connection resets and protocol errors; a bad URL (4xx) is not in it and must +# surface immediately. with_retry adds 5xx responses on top of this tuple. +DOWNLOAD_RETRY_ON = (httpx.TransportError, _RangeRestart) + + +def _hash_prefix(sha, path: Path, length: int) -> None: + """Feed the first ``length`` bytes of ``path`` into ``sha``. + + Runs in a worker thread — the prefix of a resumed model download can be + tens of gigabytes and must never block the event loop. + """ + remaining = length + with open(path, "rb") as f: + while remaining > 0: + chunk = f.read(min(1024 * 1024, remaining)) + if not chunk: + break + sha.update(chunk) + remaining -= len(chunk) + + +def _content_range_total(header: str | None) -> int: + """Full resource size from a ``Content-Range: bytes X-Y/Z`` header. + + Returns 0 when the header is absent or the total is the unknown ``*``, + which leaves _validate_download's size check disabled rather than + comparing against a wrong number. + """ + if not header or "/" not in header: + return 0 + total = header.rsplit("/", 1)[1].strip() + return int(total) if total.isdigit() else 0 + @dataclass class DownloadTask: @@ -81,6 +150,7 @@ def start_download( HuggingFace resolve URLs) that the swarm rides as a correctness fallback. They are passed straight through to the torrent path. """ + self._prune_tasks() task = DownloadTask(id=download_id, url=url, dest=dest) self._tasks[download_id] = task self._running[download_id] = asyncio.create_task( @@ -108,6 +178,7 @@ def start_installer_task(self, download_id: str, coro) -> DownloadTask: that streams incremental progress (e.g. rkllama's ndjson ``/api/pull``) update this task's ``downloaded_bytes``/``total_bytes`` as it goes. """ + self._prune_tasks() task = DownloadTask(id=download_id, url="", dest=Path()) self._tasks[download_id] = task @@ -263,34 +334,35 @@ def _progress(t): await self._download(task, expected_sha256) async def _download(self, task: DownloadTask, expected_sha256: str | None = None): + """HTTP transfer, staged through a ``.part`` file. + + Nothing is ever written to ``task.dest`` until a complete, validated + body is on disk, so a failure can never leave a corrupt weight sitting + at the canonical path where every later "is this model installed?" + existence check would take it for the real thing. The stage file + survives a failure on purpose: it is what the next attempt — this + process or the next boot — resumes from. + """ task.status = "downloading" task.started_at = time.time() - sha = hashlib.sha256() + part = task.dest.with_name(task.dest.name + ".part") + promoted = False try: task.dest.parent.mkdir(parents=True, exist_ok=True) - async with httpx.AsyncClient(timeout=None, follow_redirects=True) as client: - async with client.stream("GET", task.url) as resp: - resp.raise_for_status() - # Content-Length is the size of the ON-THE-WIRE body. When - # the response is content-encoded (gzip/br/deflate/zstd), - # httpx's aiter_bytes() transparently decompresses, so the - # bytes we write to disk are LARGER than Content-Length. - # Treating that as the expected on-disk size would make - # _validate_download flag a perfectly good download as a - # "size mismatch" and delete it, so leave total_bytes at 0 - # (unknown) for encoded responses and rely on the SHA check. - total = resp.headers.get("content-length") - if total and not resp.headers.get("content-encoding"): - task.total_bytes = int(total) - else: - task.total_bytes = 0 - with open(task.dest, "wb") as f: - async for chunk in resp.aiter_bytes(chunk_size=65536): - f.write(chunk) - sha.update(chunk) - task.downloaded_bytes += len(chunk) - error = await self._validate_download(task, expected_sha256, computed_sha256=sha.hexdigest()) + digest = await with_retry( + lambda: self._stream_to_part(task, part), + max_attempts=DOWNLOAD_MAX_ATTEMPTS, + base_delay=DOWNLOAD_RETRY_BASE_DELAY, + max_delay=DOWNLOAD_RETRY_MAX_DELAY, + retry_on=DOWNLOAD_RETRY_ON, + ) + part.replace(task.dest) + promoted = True + error = await self._validate_download(task, expected_sha256, computed_sha256=digest) if error: + # The bytes are wrong, so the promoted file is worthless: this + # attempt just replaced task.dest with a bad copy, and that + # copy (not the now-gone .part) is what has to go. task.dest.unlink(missing_ok=True) task.status = "error" task.error = error @@ -298,6 +370,84 @@ async def _download(self, task: DownloadTask, expected_sha256: str | None = None task.status = "complete" task.completed_at = time.time() except Exception as e: + # Only clean up task.dest when THIS attempt is the one that put a + # (possibly bad) file there. _stream_to_part writes exclusively to + # `part`, so a failure before promotion leaves task.dest exactly as + # it was -- which, for a re-download of an already-installed model, + # is a perfectly good file. Deleting it here would destroy a valid + # install over a transient failure of the NEW attempt. + if promoted: + task.dest.unlink(missing_ok=True) task.status = "error" task.error = str(e) - logger.error(f"Download failed for {task.id}: {e}") + logger.error("Download failed for %s: %s", task.id, e) + + async def _stream_to_part(self, task: DownloadTask, part: Path) -> str: + """One HTTP attempt, appending into the ``.part`` stage file. + + Returns the SHA-256 hex digest of everything the stage file holds. + Whatever an interrupted attempt already wrote is asked for with a + ``Range`` header rather than thrown away, so a 40 GB model that died + at 39 GB does not restart from zero. A server that ignores the header + (answering 200 instead of 206) restarts the file cleanly instead of + concatenating a second copy onto the first. + """ + resume_from = part.stat().st_size if part.exists() else 0 + headers = {"Range": f"bytes={resume_from}-"} if resume_from else {} + sha = hashlib.sha256() + async with httpx.AsyncClient(timeout=DOWNLOAD_TIMEOUT, follow_redirects=True) as client: + async with client.stream("GET", task.url, headers=headers) as resp: + if resp.status_code == 416 and resume_from: + part.unlink(missing_ok=True) + raise _RangeRestart( + f"server rejected resume offset {resume_from}; restarting" + ) + resp.raise_for_status() + resumed = resume_from > 0 and resp.status_code == 206 + if not resumed: + resume_from = 0 + else: + # Re-hash the prefix already on disk so the streamed digest + # covers the whole file, not just this attempt's tail. + await asyncio.to_thread(_hash_prefix, sha, part, resume_from) + # Content-Length is the size of the ON-THE-WIRE body — and on a + # 206 that is only the REMAINING bytes, with the full size in + # the Content-Range total. When the response is content-encoded + # (gzip/br/deflate/zstd), httpx's aiter_bytes() transparently + # decompresses, so the bytes we write to disk are LARGER than + # Content-Length. Treating that as the expected on-disk size + # would make _validate_download flag a perfectly good download + # as a "size mismatch" and delete it, so leave total_bytes at 0 + # (unknown) for encoded responses and rely on the SHA check. + if resp.headers.get("content-encoding"): + task.total_bytes = 0 + elif resumed: + task.total_bytes = _content_range_total(resp.headers.get("content-range")) + else: + total = resp.headers.get("content-length") + task.total_bytes = int(total) if total else 0 + task.downloaded_bytes = resume_from + with open(part, "ab" if resumed else "wb") as f: + async for chunk in resp.aiter_bytes(chunk_size=65536): + f.write(chunk) + sha.update(chunk) + task.downloaded_bytes += len(chunk) + return sha.hexdigest() + + def _prune_tasks(self) -> None: + """Drop finished tasks older than the retention window. + + Without this every download ever started stays in ``_tasks`` (and its + asyncio.Task in ``_running``) for the lifetime of the process, and + /api/models/downloads returns the lot. Tasks still pending or + downloading are never pruned however long they have been running. + """ + cutoff = time.time() - TASK_RETENTION_SECONDS + for download_id, task in list(self._tasks.items()): + if task.status not in _TERMINAL_STATUSES: + continue + finished_at = task.completed_at or task.started_at + if finished_at and finished_at > cutoff: + continue + del self._tasks[download_id] + self._running.pop(download_id, None)