ocr: size a VLM engine's vLLM against free VRAM, and never answer 503 (v0.2.25)

surya-on-vLLM was failing every ~20 minutes in production with
"Available KV cache memory: -2.81 GiB", and the cause was the 0.35 share itself.
vLLM's --gpu-memory-utilization is a fraction of the card's TOTAL memory and every
other process' allocation counts against it, so with 3.6 GB of resident bge-m3 on a
24 GB 3090 a 0.35 budget left surya about 5 GB — less than its weights (1.37) plus
activation peak at max_model_len=18432 plus CUDA graphs (0.30). Each attempt burned
~3 minutes of boot and profiling, then died, and the next OCR request started over.

- vlm_gpu_memory_utilization is now the engine's OWN share and is translated at launch
  against what else holds the card (effective_gmu): 0.35 with 3.6 GB resident becomes
  0.50, i.e. a 12 GB budget of which 8.4 GB is the engine's. Capped at 0.92 so the
  driver keeps its room.
- The evict-first demand (prelaunch_free_gb) now covers the share plus that headroom,
  so a genuinely crowded card gets models evicted rather than the share clamped.
- stop_service_for reports what the instance actually reserved, recorded at launch,
  instead of the caller's share — under-reporting sent the evictor hunting for VRAM
  that was already free.
- Stopping the managed vLLM is no longer gated on the engine's current serve mode: a
  mode changed while a service was up leaked its share of the card until a restart.
- ocr.fallback_engines (default "auto"): an engine that cannot serve hands the document
  to the other enabled engines instead of returning 503, fewest-ways-to-fail first. The
  response's "engine" field names whoever actually read the pages. This also covers a
  pool evicted mid-document — the failed pages are re-read rather than failing.

21 new tests; the two older vLLM-worker tests now patch the package attribute as well
as sys.modules, without which the fake worker silently did not apply.
Co-Authored-By: 's avatarClaude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0182DhR1eFNPQyGmDrpHbedY
parent e6c66928
...@@ -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.24" __version__ = "0.2.25"
__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.
......
...@@ -3872,6 +3872,7 @@ def build_settings_dict(c, gpu_cards): ...@@ -3872,6 +3872,7 @@ def build_settings_dict(c, gpu_cards):
"surya_server_url": c.ocr.surya_server_url, "surya_server_url": c.ocr.surya_server_url,
"build_retry_cooldown_s": c.ocr.build_retry_cooldown_s, "build_retry_cooldown_s": c.ocr.build_retry_cooldown_s,
"vlm_gpu_memory_utilization": c.ocr.vlm_gpu_memory_utilization, "vlm_gpu_memory_utilization": c.ocr.vlm_gpu_memory_utilization,
"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,
"olmocr_serve": c.ocr.olmocr_serve, "olmocr_serve": c.ocr.olmocr_serve,
...@@ -4603,6 +4604,7 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad ...@@ -4603,6 +4604,7 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad
if "surya_server_url" in d: o.surya_server_url = (d.get("surya_server_url") or "").strip() if "surya_server_url" in d: o.surya_server_url = (d.get("surya_server_url") or "").strip()
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 "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)
if "olmocr_serve" in d: o.olmocr_serve = (d.get("olmocr_serve") or "model").strip().lower() if "olmocr_serve" in d: o.olmocr_serve = (d.get("olmocr_serve") or "model").strip().lower()
......
...@@ -1074,6 +1074,12 @@ ...@@ -1074,6 +1074,12 @@
<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:170px">
<label class="form-label">Fallback engines</label>
<input type="text" id="s-ocr-fallback" class="form-input" placeholder="auto">
<div class="form-hint">Tried when the chosen engine can't serve, so a document
is never refused: <code>auto</code>, <code>off</code>, or an explicit order.</div>
</div>
</div> </div>
<div class="card-title" style="font-size:13px;margin-top:.4rem">PaddleOCR</div> <div class="card-title" style="font-size:13px;margin-top:.4rem">PaddleOCR</div>
...@@ -2004,6 +2010,7 @@ async function loadSettings(){ ...@@ -2004,6 +2010,7 @@ async function loadSettings(){
document.getElementById('s-ocr-lang').value = ocr.lang ?? 'it'; document.getElementById('s-ocr-lang').value = ocr.lang ?? 'it';
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-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;
document.getElementById('s-ocr-paddle-structure').checked = ocr.paddle_structure !== false; document.getElementById('s-ocr-paddle-structure').checked = ocr.paddle_structure !== false;
...@@ -2231,6 +2238,7 @@ async function saveSettings(silent){ ...@@ -2231,6 +2238,7 @@ async function saveSettings(silent){
lang: document.getElementById('s-ocr-lang').value.trim() || 'it', lang: document.getElementById('s-ocr-lang').value.trim() || 'it',
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',
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,
paddle_structure: document.getElementById('s-ocr-paddle-structure').checked, paddle_structure: document.getElementById('s-ocr-paddle-structure').checked,
......
...@@ -125,9 +125,7 @@ def _launch_cmd(py, cfg, host: str, port: int, model_path: str, ...@@ -125,9 +125,7 @@ def _launch_cmd(py, cfg, host: str, port: int, model_path: str,
ctx = int(getattr(cfg, "ctx", 0) or 0) ctx = 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)]
# A side job that serves its OWN small model on a shared card (the OCR VLM engines) # Already absolute (resolve_gmu translated a side job's own-share if there was one).
# passes its own share here, so it does not have to claim as much of the GPU as the
# LLM instance wants.
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)]
...@@ -302,9 +300,13 @@ def ensure_service(cfg, model_path: Optional[str] = None, ...@@ -302,9 +300,13 @@ def ensure_service(cfg, model_path: Optional[str] = None,
port = int(getattr(cfg, "port", 0) or 0) or _free_port() 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_host = "127.0.0.1" if host in ("0.0.0.0", "") else host
url = f"http://{url_host}:{port}" 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)
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=gpu_memory_utilization) gpu_memory_utilization=gmu)
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
...@@ -340,7 +342,8 @@ def ensure_service(cfg, model_path: Optional[str] = None, ...@@ -340,7 +342,8 @@ def ensure_service(cfg, model_path: Optional[str] = None,
ray.stop() ray.stop()
raise raise
threading.Thread(target=_pump_logs, args=(proc, tail), daemon=True).start() threading.Thread(target=_pump_logs, args=(proc, tail), daemon=True).start()
_services[svc_key] = {"proc": proc, "port": port, "url": url, "ray": ray} _services[svc_key] = {"proc": proc, "port": port, "url": url, "ray": ray,
"gmu": gmu}
def _tail_msg(): def _tail_msg():
# The last six lines of a vLLM crash are the wrapper's own traceback # The last six lines of a vLLM crash are the wrapper's own traceback
...@@ -412,6 +415,94 @@ def is_running(cfg, model_path: Optional[str] = None, served_name: Optional[str] ...@@ -412,6 +415,94 @@ def is_running(cfg, model_path: Optional[str] = None, served_name: Optional[str]
return bool(svc and svc["proc"].poll() is None) return bool(svc and svc["proc"].poll() is None)
# vLLM must not be told to budget the whole card: the driver context, fragmentation and
# any allocation made between our measurement and its own need somewhere to live.
GMU_CEILING = 0.92
def _used_by_others_gb() -> float:
"""VRAM (GB) currently held on this box by anything that is NOT the vLLM we are
about to launch — resident embedders, an LLM, another process' context."""
try:
from codai.models.manager import multi_model_manager
total = float(multi_model_manager._total_vram_gb() or 0.0)
free = float(multi_model_manager._get_free_vram_gb() or 0.0)
if total <= 0 or free >= 999.0: # 999 = "unknown, assume enough"
return 0.0
return max(0.0, total - free)
except Exception:
return 0.0
def resolve_gmu(cfg, gpu_memory_utilization: Optional[float] = None) -> float:
"""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
translated against what else is on the card; the backend's own configured value is
already absolute and is used as-is."""
share = float(gpu_memory_utilization or 0)
if share > 0:
return effective_gmu(share)
return float(getattr(cfg, "gpu_memory_utilization", 0) or 0)
def effective_gmu(self_share: float) -> float:
"""Translate "the share of the card I want for MYSELF" into the absolute
``--gpu-memory-utilization`` vLLM wants.
vLLM's flag is a fraction of the card's TOTAL memory and every other process'
allocation counts against it: with 3.6 GB of bge-m3 resident on a 24 GB 3090,
``0.35`` left surya-2 about 5 GB, and after weights (1.37) + activation peak at
``max_model_len=18432`` + CUDA graphs (0.3) it reported ``Available KV cache memory:
-2.81 GiB`` and died — every 20 minutes, for hours. So add back what others hold:
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
is actually ours. Capped at :data:`GMU_CEILING`; when the cap bites, the caller's
evict-first pass (``prelaunch_vram_gb``) is what makes room."""
try:
share = max(0.0, float(self_share or 0.0))
except Exception:
return 0.0
if share <= 0:
return 0.0
try:
from codai.models.manager import multi_model_manager
total = float(multi_model_manager._total_vram_gb() or 0.0)
except Exception:
total = 0.0
if total <= 0:
return min(share, GMU_CEILING)
used = _used_by_others_gb()
eff = share + (used / total)
if eff > GMU_CEILING:
print(f"[vllm] share {share:.2f} + {used:.1f} GB held by others exceeds the "
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
def prelaunch_free_gb(cfg, gpu_memory_utilization: Optional[float] = None) -> float:
"""How much VRAM must be FREE before launching this vLLM for the side-job share to
fit under :data:`GMU_CEILING` — i.e. what to ask the evictor for.
It is the engine's own need plus the headroom the ceiling reserves, so a crowded
card gets models evicted instead of :func:`effective_gmu` clamping and vLLM dying
on KV cache again."""
gmu = float(gpu_memory_utilization or 0) or float(getattr(cfg, "gpu_memory_utilization", 0) or 0)
if gmu <= 0:
return planned_vram_gb(cfg, gpu_memory_utilization)
try:
from codai.models.manager import multi_model_manager
total = float(multi_model_manager._total_vram_gb() or 24.0)
except Exception:
total = 24.0
total = total or 24.0
return max(0.5, (min(gmu, GMU_CEILING) + (1.0 - GMU_CEILING)) * total)
def planned_vram_gb(cfg, gpu_memory_utilization: Optional[float] = None) -> float: def planned_vram_gb(cfg, gpu_memory_utilization: Optional[float] = None) -> float:
"""How much VRAM a vLLM instance started from ``cfg`` will claim on this box. """How much VRAM a vLLM instance started from ``cfg`` will claim on this box.
...@@ -446,9 +537,21 @@ def stop_service_for(cfg, model_path: Optional[str] = None, ...@@ -446,9 +537,21 @@ def stop_service_for(cfg, model_path: Optional[str] = None,
with _lock: with _lock:
svc = _services.get(svc_key) svc = _services.get(svc_key)
running = bool(svc and svc["proc"].poll() is None) running = bool(svc and svc["proc"].poll() is None)
launched_gmu = float((svc or {}).get("gmu") or 0.0)
if not running: if not running:
return 0.0 return 0.0
stop_service(svc_key) stop_service(svc_key)
# What it reserved is what we just freed. The share the caller passes is the engine's
# OWN share; the instance was launched with that plus whatever else held the card
# (effective_gmu), so reporting the caller's figure would under-count the release and
# send the eviction path looking for memory that is already free.
if launched_gmu > 0:
try:
from codai.models.manager import multi_model_manager
total = float(multi_model_manager._total_vram_gb() or 24.0)
except Exception:
total = 24.0
return max(0.5, launched_gmu * (total or 24.0))
return planned_vram_gb(cfg, gpu_memory_utilization) return planned_vram_gb(cfg, gpu_memory_utilization)
......
...@@ -787,6 +787,13 @@ class OcrConfig: ...@@ -787,6 +787,13 @@ class OcrConfig:
# utilization". 0.35 of a 24 GB card (8.4 GB) fits surya-2 / olmOCR-2 with their KV # utilization". 0.35 of a 24 GB card (8.4 GB) fits surya-2 / olmOCR-2 with their KV
# cache and leaves the resident engines where they are. # cache and leaves the resident engines where they are.
vlm_gpu_memory_utilization: float = 0.35 vlm_gpu_memory_utilization: float = 0.35
# 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
# worse than a transcription from the second-best engine. "auto" = every other
# ENABLED engine, fewest-ways-to-fail first (paddle, doctr, olmocr, surya);
# "off"/"none" = no fallback (fail with the original error); or an explicit
# comma/space separated order.
fallback_engines: str = "auto"
# --- PaddleOCR / PP-Structure --- # --- PaddleOCR / PP-Structure ---
paddle_enabled: bool = True paddle_enabled: bool = True
...@@ -1417,6 +1424,7 @@ class ConfigManager: ...@@ -1417,6 +1424,7 @@ class ConfigManager:
"surya_server_url": self.config.ocr.surya_server_url, "surya_server_url": self.config.ocr.surya_server_url,
"build_retry_cooldown_s": self.config.ocr.build_retry_cooldown_s, "build_retry_cooldown_s": self.config.ocr.build_retry_cooldown_s,
"vlm_gpu_memory_utilization": self.config.ocr.vlm_gpu_memory_utilization, "vlm_gpu_memory_utilization": self.config.ocr.vlm_gpu_memory_utilization,
"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,
"olmocr_serve": self.config.ocr.olmocr_serve, "olmocr_serve": self.config.ocr.olmocr_serve,
......
...@@ -282,8 +282,12 @@ class OcrManager: ...@@ -282,8 +282,12 @@ class OcrManager:
def _stop_surya_vllm(self) -> float: def _stop_surya_vllm(self) -> float:
"""Stop the managed Surya-2 vLLM subprocess (if this build serves Surya via vLLM). """Stop the managed Surya-2 vLLM subprocess (if this build serves Surya via vLLM).
Returns estimated GB freed.""" Returns estimated GB freed."""
# NOT gated on surya_serve being "vllm" right now: if the setting was changed
# while a service was up, the old subprocess is still holding its share of the
# card and skipping the stop would leak it until a restart. stop_service_for() is
# a no-op returning 0.0 when nothing is running, so asking is free.
cfg = self._cfg cfg = self._cfg
if cfg is None or (getattr(cfg, "surya_serve", "") or "").strip().lower() != "vllm": if cfg is None:
return 0.0 return 0.0
try: try:
from codai.api import vllm_worker from codai.api import vllm_worker
...@@ -302,8 +306,8 @@ class OcrManager: ...@@ -302,8 +306,8 @@ class OcrManager:
def _stop_olmocr_vllm(self) -> float: def _stop_olmocr_vllm(self) -> float:
"""Stop the managed olmOCR vLLM instance (olmocr_serve = "vllm"), for the same """Stop the managed olmOCR vLLM instance (olmocr_serve = "vllm"), for the same
reason Surya's needs stopping: its VRAM is the subprocess's, not a pool's.""" reason Surya's needs stopping: its VRAM is the subprocess's, not a pool's."""
cfg = self._cfg cfg = self._cfg # unconditional, for the reason in _stop_surya_vllm
if cfg is None or (getattr(cfg, "olmocr_serve", "") or "").strip().lower() != "vllm": if cfg is None:
return 0.0 return 0.0
try: try:
from codai.api import vllm_worker from codai.api import vllm_worker
...@@ -361,9 +365,27 @@ class OcrManager: ...@@ -361,9 +365,27 @@ class OcrManager:
self._pools[name] = pool self._pools[name] = pool
return pool return pool
async def ocr_pages(self, images: List, engine: Optional[str] = None) -> (str, List[OcrPage]): def _fallback_order(self, failed: str) -> List[str]:
"""OCR a list of PIL page images. Returns (engine_name, [OcrPage]).""" """Enabled engines to try after ``failed`` could not serve the document.
name = self.resolve_engine(engine)
An OCR request that comes back 503 is a document nobody read — and the whole
point of having four engines is that a page can still be read when one of them
will not come up (a VLM engine whose vLLM cannot fit on the card, a venv that
needs rebuilding, a license flag). ``auto`` prefers the engines with the fewest
ways to fail over the ones with the best output, because the alternative to a
worse transcription here is no transcription."""
cfg = self._require_cfg()
raw = str(getattr(cfg, "fallback_engines", "auto") or "auto").strip().lower()
if raw in ("off", "none", "no", "false"):
return []
if raw in ("auto", "", "on", "true"):
names = ["paddle", "doctr", "olmocr", "surya"]
else:
names = [n.strip() for n in raw.replace(",", " ").split() if n.strip()]
return [n for n in names
if n != failed and n in _ENGINE_FACTORIES and _engine_enabled(cfg, n)]
async def _pages_with(self, name: str, images: List) -> List[OcrPage]:
pool = await self._get_pool(name) pool = await self._get_pool(name)
sem = self._sem or asyncio.Semaphore(4) sem = self._sem or asyncio.Semaphore(4)
...@@ -374,7 +396,35 @@ class OcrManager: ...@@ -374,7 +396,35 @@ class OcrManager:
return page return page
pages = await asyncio.gather(*[_one(i, im) for i, im in enumerate(images)]) pages = await asyncio.gather(*[_one(i, im) for i, im in enumerate(images)])
return name, list(pages) return list(pages)
async def ocr_pages(self, images: List, engine: Optional[str] = None) -> (str, List[OcrPage]):
"""OCR a list of PIL page images. Returns (engine_name, [OcrPage]).
The returned name is the engine that ACTUALLY read the pages, which is not
necessarily the one asked for: when an engine cannot serve, the remaining
enabled engines are tried in turn (see :meth:`_fallback_order`) rather than
failing the document. Callers surface the name, so a caller that cares can see
it fell back."""
name = self.resolve_engine(engine)
try:
return name, await self._pages_with(name, images)
except Exception as first:
chain = self._fallback_order(name)
if not chain:
raise
print(f"[ocr] engine '{name}' could not serve ({first}); "
f"falling back to {' -> '.join(chain)}")
for alt in chain:
try:
pages = await self._pages_with(alt, images)
except Exception as e:
print(f"[ocr] fallback engine '{alt}' also failed: {e}")
continue
print(f"[ocr] served by fallback engine '{alt}' instead of '{name}'")
return alt, pages
# Everything enabled is down: the original failure is the useful one.
raise first
async def ocr_document(self, data: bytes, filename: str = "", content_type: str = "", async def ocr_document(self, data: bytes, filename: str = "", content_type: str = "",
engine: Optional[str] = None, dpi: Optional[int] = None, engine: Optional[str] = None, dpi: Optional[int] = None,
......
...@@ -276,11 +276,14 @@ class OlmOcrEngine(OcrEngine): ...@@ -276,11 +276,14 @@ class OlmOcrEngine(OcrEngine):
return v or None return v or None
def _vllm_need(self) -> float: def _vllm_need(self) -> float:
# The engine's own share PLUS the headroom effective_gmu's ceiling reserves: our
# share is translated against whatever else holds the card, so a crowded card has
# to be cleared rather than clamped.
try: try:
from codai.api.vllm_worker import planned_vram_gb from codai.api.vllm_worker import prelaunch_free_gb
from codai.models.manager import get_active_vllm_config from codai.models.manager import get_active_vllm_config
vcfg = get_active_vllm_config() vcfg = get_active_vllm_config()
return planned_vram_gb(vcfg, self._vlm_gmu()) if vcfg is not None else 0.0 return prelaunch_free_gb(vcfg, self._vlm_gmu()) if vcfg is not None else 0.0
except Exception: except Exception:
return 0.0 return 0.0
......
...@@ -53,15 +53,18 @@ class SuryaEngine(SubprocessOcrEngine): ...@@ -53,15 +53,18 @@ class SuryaEngine(SubprocessOcrEngine):
def prelaunch_vram_gb(self) -> float: def prelaunch_vram_gb(self) -> float:
"""In the served modes Surya rides a vLLM instance that grabs """In the served modes Surya rides a vLLM instance that grabs
gpu_memory_utilization × the whole card before it will start, so that much has to gpu_memory_utilization × the whole card before it will start, so that much has to
be free FIRST. (A 3090 with 11.3 GB free against a 14.1 GB demand crash-looped be free FIRST — plus the ceiling's headroom, since our share is translated against
Surya-2 for hours: vLLM died on startup, the next OCR request tried again.)""" what others hold (see vllm_worker.effective_gmu). (A 3090 with 11.3 GB free against
a 14.1 GB demand crash-looped Surya-2 for hours: vLLM died on startup, the next OCR
request tried again. Then 3.6 GB of resident embedder did the same to a 0.35
share.)"""
if self._serve_mode() != "vllm": if self._serve_mode() != "vllm":
return 0.0 return 0.0
try: try:
from codai.api.vllm_worker import planned_vram_gb from codai.api.vllm_worker import prelaunch_free_gb
from codai.models.manager import get_active_vllm_config from codai.models.manager import get_active_vllm_config
vcfg = get_active_vllm_config() vcfg = get_active_vllm_config()
return planned_vram_gb(vcfg, self._vlm_gmu()) if vcfg is not None else 0.0 return prelaunch_free_gb(vcfg, self._vlm_gmu()) if vcfg is not None else 0.0
except Exception: except Exception:
return 0.0 return 0.0
......
...@@ -121,11 +121,21 @@ thing out there. ...@@ -121,11 +121,21 @@ thing out there.
The VLM engines (`surya` in `vllm` mode, `olmocr` in `vllm` mode) boot a vLLM instance The VLM engines (`surya` in `vllm` mode, `olmocr` in `vllm` mode) boot a vLLM instance
that claims `gpu_memory_utilization` × the **whole card** before it will serve, and that claims `gpu_memory_utilization` × the **whole card** before it will serve, and
refuses to start when that much is not free. Two rules follow, both enforced by the pool: refuses to start when that much is not free. Three rules follow, all enforced by the pool:
- **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`
is the slice the OCR engine wants for *itself*, but vLLM's own flag is a fraction of
**total** card memory and every other process' allocation counts against it. So the
share is translated at launch: 0.35 on a 24 GB card with 3.6 GB of resident embedder
becomes `--gpu-memory-utilization 0.50`, a 12 GB budget of which 8.4 GB is actually the
engine's. Without the translation, 0.35 left surya-2 about 5 GB, and after weights,
activation peak at `max_model_len=18432` and CUDA graphs it reported
`Available KV cache memory: -2.81 GiB` and died — once every 20 minutes, for hours.
The translation is capped (0.92) so the driver keeps its room; when the cap would bite,
the evict-first pass above is what clears the card instead.
- **A failed build is not retried on every request.** For `ocr.build_retry_cooldown_s` - **A failed build is not retried on every request.** For `ocr.build_retry_cooldown_s`
(default 60 s) that engine fails fast with the reason the boot gave, instead of spending (default 60 s) that engine fails fast with the reason the boot gave, instead of spending
~35 s booting and dying again per request. Without this, a card 11 GB short of Surya-2's ~35 s booting and dying again per request. Without this, a card 11 GB short of Surya-2's
...@@ -140,11 +150,23 @@ of 0.9 (or 0.6) can be more than the card has left. Otherwise: give the card les ...@@ -140,11 +150,23 @@ of 0.9 (or 0.6) can be more than the card has left. Otherwise: give the card les
resident work, or serve the engine with `model`/`server` mode, which needs no instance of resident work, or serve the engine with `model`/`server` mode, which needs no instance of
its own at all. its own at all.
### A document is never refused because one engine is down
`ocr.fallback_engines` (default `auto`) makes a failing engine hand the document to the
other **enabled** engines rather than returning 503 — fewest-ways-to-fail first
(`paddle`, `doctr`, `olmocr`, `surya`), since a transcription from the second-best engine
beats no transcription. The response's `engine` field names the engine that actually read
the pages, so a fallback is visible to the caller, and the reason the first choice failed
is logged. This also covers a pool that is evicted *mid-document* by a model load: the
failed pages are re-read by the next engine instead of failing the request. Set an
explicit order (`"doctr, paddle"`) to override it, or `off` to get the original error
back.
## Enable & configure ## Enable & configure
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`; per-engine `*_enabled` / `vlm_gpu_memory_utilization`, `fallback_engines`; 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`.
......
...@@ -36,6 +36,18 @@ def _cfg(**kw): ...@@ -36,6 +36,18 @@ def _cfg(**kw):
return c return c
def _patch_vllm_worker(monkeypatch, fake):
"""Swap in a fake vLLM worker for code that does ``from codai.api import vllm_worker``.
Patching sys.modules alone is not enough: once the real module has been imported by
anything (another test file, for instance), it is bound as an attribute of the
``codai.api`` package and the from-import takes that attribute without consulting
sys.modules. Patch both, or the fake silently does not apply."""
import codai.api
monkeypatch.setitem(sys.modules, "codai.api.vllm_worker", fake)
monkeypatch.setattr(codai.api, "vllm_worker", fake, raising=False)
# ------------------------------------------------------------ front matter # ------------------------------------------------------------ front matter
def test_front_matter_split(): def test_front_matter_split():
meta, body = parse_front_matter( meta, body = parse_front_matter(
...@@ -312,19 +324,25 @@ def test_vllm_mode_serves_with_the_ocr_side_gpu_share(monkeypatch): ...@@ -312,19 +324,25 @@ def test_vllm_mode_serves_with_the_ocr_side_gpu_share(monkeypatch):
"""The OCR VLM gets its own share of the card, not the one sized for an LLM.""" """The OCR VLM gets its own share of the card, not the one sized for an LLM."""
seen = {} seen = {}
from codai.api.vllm_worker import GMU_CEILING
_headroom = 1.0 - GMU_CEILING # the ceiling keeps room for the driver/fragmentation
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: (
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), "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,
) )
monkeypatch.setitem(sys.modules, "codai.api.vllm_worker", fake_worker) _patch_vllm_worker(monkeypatch, fake_worker)
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.35)
eng = OlmOcrEngine(cfg) eng = OlmOcrEngine(cfg)
assert eng.prelaunch_vram_gb() == pytest.approx(8.4) # 0.35 of the card, not 0.9 # 0.35 of the card for itself (not the LLM's 0.9) PLUS the ceiling's headroom, because
# the share is translated against whatever else holds the card at launch.
assert eng.prelaunch_vram_gb() == pytest.approx((0.35 + _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.35
...@@ -335,7 +353,7 @@ def test_vllm_mode_serves_with_the_ocr_side_gpu_share(monkeypatch): ...@@ -335,7 +353,7 @@ def test_vllm_mode_serves_with_the_ocr_side_gpu_share(monkeypatch):
assert OcrConfig().vlm_gpu_memory_utilization == 0.35 assert OcrConfig().vlm_gpu_memory_utilization == 0.35
# 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(21.6) assert eng2.prelaunch_vram_gb() == pytest.approx((0.9 + _headroom) * 24.0)
def test_a_shared_server_is_reserved_for_ONCE_not_per_instance(monkeypatch): def test_a_shared_server_is_reserved_for_ONCE_not_per_instance(monkeypatch):
...@@ -423,7 +441,7 @@ def test_the_model_manager_can_evict_ocr_including_both_vlm_servers(monkeypatch) ...@@ -423,7 +441,7 @@ def test_the_model_manager_can_evict_ocr_including_both_vlm_servers(monkeypatch)
planned_vram_gb=lambda vcfg, share=None: 8.4, planned_vram_gb=lambda vcfg, share=None: 8.4,
ensure_service=lambda *a, **k: "http://127.0.0.1:1/", ensure_service=lambda *a, **k: "http://127.0.0.1:1/",
) )
monkeypatch.setitem(sys.modules, "codai.api.vllm_worker", fake_worker) _patch_vllm_worker(monkeypatch, fake_worker)
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))
......
This diff is collapsed.
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