From 2c95fbe1d7677dd2618eba86b8f6be2cff2ac49c Mon Sep 17 00:00:00 2001 From: Eric Curtin Date: Sun, 30 Aug 2026 21:51:15 +0100 Subject: [PATCH] feat: support oci:// model source via llmman serve Adds OCI as a URI protocol so a model published as a CNCF ModelPack artifact can be a model source: source: uri: oci://ghcr.io/org/model:tag This reuses the registry, credentials and mirroring a cluster already has for container images, which is often easier to run air-gapped than a model hub. Controller side: a new OCI protocol on URIProvider, injecting MODEL_SOURCE_TYPE=oci, the scheme-stripped reference, and the address of the llmman daemon the loader pulls through (LLMAZ_LLMMAN_HOST on the controller points every loader at one shared daemon). ModelPath returns the model directory itself: an OCI artifact unpacks as a whole and its layer filepaths already name the files, so there is no bucket key to derive models-- or a .gguf filename from. Loader side: a new llmaz.model_loader.oci module that delegates 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 Python dependency. Files are hard-linked out of llmman's store where possible, falling back to a copy across filesystems, so a model shared with llmman costs its bytes once rather than twice. Registry credentials live on the daemon rather than being injected per model, so one place covers every model pulled through it. Signed-off-by: Eric Curtin --- docs/examples/oci/README.md | 40 +++ docs/examples/oci/playground.yaml | 25 ++ llmaz/main.py | 9 + llmaz/model_loader/constant.py | 2 + llmaz/model_loader/oci/__init__.py | 0 llmaz/model_loader/oci/llmman.py | 249 ++++++++++++++++++ llmaz/model_loader/oci/oci.py | 74 ++++++ llmaz/tests/test_llmman.py | 208 +++++++++++++++ llmaz/tests/test_oci_loader.py | 114 ++++++++ .../modelsource/modelsource.go | 12 + .../modelsource/modelsource_test.go | 59 +++++ pkg/controller_helper/modelsource/uri.go | 31 +++ pkg/webhook/openmodel_webhook.go | 1 + pyproject.toml | 1 + 14 files changed, 825 insertions(+) create mode 100644 docs/examples/oci/README.md create mode 100644 docs/examples/oci/playground.yaml create mode 100644 llmaz/model_loader/oci/__init__.py create mode 100644 llmaz/model_loader/oci/llmman.py create mode 100644 llmaz/model_loader/oci/oci.py create mode 100644 llmaz/tests/test_llmman.py create mode 100644 llmaz/tests/test_oci_loader.py diff --git a/docs/examples/oci/README.md b/docs/examples/oci/README.md new file mode 100644 index 00000000..4b2f3e4e --- /dev/null +++ b/docs/examples/oci/README.md @@ -0,0 +1,40 @@ +# Load a model from an OCI registry + +A model published as a [CNCF ModelPack](https://github.com/modelpack/model-spec) +artifact can be used as a model source with the `oci://` protocol: + +```yaml +source: + uri: oci://ghcr.io/inftyai/qwen2-0.5b:latest +``` + +The model-loader initContainer pulls the artifact through a running +[`llmman serve`](https://github.com/llmmanorg/llmman) daemon and places the +files in the model directory, so any inference backend loads it the same way it +would a model fetched from a model hub or an object store. + +This reuses the registry, credentials and mirroring a cluster already has for +container images, which is often easier to run air-gapped than a model hub. + +## Requirements + +The loader talks to an `llmman serve` daemon. By default it uses llmman's own +default address, `127.0.0.1:17434`. To point every loader at one shared daemon +instead, set `LLMAZ_LLMMAN_HOST` on the controller: + +```yaml +env: + - name: LLMAZ_LLMMAN_HOST + value: llmman.llmaz-system.svc:17434 +``` + +## Private registries + +Registry credentials are configured on the llmman daemon rather than on the +model, so one place covers every model pulled through it. + +## How to use + +```bash +kubectl apply -f playground.yaml +``` diff --git a/docs/examples/oci/playground.yaml b/docs/examples/oci/playground.yaml new file mode 100644 index 00000000..dc9fce54 --- /dev/null +++ b/docs/examples/oci/playground.yaml @@ -0,0 +1,25 @@ +apiVersion: llmaz.io/v1alpha1 +kind: OpenModel +metadata: + name: qwen2-0--5b +spec: + familyName: qwen2 + source: + # A model published as a CNCF ModelPack artifact (https://github.com/modelpack/model-spec) + # in any OCI registry, following the protocol: + # oci:///[:|@] + uri: oci://ghcr.io/inftyai/qwen2-0.5b:latest + inferenceConfig: + flavors: + - name: t4 # GPU type + limits: + nvidia.com/gpu: 1 +--- +apiVersion: inference.llmaz.io/v1alpha1 +kind: Playground +metadata: + name: qwen2-0--5b +spec: + replicas: 1 + modelClaim: + modelName: qwen2-0--5b diff --git a/llmaz/main.py b/llmaz/main.py index 186164c3..96db58f7 100644 --- a/llmaz/main.py +++ b/llmaz/main.py @@ -20,6 +20,7 @@ from llmaz.model_loader.constant import * from llmaz.model_loader.objstore.objstore import model_download +from llmaz.model_loader.oci.oci import model_download as oci_model_download from llmaz.model_loader.model_hub.hub_factory import HubFactory from llmaz.model_loader.model_hub.huggingface import HUB_HUGGING_FACE from llmaz.util.logger import Logger @@ -58,6 +59,14 @@ src = os.getenv(ENV_OBJ_MODEL_PATH) model_download(provider=provider, endpoint=endpoint, bucket=bucket, src=src) + elif model_source_type == "oci": + reference = os.getenv(ENV_OCI_REFERENCE) + if not reference: + raise EnvironmentError( + f"Environment variable '{ENV_OCI_REFERENCE}' not found." + ) + + oci_model_download(reference=reference) else: raise EnvironmentError(f"unknown model source type {model_source_type}") diff --git a/llmaz/model_loader/constant.py b/llmaz/model_loader/constant.py index 08674898..7efa3629 100644 --- a/llmaz/model_loader/constant.py +++ b/llmaz/model_loader/constant.py @@ -14,3 +14,5 @@ ENV_OBJ_ENDPOINT = "ENDPOINT" ENV_OBJ_BUCKET = "BUCKET" ENV_OBJ_MODEL_PATH = "MODEL_PATH" + +ENV_OCI_REFERENCE = "OCI_REFERENCE" diff --git a/llmaz/model_loader/oci/__init__.py b/llmaz/model_loader/oci/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/llmaz/model_loader/oci/llmman.py b/llmaz/model_loader/oci/llmman.py new file mode 100644 index 00000000..efe5c1a4 --- /dev/null +++ b/llmaz/model_loader/oci/llmman.py @@ -0,0 +1,249 @@ +""" +Copyright 2024. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""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 llmaz.util.logger import Logger + +HOST_ENV = "LLMMAN_HOST" +BIN_ENV = "LLMAZ_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], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + 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(f"Pulling {reference} via llmman daemon at {base}") + pull(base, reference, progress) + return resolve(reference) diff --git a/llmaz/model_loader/oci/oci.py b/llmaz/model_loader/oci/oci.py new file mode 100644 index 00000000..1d05b99f --- /dev/null +++ b/llmaz/model_loader/oci/oci.py @@ -0,0 +1,74 @@ +""" +Copyright 2024. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import os +import shutil + +from llmaz.model_loader.constant import MODEL_LOCAL_DIR +from llmaz.model_loader.oci import llmman +from llmaz.util.logger import Logger + + +def model_download(reference: str, out_dir: str = MODEL_LOCAL_DIR): + """Acquire a CNCF ModelPack artifact through a running `llmman serve`. + + The daemon does the pull (POST /api/pull, streamed so a multi-gigabyte + fetch is not silent) but deliberately exposes no local path, so + `llmman resolve --no-pull` reports where the bytes landed. + """ + if not reference or not reference.strip(): + raise ValueError("OCI reference cannot be empty") + reference = reference.strip() + + def _progress(status, completed, total): + if total: + Logger.info(f"llmman: {status} ({completed}/{total} bytes)") + else: + Logger.info(f"llmman: {status}") + + resolved = llmman.pull_and_resolve(reference, progress=_progress) + materialize(resolved, out_dir) + Logger.info(f"placed {reference} at {out_dir}") + + +def materialize(src: str, out_dir: str): + """Place llmman's extracted model at out_dir. + + Files are hard-linked where possible so a model shared with llmman's store + costs its bytes once, falling back to a copy across filesystems. + """ + os.makedirs(out_dir, exist_ok=True) + + if not os.path.isdir(src): + link_or_copy(src, os.path.join(out_dir, os.path.basename(src))) + return + + for root, _, files in os.walk(src): + rel = os.path.relpath(root, src) + dest_root = out_dir if rel == "." else os.path.join(out_dir, rel) + os.makedirs(dest_root, exist_ok=True) + for name in files: + link_or_copy(os.path.join(root, name), os.path.join(dest_root, name)) + + +def link_or_copy(src: str, dest: str): + if os.path.lexists(dest): + os.remove(dest) + try: + os.link(src, dest) + except OSError: + # Different filesystem, or one without hard links. + shutil.copy2(src, dest) diff --git a/llmaz/tests/test_llmman.py b/llmaz/tests/test_llmman.py new file mode 100644 index 00000000..b15b7f67 --- /dev/null +++ b/llmaz/tests/test_llmman.py @@ -0,0 +1,208 @@ +""" +Copyright 2024. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""The `llmman serve` client: endpoint resolution and the daemon protocol. + +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 tempfile +import threading + +import pytest + +from llmaz.model_loader.oci import llmman + + +@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", "http://1.2.3.4:9999"), + ("http://1.2.3.4:9999/ignored", "http://1.2.3.4:9999"), + ('"1.2.3.4:9999"', "http://1.2.3.4:9999"), + # A wildcard bind is meaningful to the server but not to a client, + # which cannot connect to "every interface". + ("0.0.0.0:9999", "http://127.0.0.1:9999"), + ("[::]:9999", "http://[::1]:9999"), + ], +) +def test_endpoint_parsing(monkeypatch, host, want): + monkeypatch.setenv(llmman.HOST_ENV, host) + assert llmman.endpoint() == want + + +def test_binary_default_and_override(monkeypatch): + monkeypatch.delenv(llmman.BIN_ENV, raising=False) + assert llmman.llmman_bin() == "llmman" + monkeypatch.setenv(llmman.BIN_ENV, "/opt/bin/llmman") + assert llmman.llmman_bin() == "/opt/bin/llmman" + # An empty override is a mistake, not a request to run the empty string. + monkeypatch.setenv(llmman.BIN_ENV, " ") + assert llmman.llmman_bin() == "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): + assert self.path == "/api/version" + self._send(200, json.dumps(daemon.version), "application/json") + + def do_POST(self): + assert self.path == "/api/pull" + 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_check_daemon_accepts_a_llmman_daemon(daemon): + llmman.check_daemon(daemon.url) + + +def test_check_daemon_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_check_daemon_reports_nothing_listening(): + 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_pull_reports_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_pull_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_pull_tolerates_a_non_json_diagnostic_line(daemon): + daemon.pull_body = "not json\n" + _ndjson({"status": "success"}) + llmman.pull(daemon.url, "ref") + + +def test_pull_reports_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_parse_resolve_output_accepts_the_documented_contract(): + with tempfile.TemporaryDirectory() as path: + line = json.dumps({"reference": "r", "path": path, "format": "safetensors"}) + assert llmman.parse_resolve_output(line, "r") == path + + +def test_parse_resolve_output_tolerates_leaked_diagnostics(): + with tempfile.TemporaryDirectory() as path: + out = "pulling...\n" + json.dumps({"path": path}) + "\n" + assert llmman.parse_resolve_output(out, "r") == path + + +def test_parse_resolve_output_ignores_unknown_fields(): + with tempfile.TemporaryDirectory() as path: + line = json.dumps({"path": path, "format": "gguf", "mmproj": "/x", "new": 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_parse_resolve_output_rejects_malformed_output(bad): + with pytest.raises(RuntimeError): + llmman.parse_resolve_output(bad, "r") + + +def test_resolve_reports_a_missing_binary(monkeypatch): + monkeypatch.setenv(llmman.BIN_ENV, "/definitely/not/here/llmman") + with pytest.raises(RuntimeError, match="not found"): + llmman.resolve("ref") diff --git a/llmaz/tests/test_oci_loader.py b/llmaz/tests/test_oci_loader.py new file mode 100644 index 00000000..54482aa0 --- /dev/null +++ b/llmaz/tests/test_oci_loader.py @@ -0,0 +1,114 @@ +""" +Copyright 2024. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import os +from unittest import mock + +import pytest + +from llmaz.model_loader.oci.oci import link_or_copy, materialize, model_download + + +class TestModelDownload: + def test_rejects_an_empty_reference(self): + for ref in ("", " "): + with pytest.raises(ValueError, match="cannot be empty"): + model_download(ref, out_dir="/tmp/unused") + + def test_pulls_through_the_daemon_then_materializes(self, tmp_path): + store = tmp_path / "store" + store.mkdir() + (store / "model.safetensors").write_text("w") + out = str(tmp_path / "out") + + with mock.patch( + "llmaz.model_loader.oci.oci.llmman.pull_and_resolve", + return_value=str(store), + ) as acquire: + model_download("ghcr.io/org/model:tag", out_dir=out) + + # The daemon receives the bare reference; progress is forwarded. + assert acquire.call_args[0][0] == "ghcr.io/org/model:tag" + assert acquire.call_args[1]["progress"] is not None + assert open(os.path.join(out, "model.safetensors")).read() == "w" + + +class TestMaterialize: + def test_hard_links_a_directory(self, tmp_path): + src = tmp_path / "store" + (src / "sub").mkdir(parents=True) + (src / "config.json").write_text("{}") + (src / "sub" / "model.safetensors").write_text("w") + out = str(tmp_path / "out") + + materialize(str(src), out) + + assert open(os.path.join(out, "sub", "model.safetensors")).read() == "w" + # A model shared with llmman's store should cost its bytes once. + assert ( + os.stat(os.path.join(out, "config.json")).st_ino + == os.stat(str(src / "config.json")).st_ino + ) + + def test_handles_a_single_file_payload(self, tmp_path): + # A GGUF payload resolves to the file itself, not a directory. + src = tmp_path / "model.gguf" + src.write_text("gguf") + out = str(tmp_path / "out") + + materialize(str(src), out) + + assert open(os.path.join(out, "model.gguf")).read() == "gguf" + + def test_overwrites_a_stale_destination(self, tmp_path): + src = tmp_path / "store" + src.mkdir() + (src / "config.json").write_text("new") + out = tmp_path / "out" + out.mkdir() + (out / "config.json").write_text("stale") + + materialize(str(src), str(out)) + + assert (out / "config.json").read_text() == "new" + + def test_falls_back_to_copy_across_filesystems(self, tmp_path): + src = tmp_path / "store" + src.mkdir() + (src / "config.json").write_text("{}") + out = str(tmp_path / "out") + + with mock.patch( + "llmaz.model_loader.oci.oci.os.link", side_effect=OSError("EXDEV") + ): + materialize(str(src), out) + + assert open(os.path.join(out, "config.json")).read() == "{}" + assert ( + os.stat(os.path.join(out, "config.json")).st_ino + != os.stat(str(src / "config.json")).st_ino + ) + + +def test_link_or_copy_replaces_an_existing_target(tmp_path): + src = tmp_path / "a" + src.write_text("new") + dest = tmp_path / "b" + dest.write_text("stale") + + link_or_copy(str(src), str(dest)) + + assert dest.read_text() == "new" diff --git a/pkg/controller_helper/modelsource/modelsource.go b/pkg/controller_helper/modelsource/modelsource.go index 7225972e..e1fb35e0 100644 --- a/pkg/controller_helper/modelsource/modelsource.go +++ b/pkg/controller_helper/modelsource/modelsource.go @@ -39,6 +39,7 @@ const ( // model source type MODEL_SOURCE_MODELHUB = "modelhub" MODEL_SOURCE_MODEL_OBJ_STORE = "objstore" + MODEL_SOURCE_OCI = "oci" // secrets MODELHUB_SECRET_NAME = "modelhub-secret" @@ -52,6 +53,13 @@ const ( AWS_ACCESS_SECRET_NAME = "aws-access-secret" AWS_ACCESS_KEY_ID = "AWS_ACCESS_KEY_ID" AWS_ACCESS_KEY_SECRET = "AWS_SECRET_ACCESS_KEY" + + // Address of the `llmman serve` daemon the model loader pulls oci:// + // sources through. Defaults to llmman's own default; override with the + // LLMAZ_LLMMAN_HOST env var on the controller to point at a shared + // daemon (a DaemonSet Service, say) instead of a pod-local one. + LLMMAN_HOST_ENV = "LLMMAN_HOST" + DEFAULT_LLMMAN_HOST = "127.0.0.1:17434" ) type ModelSourceProvider interface { @@ -92,6 +100,10 @@ func NewModelSourceProvider(model *coreapi.OpenModel) ModelSourceProvider { provider.modelPath = value case Ollama: provider.modelPath = value + case OCI: + // The whole address is the registry reference; there is no bucket or + // endpoint to split out. + provider.modelPath = value default: // This should be validated at webhooks. panic("protocol not supported") diff --git a/pkg/controller_helper/modelsource/modelsource_test.go b/pkg/controller_helper/modelsource/modelsource_test.go index 9a2e54b1..4d9102db 100644 --- a/pkg/controller_helper/modelsource/modelsource_test.go +++ b/pkg/controller_helper/modelsource/modelsource_test.go @@ -77,6 +77,30 @@ func TestModelSourceProvider(t *testing.T) { wantModelPath: "/workspace/models/weight.gguf", skipModelLoader: false, }, + { + // An OCI artifact unpacks as a whole: its layer filepaths name the + // files, so no models-- directory is derived from the reference. + name: "model with OCI URI configured", + model: wrapper.MakeModel("test-7b").FamilyName("test").ModelSourceWithURI("oci://ghcr.io/org/model:tag").Obj(), + wantModelName: "test-7b", + wantModelPath: "/workspace/models/", + skipModelLoader: false, + }, + { + name: "model with OCI URI configured and skipModelLoader is true", + model: wrapper.MakeModel("test-7b").FamilyName("test").ModelSourceWithURI("oci://ghcr.io/org/model:tag").Obj(), + wantModelName: "test-7b", + wantModelPath: "oci://ghcr.io/org/model:tag", + skipModelLoader: true, + }, + { + // A tag containing a dot must not be mistaken for a GGUF filename. + name: "OCI URI whose tag contains a dot", + model: wrapper.MakeModel("test-7b").FamilyName("test").ModelSourceWithURI("oci://ghcr.io/org/model:v1.2.gguf").Obj(), + wantModelName: "test-7b", + wantModelPath: "/workspace/models/", + skipModelLoader: false, + }, } for _, tc := range testCases { @@ -226,3 +250,38 @@ func TestInjectModelEnvVars(t *testing.T) { }) } } + +func TestOCIInjectModelLoader(t *testing.T) { + provider := NewModelSourceProvider( + wrapper.MakeModel("test-7b").FamilyName("test").ModelSourceWithURI("oci://ghcr.io/org/model:tag").Obj(), + ) + + template := coreapplyv1.PodTemplateSpec().WithSpec( + coreapplyv1.PodSpec().WithContainers( + coreapplyv1.Container().WithName(MODEL_RUNNER_CONTAINER_NAME).WithImage("vllm:test"), + ), + ) + provider.InjectModelLoader(template, 0, "loader:test") + + assert.Len(t, template.Spec.InitContainers, 1) + envs := map[string]string{} + for _, env := range template.Spec.InitContainers[0].Env { + if env.Value != nil { + envs[*env.Name] = *env.Value + } + } + assert.Equal(t, MODEL_SOURCE_OCI, envs["MODEL_SOURCE_TYPE"]) + // The scheme is stripped: the daemon receives the bare registry reference. + assert.Equal(t, "ghcr.io/org/model:tag", envs["OCI_REFERENCE"]) + // The loader pulls through an llmman daemon, so it must know where it is. + assert.Equal(t, DEFAULT_LLMMAN_HOST, envs[LLMMAN_HOST_ENV]) +} + +func TestOCILlmmanHostIsOverridable(t *testing.T) { + // A cluster can point every loader at one shared daemon. + t.Setenv("LLMAZ_LLMMAN_HOST", "llmman.llmaz-system.svc:17434") + assert.Equal(t, "llmman.llmaz-system.svc:17434", llmmanHost()) + + t.Setenv("LLMAZ_LLMMAN_HOST", " ") + assert.Equal(t, DEFAULT_LLMMAN_HOST, llmmanHost()) +} diff --git a/pkg/controller_helper/modelsource/uri.go b/pkg/controller_helper/modelsource/uri.go index c60812ea..7cd49112 100644 --- a/pkg/controller_helper/modelsource/uri.go +++ b/pkg/controller_helper/modelsource/uri.go @@ -17,6 +17,7 @@ limitations under the License. package modelSource import ( + "os" "strconv" "strings" @@ -31,6 +32,9 @@ const ( S3 = "S3" Ollama = "OLLAMA" HostPath = "HOST" + // OCI addresses a model published as a CNCF ModelPack artifact in a + // container registry, e.g. oci://ghcr.io/org/model:tag. + OCI = "OCI" ) type URIProvider struct { @@ -69,6 +73,13 @@ func (p *URIProvider) ModelPath(skipModelLoader bool) string { return p.uri } + // An OCI artifact is unpacked as a whole into the model directory: its layer + // filepaths already name the files, so there is no bucket key to derive a + // models-- directory or a .gguf filename from. + if p.protocol == OCI { + return CONTAINER_MODEL_PATH + } + // protocol is oss. splits := strings.Split(p.modelPath, "/") @@ -135,11 +146,31 @@ func (p *URIProvider) InjectModelLoader(template *coreapplyv1.PodTemplateSpecApp coreapplyv1.EnvVar().WithName(OSS_ACCESS_KEY_ID).WithValueFrom(coreapplyv1.EnvVarSource().WithSecretKeyRef(coreapplyv1.SecretKeySelector().WithName(OSS_ACCESS_SECRET_NAME).WithKey(OSS_ACCESS_KEY_ID).WithOptional(true))), coreapplyv1.EnvVar().WithName(OSS_ACCESS_KEY_SECRET).WithValueFrom(coreapplyv1.EnvVarSource().WithSecretKeyRef(coreapplyv1.SecretKeySelector().WithName(OSS_ACCESS_SECRET_NAME).WithKey(OSS_ACCESS_KEY_SECRET).WithOptional(true))), ) + case OCI: + initContainer.WithEnv( + coreapplyv1.EnvVar().WithName("MODEL_SOURCE_TYPE").WithValue(MODEL_SOURCE_OCI), + // The loader receives the reference without the scheme, matching how + // the other protocols pass an already-parsed address. + coreapplyv1.EnvVar().WithName("OCI_REFERENCE").WithValue(p.modelPath), + // Registry work is done by an llmman daemon, not the loader itself, + // so it needs to know where that daemon is. Registry credentials + // are configured on the daemon rather than injected here. + coreapplyv1.EnvVar().WithName(LLMMAN_HOST_ENV).WithValue(llmmanHost()), + ) } template.Spec.WithInitContainers(initContainer) } +// llmmanHost is the address the model loader talks to, overridable so a +// cluster can point every loader at one shared daemon. +func llmmanHost() string { + if host := strings.TrimSpace(os.Getenv("LLMAZ_LLMMAN_HOST")); host != "" { + return host + } + return DEFAULT_LLMMAN_HOST +} + func (p *URIProvider) InjectModelEnvVars(template *coreapplyv1.PodTemplateSpecApplyConfiguration) { switch p.protocol { case S3, GCS: diff --git a/pkg/webhook/openmodel_webhook.go b/pkg/webhook/openmodel_webhook.go index 9e3176b0..1cf4576c 100644 --- a/pkg/webhook/openmodel_webhook.go +++ b/pkg/webhook/openmodel_webhook.go @@ -54,6 +54,7 @@ var SUPPORTED_OBJ_STORES = map[string]struct{}{ modelSource.S3: {}, modelSource.Ollama: {}, modelSource.HostPath: {}, + modelSource.OCI: {}, } // Default implements webhook.Defaulter so a webhook will be registered for the type diff --git a/pyproject.toml b/pyproject.toml index eeed69b3..1b899584 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ modelscope = "^1.17.0" omnistore = "^0.0.4" + [tool.poetry.group.dev.dependencies] black = "^24.4.2" pytest = "^8.3.2"