[CI][AMD] Support build-scoped ROCm images

buildkite/bootstrap-amd.sh

@@ -136,7 +136,16 @@ get_diff_main() {
 # upload_pipeline: render and upload the Buildkite pipeline YAML
 # ---------------------------------------------------------------------------
 upload_pipeline() {
+    local rocm_stable_image_promotion=0
+
+    if [[ -f ".buildkite/scripts/rocm/promote-stable-images.sh" ]] \
+        && grep -q 'CI_BASE_IMAGE_TAG_BUILD_REF' \
+            .buildkite/scripts/ci-bake-rocm.sh 2>/dev/null; then
+        rocm_stable_image_promotion=1
+    fi
+
     echo "Uploading pipeline..."
+    echo "ROCm stable image promotion: ${rocm_stable_image_promotion}"
     # Install minijinja
     ls .buildkite || buildkite-agent annotate --style error 'Please merge upstream main branch for buildkite CI'
     curl -sSfL https://github.com/mitsuhiko/minijinja/releases/download/2.3.1/minijinja-cli-installer.sh | sh
@@ -181,6 +190,7 @@ upload_pipeline() {
             -D vllm_ci_branch="$VLLM_CI_BRANCH" \
             -D rocm_base_refresh_skip="${ROCM_BASE_REFRESH_SKIP:-0}" \
             -D rocm_base_refresh_force="${ROCM_BASE_REFRESH_FORCE:-0}" \
+            -D rocm_stable_image_promotion="${rocm_stable_image_promotion}" \
             | sed '/^[[:space:]]*$/d' \
             > pipeline.yaml
     )

buildkite/ci_selector/ci_selector/handwritten.py

@@ -410,6 +410,9 @@
     "agent_pool",
     "concurrency",
     "concurrency_group",
+    # Emitted as Buildkite's `if`: evaluated on upload, after step selection.
+    # It can only restrict execution, so it does not force-select a step.
+    "if_condition",
     # Buildkite built-in, absent from the generator's Step model, so it is
     # dropped before reaching a command. Governs ordering after a failure.
     "allow_dependency_failure",

buildkite/ci_selector/tests/test_buildkite.py

@@ -11,6 +11,7 @@
 
 import pytest
 import regex as re
+from amd import AMD_STABLE_IMAGE_PROMOTION_STEP_KEY, supports_stable_image_promotion
 from ci_selector.codemap.pipeline.buildkite import (
     _expand_mirror,
     load_pipeline_configs,
@@ -121,8 +122,36 @@ def test_a_mirror_publishes_the_hardware_as_a_prefix(loaded):
     assert plain.buildkite_key == "lora"
 
 
+def _assert_amd_auto_run_semantics(rocm, promotion_supported):
+    # Only the promotion step is optional across companion vLLM checkouts.
+    # Existing keys still have to exist, and a capable checkout must include
+    # promotion. Whenever promotion is present, verify its always-run shortcut.
+    for key in AMD_ALWAYS_RUN_STEP_KEYS:
+        if (
+            key == AMD_STABLE_IMAGE_PROMOTION_STEP_KEY
+            and key not in rocm
+            and not promotion_supported
+        ):
+            continue
+        assert key in rocm, drift_message(
+            f"AMD_ALWAYS_RUN_STEP_KEYS names a step that no longer exists: {key}",
+            "These steps declare no dependencies, so always_runs is the only "
+            "thing that ever selects them. A renamed key means we stop naming "
+            "the step and CI stops running it.",
+            "we re-export this list from the generator's amd.py, so there is "
+            "nothing to edit here: fix the key in amd.py, or in vLLM's rocm "
+            "yaml if that is where it was renamed",
+        )
+        assert rocm[key].always_runs, drift_message(
+            f"{key} is in AMD_ALWAYS_RUN_STEP_KEYS but does not read as always-run.",
+            "It would only be selected if some rule happened to pick it, and "
+            "these steps have no dependencies for a rule to match on.",
+            "check Step.always_runs against ci-infra's _step_should_run",
+        )
+
+
 @pytest.mark.drift
-def test_auto_run_semantics(loaded):
+def test_auto_run_semantics(loaded, vllm_repo, monkeypatch):
     """always_runs is the generator's key shortcut (image-build*/AMD list),
     NOT "has no deps": a no-deps step without the shortcut never auto-runs."""
     _, steps, _ = loaded
@@ -153,27 +182,12 @@ def test_auto_run_semantics(loaded):
     assert by_key["cpu-arm64-image-build"].source_file_dependencies is None, prefix
     assert not by_key["cpu-arm64-image-build"].always_runs, prefix
     assert not by_key["arm64-image-build"].always_runs, prefix
-    # Membership asserted: these steps declare no deps, so always_runs
-    # is the only thing that ever selects them, and this is the
-    # only guard tying the generator's list to vLLM's yaml. Under `if key in
-    # rocm` a step vanishing from the yaml would skip the check in silence.
+    # Read capability from the vLLM checkout, not the ci-infra working directory.
     rocm = {s.key: s for s in steps["vllm_rocm_ci"] if s.key}
-    for key in AMD_ALWAYS_RUN_STEP_KEYS:
-        assert key in rocm, drift_message(
-            f"AMD_ALWAYS_RUN_STEP_KEYS names a step that no longer exists: {key}",
-            "These steps declare no dependencies, so always_runs is the only "
-            "thing that ever selects them. A renamed key means we stop naming "
-            "the step and CI stops running it.",
-            "we re-export this list from the generator's amd.py, so there is "
-            "nothing to edit here: fix the key in amd.py, or in vLLM's rocm "
-            "yaml if that is where it was renamed",
-        )
-        assert rocm[key].always_runs, drift_message(
-            f"{key} is in AMD_ALWAYS_RUN_STEP_KEYS but does not read as always-run.",
-            "It would only be selected if some rule happened to pick it, and "
-            "these steps have no dependencies for a rule to match on.",
-            "check Step.always_runs against ci-infra's _step_should_run",
-        )
+    with monkeypatch.context() as checkout:
+        checkout.chdir(vllm_repo)
+        promotion_supported = supports_stable_image_promotion()
+    _assert_amd_auto_run_semantics(rocm, promotion_supported)
     no_deps_no_shortcut = [
         s
         for s in steps["vllm_ci"]
@@ -367,6 +381,51 @@ def _parent_step(**kw):
     return Step(**{**base, **kw})
 
 
+def _amd_always_run_steps():
+    return {key: _parent_step(key=key) for key in AMD_ALWAYS_RUN_STEP_KEYS}
+
+
+@pytest.mark.parametrize("promotion_supported", [False, True])
+@pytest.mark.parametrize("promotion_present", [False, True])
+def test_amd_promotion_membership_tracks_checkout_capability(
+    promotion_supported, promotion_present
+):
+    rocm = _amd_always_run_steps()
+    if not promotion_present:
+        del rocm[AMD_STABLE_IMAGE_PROMOTION_STEP_KEY]
+    if promotion_supported and not promotion_present:
+        with pytest.raises(AssertionError, match=AMD_STABLE_IMAGE_PROMOTION_STEP_KEY):
+            _assert_amd_auto_run_semantics(rocm, promotion_supported)
+    else:
+        _assert_amd_auto_run_semantics(rocm, promotion_supported)
+
+
+@pytest.mark.parametrize(
+    "missing_key",
+    sorted(AMD_ALWAYS_RUN_STEP_KEYS - {AMD_STABLE_IMAGE_PROMOTION_STEP_KEY}),
+)
+def test_old_checkout_still_requires_existing_amd_steps(missing_key):
+    rocm = _amd_always_run_steps()
+    del rocm[AMD_STABLE_IMAGE_PROMOTION_STEP_KEY]
+    del rocm[missing_key]
+    with pytest.raises(AssertionError, match=missing_key):
+        _assert_amd_auto_run_semantics(rocm, promotion_supported=False)
+
+
+@pytest.mark.parametrize("promotion_supported", [False, True])
+def test_present_amd_promotion_must_read_as_always_run(
+    monkeypatch, promotion_supported
+):
+    rocm = _amd_always_run_steps()
+    monkeypatch.setattr(
+        Step,
+        "always_runs",
+        property(lambda step: step.key != AMD_STABLE_IMAGE_PROMOTION_STEP_KEY),
+    )
+    with pytest.raises(AssertionError, match="does not read as always-run"):
+        _assert_amd_auto_run_semantics(rocm, promotion_supported)
+
+
 def test_a_stray_mirror_key_is_recorded_against_the_mirror_not_the_parent():
     """`_expand_mirror` files unknown keys under the variant's id, so an AMD-only
     typo forces the AMD step and leaves its NVIDIA parent free."""

buildkite/ci_selector/tests/test_guards.py

@@ -123,6 +123,35 @@ def test_unknown_field_force_selected():
     assert step.step_id in pf.force_select
 
 
+@pytest.mark.parametrize(
+    "condition", ["true", "false", "build.branch == pipeline.default_branch"]
+)
+def test_buildkite_condition_does_not_force_selection(tmp_path, condition):
+    from ci_selector.codemap.pipeline.buildkite import load_steps
+
+    (tmp_path / "jobs").mkdir()
+    (tmp_path / "jobs" / "test.yaml").write_text(
+        "steps:\n"
+        "  - label: Conditional test\n"
+        "    key: conditional-test\n"
+        "    commands: [pytest tests/x]\n"
+        f"    if_condition: '{condition}'\n"
+    )
+    config = PipelineConfig("p", ".buildkite/x.yaml", ["jobs"], [], [])
+    report = LoadReport()
+    steps = load_steps(tmp_path, config, report)
+    (step,) = steps
+    pf = run_preflight(
+        tmp_path,
+        [_pipe(steps, _covered_targets(step))],
+        _healthy_full(),
+        report,
+    )
+
+    assert report.unknown_fields == {}
+    assert step.step_id not in pf.force_select
+
+
 def test_duplicate_id_force_selected():
     a, b = _step("dup"), _step("dup")
     report = LoadReport()

buildkite/pipeline_generator/amd.py

@@ -1,6 +1,7 @@
 import os
 from copy import deepcopy
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Dict, List, Mapping, Optional, TypedDict
 
 from constants import AgentQueue, DeviceType
@@ -9,6 +10,16 @@
 AMD_STABLE_CI_BASE_IMAGE = "rocm/vllm-dev:ci_base"
 AMD_NATIVE_BASE_IMAGE = "rocm/vllm-dev:ci_base-build-$BUILDKITE_BUILD_ID"
 AMD_CI_IMAGE = "rocm/vllm-ci:build-$BUILDKITE_BUILD_ID"
+AMD_ROCM_PROMOTION_SCRIPT = Path(".buildkite/scripts/rocm/promote-stable-images.sh")
+AMD_ROCM_BASE_PRODUCER = Path(".buildkite/scripts/ci-bake-rocm.sh")
+
+
+def supports_stable_image_promotion() -> bool:
+    return (
+        AMD_ROCM_PROMOTION_SCRIPT.is_file()
+        and AMD_ROCM_BASE_PRODUCER.is_file()
+        and "CI_BASE_IMAGE_TAG_BUILD_REF" in AMD_ROCM_BASE_PRODUCER.read_text()
+    )
 
 
 AMD_ARTIFACT_GLOB = "artifacts/vllm-rocm-install/vllm-rocm-install.tar.gz"
@@ -20,7 +31,7 @@
 AMD_HF_HOME = "/home/buildkite-agent/huggingface"
 AMD_NATIVE_WORKSPACE = "/vllm-workspace"
 AMD_NATIVE_WORKSPACE_VOLUME = "vllm-workspace"
-AMD_NATIVE_SHM_SIZE = "16Gi"
+AMD_NATIVE_SHM_SIZE = "32Gi"
 AMD_NATIVE_POD_IDENTITY_ENV = {
     "VLLM_CI_K8S_POD_NAME": "metadata.name",
     "VLLM_CI_K8S_NAMESPACE": "metadata.namespace",
@@ -31,13 +42,15 @@
 )
 AMD_STANDARD_TIMEOUT_MINUTES = 3 * 60
 AMD_ROCM_BASE_REFRESH_STEP_KEY = "refresh-rocm-base-amd"
+AMD_STABLE_IMAGE_PROMOTION_STEP_KEY = "promote-stable-rocm-images-amd"
 AMD_ROCM_BASE_DOCKERFILE = "docker/Dockerfile.rocm_base"
 AMD_ROCM_BASE_REFRESH_TIMEOUT_MINUTES = 9 * 60
 AMD_ROCM_BASE_REFRESH_NOOP_TIMEOUT_MINUTES = 15
 AMD_ALWAYS_RUN_STEP_KEYS = frozenset(
     {
         "ensure-ci-base-amd",
         AMD_ROCM_BASE_REFRESH_STEP_KEY,
+        AMD_STABLE_IMAGE_PROMOTION_STEP_KEY,
     }
 )
 AMD_STACK_ERROR_RETRY = {"signal_reason": "stack_error", "limit": 1}
@@ -57,6 +70,12 @@
 ROCM_DEBUG_AGENT_LIB = "/opt/rocm/lib/librocm-debug-agent.so.2"
 
 
+def get_amd_ci_base_image(*, dind: bool) -> str:
+    if dind and supports_stable_image_promotion():
+        return AMD_NATIVE_BASE_IMAGE
+    return AMD_STABLE_CI_BASE_IMAGE if dind else AMD_NATIVE_BASE_IMAGE
+
+
 def get_amd_timeout_in_minutes(timeout_in_minutes: Optional[int]) -> int:
     """Default and cap standard AMD jobs at three hours."""
     requested_timeout = (
@@ -292,6 +311,7 @@ def _get_amd_env(
     env["BUILDKITE_ARTIFACT_UPLOAD_SKIP_SYMLINKS"] = "true"
     env["VLLM_CI_DIAGNOSTICS_DIR"] = AMD_DIAGNOSTICS_DIR
     env["VLLM_CI_EXPECTED_GPU_COUNT"] = str(gpu_count)
+    ci_base_image = get_amd_ci_base_image(dind=dind)
     if not dind:
         # Native agents have no DinD sidecar, so Docker hook inputs must not
         # escape into this execution mode.
@@ -306,7 +326,7 @@ def _get_amd_env(
         env.setdefault("HF_HUB_ETAG_TIMEOUT", "60")
         env.update(
             {
-                "VLLM_CI_BASE_IMAGE": AMD_NATIVE_BASE_IMAGE,
+                "VLLM_CI_BASE_IMAGE": ci_base_image,
                 "VLLM_CI_USE_ARTIFACTS": "1",
                 "VLLM_CI_ARTIFACT_GLOB": AMD_ARTIFACT_GLOB,
                 "VLLM_CI_RESULTS_ROOT": AMD_RESULTS_ROOT,
@@ -323,8 +343,8 @@ def _get_amd_env(
         env.update(
             {
                 "DOCKER_BUILDKIT": "1",
-                "DOCKER_IMAGE_NAME": AMD_STABLE_CI_BASE_IMAGE,
-                "VLLM_CI_BASE_IMAGE": AMD_STABLE_CI_BASE_IMAGE,
+                "DOCKER_IMAGE_NAME": ci_base_image,
+                "VLLM_CI_BASE_IMAGE": ci_base_image,
                 "VLLM_CI_FALLBACK_IMAGE": AMD_CI_IMAGE,
                 "VLLM_CI_USE_ARTIFACTS": "1",
                 "VLLM_CI_ARTIFACT_GLOB": AMD_ARTIFACT_GLOB,
@@ -441,6 +461,7 @@ def build_amd_step_options(
     gpu_count = resolve_amd_gpu_count(device, num_devices, no_gpu)
     plugins = None
     if not dind:
+        ci_base_image = get_amd_ci_base_image(dind=False)
         container_env = {
             "AMD_CI_RUNTIME": "native",
             "NATIVE_CI": "true",
@@ -452,7 +473,7 @@ def build_amd_step_options(
         }
         plugins = [
             get_amd_k8s_plugin(
-                image=AMD_NATIVE_BASE_IMAGE,
+                image=ci_base_image,
                 gpu_count=gpu_count,
                 workspace=AMD_NATIVE_WORKSPACE,
                 workspace_volume_name=AMD_NATIVE_WORKSPACE_VOLUME,

buildkite/pipeline_generator/buildkite_step.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, Field, field_validator, model_validator
 from typing import Dict, List, Optional, Any, Union, Literal
 from copy import deepcopy
 import os
@@ -226,11 +226,27 @@ class BuildkiteCommandStep(BaseModel):
     env: Optional[Dict[str, str]] = None
     artifact_paths: Optional[List[str]] = None
     parallelism: Optional[int] = None
-    concurrency: Optional[int] = None
-    concurrency_group: Optional[str] = None
+    concurrency: Optional[int] = Field(default=None, gt=0, strict=True)
+    concurrency_group: Optional[str] = Field(default=None, min_length=1)
+    if_condition: Optional[str] = Field(default=None, serialization_alias="if")
     timeout_in_minutes: Optional[int] = None
     priority: Optional[int] = None
 
+    @field_validator("concurrency_group")
+    @classmethod
+    def validate_concurrency_group(cls, value: Optional[str]) -> Optional[str]:
+        if value is not None and not value.strip():
+            raise ValueError("'concurrency_group' must be a nonempty string.")
+        return value
+
+    @model_validator(mode="after")
+    def validate_concurrency_pair(self):
+        if (self.concurrency is None) != (self.concurrency_group is None):
+            raise ValueError(
+                "'concurrency' and 'concurrency_group' must be defined together."
+            )
+        return self
+
     def to_yaml(self):
         return {
             "label": self.label,
@@ -247,6 +263,7 @@ def to_yaml(self):
             "parallelism": self.parallelism,
             "concurrency": self.concurrency,
             "concurrency_group": self.concurrency_group,
+            "if": self.if_condition,
             "timeout_in_minutes": self.timeout_in_minutes,
             "priority": self.priority,
         }
@@ -669,6 +686,7 @@ def convert_group_step_to_buildkite_step(
                     parallelism=step.parallelism,
                     concurrency=step.concurrency,
                     concurrency_group=step.concurrency_group,
+                    if_condition=step.if_condition,
                     timeout_in_minutes=step.timeout_in_minutes,
                     agent_tags=step.agent_tags,
                 )
@@ -696,6 +714,7 @@ def convert_group_step_to_buildkite_step(
                 priority=1000 if os.getenv("PRIORITY", "") == "HIGH" else 0,
                 concurrency=step.concurrency,
                 concurrency_group=step.concurrency_group,
+                if_condition=step.if_condition,
             )
 
             if step.env:
@@ -823,6 +842,7 @@ def convert_group_step_to_buildkite_step(
                     concurrency_group=amd.get(
                         "concurrency_group", step.concurrency_group
                     ),
+                    if_condition=step.if_condition,
                     timeout_in_minutes=amd.get("timeout_in_minutes"),
                     agent_tags=amd.get("agent_tags"),
                     display_label=amd.get("label"),
@@ -963,6 +983,7 @@ def _create_amd_step(
     parallelism: Optional[int],
     concurrency: Optional[int],
     concurrency_group: Optional[str],
+    if_condition: Optional[str],
     key: str,
     timeout_in_minutes: Optional[int] = None,
     agent_tags: Optional[Dict[str, str]] = None,
@@ -991,6 +1012,7 @@ def _create_amd_step(
         parallelism=parallelism,
         concurrency=concurrency,
         concurrency_group=concurrency_group,
+        if_condition=if_condition,
         timeout_in_minutes=_get_timeout_in_minutes(
             get_amd_timeout_in_minutes(timeout_in_minutes)
         ),