diff --git a/buildkite/pipeline_generator/buildkite_step.py b/buildkite/pipeline_generator/buildkite_step.py index 217f7ae0..12896b7e 100644 --- a/buildkite/pipeline_generator/buildkite_step.py +++ b/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 diff --git a/buildkite/tests/pipeline_generator/test_step.py b/buildkite/tests/pipeline_generator/test_step.py index b4227e68..4a18ced5 100644 --- a/buildkite/tests/pipeline_generator/test_step.py +++ b/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",