diff --git a/tests/v1/worker/test_gpu_worker.py b/tests/v1/worker/test_gpu_worker.py index 375e43810f..722932636b 100644 --- a/tests/v1/worker/test_gpu_worker.py +++ b/tests/v1/worker/test_gpu_worker.py @@ -258,3 +258,32 @@ def test_execute_model_waits_previous_pp_send_before_forward( assert log == ["wait:prev-tensor", "forward", "isend"] assert worker._pp_send_work == [tensor_handle] + + +def test_jit_monitor_activation_follows_enable_jit_warmup( + monkeypatch: pytest.MonkeyPatch, +): + """The post-warmup JIT monitor must stay off when JIT warmup is disabled + (e.g. by enforce_eager): runtime compilation is then expected, and + warning/erroring on it would be noise.""" + from vllm.utils import jit_monitor + + calls = [] + monkeypatch.setattr(jit_monitor, "activate", lambda **kwargs: calls.append(kwargs)) + + def worker(enable_jit_warmup): + return SimpleNamespace( + vllm_config=SimpleNamespace( + kernel_config=SimpleNamespace(enable_jit_warmup=enable_jit_warmup) + ), + observability_config=SimpleNamespace( + jit_monitor_mode="warn", jit_monitor_verbose=False + ), + ) + + gpu_worker.Worker._maybe_activate_jit_monitor(worker(True)) + assert calls == [{"mode": "warn", "verbose": False}] + + calls.clear() + gpu_worker.Worker._maybe_activate_jit_monitor(worker(False)) + assert calls == [] diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 7e4831447c..35c986b131 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -966,12 +966,7 @@ class Worker(WorkerBase): # All warmup is done — start monitoring for unexpected JIT # compilations that would cause latency spikes during inference. - from vllm.utils.jit_monitor import activate as activate_jit_monitor - - activate_jit_monitor( - mode=self.observability_config.jit_monitor_mode, - verbose=self.observability_config.jit_monitor_verbose, - ) + self._maybe_activate_jit_monitor() # Freeze the worker heap so the GC won't scan static objects # (model weights, KV caches, CUDA graphs) during inference. @@ -987,6 +982,20 @@ class Worker(WorkerBase): encoder=self.compilation_config.encoder_compilation_time, ) + def _maybe_activate_jit_monitor(self) -> None: + # When JIT warmup is disabled (e.g. enforce_eager), runtime JIT + # compilation is expected, so monitoring would only produce noise + # (or spurious errors in "error" mode). + if not self.vllm_config.kernel_config.enable_jit_warmup: + return + + from vllm.utils.jit_monitor import activate as activate_jit_monitor + + activate_jit_monitor( + mode=self.observability_config.jit_monitor_mode, + verbose=self.observability_config.jit_monitor_verbose, + ) + def _get_cudagraph_capture_context(self) -> AbstractContextManager[None]: """Let the configured profiler observe CUDA graph capture.""" if not self.use_v2_model_runner: