frontproxy: fix leaked GPU-swap slot on request cancellation (queued swap never fired)

The GpuSwapGate.release() was async and awaited in the dispatch `finally`
blocks. When a request was cancelled/interrupted mid-flight (client disconnect,
an interrupted text generation), `await self._swap_release(...)` in the finally
could itself be cancelled BEFORE it decremented the running counter — stranding
a gate slot. With `running` stuck > 0, `_pump()` never ran, so a video request
queued behind the interrupted text request was never granted: the GPU never
swapped even though the text engine had gone idle (observed: video stuck ~37min
while the text engine was idle for 11).

Fix: make release()/_pump() SYNCHRONOUS and drop the asyncio.Lock. Every critical
section is straight-line (no await between read and write), so under asyncio's
single thread they're already atomic — and a synchronous release from a `finally`
always completes even while the coroutine is being cancelled. acquire() keeps its
one `await` (the event wait) with synchronous cancel-cleanup. All release call
sites are now non-awaited. Added `[gpu-swap] queued/swapping` logging so the
owner/queue/swap transitions are visible in debug.log.

Validated: cancelling a text request whose slot a queued video is waiting behind
now frees the slot and grants the video; a cancelled queued waiter leaks nothing.
Co-Authored-By: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RdMufYvtTbtGDWsiZVoXce
parent c0a970b0
...@@ -345,7 +345,7 @@ class FrontProxy: ...@@ -345,7 +345,7 @@ class FrontProxy:
rid=engine.name + ":" + (model or ""), model=model or "", rid=engine.name + ":" + (model or ""), model=model or "",
engine=engine.name) engine=engine.name)
except QueueFull: except QueueFull:
await self._swap_release(_swap_tok) self._swap_release(_swap_tok)
return {"status_code": 503, return {"status_code": 503,
"headers": {"content-type": "application/json"}, "headers": {"content-type": "application/json"},
"body": b'{"error":"Server busy: the generation queue is ' "body": b'{"error":"Server busy: the generation queue is '
...@@ -391,7 +391,7 @@ class FrontProxy: ...@@ -391,7 +391,7 @@ class FrontProxy:
engine.exit_request(_rid) engine.exit_request(_rid)
if _qkey is not None: if _qkey is not None:
await self.reqqueue.release(_qkey) await self.reqqueue.release(_qkey)
await self._swap_release(_swap_tok) self._swap_release(_swap_tok)
if _router.is_inference_path(path): if _router.is_inference_path(path):
self._record_activity(model, self._task_kind(path), _status, _started) self._record_activity(model, self._task_kind(path), _status, _started)
# Surface the engine's actual reply so a brokered request that "doesn't get # Surface the engine's actual reply so a brokered request that "doesn't get
...@@ -487,7 +487,7 @@ class FrontProxy: ...@@ -487,7 +487,7 @@ class FrontProxy:
rid=engine.name + ":" + (model or ""), model=model or "", rid=engine.name + ":" + (model or ""), model=model or "",
engine=engine.name) engine=engine.name)
except QueueFull: except QueueFull:
await self._swap_release(_swap_tok) self._swap_release(_swap_tok)
yield ('data: {"error":"Server busy: the generation queue is full, ' yield ('data: {"error":"Server busy: the generation queue is full, '
'please retry shortly."}\n\n') 'please retry shortly."}\n\n')
return return
...@@ -565,7 +565,7 @@ class FrontProxy: ...@@ -565,7 +565,7 @@ class FrontProxy:
engine.exit_request(_rid) engine.exit_request(_rid)
if _qkey is not None: if _qkey is not None:
await self.reqqueue.release(_qkey) await self.reqqueue.release(_qkey)
await self._swap_release(_swap_tok) self._swap_release(_swap_tok)
if _is_infer: if _is_infer:
self._record_activity(model, self._task_kind(path), _status, _started) self._record_activity(model, self._task_kind(path), _status, _started)
...@@ -683,7 +683,9 @@ class FrontProxy: ...@@ -683,7 +683,9 @@ class FrontProxy:
gate = self._swap_gates.get(gkey) gate = self._swap_gates.get(gkey)
if gate is None: if gate is None:
from codai.frontproxy.reqqueue import GpuSwapGate from codai.frontproxy.reqqueue import GpuSwapGate
gate = GpuSwapGate(cap=self._swap_cap()) gate = GpuSwapGate(
cap=self._swap_cap(),
log=lambda m: print("[front] " + m, flush=True))
self._swap_gates[gkey] = gate self._swap_gates[gkey] = gate
return gate return gate
except Exception: except Exception:
...@@ -708,10 +710,14 @@ class FrontProxy: ...@@ -708,10 +710,14 @@ class FrontProxy:
await gate.acquire(key) await gate.acquire(key)
return (gate, key) return (gate, key)
async def _swap_release(self, token) -> None: def _swap_release(self, token) -> None:
"""Synchronous on purpose: called from `finally` blocks that may run while
the request coroutine is being cancelled — a sync release always completes
and frees the slot, where an awaited one could be cancelled mid-way and
strand the slot (blocking every queued swap behind it)."""
if token is not None: if token is not None:
try: try:
await token[0].release(token[1]) token[0].release(token[1])
except Exception: except Exception:
pass pass
...@@ -1447,7 +1453,7 @@ class FrontProxy: ...@@ -1447,7 +1453,7 @@ class FrontProxy:
# Release / cancel the shared-GPU swap slot. # Release / cancel the shared-GPU swap slot.
if _swap_acq is not None and not _swap_acq.done(): if _swap_acq is not None and not _swap_acq.done():
_swap_acq.cancel() _swap_acq.cancel()
await self._swap_release(_swap_tok) self._swap_release(_swap_tok)
if _router.is_inference_path(path): if _router.is_inference_path(path):
self._record_activity(model, self._task_kind(path), _status, _started) self._record_activity(model, self._task_kind(path), _status, _started)
...@@ -1789,7 +1795,7 @@ class FrontProxy: ...@@ -1789,7 +1795,7 @@ class FrontProxy:
engine.exit_request(_rid) engine.exit_request(_rid)
if _qkey is not None: if _qkey is not None:
await self.reqqueue.release(_qkey) await self.reqqueue.release(_qkey)
await self._swap_release(_swap_tok) self._swap_release(_swap_tok)
return JSONResponse( return JSONResponse(
{"error": f"Engine#{engine.id} unreachable: {exc}"}, status_code=502) {"error": f"Engine#{engine.id} unreachable: {exc}"}, status_code=502)
...@@ -1800,7 +1806,7 @@ class FrontProxy: ...@@ -1800,7 +1806,7 @@ class FrontProxy:
engine.exit_request(_rid) engine.exit_request(_rid)
if _qkey is not None: if _qkey is not None:
await self.reqqueue.release(_qkey) await self.reqqueue.release(_qkey)
await self._swap_release(_swap_tok) self._swap_release(_swap_tok)
if _meta is not None: if _meta is not None:
self._record_activity(model, self._task_kind(path), self._record_activity(model, self._task_kind(path),
rp_resp.status_code, _started) rp_resp.status_code, _started)
......
...@@ -157,19 +157,25 @@ class GpuSwapGate: ...@@ -157,19 +157,25 @@ class GpuSwapGate:
model identity that determines residency. Cancelling a pending acquire (client model identity that determines residency. Cancelling a pending acquire (client
disconnect) drops the waiter with no slot leak.""" disconnect) drops the waiter with no slot leak."""
def __init__(self, cap: int = 10): def __init__(self, cap: int = 10, log=None):
self._lock = asyncio.Lock() # NO asyncio.Lock on purpose: every critical section below is straight-line
# (no `await` between reading and writing state), so under asyncio's single
# thread they are already atomic. Crucially this makes release()/_pump()
# SYNCHRONOUS and therefore uncancellable — a request cancelled mid-flight
# (client disconnect / interrupt) still frees its slot from a `finally`,
# instead of an `await release()` being cancelled before it decrements the
# counter and permanently stranding a slot (→ a queued swap never fires).
self.cap = max(1, int(cap)) self.cap = max(1, int(cap))
self._owner = None # model key currently allowed to run GPU work self._owner = None # model key currently allowed to run GPU work
self._running = 0 # in-flight granted requests (all for _owner) self._running = 0 # in-flight granted requests (all for _owner)
self._served = 0 # grants since _owner took the GPU (batch-cap counter) self._served = 0 # grants since _owner took the GPU (batch-cap counter)
self._waiters = [] # FIFO list[_SwapWaiter] self._waiters = [] # FIFO list[_SwapWaiter]
self._log = log
def _other_waiting(self) -> bool: def _other_waiting(self) -> bool:
return any(w.key != self._owner for w in self._waiters) return any(w.key != self._owner for w in self._waiters)
async def acquire(self, key) -> None: async def acquire(self, key) -> None:
async with self._lock:
if self._owner is None: if self._owner is None:
self._owner = key self._owner = key
self._served = 0 self._served = 0
...@@ -182,6 +188,13 @@ class GpuSwapGate: ...@@ -182,6 +188,13 @@ class GpuSwapGate:
return return
w = _SwapWaiter(key) w = _SwapWaiter(key)
self._waiters.append(w) self._waiters.append(w)
if self._log:
try:
self._log(f"[gpu-swap] queued {key!r} behind owner {self._owner!r} "
f"(running={self._running}, served={self._served}/{self.cap}, "
f"waiting={len(self._waiters)})")
except Exception:
pass
# If the GPU is idle right now nothing will pump this waiter later, so # If the GPU is idle right now nothing will pump this waiter later, so
# process it immediately (may swap the owner to `key`). # process it immediately (may swap the owner to `key`).
if self._running == 0: if self._running == 0:
...@@ -189,8 +202,8 @@ class GpuSwapGate: ...@@ -189,8 +202,8 @@ class GpuSwapGate:
try: try:
await w.event.wait() await w.event.wait()
except BaseException: except BaseException:
# Cancelled/errored while queued OR just after being granted. # Cancelled/errored while queued OR just after being granted. Synchronous
async with self._lock: # cleanup (no await) so it always runs to completion even under cancel.
if w in self._waiters: if w in self._waiters:
self._waiters.remove(w) # never held a slot self._waiters.remove(w) # never held a slot
elif w.granted: elif w.granted:
...@@ -200,8 +213,9 @@ class GpuSwapGate: ...@@ -200,8 +213,9 @@ class GpuSwapGate:
self._pump() self._pump()
raise raise
async def release(self, key) -> None: def release(self, key) -> None:
async with self._lock: """SYNCHRONOUS — safe to call from a `finally` even while the surrounding
coroutine is being cancelled (no await to interrupt)."""
self._running -= 1 self._running -= 1
if self._running <= 0: if self._running <= 0:
self._running = 0 self._running = 0
...@@ -215,6 +229,12 @@ class GpuSwapGate: ...@@ -215,6 +229,12 @@ class GpuSwapGate:
w.event.set() w.event.set()
def _swap_to(self, key) -> None: def _swap_to(self, key) -> None:
if self._log and key != self._owner:
try:
self._log(f"[gpu-swap] swapping GPU owner {self._owner!r} -> {key!r} "
f"(waiting={len(self._waiters)})")
except Exception:
pass
self._owner = key self._owner = key
self._served = 0 self._served = 0
for w in [x for x in self._waiters if x.key == key]: for w in [x for x in self._waiters if x.key == key]:
......
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