| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import os |
| 4 | from contextlib import nullcontext |
| 5 | from types import SimpleNamespace |
| 6 | from unittest.mock import MagicMock, Mock, call, patch |
| 7 | from uuid import UUID |
| 8 | |
| 9 | import pytest |
| 10 | import torch |
| 11 | from pydantic import ValidationError |
| 12 | |
| 13 | from vllm.config import ( |
| 14 | CompilationConfig, |
| 15 | CUDAGraphMode, |
| 16 | ProfilerConfig, |
| 17 | VllmConfig, |
| 18 | ) |
| 19 | from vllm.config.profiler import _is_uri_path |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.profiler.wrapper import ( |
| 22 | ProtonProfilerWrapper, |
| 23 | TorchProfilerWrapper, |
| 24 | WorkerProfiler, |
| 25 | create_worker_profiler, |
| 26 | validate_worker_profiler_config, |
| 27 | ) |
| 28 | from vllm.v1.core.sched.output import CachedRequestData |
| 29 | from vllm.v1.worker.gpu_model_runner import GPUModelRunner |
| 30 | from vllm.v1.worker.gpu_worker import Worker |
| 31 | from vllm.v1.worker.xpu_worker import XPUWorker |
| 32 | |
| 33 | |
| 34 | class ConcreteWorkerProfiler(WorkerProfiler): |
| 35 | """A basic implementation of a worker profiler for testing purposes.""" |
| 36 | |
| 37 | def __init__(self, profiler_config: ProfilerConfig): |
| 38 | self.start_call_count = 0 |
| 39 | self.stop_call_count = 0 |
| 40 | self.should_fail_start = False |
| 41 | super().__init__(profiler_config) |
| 42 | |
| 43 | def _start(self) -> None: |
| 44 | if self.should_fail_start: |
| 45 | raise RuntimeError("Simulated start failure") |
| 46 | self.start_call_count += 1 |
| 47 | |
| 48 | def _stop(self) -> None: |
| 49 | self.stop_call_count += 1 |
| 50 | |
| 51 | |
| 52 | @pytest.fixture |
| 53 | def default_profiler_config(): |
| 54 | return ProfilerConfig( |
| 55 | profiler="torch", |
| 56 | torch_profiler_dir="/tmp/mock", |
| 57 | delay_iterations=0, |
| 58 | max_iterations=0, |
| 59 | ) |
| 60 | |
| 61 | |
| 62 | def test_torch_profiler_rebuilds_one_shot_profiler_each_round(tmp_path): |
| 63 | config = ProfilerConfig( |
| 64 | profiler="torch", |
| 65 | torch_profiler_dir=str(tmp_path), |
| 66 | torch_profiler_activities=["CUDA"], |
| 67 | torch_profiler_dump_cuda_time_total=False, |
| 68 | warmup_iterations=2, |
| 69 | ) |
| 70 | profilers = [MagicMock(), MagicMock()] |
| 71 | for profiler in profilers: |
| 72 | profiler.profiler = MagicMock() |
| 73 | |
| 74 | with patch( |
| 75 | "vllm.profiler.wrapper.torch.profiler.profile", side_effect=profilers |
| 76 | ) as profile: |
| 77 | wrapper = TorchProfilerWrapper( |
| 78 | config, worker_name="worker", local_rank=0, activities=["CUDA"] |
| 79 | ) |
| 80 | profile.assert_not_called() |
| 81 | |
| 82 | wrapper.start() |
| 83 | assert not wrapper.should_annotate |
| 84 | assert not wrapper._profiler_step() |
| 85 | wrapper.stop() |
| 86 | wrapper.start() |
| 87 | assert not wrapper._profiler_step() |
| 88 | wrapper.stop() |
| 89 | |
| 90 | assert profile.call_count == 2 |
| 91 | assert profile.call_args.kwargs["activities"] == [ |
| 92 | torch.profiler.ProfilerActivity.CUDA |
| 93 | ] |
| 94 | for profiler in profilers: |
| 95 | profiler.start.assert_called_once_with() |
| 96 | profiler.step.assert_called_once_with() |
| 97 | profiler.stop.assert_called_once_with() |
| 98 | |
| 99 | |
| 100 | def test_torch_profiler_records_each_profile_round(tmp_path): |
| 101 | traces: list[torch.profiler.profile] = [] |
| 102 | wrapper = TorchProfilerWrapper( |
| 103 | ProfilerConfig( |
| 104 | profiler="torch", |
| 105 | torch_profiler_dir=str(tmp_path), |
| 106 | torch_profiler_dump_cuda_time_total=False, |
| 107 | ), |
| 108 | worker_name="worker", |
| 109 | local_rank=1, |
| 110 | activities=["CPU"], |
| 111 | on_trace_ready=traces.append, |
| 112 | ) |
| 113 | |
| 114 | for run in range(2): |
| 115 | wrapper.start() |
| 116 | with torch.profiler.record_function(f"run_{run}"): |
| 117 | pass |
| 118 | wrapper.stop() |
| 119 | |
| 120 | assert len(traces) == 2 |
| 121 | for run, trace in enumerate(traces): |
| 122 | assert { |
| 123 | event.name for event in trace.events() if event.name.startswith("run_") |
| 124 | } == {f"run_{run}"} |
| 125 | |
| 126 | |
| 127 | @pytest.mark.parametrize( |
| 128 | "activities", [["CPU", "CUDA"], ["CUDA"], ["CPU", "XPU"], ["XPU"]] |
| 129 | ) |
| 130 | @pytest.mark.parametrize("dump_device_time", [True, False]) |
| 131 | def test_torch_profiler_device_summary(tmp_path, capsys, activities, dump_device_time): |
| 132 | """Device summaries honor the dump option on both CUDA and XPU.""" |
| 133 | config = ProfilerConfig( |
| 134 | profiler="torch", |
| 135 | torch_profiler_dir=str(tmp_path), |
| 136 | torch_profiler_dump_cuda_time_total=dump_device_time, |
| 137 | ) |
| 138 | with patch("vllm.profiler.wrapper.torch.profiler.profile") as profile: |
| 139 | profile.return_value.key_averages.return_value.table.return_value = ( |
| 140 | "device times" |
| 141 | ) |
| 142 | wrapper = TorchProfilerWrapper( |
| 143 | config, worker_name="worker", local_rank=0, activities=activities |
| 144 | ) |
| 145 | wrapper.start() |
| 146 | wrapper.stop() |
| 147 | |
| 148 | summary = tmp_path / "profiler_out_0.txt" |
| 149 | assert summary.exists() == dump_device_time |
| 150 | assert ("device times" in capsys.readouterr().out) == dump_device_time |
| 151 | if dump_device_time: |
| 152 | assert summary.read_text() == "device times\n" |
| 153 | |
| 154 | |
| 155 | @pytest.mark.parametrize( |
| 156 | "activities", |
| 157 | [[], ["CPU", "CPU"], ["INVALID"]], |
| 158 | ) |
| 159 | def test_torch_profiler_activities_reject_invalid_values(activities): |
| 160 | with pytest.raises(ValueError, match="torch_profiler_activities"): |
| 161 | ProfilerConfig( |
| 162 | profiler="torch", |
| 163 | torch_profiler_dir="/tmp/mock", |
| 164 | torch_profiler_activities=activities, |
| 165 | ) |
| 166 | |
| 167 | |
| 168 | def test_torch_profiler_activities_require_torch_profiler(): |
| 169 | with pytest.raises(ValueError, match="only applicable"): |
| 170 | ProfilerConfig(torch_profiler_activities=["CPU"]) |
| 171 | |
| 172 | |
| 173 | @pytest.mark.parametrize( |
| 174 | ("device_type", "activities"), |
| 175 | [ |
| 176 | ("cuda", ["XPU"]), |
| 177 | ("xpu", ["CUDA"]), |
| 178 | ("cpu", ["CUDA"]), |
| 179 | ("cpu", ["XPU"]), |
| 180 | ], |
| 181 | ) |
| 182 | def test_worker_rejects_unsupported_activities_at_startup(device_type, activities): |
| 183 | config = ProfilerConfig( |
| 184 | profiler="torch", |
| 185 | torch_profiler_dir="/tmp/mock", |
| 186 | torch_profiler_activities=activities, |
| 187 | ) |
| 188 | |
| 189 | with ( |
| 190 | patch.object(current_platform, "device_type", device_type), |
| 191 | pytest.raises(ValueError, match="Unsupported torch profiler activities"), |
| 192 | ): |
| 193 | validate_worker_profiler_config(config) |
| 194 | |
| 195 | |
| 196 | @pytest.mark.parametrize("device_type", ["cpu", "xpu"]) |
| 197 | def test_worker_rejects_cuda_profiler_on_other_devices(device_type): |
| 198 | with ( |
| 199 | patch.object(current_platform, "device_type", device_type), |
| 200 | pytest.raises(ValueError, match="Unsupported profiler type"), |
| 201 | ): |
| 202 | validate_worker_profiler_config(ProfilerConfig(profiler="cuda")) |
| 203 | |
| 204 | |
| 205 | @pytest.mark.parametrize( |
| 206 | ("device_type", "activities", "expected"), |
| 207 | [ |
| 208 | ("cuda", None, ("CPU", "CUDA")), |
| 209 | ("xpu", None, ("CPU", "XPU")), |
| 210 | ("cpu", None, ("CPU",)), |
| 211 | ("cuda", ["CUDA"], ("CUDA",)), |
| 212 | ("xpu", ["XPU"], ("XPU",)), |
| 213 | ("cuda", ["CPU"], ("CPU",)), |
| 214 | ("xpu", ["CPU"], ("CPU",)), |
| 215 | ("cpu", ["CPU"], ("CPU",)), |
| 216 | ], |
| 217 | ) |
| 218 | def test_worker_creates_platform_torch_profiler(device_type, activities, expected): |
| 219 | config = ProfilerConfig( |
| 220 | profiler="torch", |
| 221 | torch_profiler_dir="/tmp/mock", |
| 222 | torch_profiler_activities=activities, |
| 223 | ) |
| 224 | |
| 225 | with ( |
| 226 | patch.object(current_platform, "device_type", device_type), |
| 227 | patch("vllm.profiler.wrapper.TorchProfilerWrapper") as wrapper, |
| 228 | ): |
| 229 | validate_worker_profiler_config(config) |
| 230 | profiler = create_worker_profiler( |
| 231 | config, |
| 232 | worker_name="rank0", |
| 233 | local_rank=0, |
| 234 | ) |
| 235 | |
| 236 | assert profiler is wrapper.return_value |
| 237 | wrapper.assert_called_once_with( |
| 238 | config, |
| 239 | worker_name="rank0", |
| 240 | local_rank=0, |
| 241 | activities=expected, |
| 242 | ) |
| 243 | |
| 244 | |
| 245 | @pytest.mark.parametrize("worker_type", [Worker, XPUWorker]) |
| 246 | def test_worker_reuses_torch_wrapper_across_profile_rounds(worker_type): |
| 247 | worker = object.__new__(worker_type) |
| 248 | worker.rank = 0 |
| 249 | worker.local_rank = 0 |
| 250 | worker.profiler = None |
| 251 | worker.profiler_config = ProfilerConfig( |
| 252 | profiler="torch", torch_profiler_dir="/tmp/mock" |
| 253 | ) |
| 254 | |
| 255 | with ( |
| 256 | patch("vllm.distributed.utils.get_worker_rank_suffix", return_value="rank0"), |
| 257 | patch("vllm.profiler.wrapper.TorchProfilerWrapper") as wrapper, |
| 258 | ): |
| 259 | worker.profile() |
| 260 | worker.profile(is_start=False) |
| 261 | worker.profile() |
| 262 | |
| 263 | assert worker.profiler is wrapper.return_value |
| 264 | wrapper.assert_called_once() |
| 265 | assert wrapper.return_value.start.call_count == 2 |
| 266 | wrapper.return_value.stop.assert_called_once_with() |
| 267 | |
| 268 | |
| 269 | def test_immediate_start_stop(default_profiler_config): |
| 270 | """Test standard start without delay.""" |
| 271 | profiler = ConcreteWorkerProfiler(default_profiler_config) |
| 272 | profiler.start() |
| 273 | assert profiler._running is True |
| 274 | assert profiler._active is True |
| 275 | assert profiler.start_call_count == 1 |
| 276 | |
| 277 | profiler.stop() |
| 278 | assert profiler._running is False |
| 279 | assert profiler._active is False |
| 280 | assert profiler.stop_call_count == 1 |
| 281 | |
| 282 | |
| 283 | def test_delayed_start(default_profiler_config): |
| 284 | """Test that profiler waits for N steps before actually starting.""" |
| 285 | default_profiler_config.delay_iterations = 2 |
| 286 | profiler = ConcreteWorkerProfiler(default_profiler_config) |
| 287 | |
| 288 | # User requests start |
| 289 | profiler.start() |
| 290 | |
| 291 | # Should be active (request accepted) but not running (waiting for delay) |
| 292 | assert profiler._active is True |
| 293 | assert profiler._running is False |
| 294 | assert profiler.start_call_count == 0 |
| 295 | |
| 296 | # Step 1 |
| 297 | profiler.step() |
| 298 | assert profiler._running is False |
| 299 | |
| 300 | # Step 2 (Threshold reached) |
| 301 | profiler.step() |
| 302 | assert profiler._running is True |
| 303 | assert profiler.start_call_count == 1 |
| 304 | |
| 305 | |
| 306 | def test_max_iterations(default_profiler_config): |
| 307 | """Test that profiler stops automatically after max iterations.""" |
| 308 | default_profiler_config.max_iterations = 2 |
| 309 | profiler = ConcreteWorkerProfiler(default_profiler_config) |
| 310 | |
| 311 | profiler.start() |
| 312 | assert profiler._running is True |
| 313 | |
| 314 | # Iteration 1 |
| 315 | profiler.step() # profiling_count becomes 1 |
| 316 | assert profiler._running is True |
| 317 | |
| 318 | # Iteration 2 |
| 319 | profiler.step() # profiling_count becomes 2 |
| 320 | assert profiler._running is True |
| 321 | |
| 322 | # Iteration 3 (Exceeds max) |
| 323 | profiler.step() # profiling_count becomes 3 |
| 324 | |
| 325 | # Should have stopped now |
| 326 | assert profiler._running is False |
| 327 | assert profiler.stop_call_count == 1 |
| 328 | # And fully reset, not just paused -- a later start_profile must not be a |
| 329 | # permanent no-op just because max_iterations already fired once. |
| 330 | assert profiler._active is False |
| 331 | assert profiler._active_iteration_count == 0 |
| 332 | assert profiler._profiling_for_iters == 0 |
| 333 | |
| 334 | |
| 335 | def test_restart_after_max_iterations(default_profiler_config): |
| 336 | """A start_profile after an auto-stop must actually restart, not be |
| 337 | silently ignored (regression test: auto-stop used to leave _active |
| 338 | True forever, so start() always bailed out early).""" |
| 339 | default_profiler_config.max_iterations = 2 |
| 340 | profiler = ConcreteWorkerProfiler(default_profiler_config) |
| 341 | |
| 342 | profiler.start() |
| 343 | profiler.step() |
| 344 | profiler.step() |
| 345 | profiler.step() # exceeds max, auto-stops |
| 346 | assert profiler._running is False |
| 347 | assert profiler.start_call_count == 1 |
| 348 | |
| 349 | profiler.start() |
| 350 | assert profiler._active is True |
| 351 | assert profiler._running is True |
| 352 | assert profiler.start_call_count == 2 |
| 353 | |
| 354 | profiler.step() |
| 355 | assert profiler._running is True |
| 356 | |
| 357 | |
| 358 | def test_delayed_start_and_max_iters(default_profiler_config): |
| 359 | """Test combined delayed start and max iterations.""" |
| 360 | default_profiler_config.delay_iterations = 2 |
| 361 | default_profiler_config.max_iterations = 2 |
| 362 | profiler = ConcreteWorkerProfiler(default_profiler_config) |
| 363 | profiler.start() |
| 364 | |
| 365 | # Step 1 |
| 366 | profiler.step() |
| 367 | assert profiler._running is False |
| 368 | assert profiler._active is True |
| 369 | |
| 370 | # Step 2 (Starts now) |
| 371 | profiler.step() |
| 372 | assert profiler._profiling_for_iters == 1 |
| 373 | assert profiler._running is True |
| 374 | assert profiler._active is True |
| 375 | |
| 376 | # Next iteration |
| 377 | profiler.step() |
| 378 | assert profiler._profiling_for_iters == 2 |
| 379 | assert profiler._running is True |
| 380 | |
| 381 | # Iteration 2 (exceeds max) |
| 382 | profiler.step() |
| 383 | |
| 384 | # Should have stopped now |
| 385 | assert profiler._running is False |
| 386 | assert profiler.stop_call_count == 1 |
| 387 | |
| 388 | |
| 389 | def test_idempotency(default_profiler_config): |
| 390 | """Test that calling start/stop multiple times doesn't break logic.""" |
| 391 | profiler = ConcreteWorkerProfiler(default_profiler_config) |
| 392 | |
| 393 | # Double Start |
| 394 | profiler.start() |
| 395 | profiler.start() |
| 396 | assert profiler.start_call_count == 1 # Should only start once |
| 397 | |
| 398 | # Double Stop |
| 399 | profiler.stop() |
| 400 | profiler.stop() |
| 401 | assert profiler.stop_call_count == 1 # Should only stop once |
| 402 | |
| 403 | |
| 404 | def test_step_inactive(default_profiler_config): |
| 405 | """Test that stepping while inactive does nothing.""" |
| 406 | default_profiler_config.delay_iterations = 2 |
| 407 | profiler = ConcreteWorkerProfiler(default_profiler_config) |
| 408 | |
| 409 | # Not started yet |
| 410 | profiler.step() |
| 411 | profiler.step() |
| 412 | |
| 413 | # Even though we stepped 2 times, start shouldn't happen because active=False |
| 414 | assert profiler.start_call_count == 0 |
| 415 | |
| 416 | |
| 417 | def test_start_failure(default_profiler_config): |
| 418 | """Test behavior when the underlying _start method raises exception.""" |
| 419 | profiler = ConcreteWorkerProfiler(default_profiler_config) |
| 420 | profiler.should_fail_start = True |
| 421 | |
| 422 | profiler.start() |
| 423 | |
| 424 | # Exception caught in _call_start |
| 425 | assert profiler._running is False # Should not mark as running |
| 426 | assert profiler._active is True # Request is still considered active |
| 427 | assert profiler.start_call_count == 0 # Logic failed inside start |
| 428 | |
| 429 | |
| 430 | def test_shutdown(default_profiler_config): |
| 431 | """Test that shutdown calls stop only if running.""" |
| 432 | profiler = ConcreteWorkerProfiler(default_profiler_config) |
| 433 | |
| 434 | # Case 1: Not running |
| 435 | profiler.shutdown() |
| 436 | assert profiler.stop_call_count == 0 |
| 437 | |
| 438 | # Case 2: Running |
| 439 | profiler.start() |
| 440 | profiler.shutdown() |
| 441 | assert profiler.stop_call_count == 1 |
| 442 | |
| 443 | |
| 444 | def test_mixed_delay_and_stop(default_profiler_config): |
| 445 | """Test manual stop during the delay period.""" |
| 446 | default_profiler_config.delay_iterations = 5 |
| 447 | profiler = ConcreteWorkerProfiler(default_profiler_config) |
| 448 | |
| 449 | profiler.start() |
| 450 | profiler.step() |
| 451 | profiler.step() |
| 452 | |
| 453 | # User cancels before delay finishes |
| 454 | profiler.stop() |
| 455 | assert profiler._active is False |
| 456 | |
| 457 | # Further steps should not trigger start |
| 458 | profiler.step() |
| 459 | profiler.step() |
| 460 | profiler.step() |
| 461 | |
| 462 | assert profiler.start_call_count == 0 |
| 463 | |
| 464 | |
| 465 | class TestIsUriPath: |
| 466 | """Tests for the _is_uri_path helper function.""" |
| 467 | |
| 468 | @pytest.mark.parametrize( |
| 469 | "path,expected", |
| 470 | [ |
| 471 | # Valid URI schemes - should return True |
| 472 | ("gs://bucket/path", True), |
| 473 | ("s3://bucket/path", True), |
| 474 | ("hdfs://cluster/path", True), |
| 475 | ("abfs://container/path", True), |
| 476 | ("http://example.com/path", True), |
| 477 | ("https://example.com/path", True), |
| 478 | # Local paths - should return False |
| 479 | ("/tmp/local/path", False), |
| 480 | ("./relative/path", False), |
| 481 | ("relative/path", False), |
| 482 | ("/absolute/path", False), |
| 483 | # Windows drive letters - should return False (single char scheme) |
| 484 | ("C://windows/path", False), |
| 485 | ("D://drive/path", False), |
| 486 | # Edge cases |
| 487 | ("", False), |
| 488 | ("no-scheme", False), |
| 489 | ("scheme-no-slashes:", False), |
| 490 | ("://no-scheme", False), |
| 491 | ], |
| 492 | ) |
| 493 | def test_is_uri_path(self, path, expected): |
| 494 | """Test that _is_uri_path correctly identifies URI vs local paths.""" |
| 495 | assert _is_uri_path(path) == expected |
| 496 | |
| 497 | |
| 498 | class TestAnnotateProfile: |
| 499 | """Tests for Worker.annotate_profile() annotation string formatting.""" |
| 500 | |
| 501 | def _annotate(self, detailed: bool) -> str: |
| 502 | worker = MagicMock() |
| 503 | worker.vllm_config.profiler_config.detailed_trace_annotation = detailed |
| 504 | worker.profiler = MagicMock() |
| 505 | |
| 506 | ctx_req = MagicMock(req_id="ctx1", num_computed_tokens=0) |
| 507 | cached = CachedRequestData( |
| 508 | req_ids=["gen1"], |
| 509 | resumed_req_ids=set(), |
| 510 | new_token_ids=[], |
| 511 | all_token_ids={}, |
| 512 | new_block_ids=[], |
| 513 | num_computed_tokens=[10], |
| 514 | num_output_tokens=[1], |
| 515 | ) |
| 516 | sched = MagicMock( |
| 517 | scheduled_new_reqs=[ctx_req], |
| 518 | scheduled_cached_reqs=cached, |
| 519 | num_scheduled_tokens={"ctx1": 4, "gen1": 1}, |
| 520 | ) |
| 521 | |
| 522 | Worker.annotate_profile(worker, sched) |
| 523 | return worker.profiler.annotate_context_manager.call_args[0][0] |
| 524 | |
| 525 | def test_simple_format_mixed(self): |
| 526 | assert self._annotate(detailed=False) == ( |
| 527 | "execute_context_1(4)_generation_1(1)" |
| 528 | ) |
| 529 | |
| 530 | def test_detailed_format_mixed(self): |
| 531 | # ctx1: sq=4, sk=4, sqsq=16, sqsk=16 | gen1: sq=1, sk=11, sqsq=1, sqsk=11 | bs=5 |
| 532 | assert self._annotate(detailed=True) == ( |
| 533 | "execute_5_context_1(sq4sk4sqsq16sqsk16)_generation_1(sq1sk11sqsq1sqsk11)" |
| 534 | ) |
| 535 | |
| 536 | def test_skips_annotation_work_when_profiler_does_not_annotate(self): |
| 537 | worker = MagicMock() |
| 538 | worker.profiler.should_annotate = False |
| 539 | |
| 540 | with patch( |
| 541 | "vllm.v1.worker.gpu_worker.compute_iteration_details" |
| 542 | ) as compute_iteration_details: |
| 543 | context = Worker.annotate_profile(worker, scheduler_output=None) |
| 544 | |
| 545 | worker.profiler.step.assert_called_once_with() |
| 546 | compute_iteration_details.assert_not_called() |
| 547 | worker.profiler.annotate_context_manager.assert_not_called() |
| 548 | assert isinstance(context, nullcontext) |
| 549 | |
| 550 | |
| 551 | def test_profiler_entered_during_capture(): |
| 552 | """Profiler is used as a context manager in _warmup_and_capture, |
| 553 | confirming it is active during the actual graph capture run.""" |
| 554 | runner = MagicMock() |
| 555 | runner.compilation_config.cudagraph_num_of_warmups = 0 |
| 556 | mock_profiler = MagicMock() |
| 557 | |
| 558 | GPUModelRunner._warmup_and_capture( |
| 559 | runner, |
| 560 | desc=MagicMock(num_tokens=4, uniform=True), |
| 561 | cudagraph_runtime_mode=CUDAGraphMode.FULL, |
| 562 | profiler=mock_profiler, |
| 563 | ) |
| 564 | |
| 565 | mock_profiler.__enter__.assert_called_once() |
| 566 | mock_profiler.__exit__.assert_called_once() |
| 567 | |
| 568 | |
| 569 | def make_proton(session_id: int | None = 7): |
| 570 | data = SimpleNamespace( |
| 571 | advance_phase=Mock(side_effect=range(1, 100)), |
| 572 | clear=Mock(), |
| 573 | get=Mock(return_value={"traceEvents": []}), |
| 574 | get_msgpack=Mock(return_value=b"profile"), |
| 575 | ) |
| 576 | return SimpleNamespace( |
| 577 | start=Mock(return_value=session_id), |
| 578 | activate=Mock(), |
| 579 | deactivate=Mock(), |
| 580 | finalize=Mock(), |
| 581 | scope=Mock(return_value=nullcontext()), |
| 582 | data=data, |
| 583 | ) |
| 584 | |
| 585 | |
| 586 | def make_proton_wrapper( |
| 587 | tmp_path, proton=None, triton_version="3.7.0", **config_overrides |
| 588 | ): |
| 589 | proton = proton or make_proton() |
| 590 | config = ProfilerConfig( |
| 591 | profiler="proton", |
| 592 | proton_profiler_dir=str(tmp_path), |
| 593 | **config_overrides, |
| 594 | ) |
| 595 | |
| 596 | def import_module(name): |
| 597 | if name == "triton.profiler": |
| 598 | return proton |
| 599 | assert name == "triton" |
| 600 | return SimpleNamespace(__version__=triton_version) |
| 601 | |
| 602 | with patch( |
| 603 | "vllm.profiler.wrapper.importlib.import_module", side_effect=import_module |
| 604 | ): |
| 605 | wrapper = ProtonProfilerWrapper(config, worker_name="rank_3") |
| 606 | return wrapper, proton |
| 607 | |
| 608 | |
| 609 | _requires_cuda_for_proton = pytest.mark.skipif( |
| 610 | not current_platform.is_cuda(), |
| 611 | reason="Proton profiling tests require an NVIDIA CUDA platform.", |
| 612 | ) |
| 613 | |
| 614 | # Proton subscribes to CUPTI itself, and a process gets one CUPTI subscriber. |
| 615 | # CI's kernel-launch recorder is another, injected through CUDA_INJECTION64_PATH |
| 616 | # on recording runs, and there Proton's cuptiSubscribe fails with error 39. |
| 617 | _requires_no_injected_cupti_tool = pytest.mark.skipif( |
| 618 | bool(os.environ.get("CUDA_INJECTION64_PATH")), |
| 619 | reason="Another CUPTI tool is injected (CUDA_INJECTION64_PATH); " |
| 620 | "Proton has to be the only one.", |
| 621 | ) |
| 622 | |
| 623 | |
| 624 | @_requires_cuda_for_proton |
| 625 | class TestProtonConfig: |
| 626 | def test_normalizes_local_output_directory(self, tmp_path, monkeypatch): |
| 627 | monkeypatch.chdir(tmp_path) |
| 628 | config = ProfilerConfig(profiler="proton", proton_profiler_dir="profiles") |
| 629 | assert config.proton_profiler_dir == os.path.join(tmp_path, "profiles") |
| 630 | |
| 631 | @pytest.mark.parametrize( |
| 632 | ("options", "message"), |
| 633 | [ |
| 634 | ({"proton_profiler_dir": ""}, "must be set"), |
| 635 | ({"proton_profiler_dir": "s3://bucket/profiles"}, "local directory"), |
| 636 | ( |
| 637 | {"proton_data": "tree", "proton_output_format": "chrome_trace"}, |
| 638 | "requires proton_data", |
| 639 | ), |
| 640 | ( |
| 641 | {"proton_data": "trace", "proton_output_format": "hatchet"}, |
| 642 | "requires proton_data", |
| 643 | ), |
| 644 | ( |
| 645 | { |
| 646 | "proton_data": "trace", |
| 647 | "proton_output_format": "hatchet_msgpack", |
| 648 | }, |
| 649 | "requires proton_data", |
| 650 | ), |
| 651 | ( |
| 652 | { |
| 653 | "proton_data": "trace", |
| 654 | "proton_graph_attribution": True, |
| 655 | }, |
| 656 | "requires proton_data='tree'", |
| 657 | ), |
| 658 | ], |
| 659 | ) |
| 660 | def test_rejects_invalid_option_combinations(self, tmp_path, options, message): |
| 661 | kwargs = {"proton_profiler_dir": str(tmp_path), **options} |
| 662 | with pytest.raises(ValueError, match=message): |
| 663 | ProfilerConfig(profiler="proton", **kwargs) |
| 664 | |
| 665 | @pytest.mark.parametrize( |
| 666 | "field", |
| 667 | [ |
| 668 | "proton_context", |
| 669 | "proton_data", |
| 670 | "proton_backend", |
| 671 | "proton_hook", |
| 672 | "proton_output_format", |
| 673 | ], |
| 674 | ) |
| 675 | def test_rejects_invalid_typed_options(self, field, tmp_path): |
| 676 | with pytest.raises(ValidationError): |
| 677 | ProfilerConfig( |
| 678 | profiler="proton", |
| 679 | proton_profiler_dir=str(tmp_path), |
| 680 | **{field: "invalid"}, |
| 681 | ) |
| 682 | |
| 683 | @pytest.mark.parametrize( |
| 684 | ("field", "value"), |
| 685 | [ |
| 686 | ("proton_profiler_dir", "profiles"), |
| 687 | ("proton_context", "python"), |
| 688 | ("proton_data", "trace"), |
| 689 | ("proton_backend", "cupti"), |
| 690 | ("proton_mode", "pcsampling"), |
| 691 | ("proton_hook", "triton"), |
| 692 | ("proton_output_format", "chrome_trace"), |
| 693 | ("proton_graph_attribution", True), |
| 694 | ], |
| 695 | ) |
| 696 | def test_rejects_proton_options_for_other_profilers(self, field, value): |
| 697 | with pytest.raises(ValueError, match=f"{field} only applicable"): |
| 698 | ProfilerConfig(**{field: value}) |
| 699 | |
| 700 | def test_allows_proton_when_cuda_graphs_are_disabled(self, tmp_path): |
| 701 | config = VllmConfig( |
| 702 | profiler_config=ProfilerConfig( |
| 703 | profiler="proton", proton_profiler_dir=str(tmp_path) |
| 704 | ), |
| 705 | compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE), |
| 706 | ) |
| 707 | |
| 708 | assert config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE |
| 709 | |
| 710 | def test_rejects_proton_on_non_cuda_platforms(self, tmp_path): |
| 711 | with ( |
| 712 | patch("vllm.platforms.current_platform.is_cuda", return_value=False), |
| 713 | pytest.raises(ValueError, match="supports NVIDIA CUDA only"), |
| 714 | ): |
| 715 | VllmConfig( |
| 716 | profiler_config=ProfilerConfig( |
| 717 | profiler="proton", proton_profiler_dir=str(tmp_path) |
| 718 | ), |
| 719 | compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE), |
| 720 | ) |
| 721 | |
| 722 | @pytest.mark.parametrize("encoder_only", [False, True]) |
| 723 | @pytest.mark.parametrize("attribution", [False, True]) |
| 724 | def test_cuda_graphs_require_attribution(self, tmp_path, encoder_only, attribution): |
| 725 | # Encoder graphs are independent of the decoder's cudagraph_mode. |
| 726 | expected = ( |
| 727 | nullcontext() |
| 728 | if attribution |
| 729 | else pytest.raises( |
| 730 | ValueError, match="requires proton_graph_attribution=True" |
| 731 | ) |
| 732 | ) |
| 733 | with expected: |
| 734 | VllmConfig( |
| 735 | profiler_config=ProfilerConfig( |
| 736 | profiler="proton", |
| 737 | proton_profiler_dir=str(tmp_path), |
| 738 | proton_graph_attribution=attribution, |
| 739 | ), |
| 740 | compilation_config=CompilationConfig( |
| 741 | cudagraph_mode=( |
| 742 | CUDAGraphMode.NONE if encoder_only else CUDAGraphMode.FULL |
| 743 | ), |
| 744 | cudagraph_mm_encoder=encoder_only, |
| 745 | ), |
| 746 | ) |
| 747 | |
| 748 | def test_validates_default_cuda_graph_mode_after_resolution(self, tmp_path): |
| 749 | with pytest.raises(ValueError, match="requires proton_graph_attribution=True"): |
| 750 | VllmConfig( |
| 751 | profiler_config=ProfilerConfig( |
| 752 | profiler="proton", proton_profiler_dir=str(tmp_path) |
| 753 | ), |
| 754 | ) |
| 755 | |
| 756 | @pytest.mark.parametrize( |
| 757 | "mode", |
| 758 | [ |
| 759 | "periodic_flushing", |
| 760 | "periodic_flushing:format=hatchet", |
| 761 | "PERIODIC_FLUSHING:format=hatchet", |
| 762 | ], |
| 763 | ) |
| 764 | def test_graph_attribution_rejects_periodic_flushing(self, tmp_path, mode): |
| 765 | # Reject before Proton's native phase manager can abort the worker. |
| 766 | with pytest.raises(ValueError, match="incompatible with periodic_flushing"): |
| 767 | ProfilerConfig( |
| 768 | profiler="proton", |
| 769 | proton_profiler_dir=str(tmp_path), |
| 770 | proton_graph_attribution=True, |
| 771 | proton_mode=mode, |
| 772 | ) |
| 773 | |
| 774 | @pytest.mark.parametrize("attribution", [False, True]) |
| 775 | @pytest.mark.parametrize( |
| 776 | "mode", ["pcsampling", "pcsampling:interval=100", "PcSampling:interval=100"] |
| 777 | ) |
| 778 | @pytest.mark.parametrize("encoder_only", [False, True]) |
| 779 | def test_pcsampling_requires_graphs_disabled( |
| 780 | self, tmp_path, attribution, mode, encoder_only |
| 781 | ): |
| 782 | with pytest.raises(ValueError, match="PC sampling requires CUDA graphs"): |
| 783 | VllmConfig( |
| 784 | profiler_config=ProfilerConfig( |
| 785 | profiler="proton", |
| 786 | proton_profiler_dir=str(tmp_path), |
| 787 | proton_graph_attribution=attribution, |
| 788 | proton_mode=mode, |
| 789 | ), |
| 790 | compilation_config=CompilationConfig( |
| 791 | cudagraph_mode=CUDAGraphMode.NONE |
| 792 | if encoder_only |
| 793 | else CUDAGraphMode.FULL, |
| 794 | cudagraph_mm_encoder=encoder_only, |
| 795 | ), |
| 796 | ) |
| 797 | |
| 798 | def test_ordinary_proton_keeps_mrv1_graph_restriction(self, tmp_path, monkeypatch): |
| 799 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") |
| 800 | with pytest.raises(ValueError, match="requires proton_graph_attribution=True"): |
| 801 | VllmConfig( |
| 802 | profiler_config=ProfilerConfig( |
| 803 | profiler="proton", proton_profiler_dir=str(tmp_path) |
| 804 | ), |
| 805 | compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.FULL), |
| 806 | ) |
| 807 | |
| 808 | def test_graph_attribution_rejects_mrv1(self, tmp_path, monkeypatch): |
| 809 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") |
| 810 | with pytest.raises(ValueError, match="requires the V2 model runner"): |
| 811 | VllmConfig( |
| 812 | profiler_config=ProfilerConfig( |
| 813 | profiler="proton", |
| 814 | proton_profiler_dir=str(tmp_path), |
| 815 | proton_graph_attribution=True, |
| 816 | ) |
| 817 | ) |
| 818 | |
| 819 | |
| 820 | @_requires_cuda_for_proton |
| 821 | class TestProtonProfilerWrapper: |
| 822 | def test_passes_config_and_global_rank_name_to_proton(self, tmp_path): |
| 823 | wrapper, proton = make_proton_wrapper( |
| 824 | tmp_path, |
| 825 | proton_context="python", |
| 826 | proton_data="trace", |
| 827 | proton_backend="cupti", |
| 828 | proton_mode="pcsampling", |
| 829 | proton_hook="triton", |
| 830 | proton_output_format="chrome_trace", |
| 831 | ) |
| 832 | |
| 833 | wrapper.start() |
| 834 | |
| 835 | start_args = proton.start.call_args.kwargs |
| 836 | assert start_args["name"].startswith(os.path.join(tmp_path, "proton_rank_3_")) |
| 837 | assert start_args["name"].endswith("_run0") |
| 838 | assert start_args | {"name": None} == { |
| 839 | "name": None, |
| 840 | "context": "python", |
| 841 | "data": "trace", |
| 842 | "backend": "cupti", |
| 843 | "mode": "pcsampling", |
| 844 | "hook": "triton", |
| 845 | } |
| 846 | wrapper.stop() |
| 847 | proton.finalize.assert_called_once_with(session=7, output_format="chrome_trace") |
| 848 | assert tmp_path.is_dir() |
| 849 | |
| 850 | def test_finalizes_each_profile_with_unique_output_names(self, tmp_path): |
| 851 | proton = make_proton() |
| 852 | proton.start.side_effect = [7, 8] |
| 853 | wrapper, proton = make_proton_wrapper(tmp_path, proton) |
| 854 | |
| 855 | wrapper.start() |
| 856 | wrapper.start() |
| 857 | wrapper.stop() |
| 858 | wrapper.start() |
| 859 | wrapper.stop() |
| 860 | |
| 861 | names = [c.kwargs["name"] for c in proton.start.call_args_list] |
| 862 | assert len(names) == len(set(names)) == 2 |
| 863 | assert names[0].endswith("_run0") |
| 864 | assert names[1].endswith("_run1") |
| 865 | assert proton.deactivate.call_count == 2 |
| 866 | assert proton.finalize.call_args_list == [call(session=7), call(session=8)] |
| 867 | |
| 868 | def test_output_names_are_unique_across_worker_restarts(self, tmp_path): |
| 869 | with patch( |
| 870 | "vllm.profiler.wrapper.uuid4", |
| 871 | side_effect=[UUID(int=1), UUID(int=2)], |
| 872 | ): |
| 873 | first, first_proton = make_proton_wrapper(tmp_path) |
| 874 | second, second_proton = make_proton_wrapper(tmp_path) |
| 875 | |
| 876 | first.start() |
| 877 | second.start() |
| 878 | |
| 879 | first_name = first_proton.start.call_args.kwargs["name"] |
| 880 | second_name = second_proton.start.call_args.kwargs["name"] |
| 881 | assert first_name != second_name |
| 882 | assert first_name.endswith(f"_{UUID(int=1).hex}_run0") |
| 883 | assert second_name.endswith(f"_{UUID(int=2).hex}_run0") |
| 884 | |
| 885 | @pytest.mark.parametrize( |
| 886 | ("option", "value", "feature"), |
| 887 | [ |
| 888 | ("proton_output_format", "hatchet_msgpack", "hatchet_msgpack"), |
| 889 | ("proton_mode", "periodic_flushing", "periodic flushing"), |
| 890 | ], |
| 891 | ) |
| 892 | def test_newer_features_reject_triton_3_6(self, tmp_path, option, value, feature): |
| 893 | with pytest.raises(RuntimeError, match=feature): |
| 894 | make_proton_wrapper(tmp_path, triton_version="3.6.0", **{option: value}) |
| 895 | |
| 896 | @pytest.mark.parametrize("version", ["3.6.0", "unknown"]) |
| 897 | def test_graph_attribution_requires_phase_api(self, tmp_path, version): |
| 898 | with pytest.raises(RuntimeError, match="requires Triton >= 3.7"): |
| 899 | make_proton_wrapper( |
| 900 | tmp_path, triton_version=version, proton_graph_attribution=True |
| 901 | ) |
| 902 | |
| 903 | def test_ordinary_profiling_still_supports_triton_3_6(self, tmp_path): |
| 904 | wrapper, proton = make_proton_wrapper(tmp_path, triton_version="3.6.0") |
| 905 | wrapper.start() |
| 906 | wrapper.stop() |
| 907 | proton.finalize.assert_called_once_with(session=7) |
| 908 | |
| 909 | @pytest.mark.parametrize( |
| 910 | ("option", "value"), |
| 911 | [ |
| 912 | ("proton_output_format", "hatchet_msgpack"), |
| 913 | ("proton_mode", "periodic_flushing"), |
| 914 | ], |
| 915 | ) |
| 916 | def test_triton_3_7_features(self, tmp_path, option, value): |
| 917 | make_proton_wrapper(tmp_path, triton_version="3.7.0", **{option: value}) |
| 918 | |
| 919 | def test_rejects_output_format_when_finalize_lacks_capability(self, tmp_path): |
| 920 | proton = make_proton() |
| 921 | proton.finalize = lambda session=None: None |
| 922 | |
| 923 | with pytest.raises(RuntimeError, match="does not support selecting"): |
| 924 | make_proton_wrapper(tmp_path, proton, proton_output_format="hatchet") |
| 925 | |
| 926 | def test_cuda_graph_tree_phase_is_written_at_stop(self, tmp_path): |
| 927 | wrapper, proton = make_proton_wrapper( |
| 928 | tmp_path, |
| 929 | proton_context="python", |
| 930 | proton_graph_attribution=True, |
| 931 | ) |
| 932 | with wrapper.capture_cuda_graphs(): |
| 933 | proton.start.assert_called_once() |
| 934 | proton.deactivate.assert_not_called() |
| 935 | |
| 936 | capture_args = proton.start.call_args.kwargs |
| 937 | assert capture_args["context"] == "python" |
| 938 | assert capture_args["data"] == "tree" |
| 939 | proton.data.advance_phase.assert_called_once_with(7) |
| 940 | proton.deactivate.assert_called_once_with(session=7, flushing=True) |
| 941 | proton.data.clear.assert_called_once_with(7, 0) |
| 942 | |
| 943 | wrapper.start() |
| 944 | wrapper.stop() |
| 945 | wrapper.start() |
| 946 | wrapper.stop() |
| 947 | |
| 948 | assert proton.data.get.call_args_list == [call(7, 1), call(7, 2)] |
| 949 | proton.data.clear.assert_has_calls([call(7, 0), call(7, 1), call(7, 2)]) |
| 950 | output_names = sorted(tmp_path.glob("proton_rank_3_*.hatchet")) |
| 951 | assert len(output_names) == 2 |
| 952 | assert output_names[0].name.endswith("_run0.hatchet") |
| 953 | assert output_names[1].name.endswith("_run1.hatchet") |
| 954 | |
| 955 | def test_capture_is_noop_without_opt_in(self, tmp_path): |
| 956 | wrapper, proton = make_proton_wrapper(tmp_path) |
| 957 | with wrapper.capture_cuda_graphs(): |
| 958 | pass |
| 959 | proton.start.assert_not_called() |
| 960 | |
| 961 | @pytest.mark.parametrize("delay", [0, 2]) |
| 962 | def test_duplicate_start_preserves_output_prefix(self, tmp_path, delay): |
| 963 | wrapper, proton = make_proton_wrapper( |
| 964 | tmp_path, proton_graph_attribution=True, delay_iterations=delay |
| 965 | ) |
| 966 | with wrapper.capture_cuda_graphs(): |
| 967 | pass |
| 968 | wrapper.set_output_name("first") |
| 969 | wrapper.start() |
| 970 | wrapper.set_output_name("duplicate") |
| 971 | wrapper.start() |
| 972 | for _ in range(delay): |
| 973 | wrapper.step() |
| 974 | wrapper.stop() |
| 975 | assert len(list(tmp_path.glob("proton_first_*.hatchet"))) == 1 |
| 976 | assert not list(tmp_path.glob("proton_duplicate_*")) |
| 977 | wrapper.set_output_name("second") |
| 978 | wrapper.start() |
| 979 | for _ in range(delay): |
| 980 | wrapper.step() |
| 981 | wrapper.stop() |
| 982 | assert len(list(tmp_path.glob("proton_second_*.hatchet"))) == 1 |
| 983 | |
| 984 | def test_failed_export_clears_activity_and_allows_next_interval(self, tmp_path): |
| 985 | wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) |
| 986 | with wrapper.capture_cuda_graphs(): |
| 987 | pass |
| 988 | proton.data.get.side_effect = [OSError("disk full"), []] |
| 989 | wrapper.start() |
| 990 | wrapper.stop() |
| 991 | proton.data.clear.assert_called_with(7, 1) |
| 992 | wrapper.start() |
| 993 | wrapper.stop() |
| 994 | proton.data.clear.assert_called_with(7, 2) |
| 995 | assert len(list(tmp_path.glob("*_run1.hatchet"))) == 1 |
| 996 | |
| 997 | def test_cuda_graph_context_deactivates_after_capture_error(self, tmp_path): |
| 998 | wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) |
| 999 | |
| 1000 | with ( |
| 1001 | pytest.raises(RuntimeError, match="capture failed"), |
| 1002 | wrapper.capture_cuda_graphs(), |
| 1003 | ): |
| 1004 | raise RuntimeError("capture failed") |
| 1005 | |
| 1006 | proton.deactivate.assert_called_once_with(session=7, flushing=True) |
| 1007 | |
| 1008 | def test_cuda_graph_capture_deactivates_when_phase_advance_fails(self, tmp_path): |
| 1009 | wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) |
| 1010 | proton.data.advance_phase.side_effect = RuntimeError("advance failed") |
| 1011 | |
| 1012 | with ( |
| 1013 | pytest.raises(RuntimeError, match="advance failed"), |
| 1014 | wrapper.capture_cuda_graphs(), |
| 1015 | ): |
| 1016 | pass |
| 1017 | |
| 1018 | proton.deactivate.assert_called_once_with(session=7, flushing=True) |
| 1019 | proton.data.clear.assert_not_called() |
| 1020 | |
| 1021 | def test_cuda_graph_stop_deactivates_when_phase_advance_fails(self, tmp_path): |
| 1022 | wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) |
| 1023 | with wrapper.capture_cuda_graphs(): |
| 1024 | pass |
| 1025 | proton.data.clear.reset_mock() |
| 1026 | proton.data.advance_phase.side_effect = RuntimeError("advance failed") |
| 1027 | |
| 1028 | wrapper._start() |
| 1029 | with pytest.raises(RuntimeError, match="advance failed"): |
| 1030 | wrapper._stop() |
| 1031 | |
| 1032 | proton.deactivate.assert_called_with(session=7, flushing=True) |
| 1033 | proton.data.clear.assert_not_called() |
| 1034 | |
| 1035 | def test_shutdown_finalizes_cuda_graph_capture_session(self, tmp_path): |
| 1036 | wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True) |
| 1037 | with wrapper.capture_cuda_graphs(): |
| 1038 | pass |
| 1039 | |
| 1040 | wrapper.shutdown() |
| 1041 | wrapper.shutdown() |
| 1042 | |
| 1043 | proton.finalize.assert_called_once_with(session=7) |
| 1044 | |
| 1045 | def test_missing_proton_has_actionable_error(self, tmp_path): |
| 1046 | config = ProfilerConfig(profiler="proton", proton_profiler_dir=str(tmp_path)) |
| 1047 | with ( |
| 1048 | patch( |
| 1049 | "vllm.profiler.wrapper.importlib.import_module", |
| 1050 | side_effect=ImportError, |
| 1051 | ), |
| 1052 | pytest.raises(RuntimeError, match="requires a Triton installation"), |
| 1053 | ): |
| 1054 | ProtonProfilerWrapper(config, worker_name="rank_0") |
| 1055 | |
| 1056 | def test_scope_annotations_delegate_to_proton(self, tmp_path): |
| 1057 | wrapper, proton = make_proton_wrapper(tmp_path) |
| 1058 | wrapper.start() |
| 1059 | |
| 1060 | context = wrapper.annotate_context_manager("decode") |
| 1061 | |
| 1062 | proton.scope.assert_called_once_with("decode") |
| 1063 | assert context is not None |
| 1064 | |
| 1065 | |
| 1066 | @_requires_cuda_for_proton |
| 1067 | def test_gpu_worker_creates_proton_profiler(): |
| 1068 | worker = object.__new__(Worker) |
| 1069 | worker.rank = 1 |
| 1070 | worker.local_rank = 1 |
| 1071 | worker.profiler = None |
| 1072 | worker.profiler_config = MagicMock(profiler="proton") |
| 1073 | |
| 1074 | with ( |
| 1075 | patch( |
| 1076 | "vllm.distributed.utils.get_worker_rank_suffix", |
| 1077 | return_value="rank1", |
| 1078 | ), |
| 1079 | patch("vllm.profiler.wrapper.ProtonProfilerWrapper") as wrapper, |
| 1080 | ): |
| 1081 | Worker.profile(worker) |
| 1082 | |
| 1083 | wrapper.assert_called_once_with(worker.profiler_config, worker_name="rank1") |
| 1084 | worker.profiler.start.assert_called_once_with() |
| 1085 | |
| 1086 | |
| 1087 | @_requires_cuda_for_proton |
| 1088 | def test_gpu_worker_recreates_proton_profiler_for_each_run(): |
| 1089 | worker = object.__new__(Worker) |
| 1090 | worker.rank = 1 |
| 1091 | worker.local_rank = 1 |
| 1092 | worker.profiler = None |
| 1093 | worker.profiler_config = MagicMock(profiler="proton") |
| 1094 | |
| 1095 | with ( |
| 1096 | patch( |
| 1097 | "vllm.distributed.utils.get_worker_rank_suffix", |
| 1098 | return_value="rank1", |
| 1099 | ), |
| 1100 | patch("vllm.profiler.wrapper.ProtonProfilerWrapper") as wrapper, |
| 1101 | ): |
| 1102 | wrapper.return_value.has_cuda_graph_session = False |
| 1103 | Worker.profile(worker, profile_prefix="first") |
| 1104 | Worker.profile(worker, is_start=False) |
| 1105 | Worker.profile(worker, profile_prefix="second") |
| 1106 | |
| 1107 | assert wrapper.call_args_list == [ |
| 1108 | call(worker.profiler_config, worker_name="first_rank1"), |
| 1109 | call(worker.profiler_config, worker_name="second_rank1"), |
| 1110 | ] |
| 1111 | assert wrapper.return_value.start.call_count == 2 |
| 1112 | |
| 1113 | |
| 1114 | @_requires_cuda_for_proton |
| 1115 | def test_gpu_worker_reuses_cuda_graph_proton_session(): |
| 1116 | worker = MagicMock() |
| 1117 | worker.rank = 1 |
| 1118 | worker.profiler = MagicMock(spec=ProtonProfilerWrapper) |
| 1119 | worker.profiler.has_cuda_graph_session = True |
| 1120 | worker.profiler_config.profiler = "proton" |
| 1121 | |
| 1122 | with patch( |
| 1123 | "vllm.distributed.utils.get_worker_rank_suffix", |
| 1124 | return_value="rank1", |
| 1125 | ): |
| 1126 | Worker.profile(worker, profile_prefix="first") |
| 1127 | Worker.profile(worker, is_start=False) |
| 1128 | |
| 1129 | worker.profiler.set_output_name.assert_called_once_with("first_rank1") |
| 1130 | worker.profiler.start.assert_called_once_with() |
| 1131 | worker.profiler.stop.assert_called_once_with() |
| 1132 | |
| 1133 | |
| 1134 | @_requires_cuda_for_proton |
| 1135 | @pytest.mark.parametrize("runner", ["attribution_off", "v1", "no_capture"]) |
| 1136 | def test_proton_not_initialized_without_capture(runner): |
| 1137 | worker = MagicMock() |
| 1138 | worker.profiler = None |
| 1139 | worker.profiler_config.profiler = "proton" |
| 1140 | worker.profiler_config.proton_graph_attribution = runner != "attribution_off" |
| 1141 | worker.use_v2_model_runner = runner != "v1" |
| 1142 | worker.model_runner.needs_cudagraph_capture.return_value = runner != "no_capture" |
| 1143 | |
| 1144 | context = Worker._get_cudagraph_capture_context(worker) |
| 1145 | |
| 1146 | assert worker.profiler is None |
| 1147 | with context: |
| 1148 | pass |
| 1149 | |
| 1150 | |
| 1151 | @_requires_cuda_for_proton |
| 1152 | def test_proton_initializes_before_cuda_graph_capture(): |
| 1153 | class FakeProtonProfiler: |
| 1154 | def __init__(self, config, worker_name): |
| 1155 | self.config = config |
| 1156 | self.worker_name = worker_name |
| 1157 | self.capture_context = nullcontext() |
| 1158 | |
| 1159 | def capture_cuda_graphs(self): |
| 1160 | return self.capture_context |
| 1161 | |
| 1162 | worker = MagicMock() |
| 1163 | worker.rank = 2 |
| 1164 | worker.profiler = None |
| 1165 | worker.profiler_config.profiler = "proton" |
| 1166 | worker.profiler_config.proton_graph_attribution = True |
| 1167 | worker.use_v2_model_runner = True |
| 1168 | worker.model_runner.needs_cudagraph_capture.return_value = True |
| 1169 | |
| 1170 | with ( |
| 1171 | patch( |
| 1172 | "vllm.distributed.utils.get_worker_rank_suffix", |
| 1173 | return_value="rank2", |
| 1174 | ), |
| 1175 | patch( |
| 1176 | "vllm.profiler.wrapper.ProtonProfilerWrapper", |
| 1177 | FakeProtonProfiler, |
| 1178 | ), |
| 1179 | ): |
| 1180 | context = Worker._get_cudagraph_capture_context(worker) |
| 1181 | |
| 1182 | assert worker.profiler.config is worker.profiler_config |
| 1183 | assert worker.profiler.worker_name == "rank2" |
| 1184 | assert context is worker.profiler.capture_context |
| 1185 | |
| 1186 | |
| 1187 | @_requires_cuda_for_proton |
| 1188 | @_requires_no_injected_cupti_tool |
| 1189 | @pytest.mark.parametrize("context", ["shadow", "python"]) |
| 1190 | @pytest.mark.parametrize("output_format", ["hatchet", "hatchet_msgpack"]) |
| 1191 | def test_proton_cuda_graph_replay_attribution_on_gpu(tmp_path, context, output_format): |
| 1192 | """Both intervals contain replay kernels, without capture-only activity.""" |
| 1193 | import json |
| 1194 | |
| 1195 | import torch |
| 1196 | import triton |
| 1197 | import triton.profiler as proton |
| 1198 | from packaging.version import Version |
| 1199 | |
| 1200 | if Version(triton.__version__) < Version("3.7"): |
| 1201 | pytest.skip("Graph attribution requires Triton >= 3.7") |
| 1202 | |
| 1203 | wrapper = ProtonProfilerWrapper( |
| 1204 | ProfilerConfig( |
| 1205 | profiler="proton", |
| 1206 | proton_profiler_dir=str(tmp_path), |
| 1207 | proton_graph_attribution=True, |
| 1208 | proton_context=context, |
| 1209 | proton_output_format=output_format, |
| 1210 | ), |
| 1211 | worker_name="gpu", |
| 1212 | ) |
| 1213 | worker = SimpleNamespace( |
| 1214 | profiler=wrapper, |
| 1215 | use_v2_model_runner=True, |
| 1216 | ) |
| 1217 | x = torch.ones(1024, device="cuda") |
| 1218 | graph = torch.cuda.CUDAGraph() |
| 1219 | |
| 1220 | def capture_only(): |
| 1221 | x.add_(1) |
| 1222 | |
| 1223 | def captured_add(): |
| 1224 | x.add_(1) |
| 1225 | |
| 1226 | def kernel_metrics(value): |
| 1227 | if isinstance(value, list): |
| 1228 | return [metric for child in value for metric in kernel_metrics(child)] |
| 1229 | if isinstance(value, dict): |
| 1230 | metrics = value.get("metrics", {}) |
| 1231 | return ([metrics] if metrics.get("count", 0) else []) + [ |
| 1232 | metric for child in value.values() for metric in kernel_metrics(child) |
| 1233 | ] |
| 1234 | return [] |
| 1235 | |
| 1236 | try: |
| 1237 | with Worker._get_cudagraph_capture_context(worker): |
| 1238 | with proton.scope("capture_only"): |
| 1239 | capture_only() |
| 1240 | torch.accelerator.synchronize() |
| 1241 | with torch.cuda.graph(graph), proton.scope("captured_add"): |
| 1242 | captured_add() |
| 1243 | for run in range(2): |
| 1244 | wrapper.start() |
| 1245 | with wrapper.annotate_context_manager(f"replay_{run}"): |
| 1246 | graph.replay() |
| 1247 | wrapper.stop() |
| 1248 | (path,) = tmp_path.glob(f"*_run{run}.{output_format}") |
| 1249 | if output_format == "hatchet_msgpack": |
| 1250 | import msgpack |
| 1251 | |
| 1252 | data = msgpack.unpackb(path.read_bytes()) |
| 1253 | else: |
| 1254 | data = json.loads(path.read_text()) |
| 1255 | serialized = json.dumps(data) |
| 1256 | assert "<captured_at>" in serialized |
| 1257 | assert "captured_add" in serialized |
| 1258 | assert "capture_only" not in serialized |
| 1259 | assert f"replay_{1 - run}" not in serialized |
| 1260 | metrics = kernel_metrics(data) |
| 1261 | assert sum(metric["count"] for metric in metrics) == 1 |
| 1262 | assert sum(metric["time (ns)"] for metric in metrics) > 0 |
| 1263 | torch.testing.assert_close(x, torch.full_like(x, 4)) |
| 1264 | finally: |
| 1265 | wrapper.shutdown() |
| 1266 | assert not list(tmp_path.glob(".proton_cuda_graph_session*")) |