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"