From ffe5a031b4c6277c19c0952e4ac18792f982f431 Mon Sep 17 00:00:00 2001 From: Leonid Pershin Date: Fri, 21 Aug 2026 20:47:14 +0300 Subject: [PATCH] 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 --- docs/llm.md | 4 +- src/gpu_rent/llm_runtime.py | 88 +++++++++++++ src/gpu_rent/perf_tiers.py | 4 +- src/gpu_rent/provision.py | 27 ++++ src/gpu_rent/remote/install_ollama.sh | 2 +- src/gpu_rent/remote/ollama_warmup.py | 175 ++++++++++++++++++++++++++ src/gpu_rent/tunnel.py | 4 + tests/test_llm_runtime.py | 59 +++++++++ tests/test_ollama_warmup.py | 101 +++++++++++++++ tests/test_perf_tiers.py | 2 + tests/test_tunnel_watch.py | 1 + tests/test_warm_up_skips.py | 2 + 12 files changed, 465 insertions(+), 4 deletions(-) create mode 100644 src/gpu_rent/remote/ollama_warmup.py create mode 100644 tests/test_ollama_warmup.py diff --git a/docs/llm.md b/docs/llm.md index 7fb7d2e..3de27fc 100644 --- a/docs/llm.md +++ b/docs/llm.md @@ -79,6 +79,8 @@ $env:OLLAMA_HOST = "http://127.0.0.1:17811" На `up` — `ollama pull` по списку, сверка с **`/api/tags`**. Слой 100% без `status=success` не считается успехом (кэш ≠ модель в Assistent). Если тега всё ещё нет — warning, **GPU не гасим**. Лишнее на диске не удаляет. +После успешного тега — **warmup**: 1-token `POST /api/chat`, чтобы веса (~6 GB VL) легли в VRAM до первого сообщения в Assistent. То же при `gpu-rent tunnel`, если `/api/ps` пуст. Промах warmup — warning, GPU не гасим. + ### Пресеты (меню) | # | ключ | tag / смысл | @@ -103,7 +105,7 @@ Unit `gpu-rent-ollama` читает `/mnt/swarm_data/.gpu-rent-gpu.json`: | Tier (VRAM) | Flash Attn | KEEP_ALIVE | KV cache | GPU_OVERHEAD | CONTEXT | | --- | --- | --- | --- | --- | --- | | low (<16 GiB) | off | 2m | q4_0 | 6 GiB | 8k | -| mid (16–23) | on* | 5m | q8_0 | 10 GiB | 16k | +| mid (16–23) | on* | 15m | q8_0 | 10 GiB | 16k | | high (24–47) | on* | 15m | q8_0 | 14 GiB | 16k | | ultra (≥48) | on* | 30m | q8_0 | 20 GiB | 32k | diff --git a/src/gpu_rent/llm_runtime.py b/src/gpu_rent/llm_runtime.py index b2ddce6..d2ace44 100644 --- a/src/gpu_rent/llm_runtime.py +++ b/src/gpu_rent/llm_runtime.py @@ -2,10 +2,13 @@ from __future__ import annotations +import time +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path from typing import Any +import httpx import yaml from gpu_rent.paths import ( @@ -138,6 +141,91 @@ def already_have_ollama_tag(have: set[str], wanted: str) -> bool: return False +def preferred_ollama_model(path: Path) -> str | None: + """Manifest default, else first tag.""" + entries = parse_ollama_models(path) + for entry in entries: + if entry.default: + return entry.name + return entries[0].name if entries else None + + +def _ollama_ps_names(payload: object) -> set[str]: + names: set[str] = set() + if not isinstance(payload, dict): + return names + for item in payload.get("models") or []: + if isinstance(item, dict): + for key in ("name", "model"): + val = item.get(key) + if val: + names.add(str(val)) + elif isinstance(item, str) and item.strip(): + names.add(item.strip()) + return names + + +def warmup_ollama_http( + base_url: str, + model: str, + *, + keep_alive: str = "15m", + num_ctx: int = 16384, + timeout: float = 180.0, +) -> str: + """Load ``model`` into VRAM via a 1-token /api/chat. Skip if /api/ps already has it.""" + base = str(base_url).rstrip("/") + t0 = time.monotonic() + timeout_cfg = httpx.Timeout(timeout, connect=8.0) + with httpx.Client(timeout=timeout_cfg) as client: + try: + ps = client.get(f"{base}/api/ps", timeout=8.0) + ps.raise_for_status() + if already_have_ollama_tag(_ollama_ps_names(ps.json()), model): + return f"Ollama warmup skip — {model} уже в VRAM" + except (httpx.HTTPError, ValueError, TypeError): + pass + body = { + "model": model, + "stream": False, + "keep_alive": keep_alive, + "options": { + "num_predict": 1, + "num_ctx": int(num_ctx), + "temperature": 0, + }, + "messages": [{"role": "user", "content": "ok"}], + } + resp = client.post(f"{base}/api/chat", json=body) + resp.raise_for_status() + elapsed = time.monotonic() - t0 + return f"Ollama warmup {model} ok ({elapsed:.1f}s)" + + +def maybe_warmup_ollama_local(cfg: Any, log: Callable[[str], None]) -> None: + """After the laptop tunnel is up: load preferred tag if VRAM is cold.""" + try: + if normalize_runtime(getattr(cfg, "llm_runtime", "none")) != "ollama": + return + except ValueError: + return + manifest = Path(getattr(cfg, "ollama_models_manifest", "") or "") + model = preferred_ollama_model(manifest) + if not model: + return + port = int(getattr(cfg, "ollama_local_port", 17811) or 17811) + log(f"Ollama warmup {model} (VRAM)") + try: + log( + warmup_ollama_http( + f"http://127.0.0.1:{port}", + model, + ) + ) + except Exception as exc: + log(f"⚠ Ollama warmup: {exc}") + + def write_ollama_models_preset(path: Path, preset: str) -> None: key = (preset or "recommended").strip().lower() if key not in OLLAMA_PRESETS: diff --git a/src/gpu_rent/perf_tiers.py b/src/gpu_rent/perf_tiers.py index 9f3877f..94acb9f 100644 --- a/src/gpu_rent/perf_tiers.py +++ b/src/gpu_rent/perf_tiers.py @@ -83,10 +83,10 @@ def ollama_tune_for(info: GpuInfo) -> OllamaTune: note = "high: flash+q8 KV, 14GiB reserved for Swarm, ctx 16k, keep 15m" elif info.tier == TIER_MID: overhead = 10 * 1024**3 - keep = "5m" + keep = "15m" kv = "q8_0" ctx = 16384 - note = "mid: flash+q8 KV, 10GiB reserved for Swarm, ctx 16k, keep 5m" + note = "mid: flash+q8 KV, 10GiB reserved for Swarm, ctx 16k, keep 15m" else: overhead = 6 * 1024**3 keep = "2m" diff --git a/src/gpu_rent/provision.py b/src/gpu_rent/provision.py index 6a2447f..b1c6f96 100644 --- a/src/gpu_rent/provision.py +++ b/src/gpu_rent/provision.py @@ -1045,6 +1045,33 @@ def provision_llm(cfg: Config, host: str, log: Log) -> None: + f" (есть: {sorted(have) or 'пусто'}). " "SwarmUI ок — GPU не гасим; Assistent будет пустой." ) + warm = next( + ( + n + for n in list(defaults) + names + if already_have_ollama_tag(have, n) + ), + "", + ) + if warm: + log(f"Ollama warmup {warm} (гружу в VRAM)") + put_text( + cfg, + host, + "/tmp/gpu-rent-ollama-warmup.json", + json.dumps({"model": warm}, indent=2), + ) + try: + run_python( + cfg, + host, + _pkg_text("ollama_warmup.py"), + remote_path="/tmp/gpu-rent-ollama_warmup.py", + timeout=240, + log=log, + ) + except Exception as exc: + log(f"⚠ Ollama warmup: {exc}") else: raise CloudError(f"неизвестный LLM_RUNTIME={runtime!r}") st = load_state() diff --git a/src/gpu_rent/remote/install_ollama.sh b/src/gpu_rent/remote/install_ollama.sh index 0b438af..9b57024 100644 --- a/src/gpu_rent/remote/install_ollama.sh +++ b/src/gpu_rent/remote/install_ollama.sh @@ -49,7 +49,7 @@ if gib >= 48: elif gib >= 24: tier, overhead, keep, kv, flash, ctx = "high", 14 * 1024**3, "15m", "q8_0", True, 16384 elif gib >= 16: - tier, overhead, keep, kv, flash, ctx = "mid", 10 * 1024**3, "5m", "q8_0", True, 16384 + tier, overhead, keep, kv, flash, ctx = "mid", 10 * 1024**3, "15m", "q8_0", True, 16384 else: tier, overhead, keep, kv, flash, ctx = "low", 6 * 1024**3, "2m", "q4_0", False, 8192 flash = bool(flash and (ampere or "A100" in name.upper() or "H100" in name.upper() or gib >= 16)) diff --git a/src/gpu_rent/remote/ollama_warmup.py b/src/gpu_rent/remote/ollama_warmup.py new file mode 100644 index 0000000..640655c --- /dev/null +++ b/src/gpu_rent/remote/ollama_warmup.py @@ -0,0 +1,175 @@ +#!/usr/bin/env python3 +"""Load the preferred Ollama model into VRAM. Stdlib only. Runs on the VM. + +First Assistent /api/chat otherwise waits on a ~6GB VL mmap. A 1-token ping +after pull (and again if /api/ps is empty) keeps that wait off the chat path. +Never exits non-zero: warmup miss must not trip UP_STOP_ON_FAIL. +""" +from __future__ import annotations + +import json +import os +import sys +import time +import urllib.error +import urllib.request +from pathlib import Path + +JOBS = Path("/tmp/gpu-rent-ollama-warmup.json") +ENV_FILE = Path("/mnt/swarm_data/.gpu-rent-ollama.env") +OLLAMA = "http://127.0.0.1:11434" + + +def already_have(have: set[str], wanted: str) -> bool: + if wanted in have: + return True + if ":" not in wanted and f"{wanted}:latest" in have: + return True + if wanted.endswith(":latest") and wanted.rsplit(":", 1)[0] in have: + return True + return False + + +def parse_env_file(path: Path) -> dict[str, str]: + out: dict[str, str] = {} + if not path.is_file(): + return out + try: + text = path.read_text(encoding="utf-8") + except OSError: + return out + for line in text.splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, val = line.split("=", 1) + out[key.strip()] = val.strip() + return out + + +def read_jobs() -> dict: + if not JOBS.is_file(): + return {} + raw = json.loads(JOBS.read_text(encoding="utf-8")) + if isinstance(raw, str): + return {"model": raw.strip()} + if isinstance(raw, dict): + return raw + return {} + + +def names_from_models(payload: object) -> set[str]: + names: set[str] = set() + if not isinstance(payload, dict): + return names + for m in payload.get("models") or []: + if isinstance(m, dict): + for key in ("name", "model"): + val = m.get(key) + if val: + names.add(str(val)) + elif isinstance(m, str) and m.strip(): + names.add(m.strip()) + return names + + +def api_json(url: str, *, data: bytes | None = None, timeout: float) -> dict: + headers = {"Content-Type": "application/json"} if data is not None else {} + req = urllib.request.Request( + url, + data=data, + headers=headers, + method="POST" if data is not None else "GET", + ) + with urllib.request.urlopen(req, timeout=timeout) as resp: + raw = resp.read().decode("utf-8") + if not raw.strip(): + return {} + parsed = json.loads(raw) + return parsed if isinstance(parsed, dict) else {} + + +def loaded_names(*, timeout: float = 8.0) -> set[str]: + try: + data = api_json(f"{OLLAMA}/api/ps", timeout=timeout) + except ( + urllib.error.URLError, + urllib.error.HTTPError, + OSError, + TimeoutError, + json.JSONDecodeError, + ): + return set() + return names_from_models(data) + + +def warmup_body(model: str, *, keep_alive: str, num_ctx: int) -> bytes: + return json.dumps( + { + "model": model, + "stream": False, + "keep_alive": keep_alive, + "options": { + "num_predict": 1, + "num_ctx": int(num_ctx), + "temperature": 0, + }, + "messages": [{"role": "user", "content": "ok"}], + } + ).encode("utf-8") + + +def warmup(model: str, *, keep_alive: str, num_ctx: int, timeout: float) -> str: + have = loaded_names() + if already_have(have, model): + return f"skip — {model} уже в VRAM" + t0 = time.monotonic() + api_json( + f"{OLLAMA}/api/chat", + data=warmup_body(model, keep_alive=keep_alive, num_ctx=num_ctx), + timeout=timeout, + ) + elapsed = time.monotonic() - t0 + return f"ok {model} ({elapsed:.1f}s)" + + +def main() -> int: + jobs = read_jobs() + model = str(jobs.get("model") or "").strip() + if not model: + print("ollama warmup: нет model — skip") + return 0 + env = parse_env_file(ENV_FILE) + keep = str( + jobs.get("keep_alive") + or env.get("OLLAMA_KEEP_ALIVE") + or os.environ.get("OLLAMA_KEEP_ALIVE") + or "15m" + ).strip() + try: + ctx = int( + jobs.get("num_ctx") + or env.get("OLLAMA_CONTEXT_LENGTH") + or os.environ.get("OLLAMA_CONTEXT_LENGTH") + or 16384 + ) + except (TypeError, ValueError): + ctx = 16384 + timeout = float(jobs.get("timeout") or 180.0) + print(f"ollama warmup {model} keep_alive={keep} num_ctx={ctx}", flush=True) + try: + print(warmup(model, keep_alive=keep, num_ctx=ctx, timeout=timeout), flush=True) + except ( + urllib.error.URLError, + urllib.error.HTTPError, + OSError, + TimeoutError, + json.JSONDecodeError, + RuntimeError, + ) as exc: + print(f"FAIL warmup {model}: {exc}", flush=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/gpu_rent/tunnel.py b/src/gpu_rent/tunnel.py index c6bff95..31ccb41 100644 --- a/src/gpu_rent/tunnel.py +++ b/src/gpu_rent/tunnel.py @@ -299,6 +299,10 @@ def run_tunnel( log(f"проверка туннеля: {exc}") raise + from gpu_rent.llm_runtime import maybe_warmup_ollama_local + + maybe_warmup_ollama_local(cfg, log) + from gpu_rent.access_card import print_access_card print_access_card(cfg, tunneled=True, host=current_host) diff --git a/tests/test_llm_runtime.py b/tests/test_llm_runtime.py index ad71dcd..2380ff4 100644 --- a/tests/test_llm_runtime.py +++ b/tests/test_llm_runtime.py @@ -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 == [] diff --git a/tests/test_ollama_warmup.py b/tests/test_ollama_warmup.py new file mode 100644 index 0000000..cb667cd --- /dev/null +++ b/tests/test_ollama_warmup.py @@ -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 diff --git a/tests/test_perf_tiers.py b/tests/test_perf_tiers.py index 68c6789..a52622a 100644 --- a/tests/test_perf_tiers.py +++ b/tests/test_perf_tiers.py @@ -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)) diff --git a/tests/test_tunnel_watch.py b/tests/test_tunnel_watch.py index aa121cd..f063170 100644 --- a/tests/test_tunnel_watch.py +++ b/tests/test_tunnel_watch.py @@ -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) diff --git a/tests/test_warm_up_skips.py b/tests/test_warm_up_skips.py index 93f5533..bb5f6f1 100644 --- a/tests/test_warm_up_skips.py +++ b/tests/test_warm_up_skips.py @@ -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