[Profiler] Add Nsight Systems support for serving (#1098)

Signed-off-by: Jinheng Li <ahengljh@gmail.com>
Signed-off-by: Canlin Guo <961750412@qq.com>
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Co-authored-by: Canlin Guo <961750412@qq.com>

tests/diffusion/test_diffusion_worker_cuda_profiler.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import MagicMock
+
+import pytest
+from pytest_mock import MockerFixture
+
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture
+def mock_od_config(mocker: MockerFixture):
+    """Create a mock OmniDiffusionConfig with a CUDA profiler backend."""
+    config = mocker.Mock()
+    config.profiler_config = mocker.Mock()
+    config.profiler_config.profiler = "cuda"
+    config.diffusion_load_format = "default"
+    return config
+
+
+@pytest.fixture
+def mock_diffusion_worker_dependencies(mocker: MockerFixture):
+    """Patch heavy worker dependencies for focused profiler tests."""
+    mocker.patch.object(DiffusionWorker, "init_device")
+    mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.DiffusionModelRunner")
+
+
+class TestDiffusionWorkerCudaProfiler:
+    def test_creates_cuda_profiler_wrapper(
+        self,
+        mocker: MockerFixture,
+        mock_od_config,
+        mock_diffusion_worker_dependencies,
+    ):
+        fake_profiler = mocker.Mock()
+        cuda_profiler = mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.CudaProfilerWrapper",
+            return_value=fake_profiler,
+        )
+        create_omni_profiler = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.create_omni_profiler")
+
+        worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config, skip_load_model=True)
+
+        cuda_profiler.assert_called_once_with(mock_od_config.profiler_config)
+        create_omni_profiler.assert_not_called()
+        assert worker.profiler is fake_profiler
+
+    def test_profile_start_stop_delegates_to_cuda_profiler(
+        self,
+        mocker: MockerFixture,
+        mock_od_config,
+        mock_diffusion_worker_dependencies,
+    ):
+        fake_profiler = mocker.Mock()
+        fake_profiler.start = MagicMock()
+        fake_profiler.stop = MagicMock()
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.CudaProfilerWrapper",
+            return_value=fake_profiler,
+        )
+
+        worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config, skip_load_model=True)
+
+        assert worker.profile(is_start=True) is None
+        assert worker.profile(is_start=False) is None
+
+        fake_profiler.start.assert_called_once_with()
+        fake_profiler.stop.assert_called_once_with()
+
+    def test_returns_none_when_profiler_config_is_missing(
+        self,
+        mocker: MockerFixture,
+        mock_od_config,
+        mock_diffusion_worker_dependencies,
+    ):
+        mock_od_config.profiler_config = None
+        cuda_profiler = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.CudaProfilerWrapper")
+        create_omni_profiler = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.create_omni_profiler")
+
+        worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config, skip_load_model=True)
+
+        cuda_profiler.assert_not_called()
+        create_omni_profiler.assert_not_called()
+        assert worker.profiler is None
+
+    def test_cuda_backend_does_not_use_torch_profiler_factory(
+        self,
+        mocker: MockerFixture,
+        mock_od_config,
+        mock_diffusion_worker_dependencies,
+    ):
+        mocker.patch(
+            "vllm_omni.diffusion.worker.diffusion_worker.CudaProfilerWrapper",
+            return_value=mocker.Mock(),
+        )
+        create_omni_profiler = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.create_omni_profiler")
+
+        DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config, skip_load_model=True)
+
+        create_omni_profiler.assert_not_called()