Validate resumed HTTP transfers and cover interrupted downloads

docs/lama-runtime.md

@@ -23,7 +23,7 @@ python -m venv .venv-app
 # Windows PowerShell: .venv-app\Scripts\Activate.ps1
 python -m pip install -r requirements.txt
 python -m pip check
-python -m lama_inpaint download
+python -m lama_inpaint download --verbose
 python remwm.py --help
 ```
 
@@ -54,7 +54,7 @@ An existing valid IOPaint cache is reused without downloading it again. A custom
 `torch.hub.set_dir()` made by unrelated Python code is not consulted.
 
 ```sh
-python -m lama_inpaint download
+python -m lama_inpaint download --verbose
 python -m lama_inpaint download --offline
 python -m lama_inpaint download --cache-dir /path/to/checkpoints
 ```
@@ -67,8 +67,14 @@ can use `LamaInpaint(device="cpu", download=False)` for strict offline loading.
 The artifact URL, exact byte count and SHA-256 are pinned in
 `lama_inpaint/model_store.py`; see [provenance](../THIRD_PARTY_NOTICES.md).
 Cached weights are verified before deserialization. Downloads have read/time/size
-limits, a cross-process lock, a unique temporary file and atomic replacement after
-verification. Interrupted or invalid downloads never become the cached model and
+limits, a cross-process lock, a retained `.big-lama.pt.part` file and atomic
+replacement after verification. Retrying resumes a partial transfer with HTTP
+Range. A server that ignores or rejects Range falls back to a full download.
+Response ranges and declared sizes are checked before changing the partial file.
+A complete partial is verified locally; corrupt or oversized partials are discarded
+on the next attempt. `--verbose` prints bytes, transfer speed, ETA and verification
+status to stderr. Both setup scripts enable this output. Interrupted or invalid
+downloads never become the cached model and
 do not delete an existing artifact. An invalid cache is replaced only after a new
 download passes verification. No user-provided URL or checksum environment variable
 can bypass the pinned artifact check.

lama_inpaint/model_store.py

@@ -4,8 +4,10 @@ import hashlib
 from http.client import HTTPException
 import os
 from pathlib import Path
+import re
 import sys
 import time
+from urllib.error import HTTPError
 from urllib.request import Request, urlopen
 
 from filelock import FileLock, Timeout
@@ -63,7 +65,7 @@ def _is_valid(path):
 def ensure_model(cache_dir=None, *, download=True, progress=None, verbose=False):
     """Return verified weights, downloading atomically when allowed.
 
-    Downloads use a unique temporary file and a per-artifact cross-process lock.
+    Downloads retain a partial file under a per-artifact cross-process lock.
     A failed transfer never replaces an existing file. Offline mode never opens
     a network connection. The URL and expected digest are deliberately fixed.
     """
@@ -133,11 +135,63 @@ def _response_header(response, name):
     return headers.get(name)
 
 
+def _open_transfer(existing, verbose):
+    """Resume when supported; reuse a full response if the server ignores Range."""
+    if existing:
+        _verbose(verbose, f"Resuming from {_format_bytes(existing)} using HTTP Range...")
+        request = Request(MODEL_URL, headers={"Range": f"bytes={existing}-"})
+        try:
+            response = urlopen(request, timeout=DOWNLOAD_TIMEOUT)
+        except HTTPError as exc:
+            if exc.code != 416:
+                exc.close()
+                raise
+            exc.close()
+            _verbose(verbose, "Server rejected the range; requesting the full artifact.")
+        else:
+            if _response_status(response) == 200:
+                _verbose(verbose, "Server ignored Range; restarting from its full response.")
+                return response, 0
+            return response, existing
+    _verbose(verbose, f"Downloading from {MODEL_URL}")
+    try:
+        return urlopen(MODEL_URL, timeout=DOWNLOAD_TIMEOUT), 0
+    except HTTPError as exc:
+        exc.close()
+        raise
+
+
+def _transfer_size(response, existing):
+    """Validate HTTP framing before opening or truncating the retained partial."""
+    status = _response_status(response)
+    expected = MODEL_SIZE
+    if existing:
+        content_range = _response_header(response, "Content-Range") or ""
+        match = re.fullmatch(r"bytes ([0-9]+)-([0-9]+)/([0-9]+)", content_range)
+        if status != 206 or match is None:
+            raise ModelError("Server returned a missing or invalid download range.")
+        start, end, total = map(int, match.groups())
+        if start != existing or not start <= end < total or total != MODEL_SIZE:
+            raise ModelError("Server returned an unexpected download range.")
+        expected = end - start + 1
+    elif status != 200:
+        raise ModelError(f"Unexpected HTTP status for a full download: {status}.")
+    encoding = _response_header(response, "Content-Encoding")
+    if encoding and encoding.lower() != "identity":
+        raise ModelError("Unexpected content encoding in model download.")
+    length = _response_header(response, "Content-Length")
+    if length is not None and (
+        re.fullmatch(r"[0-9]+", length) is None or int(length) != expected
+    ):
+        raise ModelError("Download Content-Length does not match the expected size.")
+    return expected
+
+
 def _download(target, progress=None, *, verbose=False):
     """Download, resume and verify the model before publishing it atomically."""
     partial = target.with_name(f".{target.name}.part")
     started_at = time.monotonic()
-    last_progress = 0.0
+    progress_line = False
 
     try:
         existing = partial.stat().st_size if partial.exists() else 0
@@ -162,43 +216,7 @@ def _download(target, progress=None, *, verbose=False):
             f"Expected size: {_format_bytes(MODEL_SIZE)}",
         )
 
-        response = None
-        resumed = existing > 0
-        if resumed:
-            _verbose(
-                verbose,
-                f"Resuming from {_format_bytes(existing)} "
-                f"({existing / MODEL_SIZE:.1%}) using HTTP Range...",
-            )
-            request = Request(MODEL_URL, headers={"Range": f"bytes={existing}-"})
-            response = urlopen(request, timeout=DOWNLOAD_TIMEOUT)
-            status = _response_status(response)
-            if status != 206:
-                _verbose(
-                    verbose,
-                    f"Server did not accept resume (HTTP {status or 'unknown'}); "
-                    "restarting the transfer.",
-                )
-                response.close()
-                response = None
-                partial.unlink()
-                existing = 0
-                resumed = False
-
-        if response is None:
-            _verbose(verbose, f"Downloading from {MODEL_URL}")
-            response = urlopen(MODEL_URL, timeout=DOWNLOAD_TIMEOUT)
-
-        content_range = _response_header(response, "Content-Range")
-        if resumed and content_range:
-            expected_prefix = f"bytes {existing}-"
-            if not content_range.startswith(expected_prefix):
-                response.close()
-                raise ModelError(
-                    f"Server returned an unexpected range ({content_range}); "
-                    f"partial file retained at {partial} for retry."
-                )
-
+        response, existing = _open_transfer(existing, verbose)
         mode = "ab" if existing else "wb"
         downloaded = existing
         previous_downloaded = downloaded
@@ -206,11 +224,12 @@ def _download(target, progress=None, *, verbose=False):
         last_progress = 0.0
         _verbose(
             verbose,
-            f"Transfer {'resumed' if resumed else 'started'}; "
+            f"Transfer {'resumed' if existing else 'started'}; "
             "interrupted transfers will retain their progress.",
         )
 
         with response:
+            response_end = existing + _transfer_size(response, existing)
             with partial.open(mode) as output:
                 while True:
                     if time.monotonic() - started_at > DOWNLOAD_DEADLINE:
@@ -219,16 +238,18 @@ def _download(target, progress=None, *, verbose=False):
                             f"{_format_bytes(MODEL_SIZE)}; partial file retained at {partial}."
                         )
                     chunk = response.read(DOWNLOAD_CHUNK_SIZE)
+                    if time.monotonic() - started_at > DOWNLOAD_DEADLINE:
+                        raise ModelError("Download deadline reached; partial file retained for retry.")
                     if not chunk:
                         break
-                    output.write(chunk)
-                    downloaded += len(chunk)
-                    if downloaded > MODEL_SIZE:
+                    if downloaded + len(chunk) > response_end:
                         raise ModelError(
                             f"Download exceeds expected size "
-                            f"({_format_bytes(downloaded)} > {_format_bytes(MODEL_SIZE)}); "
+                            f"({_format_bytes(downloaded + len(chunk))} > {_format_bytes(response_end)}); "
                             f"partial file retained at {partial}."
                         )
+                    output.write(chunk)
+                    downloaded += len(chunk)
                     now = time.monotonic()
                     if verbose and (
                         now - last_progress >= PROGRESS_INTERVAL
@@ -247,6 +268,7 @@ def _download(target, progress=None, *, verbose=False):
                             file=sys.stderr,
                             flush=True,
                         )
+                        progress_line = True
                         last_progress = now
                     if progress:
                         progress(
@@ -260,8 +282,9 @@ def _download(target, progress=None, *, verbose=False):
                 output.flush()
                 os.fsync(output.fileno())
 
-        if verbose:
+        if progress_line:
             print(file=sys.stderr, flush=True)
+            progress_line = False
         if downloaded != MODEL_SIZE:
             raise ModelError(
                 f"Incomplete download: {_format_bytes(downloaded)}/"
@@ -273,10 +296,11 @@ def _download(target, progress=None, *, verbose=False):
             verify_stream(stream)
         os.replace(partial, target)
         _verbose(verbose, f"Model verified and saved to {target}")
-    except ModelError:
-        raise
     except (OSError, ValueError, HTTPException) as exc:
         raise ModelError(
             f"Could not download LaMA: {exc}. "
             f"Partial file retained at {partial}; run the command again."
         ) from exc
+    finally:
+        if progress_line:
+            print(file=sys.stderr, flush=True)

tests/test_lama_model_store.py

@@ -10,6 +10,13 @@ import pytest
 from lama_inpaint import model_store as store
 
 
+class Response(io.BytesIO):
+    """A controlled successful HTTP response with an optional header mapping."""
+
+    status = 200
+    headers = {}
+
+
 @pytest.fixture
 def artifact(monkeypatch):
     """Pin verification to tiny controlled bytes so download tests need no real weights."""
@@ -27,7 +34,7 @@ def test_download_then_offline_cache_reuse(tmp_path, monkeypatch, artifact):
         """Record a network attempt and supply the controlled artifact bytes."""
         assert url == store.MODEL_URL and timeout == 30
         calls.append(url)
-        return io.BytesIO(artifact)
+        return Response(artifact)
 
     monkeypatch.setattr(store, "urlopen", response)
     target = store.ensure_model(tmp_path)
@@ -41,7 +48,7 @@ def test_download_then_offline_cache_reuse(tmp_path, monkeypatch, artifact):
 @pytest.mark.parametrize("bad", [b"", b"truncated", b"x" * 200])
 def test_bad_download_never_published(tmp_path, monkeypatch, artifact, bad):
     """Reject bad transfers without publishing them, while retaining resumable data."""
-    monkeypatch.setattr(store, "urlopen", lambda *a, **k: io.BytesIO(bad))
+    monkeypatch.setattr(store, "urlopen", lambda *a, **k: Response(bad))
     with pytest.raises(store.ModelError):
         store.ensure_model(tmp_path)
     assert not (tmp_path / store.MODEL_NAME).exists()
@@ -72,7 +79,7 @@ def test_interrupted_download_preserves_old_file(tmp_path, monkeypatch, artifact
     old = tmp_path / store.MODEL_NAME
     old.write_bytes(b"old invalid file")
 
-    class Interrupted(io.BytesIO):
+    class Interrupted(Response):
         def read(self, size=-1):
             """Supply one partial chunk, then simulate an interrupted connection."""
             if self.tell():
@@ -91,7 +98,7 @@ def test_interrupted_download_resumes_with_http_range(tmp_path, monkeypatch, art
     partial = tmp_path / f".{store.MODEL_NAME}.part"
     partial.write_bytes(artifact[:5])
 
-    class Resumed(io.BytesIO):
+    class Resumed(Response):
         status = 206
         headers = {"Content-Range": f"bytes 5-{len(artifact) - 1}/{len(artifact)}"}
 
@@ -111,7 +118,7 @@ def test_corrupt_cache_replaced_only_after_verification(
     """Replace an invalid existing artifact only with the complete verified fixture."""
     target = tmp_path / store.MODEL_NAME
     target.write_bytes(b"bad")
-    monkeypatch.setattr(store, "urlopen", lambda *a, **k: io.BytesIO(artifact))
+    monkeypatch.setattr(store, "urlopen", lambda *a, **k: Response(artifact))
     assert store.ensure_model(tmp_path).read_bytes() == artifact
 
 
@@ -122,7 +129,7 @@ def test_concurrent_downloads_share_verified_artifact(tmp_path, monkeypatch, art
     def response(*a, **k):
         """Record a network attempt and supply the controlled artifact bytes."""
         calls.append(1)
-        return io.BytesIO(artifact)
+        return Response(artifact)
 
     monkeypatch.setattr(store, "urlopen", response)
     with ThreadPoolExecutor(max_workers=4) as pool:
@@ -142,7 +149,7 @@ def test_torch_cache_environment(monkeypatch, tmp_path):
 
 def test_verification_rewinds_stream(artifact):
     """Leave a verified open stream ready for deserialization from its beginning."""
-    stream = io.BytesIO(artifact)
+    stream = Response(artifact)
     stream.seek(4)
     store.verify_stream(stream)
     assert stream.tell() == 0

tests/test_lama_resume.py

@@ -0,0 +1,236 @@
+"""Exercise resumable transfer boundaries without downloading model weights."""
+
+import hashlib
+import io
+from urllib.error import HTTPError
+
+import pytest
+
+from lama_inpaint import model_store as store
+
+
+class Response(io.BytesIO):
+    """HTTP-like stream whose closure and headers can be checked."""
+
+    def __init__(self, data, status=200, headers=None):
+        super().__init__(data)
+        self.status = status
+        self.headers = headers or {}
+
+
+@pytest.fixture
+def transfer(tmp_path, monkeypatch):
+    """Use a small pinned artifact and retain an existing invalid cache file."""
+    data = b"controlled model fixture"
+    monkeypatch.setattr(store, "MODEL_SIZE", len(data))
+    monkeypatch.setattr(store, "MODEL_SHA256", hashlib.sha256(data).hexdigest())
+    target = tmp_path / store.MODEL_NAME
+    target.write_bytes(b"old cache")
+    partial = tmp_path / f".{store.MODEL_NAME}.part"
+    partial.write_bytes(data[:5])
+    return data, target, partial
+
+
+@pytest.mark.parametrize("content_range", [
+    None, "bytes 0-23/24", "bytes 5-23/99", "bytes 5-999/24",
+    "bytes 5-3/24", "bytes 5-garbage", "bytes 5-23/*",
+])
+def test_invalid_range_preserves_partial(transfer, monkeypatch, content_range):
+    """Never append a response whose range does not describe the pinned artifact."""
+    data, target, partial = transfer
+    headers = {} if content_range is None else {"Content-Range": content_range}
+    response = Response(data[5:], 206, headers)
+    monkeypatch.setattr(store, "urlopen", lambda *a, **k: response)
+    with pytest.raises(store.ModelError, match="range"):
+        store.ensure_model(target.parent)
+    assert partial.read_bytes() == data[:5]
+    assert target.read_bytes() == b"old cache"
+    assert response.closed
+
+
+def test_server_ignoring_range_reuses_full_response(transfer, monkeypatch):
+    """A 200 response restarts safely with one request, without appending old bytes."""
+    data, target, partial = transfer
+    calls = []
+    def open_response(request, **kwargs):
+        calls.append(request)
+        return Response(data, headers={"Content-Length": str(len(data))})
+    monkeypatch.setattr(store, "urlopen", open_response)
+    assert store.ensure_model(target.parent).read_bytes() == data
+    assert len(calls) == 1 and not partial.exists()
+
+
+def test_range_not_satisfiable_retries_from_zero(transfer, monkeypatch):
+    """An HTTP 416 permits a clean full request without accepting unchecked bytes."""
+    data, target, partial = transfer
+    calls = []
+    error_stream = io.BytesIO()
+    def open_response(request, **kwargs):
+        calls.append(request)
+        if len(calls) == 1:
+            raise HTTPError(store.MODEL_URL, 416, "range rejected", {}, error_stream)
+        return Response(data)
+    monkeypatch.setattr(store, "urlopen", open_response)
+    assert store.ensure_model(target.parent).read_bytes() == data
+    assert len(calls) == 2 and error_stream.closed and not partial.exists()
+
+
+@pytest.mark.parametrize("headers", [
+    {"Content-Length": "garbage"}, {"Content-Length": "999"},
+    {"Content-Length": "0"}, {"Content-Encoding": "gzip"},
+])
+def test_invalid_full_headers_keep_partial(transfer, monkeypatch, headers):
+    """Validate a full response before truncating previously downloaded bytes."""
+    data, target, partial = transfer
+    response = Response(data, headers=headers)
+    monkeypatch.setattr(store, "urlopen", lambda *a, **k: response)
+    with pytest.raises(store.ModelError):
+        store.ensure_model(target.parent)
+    assert partial.read_bytes() == data[:5]
+    assert target.read_bytes() == b"old cache" and response.closed
+
+
+def test_oversized_chunk_not_written(transfer, monkeypatch):
+    """Reject excessive bytes before writing them or replacing the old artifact."""
+    data, target, partial = transfer
+    response = Response(data[5:] + b"unexpected", 206,
+                        {"Content-Range": f"bytes 5-{len(data)-1}/{len(data)}"})
+    monkeypatch.setattr(store, "urlopen", lambda *a, **k: response)
+    with pytest.raises(store.ModelError, match="size"):
+        store.ensure_model(target.parent)
+    assert partial.read_bytes() == data[:5]
+    assert target.read_bytes() == b"old cache"
+
+
+def test_deadline_during_last_read_prevents_publication(transfer, monkeypatch):
+    """A final read that crosses the deadline must not publish a model."""
+    data, target, partial = transfer
+    now = [0.0]
+    monkeypatch.setattr(store.time, "monotonic", lambda: now[0])
+    class LateResponse(Response):
+        def read(self, size=-1):
+            chunk = super().read(size)
+            if not chunk:
+                now[0] = store.DOWNLOAD_DEADLINE + 1
+            return chunk
+    response = LateResponse(data)
+    monkeypatch.setattr(store, "urlopen", lambda *a, **k: response)
+    with pytest.raises(store.ModelError, match="deadline"):
+        store.ensure_model(target.parent)
+    assert target.read_bytes() == b"old cache"
+
+
+@pytest.mark.parametrize("partial_data", [b"complete", b"corrupt", b"oversized"])
+def test_complete_partial_reuse_or_repair(transfer, monkeypatch, partial_data):
+    """Publish a verified complete partial offline; redownload unusable partials."""
+    data, target, partial = transfer
+    partial.write_bytes(data if partial_data == b"complete" else
+                        b"x" * (len(data) + (partial_data == b"oversized")))
+    calls = []
+    def open_response(*args, **kwargs):
+        calls.append(1)
+        return Response(data)
+    monkeypatch.setattr(store, "urlopen", open_response)
+    assert store.ensure_model(target.parent).read_bytes() == data
+    assert len(calls) == (partial_data != b"complete")
+    assert not partial.exists()
+
+
+def test_verbose_resume_keeps_structured_progress(transfer, monkeypatch, capsys):
+    """Human progress goes to stderr while callbacks still report absolute bytes."""
+    data, target, partial = transfer
+    response = Response(data[5:], 206,
+                        {"Content-Range": f"bytes 5-{len(data)-1}/{len(data)}"})
+    monkeypatch.setattr(store, "urlopen", lambda *a, **k: response)
+    events = []
+    store.ensure_model(target.parent, verbose=True, progress=events.append)
+    captured = capsys.readouterr()
+    assert not captured.out
+    assert "Resuming" in captured.err and "ETA" in captured.err and "SHA-256" in captured.err
+    assert events[-1]["current"] == len(data) == events[-1]["total"]
+
+
+@pytest.mark.parametrize("resume_policy", ["range", "ignore", "reject"])
+def test_interrupted_real_http_download_resumes(tmp_path, monkeypatch, resume_policy):
+    """Exercise urllib and real HTTP framing across interrupted and resumed calls."""
+    from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+    import threading
+
+    data = b"controlled model fixture"
+    monkeypatch.setattr(store, "MODEL_SIZE", len(data))
+    monkeypatch.setattr(store, "MODEL_SHA256", hashlib.sha256(data).hexdigest())
+    requests = []
+
+    class Handler(BaseHTTPRequestHandler):
+        def log_message(self, *args):
+            """Keep local fixture requests out of test output."""
+
+        def do_GET(self):
+            """Drop the first response, then follow the selected range policy."""
+            requested_range = self.headers.get("Range")
+            requests.append(requested_range)
+            if len(requests) == 1:
+                self.send_response(200)
+                self.send_header("Content-Length", str(len(data)))
+                self.end_headers()
+                self.wfile.write(data[:9])
+                self.wfile.flush()
+                self.close_connection = True
+                return
+            if requested_range and resume_policy == "reject":
+                self.send_response(416)
+                self.send_header("Content-Range", f"bytes */{len(data)}")
+                self.send_header("Content-Length", "0")
+                self.end_headers()
+                return
+            start = 9 if requested_range and resume_policy == "range" else 0
+            self.send_response(206 if start else 200)
+            if start:
+                self.send_header("Content-Range", f"bytes {start}-{len(data)-1}/{len(data)}")
+            self.send_header("Content-Length", str(len(data)-start))
+            self.end_headers()
+            self.wfile.write(data[start:])
+
+    server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    monkeypatch.setattr(store, "MODEL_URL", f"http://127.0.0.1:{server.server_port}/model")
+    try:
+        with pytest.raises(store.ModelError, match="Incomplete"):
+            store.ensure_model(tmp_path)
+        partial = tmp_path / f".{store.MODEL_NAME}.part"
+        assert partial.read_bytes() == data[:9]
+        assert store.ensure_model(tmp_path).read_bytes() == data
+        assert requests == [None, "bytes=9-"] + ([None] if resume_policy == "reject" else [])
+        assert not partial.exists()
+    finally:
+        server.shutdown()
+        server.server_close()
+        thread.join(timeout=5)
+        assert not thread.is_alive()
+
+
+def test_transient_http_error_preserves_partial(transfer, monkeypatch):
+    """A service failure retains progress and closes the HTTP error response."""
+    data, target, partial = transfer
+    error_stream = io.BytesIO()
+    def unavailable(*args, **kwargs):
+        raise HTTPError(store.MODEL_URL, 503, "unavailable", {}, error_stream)
+    monkeypatch.setattr(store, "urlopen", unavailable)
+    with pytest.raises(store.ModelError, match="503"):
+        store.ensure_model(target.parent)
+    assert error_stream.closed and partial.read_bytes() == data[:5]
+    assert target.read_bytes() == b"old cache"
+
+
+
+def test_initial_http_error_closes_response(tmp_path, monkeypatch):
+    """A failed initial request closes its response and creates no partial file."""
+    error_stream = io.BytesIO()
+    def unavailable(*args, **kwargs):
+        raise HTTPError(store.MODEL_URL, 503, "unavailable", {}, error_stream)
+    monkeypatch.setattr(store, "urlopen", unavailable)
+    with pytest.raises(store.ModelError, match="503"):
+        store.ensure_model(tmp_path)
+    assert error_stream.closed
+    assert not (tmp_path / f".{store.MODEL_NAME}.part").exists()