Warm Ollama into VRAM on up and tunnel so Assistent chat is not cold.
A 1-token /api/chat after tags (and again if /api/ps is empty) loads VL weights before the first message. Mid KEEP_ALIVE is 15m so a short image-gen burst does not unload the model. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -7,6 +7,7 @@ from gpu_rent.llm_runtime import (
|
||||
decide_runtime,
|
||||
normalize_runtime,
|
||||
parse_ollama_models,
|
||||
preferred_ollama_model,
|
||||
write_ollama_models_preset,
|
||||
)
|
||||
|
||||
@@ -68,3 +69,61 @@ def test_already_have_ollama_tag_exact_only():
|
||||
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")
|
||||
|
||||
|
||||
def test_preferred_ollama_model(tmp_path: Path):
|
||||
path = tmp_path / "m.yaml"
|
||||
path.write_text(
|
||||
"models:\n - name: a:3b\n - name: b:7b\n default: true\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
assert preferred_ollama_model(path) == "b:7b"
|
||||
path.write_text('models:\n - "only:7b"\n', encoding="utf-8")
|
||||
assert preferred_ollama_model(path) == "only:7b"
|
||||
assert preferred_ollama_model(tmp_path / "missing.yaml") is None
|
||||
|
||||
|
||||
def test_warmup_ollama_http_skips_loaded(monkeypatch):
|
||||
from gpu_rent.llm_runtime import warmup_ollama_http
|
||||
|
||||
class Resp:
|
||||
def raise_for_status(self) -> None:
|
||||
return None
|
||||
|
||||
def json(self):
|
||||
return {"models": [{"name": "foo:7b"}]}
|
||||
|
||||
class Client:
|
||||
def __init__(self, *a, **k):
|
||||
pass
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
def get(self, url, timeout=None):
|
||||
assert url.endswith("/api/ps")
|
||||
return Resp()
|
||||
|
||||
def post(self, *_a, **_k):
|
||||
raise AssertionError("must not chat when already loaded")
|
||||
|
||||
monkeypatch.setattr("gpu_rent.llm_runtime.httpx.Client", Client)
|
||||
msg = warmup_ollama_http("http://127.0.0.1:17811", "foo:7b")
|
||||
assert "skip" in msg
|
||||
|
||||
|
||||
def test_maybe_warmup_skips_when_runtime_none():
|
||||
from gpu_rent.llm_runtime import maybe_warmup_ollama_local
|
||||
|
||||
logs: list[str] = []
|
||||
|
||||
class Cfg:
|
||||
llm_runtime = "none"
|
||||
ollama_models_manifest = Path("missing.yaml")
|
||||
ollama_local_port = 17811
|
||||
|
||||
maybe_warmup_ollama_local(Cfg(), logs.append)
|
||||
assert logs == []
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
"""Unit tests for remote ollama_warmup (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_warmup_remote",
|
||||
_ROOT / "src" / "gpu_rent" / "remote" / "ollama_warmup.py",
|
||||
)
|
||||
assert _SPEC and _SPEC.loader
|
||||
_mod = module_from_spec(_SPEC)
|
||||
_SPEC.loader.exec_module(_mod)
|
||||
|
||||
|
||||
class FakeResp:
|
||||
def __init__(self, body: str):
|
||||
self._body = body.encode("utf-8")
|
||||
|
||||
def read(self) -> bytes:
|
||||
return self._body
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
|
||||
def test_parse_env_file(tmp_path):
|
||||
path = tmp_path / ".gpu-rent-ollama.env"
|
||||
path.write_text(
|
||||
"# comment\nOLLAMA_KEEP_ALIVE=15m\nOLLAMA_CONTEXT_LENGTH=16384\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
env = _mod.parse_env_file(path)
|
||||
assert env["OLLAMA_KEEP_ALIVE"] == "15m"
|
||||
assert env["OLLAMA_CONTEXT_LENGTH"] == "16384"
|
||||
assert _mod.parse_env_file(tmp_path / "missing") == {}
|
||||
|
||||
|
||||
def test_warmup_skips_when_ps_has_model(monkeypatch):
|
||||
def fake_urlopen(req, timeout=None):
|
||||
url = getattr(req, "full_url", str(req))
|
||||
assert "/api/ps" in url
|
||||
return FakeResp('{"models":[{"name":"huihui_ai/qwen2.5-vl-abliterated:7b"}]}')
|
||||
|
||||
monkeypatch.setattr(_mod.urllib.request, "urlopen", fake_urlopen)
|
||||
msg = _mod.warmup("huihui_ai/qwen2.5-vl-abliterated:7b", keep_alive="15m", num_ctx=16384, timeout=5)
|
||||
assert "skip" in msg
|
||||
assert "VRAM" in msg
|
||||
|
||||
|
||||
def test_warmup_posts_one_token_chat(monkeypatch):
|
||||
seen: dict = {}
|
||||
|
||||
def fake_urlopen(req, timeout=None):
|
||||
url = getattr(req, "full_url", str(req))
|
||||
method = getattr(req, "method", "GET")
|
||||
if "/api/ps" in url:
|
||||
return FakeResp('{"models":[]}')
|
||||
seen["url"] = url
|
||||
seen["method"] = method
|
||||
seen["body"] = req.data
|
||||
return FakeResp('{"message":{"role":"assistant","content":"ok"}}')
|
||||
|
||||
monkeypatch.setattr(_mod.urllib.request, "urlopen", fake_urlopen)
|
||||
msg = _mod.warmup("foo:7b", keep_alive="15m", num_ctx=16384, timeout=5)
|
||||
assert msg.startswith("ok foo:7b")
|
||||
assert "/api/chat" in seen["url"]
|
||||
assert seen["method"] == "POST"
|
||||
import json
|
||||
|
||||
body = json.loads(seen["body"].decode())
|
||||
assert body["model"] == "foo:7b"
|
||||
assert body["stream"] is False
|
||||
assert body["keep_alive"] == "15m"
|
||||
assert body["options"]["num_predict"] == 1
|
||||
assert body["options"]["num_ctx"] == 16384
|
||||
assert body["messages"][0]["content"] == "ok"
|
||||
|
||||
|
||||
def test_main_never_fails_on_http_error(monkeypatch, tmp_path):
|
||||
jobs = tmp_path / "jobs.json"
|
||||
jobs.write_text('{"model":"foo:7b"}', encoding="utf-8")
|
||||
monkeypatch.setattr(_mod, "JOBS", jobs)
|
||||
monkeypatch.setattr(_mod, "ENV_FILE", tmp_path / "missing.env")
|
||||
|
||||
def boom(*_a, **_k):
|
||||
raise _mod.urllib.error.URLError("down")
|
||||
|
||||
monkeypatch.setattr(_mod.urllib.request, "urlopen", boom)
|
||||
assert _mod.main() == 0
|
||||
|
||||
|
||||
def test_warmup_body_matches_assistent_ctx():
|
||||
import json
|
||||
|
||||
body = json.loads(_mod.warmup_body("m:7b", keep_alive="15m", num_ctx=16384))
|
||||
assert body["options"]["num_ctx"] == 16384
|
||||
assert body["options"]["num_predict"] == 1
|
||||
@@ -82,4 +82,6 @@ def test_ollama_mid_4090_context_16k():
|
||||
)
|
||||
tune = ollama_tune_for(info)
|
||||
assert tune.context_length == 16384
|
||||
assert tune.keep_alive == "15m"
|
||||
assert "OLLAMA_CONTEXT_LENGTH=16384" in "\n".join(ollama_env_lines(tune))
|
||||
assert "OLLAMA_KEEP_ALIVE=15m" in "\n".join(ollama_env_lines(tune))
|
||||
|
||||
@@ -141,6 +141,7 @@ def _stub_tunnel(monkeypatch) -> None:
|
||||
monkeypatch.setattr("gpu_rent.tunnel._ssh_tunnel_forwarder", lambda: object)
|
||||
monkeypatch.setattr("gpu_rent.tunnel._start_forwarder", lambda *a, **k: _Fwd())
|
||||
monkeypatch.setattr("gpu_rent.ready.verify_stack_local", lambda *a, **k: [])
|
||||
monkeypatch.setattr("gpu_rent.llm_runtime.maybe_warmup_ollama_local", lambda *a, **k: None)
|
||||
monkeypatch.setattr("gpu_rent.tunnel.load_state", lambda: st)
|
||||
monkeypatch.setattr("gpu_rent.tunnel.save_state", lambda s: None)
|
||||
monkeypatch.setattr("gpu_rent.access_card.print_access_card", lambda *a, **k: None)
|
||||
|
||||
@@ -86,6 +86,7 @@ def test_install_ollama_skips_restart_when_unit_unchanged():
|
||||
text = files("gpu_rent.remote").joinpath("install_ollama.sh").read_text(encoding="utf-8")
|
||||
assert "cmp -s" in text
|
||||
assert "skip restart" in text
|
||||
assert '"mid", 10 * 1024**3, "15m"' in text
|
||||
|
||||
|
||||
def test_provision_llm_skips_on_api_tags_not_cli_list():
|
||||
@@ -97,6 +98,7 @@ def test_provision_llm_skips_on_api_tags_not_cli_list():
|
||||
assert "_ollama_api_tags" in text
|
||||
assert "awk 'NR>1" not in text
|
||||
assert "GPU не гасим" in text
|
||||
assert "ollama_warmup.py" in text
|
||||
assert "без моделей из ollama-models.yaml" not in text
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user