[Bugfix] Keep the frontend profiler on one thread, export off it

tests/v1/engine/test_async_llm_profiler.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""The AsyncLLM frontend profiler has to start and stop on one thread.
+
+Kineto binds its client to the thread that starts the profiler. Starting it from
+a worker thread silently records nothing, and stopping it from a different
+thread than it was started on crashes the process, so the lifecycle must stay on
+the event loop thread. Exporting the trace has no such constraint and must not
+stay there: it serializes every recorded event.
+"""
+
+import gzip
+import json
+import threading
+from functools import partial
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+import torch
+
+from vllm.v1.engine.async_llm import AsyncLLM
+
+NUM_SPANS = 64
+
+
+def make_engine(profiler, trace_handler=None):
+    engine = MagicMock(spec=AsyncLLM)
+    engine.engine_core = MagicMock()
+    engine.engine_core.profile_async = AsyncMock()
+    engine.profiler = profiler
+    engine.frontend_trace_handler = trace_handler
+    # The mock would otherwise swallow the helpers these tests are about.
+    engine._start_frontend_profiler = partial(AsyncLLM._start_frontend_profiler, engine)
+    engine._stop_frontend_profiler = partial(AsyncLLM._stop_frontend_profiler, engine)
+    return engine
+
+
+def read_spans(trace_dir):
+    traces = list(trace_dir.glob("*.pt.trace.json*"))
+    assert len(traces) == 1, f"expected one trace, got {traces}"
+    opener = gzip.open if traces[0].suffix == ".gz" else open
+    with opener(traces[0], "rt") as f:
+        events = json.load(f)["traceEvents"]
+    return [e for e in events if e.get("name") == "frontend_span"]
+
+
+@pytest.mark.asyncio
+async def test_frontend_profiler_records_what_happened(tmp_path):
+    """The exported trace must actually contain the frontend's spans."""
+    profiler = torch.profiler.profile(
+        activities=[torch.profiler.ProfilerActivity.CPU],
+    )
+    handler = torch.profiler.tensorboard_trace_handler(
+        str(tmp_path), worker_name="async_llm", use_gzip=True
+    )
+    engine = make_engine(profiler, handler)
+
+    await AsyncLLM.start_profile(engine, "trace")
+    for _ in range(NUM_SPANS):
+        with torch.profiler.record_function("frontend_span"):
+            pass
+    await AsyncLLM.stop_profile(engine)
+
+    assert len(read_spans(tmp_path)) == NUM_SPANS
+    engine.engine_core.profile_async.assert_any_await(True, "trace")
+    engine.engine_core.profile_async.assert_any_await(False)
+
+
+@pytest.mark.asyncio
+async def test_profiler_lifecycle_stays_on_the_calling_thread():
+    """Both calls must run on the caller's thread, not a worker thread."""
+    profiler = MagicMock()
+    threads: list[int] = []
+    profiler.start.side_effect = lambda: threads.append(threading.get_ident())
+    profiler.stop.side_effect = lambda: threads.append(threading.get_ident())
+    engine = make_engine(profiler, MagicMock())
+
+    await AsyncLLM.start_profile(engine)
+    await AsyncLLM.stop_profile(engine)
+
+    assert threads == [threading.get_ident(), threading.get_ident()]
+
+
+@pytest.mark.asyncio
+async def test_trace_export_does_not_run_on_the_event_loop():
+    """Exporting blocks for as long as it takes to serialize the trace.
+
+    On the event loop that stalls every other request for its whole duration,
+    so it has to happen off-thread even though stop() cannot.
+    """
+    profiler = MagicMock()
+    ran_on: dict[str, object] = {}
+    profiler.stop.side_effect = lambda: ran_on.update(stop=threading.get_ident())
+
+    def handler(prof):
+        ran_on.update(export=threading.get_ident(), exported=prof)
+
+    engine = make_engine(profiler, handler)
+
+    await AsyncLLM.stop_profile(engine)
+
+    assert ran_on["stop"] == threading.get_ident()
+    assert ran_on["export"] != threading.get_ident()
+    assert ran_on["exported"] is profiler
+
+
+@pytest.mark.asyncio
+async def test_engine_core_is_told_even_when_the_trace_export_fails():
+    """A frontend trace-export failure must not leave EngineCore profiling."""
+    profiler = MagicMock()
+
+    def handler(prof):
+        raise OSError("no space left on device")
+
+    engine = make_engine(profiler, handler)
+
+    with pytest.raises(OSError):
+        await AsyncLLM.stop_profile(engine)
+
+    engine.engine_core.profile_async.assert_any_await(False)
+
+
+@pytest.mark.asyncio
+async def test_profile_endpoints_still_work_without_a_frontend_profiler():
+    engine = make_engine(None)
+
+    await AsyncLLM.start_profile(engine)
+    await AsyncLLM.stop_profile(engine)
+
+    engine.engine_core.profile_async.assert_any_await(True, None)
+    engine.engine_core.profile_async.assert_any_await(False)

vllm/v1/engine/async_llm.py

@@ -188,19 +188,22 @@ def __init__(
                 profiler_dir,
             )
             worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm"
+            # Held rather than passed as on_trace_ready so that stop() only
+            # ends the trace and the export can run off the event loop.
+            self.frontend_trace_handler = torch.profiler.tensorboard_trace_handler(
+                profiler_dir,
+                worker_name=worker_name,
+                use_gzip=vllm_config.profiler_config.torch_profiler_use_gzip,
+            )
             self.profiler = torch.profiler.profile(
                 activities=[
                     torch.profiler.ProfilerActivity.CPU,
                 ],
                 with_stack=vllm_config.profiler_config.torch_profiler_with_stack,
-                on_trace_ready=torch.profiler.tensorboard_trace_handler(
-                    profiler_dir,
-                    worker_name=worker_name,
-                    use_gzip=vllm_config.profiler_config.torch_profiler_use_gzip,
-                ),
             )
         else:
             self.profiler = None
+            self.frontend_trace_handler = None
 
     @classmethod
     def from_vllm_config(
@@ -943,16 +946,40 @@ async def check_health(self) -> None:
         if self.errored:
             raise self.dead_error
 
+    async def _start_frontend_profiler(self) -> None:
+        """Start the frontend profiler on the event loop thread.
+
+        Kineto binds its client to the thread that starts the profiler, and
+        records only that thread's work. Started from a worker thread it
+        captures nothing, and a stop() landing on a different pool thread
+        crashes the process. See
+        https://github.com/vllm-project/vllm/issues/39603
+
+        A coroutine so it can join the gather below, keeping the EngineCore
+        request concurrent with it.
+        """
+        self.profiler.start()
+
+    async def _stop_frontend_profiler(self) -> None:
+        """Stop the frontend profiler and export its trace off the loop.
+
+        stop() has to stay on the thread that started the profiler, but the
+        export does not, and the export is the expensive half: it serializes
+        and compresses every recorded event.
+        """
+        self.profiler.stop()
+        await asyncio.to_thread(self.frontend_trace_handler, self.profiler)
+
     async def start_profile(self, profile_prefix: str | None = None) -> None:
         coros = [self.engine_core.profile_async(True, profile_prefix)]
         if self.profiler is not None:
-            coros.append(asyncio.to_thread(self.profiler.start))
+            coros.append(self._start_frontend_profiler())
         await asyncio.gather(*coros)
 
     async def stop_profile(self) -> None:
         coros = [self.engine_core.profile_async(False)]
         if self.profiler is not None:
-            coros.append(asyncio.to_thread(self.profiler.stop))
+            coros.append(self._stop_frontend_profiler())
         await asyncio.gather(*coros)
 
     async def reset_mm_cache(self) -> None: