fix(vllm): fail one request, not the engine, when the worker runs out of memory

interp_engine/vllm_backend.py

@@ -21,8 +21,8 @@
 import importlib.util
 import logging
 import os
-from collections.abc import Iterable, Sequence
-from typing import Any
+from collections.abc import Awaitable, Callable, Iterable, Sequence
+from typing import Any, TypeVar
 
 import torch
 
@@ -72,6 +72,8 @@
 
 logger = logging.getLogger(__name__)
 
+_T = TypeVar("_T")
+
 
 def _device_capability() -> tuple[int, int] | None:
     """This GPU's compute capability, or None when there is no CUDA device to ask.
@@ -709,6 +711,24 @@ def _decode_rank0(payloads: object) -> dict[Address, torch.Tensor]:
     return decode_capture_payload(payloads[0] if isinstance(payloads, list | tuple) else payloads)  # type: ignore[index]
 
 
+async def _settle(collect: Callable[[], Awaitable[_T]], release: Callable[[], Awaitable[None]]) -> _T:
+    """A request's last worker calls: ``collect`` its rows, then ``release`` its registration.
+
+    ``release`` runs even when ``collect`` raises, which a capture the worker dropped for lack of
+    memory does on purpose. Both are shielded: a server cancels a request's task when its client
+    goes, and under anyio every await in a cancelled scope is cancelled too, so an unshielded
+    cleanup never reaches the worker and the request's rows stay there for the process lifetime.
+    """
+
+    async def run() -> _T:
+        try:
+            return await collect()
+        finally:
+            await release()
+
+    return await asyncio.shield(run())
+
+
 def _step_logprobs(per_position: object, index: int, n_logprobs: int) -> list[dict[str, float | int]] | None:
     """One generated position's top-n, in the shape :class:`~interp_engine.steer.GenStep` carries.
 
@@ -1629,6 +1649,13 @@ async def _register_static_write(
     async def _unregister_static_write(self, rid: str) -> None:
         await self.engine.collective_rpc("unregister_static_write", args=(rid,))
 
+    async def _release_writes(self, rid: str, *, static_write: bool, steered: bool) -> None:
+        """Drop ``rid``'s steering or lens writes from the worker, whichever it registered."""
+        if static_write:
+            await self._unregister_static_write(rid)
+        elif steered:
+            await self.engine.collective_rpc("unregister_steering", args=(rid,))
+
     def _lens_scope(self, lens_intervention: dict) -> dict:
         return {
             "steer_generated": bool(lens_intervention.get("steer_generated", False)),
@@ -1810,18 +1837,19 @@ async def generate_steered(
             )
 
         async def _finish() -> None:
+            async def collect() -> object:
+                if not (capturing and capture_out is not None):
+                    return None
+                method = "collect_static" if static_cap else "collect_request"
+                return await self.engine.collective_rpc(method, args=(rid,))
+
+            payloads = await _settle(
+                collect, lambda: self._release_writes(rid, static_write=static_write, steered=steered)
+            )
             if capturing and capture_out is not None:
-                if static_cap:
-                    payloads = await self.engine.collective_rpc("collect_static", args=(rid,))
-                else:
-                    payloads = await self.engine.collective_rpc("collect_request", args=(rid,))
                 _merge_captures(capture_out, _decode_rank0(payloads))
                 _assert_points_captured(capture_out, pts)
                 _assert_full_width_captured(capture_out, self._hidden_size)
-            if static_write:
-                await self._unregister_static_write(rid)
-            elif steered:
-                await self.engine.collective_rpc("unregister_steering", args=(rid,))
 
         if stream:
 
@@ -1901,10 +1929,7 @@ async def _generate_request_outputs(
             async for out in self.engine.generate(prompt, sampling_params, rid):
                 yield out
         finally:
-            if static_write:
-                await self._unregister_static_write(rid)
-            elif steered:
-                await self.engine.collective_rpc("unregister_steering", args=(rid,))
+            await asyncio.shield(self._release_writes(rid, static_write=static_write, steered=steered))
 
     async def _drain_into(self, rid: str, capture_out: dict[Address, torch.Tensor], *, static: bool = False) -> None:
         """Move ``rid``'s captured rows so far to the host, leaving the hooks / static taps installed."""
@@ -2238,14 +2263,10 @@ async def capture(
                 request_id=rid,
             )
         finally:
-            if static_cap:
-                payloads = await self.engine.collective_rpc("collect_static", args=(rid,))
-            else:
-                payloads = await self.engine.collective_rpc("collect_request", args=(rid,))
-            if static_write:
-                await self._unregister_static_write(rid)
-            elif steered:
-                await self.engine.collective_rpc("unregister_steering", args=(rid,))
+            payloads = await _settle(
+                lambda: self.engine.collective_rpc("collect_static" if static_cap else "collect_request", args=(rid,)),
+                lambda: self._release_writes(rid, static_write=static_write, steered=steered),
+            )
         out = decode_capture_payload(payloads[0] if isinstance(payloads, list | tuple) else payloads)
         _assert_points_captured(out, pts)
         _assert_full_prompt_captured(out, len(prompt_token_ids))
@@ -2329,14 +2350,10 @@ async def capture_generation(
                 request_id=rid,
             )
         finally:
-            if static_cap:
-                payloads = await self.engine.collective_rpc("collect_static", args=(rid,))
-            else:
-                payloads = await self.engine.collective_rpc("collect_request", args=(rid,))
-            if static_write:
-                await self._unregister_static_write(rid)
-            elif steered or lens:
-                await self.engine.collective_rpc("unregister_steering", args=(rid,))
+            payloads = await _settle(
+                lambda: self.engine.collective_rpc("collect_static" if static_cap else "collect_request", args=(rid,)),
+                lambda: self._release_writes(rid, static_write=static_write, steered=bool(steered or lens)),
+            )
         caps = decode_capture_payload(payloads[0] if isinstance(payloads, list | tuple) else payloads)
         _assert_points_captured(caps, pts)
         _assert_full_width_captured(caps, self._hidden_size)
@@ -2428,15 +2445,11 @@ async def capture_generation_stream(
         finally:
             # Deregister on every exit path (including client disconnect), and keep whatever
             # the last forward appended after the final drain so no position is dropped.
-            if static_cap:
-                payloads = await self.engine.collective_rpc("collect_static", args=(rid,))
-            else:
-                payloads = await self.engine.collective_rpc("collect_request", args=(rid,))
+            payloads = await _settle(
+                lambda: self.engine.collective_rpc("collect_static" if static_cap else "collect_request", args=(rid,)),
+                lambda: self._release_writes(rid, static_write=static_write, steered=bool(lens)),
+            )
             tail = decode_capture_payload(payloads[0] if isinstance(payloads, list | tuple) else payloads)
-            if static_write:
-                await self._unregister_static_write(rid)
-            elif lens:
-                await self.engine.collective_rpc("unregister_steering", args=(rid,))
         if completed and (tail or len(token_ids) > reported):
             seen.update(tail)
             yield tail, token_ids
@@ -2596,11 +2609,10 @@ async def readout(final: bool) -> tuple[int, int, list[torch.Tensor], list[torch
         finally:
             # Deregister on every exit path (including client disconnect), and read out whatever
             # the last forward appended after the final drain so no position is dropped.
-            tail = await readout(final=True)
-            if static_write:
-                await self._unregister_static_write(rid)
-            elif lens:
-                await self.engine.collective_rpc("unregister_steering", args=(rid,))
+            tail = await _settle(
+                lambda: readout(final=True),
+                lambda: self._release_writes(rid, static_write=static_write, steered=bool(lens)),
+            )
         if completed and (tail[1] or len(token_ids) > reported):
             yield tail[0], tail[2], tail[3], token_ids
         if completed and captured_rows == 0:
@@ -2800,10 +2812,9 @@ async def capture_attention(
                 request_id=rid,
             )
         finally:
-            if static_attn:
-                payloads = await self.engine.collective_rpc("collect_static", args=(rid,))
-            else:
-                payloads = await self.engine.collective_rpc("collect_attn_request", args=(rid,))
+            payloads = await asyncio.shield(
+                self.engine.collective_rpc("collect_static" if static_attn else "collect_attn_request", args=(rid,))
+            )
         return recompute_attn_from_payloads(payloads, layers, self._attn_dims, self.tensor_parallel_size)
 
     async def set_steering(self, specs: list[dict]) -> None:

interp_engine/vllm_capture/_tp.py

@@ -54,6 +54,22 @@ def _all_gather(tensor: torch.Tensor, dim: int) -> torch.Tensor:
     return tensor_model_parallel_all_gather(tensor.contiguous(), dim)
 
 
+def raise_if_any_rank_failed(failure: str | None, what: str) -> None:
+    """Raise when ``failure`` is set on this rank or any other, so every rank raises together.
+
+    A collective, called before the next one at the same point on every rank. One rank raising
+    while the rest go on to a gather is a hang, not an error: they wait for it forever.
+    """
+    failed = failure is not None
+    if tp_size() > 1:
+        from vllm.distributed import tensor_model_parallel_all_reduce  # pyright: ignore[reportMissingImports]
+
+        flag = torch.tensor([1.0 if failed else 0.0], device=torch.device("cuda", torch.cuda.current_device()))
+        failed = bool(tensor_model_parallel_all_reduce(flag).item() > 0)
+    if failed:
+        raise RuntimeError(f"{what} failed: {failure or 'on another tensor-parallel rank'}")
+
+
 def _head_counts(layer: torch.nn.Module) -> tuple[int, int, int | None, int] | None:
     """``(heads_per_rank, kv_heads_per_rank, total_kv_heads, head_dim)`` off the attention module.
 

interp_engine/vllm_capture/lens/readout.py

@@ -26,12 +26,14 @@
     encode_tensor_payload,
     hook_site,
 )
+from interp_engine.vllm_capture._tp import raise_if_any_rank_failed
 from interp_engine.vllm_capture._tree import _worker_model, scale_capture
 from interp_engine.vllm_capture.lens.unembed import (
     _assert_applied_logit_scale_agrees,
     _worker_applied_softcap,
     _worker_final_norm,
 )
+from interp_engine.vllm_capture.static import _fail_capture, take_static_rows
 from interp_engine.vllm_capture.static import _state as _static_state
 
 
@@ -190,18 +192,21 @@ def worker_set_lens_jacobians(worker: object, payloads: dict[str, tuple] | None)
     return {"layers": len(jacobians), "bytes": total}
 
 
-def _static_harvest_rows(worker: object, static: Any, req_id: str) -> dict[str, torch.Tensor]:
-    """Concatenated static harvest for ``req_id``, scaled like :func:`worker_collect_static`."""
-    harvest = static.harvest.get(req_id) or {}
-    if not harvest:
-        return {}
-    model = _worker_model(worker)
+def _concat_rows(
+    model: torch.nn.Module, taken: dict[str, list[torch.Tensor]], *, scale: bool
+) -> dict[str, torch.Tensor]:
+    """One tensor per key from rows the caller has taken out of the capture store.
+
+    ``scale`` corrects a static harvest like :func:`worker_collect_static` does. Each key's pieces
+    are freed once joined, so the peak is the rows plus one key's copy rather than twice the rows.
+    """
     rows: dict[str, torch.Tensor] = {}
-    for key, chunks in harvest.items():
+    for key, chunks in taken.items():
         if not chunks:
             continue
         tensor = chunks[0] if len(chunks) == 1 else torch.cat(chunks, dim=0)
-        rows[key] = scale_capture(model, key, tensor)
+        chunks.clear()
+        rows[key] = scale_capture(model, key, tensor) if scale else tensor
     return rows
 
 
@@ -274,50 +279,30 @@ def worker_lens_capture_readout(
     demux = _get_demux(worker)
     static = _static_state(worker)
     static_active = static is not None and req_id in static.cap_points
-    caps: dict = {}
+    # The rows leave the capture store before anything is allocated from them, and a final call
+    # deregisters first. Joining them can run out of memory; when it does, the request fails and
+    # its rows are freed, rather than staying on the worker for every later step to step around.
+    failure: str | None = None
     if static_active:
         assert static is not None
-        rows_by_key = _static_harvest_rows(worker, static, req_id)
         cursor = static.lens_cursor.get(req_id, 0)
+        taken, failure = take_static_rows(static, req_id, final=final)
+        if final:
+            static.registered.discard(req_id)
+            demux.registered.discard(req_id)
     else:
         caps = demux.captures.get(req_id) or {}
-        rows_by_key = {key: torch.cat(tensors, dim=0) for key, tensors in caps.items() if tensors}
         cursor = demux.lens_cursor.get(req_id, 0)
-    if final:
-        if static_active:
-            assert static is not None
-            static.harvest.pop(req_id, None)
-            static.cap_points.pop(req_id, None)
-            static.lens_cursor.pop(req_id, None)
-            static.registered.discard(req_id)
-            demux.registered.discard(req_id)
-        else:
+        taken = {key: list(tensors) for key, tensors in caps.items() if tensors}
+        for tensors in caps.values():
+            tensors.clear()
+        if final:
             demux.captures.pop(req_id, None)
             demux.lens_cursor.pop(req_id, None)
             points = demux.cap_points.pop(req_id, set())
             for site in {hook_site(a) for a in points}:
                 _release_hook(demux, site)
             _maybe_unregister(demux, req_id)
-    elif static_active:
-        assert static is not None
-        static.harvest.pop(req_id, None)
-    else:
-        for tensors in caps.values():
-            tensors.clear()
-
-    empty: dict[str, Any] = {
-        "first_position": cursor,
-        "n_positions": 0,
-        # Rows taken, which `n_positions` does not report once `skip_before` drops some. The
-        # caller distinguishes "captured nothing" (hooks never fired) from "read out nothing".
-        "n_rows": 0,
-        "results": [{"top_idx": None, "top_probs": None} for _ in specs],
-    }
-    # Nothing captured since the last call is ordinary mid-stream: the engine may not have
-    # run a forward for this request yet. A point missing while others are present is not,
-    # and `rows_for` raises for that below.
-    if not rows_by_key:
-        return empty
 
     model = _worker_model(worker)
     param = next(model.parameters())
@@ -346,7 +331,40 @@ def rows_for(layer: int) -> torch.Tensor:
         return rows
 
     wanted = sorted({int(layer) for spec in specs for layer in spec["layers"]})
-    n_rows = min((rows_for(layer).shape[0] for layer in wanted), default=0)
+    rows_by_key: dict[str, torch.Tensor] = {}
+    n_rows = 0
+    if failure is None:
+        try:
+            rows_by_key = _concat_rows(model, taken, scale=static_active)
+            if rows_by_key:
+                n_rows = min((rows_for(layer).shape[0] for layer in wanted), default=0)
+        except torch.cuda.OutOfMemoryError as exc:
+            failure = f"out of memory joining the captured rows: {exc}"
+            rows_by_key.clear()
+            reduced.clear()
+            if static_active and not final:
+                assert static is not None
+                _fail_capture(static, req_id, failure)
+    taken.clear()
+    # Every wanted layer is memoized now. On a stacked trunk the full-width rows are n_streams
+    # times the reduced ones, and nothing reads them again.
+    rows_by_key.clear()
+    # Every rank reaches this on every call, before the unembed's gather and before any return.
+    raise_if_any_rank_failed(failure, f"lens read-out for {req_id}")
+
+    empty: dict[str, Any] = {
+        "first_position": cursor,
+        "n_positions": 0,
+        # Rows taken, which `n_positions` does not report once `skip_before` drops some. The
+        # caller distinguishes "captured nothing" (hooks never fired) from "read out nothing".
+        "n_rows": 0,
+        "results": [{"top_idx": None, "top_probs": None} for _ in specs],
+    }
+    # Nothing captured since the last call is ordinary mid-stream: the engine may not have
+    # run a forward for this request yet. A point missing while others are present is not,
+    # and `rows_for` raises for that above.
+    if not reduced:
+        return empty
     if not final:
         if static_active:
             assert static is not None