ocr: hand the card over at a page boundary, so OCR takes turns like any model

Eviction now does the same clean swap the model manager does for a busy model
(release_idle_vram waits for a request boundary rather than evicting mid-request).
OcrManager._release_vram tore pools down immediately, which lost whatever page was
being recognised — and with it the caller's document. Taking turns on one card means
the handover waits for the handover point.

- _Pool tracks in-flight pages (threading-guarded: incremented on the event loop, read
  from the manager's eviction thread) and drain_and_release() waits for them to finish
  before tearing down, bounded by ocr.evict_drain_timeout_s (default 60 s) so a wedged
  engine cannot block a model load forever.
- A rebuild waits out an in-progress drain, instead of re-evicting the model that just
  took its turn and loading into VRAM the drain is about to free.
- Evicted stays temporary: the next OCR request rebuilds and asks for room again, so
  OCR and models can alternate as many times as the traffic asks for.

4 new tests, including the two that matter: an in-flight page survives an eviction that
arrives mid-recognition, and three consecutive OCR/model turns each rebuild and
re-reserve.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0182DhR1eFNPQyGmDrpHbedY

codai/admin/routes.py

@@ -3871,6 +3871,7 @@ def build_settings_dict(c, gpu_cards):
             "surya_model": c.ocr.surya_model,
             "surya_server_url": c.ocr.surya_server_url,
             "build_retry_cooldown_s": c.ocr.build_retry_cooldown_s,
+            "evict_drain_timeout_s": c.ocr.evict_drain_timeout_s,
             "vlm_gpu_memory_utilization": c.ocr.vlm_gpu_memory_utilization,
             "fallback_engines": c.ocr.fallback_engines,
             "olmocr_enabled": c.ocr.olmocr_enabled,
@@ -4605,6 +4606,7 @@ async def api_save_settings(request: Request, username: str = Depends(require_ad
         if "build_retry_cooldown_s" in d: o.build_retry_cooldown_s = _f("build_retry_cooldown_s", o.build_retry_cooldown_s, 0.0)
         if "vlm_gpu_memory_utilization" in d: o.vlm_gpu_memory_utilization = _f("vlm_gpu_memory_utilization", o.vlm_gpu_memory_utilization, 0.0)
         if "fallback_engines" in d: o.fallback_engines = (d.get("fallback_engines") or "auto").strip()
+        if "evict_drain_timeout_s" in d: o.evict_drain_timeout_s = _f("evict_drain_timeout_s", o.evict_drain_timeout_s, 0.0)
         if "olmocr_enabled" in d: o.olmocr_enabled = bool(d["olmocr_enabled"])
         if "olmocr_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()

codai/admin/templates/settings.html

@@ -1074,6 +1074,12 @@
         <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 class="form-group" style="flex:1;min-width:150px">
+        <label class="form-label">Evict drain (s)</label>
+        <input type="number" id="s-ocr-evict-drain" class="form-input" min="0" placeholder="60">
+        <div class="form-hint">How long a model load waits for in-flight pages to finish
+          before taking OCR's VRAM — a clean swap, not a kill.</div>
+      </div>
       <div class="form-group" style="flex:1;min-width:170px">
         <label class="form-label">Fallback engines</label>
         <input type="text" id="s-ocr-fallback" class="form-input" placeholder="auto">
@@ -2011,6 +2017,7 @@ async function loadSettings(){
     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-fallback').value = ocr.fallback_engines ?? 'auto';
+    document.getElementById('s-ocr-evict-drain').value = ocr.evict_drain_timeout_s ?? 60;
     document.getElementById('s-ocr-paddle-enabled').checked = ocr.paddle_enabled !== false;
     document.getElementById('s-ocr-paddle-use-gpu').checked = ocr.paddle_use_gpu !== false;
     document.getElementById('s-ocr-paddle-structure').checked = ocr.paddle_structure !== false;
@@ -2239,6 +2246,7 @@ async function saveSettings(silent){
       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,
       fallback_engines: document.getElementById('s-ocr-fallback').value.trim() || 'auto',
+      evict_drain_timeout_s: parseFloat(document.getElementById('s-ocr-evict-drain').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,

codai/config.py

@@ -780,6 +780,11 @@ class OcrConfig:
     # 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
+    # When a model load evicts OCR, wait up to this long for in-flight pages to finish
+    # before tearing the pool down — the clean swap the model manager does for a busy
+    # model, so OCR and models take turns at a page boundary instead of OCR losing the
+    # page it was mid-way through. 0 = tear down immediately.
+    evict_drain_timeout_s: float = 60.0
     # The share of the card a VLM OCR engine's own vLLM instance claims (surya in vllm
     # mode, olmocr in vllm mode). 0 = use vllm.gpu_memory_utilization, which is sized for
     # serving an LLM and is far more than a 7B OCR model needs — on a shared card that is
@@ -1423,6 +1428,7 @@ class ConfigManager:
                 "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,
+                "evict_drain_timeout_s": self.config.ocr.evict_drain_timeout_s,
                 "vlm_gpu_memory_utilization": self.config.ocr.vlm_gpu_memory_utilization,
                 "fallback_engines": self.config.ocr.fallback_engines,
                 "olmocr_enabled": self.config.ocr.olmocr_enabled,

codai/ocr/manager.py

@@ -28,6 +28,7 @@ bounds the whole subsystem.
 """
 
 import asyncio
+import threading
 import time
 from typing import Dict, List, Optional
 
@@ -121,6 +122,58 @@ class _Pool:
         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)
+        # In-flight page accounting, for the clean swap in drain_and_release(). Touched
+        # from the event loop (recognize) and READ from the model manager's eviction
+        # thread, so it is guarded by a threading primitive rather than an asyncio one.
+        self._inflight = 0
+        self._inflight_lock = threading.Lock()
+        self._idle = threading.Event()
+        self._idle.set()
+        self._draining = False
+
+    def _enter_page(self) -> None:
+        with self._inflight_lock:
+            self._inflight += 1
+            self._idle.clear()
+
+    def _leave_page(self) -> None:
+        with self._inflight_lock:
+            self._inflight = max(0, self._inflight - 1)
+            if self._inflight == 0:
+                self._idle.set()
+
+    def _drain_timeout_s(self) -> float:
+        try:
+            return max(0.0, float(getattr(self.cfg, "evict_drain_timeout_s", 60.0)))
+        except Exception:
+            return 60.0
+
+    def drain_and_release(self) -> float:
+        """Let in-flight pages finish, THEN tear the pool down. Returns GB freed.
+
+        This is the OCR side of the clean swap the model manager already does for a busy
+        model (``release_idle_vram`` waits for a request boundary rather than evicting
+        mid-request). Tearing the pool down under a page that is being recognised loses
+        that page — and the caller's whole document — which is exactly what must not
+        happen when OCR and a model take turns on one card. Bounded by
+        ``ocr.evict_drain_timeout_s`` so a wedged engine cannot block the load forever."""
+        self._draining = True
+        try:
+            budget = self._drain_timeout_s()
+            with self._inflight_lock:
+                busy = self._inflight
+            if busy and budget > 0:
+                print(f"[ocr] engine '{self.name}': {busy} page(s) in flight — waiting up "
+                      f"to {budget:.0f}s for a clean swap before releasing VRAM")
+                if not self._idle.wait(budget):
+                    with self._inflight_lock:
+                        stuck = self._inflight
+                    print(f"[ocr] engine '{self.name}': {stuck} page(s) still in flight "
+                          f"after {budget:.0f}s — releasing anyway (those pages will be "
+                          f"re-read by the fallback engine)")
+            return self.release_sync()
+        finally:
+            self._draining = False
 
     def _cooldown_s(self) -> float:
         try:
@@ -134,6 +187,13 @@ class _Pool:
         async with self._build_lock:
             if self._built:
                 return
+            # A release is draining right now: let it finish before rebuilding, or we
+            # would re-evict the model that just took its turn and load into VRAM the
+            # drain is still about to free. Taking turns means waiting for the handover.
+            _waited = 0.0
+            while self._draining and _waited < self._drain_timeout_s() + 5.0:
+                await asyncio.sleep(0.05)
+                _waited += 0.05
             # 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
@@ -207,6 +267,7 @@ class _Pool:
     async def recognize(self, image) -> OcrPage:
         await self._ensure_built()
         eng: OcrEngine = await self._q.get()
+        self._enter_page()
         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
@@ -217,6 +278,7 @@ class _Pool:
                 pass
             return await asyncio.to_thread(eng.recognize_image, image)
         finally:
+            self._leave_page()
             await self._q.put(eng)
 
 
@@ -261,12 +323,17 @@ class OcrManager:
     def _release_vram(self, needed_gb: float = 999.0) -> float:
         """External VRAM releaser (called as ``fn(needed_gb)`` from the model manager's
         eviction path, SYNC). Tears down every built OCR pool and returns the estimated
-        GB freed. Pools rebuild lazily on the next OCR request. ``needed_gb`` is advisory
-        — OCR pools are all-or-nothing per engine, so we release everything held."""
+        GB freed. Pools rebuild lazily on the next OCR request, so OCR takes its turn on
+        the card again like any other model rather than staying dead for the rest of the
+        process' life. ``needed_gb`` is advisory — OCR pools are all-or-nothing per
+        engine, so we release everything held.
+
+        In-flight pages are allowed to FINISH first (see ``_Pool.drain_and_release``):
+        the handover is a clean swap at a page boundary, not a kill."""
         freed = 0.0
         for pool in list(self._pools.values()):
             try:
-                freed += pool.release_sync()
+                freed += pool.drain_and_release()
             except Exception:
                 pass
         # The managed Surya-2 vLLM subprocess isn't a manager-tracked model and its VRAM

tests/test_ocr_vram_share_and_fallback.py

@@ -264,3 +264,155 @@ def test_a_healthy_engine_is_not_second_guessed(monkeypatch):
 def test_the_fallback_is_on_by_default():
     assert OcrConfig().fallback_engines == "auto"
     assert OcrConfig().vlm_gpu_memory_utilization == 0.35
+    assert OcrConfig().evict_drain_timeout_s == 60.0
+
+
+# ---------------------------------------------------- taking turns on the card
+def test_eviction_waits_for_an_in_flight_page_instead_of_killing_it(monkeypatch):
+    """A model load must not take OCR's VRAM out from under a page being recognised —
+    the handover happens at a page boundary, like the model manager's own clean swap."""
+    import threading
+    started = threading.Event()
+    finish = threading.Event()
+    torn = []
+
+    class _Slow:
+        def __init__(self, cfg):
+            pass
+
+        def prelaunch_vram_gb(self):
+            return 0.0
+
+        def vram_gb(self):
+            return 8.4
+
+        def ensure_loaded(self):
+            pass
+
+        def recognize_image(self, image):
+            started.set()
+            finish.wait(5.0)            # a page mid-recognition
+            return OcrPage(index=0, text="finished anyway")
+
+        def cleanup(self):
+            torn.append(1)
+
+    monkeypatch.setitem(ocr_manager_mod._ENGINE_FACTORIES, "paddle", _Slow)
+    monkeypatch.setattr(ocr_manager_mod, "_evict_for_ocr", lambda gb: None)
+    monkeypatch.setattr(
+        "codai.models.manager.multi_model_manager.register_external_vram_releaser",
+        lambda fn: None)
+
+    m = ocr_manager_mod.OcrManager()
+    m.configure(_cfg(default_engine="paddle", evict_drain_timeout_s=5.0))
+
+    async def _go():
+        task = asyncio.ensure_future(m.ocr_pages([_page()], engine="paddle"))
+        await asyncio.to_thread(started.wait, 5.0)
+        # The model manager's eviction thread arrives mid-page.
+        releaser = asyncio.ensure_future(asyncio.to_thread(m._release_vram, 999.0))
+        await asyncio.sleep(0.2)
+        assert not releaser.done()          # it is WAITING, not tearing down
+        assert not torn                     # nothing released yet
+        finish.set()                        # the page completes
+        freed = await releaser
+        name, pages = await task
+        return freed, name, pages
+
+    freed, name, pages = asyncio.run(_go())
+    assert pages[0].text == "finished anyway"   # the page was NOT lost
+    assert name == "paddle"                     # and no fallback was needed
+    assert torn and freed > 0                   # VRAM released after the drain
+
+
+def test_the_drain_has_a_budget_so_a_wedged_engine_cannot_block_a_load(monkeypatch):
+    import threading
+    started = threading.Event()
+    stuck = threading.Event()
+
+    class _Wedged:
+        def __init__(self, cfg):
+            pass
+
+        def prelaunch_vram_gb(self):
+            return 0.0
+
+        def vram_gb(self):
+            return 8.4
+
+        def ensure_loaded(self):
+            pass
+
+        def recognize_image(self, image):
+            started.set()
+            stuck.wait(30.0)
+            return OcrPage(index=0, text="late")
+
+        def cleanup(self):
+            pass
+
+    monkeypatch.setitem(ocr_manager_mod._ENGINE_FACTORIES, "paddle", _Wedged)
+    monkeypatch.setattr(ocr_manager_mod, "_evict_for_ocr", lambda gb: None)
+    monkeypatch.setattr(
+        "codai.models.manager.multi_model_manager.register_external_vram_releaser",
+        lambda fn: None)
+
+    m = ocr_manager_mod.OcrManager()
+    m.configure(_cfg(default_engine="paddle", evict_drain_timeout_s=0.3))
+
+    async def _go():
+        task = asyncio.ensure_future(m.ocr_pages([_page()], engine="paddle"))
+        await asyncio.to_thread(started.wait, 5.0)
+        freed = await asyncio.to_thread(m._release_vram, 999.0)
+        stuck.set()
+        try:
+            await task
+        except Exception:
+            pass
+        return freed
+
+    assert asyncio.run(_go()) > 0        # gave up on the wedged page and freed the card
+
+
+def test_ocr_rebuilds_and_takes_its_turn_again_after_being_evicted(monkeypatch):
+    """Evicted is not dead: the next request rebuilds, so OCR and models can alternate
+    as many times as the traffic asks for."""
+    builds = []
+
+    class _Eng:
+        def __init__(self, cfg):
+            pass
+
+        def prelaunch_vram_gb(self):
+            return 0.0
+
+        def vram_gb(self):
+            return 8.4
+
+        def ensure_loaded(self):
+            builds.append(1)
+
+        def recognize_image(self, image):
+            return OcrPage(index=0, text="ok")
+
+        def cleanup(self):
+            pass
+
+    monkeypatch.setitem(ocr_manager_mod._ENGINE_FACTORIES, "paddle", _Eng)
+    evicted_for = []
+    monkeypatch.setattr(ocr_manager_mod, "_evict_for_ocr",
+                        lambda gb: evicted_for.append(gb))
+    monkeypatch.setattr(
+        "codai.models.manager.multi_model_manager.register_external_vram_releaser",
+        lambda fn: None)
+
+    m = ocr_manager_mod.OcrManager()
+    m.configure(_cfg(default_engine="paddle", paddle_instances=1))
+
+    for _turn in range(3):
+        name, pages = asyncio.run(m.ocr_pages([_page()], engine="paddle"))
+        assert (name, pages[0].text) == ("paddle", "ok")
+        assert m._release_vram(999.0) > 0        # a model takes the card back
+
+    assert len(builds) == 3          # rebuilt for every turn, never permanently dead
+    assert len(evicted_for) == 3     # and asked for room each time, like any model