From 8784876370e9f3293d4f2367af2dfb0fe923a960 Mon Sep 17 00:00:00 2001 From: Eric Curtin Date: Sun, 30 Aug 2026 21:58:01 +0100 Subject: [PATCH] feat: resolve oci:// model references via llmman serve Lets --model and --tokenizer point at a model published as a CNCF ModelPack OCI artifact: aphrodite run oci://ghcr.io/org/model:tag Model distribution is increasingly moving to OCI registries, which lets a deployment reuse the registry, credentials, mirroring and air-gap tooling it already has for container images. Acquisition is delegated to a running `llmman serve`, which already implements the ModelPack media types, registry auth, resumable blob download and a content-addressed store. The daemon does the pull (POST /api/pull, streamed as NDJSON so a multi-gigabyte fetch is not silent) but deliberately exposes no local path, so `llmman resolve --no-pull` reports where the bytes landed. The client is stdlib-only, so no new dependency. maybe_pull_model_tokenizer_for_runai is the existing rewrite hook that runs before anything else touches model/tokenizer, so the oci:// branch goes there. The two schemes are disjoint, so it returns early and the object-storage path is untouched. A tokenizer naming the same reference reuses the pull rather than fetching twice. An explicit oci:// scheme is required rather than sniffing a bare registry/name:tag: that shape is indistinguishable from a HuggingFace repo id, so guessing would silently hijack existing deployments. Signed-off-by: Eric Curtin --- aphrodite/config/model.py | 31 +++ aphrodite/transformers_utils/llmman.py | 212 +++++++++++++++++++++ aphrodite/transformers_utils/oci_utils.py | 65 +++++++ tests/transformers_utils/test_llmman.py | 117 ++++++++++++ tests/transformers_utils/test_oci_utils.py | 133 +++++++++++++ 5 files changed, 558 insertions(+) create mode 100644 aphrodite/transformers_utils/llmman.py create mode 100644 aphrodite/transformers_utils/oci_utils.py create mode 100644 tests/transformers_utils/test_llmman.py create mode 100644 tests/transformers_utils/test_oci_utils.py diff --git a/aphrodite/config/model.py b/aphrodite/config/model.py index d8dc40d49a..8929148a21 100644 --- a/aphrodite/config/model.py +++ b/aphrodite/config/model.py @@ -47,6 +47,7 @@ MODEL_ARCH_CONFIG_CONVERTORS, ModelArchConfigConvertorBase, ) +from aphrodite.transformers_utils.oci_utils import is_oci_uri, resolve_oci_model from aphrodite.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri from aphrodite.transformers_utils.utils import maybe_model_redirect from aphrodite.utils.import_utils import LazyLoader @@ -942,6 +943,14 @@ def maybe_pull_model_tokenizer_for_runai(self, model: str, tokenizer: str) -> No if self.model_weights: return + # A CNCF ModelPack artifact is pulled from a container registry and + # extracted to a local directory, which the default HuggingFace-format + # loading then sees. Handled before the object-storage path since the + # two schemes are disjoint. + if is_oci_uri(model) or is_oci_uri(tokenizer): + self.maybe_pull_model_tokenizer_for_oci(model, tokenizer) + return + if not (is_runai_obj_uri(model) or is_runai_obj_uri(tokenizer)): return @@ -975,6 +984,28 @@ def maybe_pull_model_tokenizer_for_runai(self, model: str, tokenizer: str) -> No ) self.tokenizer = object_storage_tokenizer.dir + def maybe_pull_model_tokenizer_for_oci(self, model: str, tokenizer: str) -> None: + """Pull a CNCF ModelPack artifact from an OCI registry. + + The whole image is extracted to one directory, so a tokenizer naming + the same reference reuses it rather than pulling twice. + + Args: + model: Model name or path + tokenizer: Tokenizer name or path + """ + resolved: dict[str, str] = {} + + if is_oci_uri(model): + resolved[model] = resolve_oci_model(model) + self.model_weights = model + self.model = resolved[model] + + if is_oci_uri(tokenizer): + if tokenizer not in resolved: + resolved[tokenizer] = resolve_oci_model(tokenizer) + self.tokenizer = resolved[tokenizer] + def _get_encoder_config(self) -> dict[str, Any] | None: return get_sentence_transformer_tokenizer_config(self.model, self.revision) diff --git a/aphrodite/transformers_utils/llmman.py b/aphrodite/transformers_utils/llmman.py new file mode 100644 index 0000000000..8f5e1c2989 --- /dev/null +++ b/aphrodite/transformers_utils/llmman.py @@ -0,0 +1,212 @@ +"""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 `` -> one line of JSON carrying ``path``. +""" + +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()}" + ) + 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) diff --git a/aphrodite/transformers_utils/oci_utils.py b/aphrodite/transformers_utils/oci_utils.py new file mode 100644 index 0000000000..031437b1df --- /dev/null +++ b/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 `` 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) diff --git a/tests/transformers_utils/test_llmman.py b/tests/transformers_utils/test_llmman.py new file mode 100644 index 0000000000..2f55665fb9 --- /dev/null +++ b/tests/transformers_utils/test_llmman.py @@ -0,0 +1,117 @@ +"""The `llmman serve` client: the daemon protocol behind oci:// model paths. + +Exercised against a real HTTP server on a loopback port rather than mocks, so +the NDJSON streaming contract is genuinely tested. +""" + +import http.server +import json +import socketserver +import threading + +import pytest + +from aphrodite.transformers_utils import llmman + + +def _ndjson(*objs): + return "".join(json.dumps(o) + "\n" for o in objs) + + +class _FakeDaemon: + """A minimal stand-in for `llmman serve`, on a real loopback port.""" + + def __init__(self): + self.version = {"version": "0.1.0", "pid": 1} + self.pull_body = _ndjson({"status": "success"}) + self.pull_status = 200 + self.last_request = None + daemon = self + + class Handler(http.server.BaseHTTPRequestHandler): + def log_message(self, *args): + pass + + def _send(self, status, body, ctype): + raw = body.encode() + self.send_response(status) + self.send_header("Content-Type", ctype) + self.send_header("Content-Length", str(len(raw))) + self.end_headers() + self.wfile.write(raw) + + def do_GET(self): + self._send(200, json.dumps(daemon.version), "application/json") + + def do_POST(self): + length = int(self.headers.get("Content-Length", 0)) + daemon.last_request = json.loads(self.rfile.read(length)) + self._send(daemon.pull_status, daemon.pull_body, "application/x-ndjson") + + self._server = socketserver.TCPServer(("127.0.0.1", 0), Handler) + self.url = f"http://127.0.0.1:{self._server.server_address[1]}" + threading.Thread(target=self._server.serve_forever, daemon=True).start() + + def close(self): + self._server.shutdown() + self._server.server_close() + + +@pytest.fixture +def daemon(): + d = _FakeDaemon() + yield d + d.close() + + +def test_accepts_a_llmman_daemon(daemon): + llmman.check_daemon(daemon.url) + + +def test_rejects_a_non_llmman_server(daemon): + daemon.version = {"hello": "world"} + with pytest.raises(RuntimeError, match="not an llmman daemon"): + llmman.check_daemon(daemon.url) + + +def test_reports_nothing_listening_actionably(): + with pytest.raises(RuntimeError, match="llmman serve"): + llmman.check_daemon("http://127.0.0.1:1") + + +def test_pull_succeeds_and_forwards_progress(daemon): + daemon.pull_body = _ndjson( + {"status": "pulling manifest"}, + {"status": "pulling blobs", "completed": 50, "total": 100}, + {"status": "success"}, + ) + seen = [] + llmman.pull(daemon.url, "ghcr.io/org/model:tag", lambda *a: seen.append(a)) + + assert daemon.last_request == {"model": "ghcr.io/org/model:tag"} + assert seen == [("pulling manifest", 0, 0), ("pulling blobs", 50, 100)] + + +def test_reports_an_in_band_error_at_http_200(daemon): + # The daemon streams errors in-band, so a 200 does not mean success. + daemon.pull_body = _ndjson({"status": "pulling"}, {"error": "unauthorized"}) + with pytest.raises(RuntimeError, match="unauthorized"): + llmman.pull(daemon.url, "ref") + + +def test_rejects_a_stream_that_ends_without_success(daemon): + daemon.pull_body = _ndjson({"status": "pulling blobs"}) + with pytest.raises(RuntimeError, match="without reporting success"): + llmman.pull(daemon.url, "ref") + + +def test_reports_a_non_ok_status(daemon): + daemon.pull_status = 400 + daemon.pull_body = '{"error":"bad request"}' + with pytest.raises(RuntimeError): + llmman.pull(daemon.url, "ref") + + +def test_tolerates_a_non_json_diagnostic_line(daemon): + daemon.pull_body = "not json\n" + _ndjson({"status": "success"}) + llmman.pull(daemon.url, "ref") diff --git a/tests/transformers_utils/test_oci_utils.py b/tests/transformers_utils/test_oci_utils.py new file mode 100644 index 0000000000..ee60d81a2c --- /dev/null +++ b/tests/transformers_utils/test_oci_utils.py @@ -0,0 +1,133 @@ +"""``oci://`` model references resolve to a local path. + +The scheme is explicit on purpose: a bare ``registry/name:tag`` is the same +shape as a HuggingFace repo id, so sniffing would hijack existing deployments. +""" + +import json +import os +import tempfile +from pathlib import Path +from unittest import mock + +import pytest + +from aphrodite.transformers_utils import llmman +from aphrodite.transformers_utils.oci_utils import ( + is_oci_uri, + resolve_oci_model, + strip_oci_scheme, +) + + +class TestScheme: + def test_recognizes_the_oci_scheme(self): + assert is_oci_uri("oci://ghcr.io/org/model:tag") + assert is_oci_uri("OCI://ghcr.io/org/model:tag") + + @pytest.mark.parametrize( + "value", + [ + "meta-llama/Llama-3-8B", + "ghcr.io/org/model:tag", + "/local/path/to/model", + "s3://bucket/key", + "gs://bucket/key", + "az://container/key", + "", + None, + ], + ) + def test_leaves_every_other_shape_alone(self, value): + assert not is_oci_uri(value) + + def test_accepts_pathlib_input(self): + assert not is_oci_uri(Path("/local/model")) + + def test_strips_the_scheme_only_when_present(self): + assert strip_oci_scheme("oci://ghcr.io/org/model:tag") == "ghcr.io/org/model:tag" + assert strip_oci_scheme("OCI://ghcr.io/org/model:tag") == "ghcr.io/org/model:tag" + assert strip_oci_scheme("meta-llama/Llama-3-8B") == "meta-llama/Llama-3-8B" + + +class TestResolveContract: + """`llmman resolve --no-pull` reports where the daemon's pull landed.""" + + def test_parses_the_documented_contract(self): + with tempfile.TemporaryDirectory() as path: + line = json.dumps({"reference": "r", "path": path, "format": "safetensors"}) + assert llmman.parse_resolve_output(line, "r") == path + + def test_tolerates_trailing_newline_and_leaked_diagnostics(self): + with tempfile.TemporaryDirectory() as path: + out = "pulling blobs...\n" + json.dumps({"path": path}) + "\n" + assert llmman.parse_resolve_output(out, "r") == path + + def test_ignores_unknown_fields_so_the_contract_can_grow(self): + with tempfile.TemporaryDirectory() as path: + line = json.dumps({"path": path, "format": "gguf", "mmproj": "/x", "future": 1}) + assert llmman.parse_resolve_output(line, "r") == path + + @pytest.mark.parametrize( + "bad", + [ + "", + " \n\n", + "not json", + '["a", "list"]', + '{"no_path": 1}', + '{"path": ""}', + '{"path": 3}', + '{"path": "/nonexistent/xyzzy"}', + ], + ) + def test_rejects_malformed_output(self, bad): + with pytest.raises(RuntimeError): + llmman.parse_resolve_output(bad, "r") + + +class TestEndpoint: + @pytest.mark.parametrize( + "host,want", + [ + ("", "http://127.0.0.1:17434"), + ("1.2.3.4:9999", "http://1.2.3.4:9999"), + ("1.2.3.4", "http://1.2.3.4:17434"), + ("http://1.2.3.4:9999/ignored", "http://1.2.3.4:9999"), + # A wildcard bind is meaningful to the server but not to a client. + ("0.0.0.0:9999", "http://127.0.0.1:9999"), + ("[::]:9999", "http://[::1]:9999"), + ], + ) + def test_parses_every_llmman_host_form(self, host, want): + with mock.patch.dict(os.environ, {llmman.HOST_ENV: host}): + assert llmman.endpoint() == want + + def test_binary_default_and_override(self): + with mock.patch.dict(os.environ, {llmman.BIN_ENV: ""}): + assert llmman.llmman_bin() == "llmman" + with mock.patch.dict(os.environ, {llmman.BIN_ENV: "/opt/llmman"}): + assert llmman.llmman_bin() == "/opt/llmman" + + +class TestResolveOciModel: + def test_rejects_an_empty_reference_without_touching_the_daemon(self): + for ref in ("oci://", "oci:// "): + with pytest.raises(ValueError): + resolve_oci_model(ref) + + def test_strips_the_scheme_before_handing_off_to_llmman(self): + with mock.patch( + "aphrodite.transformers_utils.oci_utils.llmman.pull_and_resolve", + return_value="/resolved", + ) as acquire: + assert resolve_oci_model("oci://ghcr.io/org/model:tag") == "/resolved" + assert acquire.call_args[0][0] == "ghcr.io/org/model:tag" + assert acquire.call_args[1]["progress"] is not None + + def test_reports_a_missing_binary(self): + with ( + mock.patch.dict(os.environ, {llmman.BIN_ENV: "/definitely/not/here"}), + pytest.raises(RuntimeError, match="not found"), + ): + llmman.resolve("ref")