vllm: discover the share an engine needs instead of trusting a constant (v0.2.27)

The 0.60 shipped in 0.2.26 was measured from one model on one card, which is the same
mistake as the 0.35 before it: a footprint is not knowable up front. It depends on the
model, the card's size, the context and the multimodal limits vLLM profiles for — and
the 0.2.26 boot proved the point, since capping max_num_batched_tokens moved the
footprint only 11.13 -> 10.88 GiB (all of it CUDA graphs, 0.30 -> 0.05) while the
~9.4 GiB peak stayed put: vLLM sizes surya-ocr-2's encoder cache at 114,688 tokens and
profiles it with one VIDEO item at maximum feature size.

So the share is now an opening bid and the launcher measures the rest:

- ensure_service escalates on the KV-exhaustion signature (_KV_EXHAUSTED: "No available
  memory for the cache blocks", a negative "Available KV cache memory", or vLLM's
  free-memory complaint), +0.12 per attempt to at most 4 attempts, capped at the 0.92
  ceiling, asking the evictor for the larger prelaunch figure before each retry.
- An unrelated failure (missing model, broken venv) is NOT retried — escalating past it
  would burn minutes of boot per attempt for nothing.
- What booted is persisted to ~/.coderai/vllm_gmu_learned.json keyed by
  model|card-total|ctx|max_num_batched_tokens, so later starts skip rediscovery, a
  different model/card/limit set discovers its own, and a configured share LOWER than
  what is known to work is ignored rather than reintroducing a dead budget.
- A read-only config dir costs rediscovery, not a failure.

The multimodal limits are deliberately NOT capped by default: doing so would stop that
instance serving video at all, and the point is that any request type keeps working.
docs/ocr.md records the measurement and how to opt in via vllm.extra_args.

10 new tests covering escalation, the step, the ceiling, non-retryable failures,
persistence, scoping and an unwritable store.
Co-Authored-By: 's avatarClaude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0182DhR1eFNPQyGmDrpHbedY
parent f25592dc
......@@ -16,7 +16,7 @@
# Canonical product version for CoderAI — single source of truth. Both the API
# metadata and the admin web UI read from here.
__version__ = "0.2.26"
__version__ = "0.2.27"
__tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
# Configure the CUDA caching allocator BEFORE torch is imported anywhere.
......
......@@ -20,7 +20,9 @@ the OCR subsystem to serve Surya2 (a VLM) for :mod:`codai.ocr.surya`.
"""
import collections
import json
import os
import re
import shlex
import socket
import subprocess
......@@ -264,13 +266,13 @@ def _start_ray(cfg, py: str, env: dict):
return cluster, address
def ensure_service(cfg, model_path: Optional[str] = None,
served_name: Optional[str] = None,
ready_timeout: float = 3600.0,
gpu_memory_utilization: Optional[float] = None,
max_model_len: Optional[int] = None,
max_num_batched_tokens: Optional[int] = None,
max_num_seqs: Optional[int] = None) -> str:
def _ensure_service_once(cfg, model_path: Optional[str] = None,
served_name: Optional[str] = None,
ready_timeout: float = 3600.0,
gmu_absolute: float = 0.0,
max_model_len: Optional[int] = None,
max_num_batched_tokens: Optional[int] = None,
max_num_seqs: Optional[int] = None) -> str:
"""Launch (or reuse) a vLLM OpenAI server for a model; return its base URL.
``model_path``/``served_name`` override the config (used by the OCR subsystem to serve
......@@ -316,10 +318,11 @@ def ensure_service(cfg, model_path: Optional[str] = None,
port = int(getattr(cfg, "port", 0) or 0) or _free_port()
url_host = "127.0.0.1" if host in ("0.0.0.0", "") else host
url = f"http://{url_host}:{port}"
# Resolve the share HERE and remember it: what this instance actually reserved is
# what stopping it frees, and the eviction path has to be told the truth or it
# over- or under-evicts for the next model.
gmu = resolve_gmu(cfg, gpu_memory_utilization)
# The share is resolved by the caller (ensure_service, which may be escalating
# after a too-small budget) and remembered here: what this instance actually
# reserved is what stopping it frees, and the eviction path has to be told the
# truth or it over- or under-evicts for the next model.
gmu = float(gmu_absolute or 0.0)
cmd = _launch_cmd(py, cfg, host, port, model, served_name,
model_config=_model_config_for(resolved or model),
gpu_memory_utilization=gmu,
......@@ -394,6 +397,150 @@ def ensure_service(cfg, model_path: Optional[str] = None,
raise RuntimeError(f"vLLM for {svc_key} did not become ready in time" + _tail_msg())
# A budget too small for the engine's own footprint. vLLM says this whatever the model,
# card or context is, which is why the response is to measure rather than to guess.
_KV_EXHAUSTED = re.compile(
r"No available memory for the cache blocks"
r"|Available KV cache memory: *-"
r"|less than desired GPU memory utilization",
re.IGNORECASE)
# How much to add per attempt, and the most to try. Each attempt costs a real boot
# (minutes for a VLM), so the steps are coarse.
_GMU_ESCALATION_STEP = 0.12
_GMU_MAX_ATTEMPTS = 4
def _learned_gmu_path():
try:
from codai.platform_paths import legacy_style_config_dir
return legacy_style_config_dir() / "vllm_gmu_learned.json"
except Exception:
return None
def _learned_gmu_key(cfg, model: str, ctx, mnbt) -> str:
"""What a learned share is valid FOR.
The footprint depends on the model, how much memory the card has, and the limits it
was profiled under — so a value learned for surya-ocr-2 on a 24 GB card at ctx 18432
says nothing about olmOCR-2, or about the same model on a 48 GB card. Keying on all
of it is what keeps this generic instead of a number that happened to fit once."""
try:
from codai.models.manager import multi_model_manager
total = round(float(multi_model_manager._total_vram_gb() or 0.0), 1)
except Exception:
total = 0.0
return f"{model}|{total}|{int(ctx or 0)}|{int(mnbt or 0)}"
def _load_learned_gmu(key: str) -> float:
p = _learned_gmu_path()
if not p:
return 0.0
try:
with open(p) as f:
return float((json.load(f) or {}).get(key) or 0.0)
except Exception:
return 0.0
def _save_learned_gmu(key: str, gmu: float) -> None:
"""Remember the share that actually booted, so the next start does not re-discover it.
Best-effort: a read-only config dir costs a few minutes of re-discovery, not a
failure."""
p = _learned_gmu_path()
if not p or gmu <= 0:
return
try:
data = {}
try:
with open(p) as f:
data = json.load(f) or {}
except Exception:
data = {}
if abs(float(data.get(key) or 0.0) - gmu) < 0.005:
return
data[key] = round(float(gmu), 3)
tmp = str(p) + ".tmp"
with open(tmp, "w") as f:
json.dump(data, f, indent=2, sort_keys=True)
os.replace(tmp, p)
print(f"[vllm] learned gpu-memory-utilization {gmu:.3f} for {key}", flush=True)
except Exception as e:
print(f"[vllm] could not persist the learned share: {e}", flush=True)
def ensure_service(cfg, model_path: Optional[str] = None,
served_name: Optional[str] = None,
ready_timeout: float = 3600.0,
gpu_memory_utilization: Optional[float] = None,
max_model_len: Optional[int] = None,
max_num_batched_tokens: Optional[int] = None,
max_num_seqs: Optional[int] = None) -> str:
"""Launch (or reuse) a vLLM OpenAI server for a model; return its base URL.
The configured share is a STARTING POINT, not a verdict. An engine's real footprint
is not knowable up front — it depends on the model, the card, the context and the
multimodal limits it profiles for (surya-ocr-2 on a 24 GB 3090 needs ~10.9 GiB before
a single KV block, almost all of it the vision encoder's profiling peak; another model
or card is another number). So when vLLM reports the budget cannot hold the cache,
this escalates the share and tries again, asks the evictor for the larger figure
first, and REMEMBERS what worked, keyed by model+card+limits. That way any model,
on any card, serving any kind of request, converges on a share that fits instead of
depending on a constant somebody measured once.
``model_path``/``served_name`` override the config (used by the OCR subsystem to serve
surya-2 on its own vLLM instance alongside any LLM instance)."""
start = resolve_gmu(cfg, gpu_memory_utilization)
ctx = int(max_model_len or 0) or int(getattr(cfg, "ctx", 0) or 0)
key = _learned_gmu_key(cfg, (served_name or model_path or "vllm"), ctx,
max_num_batched_tokens)
learned = _load_learned_gmu(key)
if learned > start:
print(f"[vllm] starting at the learned share {learned:.3f} rather than "
f"{start:.3f} (it is what booted last time for {key})", flush=True)
start = learned
if start <= 0:
# No share configured anywhere: nothing to escalate, let vLLM use its own default.
return _ensure_service_once(
cfg, model_path, served_name, ready_timeout, 0.0,
max_model_len, max_num_batched_tokens, max_num_seqs)
gmu = min(start, GMU_CEILING)
last = None
for attempt in range(1, _GMU_MAX_ATTEMPTS + 1):
try:
url = _ensure_service_once(
cfg, model_path, served_name, ready_timeout, gmu,
max_model_len, max_num_batched_tokens, max_num_seqs)
except RuntimeError as exc:
last = exc
if not _KV_EXHAUSTED.search(str(exc)) or gmu >= GMU_CEILING - 1e-6:
raise
nxt = min(gmu + _GMU_ESCALATION_STEP, GMU_CEILING)
if nxt <= gmu + 1e-6:
raise
print(f"[vllm] {served_name or model_path}: a {gmu:.3f} budget cannot hold "
f"the KV cache (attempt {attempt}/{_GMU_MAX_ATTEMPTS}); retrying at "
f"{nxt:.3f} — the engine's footprint is larger than the configured "
f"share, which is what this measures", flush=True)
# The bigger budget needs the room to exist before vLLM asks for it.
try:
from codai.models.manager import multi_model_manager
need = prelaunch_free_gb(cfg, nxt)
multi_model_manager._evict_models_for_vram(float(need))
except Exception as e:
print(f"[vllm] evict-before-retry skipped: {e}", flush=True)
gmu = nxt
continue
if attempt > 1 or abs(gmu - learned) > 0.005:
_save_learned_gmu(key, gmu)
return url
raise last if last is not None else RuntimeError("vLLM did not start")
def stop_service(svc_key: str) -> None:
with _lock:
svc = _services.pop(svc_key, None)
......
......@@ -785,14 +785,18 @@ class OcrConfig:
# model, so OCR and models take turns at a page boundary instead of OCR losing the
# page it was mid-way through. 0 = tear down immediately.
evict_drain_timeout_s: float = 60.0
# The share of the card a VLM OCR engine's own vLLM instance claims (surya in vllm
# mode, olmocr in vllm mode). 0 = use vllm.gpu_memory_utilization, which is sized for
# serving an LLM. This is a fraction of the card's TOTAL memory, and it has to cover
# the engine's whole footprint, not just its weights: measured live on a 24 GB 3090,
# surya-ocr-2 needs ~11.1 GB before a single KV block (weights 1.37 + CUDA graphs 0.30
# + ~9.4 GB activation peak). 0.35 (8.4 GB) was therefore below the floor and could
# never start; 0.60 (14.4 GB) leaves ~3 GB of KV cache. The evict-first pass clears
# the card for it, so the figure does not have to allow for other residents.
# The STARTING share of the card a VLM OCR engine's own vLLM instance asks for
# (surya in vllm mode, olmocr in vllm mode); 0 = use vllm.gpu_memory_utilization.
# It is a fraction of the card's TOTAL memory and has to cover the engine's whole
# footprint, not just its weights: measured live, surya-ocr-2 on a 24 GB 3090 needs
# ~10.9 GB before a single KV block (weights 1.37 + CUDA graphs 0.05 + ~9.4 GB of
# vision-encoder profiling peak), so the 0.35 this once shipped (8.4 GB) was under
# the floor and could never start. 0.60 is a reasonable opening bid for a 7B VLM —
# but it is NOT a measurement of your model on your card, and it does not have to be:
# when vLLM reports the budget cannot hold the cache, the launcher escalates the
# share, re-evicts for the larger figure and remembers what booted, keyed by
# model+card+limits (see vllm_worker.ensure_service). Any model, any card, any
# request mix converges on a share that fits.
vlm_gpu_memory_utilization: float = 0.60
# …and these bring the footprint itself down, which is the better lever. A VLM OCR
# engine otherwise inherits the LLM backend's ctx (18432 here), and the activation
......
......@@ -126,23 +126,31 @@ refuses to start when that much is not free. Three rules follow, all enforced by
- **Room is made before the load, not after it.** The engine declares what it needs up
front (`prelaunch_vram_gb`) and the model manager evicts for it first. Evicting
afterwards is useless, because the load is what fails.
- **The share must cover the engine's whole footprint, not its weights.** Measured live
on a 24 GB 3090 (two boots of surya-ocr-2: a 9.57 GiB budget gave
`Available KV cache memory: -1.54 GiB`, a 14.31 GiB budget gave `+3.18 GiB`), the fixed
cost is **~11.1 GiB** — weights 1.37, CUDA graphs 0.30 and ~9.4 GiB of activation peak —
*before a single KV block*. `vlm_gpu_memory_utilization` is a fraction of the card's
TOTAL memory, so 0.35 of this card is 8.4 GiB: below the floor, and no amount of
eviction makes it start. Hence the **0.60** default (14.4 GiB → ~3 GiB of KV cache).
Nothing is added back for other residents — those two boots show vLLM charges only its
own allocations against `total × gmu` — so the figure is simply the share, and the
evict-first pass above is what guarantees it is free.
- **The batched-token cap is the better lever.** That 9.4 GiB peak is profiled for the
batch size, and a VLM OCR engine otherwise inherits the *LLM* backend's `ctx` (18432
here). Chunked prefill is on, so `ocr.vlm_max_num_batched_tokens` (default 4096) and
`ocr.vlm_max_num_seqs` (16) shrink the peak **without** shrinking the context a page may
use; `ocr.vlm_max_model_len` (0 = inherit) is there if you want to cut that too. Lower
the peak enough and the engine fits in a far smaller share, which is what lets OCR and
an LLM sit on one card.
- **The share is a starting point, and the launcher measures the rest.** An engine's real
footprint is not knowable up front. Measured live on a 24 GB 3090, surya-ocr-2 needs
**~10.9 GiB before a single KV block** (weights 1.37, CUDA graphs 0.05, and ~9.4 GiB of
*vision-encoder profiling* peak); `vlm_gpu_memory_utilization` is a fraction of the
card's TOTAL memory, so the 0.35 this once shipped was 8.4 GiB — under the floor, unable
to start whatever else was resident. Rather than trusting a constant, when vLLM reports
that the budget cannot hold the cache the launcher **escalates the share, re-evicts for
the larger figure, retries, and records what booted**, keyed by model + card + context +
batch limits (`~/.coderai/vllm_gmu_learned.json`). A different model, a bigger card or
different limits rediscover their own value; the recorded one is reused on later starts,
and a configured share lower than what is known to work is ignored. Nothing is added
back for other residents — vLLM charges only its own allocations against `total × gmu` —
so the evict-first pass is what guarantees the room.
- **The batch limits shrink the footprint a little; the encoder is what dominates.**
`ocr.vlm_max_num_batched_tokens` (default 4096), `ocr.vlm_max_num_seqs` (16) and
`ocr.vlm_max_model_len` (0 = inherit) stop a VLM OCR engine inheriting the *LLM*
backend's sizing, and chunked prefill means capping tokens-per-batch costs no usable
context. Measured effect on surya-ocr-2, though: the footprint fell only 11.13 → 10.88
GiB, **all of it from capturing fewer CUDA graphs** (0.30 → 0.05). The ~9.4 GiB peak did
not move, because for this model vLLM sizes its encoder cache at 114,688 tokens and
profiles it *with one video item at maximum feature size*. Capping the multimodal limits
would collapse that — and would also stop the instance serving video at all, so it is
not done by default. Set it yourself via `vllm.extra_args` if that instance only ever
sees still pages.
If you see `Free memory on device cuda:0 … is less than desired GPU memory utilization`,
set **`ocr.vlm_gpu_memory_utilization`** (Settings → OCR → "VLM engines' GPU share"): the
......
"""The configured vLLM share is a starting point, not a verdict.
An engine's real footprint is not knowable up front: surya-ocr-2 on a 24 GB 3090 needs
~10.9 GiB before a single KV block, almost all of it the vision encoder's profiling peak,
and another model or another card is another number entirely. Guessing a constant is what
produced the 0.35 that could never boot — so when vLLM reports the budget cannot hold the
cache, the launcher escalates, re-evicts for the larger figure, and remembers what worked
keyed by model+card+limits.
No vLLM, no GPU: the single-attempt launcher is stubbed and the learned-share file is
redirected into tmp_path.
"""
import sys
from pathlib import Path
import pytest
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from codai.api import vllm_worker
KV_ERR = ("vLLM exited (code 1) before becoming ready. Last output: "
"ValueError: No available memory for the cache blocks. Try increasing "
"`gpu_memory_utilization` when initializing the engine")
@pytest.fixture(autouse=True)
def _isolate(monkeypatch, tmp_path):
"""Learned shares go to tmp, the card is a 24 GB one, eviction is a no-op."""
monkeypatch.setattr(vllm_worker, "_learned_gmu_path",
lambda: tmp_path / "vllm_gmu_learned.json")
from codai.models.manager import multi_model_manager
monkeypatch.setattr(multi_model_manager, "_total_vram_gb", lambda: 24.0)
monkeypatch.setattr(multi_model_manager, "_get_free_vram_gb", lambda: 24.0)
monkeypatch.setattr(multi_model_manager, "_evict_models_for_vram", lambda gb: None)
def _cfg(**kw):
base = {"model_path": "m", "ctx": 18432, "gpu_memory_utilization": 0.0}
base.update(kw)
return type("C", (), base)()
def _launcher(monkeypatch, succeed_at: float = None, error=KV_ERR):
"""Stub the one-shot launcher: record every share tried, fail until `succeed_at`."""
tried = []
def _once(cfg, model_path=None, served_name=None, ready_timeout=3600.0,
gmu_absolute=0.0, max_model_len=None, max_num_batched_tokens=None,
max_num_seqs=None):
tried.append(round(float(gmu_absolute), 3))
if succeed_at is not None and gmu_absolute >= succeed_at - 1e-6:
return "http://127.0.0.1:1/"
raise RuntimeError(error)
monkeypatch.setattr(vllm_worker, "_ensure_service_once", _once)
return tried
# ---------------------------------------------------------------- escalation
def test_a_budget_too_small_for_the_cache_is_retried_larger(monkeypatch):
tried = _launcher(monkeypatch, succeed_at=0.72)
url = vllm_worker.ensure_service(_cfg(), model_path="m", served_name="m",
gpu_memory_utilization=0.60)
assert url == "http://127.0.0.1:1/"
assert tried[0] == 0.60 # the configured opening bid
assert tried == sorted(tried) and len(tried) >= 2
assert tried[-1] >= 0.72 # climbed until it fit
def test_it_escalates_by_the_configured_step(monkeypatch):
tried = _launcher(monkeypatch, succeed_at=0.99) # never fits below the ceiling
with pytest.raises(RuntimeError):
vllm_worker.ensure_service(_cfg(), model_path="m", served_name="m",
gpu_memory_utilization=0.50)
assert tried[0] == 0.50
assert tried[1] == pytest.approx(0.50 + vllm_worker._GMU_ESCALATION_STEP, abs=1e-3)
def test_escalation_stops_at_the_ceiling_and_reraises(monkeypatch):
tried = _launcher(monkeypatch, succeed_at=None)
with pytest.raises(RuntimeError) as e:
vllm_worker.ensure_service(_cfg(), model_path="m", served_name="m",
gpu_memory_utilization=0.60)
assert "No available memory" in str(e.value)
assert max(tried) <= vllm_worker.GMU_CEILING + 1e-6
assert len(tried) <= vllm_worker._GMU_MAX_ATTEMPTS
def test_an_unrelated_failure_is_not_retried(monkeypatch):
"""Escalating past a missing model or a broken venv would just waste boots."""
tried = _launcher(monkeypatch, succeed_at=None,
error="vLLM exited (code 1). Last output: No such file: /nope")
with pytest.raises(RuntimeError):
vllm_worker.ensure_service(_cfg(), model_path="m", served_name="m",
gpu_memory_utilization=0.60)
assert tried == [0.60] # one attempt, then give up
def test_a_first_attempt_that_works_is_left_alone(monkeypatch):
tried = _launcher(monkeypatch, succeed_at=0.0)
vllm_worker.ensure_service(_cfg(), model_path="m", served_name="m",
gpu_memory_utilization=0.60)
assert tried == [0.60]
def test_no_share_anywhere_means_one_attempt_on_vllms_own_default(monkeypatch):
tried = _launcher(monkeypatch, succeed_at=0.0)
vllm_worker.ensure_service(_cfg(), model_path="m", served_name="m")
assert tried == [0.0]
# ---------------------------------------------------------------- learning
def test_what_booted_is_remembered_and_reused(monkeypatch, tmp_path):
tried = _launcher(monkeypatch, succeed_at=0.72)
vllm_worker.ensure_service(_cfg(), model_path="m", served_name="m",
gpu_memory_utilization=0.60)
assert (tmp_path / "vllm_gmu_learned.json").exists()
# Second start: no rediscovery, straight to the share that worked.
tried2 = _launcher(monkeypatch, succeed_at=0.72)
vllm_worker.ensure_service(_cfg(), model_path="m", served_name="m",
gpu_memory_utilization=0.60)
assert len(tried2) == 1 and tried2[0] >= 0.72
def test_a_learned_share_is_scoped_to_model_card_and_limits(monkeypatch):
"""A figure learned for one model on one card says nothing about another."""
_launcher(monkeypatch, succeed_at=0.72)
vllm_worker.ensure_service(_cfg(), model_path="m", served_name="surya",
gpu_memory_utilization=0.60)
tried = _launcher(monkeypatch, succeed_at=0.72)
vllm_worker.ensure_service(_cfg(), model_path="m", served_name="olmocr",
gpu_memory_utilization=0.60)
assert tried[0] == 0.60 # a different model rediscovers from the opening bid
k_surya = vllm_worker._learned_gmu_key(_cfg(), "surya", 18432, 4096)
k_ctx = vllm_worker._learned_gmu_key(_cfg(), "surya", 8192, 4096)
k_batch = vllm_worker._learned_gmu_key(_cfg(), "surya", 18432, 2048)
assert len({k_surya, k_ctx, k_batch}) == 3 # context and batch size are part of it
def test_a_lower_configured_share_never_undoes_what_was_learned(monkeypatch):
"""Turning the knob down must not reintroduce a budget already proven too small."""
_launcher(monkeypatch, succeed_at=0.72)
vllm_worker.ensure_service(_cfg(), model_path="m", served_name="m",
gpu_memory_utilization=0.60)
tried = _launcher(monkeypatch, succeed_at=0.72)
vllm_worker.ensure_service(_cfg(), model_path="m", served_name="m",
gpu_memory_utilization=0.20)
assert tried[0] >= 0.72
def test_an_unwritable_store_costs_rediscovery_not_a_failure(monkeypatch):
monkeypatch.setattr(vllm_worker, "_learned_gmu_path",
lambda: Path("/proc/nope/vllm_gmu_learned.json"))
tried = _launcher(monkeypatch, succeed_at=0.72)
url = vllm_worker.ensure_service(_cfg(), model_path="m", served_name="m",
gpu_memory_utilization=0.60)
assert url == "http://127.0.0.1:1/" and len(tried) >= 2
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment