front: OCR takes its turn at the shared-GPU swap gate, unless it runs off-box

/v1/ocr bypassed the gate entirely because it is not in the router's _INFERENCE_PATHS
— and it should not be, since that predicate drives the per-model queue, pins,
spillover and keepalive, none of which fit OCR. But OCR is GPU work: an OCR request
could start a surya vLLM boot while an LLM was mid-generation on the same card, and
vice versa. The gate is what turns that contention into turns.

- is_gpu_inference_path() = is_inference_path() + /v1/ocr and /v1/ocr/batch, and that
  is what the swap gate keys on now. is_inference_path() is unchanged, so no other
  front machinery starts applying to OCR.
- All OCR shares ONE gate owner ("ocr"): the OCR manager holds one set of pools and
  evicts between them itself, so paddle and surya are the same tenant for the purpose
  of who owns the card. It also means the front need not read the `engine` field out of
  a multipart body carrying a 100 MB PDF.
- _computed_off_this_gpu(): the gate AND the local wait queue are skipped when the work
  is computed elsewhere — a cluster node (engine.remote), a RunPod pod (model backend),
  or an external API endpoint (every enabled OCR engine pointing off-box). Such a
  request holds no local VRAM, so a turn on this card's gate would add latency and
  swaps for nothing. Applied at all four queue sites (direct, keepalive, both brokered)
  and both gate sites.

15 new tests, the front proxy's first: path classification, the owner key, each off-box
case, and that an OCR tenant and a model tenant alternate without starving each other.
Co-Authored-By: 's avatarClaude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0182DhR1eFNPQyGmDrpHbedY
parent 4be32ede
......@@ -393,7 +393,8 @@ class FrontProxy:
# direct proxy path, so brokered and direct requests share one queue.
_qkey = None
if (method.upper() == "POST" and _router.is_inference_path(path)
and self._task_kind(path) == "text"):
and self._task_kind(path) == "text"
and not self._computed_off_this_gpu(engine, model, path)):
_qkey = self._queue_key(model)
try:
await self.reqqueue.acquire(
......@@ -546,7 +547,8 @@ class FrontProxy:
pass
_qkey = None
if (method.upper() == "POST" and _is_infer
and self._task_kind(path) == "text"):
and self._task_kind(path) == "text"
and not self._computed_off_this_gpu(engine, model, path)):
_qkey = self._queue_key(model)
try:
await self.reqqueue.acquire(
......@@ -1095,12 +1097,65 @@ class FrontProxy:
except Exception:
return None
def _swap_owner_key(self, engine, model: Optional[str]) -> str:
def _swap_owner_key(self, engine, model: Optional[str], path: Optional[str] = None) -> str:
"""The model identity that determines GPU residency for the swap gate. Same
model → same owner (runs free); different model → a swap. Falls back to the
engine name for inference without an explicit model."""
# All OCR shares ONE owner: the OCR subsystem holds one set of engine pools and
# does its own internal eviction between them, so for the purpose of "who owns
# the card" paddle and surya are the same tenant. (It also means the front need
# not read the `engine` form field out of a multipart body carrying a PDF.)
if path is not None and _router.is_ocr_path(path):
return "ocr"
return self._queue_key(model) or getattr(engine, "name", "") or "?"
def _computed_off_this_gpu(self, engine, model: Optional[str], path: str) -> bool:
"""True when this request will NOT be computed on this box's GPU — a cluster
node, a RunPod pod, or an external API endpoint does the work.
Such a request holds no local VRAM, so making it wait at the swap gate (or in
the local queue) would add latency and swaps for nothing."""
try:
if getattr(engine, "remote", False):
return True # a cluster node: its cards are its own business
if (self._model_info(model).get("backend") or "").lower() == "runpod":
return True # a rented GPU elsewhere
if _router.is_ocr_path(path) and self._ocr_is_external():
return True
except Exception:
pass
return False
def _ocr_is_external(self) -> bool:
"""True when no ENABLED OCR engine would run on this box's GPU — every one of
them is served by an endpoint somewhere else.
Decided across all enabled engines rather than the one this request asked for:
the engine arrives as a multipart form field, and buffering a 100 MB PDF in the
front just to read it would cost more than the gate saves. Erring towards
gating is the safe direction — a needless turn costs latency, a missed turn
costs VRAM contention."""
try:
o = self.config.ocr
except Exception:
return False
local = []
if getattr(o, "paddle_enabled", False):
local.append(True) # in-process / isolated venv: always this box
if getattr(o, "doctr_enabled", False):
local.append(True)
if getattr(o, "surya_enabled", False) and getattr(o, "surya_accept_license", False):
mode = (getattr(o, "surya_serve", "local") or "local").strip().lower()
external = mode == "llamacpp" and bool(
(getattr(o, "surya_server_url", "") or "").strip())
local.append(not external)
if getattr(o, "olmocr_enabled", False):
mode = (getattr(o, "olmocr_serve", "model") or "model").strip().lower()
external = mode == "server" and bool(
(getattr(o, "olmocr_server_url", "") or "").strip())
local.append(not external)
return bool(local) and not any(local)
def _resident_on(self, engine, key: str) -> bool:
"""True when the model behind `key` is already loaded on `engine` (per
its last health poll), so serving it needs NO swap/eviction."""
......@@ -1183,13 +1238,17 @@ class FrontProxy:
async def _swap_acquire(self, engine, model, path, method):
"""Acquire this engine's shared-GPU swap slot for a GPU-inference request.
Returns a (gate, key) token for _swap_release, or None when no gate applies
(single-card engine or non-inference request)."""
if str(method).upper() != "POST" or not _router.is_inference_path(path):
(single-card engine, non-inference request, or work computed off this box)."""
if str(method).upper() != "POST" or not _router.is_gpu_inference_path(path):
return None
# RunPod, an external API or a cluster node computes it elsewhere: no local VRAM
# is involved, so there is nothing to serialize.
if self._computed_off_this_gpu(engine, model, path):
return None
gate = self._swap_gate_for(engine)
if gate is None:
return None
key = self._swap_owner_key(engine, model)
key = self._swap_owner_key(engine, model, path)
# A model that is ALREADY resident on this engine needs no swap — its
# own per-model queue governs its concurrency, and serializing it
# behind the current GPU owner would block cross-model parallelism
......@@ -1953,9 +2012,10 @@ class FrontProxy:
try:
# 0. Shared-GPU swap gate: if a different model owns the card, wait
# for the swap (keepalive so the client doesn't time out).
_gate = self._swap_gate_for(engine)
_gate = (None if self._computed_off_this_gpu(engine, model, path)
else self._swap_gate_for(engine))
if _gate is not None:
_skey = self._swap_owner_key(engine, model)
_skey = self._swap_owner_key(engine, model, path)
_swap_acq = _asyncio.ensure_future(_gate.acquire(_skey))
while True:
try:
......@@ -1966,7 +2026,7 @@ class FrontProxy:
except _asyncio.TimeoutError:
yield _ka("waiting for GPU (another model is finishing)")
# 1. Front per-model queue slot (text only) — keepalive while waiting.
if is_text:
if is_text and not self._computed_off_this_gpu(engine, model, path):
_qkey = self._queue_key(model)
_slot_held = False
# `on_busy`: if there is no free local slot RIGHT NOW, burst
......@@ -2399,7 +2459,8 @@ class FrontProxy:
# audio, embeddings…) pass through unqueued.
_qkey = None
if (method == "POST" and _router.is_inference_path(path)
and self._task_kind(path) == "text"):
and self._task_kind(path) == "text"
and not self._computed_off_this_gpu(engine, model, path)):
_qkey = self._queue_key(model)
_slot_held = False
# `on_busy`: offload the moment there is no free local slot, rather
......
......@@ -48,6 +48,27 @@ def is_inference_path(path: str) -> bool:
return p in _INFERENCE_PATHS
# OCR is GPU work and contends for VRAM like any model, but it is deliberately NOT in
# _INFERENCE_PATHS: it has no `model` field, it owns its own engine pools, and the
# front's per-model queue / pin / spillover / keepalive machinery does not apply to it.
# These are the paths that still have to take a turn at the shared-GPU swap gate.
_OCR_GPU_PATHS = {
"/v1/ocr",
"/v1/ocr/batch",
}
def is_ocr_path(path: str) -> bool:
return path.split("?", 1)[0].rstrip("/") in _OCR_GPU_PATHS
def is_gpu_inference_path(path: str) -> bool:
"""POSTs that consume THIS box's GPU, and so must be serialized against whatever
model currently owns the card. A superset of :func:`is_inference_path` that adds OCR."""
p = path.split("?", 1)[0].rstrip("/")
return p in _INFERENCE_PATHS or p in _OCR_GPU_PATHS
def is_admin_path(path: str) -> bool:
p = path.split("?", 1)[0]
return (p.startswith("/admin") or p.startswith("/login") or p.startswith("/logout")
......
"""OCR takes its turn at the shared-GPU swap gate like any other GPU request — unless
the work is computed somewhere else (a cluster node, a RunPod pod, an external API).
Until now /v1/ocr bypassed the gate entirely: it is not in the router's _INFERENCE_PATHS
(it has no `model` field and none of the front's per-model queue/pin machinery applies),
so an OCR request could start a surya vLLM boot while an LLM was mid-generation on the
same card, and vice versa. The gate is what makes them alternate instead of contend.
The methods are exercised against a stub `self`: constructing a real FrontProxy would
need engines, config files and sockets, and none of the logic under test touches them.
"""
import asyncio
import sys
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.frontproxy import router as _router
from codai.frontproxy.app import FrontProxy
from codai.frontproxy.reqqueue import GpuSwapGate
# ------------------------------------------------------------------ paths
def test_ocr_is_gpu_work_but_not_a_model_inference_path():
for p in ("/v1/ocr", "/v1/ocr/batch", "/v1/ocr/"):
assert _router.is_gpu_inference_path(p), p
assert _router.is_ocr_path(p), p
# Deliberately NOT an inference path: that predicate drives the per-model
# queue, pins, spillover and keepalive, none of which fit OCR.
assert not _router.is_inference_path(p), p
def test_the_model_paths_are_still_gpu_inference_paths():
for p in ("/v1/chat/completions", "/v1/embeddings", "/v1/images/generations",
"/v1/video/generations", "/v1/audio/transcriptions"):
assert _router.is_inference_path(p) and _router.is_gpu_inference_path(p), p
assert not _router.is_ocr_path(p), p
def test_schema_endpoints_are_not_gpu_work():
for p in ("/v1/ocr/schemas", "/v1/ocr/schemas/foo", "/v1/models", "/admin"):
assert not _router.is_ocr_path(p), p
assert not _router.is_gpu_inference_path(p), p
# ------------------------------------------------------------------ owner key
def _stub(ocr=None, model_info=None):
s = types.SimpleNamespace()
s.config = types.SimpleNamespace(ocr=ocr if ocr is not None else OcrConfig())
s._model_info = lambda m: (model_info or {})
s._queue_key = lambda m: (m or "").lower()
s._ocr_is_external = lambda: FrontProxy._ocr_is_external(s)
return s
def test_all_ocr_requests_share_one_owner_on_the_card():
"""paddle and surya are one tenant as far as "who owns the GPU" goes — the OCR
manager does its own eviction between its pools."""
s = _stub()
eng = types.SimpleNamespace(name="nvidia", remote=False)
assert FrontProxy._swap_owner_key(s, eng, None, "/v1/ocr") == "ocr"
assert FrontProxy._swap_owner_key(s, eng, "surya", "/v1/ocr/batch") == "ocr"
# A model request still keys on the model, as before.
assert FrontProxy._swap_owner_key(s, eng, "Qwen/Qwen3.5-9B", "/v1/chat/completions") \
== "qwen/qwen3.5-9b"
# No path given (older call site) → unchanged behaviour.
assert FrontProxy._swap_owner_key(s, eng, "lisa", None) == "lisa"
# ------------------------------------------------------------------ off-box work
def test_a_cluster_node_computes_it_on_its_own_card():
s = _stub()
node = types.SimpleNamespace(name="node-b", remote=True)
assert FrontProxy._computed_off_this_gpu(s, node, "lisa", "/v1/chat/completions")
assert FrontProxy._computed_off_this_gpu(s, node, None, "/v1/ocr")
def test_a_runpod_model_uses_no_local_vram():
s = _stub(model_info={"backend": "runpod"})
eng = types.SimpleNamespace(name="nvidia", remote=False)
assert FrontProxy._computed_off_this_gpu(s, eng, "big-model", "/v1/chat/completions")
def test_local_work_is_gated():
s = _stub(model_info={"backend": "nvidia"})
eng = types.SimpleNamespace(name="nvidia", remote=False)
assert not FrontProxy._computed_off_this_gpu(s, eng, "lisa", "/v1/chat/completions")
assert not FrontProxy._computed_off_this_gpu(s, eng, None, "/v1/ocr")
def test_ocr_served_entirely_by_external_endpoints_is_not_gated():
"""Every enabled engine points off-box, so no OCR request can touch this card."""
o = OcrConfig(
enabled=True, paddle_enabled=False, doctr_enabled=False,
surya_enabled=True, surya_accept_license=True,
surya_serve="llamacpp", surya_server_url="http://other-box:8080/v1",
olmocr_enabled=True, olmocr_serve="server",
olmocr_server_url="https://api.example/v1")
s = _stub(ocr=o)
eng = types.SimpleNamespace(name="nvidia", remote=False)
assert FrontProxy._ocr_is_external(s)
assert FrontProxy._computed_off_this_gpu(s, eng, None, "/v1/ocr")
def test_one_local_engine_is_enough_to_keep_ocr_gated():
o = OcrConfig(
enabled=True, paddle_enabled=True, # local
surya_enabled=True, surya_accept_license=True,
surya_serve="llamacpp", surya_server_url="http://other-box:8080/v1")
s = _stub(ocr=o)
eng = types.SimpleNamespace(name="nvidia", remote=False)
assert not FrontProxy._ocr_is_external(s)
assert not FrontProxy._computed_off_this_gpu(s, eng, None, "/v1/ocr")
def test_surya_on_the_local_vllm_is_local_work():
"""vllm mode boots an instance ON THIS CARD — the case the whole fix is about."""
o = OcrConfig(enabled=True, paddle_enabled=False, doctr_enabled=False,
surya_enabled=True, surya_accept_license=True, surya_serve="vllm")
s = _stub(ocr=o)
assert not FrontProxy._ocr_is_external(s)
def test_olmocr_in_model_mode_is_local_work():
o = OcrConfig(enabled=True, paddle_enabled=False, doctr_enabled=False,
olmocr_enabled=True, olmocr_serve="model")
s = _stub(ocr=o)
assert not FrontProxy._ocr_is_external(s)
def test_no_enabled_engine_at_all_is_not_claimed_external():
s = _stub(ocr=OcrConfig(enabled=True, paddle_enabled=False, doctr_enabled=False))
assert not FrontProxy._ocr_is_external(s)
def test_a_license_gated_surya_does_not_count_as_a_local_engine():
o = OcrConfig(enabled=True, paddle_enabled=False, doctr_enabled=False,
surya_enabled=True, surya_accept_license=False, # not usable
olmocr_enabled=True, olmocr_serve="server",
olmocr_server_url="https://api.example/v1")
s = _stub(ocr=o)
assert FrontProxy._ocr_is_external(s)
# ------------------------------------------------------------------ alternation
def test_ocr_and_a_model_alternate_on_the_card():
"""The behaviour the gate buys: OCR waits for the model's batch, takes the card,
and the model gets it back — turns, not contention."""
gate = GpuSwapGate(cap=2)
order = []
async def _run(key, n):
for _ in range(n):
await gate.acquire(key)
order.append(key)
await asyncio.sleep(0)
gate.release(key)
async def _go():
await asyncio.gather(_run("lisa", 4), _run("ocr", 4))
asyncio.run(_go())
assert len(order) == 8
assert set(order) == {"lisa", "ocr"}
# Neither side was starved: both finished all four turns, and the owner yielded
# rather than running all four of its own first.
assert order.count("lisa") == 4 and order.count("ocr") == 4
assert order[0] != order[-1] or len(set(order[:4])) == 2
def test_the_owner_keeps_the_card_when_nobody_else_wants_it():
gate = GpuSwapGate(cap=2)
async def _go():
for _ in range(5): # past the cap, with no other waiter
await gate.acquire("ocr")
gate.release("ocr")
return True
assert asyncio.run(_go())
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