Merge pull request #9 from runpod-workers/luke/prewarm-first-request

Pre-load the model into GPU memory at boot so the first request doesn't

Author: lukepiette

README.md

@@ -179,6 +179,8 @@ Non-streaming responses return Ollama's native response object:
 | `OLLAMA_KEEP_ALIVE` | `-1` (forever) | How long models stay loaded in VRAM |
 | `OLLAMA_LOAD_TIMEOUT` | `60m` | How long Ollama waits for a model to load into memory before failing the request. Ollama's own default is `5m`, which a large model on a fresh worker can exceed |
 
+**Boot warm-up:** after registering the configured model, the worker loads it into GPU memory before accepting jobs (an empty-prompt `/api/generate`, bounded by `OLLAMA_LOAD_TIMEOUT`). Registering only writes the model to disk — Ollama loads VRAM lazily on the first inference — so without the warm-up the *first user request* pays the multi-minute load and can hit the execution timeout. With `OLLAMA_KEEP_ALIVE=-1` the weights then stay resident for the worker's lifetime. If the warm-up fails, the worker logs the cause and still starts; the first request retries the load.
+
 ## Storage and disk sizing
 
 | Configuration | Ollama store | GGUF acquisition | Disk needed |

handler.py

@@ -3,6 +3,7 @@
 import os
 import re
 import subprocess
+import time
 import uuid
 
 import requests
@@ -852,6 +853,136 @@ def ensure_default_model():
     return model
 
 
+# Go-style durations, which is what Ollama's env vars use. "ms" before "m" so
+# "500ms" doesn't half-match as "500m" plus a dangling "s".
+_DURATION_RE = re.compile(r"^([+-]?)((?:\d+(?:\.\d+)?(?:ms|us|ns|h|m|s))+)$")
+_DURATION_PART_RE = re.compile(r"(\d+(?:\.\d+)?)(ms|us|ns|h|m|s)")
+_UNIT_SECONDS = {"h": 3600.0, "m": 60.0, "s": 1.0, "ms": 1e-3, "us": 1e-6, "ns": 1e-9}
+
+# Fallback and cap for the warm-up bound: never let boot hang on a load forever,
+# even when OLLAMA_LOAD_TIMEOUT is unset, unparseable, or "retry indefinitely".
+_WARMUP_FALLBACK_S = 3600.0
+# Client-side margin over Ollama's own load timeout, so the server's error — which
+# says *why* the load failed — wins the race against a blank client timeout.
+_WARMUP_MARGIN_S = 60.0
+
+
+def _duration_seconds(value, default):
+    """Parse a Go time.ParseDuration string ('60m', '1h30m', '90s') to seconds.
+
+    Ollama parses OLLAMA_LOAD_TIMEOUT with Go's parser; mirroring it here keeps
+    one env var meaning one thing. Unparseable values return `default` rather
+    than raising — a bad env var must not break boot.
+    """
+    match = _DURATION_RE.match((value or "").strip().lower())
+    if not match:
+        return default
+    total = sum(
+        float(number) * _UNIT_SECONDS[unit]
+        for number, unit in _DURATION_PART_RE.findall(match.group(2))
+    )
+    return -total if match.group(1) == "-" else total
+
+
+def _warmup_timeout_seconds():
+    """How long the boot warm-up may spend loading the model into memory.
+
+    Ollama itself gives up after OLLAMA_LOAD_TIMEOUT (its default is 5m; this
+    image's Hub default is 60m; non-positive means retry forever), so the client
+    bound sits just above that: Ollama's load-timeout error names the model and
+    the cause, which beats a mute client-side timeout. Non-positive or
+    unparseable values are capped at an hour so warm-up is never unbounded.
+    """
+    timeout = _duration_seconds(os.environ.get("OLLAMA_LOAD_TIMEOUT"), _WARMUP_FALLBACK_S)
+    if timeout <= 0:
+        timeout = _WARMUP_FALLBACK_S
+    return timeout + _WARMUP_MARGIN_S
+
+
+def _report_residency(model):
+    """Log where the loaded weights actually ended up (VRAM vs CPU spill).
+
+    /api/ps is the only place Ollama reports size_vram; a partial-VRAM load is
+    the "first request is mysteriously slow" case worth flagging at boot.
+    """
+    response = session.get(f"{OLLAMA_BASE_URL}/api/ps", timeout=30)
+    if not response.ok:
+        return
+    for entry in response.json().get("models", []):
+        if entry.get("name") not in (model, f"{model}:latest"):
+            continue
+        size = entry.get("size") or 0
+        vram = entry.get("size_vram") or 0
+        if size and vram < size:
+            print(
+                f"WARN: '{model}' loaded, but only {_human_size(vram)} of "
+                f"{_human_size(size)} fits in VRAM — the rest is offloaded to CPU and "
+                f"responses will be much slower. Pick a GPU with more VRAM, a smaller "
+                f"quantization via HF_QUANTIZATION, or a lower OLLAMA_CONTEXT_LENGTH.",
+                flush=True,
+            )
+        else:
+            print(f"'{model}' is resident in GPU memory ({_human_size(vram)})", flush=True)
+        return
+    print(f"WARN: '{model}' answered the warm-up but is not listed by /api/ps", flush=True)
+
+
+def warm_model(model):
+    """Load `model`'s weights into GPU memory before any user request arrives.
+
+    Registering a model (ensure_default_model) only writes it to disk — Ollama
+    loads weights into VRAM lazily, on the first inference. Without this, the
+    first user request pays the whole multi-minute load and can blow the
+    endpoint's execution timeout. An empty prompt is Ollama's documented
+    "just load it" request: no tokens are generated, and the image's
+    OLLAMA_KEEP_ALIVE=-1 default keeps the weights resident afterwards.
+    """
+    timeout = _warmup_timeout_seconds()
+    print(
+        f"Loading '{model}' into GPU memory (bounded at {int(timeout)}s "
+        f"by OLLAMA_LOAD_TIMEOUT)...",
+        flush=True,
+    )
+    started = time.monotonic()
+    response = session.post(
+        f"{OLLAMA_BASE_URL}/api/generate",
+        json={"model": model, "prompt": "", "stream": False},
+        timeout=timeout,
+    )
+    error = ollama_error(response)
+    if error:
+        raise ValueError(error)
+    print(f"Loaded '{model}' in {time.monotonic() - started:.1f}s", flush=True)
+    try:
+        _report_residency(model)
+    except (requests.RequestException, ValueError) as err:
+        print(f"WARN: could not read /api/ps after warming '{model}': {err}", flush=True)
+
+
+def warm_default_model():
+    """Best-effort boot-time warm-up of the model the endpoint is configured for.
+
+    Deliberately non-fatal: the model is already registered, so the handler still
+    works if this fails — the first request just pays the load, as it did before
+    warm-up existed. But say exactly what failed, so "too big for this GPU" or
+    "load timed out" is diagnosable from the boot log instead of from a timed-out
+    first request. Returns True when the model is warm.
+    """
+    model = resolve_default_model()
+    if not model:
+        return False
+    try:
+        warm_model(model)
+        return True
+    except (requests.RequestException, ValueError, OSError) as err:
+        print(
+            f"WARN: could not pre-load '{model}' into GPU memory — the first request "
+            f"will trigger the load instead and may be slow or time out. Cause: {err}",
+            flush=True,
+        )
+        return False
+
+
 def handler(job):
     job_input = job.get("input") or {}
 

start.sh

@@ -92,11 +92,15 @@ fi
 
 # Always runs: resolve_default_model has a fallback, so there is always a model to
 # prepare. ensure_default_model resolves HF_MODEL over OLLAMA_MODEL, reads Runpod's
-# model store first, and uses HF_TOKEN for gated repos.
-if /opt/venv/bin/python -c 'import handler; print("Model ready:", handler.ensure_default_model())'; then
+# model store first, and uses HF_TOKEN for gated repos. warm_default_model then
+# loads the weights into GPU memory: registering only writes the model to disk,
+# and without the warm-up the *first user request* pays the multi-minute VRAM
+# load and can blow the execution timeout. Warm-up failure is non-fatal (the
+# handler retries the load on the first request) and logs its own cause.
+if /opt/venv/bin/python -c 'import handler; print("Model ready:", handler.ensure_default_model()); handler.warm_default_model()'; then
     :
 else
-    echo "WARN: startup model preparation failed — the handler will retry on the first request"
+    echo "WARN: startup model preparation failed (cause in the traceback above) — the handler will retry on the first request, which will be slow and may time out"
 fi
 
 cd /

test_warmup.py

@@ -0,0 +1,219 @@
+"""Unit tests for the boot-time warm-up logic in handler.py.
+
+These need no GPU, no network and no running Ollama, which is the point: the
+warm-up only ever runs on real hardware at cold start, so getting its request
+shape, bounds or failure handling wrong is otherwise only discovered by a timed
+out first request in production.
+
+    python3 -m pytest test_warmup.py -v
+"""
+
+import json
+import sys
+import types
+
+import pytest
+
+# handler.py imports the runpod SDK at module scope; stub it when it isn't
+# installed so the pure functions stay testable from a bare checkout.
+if "runpod" not in sys.modules:
+    try:
+        import runpod  # noqa: F401
+    except ImportError:
+        sys.modules["runpod"] = types.ModuleType("runpod")
+
+import requests
+
+import handler
+
+
+class FakeResponse:
+    def __init__(self, status_code=200, payload=None, text=None, path="/api/generate"):
+        self.status_code = status_code
+        self.ok = status_code < 400
+        self._payload = payload
+        self.text = text if text is not None else json.dumps(payload or {})
+        self.request = types.SimpleNamespace(path_url=path)
+
+    def json(self):
+        if self._payload is None:
+            raise ValueError("no json")
+        return self._payload
+
+
+class FakeSession:
+    """Records every call so tests can assert on the exact request shape."""
+
+    def __init__(self, post_response, ps_response=None):
+        self.post_response = post_response
+        self.ps_response = ps_response or FakeResponse(payload={"models": []}, path="/api/ps")
+        self.posts = []
+        self.gets = []
+
+    def post(self, url, **kwargs):
+        self.posts.append((url, kwargs))
+        return self.post_response
+
+    def get(self, url, **kwargs):
+        self.gets.append((url, kwargs))
+        return self.ps_response
+
+
+def loaded_ps(model, size=1000, size_vram=None):
+    return FakeResponse(
+        payload={
+            "models": [
+                {"name": model, "size": size, "size_vram": size if size_vram is None else size_vram}
+            ]
+        },
+        path="/api/ps",
+    )
+
+
+# --- _duration_seconds: Go-style durations, as Ollama's env vars use ------------
+
+
+@pytest.mark.parametrize(
+    "value, expected",
+    [
+        ("60m", 3600.0),
+        ("1h30m", 5400.0),
+        ("90s", 90.0),
+        ("500ms", 0.5),
+        ("1.5h", 5400.0),
+        ("-5m", -300.0),
+    ],
+)
+def test_duration_seconds_parses(value, expected):
+    assert handler._duration_seconds(value, 123.0) == pytest.approx(expected)
+
+
+@pytest.mark.parametrize("value", ["", None, "garbage", "10", "5x", "m5", "1h 30m"])
+def test_duration_seconds_falls_back(value):
+    assert handler._duration_seconds(value, 123.0) == 123.0
+
+
+# --- _warmup_timeout_seconds: always bounded, tracks OLLAMA_LOAD_TIMEOUT --------
+
+
+def test_warmup_timeout_tracks_load_timeout(monkeypatch):
+    monkeypatch.setenv("OLLAMA_LOAD_TIMEOUT", "10m")
+    assert handler._warmup_timeout_seconds() == 600 + handler._WARMUP_MARGIN_S
+
+
+@pytest.mark.parametrize("value", [None, "0s", "-1s", "not-a-duration"])
+def test_warmup_timeout_never_unbounded(monkeypatch, value):
+    if value is None:
+        monkeypatch.delenv("OLLAMA_LOAD_TIMEOUT", raising=False)
+    else:
+        monkeypatch.setenv("OLLAMA_LOAD_TIMEOUT", value)
+    assert (
+        handler._warmup_timeout_seconds()
+        == handler._WARMUP_FALLBACK_S + handler._WARMUP_MARGIN_S
+    )
+
+
+# --- warm_model: the load request itself ----------------------------------------
+
+
+def test_warm_model_sends_empty_prompt_load(monkeypatch):
+    monkeypatch.setenv("OLLAMA_LOAD_TIMEOUT", "10m")
+    session = FakeSession(
+        FakeResponse(payload={"done": True}), ps_response=loaded_ps("llama3.2:3b")
+    )
+    monkeypatch.setattr(handler, "session", session)
+
+    handler.warm_model("llama3.2:3b")
+
+    url, kwargs = session.posts[0]
+    assert url.endswith("/api/generate")
+    # An empty prompt is Ollama's "load into memory, generate nothing" request.
+    assert kwargs["json"] == {"model": "llama3.2:3b", "prompt": "", "stream": False}
+    assert kwargs["timeout"] == 600 + handler._WARMUP_MARGIN_S
+
+
+def test_warm_model_raises_on_server_error(monkeypatch):
+    session = FakeSession(
+        FakeResponse(
+            status_code=500,
+            payload={"error": "timed out waiting for llama runner to start"},
+        )
+    )
+    monkeypatch.setattr(handler, "session", session)
+
+    with pytest.raises(ValueError, match="timed out waiting for llama runner"):
+        handler.warm_model("llama3.2:3b")
+
+
+def test_warm_model_warns_on_cpu_offload(monkeypatch, capsys):
+    session = FakeSession(
+        FakeResponse(payload={"done": True}),
+        ps_response=loaded_ps("big-model:latest", size=1000, size_vram=400),
+    )
+    monkeypatch.setattr(handler, "session", session)
+
+    handler.warm_model("big-model")
+
+    out = capsys.readouterr().out
+    assert "offloaded to CPU" in out
+
+
+def test_warm_model_reports_full_residency(monkeypatch, capsys):
+    session = FakeSession(
+        FakeResponse(payload={"done": True}), ps_response=loaded_ps("llama3.2:3b")
+    )
+    monkeypatch.setattr(handler, "session", session)
+
+    handler.warm_model("llama3.2:3b")
+
+    assert "resident in GPU memory" in capsys.readouterr().out
+
+
+def test_warm_model_survives_ps_failure(monkeypatch, capsys):
+    """A broken /api/ps must not turn a successful load into a failure."""
+    session = FakeSession(FakeResponse(payload={"done": True}))
+
+    def broken_get(url, **kwargs):
+        raise requests.ConnectionError("ps is down")
+
+    session.get = broken_get
+    monkeypatch.setattr(handler, "session", session)
+
+    handler.warm_model("llama3.2:3b")  # must not raise
+
+    assert "could not read /api/ps" in capsys.readouterr().out
+
+
+# --- warm_default_model: non-fatal, but says what failed -------------------------
+
+
+def test_warm_default_model_success(monkeypatch):
+    monkeypatch.setattr(handler, "resolve_default_model", lambda: "llama3.2:3b")
+    warmed = []
+    monkeypatch.setattr(handler, "warm_model", warmed.append)
+
+    assert handler.warm_default_model() is True
+    assert warmed == ["llama3.2:3b"]
+
+
+@pytest.mark.parametrize(
+    "err",
+    [
+        requests.ConnectionError("connection refused"),
+        requests.Timeout("read timed out"),
+        ValueError("HTTP 500 from /api/generate: model requires more system memory"),
+    ],
+)
+def test_warm_default_model_failure_is_nonfatal_but_loud(monkeypatch, capsys, err):
+    monkeypatch.setattr(handler, "resolve_default_model", lambda: "big-model")
+
+    def failing_warm(model):
+        raise err
+
+    monkeypatch.setattr(handler, "warm_model", failing_warm)
+
+    assert handler.warm_default_model() is False
+    out = capsys.readouterr().out
+    assert "could not pre-load 'big-model'" in out
+    # The cause must be surfaced, not swallowed behind a bare WARN.
+    assert str(err) in out