[Profiler] Export torch profiler trace off the worker's main thread

tests/v1/worker/test_gpu_profiler.py

@@ -1,12 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import threading
+import time
 from unittest.mock import MagicMock
 
 import pytest
 
 from vllm.config import CUDAGraphMode, ProfilerConfig
 from vllm.config.profiler import _is_uri_path
-from vllm.profiler.wrapper import WorkerProfiler
+from vllm.profiler.wrapper import TorchProfilerWrapper, WorkerProfiler
 from vllm.v1.core.sched.output import CachedRequestData
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 from vllm.v1.worker.gpu_worker import Worker
@@ -282,6 +284,61 @@ def test_detailed_format_mixed(self):
         )
 
 
+class TestTorchProfilerWrapperAsyncExport:
+    """Tests that trace export runs off the calling thread instead of
+    blocking stop()/step() (see TorchProfilerWrapper._async_trace_ready).
+    Uses CPU-only activities so these run without a GPU."""
+
+    def _make_wrapper(self, tmp_path, on_trace_ready):
+        config = ProfilerConfig(
+            profiler="torch",
+            torch_profiler_dir=str(tmp_path),
+            torch_profiler_dump_cuda_time_total=False,
+        )
+        return TorchProfilerWrapper(
+            profiler_config=config,
+            worker_name="test-worker",
+            local_rank=0,
+            activities=["CPU"],
+            on_trace_ready=on_trace_ready,
+        )
+
+    def test_stop_does_not_block_on_slow_export(self, tmp_path):
+        calling_thread = threading.current_thread()
+        handler_thread: dict[str, threading.Thread] = {}
+        handler_done = threading.Event()
+
+        def slow_handler(prof):
+            handler_thread["thread"] = threading.current_thread()
+            time.sleep(0.3)
+            handler_done.set()
+
+        wrapper = self._make_wrapper(tmp_path, slow_handler)
+        wrapper.start()
+
+        start = time.perf_counter()
+        wrapper.stop()
+        elapsed = time.perf_counter() - start
+
+        assert elapsed < 0.2, "stop() should not block on the trace export"
+        assert handler_done.wait(timeout=2.0), "background export never completed"
+        assert handler_thread["thread"] is not calling_thread
+
+    def test_export_errors_are_caught_not_raised(self, tmp_path):
+        handler_done = threading.Event()
+
+        def failing_handler(prof):
+            handler_done.set()
+            raise RuntimeError("boom")
+
+        wrapper = self._make_wrapper(tmp_path, failing_handler)
+        wrapper.start()
+
+        wrapper.stop()  # must not raise even though the handler will
+
+        assert handler_done.wait(timeout=2.0)
+
+
 def test_profiler_entered_during_capture():
     """Profiler is used as a context manager in _warmup_and_capture,
     confirming it is active during the actual graph capture run."""

vllm/profiler/wrapper.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import json
+import threading
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from contextlib import nullcontext
@@ -198,6 +199,36 @@ def __init__(
                 use_gzip=profiler_config.torch_profiler_use_gzip,
             )
 
+        def _async_trace_ready(prof: torch.profiler.profile) -> None:
+            """Runs the trace export (JSON serialization + optional gzip) on a
+            background thread instead of inline on whatever called
+            profiler.step()/stop() -- normally this worker's own request-
+            handling loop.
+
+            By the time this fires, Kineto has already stopped collecting and
+            holds a complete, immutable snapshot of the trace in memory (the
+            CUDA-side stop happens synchronously just before this callback, as
+            part of the same profiler.stop()/step() call) -- so none of this
+            touches CUDA or the model-execution stream anymore, and it's safe
+            to run concurrently with whatever this worker does next.
+
+            This matters because export_chrome_trace() with gzip enabled
+            writes the full uncompressed JSON to a temp file, rereads it, and
+            recompresses it with single-threaded stdlib gzip -- for a
+            multi-hundred-MB-to-GB trace this can take tens of seconds, during
+            which (without this) the worker cannot schedule its next step.
+            """
+
+            def run() -> None:
+                try:
+                    trace_handler(prof)
+                except Exception as e:
+                    logger.warning("Failed to export profiler trace: %s", e)
+
+            threading.Thread(
+                target=run, name="vllm-profiler-trace-export", daemon=True
+            ).start()
+
         self.dump_cpu_time_total = "CPU" in activities and len(activities) == 1
 
         # Create profiler schedule if warmup or wait iterations are configured
@@ -225,7 +256,7 @@ def __init__(
             profile_memory=profiler_config.torch_profiler_with_memory,
             with_stack=profiler_config.torch_profiler_with_stack,
             with_flops=profiler_config.torch_profiler_with_flops,
-            on_trace_ready=trace_handler,
+            on_trace_ready=_async_trace_ready,
         )
 
         # Track if we're using a schedule (need to call step())