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:
@@ -17,6 +17,9 @@ def test_bootstrap_script_is_native_swarmui():
|
||||
assert "GPU_RENT_BOOTSTRAP_LIGHT" in script
|
||||
assert "GPU_RENT_SKIP_SWARMUI" in script
|
||||
assert "light bootstrap — пропускаем apt-get" in script
|
||||
assert "python3.12-dev" in script
|
||||
assert "TRITON_CACHE_DIR" in script
|
||||
assert "ensure_triton_build_deps" in script
|
||||
# llm-only re-up can skip apt when data marker exists (not only Swarm boot marker)
|
||||
assert 'MARKER_DATA' in script or ".gpu-rent-ready" in script
|
||||
# .NET: chown before Swarm install script; curl fallback if wget/perms fail
|
||||
|
||||
@@ -3,6 +3,7 @@ from pathlib import Path
|
||||
import pytest
|
||||
|
||||
from gpu_rent.llm_runtime import (
|
||||
already_have_ollama_tag,
|
||||
decide_runtime,
|
||||
normalize_runtime,
|
||||
parse_ollama_models,
|
||||
@@ -59,3 +60,11 @@ def test_write_preset(tmp_path: Path):
|
||||
assert parse_ollama_models(path)[0].name == "huihui_ai/qwen2.5-vl-abliterated:32b"
|
||||
write_ollama_models_preset(path, "text")
|
||||
assert parse_ollama_models(path)[0].name == "huihui_ai/qwen2.5-abliterate:7b"
|
||||
|
||||
|
||||
def test_already_have_ollama_tag_exact_only():
|
||||
have = {"qwen2.5:7b", "foo:latest"}
|
||||
assert already_have_ollama_tag(have, "qwen2.5:7b")
|
||||
assert not already_have_ollama_tag(have, "qwen2.5:3b")
|
||||
assert already_have_ollama_tag(have, "foo")
|
||||
assert already_have_ollama_tag(have, "foo:latest")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -82,6 +82,7 @@ def test_pip_ok_patches_extra_args(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(mod, "find_pip", lambda: pip)
|
||||
monkeypatch.setattr(mod, "find_comfy_python", lambda: pip)
|
||||
monkeypatch.setattr(mod, "pip_install_sage", lambda _p: True)
|
||||
monkeypatch.setattr(mod, "triton_jit_ok", lambda _p: True)
|
||||
|
||||
assert mod.main() == 0
|
||||
marker = json.loads((data / ".gpu-rent-perf-tuned").read_text(encoding="utf-8"))
|
||||
@@ -131,6 +132,7 @@ def test_pip_fail_retries_next_run(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(mod, "find_pip", lambda: pip)
|
||||
monkeypatch.setattr(mod, "find_comfy_python", lambda: pip)
|
||||
monkeypatch.setattr(mod, "pip_install_sage", lambda _p: True)
|
||||
monkeypatch.setattr(mod, "triton_jit_ok", lambda _p: True)
|
||||
|
||||
assert mod.main() == 0
|
||||
new_m = json.loads(marker.read_text(encoding="utf-8"))
|
||||
@@ -223,3 +225,37 @@ def test_ensure_absolute_start_script(tmp_path, monkeypatch):
|
||||
text = backends.read_text(encoding="utf-8")
|
||||
assert str(main_py.resolve()) in text
|
||||
assert mod.ensure_absolute_start_script() is False
|
||||
|
||||
|
||||
def test_jit_fail_strips_sage_extra_args(tmp_path, monkeypatch):
|
||||
mod = _load()
|
||||
data = tmp_path
|
||||
backends = data / "Data" / "Backends.fds"
|
||||
backends.parent.mkdir(parents=True)
|
||||
backends.write_text("ExtraArgs: --use-sage-attention\n", encoding="utf-8")
|
||||
gpu_json = data / ".gpu-rent-gpu.json"
|
||||
gpu_json.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"vram_mib": 24576,
|
||||
"compute_cap": "8.9",
|
||||
"uuid": "gpu-1",
|
||||
"name": "RTX",
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
pip = data / "fake-pip"
|
||||
pip.write_text("#!/bin/sh\n", encoding="utf-8")
|
||||
monkeypatch.setattr(mod, "DATA", data)
|
||||
monkeypatch.setattr(mod, "GPU_JSON", gpu_json)
|
||||
monkeypatch.setattr(mod, "MARKER", data / ".gpu-rent-perf-tuned")
|
||||
monkeypatch.setattr(mod, "BACKENDS", backends)
|
||||
monkeypatch.setattr(mod, "find_comfy_python", lambda: pip)
|
||||
monkeypatch.setattr(mod, "pip_install_sage", lambda _p: True)
|
||||
monkeypatch.setattr(mod, "triton_jit_ok", lambda _p: False)
|
||||
assert mod.main() == 0
|
||||
marker = json.loads((data / ".gpu-rent-perf-tuned").read_text(encoding="utf-8"))
|
||||
assert marker["pip_ok"] is True
|
||||
assert marker["jit_ok"] is False
|
||||
assert "--use-sage-attention" not in backends.read_text(encoding="utf-8")
|
||||
|
||||
@@ -188,4 +188,59 @@ def test_stack_probe_requires_ollama_models():
|
||||
|
||||
assert "/api/tags" in _REMOTE_STACK_PROBE
|
||||
assert "0 models" in _REMOTE_STACK_PROBE
|
||||
assert 'payload.get("models")' in _REMOTE_STACK_PROBE
|
||||
assert "WANT_OLLAMA_MODELS" in _REMOTE_STACK_PROBE
|
||||
assert "retry" in _REMOTE_STACK_PROBE
|
||||
assert ".gpu-rent-ollama-pulling" in _REMOTE_STACK_PROBE
|
||||
assert "WARN 0 models" in _REMOTE_STACK_PROBE
|
||||
assert "GPU не гасим" in _REMOTE_STACK_PROBE
|
||||
|
||||
|
||||
def test_verify_stack_zero_models_does_not_fail_up(monkeypatch):
|
||||
from gpu_rent.ready import ServiceCheck, verify_stack_on_vm
|
||||
|
||||
class C:
|
||||
enable_swarmui = True
|
||||
llm_runtime = "ollama"
|
||||
|
||||
def fake_probe(_cfg, _host):
|
||||
return [
|
||||
ServiceCheck("swarmui", True, "ok", "vm", retry=False),
|
||||
ServiceCheck(
|
||||
"ollama",
|
||||
True,
|
||||
"WARN 0 models — Assistent empty (GPU не гасим)",
|
||||
"vm",
|
||||
retry=False,
|
||||
),
|
||||
]
|
||||
|
||||
monkeypatch.setattr("gpu_rent.ready._probe_vm_once", fake_probe)
|
||||
out = verify_stack_on_vm(C(), "1.2.3.4", [].append, timeout=5.0, poll_every=0.01)
|
||||
assert all(c.ok for c in out)
|
||||
|
||||
|
||||
def test_verify_stack_retries_while_pulling(monkeypatch):
|
||||
from gpu_rent.ready import ServiceCheck, verify_stack_on_vm
|
||||
|
||||
class C:
|
||||
enable_swarmui = False
|
||||
llm_runtime = "ollama"
|
||||
|
||||
n = {"i": 0}
|
||||
|
||||
def fake_probe(_cfg, _host):
|
||||
n["i"] += 1
|
||||
if n["i"] == 1:
|
||||
return [
|
||||
ServiceCheck(
|
||||
"ollama", False, "0 models — pull идёт (12s)", "vm", retry=True
|
||||
)
|
||||
]
|
||||
return [ServiceCheck("ollama", True, "1 models (qwen)", "vm", retry=False)]
|
||||
|
||||
monkeypatch.setattr("gpu_rent.ready._probe_vm_once", fake_probe)
|
||||
out = verify_stack_on_vm(
|
||||
C(), "1.2.3.4", [].append, timeout=5.0, poll_every=0.01
|
||||
)
|
||||
assert out[0].ok
|
||||
assert n["i"] == 2
|
||||
|
||||
@@ -88,6 +88,18 @@ def test_install_ollama_skips_restart_when_unit_unchanged():
|
||||
assert "skip restart" in text
|
||||
|
||||
|
||||
def test_provision_llm_skips_on_api_tags_not_cli_list():
|
||||
from pathlib import Path
|
||||
|
||||
from gpu_rent import provision
|
||||
|
||||
text = Path(provision.__file__).read_text(encoding="utf-8")
|
||||
assert "_ollama_api_tags" in text
|
||||
assert "awk 'NR>1" not in text
|
||||
assert "GPU не гасим" in text
|
||||
assert "без моделей из ollama-models.yaml" not in text
|
||||
|
||||
|
||||
def test_cli_has_update_flag():
|
||||
from gpu_rent import cli
|
||||
|
||||
|
||||
Reference in New Issue
Block a user