ocr: olmOCR-2 as an alternative engine, and stop the Surya-2 crash loop (v0.2.24)

New engine `olmocr` — allenai/olmOCR-2 (Apache-2.0), a document VLM for the pages
the detector-based engines read badly: old scans, maths, multi-column, dense
tables. Equations as LaTeX, tables as HTML, reading order sorted; no bounding
boxes, so paddle/surya stay the choice when something has to point at the page.

It is served three ways (`ocr.olmocr_serve`): "model" (default) routes through
coderai's OWN model manager, so a GGUF+mmproj or HF entry in models.json gets VRAM
accounting, eviction and the thermal governor — and it works without CUDA; "vllm"
gives it its own vLLM instance; "server" attaches to an OpenAI endpoint already
running it. olmOCR's trained no-anchoring-v4 prompt is sent verbatim, the page is
rendered at 1288 px, the YAML front matter lands in the new `OcrPage.meta` (omitted
for engines that report none), and a page the model calls rotated is turned and
asked once more.

Four fixes the production Surya failures exposed — 169 consecutive failed boots
between 08:00 and 09:35 today, every one against the same 11.34 of 23.56 GiB free
against a 14.13 GiB demand, each costing the caller a 35 s hang:

- The pool evicted for an engine AFTER loading it. For an engine that boots a
  server claiming a fixed share of the card, the load is what fails. Engines now
  declare `prelaunch_vram_gb()` and the pool frees that much first.
- That reservation is the WHOLE POOL's, asked for once. Per instance it would have
  asked the manager to free 24 x 8.4 GB on this install's 24-instance surya pool,
  and an engine that reserved up front was then evicted for a second time, after
  its server had already taken the memory.
- A failed build was retried on every single request. `ocr.build_retry_cooldown_s`
  (60 s) now fails fast with the reason the boot gave.
- `ocr.vlm_gpu_memory_utilization` (default 0.35): the OCR VLM's own share of the
  card instead of the one `vllm.gpu_memory_utilization` asks for on behalf of an
  LLM. 0.35 of 24 GB fits surya-2 / olmOCR-2 and leaves the resident engines alone;
  0.6 was more than the card had left.

Per-instance footprints corrected to match: a surya worker in a served mode speaks
only HTTP (0.05 GB, not 1 GB x 24), and an olmocr instance is one in-flight page,
holding nothing of its own. Eviction in the other direction is unchanged and
tested: the model manager's external releaser tears down the pools AND stops both
VLM engines' vLLM services, and the pools rebuild on the next request.

18 new tests (no model, no GPU, no network). Docs: docs/ocr.md.
Co-Authored-By: 's avatarClaude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01B7EnDW6RSd2nhEPZMzMkqi
parent b8b5ba3a
......@@ -275,13 +275,17 @@ registered in the embedding model category.
### Document OCR
A dedicated OCR subsystem (`/v1/ocr`) — real OCR engines, **not** a VLM prompted to read:
A dedicated OCR subsystem (`/v1/ocr`) — real OCR engines, **not** a VLM prompted to read,
plus one purpose-trained document VLM for the pages they struggle with:
- **Engines**: `doctr` (in-process), `paddle` (PaddleOCR + PP-Structure layout/tables, isolated
venv) and `surya` (opt-in, GPL; local, vLLM-served Surya-2, or an external llama-server)
venv), `surya` (opt-in, GPL; local, vLLM-served Surya-2, or an external llama-server) and
`olmocr` ([olmOCR-2](https://huggingface.co/allenai/olmOCR-2-7B-1025), Apache-2.0 — best
on old scans, maths and multi-column pages; served by one of your own models, coderai's
vLLM, or any OpenAI endpoint. Returns no bounding boxes.)
- **Input**: images or PDFs (rasterised at a configurable DPI), single or batch
- **Output**: full text plus per-page `lines[]` with bounding boxes and confidence,
`regions[]` (layout) and `tables[]`
`regions[]` (layout), `tables[]`, and `meta` where the engine reports page metadata
- **Structured extraction**: `structured=true` runs a text model against a **schema** —
a JSON document, not code — to return typed fields. Schemas are data-driven and stored in
`<config>/ocr_schemas/`; built-ins cover `italian_sentenza`, `generic_document` and an
......@@ -1168,7 +1172,9 @@ its hardware — running models that otherwise simply couldn't run on a single G
- **[PaddleOCR](https://github.com/PaddlePaddle/PaddleOCR)** (Baidu),
**[docTR](https://github.com/mindee/doctr)** (Mindee) and
**[Surya](https://github.com/VikParuchuri/surya)** — the dedicated OCR engines behind
the `/v1/ocr` document-transcription subsystem (text + layout + structured extraction).
the `/v1/ocr` document-transcription subsystem (text + layout + structured extraction),
and **[olmOCR-2](https://huggingface.co/allenai/olmOCR-2-7B-1025)** (AllenAI) — the
document VLM option on the same endpoint.
And the libraries, models and research CoderAI builds on:
......
......@@ -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.23"
__version__ = "0.2.24"
__tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
# Configure the CUDA caching allocator BEFORE torch is imported anywhere.
......
......@@ -3870,6 +3870,20 @@ def build_settings_dict(c, gpu_cards):
"surya_serve": c.ocr.surya_serve,
"surya_model": c.ocr.surya_model,
"surya_server_url": c.ocr.surya_server_url,
"build_retry_cooldown_s": c.ocr.build_retry_cooldown_s,
"vlm_gpu_memory_utilization": c.ocr.vlm_gpu_memory_utilization,
"olmocr_enabled": c.ocr.olmocr_enabled,
"olmocr_instances": c.ocr.olmocr_instances,
"olmocr_serve": c.ocr.olmocr_serve,
"olmocr_model_id": c.ocr.olmocr_model_id,
"olmocr_model": c.ocr.olmocr_model,
"olmocr_server_url": c.ocr.olmocr_server_url,
"olmocr_api_key": c.ocr.olmocr_api_key,
"olmocr_longest_side": c.ocr.olmocr_longest_side,
"olmocr_max_tokens": c.ocr.olmocr_max_tokens,
"olmocr_temperature": c.ocr.olmocr_temperature,
"olmocr_timeout": c.ocr.olmocr_timeout,
"olmocr_retry_rotation": c.ocr.olmocr_retry_rotation,
"detect_mode": c.ocr.detect_mode,
"detect_model_path": c.ocr.detect_model_path,
"detect_conf": c.ocr.detect_conf,
......@@ -4549,6 +4563,13 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad
except (TypeError, ValueError):
return cur
def _f(key, cur, lo=None):
try:
v = float(d.get(key, cur))
return max(lo, v) if lo is not None else v
except (TypeError, ValueError):
return cur
if "enabled" in d: o.enabled = bool(d["enabled"])
if "default_engine" in d: o.default_engine = (d.get("default_engine") or "paddle").strip()
if "dpi" in d: o.dpi = _i("dpi", o.dpi, 36)
......@@ -4580,6 +4601,20 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad
if "surya_serve" in d: o.surya_serve = (d.get("surya_serve") or "local").strip()
if "surya_model" in d: o.surya_model = (d.get("surya_model") 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 "vlm_gpu_memory_utilization" in d: o.vlm_gpu_memory_utilization = _f("vlm_gpu_memory_utilization", o.vlm_gpu_memory_utilization, 0.0)
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_serve" in d: o.olmocr_serve = (d.get("olmocr_serve") or "model").strip().lower()
if "olmocr_model_id" in d: o.olmocr_model_id = (d.get("olmocr_model_id") or "").strip()
if "olmocr_model" in d: o.olmocr_model = (d.get("olmocr_model") or "").strip()
if "olmocr_server_url" in d: o.olmocr_server_url = (d.get("olmocr_server_url") or "").strip()
if "olmocr_api_key" in d: o.olmocr_api_key = (d.get("olmocr_api_key") or "").strip()
if "olmocr_longest_side" in d: o.olmocr_longest_side = _i("olmocr_longest_side", o.olmocr_longest_side, 64)
if "olmocr_max_tokens" in d: o.olmocr_max_tokens = _i("olmocr_max_tokens", o.olmocr_max_tokens, 64)
if "olmocr_temperature" in d: o.olmocr_temperature = _f("olmocr_temperature", o.olmocr_temperature, 0.0)
if "olmocr_timeout" in d: o.olmocr_timeout = _f("olmocr_timeout", o.olmocr_timeout, 1.0)
if "olmocr_retry_rotation" in d: o.olmocr_retry_rotation = bool(d["olmocr_retry_rotation"])
# detection
if "detect_mode" in d: o.detect_mode = (d.get("detect_mode") or "off").strip()
if "detect_model_path" in d: o.detect_model_path = (d.get("detect_model_path") or "").strip()
......
This diff is collapsed.
......@@ -138,7 +138,7 @@ STEP_PARAMS = {
"audio_cleanup": [("audio","ref","Source audio"),("noise_reduction","checkbox","Reduce background noise"),("normalize","checkbox","Normalize levels"),("remove_hum","checkbox","Remove hum"),("repair_clicks","checkbox","Repair clicks")],
"voice_clone": [("text","textarea","Text to synthesize"),("voice_name","text","Voice profile name"),("ref_text","text","Reference transcript"),("speed","number","Speed","1.0")],
"voice_convert": [("source_audio","ref","Source audio"),("voice_name","text","Voice profile name"),("f0_condition","checkbox","Singing mode"),("pitch_shift","number","Pitch shift","0"),("diffusion_steps","number","Steps","10")],
"ocr": [("image","ref","Document image/PDF ({{stepN.url}})"),("engine","select:paddle|doctr|surya","OCR engine","paddle"),("detect","select:off|layout|detector|both","Stamp/signature detect","off"),("structured","checkbox","Structured JSON extraction")],
"ocr": [("image","ref","Document image/PDF ({{stepN.url}})"),("engine","select:paddle|doctr|surya|olmocr","OCR engine","paddle"),("detect","select:off|layout|detector|both","Stamp/signature detect","off"),("structured","checkbox","Structured JSON extraction")],
}
......
......@@ -561,7 +561,7 @@ def _remote_catalogue(url: str, api_key: str) -> list:
#: OCR engines are selected by name, not registered in models.json like a model.
#: Without this an OCR test answered "'surya' is not in the model list, so
#: nothing can say where it should run" — true of the catalogue, and useless.
_OCR_ENGINES = ("paddle", "doctr", "surya")
_OCR_ENGINES = ("paddle", "doctr", "surya", "olmocr")
def _describe(model: str):
......
......@@ -18,7 +18,8 @@
POST /v1/ocr — OCR an uploaded image or PDF with a purpose-built OCR engine
(PaddleOCR / docTR / Surya), returning faithful text plus per-line boxes and layout.
NOT a vision LLM. Structured-JSON extraction and stamp/signature detection are layered
Those three are NOT vision LLMs; ``engine=olmocr`` additionally offers olmOCR-2, a
document VLM for pages the detector-based engines read badly (it returns no boxes). Structured-JSON extraction and stamp/signature detection are layered
on in later milestones (fields already accepted, no-op until then).
"""
......@@ -60,7 +61,7 @@ def _ocr_config():
@router.post("/v1/ocr", summary="OCR an image or PDF with a dedicated OCR engine")
async def create_ocr(
file: UploadFile = File(...),
engine: Optional[str] = Form(None), # paddle|doctr|surya (default from config)
engine: Optional[str] = Form(None), # paddle|doctr|surya|olmocr (default from config)
dpi: Optional[int] = Form(None), # PDF rasterisation DPI
lang: Optional[str] = Form(None), # reserved (per-engine language override)
structured: Optional[bool] = Form(False), # structured JSON extraction via text model
......
......@@ -871,7 +871,7 @@ def pod_ocr_engine(requested: str) -> str:
request asks for over its configured default.
"""
r = (requested or "").strip().lower()
if r == "surya":
if r in ("surya", "olmocr"):
return "doctr"
return r if r in ("paddle", "doctr") else "doctr"
......@@ -962,23 +962,25 @@ def _plan(engine, image, args, mcfg, api_key, entry, served,
if cap == "ocr":
env["CODERAI_OCR_ENABLED"] = "1"
wanted = str((entry or {}).get("path") or "").strip().lower()
if wanted == "surya":
# Surya 0.22 is a VLM ("surya-ocr-2"), not a self-contained OCR
# library: it needs a server — vLLM, which it spawns IN DOCKER (a
# pod cannot), or llama-server on a GGUF, which the image does not
# ship. Locally it rides coderai's own vLLM engine; a pod has none.
# So a pod asked for surya serves the request with docTR, which is
# in the image and has passed — rather than loading surya, reaching
# inference, and dying on "docker binary not found".
print("[runpod] surya on a pod needs a VLM server the image does not "
if wanted in ("surya", "olmocr"):
# Surya 0.22 is a VLM ("surya-ocr-2") and olmOCR-2 is one outright,
# not a self-contained OCR library: both need a server — vLLM, which
# surya spawns IN DOCKER (a pod cannot), or llama-server on a GGUF,
# which the image does not ship. Locally they ride coderai's own vLLM
# engine; a pod has none. So a pod asked for either serves the request
# with docTR, which is in the image and has passed — rather than
# loading it, reaching inference, and dying on "docker binary not
# found". (An olmOCR pod would want olmocr_serve="server" pointed at a
# model server it can actually reach; nothing provisions one here.)
print(f"[runpod] {wanted} on a pod needs a VLM server the image does not "
"carry — this OCR request will be served by docTR", flush=True)
wanted = "doctr"
if wanted not in ("paddle", "doctr", "surya"):
if wanted not in ("paddle", "doctr", "surya", "olmocr"):
# The pod images ship docTR, the only engine that runs in-process.
wanted = "doctr"
env["CODERAI_OCR_DEFAULT_ENGINE"] = wanted
# The subsystem flag is not enough: each engine has a gate of its own,
# and doctr and surya both default to off.
# and all but paddle default to off.
env[f"CODERAI_OCR_{wanted.upper()}_ENABLED"] = "1"
if wanted == "paddle":
# The paddle image bakes it, because paddlepaddle brings its own
......
......@@ -115,7 +115,8 @@ def resolve_service_key(cfg, model_path: Optional[str] = None):
def _launch_cmd(py, cfg, host: str, port: int, model_path: str,
served_name: Optional[str] = None, model_config: dict = None) -> list:
served_name: Optional[str] = None, model_config: dict = None,
gpu_memory_utilization: Optional[float] = None) -> list:
mid = served_name or (getattr(cfg, "model_id", "vllm") or "vllm")
cmd = [py, "-m", "vllm.entrypoints.openai.api_server",
"--host", host, "--port", str(port),
......@@ -124,7 +125,10 @@ def _launch_cmd(py, cfg, host: str, port: int, model_path: str,
ctx = int(getattr(cfg, "ctx", 0) or 0)
if ctx > 0:
cmd += ["--max-model-len", str(ctx)]
gmu = float(getattr(cfg, "gpu_memory_utilization", 0) or 0)
# A side job that serves its OWN small model on a shared card (the OCR VLM engines)
# 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)
if gmu > 0:
cmd += ["--gpu-memory-utilization", str(gmu)]
tp = int(getattr(cfg, "tensor_parallel_size", 0) or 0)
......@@ -251,7 +255,8 @@ def _start_ray(cfg, py: str, env: dict):
def ensure_service(cfg, model_path: Optional[str] = None,
served_name: Optional[str] = None,
ready_timeout: float = 3600.0) -> str:
ready_timeout: float = 3600.0,
gpu_memory_utilization: Optional[float] = 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
......@@ -298,7 +303,8 @@ def ensure_service(cfg, model_path: Optional[str] = None,
url_host = "127.0.0.1" if host in ("0.0.0.0", "") else host
url = f"http://{url_host}:{port}"
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)
env = os.environ.copy()
# flashinfer JIT-compiles CUDA kernels with ninja at runtime, which fails on
......@@ -406,8 +412,28 @@ def is_running(cfg, model_path: Optional[str] = None, served_name: Optional[str]
return bool(svc and svc["proc"].poll() is None)
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.
vLLM allocates ``gpu_memory_utilization`` × the card's TOTAL memory up front and
refuses to start when that much is not free — so this is also what must be freed
BEFORE launching it. Callers that boot vLLM for a side job (the OCR VLM engines) use
it to ask the model manager for room first; without that, Surya-2 crash-looped on a
3090 with 11 GB free against a 14 GB demand."""
gmu = float(gpu_memory_utilization or 0) or float(getattr(cfg, "gpu_memory_utilization", 0) or 0)
if gmu <= 0:
return 4.6 # small VLM fallback estimate
try:
from codai.models.manager import multi_model_manager
total = multi_model_manager._total_vram_gb()
except Exception:
total = 24.0
return max(0.5, gmu * float(total or 24.0))
def stop_service_for(cfg, model_path: Optional[str] = None,
served_name: Optional[str] = None) -> float:
served_name: Optional[str] = None,
gpu_memory_utilization: Optional[float] = None) -> float:
"""Stop the vLLM service for ONE specific (model_path, served_name), if running.
Used by the OCR subsystem as a VRAM releaser: the managed Surya-2 vLLM subprocess
......@@ -422,16 +448,8 @@ def stop_service_for(cfg, model_path: Optional[str] = None,
running = bool(svc and svc["proc"].poll() is None)
if not running:
return 0.0
gmu = float(getattr(cfg, "gpu_memory_utilization", 0) or 0)
stop_service(svc_key)
if gmu > 0:
try:
from codai.models.manager import multi_model_manager
total = multi_model_manager._total_vram_gb()
except Exception:
total = 24.0
return max(0.5, gmu * total)
return 4.6 # small VLM fallback estimate
return planned_vram_gb(cfg, gpu_memory_utilization)
import atexit as _atexit
......
......@@ -757,6 +757,10 @@ class OcrConfig:
- ``doctr`` — Mindee docTR (Apache-2.0; pure-PyTorch; easy install)
- ``surya`` — Surya (best layout/reading-order; LICENSE-GATED — GPL/commercial:
only loaded when ``surya_accept_license`` is True)
- ``olmocr`` — olmOCR-2 (AllenAI, Apache-2.0): a document VLM, not a detector —
best transcription fidelity on hard pages (old scans, maths, tables, multi-column),
but NO bounding boxes. Served through coderai's own model manager, its vLLM
backend, or any OpenAI-compatible endpoint (``olmocr_serve``).
Heavy OCR dependencies are installed out of band (optional) so the base image
stays lean. Concurrency on a single GPU comes from loading multiple instances of
......@@ -768,10 +772,21 @@ class OcrConfig:
stamp/signature flags (``detect_*``).
"""
enabled: bool = False
default_engine: str = "paddle" # paddle|doctr|surya
default_engine: str = "paddle" # paddle|doctr|surya|olmocr
dpi: int = 200 # PDF rasterisation DPI
max_concurrency: int = 4 # global cap on in-flight OCR pages
lang: str = "it" # default document language
# After an engine fails to come up, refuse that engine FAST for this long instead of
# re-attempting the load on every request (a vLLM-backed engine burns ~35s per failed
# boot; Surya-2 once did that 338 times in a row on a card that was 11 GB short).
build_retry_cooldown_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 and is far more than a 7B OCR model needs — on a shared card that is
# the difference between booting and "Free memory … is less than desired GPU memory
# 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.
vlm_gpu_memory_utilization: float = 0.35
# --- PaddleOCR / PP-Structure ---
paddle_enabled: bool = True
......@@ -808,6 +823,29 @@ class OcrConfig:
surya_model: str = "datalab-to/surya-ocr-2" # HF checkpoint for the served (vllm) backend
surya_server_url: str = "" # external OpenAI server URL (llamacpp/manual); blank = auto
# --- olmOCR-2 (AllenAI, Apache-2.0) --- a document VLM: whole-page transcription in
# reading order, equations as LaTeX, tables as HTML, no bounding boxes. Driven over an
# OpenAI chat endpoint with olmOCR's own prompt; see codai/ocr/olmocr.py.
olmocr_enabled: bool = False
olmocr_instances: int = 1 # in-flight pages against the serving model
# How it is served:
# "model" — through coderai's OWN model manager: olmocr_model_id names a VISION
# model in models.json (GGUF+mmproj on llama.cpp, or HF/qwenvl). Gets
# VRAM accounting, eviction and the thermal governor for free, and is the
# only mode that works without CUDA.
# "vllm" — coderai's vLLM backend serves olmocr_model on its own instance.
# "server" — attach to an OpenAI-compatible server already running it.
olmocr_serve: str = "model" # model | vllm | server
olmocr_model_id: str = "" # ("model" mode) a vision model id from models.json
olmocr_model: str = "allenai/olmOCR-2-7B-1025-FP8" # ("vllm"/"server" mode) HF id / served name
olmocr_server_url: str = "" # ("server" mode) OpenAI base URL of that model
olmocr_api_key: str = "" # ("server" mode) bearer token, if it wants one
olmocr_longest_side: int = 1288 # page render size the model was trained on
olmocr_max_tokens: int = 4096 # a dense page runs long; this is the page budget
olmocr_temperature: float = 0.1
olmocr_timeout: float = 300.0 # per-page ceiling on the model call
olmocr_retry_rotation: bool = True # re-ask once on a page the model reports rotated
# --- stamp / signature detection --- [O3]
detect_mode: str = "off" # off|layout|detector|both
detect_model_path: str = "" # YOLO signature/stamp weights (detector/both modes)
......@@ -1377,6 +1415,20 @@ class ConfigManager:
"surya_serve": self.config.ocr.surya_serve,
"surya_model": self.config.ocr.surya_model,
"surya_server_url": self.config.ocr.surya_server_url,
"build_retry_cooldown_s": self.config.ocr.build_retry_cooldown_s,
"vlm_gpu_memory_utilization": self.config.ocr.vlm_gpu_memory_utilization,
"olmocr_enabled": self.config.ocr.olmocr_enabled,
"olmocr_instances": self.config.ocr.olmocr_instances,
"olmocr_serve": self.config.ocr.olmocr_serve,
"olmocr_model_id": self.config.ocr.olmocr_model_id,
"olmocr_model": self.config.ocr.olmocr_model,
"olmocr_server_url": self.config.ocr.olmocr_server_url,
"olmocr_api_key": self.config.ocr.olmocr_api_key,
"olmocr_longest_side": self.config.ocr.olmocr_longest_side,
"olmocr_max_tokens": self.config.ocr.olmocr_max_tokens,
"olmocr_temperature": self.config.ocr.olmocr_temperature,
"olmocr_timeout": self.config.ocr.olmocr_timeout,
"olmocr_retry_rotation": self.config.ocr.olmocr_retry_rotation,
"detect_mode": self.config.ocr.detect_mode,
"detect_model_path": self.config.ocr.detect_model_path,
"detect_conf": self.config.ocr.detect_conf,
......@@ -1457,7 +1509,7 @@ class ConfigManager:
# the three default to off. Enabling the subsystem alone got a pod as
# far as "OCR engine 'surya' is not enabled" — a second refusal from
# a second flag, one round later.
for name in ("paddle", "doctr", "surya"):
for name in ("paddle", "doctr", "surya", "olmocr"):
if _flag(f"CODERAI_OCR_{name.upper()}_ENABLED"):
setattr(self.config.ocr, f"{name}_enabled", True)
# Where that engine's interpreter lives. A pod image that ships
......@@ -1465,8 +1517,17 @@ class ConfigManager:
# the subprocess engine runs it instead of trying to build a venv
# on a machine rented by the second.
venv = (_os.environ.get(f"CODERAI_OCR_{name.upper()}_VENV") or "").strip()
if venv:
if venv and hasattr(self.config.ocr, f"{name}_venv"):
setattr(self.config.ocr, f"{name}_venv", venv)
# olmOCR needs to be told WHERE its model is; a pod that carries no vLLM
# is given a server URL (or the id of a vision model it does hold).
for _env, _attr in (("CODERAI_OCR_OLMOCR_SERVE", "olmocr_serve"),
("CODERAI_OCR_OLMOCR_MODEL", "olmocr_model"),
("CODERAI_OCR_OLMOCR_MODEL_ID", "olmocr_model_id"),
("CODERAI_OCR_OLMOCR_SERVER_URL", "olmocr_server_url")):
_val = (_os.environ.get(_env) or "").strip()
if _val:
setattr(self.config.ocr, _attr, _val)
if _flag("CODERAI_OCR_SURYA_ACCEPT_LICENSE"):
# Surya is GPL and gated on an explicit acceptance. The pod
# inherits the decision made here; it cannot make it itself.
......
......@@ -76,7 +76,7 @@ class ModelCapabilities:
# Retrieval / document
reranking: bool = False # cross-encoder query↔document reranking (bge-reranker, …)
ocr: bool = False # optical character recognition (paddle/doctr/surya)
ocr: bool = False # optical character recognition (paddle/doctr/surya/olmocr)
def to_list(self) -> List[str]:
out = []
......
......@@ -5595,6 +5595,8 @@ class MultiModelManager:
_ocr_engines.append("doctr")
if getattr(ocr, "surya_enabled", False) and getattr(ocr, "surya_accept_license", False):
_ocr_engines.append("surya")
if getattr(ocr, "olmocr_enabled", False):
_ocr_engines.append("olmocr")
for _eng in _ocr_engines:
_add(_eng, "ocr", {"capabilities": ["ocr"], "backend": "ocr"})
except Exception:
......
......@@ -16,7 +16,7 @@
"""OCR engine interface + shared result types + input loading.
Every OCR engine (PaddleOCR, docTR, Surya) subclasses :class:`OcrEngine` and returns
Every OCR engine (PaddleOCR, docTR, Surya, olmOCR-2) subclasses :class:`OcrEngine` and returns
:class:`OcrPage` objects. The rest of the subsystem (manager, API, detection,
extraction) speaks only these types, so engines stay interchangeable.
"""
......@@ -78,9 +78,13 @@ class OcrPage:
lines: List[OcrLine] = field(default_factory=list)
regions: List[OcrRegion] = field(default_factory=list)
tables: List[Any] = field(default_factory=list) # list of {bbox, html?, cells?}
#: Engine-reported page metadata, when the engine has any (olmOCR's front matter:
#: primary_language, is_rotation_valid, rotation_correction, is_table, is_diagram).
#: Omitted from the response when empty, so engines that report none are unchanged.
meta: dict = field(default_factory=dict)
def to_dict(self) -> dict:
return {
out = {
"index": self.index,
"text": self.text,
"width": self.width,
......@@ -89,6 +93,9 @@ class OcrPage:
"regions": [r.to_dict() for r in self.regions],
"tables": self.tables,
}
if self.meta:
out["meta"] = self.meta
return out
# ---------------------------------------------------------------------------
......@@ -131,6 +138,22 @@ class OcrEngine(ABC):
"""Rough VRAM footprint (GB) of one loaded instance, for eviction accounting."""
return 1.0
def prelaunch_vram_gb(self) -> float:
"""VRAM (GB) that must be FREE before :meth:`load` is even attempted.
Zero for engines that load weights into whatever room is left (paddle/docTR) —
they are evicted for after the fact, once their real footprint is known. Non-zero
for the ones that boot a server which claims a fixed share of the card up front
and refuses to start otherwise (the VLM engines on vLLM): for those, evicting
after the load attempt is too late, because the load attempt is what fails.
This is the total for the WHOLE POOL, not per instance: the engines that declare
it share one server between all their instances, so the pool asks for it once. (A
pool of 24 surya workers multiplied per instance would have asked for 200 GB.)
An engine that declares it is NOT also evicted for by :meth:`vram_gb` after the
load — the room has already been made."""
return 0.0
# ---------------------------------------------------------------------------
# Input loading (images + PDF) → list of PIL pages
......
......@@ -16,6 +16,10 @@
"""OCR engine manager — per-engine instance pools for concurrent GPU OCR.
An "instance" is one loaded engine; for the VLM engine (olmOCR) it is one in-flight
request against the serving model, so the pool size still bounds that engine's
concurrency.
Each engine is loaded as N resident instances (``*_instances`` in
:class:`~codai.config.OcrConfig`); pages are fanned across the pool so many documents
OCR concurrently on one GPU. An engine instance handles one page at a time (OCR runtimes
......@@ -24,6 +28,7 @@ bounds the whole subsystem.
"""
import asyncio
import time
from typing import Dict, List, Optional
from codai.ocr.base import OcrEngine, OcrPage, OcrError, load_pages
......@@ -46,10 +51,16 @@ def _surya_factory(cfg):
return SuryaEngine(cfg)
def _olmocr_factory(cfg):
from codai.ocr.olmocr import OlmOcrEngine
return OlmOcrEngine(cfg)
_ENGINE_FACTORIES = {
"paddle": _paddle_factory,
"doctr": _doctr_factory,
"surya": _surya_factory,
"olmocr": _olmocr_factory,
}
......@@ -60,6 +71,8 @@ def _engine_enabled(cfg, name: str) -> bool:
return bool(cfg.doctr_enabled)
if name == "surya":
return bool(cfg.surya_enabled and cfg.surya_accept_license)
if name == "olmocr":
return bool(getattr(cfg, "olmocr_enabled", False))
return False
......@@ -106,6 +119,14 @@ class _Pool:
self._instances = [] # every created instance (for release_sync)
self._built = False
self._build_lock = asyncio.Lock()
self._fail_at = 0.0 # when the last build attempt failed
self._fail_err = None # and why (re-raised during the cooldown)
def _cooldown_s(self) -> float:
try:
return max(0.0, float(getattr(self.cfg, "build_retry_cooldown_s", 60.0)))
except Exception:
return 60.0
async def _ensure_built(self):
if self._built:
......@@ -113,19 +134,53 @@ class _Pool:
async with self._build_lock:
if self._built:
return
# A build that just failed is very unlikely to succeed on the next request a
# second later, and retrying is not free: a vLLM-backed engine spends ~35 s
# booting before it dies. Surya-2 did that 338 times in a row against a card
# that was 11 GB short — every one of those requests hung for half a minute
# and then failed anyway. Inside the cooldown, fail FAST with the same reason.
if self._fail_err is not None and (time.time() - self._fail_at) < self._cooldown_s():
left = int(self._cooldown_s() - (time.time() - self._fail_at))
raise OcrError(
f"{self._fail_err} (retrying in {left}s — last attempt failed; "
f"OCR engine '{self.name}' is in its build cooldown)",
status=getattr(self._fail_err, "status", 503))
factory = _ENGINE_FACTORIES[self.name]
self._instances = []
for _ in range(self.size):
eng = factory(self.cfg)
# Load the first instance synchronously-in-thread so a missing
# dependency surfaces as OcrError(503) before we spawn the rest.
await asyncio.to_thread(eng.ensure_loaded)
# First instance: evict other models to make room (VRAM eviction), now
# that we know its real footprint. Mirrors the engine-load path.
if not self._instances:
await asyncio.to_thread(_evict_for_ocr, self.size * eng.vram_gb())
self._instances.append(eng)
await self._q.put(eng)
try:
for _ in range(self.size):
eng = factory(self.cfg)
first = not self._instances
need = 0.0
if first:
# Engines that boot a server claiming a fixed share of the card
# (the VLM engines on vLLM) must have the room BEFORE they load —
# the load is what fails otherwise. The figure covers the whole
# pool, because those instances share one server.
try:
need = float(eng.prelaunch_vram_gb() or 0.0)
except Exception:
need = 0.0
if need > 0:
await asyncio.to_thread(_evict_for_ocr, need)
# Load the first instance synchronously-in-thread so a missing
# dependency surfaces as OcrError(503) before we spawn the rest.
await asyncio.to_thread(eng.ensure_loaded)
# First instance of an engine that did NOT reserve up front: evict for
# its real footprint now that we know it (paddle/docTR), per instance.
# Reserving again here would evict live models for memory the engine's
# server has already taken.
if first and need <= 0:
await asyncio.to_thread(_evict_for_ocr, self.size * eng.vram_gb())
self._instances.append(eng)
await self._q.put(eng)
except Exception as e:
self.release_sync()
self._fail_at = time.time()
self._fail_err = e if isinstance(e, OcrError) else OcrError(
f"OCR engine '{self.name}' failed to load: {e}", status=503)
raise self._fail_err
self._fail_err = None
self._built = True
print(f"[ocr] engine '{self.name}': {self.size} instance(s) ready")
......@@ -153,6 +208,13 @@ class _Pool:
await self._ensure_built()
eng: OcrEngine = await self._q.get()
try:
# Engines that call back into coderai's async API (olmOCR 'model' mode) run
# in a worker thread and need a live loop to hand the coroutine to; give them
# the app's own rather than letting them spin up a second one.
try:
eng._host_loop = asyncio.get_running_loop()
except Exception:
pass
return await asyncio.to_thread(eng.recognize_image, image)
finally:
await self._q.put(eng)
......@@ -212,6 +274,7 @@ class OcrManager:
# above — it must be stopped explicitly. Do it here so on-request eviction can
# reclaim it; the next OCR request re-boots it via ensure_service.
freed += self._stop_surya_vllm()
freed += self._stop_olmocr_vllm()
if freed:
print(f"[ocr] released ~{freed:.1f} GB (pools + Surya vLLM torn down for VRAM eviction)")
return freed
......@@ -229,11 +292,34 @@ class OcrManager:
if vcfg is None:
return 0.0
model = (getattr(cfg, "surya_model", "") or "datalab-to/surya-ocr-2").strip()
return vllm_worker.stop_service_for(vcfg, model_path=model, served_name=model)
return vllm_worker.stop_service_for(
vcfg, model_path=model, served_name=model,
gpu_memory_utilization=float(getattr(cfg, "vlm_gpu_memory_utilization", 0.0) or 0.0) or None)
except Exception as e:
print(f"[ocr] Surya vLLM stop skipped: {e}")
return 0.0
def _stop_olmocr_vllm(self) -> float:
"""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."""
cfg = self._cfg
if cfg is None or (getattr(cfg, "olmocr_serve", "") or "").strip().lower() != "vllm":
return 0.0
try:
from codai.api import vllm_worker
from codai.models.manager import get_active_vllm_config
from codai.ocr.olmocr import DEFAULT_MODEL
vcfg = get_active_vllm_config()
if vcfg is None:
return 0.0
model = (getattr(cfg, "olmocr_model", "") or DEFAULT_MODEL).strip()
return vllm_worker.stop_service_for(
vcfg, model_path=model, served_name=model,
gpu_memory_utilization=float(getattr(cfg, "vlm_gpu_memory_utilization", 0.0) or 0.0) or None)
except Exception as e:
print(f"[ocr] olmOCR vLLM stop skipped: {e}")
return 0.0
@staticmethod
def _pool_params(cfg) -> tuple:
return (
......@@ -246,6 +332,11 @@ class OcrManager:
cfg.doctr_det_arch, cfg.doctr_reco_arch,
cfg.surya_enabled, cfg.surya_accept_license, cfg.surya_instances,
cfg.surya_langs, cfg.surya_venv, cfg.surya_auto_build,
getattr(cfg, "olmocr_enabled", False), getattr(cfg, "olmocr_instances", 1),
getattr(cfg, "olmocr_serve", "model"), getattr(cfg, "olmocr_model", ""),
getattr(cfg, "olmocr_model_id", ""), getattr(cfg, "olmocr_server_url", ""),
getattr(cfg, "olmocr_longest_side", 1288),
getattr(cfg, "vlm_gpu_memory_utilization", 0.0),
)
def resolve_engine(self, requested: Optional[str]) -> str:
......
This diff is collapsed.
......@@ -50,6 +50,38 @@ class SuryaEngine(SubprocessOcrEngine):
return {}
return {"SURYA_INFERENCE_BACKEND": mode, "SURYA_INFERENCE_URL": url}
def prelaunch_vram_gb(self) -> float:
"""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
be free FIRST. (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.)"""
if self._serve_mode() != "vllm":
return 0.0
try:
from codai.api.vllm_worker import planned_vram_gb
from codai.models.manager import 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
except Exception:
return 0.0
def _vlm_gmu(self):
"""The OCR-side share of the card for this engine's vLLM (0/None = the backend's)."""
try:
v = float(getattr(self.cfg, "vlm_gpu_memory_utilization", 0.0) or 0.0)
except Exception:
v = 0.0
return v or None
def vram_gb(self) -> float:
"""Per-instance footprint, for post-load eviction accounting and release reports.
In the served modes a worker only speaks HTTP to the shared vLLM — its own VRAM is
nil, and the instance's real cost is reported once through
:meth:`prelaunch_vram_gb` / the vLLM stop estimate. Counting a whole GB per
instance here asked to free 24 GB on a 24-instance pool."""
return 0.05 if self._serve_mode() in ("vllm", "llamacpp") else 1.0
def load(self) -> None:
if not self.cfg.surya_accept_license:
raise OcrError(
......@@ -67,7 +99,8 @@ class SuryaEngine(SubprocessOcrEngine):
if vcfg is None:
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()
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())
self._server_url = base.rstrip("/") + "/v1"
elif mode == "llamacpp":
url = (getattr(self.cfg, "surya_server_url", "") or "").strip()
......
# OCR subsystem (dedicated OCR engines)
CoderAI OCRs documents with **purpose-built OCR engines** (detection + recognition) —
**not** a vision LLM. That means faithful transcription with per-line bounding boxes and
layout, GPU batching, and no hallucinated text. Three engines are integrated and
selectable per request:
faithful transcription with per-line bounding boxes and layout, GPU batching, no
hallucinated text — and, for the pages those read badly, one **document VLM**. Four
engines are integrated and selectable per request:
| engine | id | license | how it runs |
|---|---|---|---|
| [docTR](https://github.com/mindee/doctr) | `doctr` | Apache-2.0 | **in-process** (uses the main venv's torch); works on GPU. Recommended on new-CUDA boxes. |
| [PaddleOCR](https://github.com/PaddlePaddle/PaddleOCR) + PP-Structure | `paddle` | Apache-2.0 | **isolated venv subprocess** (layout + tables) |
| [Surya](https://github.com/VikParuchuri/surya) | `surya` | GPL (compatible with coderai's GPLv3) | **isolated venv subprocess**; best layout/reading-order; opt-in |
| [olmOCR-2](https://huggingface.co/allenai/olmOCR-2-7B-1025) | `olmocr` | Apache-2.0 | **a document VLM**, served by a coderai model, coderai's vLLM, or any OpenAI endpoint. **No bounding boxes.** |
Per document you get three outputs: **plain text**, **structured JSON fields**, and
**stamp/signature flags**.
**Which one?** `paddle`/`doctr`/`surya` give you boxes, layout regions and tables, which
is what you want when something downstream has to point at a place on the page (redaction,
stamp detection, a table's cells). `olmocr` gives you the best *reading* of a hard page —
old scans, maths, multi-column, dense tables — and nothing positional. Stamp/signature
detection in `layout` mode has nothing to work with under `olmocr`; `detector` mode (YOLO)
still works, because it looks at the image, not the engine's output.
## Why isolated venvs for Paddle & Surya
Their dependencies conflict with the main coderai venv, so each runs in its **own
......@@ -70,10 +78,73 @@ on first use. In the OCI image these venvs are **baked in** at `/opt/coderai/pad
self-contained; the engines prefer a baked venv, then `~/.coderai/<name>_venv`, then the
config path. Surya loads only when `surya_accept_license` is set.
## olmOCR-2 (`olmocr`) — the document VLM
[olmOCR-2](https://huggingface.co/allenai/olmOCR-2-7B-1025) is AllenAI's Qwen2.5-VL-7B
fine-tune for document transcription (Apache-2.0, 82.4 on olmOCR-bench). It reads a whole
page at once and answers in its trained format: a YAML front matter block followed by the
page as markdown, with equations as LaTeX and tables as HTML. coderai sends olmOCR's own
"no-anchoring v4" prompt verbatim — paraphrasing it costs accuracy — renders the page at
`olmocr_longest_side` (1288 px, what it was trained on), parses the front matter into the
page's `meta`, and strips it from `text`. When the model reports the page rotated, coderai
turns it and asks once more (`olmocr_retry_rotation`), which is what the olmOCR pipeline
itself does.
Three ways to serve it (`ocr.olmocr_serve`):
| mode | what it means | needs |
|---|---|---|
| `model` (default) | the checkpoint is a **vision model in your models.json** and olmOCR goes through coderai's own model manager — VRAM accounting, eviction, quantisation and the thermal governor included. The only mode that works **without CUDA** (GGUF + mmproj on llama.cpp over Vulkan/ROCm). | `olmocr_model_id` |
| `vllm` | coderai's vLLM backend serves `olmocr_model` on its own instance (continuous batching). CUDA only. | the vLLM backend configured |
| `server` | attach to an OpenAI-compatible server already running it: `llama-server`, a remote vLLM, another coderai. | `olmocr_server_url` (+ `olmocr_api_key` if it wants one) |
Fields: `olmocr_enabled`, `olmocr_serve`, `olmocr_model_id`, `olmocr_model`
(default `allenai/olmOCR-2-7B-1025-FP8`), `olmocr_server_url`, `olmocr_api_key`,
`olmocr_instances` (pages in flight), `olmocr_longest_side`, `olmocr_max_tokens`,
`olmocr_temperature`, `olmocr_timeout`, `olmocr_retry_rotation`.
```bash
curl -F file=@scansione-1974.pdf -F engine=olmocr http://localhost:8776/v1/ocr
```
Each page's `meta` carries what the model reported: `primary_language`,
`is_rotation_valid`, `rotation_correction`, `is_table`, `is_diagram` (plus
`rotation_applied` when coderai re-asked a turned page). Engines that report no metadata
omit the field entirely, so nothing else in the response shape changed.
**A pod cannot serve it.** Like Surya-2, olmOCR needs a model server the RunPod OCR images
do not carry, so a pod asked for `olmocr` serves the request with docTR and says so in the
log. Point `olmocr_serve: "server"` at something the pod can reach if you want the real
thing out there.
## When an engine's server will not start
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
refuses to start when that much is not free. Two rules follow, both enforced by the pool:
- **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.
- **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
~35 s booting and dying again per request. Without this, a card 11 GB short of Surya-2's
demand produced 338 consecutive failed boots in one morning, each one a half-minute hang
for the caller.
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
OCR VLM gets its own share of the card instead of the one `vllm.gpu_memory_utilization`
asks for on behalf of an LLM — 0.35 is plenty for a 7B page model, where the LLM default
of 0.9 (or 0.6) can be more than the card has left. Otherwise: give the card less other
resident work, or serve the engine with `model`/`server` mode, which needs no instance of
its own at all.
## Enable & configure
Settings → **OCR** card, or `config.json` `"ocr"`. Key fields: `enabled`,
`default_engine`, `dpi`, `max_concurrency`, `lang`; per-engine `*_enabled` /
`default_engine`, `dpi`, `max_concurrency`, `lang`, `build_retry_cooldown_s`,
`vlm_gpu_memory_utilization`; per-engine `*_enabled` /
`*_instances` / `*_use_gpu` / lang; `detect_mode` (+ `detect_model_path`, `detect_conf`);
`extract_enabled` / `extract_model_id` / `extract_schema` / `extract_validate`.
......@@ -92,7 +163,8 @@ GET/POST/DELETE /v1/ocr/schemas[/{name}] manage extraction schemas
```
Response: `{engine, num_pages, text, pages[], stamps[], signatures[], structured}`. Each
page carries `lines[]` (text + bbox + conf), `regions[]` (layout), and `tables[]`. The
page carries `lines[]` (text + bbox + conf), `regions[]` (layout), `tables[]`, and
`meta` when the engine reports page metadata (olmOCR does; the others do not). The
`ocr` **pipeline step type** chains OCR with `text_gen` in custom pipelines.
```bash
......
This diff is collapsed.
......@@ -106,7 +106,7 @@ def _find(data: dict, model: str):
#: catalogue to pin. A temporary entry gives them the same per-model placement —
#: the alternative was moving the whole OCR capability, which is a production
#: routing switch.
_OCR_ENGINES = ("paddle", "doctr", "surya")
_OCR_ENGINES = ("paddle", "doctr", "surya", "olmocr")
def _pin_to_runpod(model: str, engine: str, vram: float, usd: float,
......
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