feat(telemetry): trace dispatch, startup, and shutdown

livekit-agents/livekit/agents/ipc/_preload.py

@@ -0,0 +1,81 @@
+"""Side-effect module: the framework's own one-time warm-up, run once per process image.
+
+Each step is lazy one-time work that would otherwise happen inside the first job, on the
+event loop, and show up as a 100-500 ms ``event_loop_blocked`` at session start that no user
+code caused. Everything here is fork-safe (imports, a ``dlopen``, model weights, a cached
+SSL context): no threads and no event loop are created.
+
+Where it runs decides how often it costs:
+
+- with the ``forkserver`` start method (Linux) the worker lists this module in
+  ``set_forkserver_preload``, ahead of ``_preload_freeze``, so it runs once in the forkserver
+  and every job process inherits the result copy-on-write;
+- with ``spawn`` (macOS, Windows) each job process imports it while it warms up, before any
+  job is assigned.
+
+The job process always imports it: under a forkserver the module is already in
+``sys.modules`` and the import is a no-op, so there is no start-method check anywhere.
+
+Failures are logged at debug level only: the first real use reports a proper error.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from typing import Any
+
+from ..log import logger
+
+
+def _step(name: str, fnc: Callable[[], Any]) -> None:
+    started = time.perf_counter()
+    try:
+        fnc()
+    except Exception:
+        logger.debug("could not preload %s", name, exc_info=True)
+        return
+    logger.debug("preloaded %s", name, extra={"elapsed": round(time.perf_counter() - started, 3)})
+
+
+def _av() -> None:
+    import av  # noqa: F401
+
+
+def _local_inference_models() -> None:
+    # the VAD and the turn detector's local end-of-turn model: constructing them later in a
+    # job is free once these singletons exist (~25 ms of GIL-held CPU otherwise)
+    import livekit.local_inference as li
+
+    li.init_vad()
+    li.init_eot()
+
+
+def _rtc_native_library() -> None:
+    # the dlopen (~150-350 ms). The runtime itself (FfiClient.instance) starts threads, so it
+    # stays per process; with the library already mapped it takes a few milliseconds
+    from livekit.rtc._ffi_client import get_ffi_lib
+
+    get_ffi_lib()
+
+
+def _openai_resources() -> None:
+    # the openai SDK, which livekit.agents.inference is built on, imports its whole resources
+    # tree on the first client attribute access (~300-550 ms)
+    import openai.resources  # noqa: F401
+
+
+def _httpx_client() -> None:
+    # the first AsyncClient in a process image pays ~40 ms (the SSL context from the CA bundle
+    # among other lazy setup); later ones take a few milliseconds. The inference LLM, STT and
+    # TTS each build one
+    import httpx
+
+    httpx.AsyncClient()
+
+
+_step("av", _av)
+_step("the local inference models", _local_inference_models)
+_step("the livekit-rtc native library", _rtc_native_library)
+_step("the openai SDK resources", _openai_resources)
+_step("the httpx client", _httpx_client)

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

@@ -18,17 +18,18 @@
 import contextlib
 import contextvars
 import socket
+import time
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 from typing import Any, cast
 
-from opentelemetry import trace
+from opentelemetry import context as otel_context, trace
 
 from livekit import rtc
 
-from ..job import JobContext, JobExecutorType, JobProcess, _JobContextVar
+from ..job import JobContext, JobExecutorType, JobProcess, RunningJobInfo, _JobContextVar
 from ..log import _add_global_log_fields, logger
-from ..telemetry import loop_monitor, trace_types, tracer
+from ..telemetry import loop_monitor, trace_types, tracer, utils as trace_utils
 from ..utils import aio, http_context, log_exceptions, shortuuid
 from .channel import Message
 from .inference_executor import InferenceExecutor
@@ -223,12 +224,17 @@ def initialize(self, init_req: InitializeRequest, client: _ProcClient) -> None:
             user_arguments=self._user_arguments,
             http_proxy=init_req.http_proxy or None,
         )
+        # the framework's warm-up, once per process image: a no-op under a forkserver, which
+        # already imported it (see the module docstring)
+        from . import _preload  # noqa: F401
+
         self._initialize_process_fnc(self._job_proc)
 
     @log_exceptions(logger=logger)
     async def entrypoint(self, cch: aio.ChanReceiver[Message]) -> None:
         self._exit_proc_flag = asyncio.Event()
         self._shutdown_fut: asyncio.Future[_ShutdownInfo] = asyncio.Future()
+        self._entrypoint_span_context: otel_context.Context | None = None
 
         @log_exceptions(logger=logger)
         async def _read_ipc_task() -> None:
@@ -322,22 +328,36 @@ def _exit_proc_cb(_: asyncio.Task[None]) -> None:
     async def _run_job_task(self) -> None:
         self._job_ctx._on_setup()
         self._job_ctx._start_log_buffering()
+        # the trace pipeline must exist before the job's first span, or that span (the job's
+        # root) is a non-recording stub and everything under it starts its own trace
+        self._job_ctx._prepare_telemetry()
 
         job_ctx_token = _JobContextVar.set(self._job_ctx)
         http_context._new_session_ctx()
 
-        @tracer.start_as_current_span("job_entrypoint")
+        # the job's root span: from the availability request to the end of the shutdown
+        # sequence, so the whole job reads as one trace (the user entrypoint returning is an
+        # event on it, most entrypoints return right after session.start())
+        job_span = _start_job_span(self._job_ctx)
+        self._entrypoint_span_context = trace.set_span_in_context(job_span)
+
         async def _traceable_entrypoint(job_ctx: JobContext) -> None:
-            job = job_ctx.job
-            current_span = trace.get_current_span()
-            current_span.set_attribute(trace_types.ATTR_JOB_ID, job.id)
-            current_span.set_attribute(trace_types.ATTR_AGENT_NAME, job.agent_name)
-            current_span.set_attribute(trace_types.ATTR_ROOM_NAME, job.room.name)
-            # blocked-loop reports emitted from the heartbeat need this job's context (for
-            # attribution) and the job_entrypoint span (as the parent when no session is up)
-            if (monitor := loop_monitor.get_monitor(asyncio.get_running_loop())) is not None:
-                monitor.set_report_context(contextvars.copy_context())
-            await self._job_entrypoint_fnc(job_ctx)
+            with tracer.use_span(
+                job_span, end_on_exit=False, record_exception=False, set_status_on_exception=False
+            ):
+                # blocked-loop reports emitted from the heartbeat need this job's context (for
+                # attribution) and the job span (as the parent when no session is up)
+                if (monitor := loop_monitor.get_monitor(asyncio.get_running_loop())) is not None:
+                    monitor.set_report_context(contextvars.copy_context())
+                try:
+                    await self._job_entrypoint_fnc(job_ctx)
+                except asyncio.CancelledError:
+                    job_span.add_event("entrypoint_cancelled")
+                    raise
+                except Exception as e:
+                    trace_utils.record_exception(job_span, e)
+                    raise
+                job_span.add_event("entrypoint_returned")
 
         job_entry_task = asyncio.create_task(
             _traceable_entrypoint(self._job_ctx), name="job_user_entrypoint"
@@ -382,6 +402,34 @@ def _on_entry_done(t: asyncio.Task[Any]) -> None:
 
         shutdown_info = await self._shutdown_fut
 
+        try:
+            # a child of job_entrypoint, like the session was: the job's trace tells the whole
+            # story from dispatch to teardown, and a viewer keyed to agent_session zooms out
+            with tracer.start_as_current_span(
+                "job_shutdown",
+                context=self._entrypoint_span_context,
+                attributes={
+                    trace_types.ATTR_SHUTDOWN_REASON: shutdown_info.reason,
+                    trace_types.ATTR_SHUTDOWN_USER_INITIATED: shutdown_info.user_initiated,
+                },
+            ):
+                await self._shutdown_job(job_entry_task, shutdown_info)
+        finally:
+            # whatever the shutdown raised (a session or room close failing), the job still
+            # ends: the root span goes out with the job's telemetry, the temp dir is removed
+            # and the per-job telemetry state is released (thread workers reuse the process)
+            job_span.end()
+
+            if tasks := self._job_ctx._pending_tasks:
+                await aio.cancel_and_wait(*tasks)
+
+            await self._job_ctx._on_cleanup()
+            await http_context._close_http_ctx()
+            _JobContextVar.reset(job_ctx_token)
+
+    async def _shutdown_job(
+        self, job_entry_task: asyncio.Task[None], shutdown_info: _ShutdownInfo
+    ) -> None:
         # wait for the entrypoint to finish, cancel if it takes too long
         if not job_entry_task.done():
             try:
@@ -405,20 +453,28 @@ def _on_entry_done(t: asyncio.Task[Any]) -> None:
                 )
 
         if self._session_end_fnc:
+            with tracer.start_as_current_span(
+                "on_session_end",
+                attributes={trace_types.ATTR_CALLBACK_NAME: _callback_name(self._session_end_fnc)},
+            ) as span:
+                try:
+                    await asyncio.wait_for(
+                        self._session_end_fnc(self._job_ctx),
+                        timeout=self._session_end_timeout,
+                    )
+                except asyncio.TimeoutError as e:
+                    trace_utils.record_exception(span, e)
+                    logger.error("on_session_end timed out after %ds", self._session_end_timeout)
+                except Exception as e:
+                    trace_utils.record_exception(span, e)
+                    logger.exception("error while executing the on_session_end callback")
+
+        with tracer.start_as_current_span("session_end_upload") as span:
             try:
-                await asyncio.wait_for(
-                    self._session_end_fnc(self._job_ctx),
-                    timeout=self._session_end_timeout,
-                )
-            except asyncio.TimeoutError:
-                logger.error("on_session_end timed out after %ds", self._session_end_timeout)
-            except Exception:
-                logger.exception("error while executing the on_session_end callback")
-
-        try:
-            await self._job_ctx._on_session_end()
-        except Exception:
-            logger.exception("error in job_ctx._on_session_end")
+                await self._job_ctx._on_session_end()
+            except Exception as e:
+                trace_utils.record_exception(span, e)
+                logger.exception("error in job_ctx._on_session_end")
 
         await self._client.send(ShuttingDown())
 
@@ -427,27 +483,117 @@ def _on_entry_done(t: asyncio.Task[Any]) -> None:
             extra={"reason": shutdown_info.reason, "user_initiated": shutdown_info.user_initiated},
         )
         await self._client.send(Exiting(reason=shutdown_info.reason))
-        await self._room.disconnect()
+        with tracer.start_as_current_span("room_disconnect"):
+            await self._room.disconnect()
+
+        async def _traced_shutdown_callback(
+            callback: Callable[[str], Awaitable[None]],
+        ) -> None:
+            if _is_framework_callback(callback):
+                # the session's own close hook, already covered by session_close; a span here
+                # would read as a second user callback
+                await callback(shutdown_info.reason)
+                return
+            # a hung callback here is why jobs hit the supervisor's shutdown deadline
+            with tracer.start_as_current_span(
+                "shutdown_callback",
+                attributes={trace_types.ATTR_CALLBACK_NAME: _callback_name(callback)},
+            ):
+                await callback(shutdown_info.reason)
 
         try:
             shutdown_tasks = []
             for callback in self._job_ctx._shutdown_callbacks:
                 shutdown_tasks.append(
                     asyncio.create_task(
-                        callback(shutdown_info.reason), name="job_shutdown_callback"
+                        _traced_shutdown_callback(callback), name="job_shutdown_callback"
                     )
                 )
 
             await asyncio.gather(*shutdown_tasks)
         except Exception:
             logger.exception("error while shutting down the job")
 
-        if tasks := self._job_ctx._pending_tasks:
-            await aio.cancel_and_wait(*tasks)
 
-        self._job_ctx._on_cleanup()
-        await http_context._close_http_ctx()
-        _JobContextVar.reset(job_ctx_token)
+def _is_framework_callback(fnc: Any) -> bool:
+    """Registered by livekit-agents itself rather than by the user's entrypoint."""
+    module = getattr(fnc, "__module__", None) or ""
+    return module == "livekit.agents" or module.startswith("livekit.agents.")
+
+
+def _callback_name(fnc: Any) -> str:
+    return str(getattr(fnc, "__qualname__", None) or getattr(fnc, "__name__", None) or repr(fnc))
+
+
+def _server_timestamp_seconds(value: int) -> float:
+    """``JobState`` timestamps are int64; the server writes unix nanoseconds. Be tolerant
+    of milliseconds/seconds should that ever change."""
+    if value > 1e17:
+        return value / 1e9
+    if value > 1e11:
+        return value / 1e3
+    return float(value)
+
+
+def _start_job_span(job_ctx: JobContext) -> trace.Span:
+    """The job's root span, ``job_entrypoint``, back-dated to the availability request.
+
+    Never made current by the caller for longer than the user entrypoint runs; ended by
+    ``_run_job_task`` after the shutdown sequence, before the telemetry release."""
+    job = job_ctx.job
+    info = job_ctx._info
+    entrypoint_started_at = time.time()
+    start_time_ns = int(info.received_at * 1e9) if info.received_at else None
+    span = tracer.start_span(
+        "job_entrypoint",
+        start_time=start_time_ns,
+        attributes={
+            trace_types.ATTR_JOB_ID: job.id,
+            trace_types.ATTR_AGENT_NAME: job.agent_name,
+            trace_types.ATTR_ROOM_NAME: job.room.name,
+            trace_types.ATTR_ROOM_SID: job.room.sid,
+            trace_types.ATTR_DISPATCH_ID: job.dispatch_id,
+            trace_types.ATTR_WORKER_ID: info.worker_id,
+            trace_types.ATTR_JOB_AGENT_ID: job.state.agent_id,
+        },
+    )
+    _record_dispatch_timeline(span, info, entrypoint_started_at)
+    return span
+
+
+def _record_dispatch_timeline(
+    span: trace.Span, info: RunningJobInfo, entrypoint_started_at: float
+) -> None:
+    """Stamp the dispatch stages on ``span``: one timestamped event per stage instant, and
+    the seconds between adjacent stages as attributes (they sum to the dispatch latency).
+
+    Timestamps travel from the worker through ``StartJobRequest``; a zero means the stage
+    is unknown (simulation, console, resumed job) and is skipped rather than guessed."""
+    stages = [
+        ("job_received", info.received_at),
+        ("job_accepted", info.accepted_at),
+        ("job_assigned", info.assigned_at),
+        ("process_assigned", info.launched_at),
+        ("entrypoint_started", entrypoint_started_at),
+    ]
+    for event_name, ts in stages:
+        if ts:
+            span.add_event(event_name, timestamp=int(ts * 1e9))
+
+    def _gap(attr: str, start: float, end: float) -> None:
+        if start and end:
+            span.set_attribute(attr, max(end - start, 0.0))
+
+    _gap(trace_types.ATTR_JOB_ACCEPT_LATENCY, info.received_at, info.accepted_at)
+    _gap(trace_types.ATTR_JOB_ASSIGNMENT_LATENCY, info.accepted_at, info.assigned_at)
+    _gap(trace_types.ATTR_JOB_LAUNCH_LATENCY, info.assigned_at, info.launched_at)
+    _gap(trace_types.ATTR_JOB_ENTRYPOINT_LATENCY, info.launched_at, entrypoint_started_at)
+    _gap(trace_types.ATTR_JOB_DISPATCH_LATENCY, info.received_at, entrypoint_started_at)
+    if (server_started := info.job.state.started_at) > 0:
+        # the server's own record of the start: the first anchor for lining the agent trace
+        # up with server-side events later
+        started = _server_timestamp_seconds(server_started)
+        span.add_event("job_started_on_server", timestamp=int(started * 1e9))
 
 
 @dataclass

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

@@ -2,6 +2,7 @@
 
 import asyncio
 import math
+import time
 from collections.abc import Awaitable, Callable
 from multiprocessing.context import BaseContext
 from typing import Any, Literal
@@ -171,6 +172,8 @@ async def launch_job(self, info: RunningJobInfo) -> None:
             finally:
                 self._jobs_waiting_for_process -= 1
 
+            # dispatch timeline: a warm process is now handling this job
+            info.launched_at = time.time()
             try:
                 await proc.launch_job(info)
                 self.emit("process_job_launched", proc)