ocr: size the VLM share to the measured footprint and cap its batch (v0.2.26)

Live measurement on the 24 GB 3090 settled what the 0.35 share could never do. Two
boots of surya-ocr-2: a 9.57 GiB budget reported `Available KV cache memory: -1.54 GiB`,
a 14.31 GiB budget reported `+3.18 GiB` — a fixed footprint of 11.11 and 11.13 GiB,
i.e. ~11.1 GiB (weights 1.37 + CUDA graphs 0.30 + ~9.4 GiB activation peak) before a
single KV block. 0.35 of this card is 8.4 GiB: under the floor, so it could not start
no matter what else was resident, and the 06:41 boot that finally succeeded did so only
because 6.1 GB held by other models inflated the old add-back to an absolute 0.604.

Those same two data points show vLLM charges only its OWN allocations against
total × gmu — if other residents counted, they would imply activation peaks of 8.3 and
3.6 GiB rather than one consistent 11.1. So:

- effective_gmu no longer adds back what others hold. It made the budget depend on
  unrelated residents and asked for share*total for this instance while `used` stayed
  held by someone else — at 0.60 that would have been 20+ GiB of a 24 GB card. The
  evict-first pass (prelaunch_free_gb) is what clears the card, which is the honest
  mechanism; the log line now states the share, its GB and what else is resident.
- vlm_gpu_memory_utilization default 0.35 -> 0.60 (14.4 GiB, ~3 GiB of KV cache).
- New ocr.vlm_max_num_batched_tokens (4096), vlm_max_num_seqs (16) and vlm_max_model_len
  (0 = inherit), plumbed through ensure_service/_launch_cmd via base.vlm_serve_limits().
  This is the better lever: the 9.4 GiB peak is profiled for the batch size, and a VLM
  OCR engine otherwise inherits the LLM backend's ctx=18432. Chunked prefill is on, so
  capping tokens per batch shrinks the peak without shrinking the usable context — lower
  it enough and OCR fits beside an LLM instead of needing 60% of the card.

9 new tests (the caps, the launch command, the footprint floor); the share tests that
asserted the add-back now assert its absence, with the reason.
Co-Authored-By: 's avatarClaude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0182DhR1eFNPQyGmDrpHbedY
parent 3cb4d23d
...@@ -16,7 +16,7 @@ ...@@ -16,7 +16,7 @@
# Canonical product version for CoderAI — single source of truth. Both the API # Canonical product version for CoderAI — single source of truth. Both the API
# metadata and the admin web UI read from here. # metadata and the admin web UI read from here.
__version__ = "0.2.25" __version__ = "0.2.26"
__tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference." __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
# Configure the CUDA caching allocator BEFORE torch is imported anywhere. # Configure the CUDA caching allocator BEFORE torch is imported anywhere.
......
...@@ -3873,6 +3873,9 @@ def build_settings_dict(c, gpu_cards): ...@@ -3873,6 +3873,9 @@ def build_settings_dict(c, gpu_cards):
"build_retry_cooldown_s": c.ocr.build_retry_cooldown_s, "build_retry_cooldown_s": c.ocr.build_retry_cooldown_s,
"evict_drain_timeout_s": c.ocr.evict_drain_timeout_s, "evict_drain_timeout_s": c.ocr.evict_drain_timeout_s,
"vlm_gpu_memory_utilization": c.ocr.vlm_gpu_memory_utilization, "vlm_gpu_memory_utilization": c.ocr.vlm_gpu_memory_utilization,
"vlm_max_num_batched_tokens": c.ocr.vlm_max_num_batched_tokens,
"vlm_max_num_seqs": c.ocr.vlm_max_num_seqs,
"vlm_max_model_len": c.ocr.vlm_max_model_len,
"fallback_engines": c.ocr.fallback_engines, "fallback_engines": c.ocr.fallback_engines,
"olmocr_enabled": c.ocr.olmocr_enabled, "olmocr_enabled": c.ocr.olmocr_enabled,
"olmocr_instances": c.ocr.olmocr_instances, "olmocr_instances": c.ocr.olmocr_instances,
...@@ -4606,6 +4609,9 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad ...@@ -4606,6 +4609,9 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad
if "build_retry_cooldown_s" in d: o.build_retry_cooldown_s = _f("build_retry_cooldown_s", o.build_retry_cooldown_s, 0.0) if "build_retry_cooldown_s" in d: o.build_retry_cooldown_s = _f("build_retry_cooldown_s", o.build_retry_cooldown_s, 0.0)
if "vlm_gpu_memory_utilization" in d: o.vlm_gpu_memory_utilization = _f("vlm_gpu_memory_utilization", o.vlm_gpu_memory_utilization, 0.0) if "vlm_gpu_memory_utilization" in d: o.vlm_gpu_memory_utilization = _f("vlm_gpu_memory_utilization", o.vlm_gpu_memory_utilization, 0.0)
if "fallback_engines" in d: o.fallback_engines = (d.get("fallback_engines") or "auto").strip() if "fallback_engines" in d: o.fallback_engines = (d.get("fallback_engines") or "auto").strip()
if "vlm_max_num_batched_tokens" in d: o.vlm_max_num_batched_tokens = _i("vlm_max_num_batched_tokens", o.vlm_max_num_batched_tokens, 0)
if "vlm_max_num_seqs" in d: o.vlm_max_num_seqs = _i("vlm_max_num_seqs", o.vlm_max_num_seqs, 0)
if "vlm_max_model_len" in d: o.vlm_max_model_len = _i("vlm_max_model_len", o.vlm_max_model_len, 0)
if "evict_drain_timeout_s" in d: o.evict_drain_timeout_s = _f("evict_drain_timeout_s", o.evict_drain_timeout_s, 0.0) if "evict_drain_timeout_s" in d: o.evict_drain_timeout_s = _f("evict_drain_timeout_s", o.evict_drain_timeout_s, 0.0)
if "olmocr_enabled" in d: o.olmocr_enabled = bool(d["olmocr_enabled"]) if "olmocr_enabled" in d: o.olmocr_enabled = bool(d["olmocr_enabled"])
if "olmocr_instances" in d: o.olmocr_instances = _i("olmocr_instances", o.olmocr_instances, 1) if "olmocr_instances" in d: o.olmocr_instances = _i("olmocr_instances", o.olmocr_instances, 1)
......
...@@ -1074,6 +1074,21 @@ ...@@ -1074,6 +1074,21 @@
<label class="form-label">VLM engines' GPU share</label> <label class="form-label">VLM engines' GPU share</label>
<input type="number" id="s-ocr-vlm-gmu" class="form-input" min="0" max="1" step="0.05" placeholder="0 = vLLM default"> <input type="number" id="s-ocr-vlm-gmu" class="form-input" min="0" max="1" step="0.05" placeholder="0 = vLLM default">
</div> </div>
<div class="form-group" style="flex:1;min-width:160px">
<label class="form-label">VLM batched tokens</label>
<input type="number" id="s-ocr-vlm-mnbt" class="form-input" min="0" placeholder="4096">
<div class="form-hint">Caps the tokens per batch for a VLM engine's own vLLM.
This is what sets its activation peak — the reason a share that looks big
enough still cannot start. 0 = inherit the vLLM backend's.</div>
</div>
<div class="form-group" style="flex:1;min-width:140px">
<label class="form-label">VLM max seqs</label>
<input type="number" id="s-ocr-vlm-mns" class="form-input" min="0" placeholder="16">
</div>
<div class="form-group" style="flex:1;min-width:150px">
<label class="form-label">VLM max model len</label>
<input type="number" id="s-ocr-vlm-mml" class="form-input" min="0" placeholder="0 = inherit">
</div>
<div class="form-group" style="flex:1;min-width:150px"> <div class="form-group" style="flex:1;min-width:150px">
<label class="form-label">Evict drain (s)</label> <label class="form-label">Evict drain (s)</label>
<input type="number" id="s-ocr-evict-drain" class="form-input" min="0" placeholder="60"> <input type="number" id="s-ocr-evict-drain" class="form-input" min="0" placeholder="60">
...@@ -2017,6 +2032,9 @@ async function loadSettings(){ ...@@ -2017,6 +2032,9 @@ async function loadSettings(){
document.getElementById('s-ocr-retry-cooldown').value = ocr.build_retry_cooldown_s ?? 60; document.getElementById('s-ocr-retry-cooldown').value = ocr.build_retry_cooldown_s ?? 60;
document.getElementById('s-ocr-vlm-gmu').value = ocr.vlm_gpu_memory_utilization ?? 0; document.getElementById('s-ocr-vlm-gmu').value = ocr.vlm_gpu_memory_utilization ?? 0;
document.getElementById('s-ocr-fallback').value = ocr.fallback_engines ?? 'auto'; document.getElementById('s-ocr-fallback').value = ocr.fallback_engines ?? 'auto';
document.getElementById('s-ocr-vlm-mnbt').value = ocr.vlm_max_num_batched_tokens ?? 4096;
document.getElementById('s-ocr-vlm-mns').value = ocr.vlm_max_num_seqs ?? 16;
document.getElementById('s-ocr-vlm-mml').value = ocr.vlm_max_model_len ?? 0;
document.getElementById('s-ocr-evict-drain').value = ocr.evict_drain_timeout_s ?? 60; document.getElementById('s-ocr-evict-drain').value = ocr.evict_drain_timeout_s ?? 60;
document.getElementById('s-ocr-paddle-enabled').checked = ocr.paddle_enabled !== false; document.getElementById('s-ocr-paddle-enabled').checked = ocr.paddle_enabled !== false;
document.getElementById('s-ocr-paddle-use-gpu').checked = ocr.paddle_use_gpu !== false; document.getElementById('s-ocr-paddle-use-gpu').checked = ocr.paddle_use_gpu !== false;
...@@ -2246,6 +2264,9 @@ async function saveSettings(silent){ ...@@ -2246,6 +2264,9 @@ async function saveSettings(silent){
build_retry_cooldown_s: parseFloat(document.getElementById('s-ocr-retry-cooldown').value) || 0, build_retry_cooldown_s: parseFloat(document.getElementById('s-ocr-retry-cooldown').value) || 0,
vlm_gpu_memory_utilization: parseFloat(document.getElementById('s-ocr-vlm-gmu').value) || 0, vlm_gpu_memory_utilization: parseFloat(document.getElementById('s-ocr-vlm-gmu').value) || 0,
fallback_engines: document.getElementById('s-ocr-fallback').value.trim() || 'auto', fallback_engines: document.getElementById('s-ocr-fallback').value.trim() || 'auto',
vlm_max_num_batched_tokens: parseInt(document.getElementById('s-ocr-vlm-mnbt').value) || 0,
vlm_max_num_seqs: parseInt(document.getElementById('s-ocr-vlm-mns').value) || 0,
vlm_max_model_len: parseInt(document.getElementById('s-ocr-vlm-mml').value) || 0,
evict_drain_timeout_s: parseFloat(document.getElementById('s-ocr-evict-drain').value) || 0, evict_drain_timeout_s: parseFloat(document.getElementById('s-ocr-evict-drain').value) || 0,
paddle_enabled: document.getElementById('s-ocr-paddle-enabled').checked, paddle_enabled: document.getElementById('s-ocr-paddle-enabled').checked,
paddle_use_gpu: document.getElementById('s-ocr-paddle-use-gpu').checked, paddle_use_gpu: document.getElementById('s-ocr-paddle-use-gpu').checked,
......
...@@ -116,19 +116,32 @@ def resolve_service_key(cfg, model_path: Optional[str] = None): ...@@ -116,19 +116,32 @@ def resolve_service_key(cfg, model_path: Optional[str] = None):
def _launch_cmd(py, cfg, host: str, port: int, model_path: str, def _launch_cmd(py, cfg, host: str, port: int, model_path: str,
served_name: Optional[str] = None, model_config: dict = None, served_name: Optional[str] = None, model_config: dict = None,
gpu_memory_utilization: Optional[float] = None) -> list: 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) -> list:
mid = served_name or (getattr(cfg, "model_id", "vllm") or "vllm") mid = served_name or (getattr(cfg, "model_id", "vllm") or "vllm")
cmd = [py, "-m", "vllm.entrypoints.openai.api_server", cmd = [py, "-m", "vllm.entrypoints.openai.api_server",
"--host", host, "--port", str(port), "--host", host, "--port", str(port),
"--model", model_path, "--model", model_path,
"--served-model-name", mid] "--served-model-name", mid]
ctx = int(getattr(cfg, "ctx", 0) or 0) # A side job serving a small model (the OCR VLM engines) overrides the context the
# LLM backend is configured with: it inherited ctx=18432, and the activation peak
# that size profiles for was 9.4 GiB of the 11.1 GiB that made surya-ocr-2 unable to
# start. A page is not an 18k-token conversation.
ctx = int(max_model_len or 0) or int(getattr(cfg, "ctx", 0) or 0)
if ctx > 0: if ctx > 0:
cmd += ["--max-model-len", str(ctx)] cmd += ["--max-model-len", str(ctx)]
# Already absolute (resolve_gmu translated a side job's own-share if there was one). # Already absolute (resolve_gmu translated a side job's own-share if there was one).
gmu = float(gpu_memory_utilization or 0) or float(getattr(cfg, "gpu_memory_utilization", 0) or 0) gmu = float(gpu_memory_utilization or 0) or float(getattr(cfg, "gpu_memory_utilization", 0) or 0)
if gmu > 0: if gmu > 0:
cmd += ["--gpu-memory-utilization", str(gmu)] cmd += ["--gpu-memory-utilization", str(gmu)]
# The profiling peak that decides whether the engine can start is driven by how many
# tokens may be in one batch, NOT by max_model_len (chunked prefill is on), so this
# is the lever that shrinks the footprint without shrinking the usable context.
mnbt = int(max_num_batched_tokens or 0)
if mnbt > 0:
cmd += ["--max-num-batched-tokens", str(mnbt)]
tp = int(getattr(cfg, "tensor_parallel_size", 0) or 0) tp = int(getattr(cfg, "tensor_parallel_size", 0) or 0)
if tp > 0: if tp > 0:
cmd += ["--tensor-parallel-size", str(tp)] cmd += ["--tensor-parallel-size", str(tp)]
...@@ -140,7 +153,7 @@ def _launch_cmd(py, cfg, host: str, port: int, model_path: str, ...@@ -140,7 +153,7 @@ def _launch_cmd(py, cfg, host: str, port: int, model_path: str,
executor = "ray" executor = "ray"
if executor: if executor:
cmd += ["--distributed-executor-backend", executor] cmd += ["--distributed-executor-backend", executor]
mns = int(getattr(cfg, "max_num_seqs", 0) or 0) mns = int(max_num_seqs or 0) or int(getattr(cfg, "max_num_seqs", 0) or 0)
if mns > 0: if mns > 0:
cmd += ["--max-num-seqs", str(mns)] cmd += ["--max-num-seqs", str(mns)]
dtype = (getattr(cfg, "dtype", "") or "").strip() dtype = (getattr(cfg, "dtype", "") or "").strip()
...@@ -254,7 +267,10 @@ def _start_ray(cfg, py: str, env: dict): ...@@ -254,7 +267,10 @@ def _start_ray(cfg, py: str, env: dict):
def ensure_service(cfg, model_path: Optional[str] = None, def ensure_service(cfg, model_path: Optional[str] = None,
served_name: Optional[str] = None, served_name: Optional[str] = None,
ready_timeout: float = 3600.0, ready_timeout: float = 3600.0,
gpu_memory_utilization: Optional[float] = None) -> str: 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. """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 ``model_path``/``served_name`` override the config (used by the OCR subsystem to serve
...@@ -306,7 +322,10 @@ def ensure_service(cfg, model_path: Optional[str] = None, ...@@ -306,7 +322,10 @@ def ensure_service(cfg, model_path: Optional[str] = None,
gmu = resolve_gmu(cfg, gpu_memory_utilization) gmu = resolve_gmu(cfg, gpu_memory_utilization)
cmd = _launch_cmd(py, cfg, host, port, model, served_name, cmd = _launch_cmd(py, cfg, host, port, model, served_name,
model_config=_model_config_for(resolved or model), model_config=_model_config_for(resolved or model),
gpu_memory_utilization=gmu) gpu_memory_utilization=gmu,
max_model_len=max_model_len,
max_num_batched_tokens=max_num_batched_tokens,
max_num_seqs=max_num_seqs)
env = os.environ.copy() env = os.environ.copy()
# flashinfer JIT-compiles CUDA kernels with ninja at runtime, which fails on # flashinfer JIT-compiles CUDA kernels with ninja at runtime, which fails on
...@@ -437,9 +456,10 @@ def _used_by_others_gb() -> float: ...@@ -437,9 +456,10 @@ def _used_by_others_gb() -> float:
def resolve_gmu(cfg, gpu_memory_utilization: Optional[float] = None) -> float: def resolve_gmu(cfg, gpu_memory_utilization: Optional[float] = None) -> float:
"""The absolute ``--gpu-memory-utilization`` to launch with. """The absolute ``--gpu-memory-utilization`` to launch with.
A side job (an OCR VLM engine) passes the share it wants for ITSELF, which has to be A side job (an OCR VLM engine) passes the share it wants for ITSELF; the backend's
translated against what else is on the card; the backend's own configured value is own configured value is used as-is. Both are fractions of the card's TOTAL memory,
already absolute and is used as-is.""" which is what the flag means — see :func:`effective_gmu` for why nothing is added
back for other residents any more."""
share = float(gpu_memory_utilization or 0) share = float(gpu_memory_utilization or 0)
if share > 0: if share > 0:
return effective_gmu(share) return effective_gmu(share)
...@@ -447,40 +467,43 @@ def resolve_gmu(cfg, gpu_memory_utilization: Optional[float] = None) -> float: ...@@ -447,40 +467,43 @@ def resolve_gmu(cfg, gpu_memory_utilization: Optional[float] = None) -> float:
def effective_gmu(self_share: float) -> float: def effective_gmu(self_share: float) -> float:
"""Translate "the share of the card I want for MYSELF" into the absolute """The share this instance may use, clamped to :data:`GMU_CEILING`, plus a log line
``--gpu-memory-utilization`` vLLM wants. naming what else is on the card.
vLLM's flag is a fraction of the card's TOTAL memory and every other process' Measured on a 24 GB 3090 (2026-10-05, two live boots of surya-ocr-2): budget
allocation counts against it: with 3.6 GB of bge-m3 resident on a 24 GB 3090, 9.57 GiB gave ``Available KV cache memory: -1.54 GiB``, budget 14.31 GiB gave
``0.35`` left surya-2 about 5 GB, and after weights (1.37) + activation peak at ``+3.18 GiB`` — a fixed footprint of 11.11 and 11.13 GiB respectively. Two things
``max_model_len=18432`` + CUDA graphs (0.3) it reported ``Available KV cache memory: follow. First, vLLM charges only its OWN allocations against ``total × gmu``; if
-2.81 GiB`` and died — every 20 minutes, for hours. So add back what others hold: other residents counted too, those two boots would imply wildly different activation
0.35 of a card with 3.6 GB in use becomes 0.50, i.e. a 12 GB budget of which 8.4 GB peaks (8.3 vs 3.6 GiB) instead of the same 11.1. So this function does NOT add back
is actually ours. Capped at :data:`GMU_CEILING`; when the cap bites, the caller's what others hold: doing so made the budget depend on unrelated residents and asked
evict-first pass (``prelaunch_vram_gb``) is what makes room.""" for ``share×total + used`` for this instance while ``used`` stayed held by someone
else. The card is cleared by the caller's evict-first pass
(:func:`prelaunch_free_gb`) instead, which is the honest mechanism.
Second, the share has to actually cover the footprint: 0.35 of this card is 8.4 GiB
against an 11.1 GiB floor, so it could never start, which is why ``0.35`` crash-looped
for hours and then appeared to "work" only when 6.1 GB of other residents inflated
the old add-back to 0.604. Hence the 0.60 default and the batched-token cap that
brings the 9.4 GiB activation peak down."""
try: try:
share = max(0.0, float(self_share or 0.0)) share = max(0.0, float(self_share or 0.0))
except Exception: except Exception:
return 0.0 return 0.0
if share <= 0: if share <= 0:
return 0.0 return 0.0
eff = min(share, GMU_CEILING)
try: try:
from codai.models.manager import multi_model_manager from codai.models.manager import multi_model_manager
total = float(multi_model_manager._total_vram_gb() or 0.0) total = float(multi_model_manager._total_vram_gb() or 0.0)
except Exception: except Exception:
total = 0.0 total = 0.0
if total <= 0: if total > 0:
return min(share, GMU_CEILING) used = _used_by_others_gb()
used = _used_by_others_gb() print(f"[vllm] gpu-memory-utilization {eff:.3f} = {eff * total:.1f} GB of a "
eff = share + (used / total) f"{total:.1f} GB card for this instance; {used:.1f} GB currently held by "
if eff > GMU_CEILING: f"others (the evict-first pass should have cleared what it needed)",
print(f"[vllm] share {share:.2f} + {used:.1f} GB held by others exceeds the " flush=True)
f"{GMU_CEILING:.2f} ceiling on a {total:.1f} GB card; clamping "
f"(eviction should have freed more)", flush=True)
eff = GMU_CEILING
if abs(eff - share) > 0.005:
print(f"[vllm] gpu-memory-utilization {share:.2f} (own share) -> {eff:.3f} "
f"absolute: {used:.1f} of {total:.1f} GB already in use", flush=True)
return eff return eff
......
...@@ -787,11 +787,21 @@ class OcrConfig: ...@@ -787,11 +787,21 @@ class OcrConfig:
evict_drain_timeout_s: float = 60.0 evict_drain_timeout_s: float = 60.0
# The share of the card a VLM OCR engine's own vLLM instance claims (surya in vllm # 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 # mode, olmocr in vllm mode). 0 = use vllm.gpu_memory_utilization, which is sized for
# serving an LLM and is far more than a 7B OCR model needs — on a shared card that is # serving an LLM. This is a fraction of the card's TOTAL memory, and it has to cover
# the difference between booting and "Free memory … is less than desired GPU memory # the engine's whole footprint, not just its weights: measured live on a 24 GB 3090,
# utilization". 0.35 of a 24 GB card (8.4 GB) fits surya-2 / olmOCR-2 with their KV # surya-ocr-2 needs ~11.1 GB before a single KV block (weights 1.37 + CUDA graphs 0.30
# cache and leaves the resident engines where they are. # + ~9.4 GB activation peak). 0.35 (8.4 GB) was therefore below the floor and could
vlm_gpu_memory_utilization: float = 0.35 # 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.
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
# peak vLLM profiles for that size IS the 9.4 GB above. Chunked prefill is on, so
# capping the tokens per batch shrinks the peak WITHOUT shrinking the usable context.
# 0 = inherit the backend's value.
vlm_max_num_batched_tokens: int = 4096
vlm_max_num_seqs: int = 16
vlm_max_model_len: int = 0 # 0 = inherit vllm.ctx (a page needs far less)
# When the chosen engine cannot serve a document, try these engines instead rather # When the chosen engine cannot serve a document, try these engines instead rather
# than failing the request: a 503 means a document nobody read, and that is strictly # than failing the request: a 503 means a document nobody read, and that is strictly
# worse than a transcription from the second-best engine. "auto" = every other # worse than a transcription from the second-best engine. "auto" = every other
...@@ -1430,6 +1440,9 @@ class ConfigManager: ...@@ -1430,6 +1440,9 @@ class ConfigManager:
"build_retry_cooldown_s": self.config.ocr.build_retry_cooldown_s, "build_retry_cooldown_s": self.config.ocr.build_retry_cooldown_s,
"evict_drain_timeout_s": self.config.ocr.evict_drain_timeout_s, "evict_drain_timeout_s": self.config.ocr.evict_drain_timeout_s,
"vlm_gpu_memory_utilization": self.config.ocr.vlm_gpu_memory_utilization, "vlm_gpu_memory_utilization": self.config.ocr.vlm_gpu_memory_utilization,
"vlm_max_num_batched_tokens": self.config.ocr.vlm_max_num_batched_tokens,
"vlm_max_num_seqs": self.config.ocr.vlm_max_num_seqs,
"vlm_max_model_len": self.config.ocr.vlm_max_model_len,
"fallback_engines": self.config.ocr.fallback_engines, "fallback_engines": self.config.ocr.fallback_engines,
"olmocr_enabled": self.config.ocr.olmocr_enabled, "olmocr_enabled": self.config.ocr.olmocr_enabled,
"olmocr_instances": self.config.ocr.olmocr_instances, "olmocr_instances": self.config.ocr.olmocr_instances,
......
...@@ -237,3 +237,25 @@ def _rasterize_pdf(data: bytes, dpi: int) -> List["Any"]: ...@@ -237,3 +237,25 @@ def _rasterize_pdf(data: bytes, dpi: int) -> List["Any"]:
return [im.convert("RGB") for im in convert_from_bytes(data, dpi=dpi)] return [im.convert("RGB") for im in convert_from_bytes(data, dpi=dpi)]
except Exception as e: except Exception as e:
raise OcrError(f"failed to rasterise PDF (pdf2image): {e}", status=400) raise OcrError(f"failed to rasterise PDF (pdf2image): {e}", status=400)
def vlm_serve_limits(cfg) -> dict:
"""The serving limits a VLM OCR engine imposes on its own vLLM instance, as kwargs
for ``vllm_worker.ensure_service``.
Without these the engine inherits the vLLM BACKEND's settings, which are sized for
serving an LLM: ``ctx=18432`` here, and the activation peak vLLM profiles for that
batch size was 9.4 GB of the 11.1 GB footprint that left surya-ocr-2 unable to start
on a 24 GB card at all. Chunked prefill is on, so capping the tokens per batch cuts
the peak without capping the context a page may use. 0/absent = inherit."""
out = {}
for key, arg in (("vlm_max_num_batched_tokens", "max_num_batched_tokens"),
("vlm_max_num_seqs", "max_num_seqs"),
("vlm_max_model_len", "max_model_len")):
try:
v = int(getattr(cfg, key, 0) or 0)
except Exception:
v = 0
if v > 0:
out[arg] = v
return out
...@@ -36,7 +36,7 @@ import json ...@@ -36,7 +36,7 @@ import json
import threading import threading
from typing import Optional from typing import Optional
from codai.ocr.base import OcrEngine, OcrPage, OcrError from codai.ocr.base import OcrEngine, OcrPage, OcrError, vlm_serve_limits
#: olmocr.prompts.build_no_anchoring_v4_yaml_prompt() — kept verbatim: the model was #: olmocr.prompts.build_no_anchoring_v4_yaml_prompt() — kept verbatim: the model was
...@@ -245,7 +245,8 @@ class OlmOcrEngine(OcrEngine): ...@@ -245,7 +245,8 @@ class OlmOcrEngine(OcrEngine):
raise OcrError("olmOCR vllm mode needs the vLLM backend configured", status=400) raise OcrError("olmOCR vllm mode needs the vLLM backend configured", status=400)
model = self._model_name() model = self._model_name()
base = vllm_worker.ensure_service(vcfg, model_path=model, served_name=model, base = vllm_worker.ensure_service(vcfg, model_path=model, served_name=model,
gpu_memory_utilization=self._vlm_gmu()) gpu_memory_utilization=self._vlm_gmu(),
**vlm_serve_limits(self.cfg))
self._base = base.rstrip("/") + "/v1" self._base = base.rstrip("/") + "/v1"
def cleanup(self) -> None: def cleanup(self) -> None:
......
...@@ -15,7 +15,7 @@ LICENSE: GPL (compatible with coderai's GPLv3). Enabled only when ...@@ -15,7 +15,7 @@ LICENSE: GPL (compatible with coderai's GPLv3). Enabled only when
import os import os
from codai.ocr.subprocess_engine import SubprocessOcrEngine, _resolve_venv_dir from codai.ocr.subprocess_engine import SubprocessOcrEngine, _resolve_venv_dir
from codai.ocr.base import OcrError from codai.ocr.base import OcrError, vlm_serve_limits
_REQ = os.path.join(os.path.dirname(__file__), "..", "..", "requirements-surya.txt") _REQ = os.path.join(os.path.dirname(__file__), "..", "..", "requirements-surya.txt")
...@@ -103,7 +103,8 @@ class SuryaEngine(SubprocessOcrEngine): ...@@ -103,7 +103,8 @@ class SuryaEngine(SubprocessOcrEngine):
raise OcrError("Surya vllm mode needs the vLLM backend configured", status=400) raise OcrError("Surya vllm mode needs the vLLM backend configured", status=400)
model = (getattr(self.cfg, "surya_model", "") or "datalab-to/surya-ocr-2").strip() model = (getattr(self.cfg, "surya_model", "") or "datalab-to/surya-ocr-2").strip()
base = vllm_worker.ensure_service(vcfg, model_path=model, served_name=model, base = vllm_worker.ensure_service(vcfg, model_path=model, served_name=model,
gpu_memory_utilization=self._vlm_gmu()) gpu_memory_utilization=self._vlm_gmu(),
**vlm_serve_limits(self.cfg))
self._server_url = base.rstrip("/") + "/v1" self._server_url = base.rstrip("/") + "/v1"
elif mode == "llamacpp": elif mode == "llamacpp":
url = (getattr(self.cfg, "surya_server_url", "") or "").strip() url = (getattr(self.cfg, "surya_server_url", "") or "").strip()
......
...@@ -126,21 +126,23 @@ refuses to start when that much is not free. Three rules follow, all enforced by ...@@ -126,21 +126,23 @@ 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 - **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 front (`prelaunch_vram_gb`) and the model manager evicts for it first. Evicting
afterwards is useless, because the load is what fails. afterwards is useless, because the load is what fails.
- **The share is relative to what is free, not to the card.** `vlm_gpu_memory_utilization` - **The share must cover the engine's whole footprint, not its weights.** Measured live
is the slice the OCR engine wants for *itself*, but vLLM's own flag is a fraction of on a 24 GB 3090 (two boots of surya-ocr-2: a 9.57 GiB budget gave
**total** card memory and every other process' allocation counts against it. So the `Available KV cache memory: -1.54 GiB`, a 14.31 GiB budget gave `+3.18 GiB`), the fixed
share is translated at launch: 0.35 on a 24 GB card with 3.6 GB of resident embedder cost is **~11.1 GiB** — weights 1.37, CUDA graphs 0.30 and ~9.4 GiB of activation peak —
becomes `--gpu-memory-utilization 0.50`, a 12 GB budget of which 8.4 GB is actually the *before a single KV block*. `vlm_gpu_memory_utilization` is a fraction of the card's
engine's. Without the translation, 0.35 left surya-2 about 5 GB, and after weights, TOTAL memory, so 0.35 of this card is 8.4 GiB: below the floor, and no amount of
activation peak at `max_model_len=18432` and CUDA graphs it reported eviction makes it start. Hence the **0.60** default (14.4 GiB → ~3 GiB of KV cache).
`Available KV cache memory: -2.81 GiB` and died — once every 20 minutes, for hours. Nothing is added back for other residents — those two boots show vLLM charges only its
The translation is capped (0.92) so the driver keeps its room; when the cap would bite, own allocations against `total × gmu` — so the figure is simply the share, and the
the evict-first pass above is what clears the card instead. evict-first pass above is what guarantees it is free.
- **A failed build is not retried on every request.** For `ocr.build_retry_cooldown_s` - **The batched-token cap is the better lever.** That 9.4 GiB peak is profiled for the
(default 60 s) that engine fails fast with the reason the boot gave, instead of spending batch size, and a VLM OCR engine otherwise inherits the *LLM* backend's `ctx` (18432
~35 s booting and dying again per request. Without this, a card 11 GB short of Surya-2's here). Chunked prefill is on, so `ocr.vlm_max_num_batched_tokens` (default 4096) and
demand produced 338 consecutive failed boots in one morning, each one a half-minute hang `ocr.vlm_max_num_seqs` (16) shrink the peak **without** shrinking the context a page may
for the caller. 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.
If you see `Free memory on device cuda:0 … is less than desired GPU memory utilization`, 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 set **`ocr.vlm_gpu_memory_utilization`** (Settings → OCR → "VLM engines' GPU share"): the
...@@ -166,7 +168,9 @@ back. ...@@ -166,7 +168,9 @@ back.
Settings → **OCR** card, or `config.json` `"ocr"`. Key fields: `enabled`, Settings → **OCR** card, or `config.json` `"ocr"`. Key fields: `enabled`,
`default_engine`, `dpi`, `max_concurrency`, `lang`, `build_retry_cooldown_s`, `default_engine`, `dpi`, `max_concurrency`, `lang`, `build_retry_cooldown_s`,
`vlm_gpu_memory_utilization`, `fallback_engines`; per-engine `*_enabled` / `vlm_gpu_memory_utilization`, `vlm_max_num_batched_tokens`, `vlm_max_num_seqs`,
`vlm_max_model_len`, `fallback_engines`, `evict_drain_timeout_s`;
per-engine `*_enabled` /
`*_instances` / `*_use_gpu` / lang; `detect_mode` (+ `detect_model_path`, `detect_conf`); `*_instances` / `*_use_gpu` / lang; `detect_mode` (+ `detect_model_path`, `detect_conf`);
`extract_enabled` / `extract_model_id` / `extract_schema` / `extract_validate`. `extract_enabled` / `extract_model_id` / `extract_schema` / `extract_validate`.
......
...@@ -329,8 +329,9 @@ def test_vllm_mode_serves_with_the_ocr_side_gpu_share(monkeypatch): ...@@ -329,8 +329,9 @@ def test_vllm_mode_serves_with_the_ocr_side_gpu_share(monkeypatch):
fake_worker = types.SimpleNamespace( fake_worker = types.SimpleNamespace(
ensure_service=lambda vcfg, model_path=None, served_name=None, ensure_service=lambda vcfg, model_path=None, served_name=None,
gpu_memory_utilization=None: ( gpu_memory_utilization=None, **kw: (
seen.update(model=model_path, share=gpu_memory_utilization), "http://127.0.0.1:1/")[1], seen.update(model=model_path, share=gpu_memory_utilization, limits=kw),
"http://127.0.0.1:1/")[1],
planned_vram_gb=lambda vcfg, share=None: (share or 0.9) * 24.0, planned_vram_gb=lambda vcfg, share=None: (share or 0.9) * 24.0,
prelaunch_free_gb=lambda vcfg, share=None: ((share or 0.9) + _headroom) * 24.0, prelaunch_free_gb=lambda vcfg, share=None: ((share or 0.9) + _headroom) * 24.0,
) )
...@@ -338,19 +339,23 @@ def test_vllm_mode_serves_with_the_ocr_side_gpu_share(monkeypatch): ...@@ -338,19 +339,23 @@ def test_vllm_mode_serves_with_the_ocr_side_gpu_share(monkeypatch):
monkeypatch.setattr("codai.models.manager.get_active_vllm_config", monkeypatch.setattr("codai.models.manager.get_active_vllm_config",
lambda: types.SimpleNamespace(gpu_memory_utilization=0.9)) lambda: types.SimpleNamespace(gpu_memory_utilization=0.9))
cfg = _cfg(olmocr_serve="vllm", vlm_gpu_memory_utilization=0.35) cfg = _cfg(olmocr_serve="vllm", vlm_gpu_memory_utilization=0.60)
eng = OlmOcrEngine(cfg) eng = OlmOcrEngine(cfg)
# 0.35 of the card for itself (not the LLM's 0.9) PLUS the ceiling's headroom, because # 0.60 of the card for itself (not the LLM's 0.9) PLUS the ceiling's headroom, which
# the share is translated against whatever else holds the card at launch. # is what must be FREE before launching — vLLM reserves it up front.
assert eng.prelaunch_vram_gb() == pytest.approx((0.35 + _headroom) * 24.0) assert eng.prelaunch_vram_gb() == pytest.approx((0.60 + _headroom) * 24.0)
eng.load() eng.load()
assert seen["model"] == "allenai/olmOCR-2-7B-1025-FP8" assert seen["model"] == "allenai/olmOCR-2-7B-1025-FP8"
assert seen["share"] == 0.35 assert seen["share"] == 0.60
# …and it caps its own instance's batch, which is what sets the activation peak.
assert seen["limits"]["max_num_batched_tokens"] == 4096
assert seen["limits"]["max_num_seqs"] == 16
assert eng._base == "http://127.0.0.1:1/v1" assert eng._base == "http://127.0.0.1:1/v1"
# 0.35 is the shipped default, because 0.9 of the card for a 7B page model is what # 0.60 is the shipped default: measured live, surya-ocr-2 needs ~11.1 GB before a
# crash-looped Surya-2 on a box that had other work resident. # single KV block, so the 0.35 this once shipped (8.4 GB of a 24 GB card) was under
assert OcrConfig().vlm_gpu_memory_utilization == 0.35 # the floor and could not start at all.
assert OcrConfig().vlm_gpu_memory_utilization == 0.60
# Set back to 0 it defers to whatever the vLLM backend itself is configured with. # Set back to 0 it defers to whatever the vLLM backend itself is configured with.
eng2 = OlmOcrEngine(_cfg(olmocr_serve="vllm", vlm_gpu_memory_utilization=0.0)) eng2 = OlmOcrEngine(_cfg(olmocr_serve="vllm", vlm_gpu_memory_utilization=0.0))
assert eng2.prelaunch_vram_gb() == pytest.approx((0.9 + _headroom) * 24.0) assert eng2.prelaunch_vram_gb() == pytest.approx((0.9 + _headroom) * 24.0)
...@@ -439,6 +444,7 @@ def test_the_model_manager_can_evict_ocr_including_both_vlm_servers(monkeypatch) ...@@ -439,6 +444,7 @@ def test_the_model_manager_can_evict_ocr_including_both_vlm_servers(monkeypatch)
stop_service_for=lambda vcfg, model_path=None, served_name=None, stop_service_for=lambda vcfg, model_path=None, served_name=None,
gpu_memory_utilization=None: (stopped.append((model_path, gpu_memory_utilization)), 8.4)[1], gpu_memory_utilization=None: (stopped.append((model_path, gpu_memory_utilization)), 8.4)[1],
planned_vram_gb=lambda vcfg, share=None: 8.4, planned_vram_gb=lambda vcfg, share=None: 8.4,
prelaunch_free_gb=lambda vcfg, share=None: 0.0,
ensure_service=lambda *a, **k: "http://127.0.0.1:1/", ensure_service=lambda *a, **k: "http://127.0.0.1:1/",
) )
_patch_vllm_worker(monkeypatch, fake_worker) _patch_vllm_worker(monkeypatch, fake_worker)
...@@ -478,8 +484,8 @@ def test_the_model_manager_can_evict_ocr_including_both_vlm_servers(monkeypatch) ...@@ -478,8 +484,8 @@ def test_the_model_manager_can_evict_ocr_including_both_vlm_servers(monkeypatch)
freed = m._release_vram(999.0) freed = m._release_vram(999.0)
assert len(released) == 2 # both instances torn down assert len(released) == 2 # both instances torn down
assert freed >= 16.0 # pool + BOTH vLLM services assert freed >= 16.0 # pool + BOTH vLLM services
assert ("datalab-to/surya-ocr-2", 0.35) in stopped assert ("datalab-to/surya-ocr-2", 0.60) in stopped
assert ("allenai/olmOCR-2-7B-1025-FP8", 0.35) in stopped assert ("allenai/olmOCR-2-7B-1025-FP8", 0.60) in stopped
# And the pool rebuilds on the next request rather than staying dead. # And the pool rebuilds on the next request rather than staying dead.
asyncio.run(m.ocr_pages([_page(64, 64)], engine="olmocr")) asyncio.run(m.ocr_pages([_page(64, 64)], engine="olmocr"))
assert len(released) == 2 assert len(released) == 2
...@@ -46,25 +46,28 @@ def test_an_empty_card_passes_the_share_through(monkeypatch): ...@@ -46,25 +46,28 @@ def test_an_empty_card_passes_the_share_through(monkeypatch):
assert vllm_worker.effective_gmu(0.35) == pytest.approx(0.35, abs=1e-3) assert vllm_worker.effective_gmu(0.35) == pytest.approx(0.35, abs=1e-3)
def test_the_share_grows_by_what_other_models_already_hold(monkeypatch): def test_other_residents_do_not_inflate_the_share(monkeypatch):
"""The actual regression: bge-m3 resident while surya tries to boot.""" """The share is what THIS instance may use, full stop.
_card(monkeypatch, total_gb=24.0, free_gb=24.0 - 3.6)
eff = vllm_worker.effective_gmu(0.35) It used to have other residents' usage added back, on the theory that vLLM charges
# 0.35 for us + 3.6/24 held by the embedder = a 12 GB budget, 8.4 GB of it ours. them against the budget. Two live boots said otherwise (the same 11.1 GB footprint at
assert eff == pytest.approx(0.35 + 3.6 / 24.0, abs=1e-3) two different budgets), and the add-back made the figure depend on unrelated models
assert eff * 24.0 - 3.6 == pytest.approx(0.35 * 24.0, abs=0.05) while asking for share*total for this instance on top of what they held. The
evict-first pass clears the card instead."""
_card(monkeypatch, total_gb=24.0, free_gb=24.0 - 6.1)
assert vllm_worker.effective_gmu(0.60) == pytest.approx(0.60, abs=1e-3)
def test_the_translation_is_capped_so_the_driver_keeps_its_room(monkeypatch): def test_the_share_is_capped_so_the_driver_keeps_its_room(monkeypatch):
_card(monkeypatch, total_gb=24.0, free_gb=2.0) # 22 GB held by others _card(monkeypatch, total_gb=24.0, free_gb=24.0)
assert vllm_worker.effective_gmu(0.35) == pytest.approx(vllm_worker.GMU_CEILING) assert vllm_worker.effective_gmu(0.99) == pytest.approx(vllm_worker.GMU_CEILING)
def test_an_unmeasurable_card_does_not_inflate_the_share(monkeypatch): def test_an_unmeasurable_card_still_passes_the_share_through(monkeypatch):
from codai.models.manager import multi_model_manager from codai.models.manager import multi_model_manager
monkeypatch.setattr(multi_model_manager, "_total_vram_gb", lambda: 24.0) monkeypatch.setattr(multi_model_manager, "_total_vram_gb", lambda: 24.0)
monkeypatch.setattr(multi_model_manager, "_get_free_vram_gb", lambda: 999.0) monkeypatch.setattr(multi_model_manager, "_get_free_vram_gb", lambda: 999.0)
assert vllm_worker.effective_gmu(0.35) == pytest.approx(0.35, abs=1e-3) assert vllm_worker.effective_gmu(0.60) == pytest.approx(0.60, abs=1e-3)
def test_no_share_means_no_flag(monkeypatch): def test_no_share_means_no_flag(monkeypatch):
...@@ -73,32 +76,83 @@ def test_no_share_means_no_flag(monkeypatch): ...@@ -73,32 +76,83 @@ def test_no_share_means_no_flag(monkeypatch):
assert vllm_worker.effective_gmu(None) == 0.0 assert vllm_worker.effective_gmu(None) == 0.0
# ---------------------------------------------------- footprint caps
def test_the_vlm_engine_caps_its_own_vllms_batch(monkeypatch):
"""The activation peak these cap IS what kept surya-ocr-2 from starting: 9.4 GB of an
11.1 GB footprint, profiled for the LLM backend's ctx=18432."""
from codai.ocr.base import vlm_serve_limits
assert vlm_serve_limits(OcrConfig()) == {
"max_num_batched_tokens": 4096, "max_num_seqs": 16} # max_model_len 0 = inherit
def test_zero_means_inherit_the_backends_value():
from codai.ocr.base import vlm_serve_limits
cfg = OcrConfig(vlm_max_num_batched_tokens=0, vlm_max_num_seqs=0, vlm_max_model_len=0)
assert vlm_serve_limits(cfg) == {}
def test_an_explicit_model_len_is_passed_too():
from codai.ocr.base import vlm_serve_limits
cfg = OcrConfig(vlm_max_model_len=8192)
assert vlm_serve_limits(cfg)["max_model_len"] == 8192
def test_the_caps_reach_the_launch_command(monkeypatch):
_card(monkeypatch)
cfg = type("C", (), {"model_id": "surya", "ctx": 18432, "max_num_seqs": 0,
"gpu_memory_utilization": 0.9})()
cmd = vllm_worker._launch_cmd(
"py", cfg, "127.0.0.1", 1, "datalab-to/surya-ocr-2", "datalab-to/surya-ocr-2",
gpu_memory_utilization=0.60, max_num_batched_tokens=4096, max_num_seqs=16)
assert "--max-num-batched-tokens" in cmd
assert cmd[cmd.index("--max-num-batched-tokens") + 1] == "4096"
assert cmd[cmd.index("--max-num-seqs") + 1] == "16"
assert cmd[cmd.index("--gpu-memory-utilization") + 1] == "0.6"
assert cmd[cmd.index("--max-model-len") + 1] == "18432" # inherited
def test_an_override_wins_over_the_backends_context(monkeypatch):
_card(monkeypatch)
cfg = type("C", (), {"model_id": "surya", "ctx": 18432, "gpu_memory_utilization": 0.9})()
cmd = vllm_worker._launch_cmd("py", cfg, "127.0.0.1", 1, "m", "m",
max_model_len=8192)
assert cmd[cmd.index("--max-model-len") + 1] == "8192"
def test_the_shipped_share_clears_the_measured_footprint():
"""0.35 of a 24 GB card is 8.4 GB against an 11.1 GB floor — below it, nothing starts.
0.60 is 14.4 GB, which left 3.18 GB of KV cache on the live card."""
share = OcrConfig().vlm_gpu_memory_utilization
assert share == 0.60
assert share * 24.0 > 11.1 + 2.0
def test_the_backends_own_setting_is_already_absolute(monkeypatch): def test_the_backends_own_setting_is_already_absolute(monkeypatch):
"""Only a side job's share is translated; vllm.gpu_memory_utilization is not.""" """Only a side job's share is translated; vllm.gpu_memory_utilization is not."""
_card(monkeypatch, total_gb=24.0, free_gb=24.0 - 8.0) _card(monkeypatch, total_gb=24.0, free_gb=24.0 - 8.0)
cfg = type("C", (), {"gpu_memory_utilization": 0.90})() cfg = type("C", (), {"gpu_memory_utilization": 0.90})()
assert vllm_worker.resolve_gmu(cfg, None) == pytest.approx(0.90) assert vllm_worker.resolve_gmu(cfg, None) == pytest.approx(0.90)
assert vllm_worker.resolve_gmu(cfg, 0.35) == pytest.approx(0.35 + 8.0 / 24.0, abs=1e-3) assert vllm_worker.resolve_gmu(cfg, 0.60) == pytest.approx(0.60, abs=1e-3)
# ---------------------------------------------------- evict-first demand # ---------------------------------------------------- evict-first demand
def test_the_evict_demand_covers_the_share_plus_the_ceilings_headroom(monkeypatch): def test_the_evict_demand_covers_the_share_plus_the_ceilings_headroom(monkeypatch):
_card(monkeypatch, total_gb=24.0, free_gb=24.0) _card(monkeypatch, total_gb=24.0, free_gb=24.0)
cfg = type("C", (), {"gpu_memory_utilization": 0.0})() cfg = type("C", (), {"gpu_memory_utilization": 0.0})()
need = vllm_worker.prelaunch_free_gb(cfg, 0.35) need = vllm_worker.prelaunch_free_gb(cfg, 0.60)
# Enough free that 0.35 + used/total still lands under the ceiling. # The share itself plus the headroom the ceiling reserves.
assert need == pytest.approx((0.35 + (1.0 - vllm_worker.GMU_CEILING)) * 24.0, abs=0.05) assert need == pytest.approx((0.60 + (1.0 - vllm_worker.GMU_CEILING)) * 24.0, abs=0.05)
assert need > vllm_worker.planned_vram_gb(cfg, 0.35) assert need > vllm_worker.planned_vram_gb(cfg, 0.60)
def test_a_crowded_card_is_cleared_rather_than_clamped(monkeypatch): def test_a_crowded_card_is_cleared_rather_than_clamped(monkeypatch):
"""When the demand is met, the translation fits — that is the point of the pairing.""" """When the demand is met, the translation fits — that is the point of the pairing."""
_card(monkeypatch, total_gb=24.0, free_gb=24.0) _card(monkeypatch, total_gb=24.0, free_gb=24.0)
cfg = type("C", (), {"gpu_memory_utilization": 0.0})() cfg = type("C", (), {"gpu_memory_utilization": 0.0})()
need = vllm_worker.prelaunch_free_gb(cfg, 0.35) need = vllm_worker.prelaunch_free_gb(cfg, 0.60)
# Evictor honoured the demand exactly: `need` free, the rest held by others. # Evictor honoured the demand exactly: `need` free, the rest held by others.
_card(monkeypatch, total_gb=24.0, free_gb=need) _card(monkeypatch, total_gb=24.0, free_gb=need)
assert vllm_worker.effective_gmu(0.35) <= vllm_worker.GMU_CEILING + 1e-6 assert need <= 24.0 and vllm_worker.effective_gmu(0.60) * 24.0 <= need
# ---------------------------------------------------- release accounting # ---------------------------------------------------- release accounting
...@@ -112,16 +166,16 @@ def test_stopping_reports_what_was_reserved_not_the_engines_share(monkeypatch): ...@@ -112,16 +166,16 @@ def test_stopping_reports_what_was_reserved_not_the_engines_share(monkeypatch):
stopped = [] stopped = []
monkeypatch.setattr(vllm_worker, "stop_service", lambda k: stopped.append(k)) monkeypatch.setattr(vllm_worker, "stop_service", lambda k: stopped.append(k))
launched = vllm_worker.effective_gmu(0.35) # 0.50 -> a 12 GB reservation launched = vllm_worker.effective_gmu(0.60) # 0.60 -> a 14.4 GB reservation
monkeypatch.setitem(vllm_worker._services, "m|m", monkeypatch.setitem(vllm_worker._services, "m|m",
{"proc": _Proc(), "port": 1, "url": "u", "ray": None, {"proc": _Proc(), "port": 1, "url": "u", "ray": None,
"gmu": launched}) "gmu": launched})
cfg = type("C", (), {"model_path": "m", "gpu_memory_utilization": 0.0})() cfg = type("C", (), {"model_path": "m", "gpu_memory_utilization": 0.0})()
freed = vllm_worker.stop_service_for(cfg, model_path="m", served_name="m", freed = vllm_worker.stop_service_for(cfg, model_path="m", served_name="m",
gpu_memory_utilization=0.35) gpu_memory_utilization=0.10)
assert stopped == ["m|m"] assert stopped == ["m|m"]
assert freed == pytest.approx(launched * 24.0, abs=0.05) # ~12 GB assert freed == pytest.approx(launched * 24.0, abs=0.05) # ~14.4 GB, what it took
assert freed > 0.35 * 24.0 # not the 8.4 GB share assert freed > 0.10 * 24.0 # not the (stale) share the caller passed
def test_a_leftover_service_is_stopped_even_after_the_mode_changed(monkeypatch): def test_a_leftover_service_is_stopped_even_after_the_mode_changed(monkeypatch):
...@@ -263,7 +317,6 @@ def test_a_healthy_engine_is_not_second_guessed(monkeypatch): ...@@ -263,7 +317,6 @@ def test_a_healthy_engine_is_not_second_guessed(monkeypatch):
def test_the_fallback_is_on_by_default(): def test_the_fallback_is_on_by_default():
assert OcrConfig().fallback_engines == "auto" assert OcrConfig().fallback_engines == "auto"
assert OcrConfig().vlm_gpu_memory_utilization == 0.35
assert OcrConfig().evict_drain_timeout_s == 60.0 assert OcrConfig().evict_drain_timeout_s == 60.0
......
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