Files
vllm-omni/.buildkite/common/scripts/upload_pipeline.py
T

883 lines
33 KiB
Python
Raw Normal View History

#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Render and optionally upload Buildkite pipeline YAML with diff-aware logic.
Bootstrap mode (``bootstrap-upload-steps.yml``):
- Hook uploads the entry YAML (``pipeline.yml`` / ``pipeline-npu.yml``) with one step that runs
``upload_pipeline.py --upload <platform>/bootstrap-upload-steps.yml``.
- Injects ``if`` by step ``key`` from skip-ci and uploads child steps (image build, L2–L5 upload).
- Detect docs-only, pytest skip-mark-only, or combined skip-ci from git diff.
- When only CI level YAML changes, enable **L2/L3** upload steps for affected levels only.
Test pipeline mode (e.g. test-merge.yml, test-nightly.yml, test-weekly.yml):
- Drop steps whose ``source_file_dependencies`` do not match changed files
(string/list keys from ci_source_file_dependencies.yml, or inline path
prefixes). Filtering applies on **PR label** uploads and on post-merge
``main`` L3 uploads. Scheduled ``main`` uploads that pass ``--all``
(NIGHTLY/WEEKLY) or ``--e2e`` keep every selected job and still strip
the field. If no job-key prefix matches, a change to the pipeline YAML
being uploaded or to a path under the ``source_filter_fallback`` registry
key keeps every job so command, env, and hardware edits can be validated
before merge. When any job-key prefix already matches, normal filtering
wins.
- Expand uploader-only ``mirror_hardwares`` into ``agents`` (+ optional ``image``
for NPU) + ``plugins`` (see ci_mirror_hardwares.yml).
- Omit ``mirror_hardwares`` to compose ``{chip}_{n}`` from pytest ``-m`` SKU
markers plus ``cards_n`` (max of several positives; ``not cards_1`` with no
positive ``cards_*`` uses that chip's highest existing preset). An explicit
preset string (``h100_2``) is used as-is (marks are ignored). Names that
are not keys in ``ci_mirror_hardwares.yml`` fail the upload.
Inferred chip: unset ``MIRROR_HW`` matches ``H100`` or ``L4`` in ``-m``
(both → H100); no match skips the step. ``MIRROR_HW=b200`` matches
``B200`` in ``-m`` (otherwise skipped). ``-m`` is not rewritten.
``MIRROR_HW`` must be empty or ``b200``; unknown values fail the upload.
A CUDA preset string such as ``h100_4`` is omitted when ``MIRROR_HW=b200``.
Usage:
python3 upload_pipeline.py [--upload] [--all | --e2e] <pipeline.yml>
Requires PyYAML (``pip install pyyaml``); installs it automatically when missing.
"""
from __future__ import annotations
import argparse
import copy
import os
import re
import subprocess
import sys
from functools import lru_cache
from pathlib import Path
from typing import Any, Literal
try:
import yaml
except ModuleNotFoundError:
subprocess.run(
[sys.executable, "-m", "pip", "install", "-q", "pyyaml"],
check=True,
)
import yaml
from skip_ci import (
ROOT,
resolve_ci_context_from_git,
)
if str(ROOT) not in sys.path:
sys.path.insert(0, str(ROOT))
from tests.helpers.mark import ( # noqa: E402
get_skus_for_platform,
get_supported_card_counts,
)
# --- Constants ---
LOG = "upload_pipeline"
BOOTSTRAP_STEPS_FILENAME = "bootstrap-upload-steps.yml"
BOOTSTRAP_IMAGE_BUILD_KEYS = frozenset(
{
"image-build",
"image-build-a2",
"image-build-a3",
"image-build-a5",
"image-build-310p",
}
)
BOOTSTRAP_UPLOAD_IF_KEYS = {
"upload-ready-pipeline": "ready",
"upload-merge-pipeline": "merge",
"upload-nightly-pipeline": "nightly",
"upload-weekly-pipeline": "weekly",
}
E2E_GROUP_MARKER = "E2E Test"
CI_MIRROR_HARDWARES_PATH = ROOT / ".buildkite/common/ci_mirror_hardwares.yml"
CI_SOURCE_FILE_DEPENDENCIES_PATH = ROOT / ".buildkite/common/ci_source_file_dependencies.yml"
# Registry key whose prefixes bypass source filtering. Paths live in
# ci_source_file_dependencies.yml; do not attach this key to a job.
SOURCE_FILTER_FALLBACK_KEY = "source_filter_fallback"
2026-09-10 12:08:28 +08:00
CUDA_HF_TOKEN_ENV = "VLLM_CI_HF_TOKEN"
CUDA_HF_TOKEN_EXPORT = f'if [ -n "$${{{CUDA_HF_TOKEN_ENV}:-}}" ]; then export HF_TOKEN="$${{{CUDA_HF_TOKEN_ENV}}}"; fi'
CUDA_RUNTIME_CHECK = "python3 .buildkite/cuda/scripts/check_vllm_runtime.py"
# Bootstrap Buildkite ``if`` expressions.
# ``*_MAIN_IF``: main + env schedule. ``*_LABEL_IF``: PR label (and/or composed with MAIN).
# ``*_UPLOAD_IF``: full gate for uploading that child pipeline.
NIGHTLY_MAIN_IF = 'build.branch == "main" && build.env("NIGHTLY") == "1"'
NIGHTLY_LABEL_IF = (
f'({NIGHTLY_MAIN_IF}) || (build.branch != "main" && build.pull_request.labels includes "nightly-test")'
)
WEEKLY_E2E_IF = 'build.branch == "main" && build.env("WEEKLY") == "1"'
WEEKLY_MAIN_IF = 'build.branch == "main" && (build.env("WEEKLY") == "1" || build.env("NON_CRITICAL") == "1")'
WEEKLY_LABEL_IF = f'({WEEKLY_MAIN_IF}) || (build.branch != "main" && build.pull_request.labels includes "weekly-test")'
READY_LABEL_IF = 'build.branch != "main" && build.pull_request.labels includes "ready"'
MERGE_LABEL_IF = 'build.branch != "main" && build.pull_request.labels includes "merge-test"'
MERGE_MAIN_IF = (
'build.branch == "main" && build.env("NIGHTLY") != "1" && '
'build.env("WEEKLY") != "1" && build.env("NON_CRITICAL") != "1"'
)
READY_UPLOAD_IF = f"({WEEKLY_E2E_IF}) || ({READY_LABEL_IF})"
MERGE_UPLOAD_IF = f"({WEEKLY_E2E_IF}) || (({MERGE_MAIN_IF}) || ({MERGE_LABEL_IF}))"
BOOTSTRAP_DISABLED_IF = "false"
BOOTSTRAP_ENABLED_IF = "true"
# --- Logging ---
def _log(message: str) -> None:
print(f"{LOG}: {message}", file=sys.stderr)
# --- Bootstrap pipeline (bootstrap-upload-steps.yml) ---
def _get_bootstrap_platform(path: Path) -> str:
parts = path.as_posix().split("/")
return "npu" if "npu" in parts else "cuda"
def _load_bootstrap_steps(path: Path) -> str:
if path.name != BOOTSTRAP_STEPS_FILENAME:
raise ValueError(f"expected {BOOTSTRAP_STEPS_FILENAME}, got {path.name}")
return path.read_text(encoding="utf-8")
def _format_bootstrap_if(expr: str) -> str:
"""Return a Buildkite ``if`` string. Buildkite rejects YAML bool ``if`` values."""
if expr in ("true", "false"):
return expr
return f"({expr})"
def _compute_bootstrap_if_exprs(*, decision, platform: str) -> dict[str, str]:
if platform == "npu":
ready_upload = READY_LABEL_IF
merge_upload = BOOTSTRAP_DISABLED_IF
weekly_label_if = BOOTSTRAP_DISABLED_IF
else:
ready_upload = READY_UPLOAD_IF
merge_upload = MERGE_UPLOAD_IF
weekly_label_if = WEEKLY_LABEL_IF
if decision.skip_all:
# Docs / skip-mark only: no PR-label escape hatch. Main scheduled
# NIGHTLY=1 still runs L4; WEEKLY=1 / NON_CRITICAL=1 still run L5.
# main+WEEKLY=1 also uploads L2/L3 (those steps then pass --e2e).
image_expr = f"({NIGHTLY_MAIN_IF}) || ({WEEKLY_MAIN_IF})" if platform == "cuda" else NIGHTLY_MAIN_IF
ready_expr = WEEKLY_E2E_IF if platform == "cuda" else BOOTSTRAP_DISABLED_IF
merge_expr = WEEKLY_E2E_IF if platform == "cuda" else BOOTSTRAP_DISABLED_IF
nightly_expr = NIGHTLY_MAIN_IF
weekly_expr = WEEKLY_MAIN_IF if platform == "cuda" else BOOTSTRAP_DISABLED_IF
elif decision.skip_l2_l3:
l2_enabled = decision.is_run("npu", "l2") if platform == "npu" else decision.is_run("cuda", "l2")
l3_enabled = platform == "cuda" and decision.is_run("cuda", "l3")
ready_expr = ready_upload if l2_enabled else BOOTSTRAP_DISABLED_IF
merge_expr = merge_upload if l3_enabled else BOOTSTRAP_DISABLED_IF
nightly_expr = NIGHTLY_LABEL_IF
weekly_expr = weekly_label_if if platform == "cuda" else BOOTSTRAP_DISABLED_IF
image_parts = [f"({NIGHTLY_LABEL_IF})"]
if platform == "cuda":
image_parts.append(f"({weekly_label_if})")
if l2_enabled:
image_parts.insert(0, f"({ready_upload})")
if l3_enabled:
image_parts.insert(1 if l2_enabled else 0, f"({merge_upload})")
image_expr = " || ".join(image_parts)
else:
image_expr = BOOTSTRAP_ENABLED_IF
ready_expr = ready_upload
merge_expr = merge_upload if platform == "cuda" else BOOTSTRAP_DISABLED_IF
nightly_expr = NIGHTLY_LABEL_IF
weekly_expr = weekly_label_if if platform == "cuda" else BOOTSTRAP_DISABLED_IF
return {
"image": _format_bootstrap_if(image_expr),
"ready": _format_bootstrap_if(ready_expr),
"merge": _format_bootstrap_if(merge_expr),
"nightly": _format_bootstrap_if(nightly_expr),
"weekly": _format_bootstrap_if(weekly_expr),
}
def _apply_bootstrap_if(steps: list[Any], if_exprs: dict[str, str]) -> list[Any]:
"""Inject ``if`` by step key; drop steps that are unconditionally disabled."""
kept: list[Any] = []
for step in steps:
if not isinstance(step, dict):
kept.append(step)
continue
nested = step.get("steps")
if isinstance(nested, list):
nested_kept = _apply_bootstrap_if(nested, if_exprs)
if nested_kept:
kept.append({**step, "steps": nested_kept})
continue
key = step.get("key")
if key in BOOTSTRAP_IMAGE_BUILD_KEYS:
step["if"] = if_exprs["image"]
elif key in BOOTSTRAP_UPLOAD_IF_KEYS:
step["if"] = if_exprs[BOOTSTRAP_UPLOAD_IF_KEYS[key]]
if_expr = step.get("if")
if if_expr == "false":
_log(f"omit disabled bootstrap step {key!r}")
continue
if if_expr == "true":
# Unconditional step: omit ``if`` (Buildkite requires string ``if``, not YAML bool).
step.pop("if", None)
kept.append(step)
return kept
def _render_bootstrap_pipeline(
steps_yaml: str,
*,
decision,
path: Path,
) -> str:
"""Load bootstrap steps YAML and inject ``if`` expressions from skip-ci decision."""
doc = yaml.safe_load(steps_yaml)
if not isinstance(doc, dict):
raise ValueError(f"invalid bootstrap steps YAML for {path}")
steps = doc.get("steps")
if not isinstance(steps, list):
raise ValueError(f"bootstrap steps YAML must contain steps: list in {path}")
platform = _get_bootstrap_platform(path)
if_exprs = _compute_bootstrap_if_exprs(decision=decision, platform=platform)
doc["steps"] = _apply_bootstrap_if(steps, if_exprs)
return yaml.safe_dump(doc, sort_keys=False)
# --- Test pipeline (test-ready.yml, test-merge.yml, test-nightly.yml) ---
@lru_cache(maxsize=1)
def _load_mirror_hardwares() -> dict[str, dict[str, Any]]:
if not CI_MIRROR_HARDWARES_PATH.is_file():
raise FileNotFoundError(f"missing CI mirror_hardwares registry: {CI_MIRROR_HARDWARES_PATH}")
doc = yaml.safe_load(CI_MIRROR_HARDWARES_PATH.read_text(encoding="utf-8"))
if not isinstance(doc, dict):
raise ValueError(f"invalid CI mirror_hardwares registry: {CI_MIRROR_HARDWARES_PATH}")
presets = doc.get("mirror_hardwares")
if not isinstance(presets, dict):
raise ValueError(f"mirror_hardwares must be a mapping in {CI_MIRROR_HARDWARES_PATH}")
return presets
@lru_cache(maxsize=1)
def _load_source_file_dependencies() -> dict[str, list[str]]:
def flatten(value: Any, *, key: str) -> list[str]:
"""Flatten YAML-anchor nested lists into a de-duplicated prefix list."""
if isinstance(value, str):
if not value.strip():
raise ValueError(f"empty path in source_file_dependencies[{key!r}]")
return [value]
if isinstance(value, list):
flattened: list[str] = []
seen: set[str] = set()
for item in value:
for path in flatten(item, key=key):
if path not in seen:
seen.add(path)
flattened.append(path)
return flattened
raise ValueError(
f"source_file_dependencies[{key!r}] must be a list of path prefixes, got {type(value).__name__}",
)
if not CI_SOURCE_FILE_DEPENDENCIES_PATH.is_file():
raise FileNotFoundError(
f"missing CI source_file_dependencies registry: {CI_SOURCE_FILE_DEPENDENCIES_PATH}",
)
doc = yaml.safe_load(CI_SOURCE_FILE_DEPENDENCIES_PATH.read_text(encoding="utf-8"))
if not isinstance(doc, dict):
raise ValueError(f"invalid CI source_file_dependencies registry: {CI_SOURCE_FILE_DEPENDENCIES_PATH}")
presets = doc.get("source_file_dependencies")
if not isinstance(presets, dict):
raise ValueError(f"source_file_dependencies must be a mapping in {CI_SOURCE_FILE_DEPENDENCIES_PATH}")
loaded: dict[str, list[str]] = {}
for key, value in presets.items():
if not isinstance(key, str) or not key.strip():
raise ValueError(
f"source_file_dependencies keys must be non-empty strings in {CI_SOURCE_FILE_DEPENDENCIES_PATH}",
)
loaded[key] = flatten(value, key=key)
return loaded
_REGISTRY_KEY_RE = re.compile(r"^[a-z][a-z0-9_]*$")
def _parse_command_text(step: dict[str, Any]) -> str:
"""Parse step ``commands`` / ``command`` into a single string."""
raw = step.get("commands", step.get("command"))
if raw is None:
return ""
if isinstance(raw, str):
return raw
if isinstance(raw, list):
return "\n".join(str(item) for item in raw if item is not None)
return str(raw)
def _resolve_source_file_dependencies(step: dict[str, Any]) -> list[str] | None:
"""Expand a registry key (or list of keys / inline path prefixes)."""
deps = step.get("source_file_dependencies")
if deps is None:
return None
def dedupe(*groups: list[str]) -> list[str]:
merged: list[str] = []
seen: set[str] = set()
for group in groups:
for path in group:
if path not in seen:
seen.add(path)
merged.append(path)
return merged
def lookup(key: str) -> list[str]:
if key == SOURCE_FILTER_FALLBACK_KEY:
raise ValueError(
f"source_file_dependencies {key!r} bypasses filtering and cannot be attached to a step "
f"({_get_step_label(step)!r})",
)
registry = _load_source_file_dependencies()
paths = registry.get(key)
if paths is None:
known = ", ".join(sorted(registry))
raise ValueError(
f"unknown source_file_dependencies {key!r} in step {_get_step_label(step)!r}; known: {known}",
)
return list(paths)
prefixes: list[str]
if isinstance(deps, str):
prefixes = lookup(deps) if _REGISTRY_KEY_RE.fullmatch(deps) else [deps]
elif isinstance(deps, list):
if not all(isinstance(item, str) for item in deps):
raise ValueError(
f"source_file_dependencies must be a string key or list of strings in step {_get_step_label(step)!r}",
)
keyish = [bool(_REGISTRY_KEY_RE.fullmatch(item)) for item in deps]
if deps and all(keyish):
prefixes = dedupe(*(lookup(key) for key in deps))
elif any(keyish):
raise ValueError(
f"source_file_dependencies in step {_get_step_label(step)!r} mixes registry keys and path prefixes",
)
else:
prefixes = list(deps)
else:
raise ValueError(
f"source_file_dependencies must be a string key or list in step {_get_step_label(step)!r}",
)
return prefixes
@lru_cache(maxsize=1)
def _cuda_mirror_chips() -> tuple[str, ...]:
"""Lowercase CUDA SKUs that have at least one ``{chip}_{n}`` preset.
Longer tokens first so ``h100`` is not matched as a prefix of ``h200``.
"""
skus = get_skus_for_platform("cuda")
known = _load_mirror_hardwares()
chips = {
sku.lower() for sku in skus if any(name == sku.lower() or name.startswith(f"{sku.lower()}_") for name in known)
}
return tuple(sorted(chips, key=lambda chip: (-len(chip), chip)))
def _parse_cards_marks(expr: str) -> tuple[set[int], bool]:
"""Parse registered ``cards_n`` / ``not cards_n`` from a pytest ``-m`` expr.
Returns positive card counts and whether any ``not cards_*`` is present.
"""
counts = get_supported_card_counts()
alt = "|".join(str(n) for n in sorted(counts, key=lambda n: (-len(str(n)), -n))) if counts else "(?!)"
positives: set[int] = set()
has_not_cards = False
for match in re.finditer(rf"\b(?:(not)\s+)?cards_({alt})\b", expr):
if match.group(1):
has_not_cards = True
else:
positives.add(int(match.group(2)))
return positives, has_not_cards
def _parse_hardware_marks(expr: str) -> set[str]:
"""Parse lowercase positive CUDA SKUs that have a mirror preset.
``not H100`` is ignored. SKUs with no ``{chip}_*`` preset (e.g. H200) are dropped.
"""
skus = sorted(get_skus_for_platform("cuda"), key=lambda s: (-len(s), s))
if not skus:
return set()
alt = "|".join(re.escape(sku) for sku in skus)
mirror = set(_cuda_mirror_chips())
chips: set[str] = set()
for match in re.finditer(rf"\b(?:(not)\s+)?({alt})\b", expr):
if match.group(1):
continue
chip = match.group(2).lower()
if chip in mirror:
chips.add(chip)
return chips
# CI policy: not pytest-mark registry.
_SUPPORTED_MIRROR_HW_SELECTORS = frozenset({"b200"})
# Unset MIRROR_HW: match H100 then L4 in -m (skip if neither).
_DEFAULT_INFER_CHIPS = ("h100", "l4")
_PYTEST_MARKER_ARG = re.compile(r"-m\s+(?:\"([^\"]*)\"|'([^']*)'|(\S+))")
_CARDS_CHIP_MAX: Literal["max"] = "max"
def _get_mirror_hw_selector() -> str:
"""Return lowercase ``MIRROR_HW``, or empty to keep string presets / marker chips.
A set value must be ``b200`` (case-insensitive). Unknown values fail closed
so a typo cannot silently drop the CUDA pipeline.
"""
selector = os.environ.get("MIRROR_HW", "").strip().lower()
if not selector:
return ""
if selector not in _SUPPORTED_MIRROR_HW_SELECTORS:
allowed = " or ".join(repr(s) for s in sorted(_SUPPORTED_MIRROR_HW_SELECTORS))
raise ValueError(f"unsupported MIRROR_HW={selector!r}; expected empty or {allowed}")
return selector
def _parse_pytest_marks(step: dict[str, Any]) -> tuple[set[str], int | Literal["max"] | None]:
"""Parse pytest ``-m`` from step commands: CUDA SKU chips and cards count.
*chips* are lowercase SKUs from positive markers (``not H100`` is ignored).
*cards* is the max positive ``cards_n``, ``"max"`` when only ``not cards_*``
is present, or ``None``.
"""
text = _parse_command_text(step)
chips: set[str] = set()
positives: set[int] = set()
has_not_cards = False
for match in _PYTEST_MARKER_ARG.finditer(text):
expr = match.group(1) or match.group(2) or match.group(3) or ""
chips |= _parse_hardware_marks(expr)
card_counts, not_cards = _parse_cards_marks(expr)
positives.update(card_counts)
has_not_cards = has_not_cards or not_cards
if positives:
cards: int | Literal["max"] | None = max(positives)
elif has_not_cards:
cards = _CARDS_CHIP_MAX
else:
cards = None
return chips, cards
def _compose_mirror_hardware_name(
chips: set[str],
cards: int | Literal["max"] | None,
*,
step_label: str,
) -> str | None:
"""Assemble ``{chip}_{n}`` from pytest ``-m`` chips/cards, or None to skip.
Unset ``MIRROR_HW``: ``H100`` if in ``-m``, else ``L4``, else skip.
Set ``MIRROR_HW``: that chip if in ``-m``, else skip.
``cards`` is a registered ``cards_n``, ``"max"`` for that chip's highest
existing preset, or None (error after a chip matched — no ``cards_*``).
"""
selector = _get_mirror_hw_selector()
if selector:
chip = selector if selector in chips else None
skip_reason = f"MIRROR_HW={selector!r} not in -m SKUs {sorted(chips)}"
else:
chip = next((c for c in _DEFAULT_INFER_CHIPS if c in chips), None)
skip_reason = f"no H100 or L4 in -m SKUs {sorted(chips)}"
if chip is None:
_log(f"skip {step_label}: {skip_reason}")
return None
if cards is None:
raise ValueError(
f"step {step_label!r} has a pytest -m SKU marker but no cards_* "
f"(or not cards_*); add cards_n or set mirror_hardwares to force a preset",
)
count_source = "not cards_*" if cards == _CARDS_CHIP_MAX else f"cards_{cards}"
count = None if cards == _CARDS_CHIP_MAX else cards
registered = get_supported_card_counts()
if count is not None and count not in registered:
supported = ", ".join(f"cards_{n}" for n in sorted(registered)) or "(none)"
raise ValueError(
f"{count_source} in step {step_label!r} is not a registered cards_* mark; supported: {supported}",
)
known = _load_mirror_hardwares()
if count is None:
available = [n for n in sorted(registered) if f"{chip}_{n}" in known]
if not available:
raise ValueError(
f"{count_source} in step {step_label!r} has no CUDA count presets for {chip}",
)
count = max(available)
chosen = f"{chip}_{count}"
if chosen not in known:
raise ValueError(
f"{count_source} in step {step_label!r} has no preset {chosen!r}",
)
_log(f"{step_label}: {count_source} → {chosen!r}")
return chosen
def _resolve_mirror_hardware_name(hardware: Any, *, step_label: str) -> str | None:
"""Map ``mirror_hardwares`` to a ``ci_mirror_hardwares.yml`` preset, or None to skip.
The value is a preset name (``h100_2``). Unknown names fail later when the
registry is loaded. CUDA names are omitted when ``MIRROR_HW`` names a
different chip; NPU strings ignore ``MIRROR_HW``.
"""
if isinstance(hardware, bool):
raise ValueError(f"mirror_hardwares must not be a boolean in step {step_label!r}")
if isinstance(hardware, str):
name = hardware.strip()
elif isinstance(hardware, int):
name = str(hardware)
else:
raise ValueError(
f"mirror_hardwares must be a preset string in step {step_label!r}",
)
if not name:
raise ValueError(f"mirror_hardwares must be a non-empty string in step {step_label!r}")
selector = _get_mirror_hw_selector()
token = name.lower()
chip = next(
(c for c in _cuda_mirror_chips() if token == c or token.startswith(f"{c}_")),
None,
)
if selector and chip is not None and chip != selector:
_log(
f"skip {step_label}: preset {name!r} is {chip} hardware; MIRROR_HW={selector!r}",
)
return None
return name
def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None:
"""Replace ``mirror_hardwares`` (or infer it from ``-m``) with preset agents/plugins.
Explicit preset string: used as-is; pytest marks are ignored.
Omitted key: match H100/L4 (or ``MIRROR_HW``) in ``-m``, then compose
``{chip}_{n}`` from ``cards_n``. No match skips the step. Steps with no
SKU/cards and no ``mirror_hardwares`` are left unchanged (CPU jobs).
Returns None when the current ``MIRROR_HW`` selector does not apply.
"""
step_label = _get_step_label(step)
has_pool = step.get("agents") is not None or step.get("plugins") is not None or step.get("image") is not None
if "mirror_hardwares" not in step:
if has_pool:
return step
chips, cards = _parse_pytest_marks(step)
if not chips and cards is None:
return step
preset_name = _compose_mirror_hardware_name(chips, cards, step_label=step_label)
else:
if has_pool:
raise ValueError(
f"step {step_label!r} sets mirror_hardwares together with agents/plugins/image; "
f"use mirror_hardwares only",
)
preset_name = _resolve_mirror_hardware_name(
step.get("mirror_hardwares"),
step_label=step_label,
)
if preset_name is None:
return None
preset = _load_mirror_hardwares().get(preset_name)
if preset is None:
known = ", ".join(sorted(_load_mirror_hardwares()))
raise ValueError(
f"unknown mirror_hardwares {preset_name!r} in step {step_label!r}; known: {known}",
)
expanded = copy.deepcopy(preset)
merged = {key: value for key, value in step.items() if key != "mirror_hardwares"} | expanded
2026-09-10 12:08:28 +08:00
# A scheduled build's HF_TOKEN may override a pod env entry with the same
# name. The preset injects the Kubernetes secret under a collision-proof
# alias; restore the conventional name after Buildkite has merged env.
if any(preset_name == chip or preset_name.startswith(f"{chip}_") for chip in _cuda_mirror_chips()):
commands = merged.get("commands")
if isinstance(commands, list):
merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK, *commands]
2026-09-10 12:08:28 +08:00
elif commands is None:
merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK]
2026-09-10 12:08:28 +08:00
else:
merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK, commands]
# Preset retry (K8S_RETRY on l4_*) must not clobber a step that opted out.
if "retry" in step:
merged["retry"] = step["retry"]
return merged
def _match_source_file(changed_files: list[str], prefixes: list[str]) -> bool:
for path in changed_files:
for prefix in prefixes:
normalized = prefix.rstrip("/")
if path == normalized or path.startswith(f"{normalized}/"):
return True
return False
def _get_step_label(step: dict[str, Any]) -> str:
return str(step.get("group") or step.get("label") or "<step>")
def _process_test_steps(
steps: list[Any],
changed_files: list[str] | None,
) -> list[Any]:
"""Drop steps by ``source_file_dependencies`` when *changed_files* is set; always strip that field."""
processed: list[Any] = []
for step in steps:
if not isinstance(step, dict):
processed.append(step)
continue
deps = _resolve_source_file_dependencies(step)
if changed_files is not None and deps is not None and not _match_source_file(changed_files, deps):
_log(f"skip {_get_step_label(step)!r} (no changes under {deps})")
continue
nested = step.get("steps")
if nested is not None:
kept_nested = _process_test_steps(nested, changed_files)
if not kept_nested:
_log(f"omit empty group {_get_step_label(step)!r}")
continue
new_step = {key: value for key, value in step.items() if key != "source_file_dependencies"}
new_step["steps"] = kept_nested
processed.append(new_step)
continue
leaf = {key: value for key, value in step.items() if key != "source_file_dependencies"}
expanded = _expand_mirror_hardwares(leaf)
if expanded is None:
continue
processed.append(expanded)
return processed
def _select_e2e_group_steps(steps: list[Any]) -> list[Any]:
"""Keep only top-level groups whose name contains ``E2E_GROUP_MARKER``."""
selected = [
step
for step in steps
if isinstance(step, dict) and isinstance(step.get("group"), str) and E2E_GROUP_MARKER in step["group"]
]
if not selected:
_log(f"no group matching {E2E_GROUP_MARKER!r} found")
else:
_log(f"keep {len(selected)} group(s) matching {E2E_GROUP_MARKER!r}")
return selected
def _any_source_dependency_match(steps: list[Any], changed_files: list[str]) -> bool:
"""True when any step with ``source_file_dependencies`` matches *changed_files*."""
for step in steps:
if not isinstance(step, dict):
continue
deps = _resolve_source_file_dependencies(step)
if deps is not None and _match_source_file(changed_files, deps):
return True
nested = step.get("steps")
if isinstance(nested, list) and _any_source_dependency_match(nested, changed_files):
return True
return False
def _render_test_pipeline(
doc: dict[str, Any],
changed_files: list[str] | None,
*,
pipeline_path: Path | None = None,
e2e_only: bool = False,
) -> dict[str, Any]:
"""Filter steps by PR diff and strip uploader-only ``source_file_dependencies`` metadata."""
# Validate MIRROR_HW once up front. Per-step skip must not run for typos
# (e.g. b20o), or a pipeline can shrink to leftover CPU-only steps and
# still upload successfully.
_get_mirror_hw_selector()
steps = doc.get("steps")
if not isinstance(steps, list):
return doc
if e2e_only:
steps = _select_e2e_group_steps(steps)
# Bypass is a fallback: only when no job-key prefix matched.
if changed_files is not None and pipeline_path is not None:
if not _any_source_dependency_match(steps, changed_files):
bypass = _source_filter_fallback_reason(changed_files, pipeline_path)
if bypass is not None:
_log(f"keep all jobs (no source_file_dependencies match; bypassed by {bypass})")
changed_files = None
steps = _process_test_steps(steps, changed_files)
return {**doc, "steps": steps}
# --- Entry (read file → bootstrap or test render → YAML string) ---
def _repo_relative_posix(path: Path) -> str:
resolved = path.resolve()
try:
return resolved.relative_to(ROOT).as_posix()
except ValueError:
return path.as_posix().replace("\\", "/")
def _source_filter_fallback_prefixes() -> list[str]:
"""Shared uploader/registry prefixes from ``source_filter_fallback``. Not a job key."""
registry = _load_source_file_dependencies()
prefixes = registry.get(SOURCE_FILTER_FALLBACK_KEY)
if not prefixes:
raise ValueError(
f"source_file_dependencies[{SOURCE_FILTER_FALLBACK_KEY!r}] must list shared uploader paths "
f"in {CI_SOURCE_FILE_DEPENDENCIES_PATH}",
)
return prefixes
def _source_filter_fallback_reason(changed_files: list[str] | None, pipeline_path: Path) -> str | None:
"""Return the changed path that should disable source filtering, if any."""
if not changed_files:
return None
changed = set(changed_files)
pipeline_rel = _repo_relative_posix(pipeline_path)
if pipeline_rel in changed:
return pipeline_rel
prefixes = _source_filter_fallback_prefixes()
for path in changed_files:
if _match_source_file([path], prefixes):
return path
return None
def _changed_files_for_source_filter(
ctx,
*,
force_all: bool,
e2e_only: bool,
) -> list[str] | None:
"""Return the diff to filter against, or None to keep every step.
``--all`` / ``--e2e`` disable filtering (scheduled NIGHTLY/WEEKLY full
uploads and weekly E2E sweeps). Post-merge ``main`` L3 and PR-label
uploads filter against the commit/PR diff. Pipeline YAML /
``source_filter_fallback`` matches are applied later in
``_render_test_pipeline`` only when no job-key prefix already matched.
"""
if force_all or e2e_only:
return None
return ctx.changed_files
def _render_pipeline(
path: Path,
*,
force_all: bool = False,
e2e_only: bool = False,
) -> str:
if path.name == BOOTSTRAP_STEPS_FILENAME:
ctx = resolve_ci_context_from_git()
continuation = _load_bootstrap_steps(path)
return _render_bootstrap_pipeline(
continuation,
decision=ctx.decision,
path=path,
)
text = path.read_text(encoding="utf-8")
ctx = resolve_ci_context_from_git()
changed_files = _changed_files_for_source_filter(ctx, force_all=force_all, e2e_only=e2e_only)
doc = yaml.safe_load(text)
if not isinstance(doc, dict):
raise ValueError(f"invalid pipeline YAML: {path}")
doc = _render_test_pipeline(
doc,
changed_files,
pipeline_path=path,
e2e_only=e2e_only,
)
return yaml.safe_dump(doc, sort_keys=False)
def _upload_to_buildkite(content: str) -> None:
subprocess.run(
["buildkite-agent", "pipeline", "upload"],
input=content,
text=True,
check=True,
)
# --- CLI ---
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"pipeline",
nargs="?",
default=".buildkite/cuda/bootstrap-upload-steps.yml",
help="Pipeline YAML path (default: .buildkite/cuda/bootstrap-upload-steps.yml)",
)
parser.add_argument(
"--upload",
action="store_true",
help="Pipe rendered YAML to buildkite-agent pipeline upload",
)
mode = parser.add_mutually_exclusive_group()
mode.add_argument(
"--all",
action="store_true",
help="Keep all steps (disable diff-aware skipping)",
)
mode.add_argument(
"--e2e",
action="store_true",
help="Keep only the E2E Test group",
)
args = parser.parse_args()
path = Path(args.pipeline)
if not path.is_absolute():
path = ROOT / path
if not path.is_file():
_log(f"missing pipeline file: {path}")
return 1
rendered = _render_pipeline(path, force_all=args.all, e2e_only=args.e2e)
if args.upload:
_upload_to_buildkite(rendered)
else:
sys.stdout.write(rendered)
return 0
if __name__ == "__main__":
raise SystemExit(main())