Deploy latest changes from main branch

.env.example

@@ -5,7 +5,14 @@ MODEL_VERSION=v3.0.0
 REDIS_URL=redis://localhost:6379/0
 
 HF_TOKEN=
+# ABMIL: either local path OR Hub (Hub wins when ABMIL_HF_REPO_ID is set)
 MODEL_CHECKPOINT_PATH=/absolute/path/to/production_model.pt
+ABMIL_HF_REPO_ID=
+ABMIL_HF_FILENAME=production_model.pt
+ABMIL_HF_REVISION=
+ABMIL_HF_SUBFOLDER=
+ABMIL_HF_LOCAL_FILES_ONLY=false
+
 DINOV3_MODEL_ID=facebook/dinov3-vitb16-pretrain-lvd1689m
 EMBEDDING_BATCH_SIZE=16
 EMBEDDING_DEVICE=

src/config.py

@@ -38,7 +38,28 @@ class Settings(BaseSettings):
     model_checkpoint_path: str = Field(
         default="",
         alias="MODEL_CHECKPOINT_PATH",
-        description="Absolute path to ABMIL production checkpoint (.pt)",
+        description="Local path to ABMIL .pt when ABMIL_HF_REPO_ID is unset; ignored when Hub repo id is set",
+    )
+
+    abmil_hf_repo_id: Optional[str] = Field(
+        default=None,
+        description="Hugging Face Hub repo id (e.g. org/private-repo). If set, ABMIL checkpoint is fetched via hf_hub_download.",
+    )
+    abmil_hf_filename: str = Field(
+        default="production_model.pt",
+        description="Checkpoint filename inside the Hub repo (top level or under ABMIL_HF_SUBFOLDER)",
+    )
+    abmil_hf_revision: Optional[str] = Field(
+        default=None,
+        description="Optional Hub revision: branch, tag, or commit SHA (pin in production)",
+    )
+    abmil_hf_subfolder: Optional[str] = Field(
+        default=None,
+        description="Optional folder inside the repo containing ABMIL_HF_FILENAME",
+    )
+    abmil_hf_local_files_only: bool = Field(
+        default=False,
+        description="If true, Hub resolution uses only local HF cache (no network)",
     )
 
     dinov3_model_id: str = Field(

src/core/services/abmil_checkpoint.py

@@ -0,0 +1,128 @@
+"""Resolve ABMIL `.pt` checkpoint path from Hugging Face Hub or local filesystem."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Union
+
+from huggingface_hub import hf_hub_download
+from huggingface_hub.errors import EntryNotFoundError, HfHubHTTPError, RepositoryNotFoundError
+
+from ...config import settings
+from ..logging import get_logger
+
+logger = get_logger()
+
+
+def _hub_auth_token() -> Union[bool, str, None]:
+    """Token for hf_hub_download: explicit string, or True to use cached CLI login."""
+    if settings.hf_token and settings.hf_token.strip():
+        return settings.hf_token.strip()
+    env_t = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+    if env_t and env_t.strip():
+        return env_t.strip()
+    return True
+
+
+def resolve_abmil_checkpoint_path() -> str:
+    """
+    Resolve the ABMIL checkpoint file path for ``torch.load``.
+
+    Precedence (mutually exclusive in practice):
+
+    1. **Hugging Face Hub** when ``ABMIL_HF_REPO_ID`` is set: ``hf_hub_download`` caches under
+       ``HF_HOME`` / Hub layout and returns the local cache path.
+    2. **Local file** when ``MODEL_CHECKPOINT_PATH`` is set: must exist and be a regular file.
+
+    Raises ``RuntimeError`` with an actionable message if neither source is valid.
+    """
+    repo_id = (settings.abmil_hf_repo_id or "").strip()
+    if repo_id:
+        filename = (settings.abmil_hf_filename or "").strip()
+        if not filename:
+            raise RuntimeError(
+                "ABMIL_HF_FILENAME must be set to a non-empty checkpoint filename when ABMIL_HF_REPO_ID is set.",
+            )
+
+        if (settings.model_checkpoint_path or "").strip():
+            logger.info(
+                "abmil_checkpoint_source_hub",
+                repo_id=repo_id,
+                filename=filename,
+                note="MODEL_CHECKPOINT_PATH is ignored while ABMIL_HF_REPO_ID is set.",
+            )
+
+        token = _hub_auth_token()
+        if token is True:
+            logger.warning(
+                "abmil_hf_download_implicit_token",
+                repo_id=repo_id,
+                message="No HF_TOKEN / hf_token in environment; using cached Hub credentials if any.",
+            )
+
+        subfolder = (settings.abmil_hf_subfolder or "").strip() or None
+        revision = (settings.abmil_hf_revision or "").strip() or None
+
+        try:
+            path = hf_hub_download(
+                repo_id=repo_id,
+                filename=filename,
+                subfolder=subfolder,
+                revision=revision,
+                token=token,
+                local_files_only=settings.abmil_hf_local_files_only,
+            )
+        except RepositoryNotFoundError as exc:
+            raise RuntimeError(
+                f"ABMIL Hub repository not found or no access: {repo_id!r}. "
+                "Verify ABMIL_HF_REPO_ID, HF_TOKEN read access, and network.",
+            ) from exc
+        except EntryNotFoundError as exc:
+            raise RuntimeError(
+                f"ABMIL checkpoint file not in Hub repo {repo_id!r}: {filename!r}"
+                + (f" (subfolder={subfolder!r})" if subfolder else "")
+                + ". Check ABMIL_HF_FILENAME / ABMIL_HF_SUBFOLDER.",
+            ) from exc
+        except HfHubHTTPError as exc:
+            raise RuntimeError(
+                f"Hugging Face Hub error while downloading ABMIL checkpoint: {exc}",
+            ) from exc
+        except OSError as exc:
+            raise RuntimeError(
+                f"Filesystem error while caching ABMIL checkpoint from Hub: {exc}",
+            ) from exc
+
+        resolved = Path(path)
+        if not resolved.is_file():
+            raise RuntimeError(f"Hub download did not produce a readable file: {path}")
+
+        logger.info(
+            "abmil_checkpoint_resolved_from_hub",
+            repo_id=repo_id,
+            filename=filename,
+            revision=revision,
+            local_files_only=settings.abmil_hf_local_files_only,
+            resolved_path=str(resolved),
+        )
+        return str(resolved)
+
+    local = (settings.model_checkpoint_path or "").strip()
+    if not local:
+        raise RuntimeError(
+            "ABMIL checkpoint not configured: set ABMIL_HF_REPO_ID and ABMIL_HF_FILENAME "
+            "(Hugging Face Hub), or set MODEL_CHECKPOINT_PATH to a local .pt file.",
+        )
+
+    p = Path(local).expanduser()
+    try:
+        p = p.resolve()
+    except OSError:
+        pass
+    if not p.is_file():
+        raise RuntimeError(
+            f"MODEL_CHECKPOINT_PATH is set but file not found or not a file: {local}",
+        )
+
+    logger.info("abmil_checkpoint_resolved_local", resolved_path=str(p))
+    return str(p)

src/core/services/ml_inference.py

@@ -16,6 +16,7 @@ from starlette.concurrency import run_in_threadpool
 from ...config import settings
 from ..exceptions import InferenceServiceError, InvalidImageError, UnsupportedImageFormatError
 from ..logging import get_logger
+from .abmil_checkpoint import resolve_abmil_checkpoint_path
 from .abmil_model import ABMILInferenceEngine
 from .dino_embedding import DINOv3Embedder
 from .preprocessing import extract_tiles_for_inference
@@ -32,11 +33,6 @@ DINO_EMBEDDER: Optional[DINOv3Embedder] = None
 ABMIL_ENGINE: Optional[ABMILInferenceEngine] = None
 
 
-def _assert_model_checkpoint_configured() -> None:
-    if not settings.model_checkpoint_path:
-        raise RuntimeError("MODEL_CHECKPOINT_PATH is required for ABMIL inference.")
-
-
 def _configure_torch_cpu_threads() -> None:
     """Optional CPU parallelism tuning (set before heavy torch work). See PyTorch threading docs."""
     try:
@@ -65,7 +61,7 @@ def initialize_models() -> None:
         logger.info("Inference models already initialized")
         return
 
-    _assert_model_checkpoint_configured()
+    abmil_checkpoint_path = resolve_abmil_checkpoint_path()
     _configure_torch_cpu_threads()
 
     embedder = DINOv3Embedder(
@@ -79,7 +75,7 @@ def initialize_models() -> None:
     embedder.load()
 
     engine = ABMILInferenceEngine(
-        checkpoint_path=settings.model_checkpoint_path,
+        checkpoint_path=abmil_checkpoint_path,
         device=settings.embedding_device,
     )
     engine.load()

tests/test_abmil_checkpoint.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from src.config import settings
+from src.core.services.abmil_checkpoint import resolve_abmil_checkpoint_path
+
+
+def test_resolve_abmil_checkpoint_local_file(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    pt = tmp_path / "production_model.pt"
+    pt.write_bytes(b"not-a-real-checkpoint")
+
+    monkeypatch.setattr(settings, "abmil_hf_repo_id", None)
+    monkeypatch.setattr(settings, "model_checkpoint_path", str(pt))
+
+    resolved = resolve_abmil_checkpoint_path()
+    assert Path(resolved) == pt.resolve()
+
+
+def test_resolve_abmil_checkpoint_local_missing(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(settings, "abmil_hf_repo_id", None)
+    monkeypatch.setattr(settings, "model_checkpoint_path", "/nonexistent/path/model.pt")
+
+    with pytest.raises(RuntimeError, match="MODEL_CHECKPOINT_PATH"):
+        resolve_abmil_checkpoint_path()
+
+
+def test_resolve_abmil_checkpoint_hub(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    cached = tmp_path / "from_hub.pt"
+    cached.write_bytes(b"x")
+
+    monkeypatch.setattr(settings, "abmil_hf_repo_id", "org/private-abmil")
+    monkeypatch.setattr(settings, "abmil_hf_filename", "weights.pt")
+    monkeypatch.setattr(settings, "abmil_hf_revision", "v1.0.0")
+    monkeypatch.setattr(settings, "abmil_hf_subfolder", "checkpoints")
+    monkeypatch.setattr(settings, "abmil_hf_local_files_only", False)
+    monkeypatch.setattr(settings, "model_checkpoint_path", "/ignored/local.pt")
+
+    captured: dict = {}
+
+    def fake_hf_hub_download(**kwargs: object) -> str:
+        captured.update(kwargs)
+        return str(cached)
+
+    monkeypatch.setattr(
+        "src.core.services.abmil_checkpoint.hf_hub_download",
+        fake_hf_hub_download,
+    )
+
+    out = resolve_abmil_checkpoint_path()
+    assert out == str(cached)
+    assert captured["repo_id"] == "org/private-abmil"
+    assert captured["filename"] == "weights.pt"
+    assert captured["revision"] == "v1.0.0"
+    assert captured["subfolder"] == "checkpoints"
+    assert captured["local_files_only"] is False
+
+
+def test_resolve_abmil_checkpoint_not_configured(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(settings, "abmil_hf_repo_id", None)
+    monkeypatch.setattr(settings, "model_checkpoint_path", "")
+
+    with pytest.raises(RuntimeError, match="ABMIL checkpoint not configured"):
+        resolve_abmil_checkpoint_path()
+
+
+def test_resolve_abmil_checkpoint_hub_empty_filename(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(settings, "abmil_hf_repo_id", "org/repo")
+    monkeypatch.setattr(settings, "abmil_hf_filename", "  ")
+
+    with pytest.raises(RuntimeError, match="ABMIL_HF_FILENAME"):
+        resolve_abmil_checkpoint_path()

tests/test_inference_smoke.py

@@ -62,6 +62,7 @@ def test_initialize_models_requires_checkpoint(monkeypatch: pytest.MonkeyPatch)
     monkeypatch.setattr(ml_inference, "DINO_EMBEDDER", None)
     monkeypatch.setattr(ml_inference, "ABMIL_ENGINE", None)
     monkeypatch.setattr(ml_inference.settings, "model_checkpoint_path", "")
+    monkeypatch.setattr(ml_inference.settings, "abmil_hf_repo_id", None)
 
-    with pytest.raises(RuntimeError, match="MODEL_CHECKPOINT_PATH"):
+    with pytest.raises(RuntimeError, match="ABMIL checkpoint not configured"):
         ml_inference.initialize_models()