Leonid PershinandCursor
ffe5a031b4
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>
102 lines
3.3 KiB
Python
102 lines
3.3 KiB
Python
"""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
|