fix(runtime): preserve completed sessions without forking

Guard terminal gRPC streams from late cancellation cleanup and replace ps-based footprint sampling so long-running decode workers cannot deadlock gRPC during memory checks.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: 2 people

inference_engine/server/grpc_app.py

@@ -240,11 +240,15 @@ def _watch_context(
         context,
         cancel_event: threading.Event,
         session_id: str,
+        completed_event: threading.Event | None = None,
     ) -> None:
         callback = getattr(context, "add_done_callback", None)
         if callback is not None:
             def done(done_context) -> None:
-                if done_context.cancelled():
+                logically_completed = (
+                    completed_event is not None and completed_event.is_set()
+                )
+                if done_context.cancelled() and not logically_completed:
                     cancel_event.set()
                     self._store.remove_session_if_present(
                         session_id, reason="client_cancelled",
@@ -305,7 +309,13 @@ async def AppendTokens(  # noqa: N802 — gRPC-generated method casing
                 f"session {request.session_id!r} already has an active operation",
             )
         cancel_event = threading.Event()
-        self._watch_context(context, cancel_event, request.session_id)
+        completed_event = threading.Event()
+        self._watch_context(
+            context,
+            cancel_event,
+            request.session_id,
+            completed_event,
+        )
         try:
             kwargs = {
                 "session_id": request.session_id,
@@ -316,6 +326,7 @@ async def AppendTokens(  # noqa: N802 — gRPC-generated method casing
             new_history_length = await asyncio.to_thread(
                 self._append.append_tokens, **kwargs,
             )
+            completed_event.set()
         except asyncio.CancelledError:
             cancel_event.set()
             self._store.remove_session_if_present(
@@ -390,7 +401,13 @@ async def Generate(  # noqa: N802 — gRPC-generated method casing
                 f"session {request.session_id!r} already has an active operation",
             )
         cancel_event = threading.Event()
-        self._watch_context(context, cancel_event, request.session_id)
+        completed_event = threading.Event()
+        self._watch_context(
+            context,
+            cancel_event,
+            request.session_id,
+            completed_event,
+        )
         kwargs = {
             "session_id": request.session_id,
             "max_tokens": request.max_tokens,
@@ -413,6 +430,7 @@ async def Generate(  # noqa: N802 — gRPC-generated method casing
                 if error is not None:
                     raise error
                 if event is _STREAM_END:
+                    completed_event.set()
                     break
                 if context.cancelled():
                     threaded_stream.cancel()
@@ -436,6 +454,7 @@ async def Generate(  # noqa: N802 — gRPC-generated method casing
                     # DoneEvent — the only remaining event type per
                     # the GenerateEvent union.
                     assert isinstance(event, DoneEvent)
+                    completed_event.set()
                     yield runtime_pb2.GenerateResponse(
                         done=runtime_pb2.GenerateDone(
                             stop_reason=_STOP_REASON_TO_PROTO[
@@ -456,9 +475,10 @@ async def Generate(  # noqa: N802 — gRPC-generated method casing
             await context.abort(grpc.StatusCode.ABORTED, str(exc))
         except asyncio.CancelledError:
             threaded_stream.cancel()
-            self._store.remove_session_if_present(
-                request.session_id, reason="client_cancelled",
-            )
+            if not completed_event.is_set():
+                self._store.remove_session_if_present(
+                    request.session_id, reason="client_cancelled",
+                )
             raise
         finally:
             threaded_stream.cancel()

inference_engine/server/runtime_health.py

@@ -5,13 +5,13 @@
 import gc
 import json
 import os
-import subprocess
 import threading
 import time
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Callable, Optional
 
+import psutil
 from fastapi import FastAPI
 from fastapi.responses import JSONResponse
 
@@ -95,15 +95,10 @@ class MemorySnapshot:
 
 
 def process_footprint_bytes(pid: Optional[int] = None) -> int:
-    """Return resident process bytes using the macOS/Linux ``ps`` contract."""
+    """Return resident process bytes without forking the threaded runtime."""
     try:
-        output = subprocess.check_output(
-            ["ps", "-o", "rss=", "-p", str(pid or os.getpid())],
-            text=True,
-            timeout=1,
-        )
-        return int(output.strip()) * 1024
-    except (OSError, ValueError, subprocess.SubprocessError):
+        return int(psutil.Process(pid or os.getpid()).memory_info().rss)
+    except (OSError, ValueError, psutil.Error):
         return 0
 
 

requirements.txt

@@ -44,6 +44,7 @@ uvicorn>=0.32,<1.0
 pydantic>=2.7,<3.0
 httpx>=0.27,<1.0
 prometheus-client>=0.20,<1.0
+psutil>=5.9,<8.0
 
 # gRPC runtime (PR-B1 of ADR 0008 Phase B; the runtime/SDK protocol).
 # grpcio-tools is needed only at build time (regenerating stubs from

tests/inference_engine/server/test_grpc_app.py

@@ -664,6 +664,50 @@ async def abort(self, code, details):  # pragma: no cover
         store.get_session(sess.session_id)
 
 
+async def test_generate_cancellation_after_done_preserves_session():
+    from inference_engine.session import (
+        DoneEvent,
+        GenerationCoordinator,
+        STOP_REASON_MAX_TOKENS,
+    )
+
+    class CompletedGeneration(GenerationCoordinator):
+        def generate(self, session_id, *, max_tokens, **kwargs):
+            del session_id, max_tokens, kwargs
+            yield DoneEvent(
+                stop_reason=STOP_REASON_MAX_TOKENS,
+                generated_token_count=0,
+                prefill_seconds=0.0,
+                total_seconds=0.0,
+            )
+
+    class Context:
+        def cancelled(self):
+            return False
+
+        async def abort(self, code, details):
+            raise AssertionError((code, details))
+
+    store = SessionStore(capacity=1)
+    session = store.create_session()
+    servicer = RuntimeServiceServicer(
+        store,
+        generation_coordinator=CompletedGeneration(store, verifier=None),
+    )
+    stream = servicer.Generate(
+        runtime_pb2.GenerateRequest(
+            session_id=session.session_id,
+            max_tokens=1,
+        ),
+        Context(),
+    )
+    response = await anext(stream)
+    assert response.WhichOneof("payload") == "done"
+    with pytest.raises(asyncio.CancelledError):
+        await stream.athrow(asyncio.CancelledError())
+    assert store.get_session(session.session_id) is session
+
+
 async def test_wire_cancellation_releases_blocked_generate_session():
     """A disconnect removes session state while model thread is still blocked."""
     from inference_engine.session import GenerationCoordinator
@@ -945,6 +989,26 @@ async def test_watch_context_removes_cancelled_session():
     assert store.active_count == 0
 
 
+async def test_watch_context_preserves_logically_completed_session():
+    store = SessionStore(capacity=1)
+    session = store.create_session()
+    servicer = RuntimeServiceServicer(store)
+    context = _DirectContext()
+    cancel_event = threading.Event()
+    completed_event = threading.Event()
+    servicer._watch_context(
+        context,
+        cancel_event,
+        session.session_id,
+        completed_event,
+    )
+    completed_event.set()
+    assert context.callback is not None
+    context.callback(type("Done", (), {"cancelled": lambda self: True})())
+    assert not cancel_event.is_set()
+    assert store.active_count == 1
+
+
 async def test_create_session_rejects_memory_drain():
     store = SessionStore(capacity=1)
     governor = type("Governor", (), {"draining": True})()

tests/inference_engine/server/test_runtime_health.py

@@ -1,7 +1,9 @@
 import json
 
+import psutil
 from fastapi.testclient import TestClient
 
+from inference_engine.server import runtime_health
 from inference_engine.server.runtime_health import (
     DecodeLiveness,
     PrimaryMemoryGovernor,
@@ -18,6 +20,29 @@ def reset(self):
         self.resets += 1
 
 
+def test_process_footprint_uses_fork_free_psutil(monkeypatch):
+    seen = []
+
+    class Process:
+        def __init__(self, pid):
+            seen.append(pid)
+
+        def memory_info(self):
+            return type("MemoryInfo", (), {"rss": 12345})()
+
+    monkeypatch.setattr(runtime_health.psutil, "Process", Process)
+    assert runtime_health.process_footprint_bytes(42) == 12345
+    assert seen == [42]
+
+
+def test_process_footprint_returns_zero_when_process_disappears(monkeypatch):
+    def missing(pid):
+        raise psutil.NoSuchProcess(pid)
+
+    monkeypatch.setattr(runtime_health.psutil, "Process", missing)
+    assert runtime_health.process_footprint_bytes(42) == 0
+
+
 def test_liveness_is_atomically_published(tmp_path):
     path = tmp_path / "live.json"
     live = DecodeLiveness(path, clock=lambda: 123.0, pid=42)