a cancelled page kept the pod slot it was granted (v0.4.23)

`placement.choose` runs in a thread and raises `pod.inflight` the moment it
picks a pod. If the awaiting coroutine was cancelled after the executor had
produced the Lease but before the value was delivered, the object was dropped
with the counter already raised -- and nothing released it, because the page
was gone and no `finally` ran.

A client that gives up on a batch cancels every page of it at once, so the leak
arrives in thousands rather than one at a time.

Measured on a live run, once `max_inflight_per_pod: 64` made it visible: the
pool reported all 52 healthy pods at their cap -- 3,328 pages "in flight" --
while `ss` inside the container showed 3 to 5 open sockets per pod, the cards
sat at `GPU KV cache usage: 0.2%` with `Running: 1 reqs`, and nothing was
preempted. Roughly 60 of every 64 pages did not exist. `_pick` then refused
every pod, correctly, reporting "all 52 serving pod(s) are at their in-flight
cap", and the run fell to 0.33 pages/s.

It had been leaking since the pod pipeline was written. With
`max_inflight_per_pod: 0` an inflated counter blocks nothing, so the only
symptom was DEMAND -- the in-flight counts of 2,700 and 4,096 that earlier runs
reported, and that the scaler dutifully rented pods to serve. The cap did not
cause this; it converted a silent accounting error into a hard stop, which is
how it was finally found. Three runs were spent on its consequences.

`_acquire_lease` shields the acquisition, so it completes even when the page
waiting on it is cancelled, and a done-callback hands the lease straight back.
Orphans are counted and said once a minute, because the entire cost of this bug
was that it was silent.

Note on the analysis: the gap between `lease = await ...` and the `try` that
releases it is NOT a window -- CancelledError is only delivered at an await
point, and there is none in between. The acquisition itself was the only one.

Proven rather than argued: the same test, same cancellation, against the
unshielded form and the shielded one -- granted=1 released=0 LEAKED=1, against
granted=1 released=1 LEAKED=0.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KnfkxBpcZ1v6ypRFFyhrVM

codai/__init__.py

@@ -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.4.22"
+__version__ = "0.4.23"
 __tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.

codai/ocr/manager.py

@@ -855,6 +855,68 @@ def _set_pod_wait_cooldown(mid: str) -> None:
         _POD_WAIT_COOLDOWN[mid] = time.time() + POD_WAIT_COOLDOWN_S
 
 
+#: Leases that were granted to a page that had already gone away. Counted, and
+#: said, because the whole cost of this bug was that it was silent.
+_ORPHAN_LEASES = [0]
+
+
+def _release_orphan(task) -> None:
+    """Give back a lease whose page was cancelled while acquiring it."""
+    try:
+        lease = task.result()
+    except BaseException:
+        return                      # it raised instead of leasing; nothing held
+    if lease is None:
+        return
+    _ORPHAN_LEASES[0] += 1
+    try:
+        lease.release()
+    except Exception:
+        pass
+    _say_rate_limited(
+        ("orphan-lease",),
+        f"[ocr] a page was cancelled while its pod lease was being granted; "
+        f"the lease was returned ({_ORPHAN_LEASES[0]} so far). A client that "
+        f"times out in bulk produces these in bulk.", every=60.0)
+
+
+async def _acquire_lease(model_id: str):
+    """Acquire a pod lease so that a CANCELLATION CANNOT LOSE ONE.
+
+    `placement.choose` runs in a thread and increments `pod.inflight` the
+    moment it picks a pod. If the awaiting coroutine is cancelled after the
+    executor has produced the Lease but before the value is delivered, the
+    object is dropped with the counter already raised -- and nothing ever
+    releases it. The page is gone, so no `finally` runs.
+
+    That is not a rare race. A client that gives up on a batch cancels every
+    page of it at once, so the leak arrives in thousands. Measured on a live
+    run: the pool reported all 52 pods at their 64-page cap -- 3,328 pages "in
+    flight" -- while `ss` showed 3 to 5 open sockets per pod and the cards sat
+    at 0.2% KV cache with `Running: 1`. About 60 of every 64 did not exist.
+    `_pick` then correctly refused every pod and the queue stopped at
+    0.33 pages/s.
+
+    It had been leaking all along. With `max_inflight_per_pod: 0` an inflated
+    counter blocked nothing, so it only showed up as demand: in-flight numbers
+    of 2,700 and 4,096 that the scaler dutifully rented pods for.
+
+    `shield` is the fix: the acquisition finishes even when the awaiting page is
+    cancelled, and the done-callback hands the lease straight back.
+    """
+    from codai.models import placement
+    task = asyncio.ensure_future(run_on_lease_pool(
+        placement.choose, model_id,
+        spill_after=placement.spill_threshold(model_id),
+        acquire_timeout=POD_LEASE_ATTEMPT_S,
+        grow_inline=False))
+    try:
+        return await asyncio.shield(task)
+    except asyncio.CancelledError:
+        task.add_done_callback(_release_orphan)
+        raise
+
+
 async def _lease_a_pod(model_id: str, policy: str, wait_s: float):
     """A lease on a pod that runs the pipeline, or None to read on this box.
 
@@ -870,27 +932,22 @@ async def _lease_a_pod(model_id: str, policy: str, wait_s: float):
     waited = 0.0
     while True:
         try:
-            # OFF THE LOOP. `placement.choose` is synchronous and ends in a
-            # pool's `acquire()`, which waits on a condition variable for up to
-            # its default 1200 s and may rent a pod INLINE -- a candidate sweep
+            # OFF THE LOOP, and cancellation-safe. `placement.choose` is
+            # synchronous and ends in a pool's `acquire()`, which waits on a
+            # condition variable and may rent a pod INLINE -- a candidate sweep
             # over every region, then a boot wait, in whichever thread asked.
             # Called straight from this coroutine it blocked the whole event
             # loop, so the process answered nothing at all -- not the client's
             # next document, not even /v1/runpod/status -- for as long as one
-            # page waited for a pod.
+            # page waited for a pod. It had been on the loop since the pod
+            # pipeline was written (0.3.70), costing nothing while acquire()
+            # returned at once, and turning fatal exactly when the pool was the
+            # thing under strain.
             #
-            # It has been on the loop since the pod pipeline was written
-            # (0.3.70) and cost nothing for as long as acquire() returned at
-            # once: a pool with a free pod answers in microseconds. It turns
-            # fatal exactly when the pool is the thing under strain, which is
-            # when it matters most. Measured on a 5,000-document run: 116 pods
-            # rented, zero documents completed in ten minutes, and a front that
-            # had stopped answering.
-            lease = await run_on_lease_pool(
-                placement.choose, model_id,
-                spill_after=placement.spill_threshold(model_id),
-                acquire_timeout=POD_LEASE_ATTEMPT_S,
-                grow_inline=False)
+            # `_acquire_lease` also makes sure a cancelled page cannot LOSE a
+            # lease it was granted -- see there; that leak is what made every
+            # pod look full while the cards idled.
+            lease = await _acquire_lease(model_id)
         except OcrError:
             raise
         except Exception as exc:

tests/test_ocr_lease_leak_on_cancel.py

+"""A cancelled page must not keep the pod slot it was granted.
+
+`placement.choose` runs in a thread and raises `pod.inflight` the moment it
+picks a pod. If the awaiting coroutine is cancelled after the executor has
+produced the Lease but before the value is delivered, the object is dropped
+with the counter already raised, and nothing releases it -- the page is gone,
+so no `finally` runs.
+
+A client that gives up on a batch cancels every page of it at once, so the leak
+arrives in thousands. Measured on a live run with `max_inflight_per_pod: 64`:
+the pool reported all 52 healthy pods at their cap -- 3,328 pages "in flight" --
+while `ss` showed 3 to 5 open sockets per pod, the cards sat at 0.2% KV cache
+with `Running: 1`, and nothing was preempted. Roughly 60 of every 64 pages did
+not exist. `_pick` then refused every pod, correctly, and the run fell to
+0.33 pages/s.
+
+It had been leaking since the pod pipeline was written. With
+`max_inflight_per_pod: 0` an inflated counter blocked nothing, so the only
+symptom was demand -- in-flight counts of 2,700 and 4,096 that the scaler
+rented pods to serve. The cap did not cause the problem; it made a silent one
+fatal, which is how it was finally found.
+"""
+
+import asyncio
+import threading
+import time
+
+import pytest
+
+from codai.ocr import manager as om
+
+
+@pytest.fixture
+def anyio_backend():
+    return "asyncio"
+
+
+@pytest.fixture(autouse=True)
+def _reset():
+    om._ORPHAN_LEASES[0] = 0
+    om._SAID_AT.clear()
+    yield
+    om._ORPHAN_LEASES[0] = 0
+    om._SAID_AT.clear()
+
+
+class _CountingPool:
+    """Stands in for a pod pool: counts slots taken and given back."""
+
+    def __init__(self):
+        self.inflight = 0
+        self.granted = 0
+        self.released = 0
+        self._lock = threading.Lock()
+
+    def take(self):
+        with self._lock:
+            self.inflight += 1
+            self.granted += 1
+
+    def give_back(self):
+        with self._lock:
+            self.inflight -= 1
+            self.released += 1
+
+
+class _Lease:
+    kind = "pod"
+    base = "https://10.0.0.1:1234/v1"
+    api_key = "k"
+
+    def __init__(self, pool):
+        self._pool = pool
+
+    def release(self):
+        self._pool.give_back()
+
+    def dead(self):
+        pass
+
+
+def _slow_choose(pool, delay=0.2):
+    """`choose` that takes a slot and then takes its time delivering it."""
+    def _choose(*_a, **_k):
+        pool.take()                 # the counter rises HERE, in the thread
+        time.sleep(delay)
+        return _Lease(pool)
+    return _choose
+
+
+# --------------------------------------------------------------- the leak #
+
+@pytest.mark.anyio
+async def test_a_cancelled_acquisition_gives_the_slot_back(monkeypatch):
+    """THE LEAK. Cancel while the lease is being granted; the slot must return."""
+    from codai.models import placement
+    pool = _CountingPool()
+    monkeypatch.setattr(placement, "choose", _slow_choose(pool))
+    monkeypatch.setattr(placement, "spill_threshold", lambda *_a, **_k: 0)
+
+    task = asyncio.ensure_future(om._acquire_lease("m"))
+    await asyncio.sleep(0.05)          # let it reach the executor
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    # The acquisition finishes in its own time; then the callback returns it.
+    for _ in range(100):
+        if pool.released:
+            break
+        await asyncio.sleep(0.02)
+    assert pool.granted == 1, "the slot was never taken; test proves nothing"
+    assert pool.released == 1, "the slot was NOT given back — still leaking"
+    assert pool.inflight == 0
+
+
+@pytest.mark.anyio
+async def test_the_orphan_is_counted_and_said(monkeypatch, capsys):
+    """The whole cost of this bug was that it was silent."""
+    from codai.models import placement
+    pool = _CountingPool()
+    monkeypatch.setattr(placement, "choose", _slow_choose(pool))
+    monkeypatch.setattr(placement, "spill_threshold", lambda *_a, **_k: 0)
+
+    task = asyncio.ensure_future(om._acquire_lease("m"))
+    await asyncio.sleep(0.05)
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    for _ in range(100):
+        if om._ORPHAN_LEASES[0]:
+            break
+        await asyncio.sleep(0.02)
+    assert om._ORPHAN_LEASES[0] == 1
+    out = capsys.readouterr().out
+    assert "cancelled while its pod lease was being granted" in out
+    assert "times out in bulk" in out
+
+
+@pytest.mark.anyio
+async def test_a_bulk_cancellation_leaks_nothing(monkeypatch):
+    """A client giving up on a batch cancels every page at once."""
+    from codai.models import placement
+    pool = _CountingPool()
+    monkeypatch.setattr(placement, "choose", _slow_choose(pool, delay=0.1))
+    monkeypatch.setattr(placement, "spill_threshold", lambda *_a, **_k: 0)
+
+    tasks = [asyncio.ensure_future(om._acquire_lease("m")) for _ in range(40)]
+    await asyncio.sleep(0.05)
+    for t in tasks:
+        t.cancel()
+    await asyncio.gather(*tasks, return_exceptions=True)
+    for _ in range(200):
+        if pool.released >= pool.granted and pool.granted:
+            break
+        await asyncio.sleep(0.02)
+    assert pool.granted > 0
+    assert pool.inflight == 0, (
+        f"{pool.granted} slots taken, {pool.released} returned — "
+        f"{pool.inflight} leaked")
+
+
+# ------------------------------------------------------- the normal paths #
+
+@pytest.mark.anyio
+async def test_an_uncancelled_acquisition_returns_the_lease(monkeypatch):
+    from codai.models import placement
+    pool = _CountingPool()
+    monkeypatch.setattr(placement, "choose", _slow_choose(pool, delay=0.01))
+    monkeypatch.setattr(placement, "spill_threshold", lambda *_a, **_k: 0)
+    lease = await om._acquire_lease("m")
+    assert lease is not None
+    assert pool.inflight == 1, "it must NOT release a lease it handed over"
+    lease.release()
+    assert pool.inflight == 0
+
+
+@pytest.mark.anyio
+async def test_nothing_is_released_when_choose_raises(monkeypatch):
+    """It raised instead of leasing, so there is no slot to give back."""
+    from codai.models import placement
+    pool = _CountingPool()
+
+    def _boom(*_a, **_k):
+        # Slow enough that the cancellation lands FIRST, which is the case
+        # under test: the page is gone and the acquisition then fails.
+        time.sleep(0.15)
+        raise RuntimeError("no upstream")
+
+    monkeypatch.setattr(placement, "choose", _boom)
+    monkeypatch.setattr(placement, "spill_threshold", lambda *_a, **_k: 0)
+    task = asyncio.ensure_future(om._acquire_lease("m"))
+    await asyncio.sleep(0.03)
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    await asyncio.sleep(0.3)
+    assert pool.released == 0
+    assert om._ORPHAN_LEASES[0] == 0
+
+
+@pytest.mark.anyio
+async def test_a_local_lease_is_not_counted_as_an_orphan(monkeypatch):
+    """choose may answer `local`, which holds no pod slot."""
+    from codai.models import placement
+
+    class _Local:
+        kind = "local"
+        base = "http://127.0.0.1:8000/v1"
+        api_key = ""
+        released = 0
+
+        def release(self):
+            type(self).released += 1
+
+    monkeypatch.setattr(placement, "choose", lambda *_a, **_k: _Local())
+    monkeypatch.setattr(placement, "spill_threshold", lambda *_a, **_k: 0)
+    lease = await om._acquire_lease("m")
+    assert lease.kind == "local"
+    assert om._ORPHAN_LEASES[0] == 0
+
+
+def test_the_release_is_wired_through_lease_a_pod():
+    """`_lease_a_pod` must use the safe acquisition, not the raw pool call."""
+    src = open(om.__file__).read()
+    i = src.index("async def _lease_a_pod(")
+    body = src[i:i + 3000]
+    assert "_acquire_lease(model_id)" in body
+    assert "run_on_lease_pool(" not in body, (
+        "it acquires directly again — a cancellation can lose that lease")