feat: resolve oci:// model references via llmman serve

aphrodite/config/model.py

@@ -48,6 +48,7 @@
     MODEL_ARCH_CONFIG_CONVERTORS,
     ModelArchConfigConvertorBase,
 )
+from aphrodite.transformers_utils.oci_utils import is_oci_uri, resolve_oci_model
 from aphrodite.transformers_utils.repo_utils import resolve_revision
 from aphrodite.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
 from aphrodite.transformers_utils.utils import maybe_model_redirect
@@ -1026,8 +1027,10 @@ def architecture(self) -> str:
         return self._architecture
 
     def maybe_pull_model_tokenizer_for_runai(self, model: str, tokenizer: str) -> None:
-        """Pull model/tokenizer from Object Storage to temporary
-        directory when needed.
+        """Pull model/tokenizer from Object Storage (s3://, gs://, az://) or
+        an OCI registry (oci://) to a local directory when needed.
+
+        Model and tokenizer are handled independently, so schemes can be mixed.
 
         Args:
             model: Model name or path
@@ -1038,10 +1041,15 @@ def maybe_pull_model_tokenizer_for_runai(self, model: str, tokenizer: str) -> No
         if self.model_weights:
             return
 
-        if not (is_runai_obj_uri(model) or is_runai_obj_uri(tokenizer)):
-            return
+        if is_oci_uri(model):
+            self.model_weights = model
+            self.model = resolve_oci_model(model)
 
-        if is_runai_obj_uri(model):
+            # A ModelPack image is pulled whole, so the tokenizer is in it
+            if model == tokenizer:
+                self.tokenizer = self.model
+                return
+        elif is_runai_obj_uri(model):
             object_storage_model = ObjectStorageModel(url=model)
             object_storage_model.pull_files(model, allow_pattern=["*.model", "*.py", "*.json"])
             self.model_weights = model
@@ -1063,7 +1071,9 @@ def maybe_pull_model_tokenizer_for_runai(self, model: str, tokenizer: str) -> No
                 return
 
         # Only download tokenizer if needed and not already handled
-        if is_runai_obj_uri(tokenizer):
+        if is_oci_uri(tokenizer):
+            self.tokenizer = resolve_oci_model(tokenizer)
+        elif is_runai_obj_uri(tokenizer):
             object_storage_tokenizer = ObjectStorageModel(url=tokenizer)
             object_storage_tokenizer.pull_files(
                 tokenizer,

aphrodite/transformers_utils/llmman.py

@@ -0,0 +1,217 @@
+"""Client for a running ``llmman serve`` daemon.
+
+Used to acquire models published as CNCF ModelPack
+(https://github.com/modelpack/model-spec) OCI artifacts. The daemon owns the
+registry work -- ModelPack media types, registry auth, resumable blob download
+and a content-addressed store -- so it is not reimplemented here.
+
+Contract (from llmman's src/cmd/serve.rs and src/daemon.rs):
+  - LLMMAN_HOST is ``[scheme://]host[:port][/path]``, default 127.0.0.1:17434.
+    A wildcard bind host (0.0.0.0, ::) is rewritten to loopback, since a client
+    cannot connect to "every interface".
+  - ``GET /api/version`` -> ``{"version":..., "exe":..., "pid":...}``.
+  - ``POST /api/pull`` ``{"model": ref}`` -> NDJSON stream of ``{"status":...}``
+    objects, terminated by ``{"status":"success"}`` or ``{"error":"..."}``.
+    An error can arrive in-band at HTTP 200.
+  - ``llmman resolve --no-pull <ref>`` -> one line of JSON carrying ``path``.
+
+Limitation: ``llmman resolve`` reads the local store (``LLMMAN_MODELS``) and
+never contacts the daemon, so the daemon must share that store with this
+process (same host and user, or ``LLMMAN_MODELS`` on shared storage).
+"""
+
+import ipaddress
+import json
+import os
+import shutil
+import subprocess
+import urllib.error
+import urllib.request
+
+from aphrodite.logger import init_logger
+
+logger = init_logger(__name__)
+
+HOST_ENV = "LLMMAN_HOST"
+BIN_ENV = "APHRODITE_LLMMAN_BIN"
+
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 17434
+
+PROBE_TIMEOUT_SECONDS = 5
+
+
+def _connectable_host(host: str) -> str:
+    """Rewrite a wildcard bind host to its loopback equivalent."""
+    try:
+        ip = ipaddress.ip_address(host.strip("[]"))
+    except ValueError:
+        return host
+    if not ip.is_unspecified:
+        return host
+    return "127.0.0.1" if ip.version == 4 else "::1"
+
+
+def endpoint() -> str:
+    """The http origin of the llmman daemon, honouring LLMMAN_HOST."""
+    raw = os.getenv(HOST_ENV, "").strip().strip("\"'")
+    if not raw:
+        return f"http://{DEFAULT_HOST}:{DEFAULT_PORT}"
+
+    if "://" in raw:
+        raw = raw.split("://", 1)[1]
+    raw = raw.split("/", 1)[0]
+
+    host, port = raw, DEFAULT_PORT
+    if raw.startswith("["):  # bracketed IPv6, optionally with :port
+        close = raw.find("]")
+        if close != -1:
+            host = raw[: close + 1]
+            rest = raw[close + 1 :]
+            if rest.startswith(":") and rest[1:].isdigit():
+                port = int(rest[1:])
+    elif raw.count(":") == 1:
+        maybe_host, maybe_port = raw.rsplit(":", 1)
+        if maybe_port.isdigit():
+            host, port = maybe_host, int(maybe_port)
+
+    host = host or DEFAULT_HOST
+    resolved = _connectable_host(host)
+    if ":" in resolved and not resolved.startswith("["):
+        resolved = f"[{resolved}]"
+    return f"http://{resolved}:{port}"
+
+
+def llmman_bin() -> str:
+    """The llmman executable name, overridable per project."""
+    return os.getenv(BIN_ENV, "").strip() or "llmman"
+
+
+def check_daemon(base: str) -> None:
+    """Confirm an llmman daemon is listening and is actually llmman."""
+    url = base + "/api/version"
+    try:
+        with urllib.request.urlopen(url, timeout=PROBE_TIMEOUT_SECONDS) as resp:
+            if resp.status != 200:
+                raise RuntimeError(f"llmman daemon at {base} answered /api/version with HTTP {resp.status}")
+            payload = json.loads(resp.read().decode("utf-8"))
+    except urllib.error.URLError as exc:
+        raise RuntimeError(
+            f"no llmman daemon reachable at {base} ({exc.reason}). Start one with "
+            f"`llmman serve`, or point {HOST_ENV} at an existing daemon."
+        ) from exc
+    except json.JSONDecodeError as exc:
+        raise RuntimeError(f"the server at {base} is not an llmman daemon (unparseable /api/version)") from exc
+
+    if not isinstance(payload, dict) or not payload.get("version"):
+        raise RuntimeError(f"the server at {base} is not an llmman daemon (no version in /api/version)")
+
+
+def pull(base: str, reference: str, progress=None) -> None:
+    """Stream POST /api/pull until the daemon reports success.
+
+    ``progress`` receives ``(status, completed, total)``. An error can arrive
+    in-band at HTTP 200, and a stream that ends without ``success`` is also a
+    failure -- neither is treated as a completed pull.
+    """
+    body = json.dumps({"model": reference}).encode("utf-8")
+    req = urllib.request.Request(
+        base + "/api/pull",
+        data=body,
+        headers={"Content-Type": "application/json"},
+        method="POST",
+    )
+
+    succeeded = False
+    try:
+        with urllib.request.urlopen(req) as resp:
+            if resp.status != 200:
+                raise RuntimeError(f"llmman pull of {reference!r} failed: HTTP {resp.status}")
+            for raw_line in resp:
+                line = raw_line.decode("utf-8").strip()
+                if not line:
+                    continue
+                try:
+                    obj = json.loads(line)
+                except json.JSONDecodeError:
+                    # Tolerate a non-JSON diagnostic rather than aborting a
+                    # pull that may still be progressing.
+                    continue
+                if not isinstance(obj, dict):
+                    continue
+                if obj.get("error"):
+                    raise RuntimeError(f"llmman pull of {reference!r} failed: {obj['error']}")
+                status = obj.get("status")
+                if status == "success":
+                    succeeded = True
+                    continue
+                if progress is not None and status:
+                    progress(status, obj.get("completed", 0), obj.get("total", 0))
+    except urllib.error.HTTPError as exc:
+        raise RuntimeError(f"llmman pull of {reference!r} failed: HTTP {exc.code}") from exc
+    except urllib.error.URLError as exc:
+        raise RuntimeError(f"llmman pull of {reference!r} failed: {exc.reason}") from exc
+
+    if not succeeded:
+        raise RuntimeError(f"llmman pull of {reference!r} ended without reporting success")
+
+
+def parse_resolve_output(stdout: str, reference: str) -> str:
+    """Parse ``llmman resolve`` stdout into the resolved local path."""
+    lines = [line.strip() for line in stdout.splitlines() if line.strip()]
+    if not lines:
+        raise RuntimeError(f"llmman resolve {reference!r}: no output on stdout")
+
+    try:
+        payload = json.loads(lines[-1])
+    except json.JSONDecodeError as exc:
+        raise RuntimeError(f"llmman resolve {reference!r}: could not parse output as JSON: {lines[-1]}") from exc
+
+    if not isinstance(payload, dict):
+        raise RuntimeError(f"llmman resolve {reference!r}: expected a JSON object, got {lines[-1]}")
+
+    path = payload.get("path")
+    if not isinstance(path, str) or not path.strip():
+        raise RuntimeError(f"llmman resolve {reference!r}: returned an empty path")
+    if not os.path.exists(path):
+        raise RuntimeError(f"llmman resolve {reference!r}: reported path {path!r} does not exist")
+    return path
+
+
+def resolve(reference: str) -> str:
+    """Ask the CLI where the daemon's pull left the model on disk.
+
+    ``--no-pull`` guarantees this only reports on bytes ``/api/pull`` already
+    fetched, so the daemon stays the only thing that touches the network.
+    """
+    binary = llmman_bin()
+    if shutil.which(binary) is None and not os.path.isfile(binary):
+        raise RuntimeError(
+            f"{binary!r} not found. Install llmman "
+            "(https://github.com/llmmanorg/llmman) and put it on PATH, or set "
+            f"{BIN_ENV} to its location."
+        )
+
+    completed = subprocess.run(
+        [binary, "resolve", "--no-pull", reference],
+        capture_output=True,
+        stdin=subprocess.DEVNULL,
+        text=True,
+        check=False,
+    )
+    if completed.returncode != 0:
+        raise RuntimeError(
+            f"`{binary} resolve --no-pull {reference}` failed with exit code "
+            f"{completed.returncode}: {completed.stderr.strip()}. The llmman daemon must "
+            "share this process's local store (same host and user, or LLMMAN_MODELS on shared storage)."
+        )
+    return parse_resolve_output(completed.stdout, reference)
+
+
+def pull_and_resolve(reference: str, progress=None) -> str:
+    """Full acquisition: probe the daemon, pull through it, report the path."""
+    base = endpoint()
+    check_daemon(base)
+    logger.info("Pulling %s via llmman daemon at %s", reference, base)
+    pull(base, reference, progress)
+    return resolve(reference)

aphrodite/transformers_utils/oci_utils.py

@@ -0,0 +1,65 @@
+"""Resolve ``oci://`` model references to a local path.
+
+A model published as a CNCF ModelPack (https://github.com/modelpack/model-spec)
+artifact lives in an ordinary container registry, so it reuses the registry,
+credentials, mirroring and air-gap tooling a deployment already has for
+container images.
+
+Registry work is delegated to the ``llmman`` CLI
+(https://github.com/llmmanorg/llmman) rather than reimplemented here: it already
+speaks the ModelPack media types, registry auth and resumable blob download, and
+keeps a content-addressed local store. ``llmman resolve <reference>`` pulls the
+image if it is not already local, extracts it, and prints one line of JSON on
+stdout::
+
+    {"reference": "ghcr.io/org/model:tag", "path": "/abs/path", "format": "safetensors"}
+
+Only ``path`` is consumed; that directory is handed to the ordinary HuggingFace
+loading path, exactly as if a local directory had been passed.
+
+An explicit ``oci://`` scheme is required rather than sniffing a bare
+``registry/name:tag``: that shape is indistinguishable from a HuggingFace repo
+id (``org/model``), so guessing would silently hijack existing deployments.
+"""
+
+from pathlib import Path
+
+from aphrodite.logger import init_logger
+from aphrodite.transformers_utils import llmman
+
+logger = init_logger(__name__)
+
+SUPPORTED_SCHEMES = ["oci://"]
+
+
+def is_oci_uri(model_or_path: str | Path | None) -> bool:
+    """Whether the reference carries the ``oci://`` scheme.
+
+    Cast to str to handle pathlib.Path inputs, mirroring is_runai_obj_uri.
+    """
+    if not model_or_path:
+        return False
+    return str(model_or_path).lower().startswith(tuple(SUPPORTED_SCHEMES))
+
+
+def strip_oci_scheme(reference: str | Path) -> str:
+    """Drop the ``oci://`` prefix, leaving the bare registry reference."""
+    text = str(reference)
+    if is_oci_uri(text):
+        return text[len(SUPPORTED_SCHEMES[0]) :]
+    return text
+
+
+def resolve_oci_model(reference: str | Path) -> str:
+    """Pull an ``oci://`` reference through llmman and return the local path."""
+    bare = strip_oci_scheme(reference)
+    if not bare.strip():
+        raise ValueError(f"empty OCI model reference: {reference!r}")
+
+    def _progress(status, completed, total):
+        if total:
+            logger.info("llmman: %s (%s/%s bytes)", status, completed, total)
+        else:
+            logger.info("llmman: %s", status)
+
+    return llmman.pull_and_resolve(bare.strip(), progress=_progress)

tests/test_config.py

@@ -1597,6 +1597,36 @@ def test_s3_url_different_model_and_tokenizer(mock_pull_files):
     assert mock_pull_files.call_args_list[1][0][0] == tokenizer_url
 
 
+@pytest.mark.parametrize(
+    ("model", "tokenizer", "expected_model", "expected_tokenizer"),
+    [
+        ("oci://ghcr.io/org/model:tag", "s3://bucket/tok/", "/oci/model", "/s3/tok"),
+        ("s3://bucket/model/", "oci://ghcr.io/org/tok:tag", "/s3/model", "/oci/tok"),
+        ("oci://ghcr.io/org/model:tag", "oci://ghcr.io/org/model:tag", "/oci/model", "/oci/model"),
+        ("oci://ghcr.io/org/model:tag", "org/tokenizer", "/oci/model", "org/tokenizer"),
+        ("org/model", "oci://ghcr.io/org/tok:tag", "org/model", "/oci/tok"),
+    ],
+)
+def test_oci_and_s3_model_tokenizer_can_be_mixed(model, tokenizer, expected_model, expected_tokenizer):
+    """oci:// and s3:// references can be mixed for model and tokenizer."""
+
+    def kind(ref):
+        return "model" if "model" in ref else "tok"
+
+    with (
+        patch("aphrodite.config.model.resolve_oci_model", side_effect=lambda ref: f"/oci/{kind(ref)}"),
+        patch(
+            "aphrodite.config.model.ObjectStorageModel",
+            side_effect=lambda url: SimpleNamespace(dir=f"/s3/{kind(url)}", pull_files=lambda *a, **k: None),
+        ),
+    ):
+        config = MockConfig(model=model, tokenizer=tokenizer)
+        ModelConfig.maybe_pull_model_tokenizer_for_runai(config, model, tokenizer)
+
+    assert config.model == expected_model
+    assert config.tokenizer == expected_tokenizer
+
+
 @pytest.mark.parametrize(
     ("model_id", "expected_attn_type", "expected_result", "reason"),
     [