From a2bd1eebfd39063e74798689491cc340efb8fa88 Mon Sep 17 00:00:00 2001 From: Divyam Talwar Date: Sun, 20 Sep 2026 03:03:16 +0530 Subject: [PATCH] fix(ollama): retain registry and namespace in model identities Address #15 with focused regression coverage. AI-assisted implementation and isolated source review; exact validation and remaining platform limitations are recorded in the draft PR. Signed-off-by: Divyam Talwar --- quantprobe/ollama.py | 30 ++++- tests/smoke.py | 11 ++ tests/test_ollama_model_identity.py | 201 ++++++++++++++++++++++++++++ 3 files changed, 240 insertions(+), 2 deletions(-) create mode 100644 tests/test_ollama_model_identity.py diff --git a/quantprobe/ollama.py b/quantprobe/ollama.py index a4ed727..ff6a889 100644 --- a/quantprobe/ollama.py +++ b/quantprobe/ollama.py @@ -26,6 +26,31 @@ MODEL_MEDIA = "application/vnd.ollama.image.model" +# The one registry/namespace pair ollama leaves implicit in a name. Everything else is +# addressable only in full, so nothing else may be stripped. +DEFAULT_REGISTRY = "registry.ollama.ai" +DEFAULT_NAMESPACE = "library" + + +def manifest_name(rel): + """Manifest path parts under manifests/ -> the name ollama itself answers to. + + rel is [, , , ]. Only `registry.ollama.ai/library` is + implicit - that is the pair `ollama pull qwen2.5:7b` writes, and expanding it would print + an identity no user recognises. Every other prefix is part of the address: keeping just + the last two parts turns `hf.co/bartowski/x:Q4_K_M` into `x:Q4_K_M`, which `ollama run` + either cannot resolve or resolves to a DIFFERENT model of that basename that happens to + be pulled - so the audit prices one blob and measures another. + """ + if len(rel) < 2: + return "/".join(rel) + parts = list(rel) + if len(parts) > 2 and parts[0] == DEFAULT_REGISTRY: + parts = parts[1:] + if len(parts) > 2 and parts[0] == DEFAULT_NAMESPACE: + parts = parts[1:] + return f"{'/'.join(parts[:-1])}:{parts[-1]}" + def store_root(explicit=None): """Where ollama keeps its blobs. OLLAMA_MODELS wins, then the per-OS default.""" @@ -42,7 +67,8 @@ def installed(root=None): Walks manifests/ rather than calling `ollama list`, so it works when the daemon is not running and it gives the blob path directly. The manifest name is its path under - manifests////, which is how ollama itself addresses it. + manifests////, which is how ollama itself addresses it - + see manifest_name() for the one prefix that is implicit and may be dropped. """ root = store_root(root) mdir = os.path.join(root, "manifests") @@ -61,7 +87,7 @@ def installed(root=None): if not layers: continue rel = os.path.relpath(f, mdir).replace(os.sep, "/").split("/") - name = f"{rel[-2]}:{rel[-1]}" if len(rel) >= 2 else "/".join(rel) + name = manifest_name(rel) digest = layers[0]["digest"].split(":")[-1] blob = os.path.join(root, "blobs", f"sha256-{digest}") if os.path.isfile(blob): diff --git a/tests/smoke.py b/tests/smoke.py index fce0397..379444c 100644 --- a/tests/smoke.py +++ b/tests/smoke.py @@ -2888,6 +2888,17 @@ def t_ollama_store_reader_survives_a_broken_store(): return None +def t_ollama_model_identity_survives_registries_and_namespaces(): + """The name read off a manifest path must be one `ollama run` can resolve. + + Lives in tests/test_ollama_model_identity.py (unittest, synthetic stores, no daemon); + this hook keeps `python tests/smoke.py` covering it too. + """ + from test_ollama_model_identity import run_smoke + + return run_smoke() + + def _business_tasks_mod(): import importlib.util root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) diff --git a/tests/test_ollama_model_identity.py b/tests/test_ollama_model_identity.py new file mode 100644 index 0000000..0f72c79 --- /dev/null +++ b/tests/test_ollama_model_identity.py @@ -0,0 +1,201 @@ +"""The name audit-ollama reads off a manifest path must be the name ollama answers to. + +`installed()` walks manifests//// and hands that name +straight to `ollama run` and `ollama stop`. Keeping only the last two path parts prints a +model ollama cannot resolve for anything that is not on the default registry under +`library` - `hf.co/bartowski/x:Q4_K_M` is reported and run as `x:Q4_K_M`, which errors or, +worse, resolves to a DIFFERENT model of that basename that happens to be pulled. The +audit then prices one file and measures another. + +Synthetic manifest+blob trees only: no ollama, no daemon, no network. +""" + +from __future__ import annotations + +import json +import os +import tempfile +import unittest +from unittest import mock + +from quantprobe.ollama import MODEL_MEDIA, installed + +BLOB = b"not a real gguf, only a file the reader must find" + + +def _store(root, entries): + """Lay out manifests/ + the blobs they point at. entries: {path: digest}.""" + for relpath, digest in entries.items(): + mp = os.path.join(root, "manifests", *relpath.split("/")) + os.makedirs(os.path.dirname(mp), exist_ok=True) + with open(mp, "w", encoding="utf-8") as fh: + json.dump( + {"layers": [{"mediaType": MODEL_MEDIA, "digest": f"sha256:{digest}", "size": 1}]}, + fh, + ) + blob = os.path.join(root, "blobs", f"sha256-{digest}") + os.makedirs(os.path.dirname(blob), exist_ok=True) + with open(blob, "wb") as fh: + fh.write(BLOB) + return root + + +def _names(root): + return [n for n, _blob, _size in installed(root)] + + +class TestManifestPathToOllamaName(unittest.TestCase): + def test_default_registry_and_library_keep_the_shorthand(self): + """`ollama pull qwen2.5:7b` lands under registry.ollama.ai/library and is addressed + by the bare name - expanding it here would print an identity no user recognises.""" + with tempfile.TemporaryDirectory() as d: + _store(d, {"registry.ollama.ai/library/qwen2.5/7b": "aaa"}) + self.assertEqual(_names(d), ["qwen2.5:7b"]) + + def test_custom_namespace_on_the_default_registry_is_kept(self): + """A namespace that is not `library` is part of the name: `ollama run mymodel:latest` + does not find `myuser/mymodel:latest`.""" + with tempfile.TemporaryDirectory() as d: + _store(d, {"registry.ollama.ai/myuser/mymodel/latest": "bbb"}) + self.assertEqual(_names(d), ["myuser/mymodel:latest"]) + + def test_non_default_registry_host_is_kept(self): + """hf.co models are addressed by their full path, including the host.""" + with tempfile.TemporaryDirectory() as d: + _store(d, {"hf.co/bartowski/Qwen2.5-7B-GGUF/Q4_K_M": "ccc"}) + self.assertEqual(_names(d), ["hf.co/bartowski/Qwen2.5-7B-GGUF:Q4_K_M"]) + + def test_library_under_a_foreign_registry_is_not_shortened(self): + """`library` is only implicit on the default registry; elsewhere it is a real + namespace and dropping it would rename someone else's model.""" + with tempfile.TemporaryDirectory() as d: + _store(d, {"registry.example.com/library/mistral/latest": "ddd"}) + self.assertEqual(_names(d), ["registry.example.com/library/mistral:latest"]) + + def test_same_basename_in_two_namespaces_stays_two_distinct_models(self): + """The collision that makes the audit lie: two different files, one printed name. + + Whichever row the user acts on, `ollama run llama3:latest` reaches at most one of + them, while the header numbers above it came from the other blob. + """ + with tempfile.TemporaryDirectory() as d: + _store( + d, + { + "registry.ollama.ai/library/llama3/latest": "1111", + "hf.co/someone/llama3/latest": "2222", + }, + ) + got = {name: os.path.basename(blob) for name, blob, _size in installed(d)} + self.assertEqual( + got, + { + "llama3:latest": "sha256-1111", + "hf.co/someone/llama3:latest": "sha256-2222", + }, + ) + + def test_tag_defaults_are_not_invented(self): + """The tag comes from the file name on disk, never from a guess.""" + with tempfile.TemporaryDirectory() as d: + _store(d, {"registry.ollama.ai/library/phi3/mini-4k-instruct-q4_0": "eee"}) + self.assertEqual(_names(d), ["phi3:mini-4k-instruct-q4_0"]) + + +class TestBrokenStoresStillDegrade(unittest.TestCase): + """Naming must not cost the reader its tolerance for a store it does not own.""" + + def test_missing_store(self): + self.assertEqual(installed(os.path.join(tempfile.gettempdir(), "no-such-store-qp")), []) + + def test_manifest_without_its_blob_is_not_listed(self): + with tempfile.TemporaryDirectory() as d: + mp = os.path.join(d, "manifests", "hf.co", "someone", "ghost") + os.makedirs(mp) + with open(os.path.join(mp, "latest"), "w", encoding="utf-8") as fh: + json.dump( + {"layers": [{"mediaType": MODEL_MEDIA, "digest": "sha256:dead", "size": 1}]}, fh + ) + self.assertEqual(installed(d), []) + + def test_junk_and_layerless_manifests_are_skipped_without_raising(self): + with tempfile.TemporaryDirectory() as d: + _store(d, {"registry.ollama.ai/library/real/latest": "fff"}) + with open(os.path.join(d, "manifests", "README.txt"), "w", encoding="utf-8") as fh: + fh.write("not json") + noml = os.path.join(d, "manifests", "hf.co", "someone", "config") + os.makedirs(noml) + with open(os.path.join(noml, "latest"), "w", encoding="utf-8") as fh: + json.dump({"layers": [{"mediaType": "application/vnd.ollama.image.license"}]}, fh) + self.assertEqual(_names(d), ["real:latest"]) + + def test_a_manifest_directly_under_manifests_does_not_crash(self): + with tempfile.TemporaryDirectory() as d: + os.makedirs(os.path.join(d, "manifests")) + os.makedirs(os.path.join(d, "blobs")) + with open(os.path.join(d, "manifests", "stray"), "w", encoding="utf-8") as fh: + json.dump( + {"layers": [{"mediaType": MODEL_MEDIA, "digest": "sha256:999", "size": 1}]}, fh + ) + with open(os.path.join(d, "blobs", "sha256-999"), "wb") as fh: + fh.write(BLOB) + self.assertEqual(_names(d), ["stray"]) + + +class TestDownstreamUsesTheFullIdentity(unittest.TestCase): + """The name is not cosmetic: measure() and unload() hand it to the ollama CLI.""" + + def _run_with_fake_cli(self, fn): + """Record the argv a call would have used. No ollama binary is ever invoked.""" + + class FakeCompleted: + stdout = "eval rate: 19.92 tokens/s\n" + stderr = "" + + calls = [] + + def fake_run(cmd, **_kw): + calls.append(list(cmd)) + return FakeCompleted() + + with ( + mock.patch("quantprobe.ollama.ollama_bin", return_value="/fake/ollama"), + mock.patch("quantprobe.ollama.subprocess.run", side_effect=fake_run), + ): + fn() + return calls + + def test_measure_runs_the_name_the_store_reader_returned(self): + from quantprobe import ollama as om + + with tempfile.TemporaryDirectory() as d: + _store(d, {"hf.co/bartowski/Qwen2.5-7B-GGUF/Q4_K_M": "ccc"}) + (name, _blob, _size) = installed(d)[0] + calls = self._run_with_fake_cli(lambda: om.measure(name, prompt="hi")) + self.assertEqual( + calls[0][:3], ["/fake/ollama", "run", "hf.co/bartowski/Qwen2.5-7B-GGUF:Q4_K_M"] + ) + + def test_unload_stops_the_name_the_store_reader_returned(self): + from quantprobe import ollama as om + + with tempfile.TemporaryDirectory() as d: + _store(d, {"registry.ollama.ai/myuser/mymodel/latest": "bbb"}) + (name, _blob, _size) = installed(d)[0] + calls = self._run_with_fake_cli(lambda: om.unload(name, tries=0)) + self.assertEqual(calls[0], ["/fake/ollama", "stop", "myuser/mymodel:latest"]) + + +def run_smoke(): + """Hook for tests/smoke.py, which has no pytest. Raises on the first failure.""" + loader = unittest.TestLoader() + suite = loader.loadTestsFromModule(__import__(__name__, fromlist=["*"])) + with open(os.devnull, "w", encoding="utf-8") as quiet: + result = unittest.TextTestRunner(stream=quiet, verbosity=0).run(suite) + if not result.wasSuccessful(): + bad = result.failures + result.errors + raise AssertionError(f"{len(bad)} ollama identity test(s) failed: {bad[0][1].strip()}") + + +if __name__ == "__main__": + unittest.main()