884 lines
33 KiB
Python
884 lines
33 KiB
Python
#!/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"
|
||
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
|
||
# 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]
|
||
elif commands is None:
|
||
merged["commands"] = [CUDA_HF_TOKEN_EXPORT, CUDA_RUNTIME_CHECK]
|
||
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())
|