support hf artifact URIs for API inference runtime

Author: Dieg0Code

pyproject.toml

@@ -56,6 +56,7 @@ quote-style = "double"
 
 [dependency-groups]
 api = [
+    "huggingface-hub>=1.4.1",
     "fastapi>=0.121.1",
     "python-jose[cryptography]>=3.5.0",
     "python-dotenv>=1.2.1",

src/api/deps/inference.py

@@ -5,6 +5,7 @@
 from fastapi import HTTPException, status
 
 from api.config import Settings, get_settings
+from api.inference_artifacts import resolve_artifact_uri
 from inference.service import InferenceService
 
 
@@ -30,9 +31,14 @@ def _build_inference_service(
 def get_inference_service_dep() -> InferenceService:
     settings: Settings = get_settings()
     try:
+        resolved_checkpoint = resolve_artifact_uri(settings.model_checkpoint_path)
+        resolved_onnx = resolve_artifact_uri(settings.model_onnx_path)
+
+        # InferenceService expects a checkpoint path value, but ONNX-only runtime is valid.
+        checkpoint_arg = resolved_checkpoint or "__missing_checkpoint__.ckpt"
         return _build_inference_service(
-            checkpoint_path=settings.model_checkpoint_path,
-            onnx_path=settings.model_onnx_path,
+            checkpoint_path=checkpoint_arg,
+            onnx_path=resolved_onnx or "",
             device=settings.inference_device,
             mcts_sims=settings.inference_mcts_sims,
             c_puct=settings.inference_c_puct,

src/api/inference_artifacts.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import importlib
+import os
+from functools import lru_cache
+from urllib.parse import parse_qs, unquote, urlparse
+
+
+def _normalize_local_artifact_path(uri: str) -> str:
+    parsed = urlparse(uri)
+    if parsed.scheme == "file":
+        if parsed.netloc not in {"", "localhost"}:
+            raise ValueError(f"Unsupported artifact URI host: {uri}")
+        path = unquote(parsed.path)
+        if path == "":
+            raise ValueError(f"Invalid artifact URI (empty path): {uri}")
+        return path
+    if parsed.scheme == "":
+        return uri
+    raise ValueError(
+        f"Unsupported artifact URI scheme '{parsed.scheme}'. Use local paths, file://, or hf:// URIs."
+    )
+
+
+def _parse_hf_uri(
+    uri: str,
+    *,
+    default_repo_id: str | None = None,
+    default_revision: str | None = None,
+) -> tuple[str, str, str | None]:
+    parsed = urlparse(uri)
+    if parsed.scheme != "hf":
+        raise ValueError(f"Invalid HF URI: {uri}")
+
+    # Expected format: hf://<owner>/<repo>/<path/to/file>[?revision=<rev>]
+    payload = f"{parsed.netloc}{parsed.path}".lstrip("/")
+    segments = [seg for seg in payload.split("/") if seg]
+    if len(segments) >= 3:
+        repo_id = f"{segments[0]}/{segments[1]}"
+        filename = "/".join(segments[2:])
+    elif len(segments) >= 1 and default_repo_id:
+        repo_id = default_repo_id
+        filename = "/".join(segments)
+    else:
+        raise ValueError(
+            f"Invalid HF URI '{uri}'. Expected hf://<owner>/<repo>/<artifact_path>."
+        )
+
+    if filename == "":
+        raise ValueError(f"Invalid HF URI '{uri}': missing artifact path.")
+
+    query = parse_qs(parsed.query)
+    revision = query.get("revision", [None])[0] or default_revision
+    return repo_id, filename, revision
+
+
+@lru_cache(maxsize=128)
+def _download_hf_artifact(
+    *,
+    repo_id: str,
+    filename: str,
+    revision: str | None,
+) -> str:
+    try:
+        hf_module = importlib.import_module("huggingface_hub")
+    except ModuleNotFoundError as exc:
+        raise ValueError(
+            "huggingface-hub is required for hf:// artifacts. "
+            "Install it in the API runtime dependencies."
+        ) from exc
+
+    hf_hub_download = getattr(hf_module, "hf_hub_download", None)
+    if hf_hub_download is None:
+        raise ValueError("huggingface_hub.hf_hub_download is unavailable in this runtime.")
+
+    token = os.getenv("HF_TOKEN")
+    return str(
+        hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            revision=revision,
+            token=token if token else None,
+        )
+    )
+
+
+def resolve_artifact_uri(
+    uri: str | None,
+    *,
+    default_repo_id: str | None = None,
+    default_revision: str | None = None,
+) -> str | None:
+    if uri is None:
+        return None
+    cleaned = uri.strip()
+    if cleaned == "":
+        return None
+
+    parsed = urlparse(cleaned)
+    if parsed.scheme in {"", "file"}:
+        return _normalize_local_artifact_path(cleaned)
+    if parsed.scheme == "hf":
+        repo_id, filename, revision = _parse_hf_uri(
+            cleaned,
+            default_repo_id=default_repo_id,
+            default_revision=default_revision,
+        )
+        return _download_hf_artifact(
+            repo_id=repo_id,
+            filename=filename,
+            revision=revision,
+        )
+    raise ValueError(
+        f"Unsupported artifact URI scheme '{parsed.scheme}'. Use local paths, file://, or hf:// URIs."
+    )

src/api/modules/matches/model_runtime.py

@@ -1,35 +1,12 @@
 from __future__ import annotations
 
 from functools import lru_cache
-from urllib.parse import unquote, urlparse
 
 from api.db.models import ModelVersion
+from api.inference_artifacts import resolve_artifact_uri
 from inference.service import InferenceService
 
 
-def _normalize_local_artifact_path(uri: str | None) -> str | None:
-    if uri is None:
-        return None
-    cleaned = uri.strip()
-    if cleaned == "":
-        return None
-
-    parsed = urlparse(cleaned)
-    if parsed.scheme in {"", "file"}:
-        if parsed.scheme == "file":
-            if parsed.netloc not in {"", "localhost"}:
-                raise ValueError(f"Unsupported artifact URI host: {cleaned}")
-            path = unquote(parsed.path)
-            if path == "":
-                raise ValueError(f"Invalid artifact URI (empty path): {cleaned}")
-            return path
-        return cleaned
-
-    raise ValueError(
-        f"Unsupported artifact URI scheme '{parsed.scheme}'. Use local paths or file:// URIs."
-    )
-
-
 def _runtime_config_from_base(base_service: InferenceService | None) -> tuple[str, int, float, bool]:
     if base_service is None:
         return "auto", 160, 1.5, True
@@ -65,11 +42,19 @@ def resolve_model_inference_service(
     version: ModelVersion,
     base_service: InferenceService | None,
 ) -> InferenceService:
-    checkpoint_path = _normalize_local_artifact_path(version.checkpoint_uri)
-    onnx_path = _normalize_local_artifact_path(version.onnx_uri)
+    checkpoint_path = resolve_artifact_uri(
+        version.checkpoint_uri,
+        default_repo_id=version.hf_repo_id,
+        default_revision=version.hf_revision,
+    )
+    onnx_path = resolve_artifact_uri(
+        version.onnx_uri,
+        default_repo_id=version.hf_repo_id,
+        default_revision=version.hf_revision,
+    )
     if checkpoint_path is None and onnx_path is None:
         raise ValueError(
-            f"Model version '{version.name}' has no local checkpoint_uri/onnx_uri configured."
+            f"Model version '{version.name}' has no usable checkpoint_uri/onnx_uri configured."
         )
 
     # Keep runtime knobs aligned with the API default inference service so

tests/test_api_inference_artifacts.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import sys
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
+
+from api.inference_artifacts import resolve_artifact_uri
+
+
+class TestApiInferenceArtifacts(unittest.TestCase):
+    def test_resolve_local_path_returns_same_value(self) -> None:
+        path = resolve_artifact_uri("checkpoints/model_iter_039.pt")
+        self.assertEqual(path, "checkpoints/model_iter_039.pt")
+
+    def test_resolve_file_uri_returns_local_path(self) -> None:
+        path = resolve_artifact_uri("file:///var/lib/ataxx/model.ckpt")
+        self.assertEqual(path, "/var/lib/ataxx/model.ckpt")
+
+    @patch("api.inference_artifacts.importlib.import_module")
+    def test_resolve_hf_uri_downloads_artifact(self, mock_download: MagicMock) -> None:
+        hf_download = MagicMock(return_value="/var/lib/ataxx/model_iter_039.pt")
+        mock_download.return_value = SimpleNamespace(hf_hub_download=hf_download)
+
+        resolved = resolve_artifact_uri("hf://dieg0code/ataxx-zero/model_iter_039.pt")
+        self.assertEqual(resolved, "/var/lib/ataxx/model_iter_039.pt")
+        hf_download.assert_called_once_with(
+            repo_id="dieg0code/ataxx-zero",
+            filename="model_iter_039.pt",
+            revision=None,
+            token=None,
+        )
+
+    def test_rejects_unknown_scheme(self) -> None:
+        with self.assertRaises(ValueError):
+            resolve_artifact_uri("s3://bucket/model.pt")
+
+
+if __name__ == "__main__":
+    unittest.main()