Enable fail-fast logit NaN detection in vLLM CI

buildkite/pipeline_generator/buildkite_step.py

@@ -35,6 +35,11 @@
 
 SKIP_TIMEOUT_ENV_VAR = "SKIP_TIMEOUT"
 EXIT_STATUS_NEGATIVE_ONE_RETRY = {"exit_status": -1, "limit": 1}
+DEFAULT_CI_ENV = {"VLLM_RAISE_ON_LOGIT_NANS": "1"}
+
+
+def _with_default_ci_env(env: Optional[Dict[str, str]]) -> Dict[str, str]:
+    return {**DEFAULT_CI_ENV, **(env or {})}
 
 
 # Self-contained poll of the pre-commit GitHub Actions check run. Baked with the
@@ -553,11 +558,10 @@ def convert_group_step_to_buildkite_step(
                 depends_on=step.depends_on,
                 soft_fail=step.soft_fail,
                 agents=_get_step_agents(step),
+                env=_with_default_ci_env(step.env),
                 priority=1000 if os.getenv("PRIORITY", "") == "HIGH" else 0,
             )
 
-            if step.env:
-                buildkite_step.env = step.env
             if step.retry:
                 buildkite_step.retry = step.retry
             buildkite_step.retry = ensure_exit_status_negative_one_retry(
@@ -810,7 +814,7 @@ def _create_amd_step(
         num_nodes=num_nodes,
         agent_tags=agent_tags,
     )
-    return BuildkiteCommandStep(
+    buildkite_step = BuildkiteCommandStep(
         **options,
         key=key,
         soft_fail=soft_fail or False,
@@ -819,3 +823,5 @@ def _create_amd_step(
             get_amd_timeout_in_minutes(timeout_in_minutes)
         ),
     )
+    buildkite_step.env = _with_default_ci_env(buildkite_step.env)
+    return buildkite_step

buildkite/tests/pipeline_generator/test_step.py

@@ -165,6 +165,56 @@ def test_generated_steps_retry_when_the_agent_is_lost():
     }
 
 
+def test_generated_steps_raise_on_logit_nans_by_default():
+    step = Step(
+        label="NaN detection",
+        group="Correctness",
+        commands=["pytest tests/basic.py"],
+    )
+
+    command_step = next(
+        rendered_step
+        for rendered_step in _render_single_step(step).steps
+        if isinstance(rendered_step, buildkite_step.BuildkiteCommandStep)
+    )
+
+    assert command_step.env["VLLM_RAISE_ON_LOGIT_NANS"] == "1"
+
+
+def test_step_environment_can_disable_raise_on_logit_nans():
+    step = Step(
+        label="Expected NaNs",
+        group="Correctness",
+        commands=["pytest tests/expected_nans.py"],
+        env={"VLLM_RAISE_ON_LOGIT_NANS": "0"},
+    )
+
+    command_step = next(
+        rendered_step
+        for rendered_step in _render_single_step(step).steps
+        if isinstance(rendered_step, buildkite_step.BuildkiteCommandStep)
+    )
+
+    assert command_step.env["VLLM_RAISE_ON_LOGIT_NANS"] == "0"
+
+
+def test_amd_steps_raise_on_logit_nans_by_default():
+    step = Step(
+        label="AMD NaN detection",
+        group="Correctness",
+        device="mi300_4",
+        commands=["pytest tests/basic.py"],
+    )
+
+    command_step = next(
+        rendered_step
+        for rendered_step in _render_single_step(step).steps
+        if isinstance(rendered_step, buildkite_step.BuildkiteCommandStep)
+    )
+
+    assert command_step.env["VLLM_RAISE_ON_LOGIT_NANS"] == "1"
+
+
 def test_agent_lost_retry_preserves_step_retry_conditions():
     step = Step(
         label="Agent retry with policy",