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()
......
......@@ -1036,7 +1036,7 @@
<!-- OCR (dedicated document OCR engines) -->
<div class="card">
<div class="card-title">OCR (document → text)</div>
<p class="form-hint" style="margin-bottom:.6rem">Dedicated OCR engines (detection + recognition), <b>not</b> a vision LLM — faithful text with bounding boxes and layout, batched on GPU. Serve <a href="https://github.com/PaddlePaddle/PaddleOCR" target="_blank" rel="noopener">PaddleOCR</a>, <a href="https://github.com/mindee/doctr" target="_blank" rel="noopener">docTR</a> or <a href="https://github.com/VikParuchuri/surya" target="_blank" rel="noopener">Surya</a>, selectable per request. Endpoints: <code>POST /v1/ocr</code>, <code>POST /v1/ocr/batch</code>. Heavy deps install out of band (<code>pip install -r requirements-ocr.txt</code> or <code>./build.sh … --ocr</code>); a missing engine returns HTTP 503.</p>
<p class="form-hint" style="margin-bottom:.6rem">Dedicated OCR engines (detection + recognition), <b>not</b> a vision LLM — faithful text with bounding boxes and layout, batched on GPU. Serve <a href="https://github.com/PaddlePaddle/PaddleOCR" target="_blank" rel="noopener">PaddleOCR</a>, <a href="https://github.com/mindee/doctr" target="_blank" rel="noopener">docTR</a> or <a href="https://github.com/VikParuchuri/surya" target="_blank" rel="noopener">Surya</a> — or <a href="https://huggingface.co/allenai/olmOCR-2-7B-1025" target="_blank" rel="noopener">olmOCR-2</a>, a document VLM, when a page beats them — selectable per request. Endpoints: <code>POST /v1/ocr</code>, <code>POST /v1/ocr/batch</code>. Heavy deps install out of band (<code>pip install -r requirements-ocr.txt</code> or <code>./build.sh … --ocr</code>); a missing engine returns HTTP 503.</p>
<div class="form-row">
<label style="display:flex;align-items:center;gap:.5rem;cursor:pointer">
<input type="checkbox" id="s-ocr-enabled" onchange="toggleOcrFields()">
......@@ -1051,6 +1051,7 @@
<option value="paddle">PaddleOCR</option>
<option value="doctr">docTR</option>
<option value="surya">Surya</option>
<option value="olmocr">olmOCR-2 (VLM)</option>
</select>
</div>
<div class="form-group" style="flex:1;min-width:110px">
......@@ -1065,6 +1066,14 @@
<label class="form-label">Language</label>
<input type="text" id="s-ocr-lang" class="form-input" placeholder="it">
</div>
<div class="form-group" style="flex:1;min-width:130px">
<label class="form-label">Retry cooldown (s)</label>
<input type="number" id="s-ocr-retry-cooldown" class="form-input" min="0" placeholder="60">
</div>
<div class="form-group" style="flex:1;min-width:150px">
<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">
</div>
</div>
<div class="card-title" style="font-size:13px;margin-top:.4rem">PaddleOCR</div>
......@@ -1110,6 +1119,36 @@
</div>
<div style="display:flex;gap:.5rem;align-items:center;flex-wrap:wrap"><button type="button" class="btn btn-sm" onclick="ocrBuildVenv('surya')">Build Surya venv now</button><span id="s-ocr-surya-build-status" class="form-hint"></span></div>
<div class="card-title" style="font-size:13px;margin-top:.4rem">olmOCR-2 <span class="muted" style="font-weight:400">(document VLM)</span></div>
<p class="form-hint" style="margin:0 0 .4rem"><a href="https://huggingface.co/allenai/olmOCR-2-7B-1025" target="_blank" rel="noopener">allenai/olmOCR-2</a> (Apache-2.0) reads a whole page at once — equations as LaTeX, tables as HTML, reading order sorted — and is the one to reach for on old scans, maths and multi-column pages. It is a <b>vision model</b>, so it returns <b>no bounding boxes</b> and no layout regions: use paddle or Surya when you need those.</p>
<div class="form-row" style="display:flex;gap:1rem;flex-wrap:wrap;align-items:center">
<label style="display:flex;align-items:center;gap:.4rem;cursor:pointer"><input type="checkbox" id="s-ocr-olmocr-enabled"><span>Enable</span></label>
<div class="form-group" style="flex:1;min-width:150px"><label class="form-label">Served by</label>
<select id="s-ocr-olmocr-serve" class="form-input">
<option value="model">coderai model (any GPU)</option>
<option value="vllm">coderai vLLM (CUDA)</option>
<option value="server">external OpenAI server</option>
</select>
</div>
<div class="form-group" style="flex:1;min-width:110px"><label class="form-label">Instances</label><input type="number" id="s-ocr-olmocr-instances" class="form-input" min="1" placeholder="1"></div>
<label style="display:flex;align-items:center;gap:.4rem;cursor:pointer"><input type="checkbox" id="s-ocr-olmocr-retry-rotation"><span>Re-ask rotated pages</span></label>
</div>
<p class="form-hint" style="margin:0 0 .3rem"><b>coderai model</b>: register an olmOCR-2 checkpoint as a <i>vision</i> model (GGUF + mmproj, or HF) and name it below — it then gets VRAM accounting, eviction and the thermal governor like any other model, and works on Vulkan/ROCm. <b>vLLM</b> serves the HF id on its own instance (CUDA only). <b>External</b> attaches to a server already running it.</p>
<div class="form-row" style="display:flex;gap:1rem;flex-wrap:wrap">
<div class="form-group" style="flex:2;min-width:200px"><label class="form-label">Model id <span class="muted">(coderai model mode)</span></label><input type="text" id="s-ocr-olmocr-model-id" class="form-input" placeholder="olmocr-2 (a vision model in models.json)"></div>
<div class="form-group" style="flex:2;min-width:200px"><label class="form-label">HF checkpoint <span class="muted">(vLLM / served name)</span></label><input type="text" id="s-ocr-olmocr-model" class="form-input" placeholder="allenai/olmOCR-2-7B-1025-FP8"></div>
</div>
<div class="form-row" style="display:flex;gap:1rem;flex-wrap:wrap">
<div class="form-group" style="flex:2;min-width:200px"><label class="form-label">Server URL <span class="muted">(external mode)</span></label><input type="text" id="s-ocr-olmocr-server-url" class="form-input" placeholder="http://127.0.0.1:8000/v1"></div>
<div class="form-group" style="flex:1;min-width:140px"><label class="form-label">API key <span class="muted">(opt)</span></label><input type="password" id="s-ocr-olmocr-api-key" class="form-input" placeholder=""></div>
</div>
<div class="form-row" style="display:flex;gap:1rem;flex-wrap:wrap">
<div class="form-group" style="flex:1;min-width:120px"><label class="form-label">Longest side (px)</label><input type="number" id="s-ocr-olmocr-longest-side" class="form-input" min="64" placeholder="1288"></div>
<div class="form-group" style="flex:1;min-width:120px"><label class="form-label">Max tokens / page</label><input type="number" id="s-ocr-olmocr-max-tokens" class="form-input" min="64" placeholder="4096"></div>
<div class="form-group" style="flex:1;min-width:110px"><label class="form-label">Temperature</label><input type="number" id="s-ocr-olmocr-temperature" class="form-input" min="0" max="2" step="0.1" placeholder="0.1"></div>
<div class="form-group" style="flex:1;min-width:110px"><label class="form-label">Timeout (s)</label><input type="number" id="s-ocr-olmocr-timeout" class="form-input" min="1" placeholder="300"></div>
</div>
<div class="card-title" style="font-size:13px;margin-top:.4rem">Stamp / signature detection</div>
<div class="form-row" style="display:flex;gap:1rem;flex-wrap:wrap">
<div class="form-group" style="flex:1;min-width:150px">
......@@ -1963,6 +2002,8 @@ async function loadSettings(){
document.getElementById('s-ocr-dpi').value = ocr.dpi ?? 200;
document.getElementById('s-ocr-max-concurrency').value = ocr.max_concurrency ?? 4;
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-vlm-gmu').value = ocr.vlm_gpu_memory_utilization ?? 0;
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-structure').checked = ocr.paddle_structure !== false;
......@@ -1983,6 +2024,18 @@ async function loadSettings(){
document.getElementById('s-ocr-surya-langs').value = ocr.surya_langs ?? 'it';
document.getElementById('s-ocr-surya-venv').value = ocr.surya_venv ?? '';
document.getElementById('s-ocr-surya-auto-build').checked = !!ocr.surya_auto_build;
document.getElementById('s-ocr-olmocr-enabled').checked = !!ocr.olmocr_enabled;
document.getElementById('s-ocr-olmocr-serve').value = ocr.olmocr_serve ?? 'model';
document.getElementById('s-ocr-olmocr-instances').value = ocr.olmocr_instances ?? 1;
document.getElementById('s-ocr-olmocr-retry-rotation').checked = ocr.olmocr_retry_rotation !== false;
document.getElementById('s-ocr-olmocr-model-id').value = ocr.olmocr_model_id ?? '';
document.getElementById('s-ocr-olmocr-model').value = ocr.olmocr_model ?? '';
document.getElementById('s-ocr-olmocr-server-url').value = ocr.olmocr_server_url ?? '';
document.getElementById('s-ocr-olmocr-api-key').value = ocr.olmocr_api_key ?? '';
document.getElementById('s-ocr-olmocr-longest-side').value = ocr.olmocr_longest_side ?? 1288;
document.getElementById('s-ocr-olmocr-max-tokens').value = ocr.olmocr_max_tokens ?? 4096;
document.getElementById('s-ocr-olmocr-temperature').value = ocr.olmocr_temperature ?? 0.1;
document.getElementById('s-ocr-olmocr-timeout').value = ocr.olmocr_timeout ?? 300;
document.getElementById('s-ocr-detect-mode').value = ocr.detect_mode ?? 'off';
document.getElementById('s-ocr-detect-model').value = ocr.detect_model_path ?? '';
document.getElementById('s-ocr-detect-conf').value = ocr.detect_conf ?? 0.35;
......@@ -2176,6 +2229,8 @@ async function saveSettings(silent){
dpi: parseInt(document.getElementById('s-ocr-dpi').value) || 200,
max_concurrency: parseInt(document.getElementById('s-ocr-max-concurrency').value) || 4,
lang: document.getElementById('s-ocr-lang').value.trim() || 'it',
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,
paddle_enabled: document.getElementById('s-ocr-paddle-enabled').checked,
paddle_use_gpu: document.getElementById('s-ocr-paddle-use-gpu').checked,
paddle_structure: document.getElementById('s-ocr-paddle-structure').checked,
......@@ -2196,6 +2251,18 @@ async function saveSettings(silent){
surya_langs: document.getElementById('s-ocr-surya-langs').value.trim() || 'it',
surya_venv: document.getElementById('s-ocr-surya-venv').value.trim(),
surya_auto_build: document.getElementById('s-ocr-surya-auto-build').checked,
olmocr_enabled: document.getElementById('s-ocr-olmocr-enabled').checked,
olmocr_serve: document.getElementById('s-ocr-olmocr-serve').value || 'model',
olmocr_instances: parseInt(document.getElementById('s-ocr-olmocr-instances').value) || 1,
olmocr_retry_rotation: document.getElementById('s-ocr-olmocr-retry-rotation').checked,
olmocr_model_id: document.getElementById('s-ocr-olmocr-model-id').value.trim(),
olmocr_model: document.getElementById('s-ocr-olmocr-model').value.trim(),
olmocr_server_url: document.getElementById('s-ocr-olmocr-server-url').value.trim(),
olmocr_api_key: document.getElementById('s-ocr-olmocr-api-key').value.trim(),
olmocr_longest_side: parseInt(document.getElementById('s-ocr-olmocr-longest-side').value) || 1288,
olmocr_max_tokens: parseInt(document.getElementById('s-ocr-olmocr-max-tokens').value) || 4096,
olmocr_temperature: parseFloat(document.getElementById('s-ocr-olmocr-temperature').value) || 0.1,
olmocr_timeout: parseFloat(document.getElementById('s-ocr-olmocr-timeout').value) || 300,
detect_mode: document.getElementById('s-ocr-detect-mode').value || 'off',
detect_model_path: document.getElementById('s-ocr-detect-model').value.trim(),
detect_conf: parseFloat(document.getElementById('s-ocr-detect-conf').value) || 0.35,
......
......@@ -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 = []
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: evict other models to make room (VRAM eviction), now
# that we know its real footprint. Mirrors the engine-load path.
if not self._instances:
# 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:
......
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
# GPLv3 - see the project LICENSE.
"""olmOCR-2 engine — the AllenAI document VLM, driven over an OpenAI chat endpoint.
Unlike paddle/docTR/Surya-local, olmOCR-2 (``allenai/olmOCR-2-7B-1025``, a Qwen2.5-VL-7B
fine-tune, Apache-2.0) is a *vision language model*: one page image in, the whole page's
text out, in reading order, with equations as LaTeX and tables as HTML. It has no text
detector, so it returns NO per-line bounding boxes — pick it for transcription fidelity on
hard documents (old scans, maths, multi-column, tables), and paddle/Surya when you need
boxes or layout regions.
Three ways to serve it (``ocr.olmocr_serve``):
- ``model`` — through coderai's OWN model manager: register the checkpoint in
models.json as a vision model (GGUF+mmproj on llama.cpp, or an HF/qwenvl entry) and put
its id in ``ocr.olmocr_model_id``. Gets VRAM accounting, eviction, quantisation and the
thermal governor for free, and is the only mode that works on a Vulkan/ROCm box.
- ``vllm`` — coderai's vLLM backend serves ``ocr.olmocr_model`` on its own instance
(continuous batching, CUDA only), exactly like the Surya-2 path.
- ``server`` — attach to an OpenAI-compatible server already running the model
(``ocr.olmocr_server_url``): llama-server, a remote vLLM, another coderai.
The prompt is olmOCR's own "no-anchoring v4 YAML" prompt, so the model answers in its
trained format: a YAML front matter block (primary_language, is_rotation_valid,
rotation_correction, is_table, is_diagram) followed by the page as markdown. The front
matter is parsed into ``OcrPage.meta`` and stripped from ``text``; when the model reports
the page as rotated we re-render it turned and ask once more (what the olmOCR toolkit's
pipeline does).
"""
import base64
import io
import json
import threading
from typing import Optional
from codai.ocr.base import OcrEngine, OcrPage, OcrError
#: olmocr.prompts.build_no_anchoring_v4_yaml_prompt() — kept verbatim: the model was
#: trained on this wording, so paraphrasing it costs accuracy. Mirrored here rather than
#: importing olmocr, which would drag the whole toolkit (and its pinned deps) in.
OLMOCR_PROMPT = (
"Attached is one page of a document that you must process. "
"Just return the plain text representation of this document as if you were reading it naturally. "
"Convert equations to LateX and tables to HTML.\n"
"If there are any figures or charts, label them with the following markdown syntax "
"![Alt text describing the contents of the figure](page_startx_starty_width_height.png)\n"
"Return your output as markdown, with a front matter section on top specifying values for the "
"primary_language, is_rotation_valid, rotation_correction, is_table, and is_diagram parameters."
)
#: The model card's rendering rule: longest side 1288 px. Off-size pages still work, but
#: this is what it was trained and benchmarked on.
DEFAULT_LONGEST_SIDE = 1288
VALID_MODES = ("model", "vllm", "server")
DEFAULT_MODEL = "allenai/olmOCR-2-7B-1025-FP8"
_FRONT_MATTER_KEYS = (
"primary_language", "is_rotation_valid", "rotation_correction",
"is_table", "is_diagram",
)
# ---------------------------------------------------------------------------
# front matter
# ---------------------------------------------------------------------------
def _coerce(value: str):
v = (value or "").strip().strip('"').strip("'")
lo = v.lower()
if lo in ("true", "yes"):
return True
if lo in ("false", "no"):
return False
if lo in ("none", "null", "~", ""):
return None
try:
return int(v)
except ValueError:
pass
try:
return float(v)
except ValueError:
pass
return v
def parse_front_matter(content: str) -> (dict, str):
"""Split olmOCR's reply into (metadata, body).
The reply normally opens with a ``---`` fenced YAML block. Models drift: the fence may
be wrapped in a markdown code block, the opening fence may be missing entirely, or the
block may not be there at all. Anything we cannot read as front matter stays in the
body — losing page text to a parse detail would be much worse than a missing flag.
"""
s = (content or "").strip()
if s.startswith("```"):
# ```markdown\n---\n… — drop the fence line, and a trailing fence if it closes.
nl = s.find("\n")
s = s[nl + 1:] if nl != -1 else s
if s.rstrip().endswith("```"):
s = s.rstrip()[:-3]
s = s.strip()
lines = s.splitlines()
start = 0
if lines and lines[0].strip() == "---":
start = 1
elif lines and lines[0].split(":", 1)[0].strip() in _FRONT_MATTER_KEYS:
start = 0 # fence omitted but the first key is one of ours
else:
return {}, s
meta, end = {}, None
for i in range(start, len(lines)):
line = lines[i]
if line.strip() == "---":
end = i
break
if not line.strip():
if start == 0:
end = i # unfenced block ends at the first blank line
break
continue
if ":" not in line:
if start == 0:
end = i
break
continue
key, val = line.split(":", 1)
key = key.strip()
if key:
meta[key] = _coerce(val)
if not meta:
return {}, s
body = "\n".join(lines[(end + 1) if end is not None else len(lines):]).strip()
return meta, body
# ---------------------------------------------------------------------------
# image prep
# ---------------------------------------------------------------------------
def prepare_image(image, longest_side: int = DEFAULT_LONGEST_SIDE) -> str:
"""Return a ``data:image/png;base64,…`` URI of ``image`` scaled to ``longest_side``.
Only ever scales DOWN past the target: upsampling a small crop buys nothing and costs
vision tokens.
"""
from PIL import Image
img = image.convert("RGB")
target = max(64, int(longest_side or DEFAULT_LONGEST_SIDE))
w, h = img.size
longest = max(w, h)
if longest > target:
scale = target / float(longest)
img = img.resize((max(1, int(round(w * scale))), max(1, int(round(h * scale)))),
Image.LANCZOS)
buf = io.BytesIO()
img.save(buf, format="PNG")
return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("ascii")
# ---------------------------------------------------------------------------
# engine
# ---------------------------------------------------------------------------
_loop = None
_loop_lock = threading.Lock()
def _own_loop():
"""A module-owned event loop, used only when no host loop was handed to us.
``recognize_image`` runs in a worker thread (the pool calls it through
``asyncio.to_thread``), so reaching coderai's async model path needs a loop. The pool
passes the app's own loop — this is the fallback for direct/unit-test calls."""
import asyncio
global _loop
with _loop_lock:
if _loop is not None and not _loop.is_closed():
return _loop
_loop = asyncio.new_event_loop()
threading.Thread(target=_loop.run_forever, name="olmocr-loop",
daemon=True).start()
return _loop
class OlmOcrEngine(OcrEngine):
name = "olmocr"
def __init__(self, cfg):
super().__init__(cfg)
self._base = "" # OpenAI base URL for the http modes
self._host_loop = None # set by the pool: the app's event loop
# -- config ------------------------------------------------------------
def _mode(self) -> str:
mode = (getattr(self.cfg, "olmocr_serve", "model") or "model").strip().lower()
if mode not in VALID_MODES:
raise OcrError(
f"invalid ocr.olmocr_serve '{mode}' (use {'/'.join(VALID_MODES)})",
status=400)
return mode
def _model_name(self) -> str:
return (getattr(self.cfg, "olmocr_model", "") or DEFAULT_MODEL).strip()
# -- lifecycle ---------------------------------------------------------
def load(self) -> None:
mode = self._mode()
self._base = ""
if mode == "model":
if not (getattr(self.cfg, "olmocr_model_id", "") or "").strip():
raise OcrError(
"olmOCR 'model' mode needs ocr.olmocr_model_id — the id of a VISION "
"model in models.json holding an olmOCR-2 checkpoint (GGUF+mmproj or "
"HF). Set it, or switch ocr.olmocr_serve to vllm/server.",
status=400)
return
if mode == "server":
url = (getattr(self.cfg, "olmocr_server_url", "") or "").strip()
if not url:
raise OcrError(
"olmOCR 'server' mode needs ocr.olmocr_server_url (an OpenAI-compatible "
"endpoint already serving olmOCR-2).", status=400)
self._base = url.rstrip("/")
if not self._base.endswith("/v1"):
self._base += "/v1"
return
# vllm: boot (or reuse) coderai's vLLM backend on the olmOCR checkpoint.
from codai.api import vllm_worker
from codai.models.manager import get_active_vllm_config
vcfg = get_active_vllm_config()
if vcfg is None:
raise OcrError("olmOCR vllm mode needs the vLLM backend configured", status=400)
model = self._model_name()
base = vllm_worker.ensure_service(vcfg, model_path=model, served_name=model,
gpu_memory_utilization=self._vlm_gmu())
self._base = base.rstrip("/") + "/v1"
def cleanup(self) -> None:
# Nothing of ours holds VRAM: in 'model' mode the model manager owns the weights,
# in 'vllm' mode the shared service does (stopped by the manager's releaser), and
# 'server' mode is someone else's process.
self._loaded = False
def vram_gb(self) -> float:
# Nothing is held PER INSTANCE: an instance is one in-flight page against a shared
# server. In 'model' mode the model manager already accounts for the weights,
# 'server' mode is another process entirely, and the 'vllm' instance is accounted
# for once through prelaunch_vram_gb() and released by the manager's stop estimate.
return 0.0
def prelaunch_vram_gb(self) -> float:
"""vLLM claims gpu_memory_utilization × the whole card before it will serve, and
refuses to start when that much is not free — so it has to be freed BEFORE the
load, not after it (see Surya's crash loop)."""
return self._vllm_need() if self._mode() == "vllm" else 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 _vllm_need(self) -> float:
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
# -- inference ---------------------------------------------------------
def recognize_image(self, image) -> OcrPage:
w, h = image.size
page = OcrPage(index=0, width=int(w), height=int(h))
content = self._ask(image)
meta, body = parse_front_matter(content)
rotation = meta.get("rotation_correction") or 0
valid = meta.get("is_rotation_valid")
if (valid is False or rotation) and bool(getattr(self.cfg, "olmocr_retry_rotation", True)):
try:
turned = int(rotation) % 360
except Exception:
turned = 0
if turned:
# PIL rotates counter-clockwise; olmOCR reports the correction to APPLY,
# so expand=True and the same sign the toolkit uses.
retry_img = image.rotate(-turned, expand=True)
retry = self._ask(retry_img)
rmeta, rbody = parse_front_matter(retry)
if rbody.strip():
meta, body = rmeta, rbody
meta["rotation_applied"] = turned
rw, rh = retry_img.size
page.width, page.height = int(rw), int(rh)
page.text = body
page.meta = {k: v for k, v in meta.items() if v is not None}
# No detector, so no lines/regions/tables boxes — the text IS the result.
return page
def _ask(self, image) -> str:
"""One page → the model's raw reply text."""
data_uri = prepare_image(
image, int(getattr(self.cfg, "olmocr_longest_side", DEFAULT_LONGEST_SIDE) or
DEFAULT_LONGEST_SIDE))
messages = [{
"role": "user",
"content": [
{"type": "text", "text": OLMOCR_PROMPT},
{"type": "image_url", "image_url": {"url": data_uri}},
],
}]
max_tokens = int(getattr(self.cfg, "olmocr_max_tokens", 4096) or 4096)
temperature = float(getattr(self.cfg, "olmocr_temperature", 0.1) or 0.0)
if self._mode() == "model":
return self._ask_local(messages, max_tokens, temperature)
return self._ask_http(messages, max_tokens, temperature)
def _ask_local(self, messages, max_tokens, temperature) -> str:
"""Through coderai's own in-process chat handler (model mode)."""
import asyncio
from codai.api.text import chat_completions
from codai.pydantic.textrequest import ChatCompletionRequest
model_id = (getattr(self.cfg, "olmocr_model_id", "") or "").strip()
req = ChatCompletionRequest(
model=model_id, messages=messages, stream=False,
temperature=temperature, max_tokens=max_tokens,
)
loop = self._host_loop
if loop is None or loop.is_closed():
loop = _own_loop()
fut = asyncio.run_coroutine_threadsafe(chat_completions(req, None), loop)
try:
result = fut.result(timeout=float(getattr(self.cfg, "olmocr_timeout", 300.0) or 300.0))
except Exception as e:
raise OcrError(f"olmOCR model '{model_id}' call failed: {e}", status=500)
content = _content_of(result)
if not content.strip():
raise OcrError(f"olmOCR model '{model_id}' returned no text", status=500)
return content
def _ask_http(self, messages, max_tokens, temperature) -> str:
"""Against an OpenAI-compatible server (vllm/server modes)."""
import requests
payload = {
"model": self._model_name(),
"messages": messages,
"max_tokens": max_tokens,
"temperature": temperature,
"stream": False,
}
headers = {"Content-Type": "application/json"}
key = (getattr(self.cfg, "olmocr_api_key", "") or "").strip()
if key:
headers["Authorization"] = f"Bearer {key}"
url = self._base.rstrip("/") + "/chat/completions"
try:
r = requests.post(url, json=payload, headers=headers,
timeout=float(getattr(self.cfg, "olmocr_timeout", 300.0) or 300.0))
except Exception as e:
raise OcrError(f"olmOCR server {url} unreachable: {e}", status=503)
if r.status_code >= 400:
raise OcrError(f"olmOCR server {url} returned HTTP {r.status_code}: "
f"{r.text[:300]}", status=502)
try:
content = _content_of(r.json())
except Exception as e:
raise OcrError(f"olmOCR server returned unparseable JSON: {e}", status=502)
if not content.strip():
raise OcrError("olmOCR server returned no text", status=502)
return content
def _content_of(result) -> str:
"""Assistant message content out of a chat-completions result, whatever shape it
arrives in (dict, pydantic model, JSONResponse)."""
r = None
if isinstance(result, dict):
r = result
elif hasattr(result, "model_dump"):
try:
r = result.model_dump()
except Exception:
r = None
if r is None and hasattr(result, "body"):
try:
r = json.loads(result.body)
except Exception:
r = None
if not isinstance(r, dict):
return ""
choices = r.get("choices") or []
if not choices:
return ""
first = choices[0]
if not isinstance(first, dict):
return ""
msg = first.get("message")
if isinstance(msg, dict):
c = msg.get("content")
if isinstance(c, list): # multipart reply — join its text parts
return "".join(p.get("text", "") for p in c if isinstance(p, dict))
return c or ""
return first.get("text") or ""
......@@ -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
......
"""olmOCR-2 as an alternative OCR engine, plus the two things the Surya crash loop
exposed in the pool: evict BEFORE a server-backed engine loads, and a cooldown after a
failed build instead of re-attempting it on every request.
No model, no GPU, no network: the engine's HTTP/model call is stubbed, images are tiny
PIL pages.
"""
import asyncio
import sys
import time
import types
from pathlib import Path
import pytest
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from codai.config import OcrConfig
from codai.ocr.base import OcrError, OcrPage
from codai.ocr.olmocr import (
DEFAULT_LONGEST_SIDE, OLMOCR_PROMPT, OlmOcrEngine, parse_front_matter, prepare_image,
)
from codai.ocr import manager as ocr_manager_mod
def _page(w=900, h=1200, color=(255, 255, 255)):
from PIL import Image
return Image.new("RGB", (w, h), color)
def _cfg(**kw):
c = OcrConfig(enabled=True, olmocr_enabled=True, default_engine="olmocr")
for k, v in kw.items():
setattr(c, k, v)
return c
# ------------------------------------------------------------ front matter
def test_front_matter_split():
meta, body = parse_front_matter(
"---\n"
"primary_language: it\n"
"is_rotation_valid: True\n"
"rotation_correction: 0\n"
"is_table: False\n"
"is_diagram: False\n"
"---\n"
"TRIBUNALE DI TORINO\n\nSentenza n. 42\n"
)
assert meta == {
"primary_language": "it", "is_rotation_valid": True,
"rotation_correction": 0, "is_table": False, "is_diagram": False,
}
assert body.startswith("TRIBUNALE DI TORINO")
assert "---" not in body
def test_front_matter_in_a_code_fence():
meta, body = parse_front_matter(
"```markdown\n---\nprimary_language: en\nis_table: True\n---\n<table></table>\n```")
assert meta["primary_language"] == "en" and meta["is_table"] is True
assert body == "<table></table>"
def test_front_matter_unfenced_and_absent():
meta, body = parse_front_matter("primary_language: en\nis_table: False\n\nHello page")
assert meta["primary_language"] == "en" and body == "Hello page"
# No front matter at all: every byte stays in the body rather than being eaten.
meta, body = parse_front_matter("Just the page text, no metadata")
assert meta == {} and body == "Just the page text, no metadata"
# ---------------------------------------------------------------- image prep
def test_prepare_image_scales_down_to_1288_and_never_up():
big = prepare_image(_page(4000, 2000))
assert big.startswith("data:image/png;base64,")
from PIL import Image
import base64, io
img = Image.open(io.BytesIO(base64.b64decode(big.split(",", 1)[1])))
assert max(img.size) == DEFAULT_LONGEST_SIDE
assert img.size == (1288, 644)
small = prepare_image(_page(300, 200))
img = Image.open(io.BytesIO(base64.b64decode(small.split(",", 1)[1])))
assert img.size == (300, 200)
# ------------------------------------------------------------------- engine
class _Stub:
"""Records the requests the engine makes and replies with canned pages."""
def __init__(self, *replies):
self.replies = list(replies)
self.calls = []
def __call__(self, messages, max_tokens, temperature):
self.calls.append({"messages": messages, "max_tokens": max_tokens,
"temperature": temperature})
return self.replies.pop(0) if len(self.replies) > 1 else self.replies[0]
def _engine(cfg, stub):
eng = OlmOcrEngine(cfg)
eng._ask_local = stub
eng._ask_http = stub
eng.ensure_loaded()
return eng
def test_recognize_uses_the_trained_prompt_and_returns_text_plus_meta():
stub = _Stub("---\nprimary_language: it\nis_table: True\n---\n# Atto\n\nTesto.")
eng = _engine(_cfg(olmocr_model_id="olmocr-2"), stub)
page = eng.recognize_image(_page(800, 1000))
msg = stub.calls[0]["messages"][0]
assert msg["role"] == "user"
assert msg["content"][0]["text"] == OLMOCR_PROMPT # verbatim, as trained
assert msg["content"][1]["image_url"]["url"].startswith("data:image/png;base64,")
assert page.text == "# Atto\n\nTesto."
assert page.meta["primary_language"] == "it" and page.meta["is_table"] is True
assert page.width == 800 and page.height == 1000
# A VLM has no detector: no boxes, and the response omits an empty meta for the
# engines that report none.
assert page.lines == [] and page.regions == []
assert "meta" in page.to_dict()
assert "meta" not in OcrPage(index=0, text="x").to_dict()
def test_rotated_page_is_asked_once_more_turned():
stub = _Stub(
"---\nis_rotation_valid: False\nrotation_correction: 90\n---\ngibberish",
"---\nis_rotation_valid: True\nrotation_correction: 0\n---\nthe real text",
)
eng = _engine(_cfg(olmocr_model_id="olmocr-2"), stub)
page = eng.recognize_image(_page(900, 1200))
assert len(stub.calls) == 2
assert page.text == "the real text"
assert page.meta["rotation_applied"] == 90
# The retry ran on the turned page, so the reported size is the turned one.
assert (page.width, page.height) == (1200, 900)
def test_rotation_retry_can_be_turned_off():
stub = _Stub("---\nis_rotation_valid: False\nrotation_correction: 90\n---\ngibberish")
eng = _engine(_cfg(olmocr_model_id="olmocr-2", olmocr_retry_rotation=False), stub)
page = eng.recognize_image(_page())
assert len(stub.calls) == 1 and page.text == "gibberish"
def test_each_mode_refuses_to_load_without_what_it_needs():
with pytest.raises(OcrError) as e:
OlmOcrEngine(_cfg(olmocr_serve="model", olmocr_model_id="")).load()
assert "olmocr_model_id" in str(e.value) and e.value.status == 400
with pytest.raises(OcrError) as e:
OlmOcrEngine(_cfg(olmocr_serve="server", olmocr_server_url="")).load()
assert "olmocr_server_url" in str(e.value)
with pytest.raises(OcrError) as e:
OlmOcrEngine(_cfg(olmocr_serve="nope")).load()
assert "olmocr_serve" in str(e.value)
def test_server_mode_normalises_the_base_url():
eng = OlmOcrEngine(_cfg(olmocr_serve="server", olmocr_server_url="http://box:8000/"))
eng.load()
assert eng._base == "http://box:8000/v1"
eng2 = OlmOcrEngine(_cfg(olmocr_serve="server", olmocr_server_url="http://box:8000/v1"))
eng2.load()
assert eng2._base == "http://box:8000/v1"
def test_http_failures_map_to_useful_statuses(monkeypatch):
cfg = _cfg(olmocr_serve="server", olmocr_server_url="http://box:8000/v1")
eng = OlmOcrEngine(cfg)
eng.load()
fake = types.SimpleNamespace()
class _Resp:
status_code = 500
text = "boom"
def json(self):
return {}
fake.post = lambda *a, **k: _Resp()
monkeypatch.setitem(sys.modules, "requests", fake)
with pytest.raises(OcrError) as e:
eng.recognize_image(_page(64, 64))
assert e.value.status == 502
def _raise(*a, **k):
raise OSError("connection refused")
fake.post = _raise
with pytest.raises(OcrError) as e:
eng.recognize_image(_page(64, 64))
assert e.value.status == 503
def test_only_the_vllm_mode_asks_for_room_up_front(monkeypatch):
monkeypatch.setattr(OlmOcrEngine, "_vllm_need", staticmethod(lambda: 14.1))
assert OlmOcrEngine(_cfg(olmocr_serve="model", olmocr_model_id="x")).prelaunch_vram_gb() == 0.0
assert OlmOcrEngine(_cfg(olmocr_serve="server", olmocr_server_url="u")).prelaunch_vram_gb() == 0.0
assert OlmOcrEngine(_cfg(olmocr_serve="vllm")).prelaunch_vram_gb() == 14.1
# --------------------------------------------------------------- the manager
def test_the_engine_is_registered_and_gated_by_its_own_flag():
m = ocr_manager_mod.OcrManager()
m.configure(_cfg())
assert m.resolve_engine(None) == "olmocr" # it is this config's default
assert m.resolve_engine("olmocr") == "olmocr"
m.configure(_cfg(olmocr_enabled=False, default_engine="paddle"))
with pytest.raises(OcrError) as e:
m.resolve_engine("olmocr")
assert "not enabled" in str(e.value)
def test_the_pool_frees_vram_before_a_server_backed_engine_loads(monkeypatch):
"""The Surya-2 bug: eviction ran AFTER the load, and the load was what failed."""
order = []
class _Eng:
def __init__(self, cfg):
pass
def prelaunch_vram_gb(self):
return 14.1
def vram_gb(self):
return 14.1
def ensure_loaded(self):
order.append("load")
def recognize_image(self, image):
return OcrPage(index=0, text="ok")
def cleanup(self):
pass
monkeypatch.setitem(ocr_manager_mod._ENGINE_FACTORIES, "olmocr", _Eng)
monkeypatch.setattr(ocr_manager_mod, "_evict_for_ocr",
lambda gb: order.append(f"evict:{gb:.1f}"))
m = ocr_manager_mod.OcrManager()
m.configure(_cfg(olmocr_instances=1))
async def _go():
_name, pages = await m.ocr_pages([_page(64, 64)], engine="olmocr")
return pages
pages = asyncio.run(_go())
assert pages[0].text == "ok"
assert order[0] == "evict:14.1" and order[1] == "load"
def test_a_failed_build_fails_fast_for_the_cooldown_then_tries_again(monkeypatch):
attempts = []
class _Broken:
def __init__(self, cfg):
pass
def prelaunch_vram_gb(self):
return 0.0
def vram_gb(self):
return 0.0
def ensure_loaded(self):
attempts.append(time.time())
raise OcrError("vLLM exited before becoming ready: 11.3 GB free, 14.1 wanted",
status=503)
def cleanup(self):
pass
monkeypatch.setitem(ocr_manager_mod._ENGINE_FACTORIES, "olmocr", _Broken)
monkeypatch.setattr(ocr_manager_mod, "_evict_for_ocr", lambda gb: None)
m = ocr_manager_mod.OcrManager()
m.configure(_cfg(build_retry_cooldown_s=30.0))
async def _one():
return await m.ocr_pages([_page(64, 64)], engine="olmocr")
for _ in range(3):
with pytest.raises(OcrError) as e:
asyncio.run(_one())
# Only the FIRST request paid for a boot attempt; the other two were refused
# immediately, with the reason and the time left.
assert len(attempts) == 1
assert "11.3 GB free" in str(e.value) and "cooldown" in str(e.value)
assert e.value.status == 503
# Once the cooldown is up the engine gets another chance.
m2 = ocr_manager_mod.OcrManager()
m2.configure(_cfg(build_retry_cooldown_s=0.0))
for _ in range(2):
with pytest.raises(OcrError):
asyncio.run(m2.ocr_pages([_page(64, 64)], engine="olmocr"))
assert len(attempts) == 3
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."""
seen = {}
fake_worker = types.SimpleNamespace(
ensure_service=lambda vcfg, model_path=None, served_name=None,
gpu_memory_utilization=None: (
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,
)
monkeypatch.setitem(sys.modules, "codai.api.vllm_worker", fake_worker)
monkeypatch.setattr("codai.models.manager.get_active_vllm_config",
lambda: types.SimpleNamespace(gpu_memory_utilization=0.9))
cfg = _cfg(olmocr_serve="vllm", vlm_gpu_memory_utilization=0.35)
eng = OlmOcrEngine(cfg)
assert eng.prelaunch_vram_gb() == pytest.approx(8.4) # 0.35 of the card, not 0.9
eng.load()
assert seen["model"] == "allenai/olmOCR-2-7B-1025-FP8"
assert seen["share"] == 0.35
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
# crash-looped Surya-2 on a box that had other work resident.
assert OcrConfig().vlm_gpu_memory_utilization == 0.35
# 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))
assert eng2.prelaunch_vram_gb() == pytest.approx(21.6)
def test_a_shared_server_is_reserved_for_ONCE_not_per_instance(monkeypatch):
"""surya_instances was 24 in production: multiplying the shared vLLM's demand by the
pool size would have asked the manager to free 200 GB of a 24 GB card."""
asked = []
class _Shared:
def __init__(self, cfg):
pass
def prelaunch_vram_gb(self):
return 8.4 # the whole pool's share of the card, not each instance's
def vram_gb(self):
return 0.05 # a worker that only speaks HTTP to that server
def ensure_loaded(self):
pass
def recognize_image(self, image):
return OcrPage(index=0, text="ok")
def cleanup(self):
pass
monkeypatch.setitem(ocr_manager_mod._ENGINE_FACTORIES, "olmocr", _Shared)
monkeypatch.setattr(ocr_manager_mod, "_evict_for_ocr", lambda gb: asked.append(gb))
m = ocr_manager_mod.OcrManager()
m.configure(_cfg(olmocr_instances=24))
asyncio.run(m.ocr_pages([_page(64, 64)], engine="olmocr"))
# Once, for the server's share — and NOT again after the load, for memory that
# server has already taken.
assert asked == [8.4]
def test_an_engine_without_a_server_is_still_evicted_for_per_instance(monkeypatch):
"""The paddle/docTR path is unchanged: nothing is reserved up front, and the real
footprint is evicted for once it is known, times the pool size."""
asked = []
class _Local:
def __init__(self, cfg):
pass
def prelaunch_vram_gb(self):
return 0.0
def vram_gb(self):
return 0.7
def ensure_loaded(self):
asked.append("load")
def recognize_image(self, image):
return OcrPage(index=0, text="ok")
def cleanup(self):
pass
monkeypatch.setitem(ocr_manager_mod._ENGINE_FACTORIES, "doctr", _Local)
monkeypatch.setattr(ocr_manager_mod, "_evict_for_ocr",
lambda gb: asked.append(round(gb, 2)))
m = ocr_manager_mod.OcrManager()
m.configure(_cfg(doctr_enabled=True, doctr_instances=3, default_engine="doctr"))
asyncio.run(m.ocr_pages([_page(64, 64)], engine="doctr"))
assert asked == ["load", 2.1, "load", "load"]
def test_the_model_manager_can_evict_ocr_including_both_vlm_servers(monkeypatch):
"""The other direction: loading a model asks OCR to give its VRAM back. Pools are torn
down AND the managed vLLM instance is stopped — tearing down the pool alone leaves
gpu_memory_utilization × the card held by a subprocess the manager cannot see."""
registered = []
monkeypatch.setattr(
"codai.models.manager.multi_model_manager.register_external_vram_releaser",
lambda fn: registered.append(fn))
stopped = []
fake_worker = types.SimpleNamespace(
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],
planned_vram_gb=lambda vcfg, share=None: 8.4,
ensure_service=lambda *a, **k: "http://127.0.0.1:1/",
)
monkeypatch.setitem(sys.modules, "codai.api.vllm_worker", fake_worker)
monkeypatch.setattr("codai.models.manager.get_active_vllm_config",
lambda: types.SimpleNamespace(gpu_memory_utilization=0.9))
released = []
class _Eng:
def __init__(self, cfg):
pass
def prelaunch_vram_gb(self):
return 0.0
def vram_gb(self):
return 0.05
def ensure_loaded(self):
pass
def recognize_image(self, image):
return OcrPage(index=0, text="ok")
def cleanup(self):
released.append(1)
monkeypatch.setitem(ocr_manager_mod._ENGINE_FACTORIES, "olmocr", _Eng)
monkeypatch.setattr(ocr_manager_mod, "_evict_for_ocr", lambda gb: None)
m = ocr_manager_mod.OcrManager()
m.configure(_cfg(olmocr_serve="vllm", olmocr_instances=2,
surya_enabled=True, surya_accept_license=True, surya_serve="vllm"))
assert m._release_vram in registered # the manager now knows how to reclaim OCR
asyncio.run(m.ocr_pages([_page(64, 64)], engine="olmocr"))
freed = m._release_vram(999.0)
assert len(released) == 2 # both instances torn down
assert freed >= 16.0 # pool + BOTH vLLM services
assert ("datalab-to/surya-ocr-2", 0.35) in stopped
assert ("allenai/olmOCR-2-7B-1025-FP8", 0.35) in stopped
# And the pool rebuilds on the next request rather than staying dead.
asyncio.run(m.ocr_pages([_page(64, 64)], engine="olmocr"))
assert len(released) == 2
......@@ -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