fix(models): race-free lazy encoder load; timeout and length check for model downloads

docs/semantic-models.md

@@ -54,6 +54,11 @@ Models live in `~/.freshdata/models/<model-id>/`; override the directory with
 > checksums as artifacts are published; a pinned checksum that does not match
 > is refused at download *and* at load time.
 
+Downloads use a network timeout of 60 seconds per connection attempt and
+socket read; override it with `FRESHDATA_MODEL_TIMEOUT` (a positive number of
+seconds). A transfer that ends before the server's advertised `Content-Length`
+raises `OSError` and nothing is installed; the partial `.part` file is removed.
+
 ### Air-gapped installs
 
 Copy the model files into the model directory by hand:

src/freshdata/models/download.py

@@ -8,8 +8,8 @@
 
 from __future__ import annotations
 
+import math
 import os
-import shutil
 import tempfile
 import urllib.request
 from pathlib import Path
@@ -20,16 +20,55 @@
 _URL_BASE_ENV = "FRESHDATA_MODEL_URL_BASE"
 #: Empty until official artifacts are hosted; see docs/semantic-models.md.
 _DEFAULT_URL_BASE = ""
+_TIMEOUT_ENV = "FRESHDATA_MODEL_TIMEOUT"
+_DEFAULT_TIMEOUT = "60"
+_CHUNK_SIZE = 1024 * 1024
 
 
 def _url_base() -> str:
     return os.environ.get(_URL_BASE_ENV, _DEFAULT_URL_BASE)
 
 
+def _timeout() -> float:
+    """Network timeout in seconds (``FRESHDATA_MODEL_TIMEOUT``, default 60)."""
+    raw = os.environ.get(_TIMEOUT_ENV, _DEFAULT_TIMEOUT)
+    try:
+        value = float(raw)
+    except ValueError:
+        value = math.nan
+    if not (math.isfinite(value) and value > 0):
+        raise ValueError(f"{_TIMEOUT_ENV} must be a positive number of seconds, got {raw!r}")
+    return value
+
+
+def _content_length(response: object) -> int | None:
+    headers = getattr(response, "headers", None)
+    raw = headers.get("Content-Length") if headers is not None else None
+    try:
+        return int(raw) if raw is not None else None
+    except ValueError:
+        return None
+
+
 def _fetch(url: str, dest: Path) -> None:
-    """Download ``url`` to ``dest``. The single network seam (mocked in tests)."""
-    with urllib.request.urlopen(url) as response, dest.open("wb") as fh:  # noqa: S310
-        shutil.copyfileobj(response, fh)
+    """Download ``url`` to ``dest``. The single network seam (mocked in tests).
+
+    Raises :class:`OSError` when the transfer ends before the advertised
+    ``Content-Length`` (the caller discards the partial file).
+    """
+    timeout = _timeout()
+    with urllib.request.urlopen(url, timeout=timeout) as response:  # noqa: S310
+        expected = _content_length(response)
+        written = 0
+        with dest.open("wb") as fh:
+            while True:
+                chunk = response.read(_CHUNK_SIZE)
+                if not chunk:
+                    break
+                fh.write(chunk)
+                written += len(chunk)
+    if expected is not None and written != expected:
+        raise OSError(f"incomplete download of {url}: got {written} of {expected} bytes")
 
 
 def pull(model_id: str, *, force: bool = False) -> Path:

src/freshdata/models/runtime.py

@@ -53,26 +53,34 @@ def __init__(self, model_id: str) -> None:
         self._tokenizer: Any = None
         self._lock = threading.Lock()
 
-    def _load(self) -> None:  # pragma: no cover - requires real model files
-        if self._session is not None:
+    def _loaded(self) -> bool:
+        return self._session is not None and self._tokenizer is not None
+
+    def _load(self) -> None:
+        # Double-checked locking. Both objects are built into locals and
+        # published tokenizer-first, session-last, so a concurrent caller on
+        # the unlocked fast path never observes a half-loaded encoder.
+        if self._loaded():
             return
         with self._lock:
-            if self._session is not None:
+            if self._loaded():
                 return
             ort = require_onnxruntime()
             tokenizers = require_tokenizers()
             verify(self.model_id)  # refuses on pinned-checksum mismatch
             base = model_dir() / self.model_id
             options = ort.SessionOptions()
             options.intra_op_num_threads = min(4, os.cpu_count() or 1)
-            self._session = ort.InferenceSession(
+            session = ort.InferenceSession(
                 str(base / "model.onnx"),
                 sess_options=options,
                 providers=["CPUExecutionProvider"],
             )
-            self._tokenizer = tokenizers.Tokenizer.from_file(str(base / "tokenizer.json"))
+            tokenizer = tokenizers.Tokenizer.from_file(str(base / "tokenizer.json"))
+            self._tokenizer = tokenizer
+            self._session = session
 
-    def encode_texts(self, texts: Sequence[str]) -> np.ndarray:  # pragma: no cover
+    def encode_texts(self, texts: Sequence[str]) -> np.ndarray:
         self._load()
         chunks: list[np.ndarray] = []
         for start in range(0, len(texts), _BATCH_SIZE):

tests/test_models_download_robustness.py

@@ -0,0 +1,277 @@
+"""Model download robustness and race-free lazy encoder loading.
+
+Covers #340 (``OnnxEncoder._load`` must never expose a half-loaded encoder to
+a concurrent caller) and #341 (``fd.models.pull`` must pass a network timeout
+and refuse truncated transfers). Everything is offline: ``onnxruntime`` and
+``tokenizers`` are replaced with fakes, and ``urllib.request.urlopen`` with an
+in-memory response.
+"""
+
+from __future__ import annotations
+
+import io
+import sys
+import threading
+import types
+import urllib.request
+from email.message import Message
+
+import numpy as np
+import pytest
+
+import freshdata as fd
+from freshdata.models import download as dl
+from freshdata.models import runtime
+
+# --------------------------------------------------------------------------
+# #340: lazy encoder load race
+# --------------------------------------------------------------------------
+
+
+class _Encoding:
+    ids = [1, 2]
+    attention_mask = [1, 1]
+
+
+class _FakeTokenizer:
+    def encode_batch(self, batch):
+        return [_Encoding() for _ in batch]
+
+
+class _FakeSession:
+    def run(self, _outputs, feeds):
+        return [np.ones(feeds["input_ids"].shape + (4,), dtype=np.float32)]
+
+
+@pytest.fixture
+def fake_runtime(monkeypatch, tmp_path):
+    """Install fake onnxruntime/tokenizers; ``hooks`` controls from_file."""
+    monkeypatch.setenv("FRESHDATA_MODEL_DIR", str(tmp_path))
+    hooks: dict = {"from_file": lambda: None, "from_file_calls": 0, "session_calls": 0}
+
+    ort = types.ModuleType("onnxruntime")
+    tok = types.ModuleType("tokenizers")
+
+    class SessionOptions:
+        intra_op_num_threads = 1
+
+    def inference_session(*_args, **_kwargs):
+        hooks["session_calls"] += 1
+        return _FakeSession()
+
+    class Tokenizer:
+        @staticmethod
+        def from_file(_path):
+            hooks["from_file_calls"] += 1
+            hooks["from_file"]()
+            return _FakeTokenizer()
+
+    ort.SessionOptions = SessionOptions
+    ort.InferenceSession = inference_session
+    tok.Tokenizer = Tokenizer
+    monkeypatch.setitem(sys.modules, "onnxruntime", ort)
+    monkeypatch.setitem(sys.modules, "tokenizers", tok)
+    monkeypatch.setattr(runtime, "verify", lambda model_id: False)
+    return hooks
+
+
+def test_concurrent_first_use_waits_for_complete_load(fake_runtime):
+    started = threading.Event()
+    release = threading.Event()
+
+    def slow_from_file():
+        started.set()
+        assert release.wait(10), "test did not release the tokenizer load"
+
+    fake_runtime["from_file"] = slow_from_file
+    enc = runtime.OnnxEncoder("fd-col-encoder-v1")
+    errors: list[str] = []
+    results: list[np.ndarray] = []
+
+    def work():
+        try:
+            results.append(enc.encode_texts(["a"]))
+        except Exception as exc:  # noqa: BLE001 - surfaced via assertion
+            errors.append(repr(exc))
+
+    first = threading.Thread(target=work)
+    first.start()
+    assert started.wait(10)
+    # The first thread is now inside Tokenizer.from_file holding the lock.
+    second = threading.Thread(target=work)
+    second.start()
+    second.join(0.3)
+    # Before the fix the second thread saw ``_session`` set and crashed on a
+    # None tokenizer; now it must still be waiting for the load to finish.
+    assert second.is_alive(), errors
+    assert errors == []
+    release.set()
+    first.join(10)
+    second.join(10)
+    assert not first.is_alive() and not second.is_alive()
+    assert errors == []
+    assert len(results) == 2
+    assert all(r.shape == (1, 4) for r in results)
+    assert fake_runtime["from_file_calls"] == 1
+    assert fake_runtime["session_calls"] == 1
+
+
+def test_many_threads_first_use_load_once(fake_runtime):
+    n = 8
+    barrier = threading.Barrier(n)
+    enc = runtime.OnnxEncoder("fd-col-encoder-v1")
+    errors: list[str] = []
+
+    def work():
+        try:
+            barrier.wait(10)
+            enc.encode_texts(["a", "b"])
+        except Exception as exc:  # noqa: BLE001 - surfaced via assertion
+            errors.append(repr(exc))
+
+    threads = [threading.Thread(target=work) for _ in range(n)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join(10)
+    assert errors == []
+    assert fake_runtime["from_file_calls"] == 1
+    assert fake_runtime["session_calls"] == 1
+
+
+def test_fast_path_requires_both_session_and_tokenizer(fake_runtime):
+    enc = runtime.OnnxEncoder("fd-col-encoder-v1")
+    enc._session = _FakeSession()  # half-loaded: session without tokenizer
+    enc._load()
+    assert enc._tokenizer is not None
+    assert fake_runtime["from_file_calls"] == 1
+
+    enc._load()  # fully loaded: fast path, no reload
+    assert fake_runtime["from_file_calls"] == 1
+
+
+def test_failed_tokenizer_load_leaves_encoder_unloaded(fake_runtime):
+    def boom():
+        raise RuntimeError("tokenizer.json unreadable")
+
+    fake_runtime["from_file"] = boom
+    enc = runtime.OnnxEncoder("fd-col-encoder-v1")
+    with pytest.raises(RuntimeError, match="unreadable"):
+        enc.encode_texts(["a"])
+    assert enc._session is None
+    assert enc._tokenizer is None
+
+    fake_runtime["from_file"] = lambda: None
+    assert enc.encode_texts(["a"]).shape == (1, 4)
+
+
+# --------------------------------------------------------------------------
+# #341: download timeout and truncation detection
+# --------------------------------------------------------------------------
+
+
+class _FakeResponse:
+    def __init__(self, body: bytes, content_length: int | None) -> None:
+        self._stream = io.BytesIO(body)
+        self.headers = Message()
+        if content_length is not None:
+            self.headers["Content-Length"] = str(content_length)
+
+    def read(self, amt: int = -1) -> bytes:
+        return self._stream.read(amt)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_exc):
+        return False
+
+
+@pytest.fixture
+def model_home(tmp_path, monkeypatch):
+    monkeypatch.setenv("FRESHDATA_MODEL_DIR", str(tmp_path))
+    monkeypatch.setenv("FRESHDATA_MODEL_URL_BASE", "https://example.test/models")
+    monkeypatch.delenv("FRESHDATA_MODEL_TIMEOUT", raising=False)
+    return tmp_path
+
+
+def _install_urlopen(monkeypatch, body: bytes, content_length: int | None):
+    calls: list[dict] = []
+
+    def fake_urlopen(url, *args, **kwargs):
+        calls.append({"url": url, "args": args, "kwargs": kwargs})
+        return _FakeResponse(body, content_length)
+
+    monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
+    return calls
+
+
+def test_truncated_download_raises_and_installs_nothing(model_home, monkeypatch):
+    _install_urlopen(monkeypatch, b"X" * 1000, content_length=1_000_000)
+    with pytest.raises(OSError, match=r"incomplete download of .*got 1000 of 1000000 bytes"):
+        fd.models.pull("fd-intent-v1")
+    target = model_home / "fd-intent-v1"
+    assert not (target / "model.onnx").exists()
+    assert not list(target.glob("*.part"))
+    assert fd.models.status()["fd-intent-v1"]["installed"] is False
+
+
+def test_fetch_truncated_leaves_partial_for_caller_cleanup(tmp_path, monkeypatch):
+    monkeypatch.delenv("FRESHDATA_MODEL_TIMEOUT", raising=False)
+    _install_urlopen(monkeypatch, b"abc", content_length=10)
+    dest = tmp_path / "x.part"
+    with pytest.raises(OSError, match="got 3 of 10 bytes"):
+        dl._fetch("https://example.test/m/x", dest)
+    assert dest.read_bytes() == b"abc"  # pull()'s finally removes it
+
+
+def test_complete_download_installs(model_home, monkeypatch):
+    payload = b"model-bytes" * 500_000  # larger than one read chunk
+    calls = _install_urlopen(monkeypatch, payload, content_length=len(payload))
+    primary = fd.models.pull("fd-intent-v1")
+    assert primary.read_bytes() == payload
+    assert not list(primary.parent.glob("*.part"))
+    assert [c["url"] for c in calls] == ["https://example.test/models/fd-intent-v1/model.onnx"]
+
+
+def test_download_without_content_length_is_accepted(model_home, monkeypatch):
+    _install_urlopen(monkeypatch, b"payload", content_length=None)
+    primary = fd.models.pull("fd-intent-v1")
+    assert primary.read_bytes() == b"payload"
+
+
+def test_default_timeout_is_passed(model_home, monkeypatch):
+    calls = _install_urlopen(monkeypatch, b"ok", content_length=2)
+    fd.models.pull("fd-intent-v1")
+    assert calls[0]["kwargs"]["timeout"] == 60.0
+
+
+def test_timeout_env_override_is_passed(model_home, monkeypatch):
+    monkeypatch.setenv("FRESHDATA_MODEL_TIMEOUT", "2.5")
+    calls = _install_urlopen(monkeypatch, b"ok", content_length=2)
+    fd.models.pull("fd-intent-v1")
+    assert calls[0]["kwargs"]["timeout"] == 2.5
+
+
+@pytest.mark.parametrize("raw", ["abc", "", "0", "-1", "nan", "inf"])
+def test_invalid_timeout_env_raises_value_error(model_home, monkeypatch, raw):
+    monkeypatch.setenv("FRESHDATA_MODEL_TIMEOUT", raw)
+    calls = _install_urlopen(monkeypatch, b"ok", content_length=2)
+    with pytest.raises(ValueError, match="FRESHDATA_MODEL_TIMEOUT must be a positive number"):
+        fd.models.pull("fd-intent-v1")
+    assert calls == []
+    target = model_home / "fd-intent-v1"
+    assert not (target / "model.onnx").exists()
+    assert not list(target.glob("*.part"))
+
+
+def test_timeout_error_discards_partial(model_home, monkeypatch):
+    def stalled(url, *args, **kwargs):
+        raise TimeoutError("timed out")
+
+    monkeypatch.setattr(urllib.request, "urlopen", stalled)
+    with pytest.raises(TimeoutError):
+        fd.models.pull("fd-intent-v1")
+    target = model_home / "fd-intent-v1"
+    assert not (target / "model.onnx").exists()
+    assert not list(target.glob("*.part"))