diff --git a/tests/worker/test_worker_base.py b/tests/worker/test_worker_base.py index 18131a31c..4570e6eda 100644 --- a/tests/worker/test_worker_base.py +++ b/tests/worker/test_worker_base.py @@ -239,3 +239,25 @@ def test_wake_up_forwards_tags_to_allocator(monkeypatch): assert OmniGPUWorkerBase.wake_up(worker, tags=["weights"]) is True assert calls["tags"] == ["weights"] + + +def test_cuda_profiler_is_wired_for_omni_worker(monkeypatch): + import vllm.profiler.wrapper as profiler_wrapper + + def fake_gpu_worker_init(worker, *args, **kwargs): + worker.vllm_config = SimpleNamespace(profiler_config=SimpleNamespace(profiler="cuda")) + worker.local_rank = 0 + worker.rank = 0 + worker.profiler = None + + class FakeCudaProfiler: + def __init__(self, profiler_config): + self.profiler_config = profiler_config + + monkeypatch.setattr(base.GPUWorker, "__init__", fake_gpu_worker_init) + monkeypatch.setattr(profiler_wrapper, "CudaProfilerWrapper", FakeCudaProfiler) + + worker = OmniGPUWorkerBase() + + assert isinstance(worker.profiler, FakeCudaProfiler) + assert worker.profiler.profiler_config.profiler == "cuda" diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py index 2a7e9c876..01a08bd72 100644 --- a/vllm_omni/worker/base.py +++ b/vllm_omni/worker/base.py @@ -58,6 +58,13 @@ class OmniGPUWorkerBase(GPUWorker): worker_name=worker_name, local_rank=self.local_rank, ) + elif profiler_config and profiler_config.profiler == "cuda": + # OmniGPUWorkerBase replaces vLLM's torch profiler above; keep the + # native CUDA profiler wrapper available for Nsight Systems capture + # through the /start_profile and /stop_profile endpoints as well. + from vllm.profiler.wrapper import CudaProfilerWrapper + + self.profiler = CudaProfilerWrapper(profiler_config) def profile(self, is_start: bool = True, profile_prefix: str | None = None): """Override to set trace filename before starting the profiler.