frontproxy: intelligent per-shared-GPU model-swap queue (batch, then swap)

Builds on the cross-engine clean-swap eviction: instead of two engines on one
shared card ever running forwards concurrently (→ VRAM contention → OOM →
disk-thrash), the front now serializes model OWNERSHIP of a shared GPU while
batching to avoid per-request thrash.

New GpuSwapGate (frontproxy/reqqueue.py), one per shared-GPU group (keyed by the
co-located engines' CODERAI_ENGINE_GPUS selector, created only when an engine has
a sibling on its card):

  * A request for the model that currently OWNS the GPU runs immediately — a swap
    isn't needed (a lone stream never stalls). Concurrency stays capped downstream
    by the existing per-model FrontQueue.
  * A request for a DIFFERENT model queues. The owner keeps being served up to
    `cap` requests (server.gpu_swap_batch, default 10) while another model waits,
    then — once the owner is fully idle (never mid-request) — the GPU SWAPS to the
    waiting model (which evicts + loads), serves it, and round-robins BACK if the
    original has requests queued. No thrash (batch), no starvation (cap).

Wired into all four dispatch paths (broker, broker-stream, direct stream with
keepalive, direct non-stream) for every GPU-inference kind (text/image/video):
acquire the swap slot before the per-model queue, release in the finalizer;
cancelling a pending acquire (client disconnect) drops the waiter with no leak.
The text-stream path emits keepalives while waiting out a swap so the client
doesn't time out.

Scheduler validated by async unit tests: cap engages at exactly N with a
competitor waiting; a lone same-model stream runs unbounded; round-robin
alternates; cancelled waiters leak no slot.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RdMufYvtTbtGDWsiZVoXce

codai/config.py

@@ -52,6 +52,11 @@ class ServerConfig:
                                         # generous enough that a GIL-busy engine's
                                         # health poll doesn't time out mid-generation
     proxy_max_inflight: int = 64        # max concurrent proxied requests through the front
+    gpu_swap_batch: int = 10            # on a shared GPU (GGUF-isolation split), serve up to
+                                        # this many requests for the model that currently owns
+                                        # the card before swapping to a queued different model
+                                        # (then round-robin back). Prevents cross-engine VRAM
+                                        # contention while avoiding per-request model thrash.
     engine_restart_drain_grace: float = 30.0  # on engine restart, wait this many seconds
                                               # for in-flight requests to finish before
                                               # killing the process (0 = bounce immediately)
@@ -659,6 +664,7 @@ class ConfigManager:
                 "engine_gpus": self.config.server.engine_gpus,
                 "proxy_status_timeout": self.config.server.proxy_status_timeout,
                 "proxy_max_inflight": self.config.server.proxy_max_inflight,
+                "gpu_swap_batch": self.config.server.gpu_swap_batch,
                 "engine_restart_drain_grace": self.config.server.engine_restart_drain_grace,
                 "isolate_gguf_engine": self.config.server.isolate_gguf_engine,
                 "engine_specs": self.config.server.engine_specs,

codai/frontproxy/app.py

@@ -92,6 +92,11 @@ class FrontProxy:
         # Front-managed generation queue: admission control + ordering + queue
         # position, sized per-model to max_instances so the engine never queues.
         self.reqqueue = FrontQueue()
+        # Per-shared-GPU model-swap scheduler (GGUF-isolation split): serialize which
+        # model owns a shared card so two forwards never contend for VRAM, batching
+        # same-model requests before swapping. Keyed by the engines' shared-GPU
+        # selector; created lazily for engines that actually have a co-located sibling.
+        self._swap_gates = {}
         # Recent inference activity (front-tracked, since the front relays every
         # request) so the Overview dashboard's activity table is served natively
         # without asking the engine. Newest first; bounded.
@@ -322,6 +327,12 @@ class FrontProxy:
         # Signed with the internal token; the engine accepts it only if it matches.
         if self.internal_token:
             send_headers["x-coderai-broker-authed"] = self.internal_token
+        # Shared-GPU swap gate (all inference kinds): wait out any in-flight swap on
+        # a shared card so this request doesn't contend for VRAM.
+        try:
+            _swap_tok = await self._swap_acquire(engine, model, path, method)
+        except Exception:
+            _swap_tok = None
         # Front-managed generation queue (text only) — same per-model gate as the
         # direct proxy path, so brokered and direct requests share one queue.
         _qkey = None
@@ -334,6 +345,7 @@ class FrontProxy:
                     rid=engine.name + ":" + (model or ""), model=model or "",
                     engine=engine.name)
             except QueueFull:
+                await self._swap_release(_swap_tok)
                 return {"status_code": 503,
                         "headers": {"content-type": "application/json"},
                         "body": b'{"error":"Server busy: the generation queue is '
@@ -379,6 +391,7 @@ class FrontProxy:
             engine.exit_request(_rid)
             if _qkey is not None:
                 await self.reqqueue.release(_qkey)
+            await self._swap_release(_swap_tok)
             if _router.is_inference_path(path):
                 self._record_activity(model, self._task_kind(path), _status, _started)
         # Surface the engine's actual reply so a brokered request that "doesn't get
@@ -460,6 +473,10 @@ class FrontProxy:
                         if k.lower() not in _DROP_REQ}
         if self.internal_token:
             send_headers["x-coderai-broker-authed"] = self.internal_token
+        try:
+            _swap_tok = await self._swap_acquire(engine, model, path, method)
+        except Exception:
+            _swap_tok = None
         _qkey = None
         if (method.upper() == "POST" and _is_infer
                 and self._task_kind(path) == "text"):
@@ -470,6 +487,7 @@ class FrontProxy:
                     rid=engine.name + ":" + (model or ""), model=model or "",
                     engine=engine.name)
             except QueueFull:
+                await self._swap_release(_swap_tok)
                 yield ('data: {"error":"Server busy: the generation queue is full, '
                        'please retry shortly."}\n\n')
                 return
@@ -547,6 +565,7 @@ class FrontProxy:
             engine.exit_request(_rid)
             if _qkey is not None:
                 await self.reqqueue.release(_qkey)
+            await self._swap_release(_swap_tok)
             if _is_infer:
                 self._record_activity(model, self._task_kind(path), _status, _started)
 
@@ -642,6 +661,60 @@ class FrontProxy:
         info = self._model_info(model)
         return (info.get("model_id") or model or "").lower()
 
+    def _swap_cap(self) -> int:
+        try:
+            return int(getattr(self.config.server, "gpu_swap_batch", 10) or 10)
+        except Exception:
+            return 10
+
+    def _swap_gate_for(self, engine):
+        """Return the shared-GPU model-swap gate for `engine`, or None when it isn't
+        on a shared card (nothing to serialize). Co-location is the same signal the
+        supervisor uses for cross-engine VRAM release: an engine with a co-located
+        sibling carries CODERAI_COSITED_URLS; engines sharing a card have the same
+        CODERAI_ENGINE_GPUS selector, which keys the gate so both share one."""
+        try:
+            if engine is None or getattr(engine, "role", "engine") == "system":
+                return None
+            env = getattr(engine, "env", None) or {}
+            if not env.get("CODERAI_COSITED_URLS"):
+                return None  # no sibling on this card → no cross-engine contention
+            gkey = env.get("CODERAI_ENGINE_GPUS") or getattr(engine, "url", "") or "shared"
+            gate = self._swap_gates.get(gkey)
+            if gate is None:
+                from codai.frontproxy.reqqueue import GpuSwapGate
+                gate = GpuSwapGate(cap=self._swap_cap())
+                self._swap_gates[gkey] = gate
+            return gate
+        except Exception:
+            return None
+
+    def _swap_owner_key(self, engine, model: Optional[str]) -> 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."""
+        return self._queue_key(model) or getattr(engine, "name", "") or "?"
+
+    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):
+            return None
+        gate = self._swap_gate_for(engine)
+        if gate is None:
+            return None
+        key = self._swap_owner_key(engine, model)
+        await gate.acquire(key)
+        return (gate, key)
+
+    async def _swap_release(self, token) -> None:
+        if token is not None:
+            try:
+                await token[0].release(token[1])
+            except Exception:
+                pass
+
     def _model_capacity(self, model: Optional[str]) -> int:
         """Per-model concurrency = its max_instances, falling back to the global
         server default. This is the number of front queue slots for the model."""
@@ -1270,7 +1343,23 @@ class FrontProxy:
             _status = 502
             _started = _t.time()
             rp_resp = None
+            _swap_tok = None
+            _swap_acq = None
             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)
+                if _gate is not None:
+                    _skey = self._swap_owner_key(engine, model)
+                    _swap_acq = _asyncio.ensure_future(_gate.acquire(_skey))
+                    while True:
+                        try:
+                            await _asyncio.wait_for(_asyncio.shield(_swap_acq),
+                                                    timeout=_KA)
+                            _swap_tok = (_gate, _skey)
+                            break
+                        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:
                     _qkey = self._queue_key(model)
@@ -1355,6 +1444,10 @@ class FrontProxy:
                         await self.reqqueue.release(_qkey)
                     except Exception:
                         pass
+                # Release / cancel the shared-GPU swap slot.
+                if _swap_acq is not None and not _swap_acq.done():
+                    _swap_acq.cancel()
+                await self._swap_release(_swap_tok)
                 if _router.is_inference_path(path):
                     self._record_activity(model, self._task_kind(path), _status, _started)
 
@@ -1671,6 +1764,13 @@ class FrontProxy:
         headers = self._filter_headers(request.headers, _DROP_REQ)
         content = body_bytes if body_bytes is not None else request.stream()
 
+        # Shared-GPU swap gate (all inference kinds, incl. image/video): wait for the
+        # card if a different model currently owns it, so this forward never contends.
+        try:
+            _swap_tok = await self._swap_acquire(engine, model, path, method)
+        except Exception:
+            _swap_tok = None
+
         rp_req = self._long.build_request(
             method, url, headers=headers, params=request.query_params,
             content=content)
@@ -1689,6 +1789,7 @@ class FrontProxy:
             engine.exit_request(_rid)
             if _qkey is not None:
                 await self.reqqueue.release(_qkey)
+            await self._swap_release(_swap_tok)
             return JSONResponse(
                 {"error": f"Engine#{engine.id} unreachable: {exc}"}, status_code=502)
 
@@ -1699,6 +1800,7 @@ class FrontProxy:
                 engine.exit_request(_rid)
                 if _qkey is not None:
                     await self.reqqueue.release(_qkey)
+                await self._swap_release(_swap_tok)
                 if _meta is not None:
                     self._record_activity(model, self._task_kind(path),
                                           rp_resp.status_code, _started)

codai/frontproxy/reqqueue.py

@@ -121,3 +121,134 @@ class FrontQueue:
                 out.append({"rid": w.rid, "model": w.model, "engine": w.engine,
                             "position": i + 1, "enqueued_at": w.enqueued_at})
         return out
+
+
+class _SwapWaiter:
+    __slots__ = ("key", "event", "granted", "enqueued_at")
+
+    def __init__(self, key):
+        import time as _t
+        self.key = key
+        self.event = asyncio.Event()
+        self.granted = False
+        self.enqueued_at = _t.time()
+
+
+class GpuSwapGate:
+    """Serialize model 'ownership' of one shared GPU across co-located engines.
+
+    On the GGUF-isolation split a torch (image/video) engine and a gguf (text)
+    engine share a single NVIDIA card and cannot hold both big models at once. This
+    gate makes at most ONE model own the GPU at a time — so two model forwards never
+    run concurrently and contend for VRAM (the OOM-then-disk-thrash failure) — while
+    keeping it efficient:
+
+      * Requests for the model that currently OWNS the GPU run immediately (a swap
+        isn't needed), concurrency still capped downstream by the per-model queue.
+      * A request for a DIFFERENT model queues. The owner keeps being served — up to
+        `cap` requests while another model is waiting — then, once the owner is fully
+        idle (its in-flight requests finished; we never swap mid-request), the GPU
+        SWAPS to the waiting model (which evicts + loads), serves it, and later swaps
+        BACK if the original model has requests queued. Round-robin with a per-turn
+        batch cap: no thrash (a lone request doesn't force a swap) and no starvation
+        (a busy model yields after `cap`).
+
+    `acquire(key)`/`release(key)` bracket each GPU-inference request; `key` is the
+    model identity that determines residency. Cancelling a pending acquire (client
+    disconnect) drops the waiter with no slot leak."""
+
+    def __init__(self, cap: int = 10):
+        self._lock = asyncio.Lock()
+        self.cap = max(1, int(cap))
+        self._owner = None      # model key currently allowed to run GPU work
+        self._running = 0       # in-flight granted requests (all for _owner)
+        self._served = 0        # grants since _owner took the GPU (batch-cap counter)
+        self._waiters = []      # FIFO list[_SwapWaiter]
+
+    def _other_waiting(self) -> bool:
+        return any(w.key != self._owner for w in self._waiters)
+
+    async def acquire(self, key) -> None:
+        async with self._lock:
+            if self._owner is None:
+                self._owner = key
+                self._served = 0
+            # Fast path: the resident model, unless its batch cap is spent AND a
+            # different model is waiting (then it must yield — fall through to queue).
+            if key == self._owner and (self._served < self.cap
+                                       or not self._other_waiting()):
+                self._running += 1
+                self._served += 1
+                return
+            w = _SwapWaiter(key)
+            self._waiters.append(w)
+            # If the GPU is idle right now nothing will pump this waiter later, so
+            # process it immediately (may swap the owner to `key`).
+            if self._running == 0:
+                self._pump()
+        try:
+            await w.event.wait()
+        except BaseException:
+            # Cancelled/errored while queued OR just after being granted.
+            async with self._lock:
+                if w in self._waiters:
+                    self._waiters.remove(w)             # never held a slot
+                elif w.granted:
+                    self._running -= 1                  # granted as we cancelled
+                    if self._running <= 0:
+                        self._running = 0
+                        self._pump()
+            raise
+
+    async def release(self, key) -> None:
+        async with self._lock:
+            self._running -= 1
+            if self._running <= 0:
+                self._running = 0
+                self._pump()
+
+    def _grant(self, w: "_SwapWaiter") -> None:
+        self._waiters.remove(w)
+        self._running += 1
+        self._served += 1
+        w.granted = True
+        w.event.set()
+
+    def _swap_to(self, key) -> None:
+        self._owner = key
+        self._served = 0
+        for w in [x for x in self._waiters if x.key == key]:
+            self._grant(w)
+
+    def _pump(self) -> None:
+        """Owner is idle (_running == 0): decide who runs next. Called under lock."""
+        if not self._waiters:
+            return  # keep _owner as the last-resident model so a repeat runs free
+        owner_w = [w for w in self._waiters if w.key == self._owner]
+        other_w = [w for w in self._waiters if w.key != self._owner]
+        if owner_w and self._served < self.cap:
+            # Keep serving the resident model. When another model is waiting, grant
+            # only up to the remaining cap so it eventually yields; otherwise all.
+            room = (self.cap - self._served) if other_w else len(owner_w)
+            granted = 0
+            for w in list(owner_w):
+                if granted >= room:
+                    break
+                self._grant(w)
+                granted += 1
+            if granted == 0 and other_w:
+                self._swap_to(other_w[0].key)   # cap spent → swap
+            return
+        if other_w:
+            self._swap_to(other_w[0].key)       # cap spent or owner drained → swap
+            return
+        # Only the owner is waiting (cap spent, nobody else): reset the turn.
+        self._served = 0
+        for w in list(owner_w):
+            self._grant(w)
+
+    def snapshot(self) -> dict:
+        return {"owner": self._owner, "running": self._running,
+                "served": self._served, "cap": self.cap,
+                "waiting": [{"key": w.key, "enqueued_at": w.enqueued_at}
+                            for w in self._waiters]}