fix: bound ProcPool process wait

livekit-agents/livekit/agents/ipc/proc_pool.py

@@ -23,6 +23,7 @@
 ]
 
 MAX_CONCURRENT_INITIALIZATIONS = min(math.ceil(get_cpu_monitor().cpu_count()), 4)
+WARMED_PROCESS_WAIT_TIMEOUT_BUFFER = 2.0
 
 
 class ProcPool(utils.EventEmitter[EventTypes]):
@@ -129,7 +130,25 @@ async def launch_job(self, info: RunningJobInfo) -> None:
                         extra={"job_id": info.job.id},
                     )
 
-                proc = await self._warmed_proc_queue.get()
+                try:
+                    proc = await asyncio.wait_for(
+                        self._warmed_proc_queue.get(),
+                        timeout=self._initialize_timeout + WARMED_PROCESS_WAIT_TIMEOUT_BUFFER,
+                    )
+                except asyncio.TimeoutError as exc:
+                    if attempt == MAX_ATTEMPTS - 1:
+                        logger.error(
+                            "failed to launch job because no process became available after %d attempts",
+                            MAX_ATTEMPTS,
+                            extra={"job_id": info.job.id},
+                        )
+                        raise RuntimeError("no process became available for the job") from exc
+
+                    logger.warning(
+                        "timed out waiting for warmed process, retrying with a new process",
+                        extra={"job_id": info.job.id, "attempt": attempt + 1},
+                    )
+                    continue
             finally:
                 self._jobs_waiting_for_process -= 1
 

tests/test_ipc.py

@@ -14,6 +14,7 @@
 from typing import ClassVar
 
 import psutil
+import pytest
 
 from livekit.agents import JobContext, JobProcess, ipc, job, utils
 from livekit.agents.ipc.log_queue import LogQueueHandler, LogQueueListener
@@ -406,6 +407,43 @@ def _process_closed(proc: ipc.job_proc_executor.ProcJobExecutor):
         assert exitcode != 0, "process should have been killed"
 
 
+async def test_proc_pool_launch_job_times_out_when_spawn_never_warms_process(monkeypatch):
+    mp_ctx = mp.get_context("spawn")
+    loop = asyncio.get_running_loop()
+    pool = ipc.proc_pool.ProcPool(
+        initialize_process_fnc=_initialize_proc,
+        job_entrypoint_fnc=_job_entrypoint,
+        session_end_fnc=None,
+        num_idle_processes=0,
+        job_executor_type=job.JobExecutorType.THREAD,
+        initialize_timeout=0.01,
+        close_timeout=20.0,
+        session_end_timeout=300.0,
+        inference_executor=None,
+        memory_warn_mb=0,
+        memory_limit_mb=0,
+        http_proxy=None,
+        mp_ctx=mp_ctx,
+        loop=loop,
+    )
+
+    spawn_attempts = 0
+
+    async def _spawn_fails_without_queue_item() -> None:
+        nonlocal spawn_attempts
+        spawn_attempts += 1
+        await asyncio.sleep(0)
+
+    monkeypatch.setattr(pool, "_proc_spawn_task", _spawn_fails_without_queue_item)
+    monkeypatch.setattr(ipc.proc_pool, "WARMED_PROCESS_WAIT_TIMEOUT_BUFFER", 0.0)
+
+    with pytest.raises(RuntimeError, match="no process became available"):
+        await pool.launch_job(_generate_fake_job())
+
+    assert spawn_attempts == 3
+    assert pool._jobs_waiting_for_process == 0
+
+
 def _create_proc(
     *,
     close_timeout: float,