feat(ipc): configure entrypoint shutdown timeout

livekit-agents/livekit/agents/ipc/job_proc_executor.py

@@ -30,6 +30,7 @@ def __init__(
         inference_executor: InferenceExecutor | None,
         initialize_timeout: float,
         close_timeout: float,
+        entrypoint_shutdown_timeout: float,
         session_end_timeout: float,
         memory_warn_mb: float,
         memory_limit_mb: float,
@@ -60,6 +61,7 @@ def __init__(
         self._job_entrypoint_fnc = job_entrypoint_fnc
         self._session_end_fnc = session_end_fnc
         self._simulation_end_fnc = simulation_end_fnc
+        self._entrypoint_shutdown_timeout = entrypoint_shutdown_timeout
         self._session_end_timeout = session_end_timeout
         self._inference_executor = inference_executor
         self._inference_tasks: set[asyncio.Task[None]] = set()
@@ -106,6 +108,7 @@ def _create_process(self, cch: socket.socket, log_cch: socket.socket) -> mp.Proc
             job_entrypoint_fnc=self._job_entrypoint_fnc,
             session_end_fnc=self._session_end_fnc,
             simulation_end_fnc=self._simulation_end_fnc,
+            entrypoint_shutdown_timeout=self._entrypoint_shutdown_timeout,
             session_end_timeout=self._session_end_timeout,
             log_cch=log_cch,
             mp_cch=cch,

livekit-agents/livekit/agents/ipc/job_proc_lazy_main.py

@@ -57,6 +57,7 @@ class ProcStartArgs:
     initialize_process_fnc: Callable[[JobProcess], Any]
     job_entrypoint_fnc: Callable[[JobContext], Any]
     session_end_fnc: Callable[[JobContext], Awaitable[None]] | None
+    entrypoint_shutdown_timeout: float
     session_end_timeout: float
     user_arguments: Any | None
     mp_cch: socket.socket
@@ -83,6 +84,7 @@ def proc_main(args: ProcStartArgs) -> None:
         args.initialize_process_fnc,
         args.job_entrypoint_fnc,
         args.session_end_fnc,
+        entrypoint_shutdown_timeout=args.entrypoint_shutdown_timeout,
         session_end_timeout=args.session_end_timeout,
         executor_type=JobExecutorType.PROCESS,
         user_arguments=args.user_arguments,
@@ -191,6 +193,7 @@ def __init__(
         job_entrypoint_fnc: Callable[[JobContext], Any],
         session_end_fnc: Callable[[JobContext], Awaitable[None]] | None,
         *,
+        entrypoint_shutdown_timeout: float,
         session_end_timeout: float,
         executor_type: JobExecutorType,
         user_arguments: Any | None = None,
@@ -202,6 +205,7 @@ def __init__(
         self._job_entrypoint_fnc = job_entrypoint_fnc
         self._session_end_fnc = session_end_fnc
         self._simulation_end_fnc = simulation_end_fnc
+        self._entrypoint_shutdown_timeout = entrypoint_shutdown_timeout
         self._session_end_timeout = session_end_timeout
         self._job_task: asyncio.Task[None] | None = None
 
@@ -379,7 +383,10 @@ def _on_entry_done(t: asyncio.Task[Any]) -> None:
         # wait for the entrypoint to finish, cancel if it takes too long
         if not job_entry_task.done():
             try:
-                await asyncio.wait_for(asyncio.shield(job_entry_task), timeout=15)
+                await asyncio.wait_for(
+                    asyncio.shield(job_entry_task),
+                    timeout=self._entrypoint_shutdown_timeout,
+                )
             except asyncio.TimeoutError:
                 logger.warning("entrypoint did not exit in time, cancelling")
                 await aio.cancel_and_wait(job_entry_task)
@@ -449,6 +456,7 @@ class ThreadStartArgs:
     initialize_process_fnc: Callable[[JobProcess], Any]
     job_entrypoint_fnc: Callable[[JobContext], Any]
     session_end_fnc: Callable[[JobContext], Awaitable[None]] | None
+    entrypoint_shutdown_timeout: float
     session_end_timeout: float
     join_fnc: Callable[[], None]
     mp_cch: socket.socket
@@ -467,6 +475,7 @@ def thread_main(
             args.initialize_process_fnc,
             args.job_entrypoint_fnc,
             args.session_end_fnc,
+            entrypoint_shutdown_timeout=args.entrypoint_shutdown_timeout,
             session_end_timeout=args.session_end_timeout,
             executor_type=JobExecutorType.THREAD,
             user_arguments=args.user_arguments,

livekit-agents/livekit/agents/ipc/job_thread_executor.py

@@ -26,6 +26,7 @@ class _ProcOpts:
     simulation_end_fnc: Callable[[Any], Any] | None
     initialize_timeout: float
     close_timeout: float
+    entrypoint_shutdown_timeout: float
     session_end_timeout: float
     ping_interval: float
     high_ping_threshold: float
@@ -43,6 +44,7 @@ def __init__(
         inference_executor: InferenceExecutor | None,
         initialize_timeout: float,
         close_timeout: float,
+        entrypoint_shutdown_timeout: float,
         session_end_timeout: float,
         ping_interval: float,
         high_ping_threshold: float,
@@ -57,6 +59,7 @@ def __init__(
             simulation_end_fnc=simulation_end_fnc,
             initialize_timeout=initialize_timeout,
             close_timeout=close_timeout,
+            entrypoint_shutdown_timeout=entrypoint_shutdown_timeout,
             session_end_timeout=session_end_timeout,
             ping_interval=ping_interval,
             high_ping_threshold=high_ping_threshold,
@@ -137,6 +140,7 @@ def _on_join() -> None:
                     job_entrypoint_fnc=self._opts.job_entrypoint_fnc,
                     session_end_fnc=self._opts.session_end_fnc,
                     simulation_end_fnc=self._opts.simulation_end_fnc,
+                    entrypoint_shutdown_timeout=self._opts.entrypoint_shutdown_timeout,
                     session_end_timeout=self._opts.session_end_timeout,
                     user_arguments=self._user_args,
                     join_fnc=_on_join,

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

@@ -36,6 +36,7 @@ def __init__(
         num_idle_processes: int,
         initialize_timeout: float,
         close_timeout: float,
+        entrypoint_shutdown_timeout: float,
         session_end_timeout: float,
         inference_executor: inference_executor.InferenceExecutor | None,
         job_executor_type: JobExecutorType,
@@ -53,6 +54,7 @@ def __init__(
         self._session_end_fnc = session_end_fnc
         self._simulation_end_fnc = simulation_end_fnc
         self._close_timeout = close_timeout
+        self._entrypoint_shutdown_timeout = entrypoint_shutdown_timeout
         self._session_end_timeout = session_end_timeout
         self._inf_executor = inference_executor
         self._initialize_timeout = initialize_timeout
@@ -209,6 +211,7 @@ async def _proc_spawn_task(self) -> None:
                 simulation_end_fnc=self._simulation_end_fnc,
                 initialize_timeout=self._initialize_timeout,
                 close_timeout=self._close_timeout,
+                entrypoint_shutdown_timeout=self._entrypoint_shutdown_timeout,
                 session_end_timeout=self._session_end_timeout,
                 inference_executor=self._inf_executor,
                 ping_interval=2.5,
@@ -224,6 +227,7 @@ async def _proc_spawn_task(self) -> None:
                 simulation_end_fnc=self._simulation_end_fnc,
                 initialize_timeout=self._initialize_timeout,
                 close_timeout=self._close_timeout,
+                entrypoint_shutdown_timeout=self._entrypoint_shutdown_timeout,
                 session_end_timeout=self._session_end_timeout,
                 inference_executor=self._inf_executor,
                 mp_ctx=self._mp_ctx,

livekit-agents/livekit/agents/worker.py

@@ -270,6 +270,8 @@ class ServerOptions:
     When set, the PROMETHEUS_MULTIPROC_DIR environment variable will be configured automatically.
     When None (default), multiprocess mode is disabled and only main process metrics are collected.
     Users can also set PROMETHEUS_MULTIPROC_DIR environment variable directly before starting the worker."""
+    entrypoint_shutdown_timeout: float = 15.0
+    """Maximum time to wait for a job entrypoint to exit before cancelling it."""
 
     def __post_init__(self) -> None:
         self.log_level = _validate_and_normalize_log_level(self.log_level)
@@ -310,6 +312,7 @@ def __init__(
         drain_timeout: int = DRAIN_TIMEOUT,
         num_idle_processes: int | ServerEnvOption[int] = _default_num_idle_processes,
         shutdown_process_timeout: float = 10.0,
+        entrypoint_shutdown_timeout: float = 15.0,
         session_end_timeout: float = 300.0,
         initialize_process_timeout: float = 10.0,
         permissions: WorkerPermissions = _default_permissions,
@@ -346,6 +349,7 @@ def __init__(
         self._drain_timeout = drain_timeout
         self._num_idle_processes = num_idle_processes
         self._shutdown_process_timeout = shutdown_process_timeout
+        self._entrypoint_shutdown_timeout = entrypoint_shutdown_timeout
         self._session_end_timeout = session_end_timeout
         self._initialize_process_timeout = initialize_process_timeout
         self._permissions = permissions
@@ -421,6 +425,7 @@ def from_server_options(cls, options: ServerOptions) -> AgentServer:
             drain_timeout=options.drain_timeout,
             num_idle_processes=options.num_idle_processes,
             shutdown_process_timeout=options.shutdown_process_timeout,
+            entrypoint_shutdown_timeout=options.entrypoint_shutdown_timeout,
             session_end_timeout=options.session_end_timeout,
             initialize_process_timeout=options.initialize_process_timeout,
             permissions=options.permissions,
@@ -628,6 +633,7 @@ async def run(self, *, devmode: bool = False, unregistered: bool = False) -> Non
                 mp_ctx=self._mp_ctx,
                 initialize_timeout=self._initialize_process_timeout,
                 close_timeout=self._shutdown_process_timeout,
+                entrypoint_shutdown_timeout=self._entrypoint_shutdown_timeout,
                 session_end_timeout=self._session_end_timeout,
                 memory_warn_mb=self._job_memory_warn_mb,
                 memory_limit_mb=self._job_memory_limit_mb,
@@ -868,6 +874,7 @@ def update_options(
         drain_timeout: NotGivenOr[int] = NOT_GIVEN,
         num_idle_processes: NotGivenOr[int] = NOT_GIVEN,
         shutdown_process_timeout: NotGivenOr[float] = NOT_GIVEN,
+        entrypoint_shutdown_timeout: NotGivenOr[float] = NOT_GIVEN,
         session_end_timeout: NotGivenOr[float] = NOT_GIVEN,
         initialize_process_timeout: NotGivenOr[float] = NOT_GIVEN,
     ) -> None:
@@ -907,6 +914,9 @@ def update_options(
         if is_given(shutdown_process_timeout):
             self._shutdown_process_timeout = shutdown_process_timeout
 
+        if is_given(entrypoint_shutdown_timeout):
+            self._entrypoint_shutdown_timeout = entrypoint_shutdown_timeout
+
         if is_given(session_end_timeout):
             self._session_end_timeout = session_end_timeout
 

tests/test_ipc.py

@@ -287,6 +287,7 @@ async def test_proc_pool():
         job_executor_type=job.JobExecutorType.PROCESS,
         initialize_timeout=20.0,
         close_timeout=20.0,
+        entrypoint_shutdown_timeout=15.0,
         session_end_timeout=300.0,
         inference_executor=None,
         memory_warn_mb=0,
@@ -370,6 +371,7 @@ async def test_slow_initialization():
         num_idle_processes=num_idle_processes,
         initialize_timeout=1.0,
         close_timeout=20.0,
+        entrypoint_shutdown_timeout=15.0,
         session_end_timeout=300.0,
         inference_executor=None,
         memory_warn_mb=0,
@@ -437,6 +439,7 @@ async def test_proc_pool_launch_job_raises_when_all_spawns_fail():
         # generous so initialize() fails via the init fn raising, not a spawn-racing timeout
         initialize_timeout=10.0,
         close_timeout=20.0,
+        entrypoint_shutdown_timeout=15.0,
         session_end_timeout=300.0,
         inference_executor=None,
         memory_warn_mb=0,
@@ -463,6 +466,7 @@ def _create_proc(
     mp_ctx: BaseContext,
     initialize_timeout: float = 20.0,
     job_entrypoint_fnc: Callable[[JobContext], object] = _job_entrypoint,
+    entrypoint_shutdown_timeout: float = 15.0,
 ) -> tuple[ipc.job_proc_executor.ProcJobExecutor, _StartArgs]:
     start_args = _new_start_args(mp_ctx)
     loop = asyncio.get_running_loop()
@@ -474,6 +478,7 @@ def _create_proc(
         initialize_timeout=initialize_timeout,
         close_timeout=close_timeout,
         session_end_timeout=300.0,
+        entrypoint_shutdown_timeout=entrypoint_shutdown_timeout,
         memory_warn_mb=0,
         memory_limit_mb=0,
         ping_interval=2.5,
@@ -521,6 +526,27 @@ async def test_shutdown_no_job():
     assert start_args.shutdown_counter.value == 0, "shutdown_cb isn't called when there is no job"
 
 
+async def test_job_entrypoint_shutdown_timeout():
+    mp_ctx = mp.get_context("spawn")
+    proc, start_args = _create_proc(
+        close_timeout=5.0,
+        mp_ctx=mp_ctx,
+        entrypoint_shutdown_timeout=0.0,
+    )
+    start_args.entrypoint_simulate_work_time = 30.0
+
+    await proc.start()
+    await proc.initialize()
+    await proc.launch_job(_generate_fake_job())
+    await _poll_until(lambda: start_args.entrypoint_counter.value >= 1)
+
+    await proc.aclose()
+
+    assert proc.exitcode == 0, "process should have exited cleanly"
+    assert not proc.killed
+    assert start_args.shutdown_counter.value == 1
+
+
 async def test_job_slow_shutdown():
     mp_ctx = mp.get_context("spawn")
     proc, start_args = _create_proc(close_timeout=0.3, mp_ctx=mp_ctx)

tests/test_worker_config.py

@@ -16,7 +16,7 @@
 
 from livekit.agents.cli import proto
 from livekit.agents.cli.cli import _run_worker
-from livekit.agents.worker import AgentServer
+from livekit.agents.worker import AgentServer, ServerOptions
 
 pytestmark = pytest.mark.unit
 
@@ -140,6 +140,37 @@ def test_empty_when_neither(self):
             assert config["agent_name"] == ""
 
 
+class TestEntrypointShutdownTimeout:
+    @staticmethod
+    async def _entrypoint(_ctx):
+        pass
+
+    def test_server_options_default_is_preserved(self):
+        options = ServerOptions(entrypoint_fnc=self._entrypoint)
+
+        server = AgentServer.from_server_options(options)
+
+        assert options.entrypoint_shutdown_timeout == 15.0
+        assert server._entrypoint_shutdown_timeout == 15.0
+
+    def test_server_options_zero_is_forwarded(self):
+        options = ServerOptions(
+            entrypoint_fnc=self._entrypoint,
+            entrypoint_shutdown_timeout=0.0,
+        )
+
+        server = AgentServer.from_server_options(options)
+
+        assert server._entrypoint_shutdown_timeout == 0.0
+
+    def test_update_options_sets_timeout(self):
+        server = _TestableServer()
+
+        server.update_options(entrypoint_shutdown_timeout=0.5)
+
+        assert server._entrypoint_shutdown_timeout == 0.5
+
+
 class TestFullPrecedenceChain:
     def test_everything_set(self):
         with patch.dict(