Leonid Pershin
5832c5cf75
- 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.
116 lines
3.5 KiB
Python
116 lines
3.5 KiB
Python
"""Unit tests for remote ollama_pull matching (stdlib helpers)."""
|
|
|
|
from importlib.util import module_from_spec, spec_from_file_location
|
|
from pathlib import Path
|
|
|
|
_ROOT = Path(__file__).resolve().parents[1]
|
|
_SPEC = spec_from_file_location(
|
|
"ollama_pull_remote",
|
|
_ROOT / "src" / "gpu_rent" / "remote" / "ollama_pull.py",
|
|
)
|
|
assert _SPEC and _SPEC.loader
|
|
_mod = module_from_spec(_SPEC)
|
|
_SPEC.loader.exec_module(_mod)
|
|
already_have = _mod.already_have
|
|
|
|
|
|
def test_exact_tag_only():
|
|
have = {"qwen2.5:3b", "qwen2.5:7b"}
|
|
assert already_have(have, "qwen2.5:7b")
|
|
assert not already_have(have, "qwen2.5:14b")
|
|
assert already_have(have, "qwen2.5:3b")
|
|
|
|
|
|
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()
|