Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions tests/v1/worker/test_gpu_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -258,3 +258,32 @@ def run_model(scheduler_output, intermediate_tensors):

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 == []
21 changes: 15 additions & 6 deletions vllm/v1/worker/gpu_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -966,12 +966,7 @@ def compile_or_warm_up_model(self) -> CompilationTimes:

# 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.
Expand All @@ -987,6 +982,20 @@ def compile_or_warm_up_model(self) -> CompilationTimes:
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:
Expand Down
Loading