Enhance Ollama model management and performance tuning
- Updated the `provision_llm` function to utilize the `/api/tags` endpoint for verifying available models, improving accuracy in model management. - Introduced a new `already_have_ollama_tag` function to ensure exact tag matching, preventing mismatches during model checks. - Enhanced the `pull_stream` function to require a successful status from the API before proceeding, ensuring reliable model downloads. - Added logic to handle unwritten blob files, improving the robustness of the model pulling process. - Updated documentation and tests to reflect these changes, ensuring clarity and reliability in Ollama model operations.
This commit is contained in:
@@ -21,7 +21,95 @@ def test_exact_tag_only():
|
||||
assert already_have(have, "qwen2.5:3b")
|
||||
|
||||
|
||||
def test_latest_alias():
|
||||
assert already_have({"foo:latest"}, "foo")
|
||||
assert already_have({"foo"}, "foo:latest")
|
||||
assert not already_have({"foo:3b"}, "foo")
|
||||
def test_pull_stream_requires_success_then_tags(monkeypatch):
|
||||
class FakeResp:
|
||||
def __init__(self, lines: list[str]):
|
||||
self._lines = [ln.encode() for ln in lines]
|
||||
self._i = 0
|
||||
|
||||
def readline(self) -> bytes:
|
||||
if self._i >= len(self._lines):
|
||||
return b""
|
||||
row = self._lines[self._i]
|
||||
self._i += 1
|
||||
return row
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
lines = [
|
||||
'{"status":"pulling manifest"}\n',
|
||||
'{"status":"downloading","total":100,"completed":100}\n',
|
||||
]
|
||||
monkeypatch.setattr(
|
||||
_mod.urllib.request, "urlopen", lambda *a, **k: FakeResp(lines)
|
||||
)
|
||||
cli = {"n": 0}
|
||||
|
||||
def fake_cli(*_a, **_k):
|
||||
cli["n"] += 1
|
||||
|
||||
monkeypatch.setattr(_mod, "pull_cli", fake_cli)
|
||||
monkeypatch.setattr(_mod, "api_tags", lambda: set())
|
||||
monkeypatch.setattr(_mod, "listed", lambda: set())
|
||||
monkeypatch.setattr(_mod.time, "sleep", lambda *_a, **_k: None)
|
||||
try:
|
||||
_mod.pull_stream("foo:7b", "x", tags_wait=0)
|
||||
assert False, "expected RuntimeError"
|
||||
except RuntimeError as exc:
|
||||
assert "/api/tags" in str(exc)
|
||||
assert cli["n"] == 2
|
||||
|
||||
|
||||
def test_pull_stream_ok_when_success_and_tags(monkeypatch):
|
||||
class FakeResp:
|
||||
def __init__(self, lines: list[str]):
|
||||
self._lines = [ln.encode() for ln in lines]
|
||||
self._i = 0
|
||||
|
||||
def readline(self) -> bytes:
|
||||
if self._i >= len(self._lines):
|
||||
return b""
|
||||
row = self._lines[self._i]
|
||||
self._i += 1
|
||||
return row
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
lines = ['{"status":"success"}\n']
|
||||
monkeypatch.setattr(
|
||||
_mod.urllib.request, "urlopen", lambda *a, **k: FakeResp(lines)
|
||||
)
|
||||
monkeypatch.setattr(_mod, "api_tags", lambda: {"foo:7b"})
|
||||
monkeypatch.setattr(_mod, "listed", lambda: set())
|
||||
_mod.pull_stream("foo:7b", "x", tags_wait=0)
|
||||
|
||||
|
||||
def test_unwritten_blob_is_all_nuls(tmp_path):
|
||||
z = tmp_path / "zeros"
|
||||
z.write_bytes(b"\x00" * 23)
|
||||
real = tmp_path / "gguf"
|
||||
real.write_bytes(b"GGUF" + b"\x00" * 12)
|
||||
empty = tmp_path / "empty"
|
||||
empty.write_bytes(b"")
|
||||
assert _mod.is_unwritten_blob(z)
|
||||
assert not _mod.is_unwritten_blob(real)
|
||||
assert not _mod.is_unwritten_blob(empty)
|
||||
|
||||
|
||||
def test_purge_nul_blobs_keeps_gguf(tmp_path):
|
||||
blobs = tmp_path / "blobs"
|
||||
blobs.mkdir()
|
||||
(blobs / "sha256-dead").write_bytes(b"\x00" * 28)
|
||||
(blobs / "sha256-gguf").write_bytes(b"GGUF\x03\x00\x00\x00")
|
||||
n = _mod.purge_nul_blobs(tmp_path)
|
||||
assert n == 1
|
||||
assert not (blobs / "sha256-dead").exists()
|
||||
assert (blobs / "sha256-gguf").exists()
|
||||
|
||||
Reference in New Issue
Block a user