fix(ipc): fail pending inference requests when the inference process dies

livekit-agents/livekit/agents/ipc/inference_proc_executor.py

@@ -46,6 +46,9 @@ def __init__(
 
         self._runners = runners
         self._active_requests: dict[str, asyncio.Future[proto.InferenceResponse]] = {}
+        # set once the IPC read loop has ended: the process is gone (or being torn
+        # down) and no inference response can ever arrive anymore
+        self._requests_closed = False
 
     @property
     def process_kind(self) -> SupervisedProcKind:
@@ -66,35 +69,75 @@ def _create_process(self, cch: socket.socket, log_cch: socket.socket) -> mp.Proc
 
     @log_exceptions(logger=logger)
     async def _main_task(self, ipc_ch: aio.ChanReceiver[channel.Message]) -> None:
-        async for msg in ipc_ch:
-            if isinstance(msg, proto.InferenceResponse):
-                fut = self._active_requests.pop(msg.request_id, None)
-                if fut is None:
-                    logger.warning(
-                        "received unexpected inference response",
-                        extra={"request_id": msg.request_id},
-                    )
-                    continue
-
-                with contextlib.suppress(asyncio.InvalidStateError):
-                    fut.set_result(msg)
+        try:
+            async for msg in ipc_ch:
+                if isinstance(msg, proto.InferenceResponse):
+                    fut = self._active_requests.pop(msg.request_id, None)
+                    if fut is None:
+                        logger.warning(
+                            "received unexpected inference response",
+                            extra={"request_id": msg.request_id},
+                        )
+                        continue
+
+                    with contextlib.suppress(asyncio.InvalidStateError):
+                        fut.set_result(msg)
+        finally:
+            # the read task is the sole reader of the IPC channel; once it ends, the
+            # inference process is gone (channel closed) or the executor is being torn
+            # down. Reject the pending requests so do_inference() callers don't wait
+            # forever for a response that can never arrive.
+            self._fail_pending_requests()
+
+    def _fail_pending_requests(self) -> None:
+        """Reject outstanding inference requests after the IPC reader stops."""
+        self._requests_closed = True
+
+        pending_requests, self._active_requests = self._active_requests, {}
+        for fut in pending_requests.values():
+            with contextlib.suppress(asyncio.InvalidStateError):
+                fut.set_exception(RuntimeError("the inference process exited before responding"))
 
     async def do_inference(self, method: str, data: bytes) -> bytes | None:
         if not self.started:
             raise RuntimeError("process not started")
 
+        if self._requests_closed:
+            # the process is gone (or shutting down), a response would never arrive
+            raise RuntimeError("the inference process exited before responding")
+
         request_id = shortuuid("inference_req_")
         fut = asyncio.Future[proto.InferenceResponse]()
         self._active_requests[request_id] = fut
 
+        def _drop_cancelled_request(fut: asyncio.Future[proto.InferenceResponse]) -> None:
+            # the caller stopped waiting (its task was cancelled while awaiting the
+            # response): forget the request so the reader loop and the
+            # pending-requests cleanup never try to complete it afterwards
+            if fut.cancelled():
+                self._active_requests.pop(request_id, None)
+
+        fut.add_done_callback(_drop_cancelled_request)
+
         try:
             await channel.asend_message(
                 self._pch,
                 proto.InferenceRequest(request_id=request_id, method=method, data=data),
             )
-        except Exception:
+        except BaseException:
+            # BaseException, not Exception: a cancellation while the send is blocked
+            # (e.g. a large request stuck in the writer's drain()) used to bypass
+            # this cleanup and leave the future registered with no waiter, so
+            # _fail_pending_requests() later completed it with an exception nobody
+            # retrieved.
             if not fut.done():
                 fut.cancel()
+            elif not fut.cancelled():
+                # the pending-requests cleanup rejected it while this send was
+                # blocked (the process died): retrieve the exception so asyncio
+                # doesn't report it as never retrieved, the send error is raised
+                # to the caller instead
+                fut.exception()
             self._active_requests.pop(request_id, None)
             raise
 

tests/test_ipc.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import ctypes
+import gc
 import io
 import logging
 import multiprocessing as mp
@@ -17,6 +19,7 @@
 import pytest
 
 from livekit.agents import JobContext, JobProcess, ipc, job, utils
+from livekit.agents.inference_runner import _InferenceRunner
 from livekit.agents.ipc.log_queue import LogQueueHandler, LogQueueListener
 from livekit.agents.utils.aio import duplex_unix
 from livekit.protocol import agent
@@ -771,3 +774,205 @@ def test_inference_executor_is_alive_after_proc_closed() -> None:
     executor._proc = proc
 
     assert executor.is_alive() is False
+
+
+class _EchoInferenceRunner(_InferenceRunner):
+    INFERENCE_METHOD = "test_echo"
+
+    def initialize(self) -> None: ...
+
+    def run(self, data: bytes) -> bytes | None:
+        return data
+
+
+class _BlockingInferenceRunner(_InferenceRunner):
+    INFERENCE_METHOD = "test_block"
+
+    def initialize(self) -> None: ...
+
+    def run(self, data: bytes) -> bytes | None:
+        # keep the request in-flight long enough for the process to be killed mid-request
+        time.sleep(10)
+        return data
+
+
+def _create_inf_executor() -> ipc.inference_proc_executor.InferenceProcExecutor:
+    return ipc.inference_proc_executor.InferenceProcExecutor(
+        runners={
+            _EchoInferenceRunner.INFERENCE_METHOD: _EchoInferenceRunner,
+            _BlockingInferenceRunner.INFERENCE_METHOD: _BlockingInferenceRunner,
+        },
+        initialize_timeout=20.0,
+        close_timeout=20.0,
+        memory_warn_mb=0,
+        memory_limit_mb=0,
+        ping_interval=2.5,
+        ping_timeout=10.0,
+        high_ping_threshold=1.0,
+        mp_ctx=mp.get_context("spawn"),
+        loop=asyncio.get_running_loop(),
+        http_proxy=None,
+    )
+
+
+async def test_inference_proc_death_fails_pending_requests():
+    """Regression test: when the inference process dies (crash, OOM kill, manual
+    kill) with a request in flight, the pending do_inference() call must reject
+    with an error instead of waiting forever for a response that can never
+    arrive."""
+    executor = _create_inf_executor()
+    try:
+        await executor.start()
+        await executor.initialize()
+
+        # round-trip one request first, so a failure below can't be blamed on startup
+        res = await asyncio.wait_for(executor.do_inference("test_echo", b"hello"), timeout=20.0)
+        assert res == b"hello"
+
+        inf_task = asyncio.create_task(executor.do_inference("test_block", b""))
+        # do_inference registers the request and completes the send without ever
+        # yielding to the event loop, so once the request is visible it was sent
+        await _poll_until(lambda: len(executor._active_requests) == 1)
+
+        await executor.kill()
+
+        with pytest.raises(RuntimeError, match="exited before responding"):
+            await asyncio.wait_for(inf_task, timeout=5.0)
+    finally:
+        with contextlib.suppress(Exception):
+            await executor.aclose()
+
+
+def _stalled_inference_proc_main(cch: socket.socket) -> None:
+    # complete the init handshake, then never read from the IPC socket again:
+    # simulates an inference process that hung with its socket undrained, so a
+    # large request from the parent blocks inside the writer's drain()
+    dplx = duplex_unix._Duplex.open(cch)
+    first_req = ipc.channel.recv_message(dplx, ipc.proto.IPC_MESSAGES)
+    assert isinstance(first_req, ipc.proto.InitializeRequest), (
+        "first message must be proto.InitializeRequest"
+    )
+    ipc.channel.send_message(dplx, ipc.proto.InitializeResponse())
+    time.sleep(3600)
+
+
+class _StalledInferenceExecutor(ipc.inference_proc_executor.InferenceProcExecutor):
+    def _create_process(self, cch: socket.socket, log_cch: socket.socket) -> mp.Process:
+        return self._mp_ctx.Process(  # type: ignore
+            target=_stalled_inference_proc_main,
+            args=(cch,),
+            name="agents_inference_process",
+        )
+
+
+def _create_stalled_inf_executor() -> ipc.inference_proc_executor.InferenceProcExecutor:
+    return _StalledInferenceExecutor(
+        runners={_EchoInferenceRunner.INFERENCE_METHOD: _EchoInferenceRunner},
+        initialize_timeout=20.0,
+        close_timeout=20.0,
+        memory_warn_mb=0,
+        memory_limit_mb=0,
+        # the stalled child never answers pings: keep the health checks from
+        # killing it while the test is running
+        ping_interval=9999.0,
+        ping_timeout=9999.0,
+        high_ping_threshold=1.0,
+        mp_ctx=mp.get_context("spawn"),
+        loop=asyncio.get_running_loop(),
+        http_proxy=None,
+    )
+
+
+class _AsyncioErrorCapture(logging.Handler):
+    """Collect asyncio error records, e.g. unretrieved future exceptions."""
+
+    def __init__(self) -> None:
+        super().__init__(level=logging.ERROR)
+        self.messages: list[str] = []
+
+    def emit(self, record: logging.LogRecord) -> None:
+        self.messages.append(record.getMessage())
+
+
+async def test_cancelled_inference_send_leaves_no_pending_request():
+    """Regression test: cancelling do_inference() while its request is blocked in
+    the IPC writer's drain() must deregister the request. The orphaned future
+    would otherwise have no waiter and _fail_pending_requests() would later
+    complete it with an exception nobody retrieves (asyncio "Future exception was
+    never retrieved" warning when the process exits)."""
+    executor = _create_stalled_inf_executor()
+    capture = _AsyncioErrorCapture()
+    logging.getLogger("asyncio").addHandler(capture)
+    task_a = task_b = None
+    try:
+        await executor.start()
+        await executor.initialize()
+
+        payload = b"x" * (4 * 1024 * 1024)
+        # the first large request fills the OS send pipeline against the child
+        # that stopped reading...
+        task_a = asyncio.create_task(executor.do_inference("test_echo", payload))
+        await _poll_until(lambda: len(executor._active_requests) == 1)
+        # ...so this send blocks inside the writer's drain()
+        task_b = asyncio.create_task(executor.do_inference("test_echo", payload))
+        await _poll_until(lambda: len(executor._active_requests) == 2)
+
+        task_b.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task_b
+
+        # the cancelled send must not leave its future registered without a waiter
+        assert len(executor._active_requests) == 1
+
+        # the process dies: the remaining request is rejected (its send had
+        # completed) or fails on the closed channel (its send was still blocked
+        # in the writer), and no unhandled future failure may surface
+        await executor.kill()
+        with contextlib.suppress(RuntimeError, duplex_unix.DuplexClosed):
+            await asyncio.wait_for(task_a, timeout=5.0)
+
+        # the cancelled task holds the do_inference frame (and with it the request
+        # future) alive through its CancelledError traceback: release the
+        # references before collecting, a pending unretrieved exception could
+        # otherwise be missed
+        task_a = task_b = None
+        gc.collect()
+        assert not [m for m in capture.messages if "was never retrieved" in m]
+    finally:
+        logging.getLogger("asyncio").removeHandler(capture)
+        for t in (task_a, task_b):
+            if t is not None and not t.done():
+                t.cancel()
+                with contextlib.suppress(asyncio.CancelledError, Exception):
+                    await t
+        with contextlib.suppress(Exception):
+            await executor.kill()
+
+
+async def test_cancelled_inference_wait_leaves_no_pending_request():
+    """Regression test: cancelling do_inference() while it waits for the response
+    must also deregister the request, the dead entry would otherwise stay in
+    _active_requests until the process exits."""
+    executor = _create_stalled_inf_executor()
+    task = None
+    try:
+        await executor.start()
+        await executor.initialize()
+
+        # small payload: the send itself completes against the stalled child, the
+        # task then waits for a response that never arrives
+        task = asyncio.create_task(executor.do_inference("test_echo", b"hello"))
+        await _poll_until(lambda: len(executor._active_requests) == 1)
+
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+        assert len(executor._active_requests) == 0
+    finally:
+        if task is not None and not task.done():
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError, Exception):
+                await task
+        with contextlib.suppress(Exception):
+            await executor.kill()