[Cleanup] Refactor profiling env vars into a CLI config (#29912)

Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

tests/v1/worker/test_gpu_profiler.py

@@ -2,8 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import pytest
 
-import vllm.envs as envs
-from vllm.profiler.gpu_profiler import WorkerProfiler
+from vllm.config import ProfilerConfig
+from vllm.profiler.wrapper import WorkerProfiler
 
 
 class ConcreteWorkerProfiler(WorkerProfiler):
@@ -11,11 +11,11 @@ class ConcreteWorkerProfiler(WorkerProfiler):
     A basic implementation of a worker profiler for testing purposes.
     """
 
-    def __init__(self):
+    def __init__(self, profiler_config: ProfilerConfig):
         self.start_call_count = 0
         self.stop_call_count = 0
         self.should_fail_start = False
-        super().__init__()
+        super().__init__(profiler_config)
 
     def _start(self) -> None:
         if self.should_fail_start:
@@ -26,17 +26,19 @@ class ConcreteWorkerProfiler(WorkerProfiler):
         self.stop_call_count += 1
 
 
-@pytest.fixture(autouse=True)
-def reset_mocks():
-    """Fixture to reset mocks and env variables before each test."""
-    envs.VLLM_PROFILER_DELAY_ITERS = 0
-    envs.VLLM_PROFILER_MAX_ITERS = 0
+@pytest.fixture
+def default_profiler_config():
+    return ProfilerConfig(
+        profiler="torch",
+        torch_profiler_dir="/tmp/mock",
+        delay_iterations=0,
+        max_iterations=0,
+    )
 
 
-def test_immediate_start_stop():
+def test_immediate_start_stop(default_profiler_config):
     """Test standard start without delay."""
-    profiler = ConcreteWorkerProfiler()
-
+    profiler = ConcreteWorkerProfiler(default_profiler_config)
     profiler.start()
     assert profiler._running is True
     assert profiler._active is True
@@ -48,10 +50,10 @@ def test_immediate_start_stop():
     assert profiler.stop_call_count == 1
 
 
-def test_delayed_start():
+def test_delayed_start(default_profiler_config):
     """Test that profiler waits for N steps before actually starting."""
-    envs.VLLM_PROFILER_DELAY_ITERS = 2
-    profiler = ConcreteWorkerProfiler()
+    default_profiler_config.delay_iterations = 2
+    profiler = ConcreteWorkerProfiler(default_profiler_config)
 
     # User requests start
     profiler.start()
@@ -71,10 +73,10 @@ def test_delayed_start():
     assert profiler.start_call_count == 1
 
 
-def test_max_iterations():
+def test_max_iterations(default_profiler_config):
     """Test that profiler stops automatically after max iterations."""
-    envs.VLLM_PROFILER_MAX_ITERS = 2
-    profiler = ConcreteWorkerProfiler()
+    default_profiler_config.max_iterations = 2
+    profiler = ConcreteWorkerProfiler(default_profiler_config)
 
     profiler.start()
     assert profiler._running is True
@@ -95,12 +97,11 @@ def test_max_iterations():
     assert profiler.stop_call_count == 1
 
 
-def test_delayed_start_and_max_iters():
+def test_delayed_start_and_max_iters(default_profiler_config):
     """Test combined delayed start and max iterations."""
-    envs.VLLM_PROFILER_DELAY_ITERS = 2
-    envs.VLLM_PROFILER_MAX_ITERS = 2
-    profiler = ConcreteWorkerProfiler()
-
+    default_profiler_config.delay_iterations = 2
+    default_profiler_config.max_iterations = 2
+    profiler = ConcreteWorkerProfiler(default_profiler_config)
     profiler.start()
 
     # Step 1
@@ -127,9 +128,9 @@ def test_delayed_start_and_max_iters():
     assert profiler.stop_call_count == 1
 
 
-def test_idempotency():
+def test_idempotency(default_profiler_config):
     """Test that calling start/stop multiple times doesn't break logic."""
-    profiler = ConcreteWorkerProfiler()
+    profiler = ConcreteWorkerProfiler(default_profiler_config)
 
     # Double Start
     profiler.start()
@@ -142,10 +143,10 @@ def test_idempotency():
     assert profiler.stop_call_count == 1  # Should only stop once
 
 
-def test_step_inactive():
+def test_step_inactive(default_profiler_config):
     """Test that stepping while inactive does nothing."""
-    envs.VLLM_PROFILER_DELAY_ITERS = 2
-    profiler = ConcreteWorkerProfiler()
+    default_profiler_config.delay_iterations = 2
+    profiler = ConcreteWorkerProfiler(default_profiler_config)
 
     # Not started yet
     profiler.step()
@@ -155,9 +156,9 @@ def test_step_inactive():
     assert profiler.start_call_count == 0
 
 
-def test_start_failure():
+def test_start_failure(default_profiler_config):
     """Test behavior when the underlying _start method raises exception."""
-    profiler = ConcreteWorkerProfiler()
+    profiler = ConcreteWorkerProfiler(default_profiler_config)
     profiler.should_fail_start = True
 
     profiler.start()
@@ -168,9 +169,9 @@ def test_start_failure():
     assert profiler.start_call_count == 0  # Logic failed inside start
 
 
-def test_shutdown():
+def test_shutdown(default_profiler_config):
     """Test that shutdown calls stop only if running."""
-    profiler = ConcreteWorkerProfiler()
+    profiler = ConcreteWorkerProfiler(default_profiler_config)
 
     # Case 1: Not running
     profiler.shutdown()
@@ -182,10 +183,10 @@ def test_shutdown():
     assert profiler.stop_call_count == 1
 
 
-def test_mixed_delay_and_stop():
+def test_mixed_delay_and_stop(default_profiler_config):
     """Test manual stop during the delay period."""
-    envs.VLLM_PROFILER_DELAY_ITERS = 5
-    profiler = ConcreteWorkerProfiler()
+    default_profiler_config.delay_iterations = 5
+    profiler = ConcreteWorkerProfiler(default_profiler_config)
 
     profiler.start()
     profiler.step()