diff --git a/pyproject.toml b/pyproject.toml index 35d333fd1..bb9cdc688 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,6 +71,11 @@ demo = [ "gradio>=6.7.0", ] +# Local-test-only extras +local = [ + "FlagEmbedding", +] + docs = [ "mkdocs>=1.5.0", "mkdocs-api-autonav", @@ -184,6 +189,7 @@ markers = [ "core_model: L1&L2 tests (run in each PR)", "advanced_model: L3 level tests (run on each merge)", "full_model: L4 level tests (run nightly)", + "local_model: tests (run locally)", # function module markers "diffusion: Diffusion model tests", "omni: Omni model tests", diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py index bf230c07f..6541866d4 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3.py +++ b/tests/e2e/accuracy/test_hunyuan_image3.py @@ -23,18 +23,31 @@ from tests.e2e.accuracy.helpers import ( download_images, model_output_dir, ) -from tests.helpers.runtime import OmniRunner +from tests.helpers.mark import hardware_marks from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA" -pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] +pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, *hardware_marks(res={"cuda": "H100"}, num_cards=8)] # ============================================================================ # Configurable Parameters # ============================================================================ -AR_DEVICES = "0,1" -DIT_DEVICES = "2,3" +# Comma-separated logical CUDA device ids per stage: split visible GPUs (0..n-1), first half -> AR, second -> DiT. + + +def _default_ar_dit_devices() -> tuple[str, str]: + """First floor(n/2) logical devices -> AR, rest -> DiT. ``device_count`` respects ``CUDA_VISIBLE_DEVICES``.""" + n = torch.accelerator.device_count() + if n < 2: + return "0,1", "2,3" + split = n // 2 + ar = ",".join(str(i) for i in range(split)) + dit = ",".join(str(i) for i in range(split, n)) + return ar, dit + + +AR_DEVICES, DIT_DEVICES = _default_ar_dit_devices() MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct" NUM_INFERENCE_STEPS = 50 GUIDANCE_SCALE = 2.5 @@ -75,7 +88,7 @@ _BASE_CONFIG = { "scheduler_cls": "vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler", "gpu_memory_utilization": 0.95, "enforce_eager": True, "trust_remote_code": True, "engine_output_type": "latent", "enable_prefix_caching": False, - "max_num_batched_tokens": 32768, "tensor_parallel_size": 2, "pipeline_parallel_size": 1, + "max_num_batched_tokens": 32768, "tensor_parallel_size": AR_TP_SIZE, "pipeline_parallel_size": 1, "hf_overrides": {"rope_parameters": {"mrope_section": [0, 32, 32], "rope_type": "default"}}, }, "is_comprehension": False, "final_output": True, "final_output_type": "text", @@ -91,7 +104,7 @@ _BASE_CONFIG = { "engine_args": { "model_stage": "dit", "model_arch": "HunyuanImage3ForCausalMM", "enforce_eager": True, "trust_remote_code": True, "distributed_executor_backend": "mp", - "parallel_config": {"tensor_parallel_size": 2, "enable_expert_parallel": True}, + "parallel_config": {"tensor_parallel_size": DIT_TP_SIZE, "enable_expert_parallel": True}, }, "engine_input_source": [0], "custom_process_input_func": "vllm_omni.model_executor.stage_input_processors.hunyuan_image3.ar2diffusion", @@ -134,12 +147,13 @@ def _make_config(enable_kv_reuse: bool, path: Path) -> None: def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Image, str, float]: from transformers import AutoTokenizer + from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType from vllm_omni.platforms import current_omni_platform build_kwargs: dict = {"task": "it2i", "bot_task": "think_recaption", "sys_type": "en_unified", "num_images": 2} - tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True) result = build_prompt_tokens( PROMPT, tokenizer, @@ -149,7 +163,7 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag system_prompt_type = result.system_prompt_type ar_stop_token_ids = resolve_stop_token_ids(task="it2i", bot_task="think_recaption", tokenizer=tokenizer) - with OmniRunner(MODEL_NAME, stage_configs_path=stage_configs_path) as runner: + with OmniRunner(MODEL_PATH, stage_configs_path=stage_configs_path) as runner: params_list = list(runner.omni.default_sampling_params_list) for sp in params_list: if isinstance(sp, OmniDiffusionSamplingParams): @@ -206,7 +220,10 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag return image, cot_text, elapsed -@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Needs 4+ GPUs (2 AR + 2 DiT)") +@pytest.mark.skipif( + torch.accelerator.device_count() < AR_TP_SIZE + DIT_TP_SIZE, + reason=f"Needs {AR_TP_SIZE + DIT_TP_SIZE}+ GPUs ({AR_TP_SIZE} AR + {DIT_TP_SIZE} DiT)", +) def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None: if importlib.util.find_spec("FlagEmbedding") is None: raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding") diff --git a/tools/nightly/run_nightly_jobs.sh b/tools/nightly/run_nightly_jobs.sh new file mode 100644 index 000000000..b115a9ee2 --- /dev/null +++ b/tools/nightly/run_nightly_jobs.sh @@ -0,0 +1,754 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Extract steps from .buildkite/test-nightly.yml that contain pytest, synthesize +# small bash wrappers (exports + pytest), run them, and tee output to logs named +# after each step's Buildkite "key" when present (otherwise a slug of the label). +# YAML steps whose label contains "Perf Test" run first, then +# python tools/nightly/generate_nightly_perf_excel.py runs (even if some perf jobs failed; +# the script only aggregates JSON already present under tests/dfx/perf/results). +# Excel and generate_nightly_perf_excel.log are written under $REPO_ROOT/logs/; +# per-job tee logs use LOG_DIR (see --log-dir). Then remaining jobs run. +# +# Test kind (--test-type / TEST_TYPE), multiple allowed (OR semantics): +# Repeat the flag and/or use comma-separated values, e.g. +# --test-type perf,acc OR --test-type perf --test-type function +# If any selected value is all, YAML steps are unconstrained except that perf/acc/function tokens +# are ignored; all,stability / all,local / all,stability,local still run those extra bundles. +# perf — label contains "Perf Test" (throughput / benchmark jobs) +# acc — label contains "Accuracy Test" (incl. GEBench / GEdit-Bench style) +# function — label has neither "Perf Test" nor "Accuracy Test" (incl. Doc, Multi-Replica, etc.) +# stability— fixed dfx stability scripts under tests/dfx/stability/scripts/ (see below) +# local — pytest -sv -m " and local_model" from repo root (no YAML step extract) +# all — no test-kind filter for YAML steps (any leaf step with pytest in test-nightly.yml) +# +# Model area (--model-type / MODEL_TYPE), multiple allowed (OR semantics): +# e.g. --model-type omni,tts +# omni — label contains "Omni ·" +# tts — label contains "TTS ·" +# diffusion— label contains "Diffusion" +# all — no model filter (e.g. Quantization-only steps match only when test-type allows) +# +# Combining stability and/or local with perf/acc/function/all runs each enabled bundle; each enabled +# mode must match at least one job or the script exits with status 2. +# +# local (when included in TEST_TYPE): +# From repo root: pytest -sv -m " and local_model" (markers from MODEL_TYPE: omni, tts, +# diffusion; all → "(omni or tts or diffusion) and local_model"). Not filtered by nightly YAML. +# +# stability (when included in TEST_TYPE): +# From repo root: pytest -s -v tests/dfx/stability/scripts/test_stability_*.py +# model_type: omni → qwen3_omni; tts → qwen3_tts; diffusion → qwen_image + wan22; all → all four +# LABEL_SUBSTR: if set, script path / job key / filename must contain it +# +# Requirements: bash, python3, PyYAML (pip install pyyaml) +# +# Usage: +# bash path/to/run_nightly_jobs.sh +# REPO_ROOT=/path/to/vllm-omni bash path/to/run_nightly_jobs.sh --test-type perf,acc --dry-run +# YML=/path/to/vllm-omni/.buildkite/test-nightly.yml bash path/to/run_nightly_jobs.sh +# +# Repository / YAML (no dependency on where this script lives): +# • Set REPO_ROOT (or pass --repo-root) — default YAML is $REPO_ROOT/.buildkite/test-nightly.yml +# • Or set YML (or --yaml) — repo root is inferred as parent of the .buildkite directory +# • Or run from inside the clone: git rev-parse --show-toplevel, else walk up from $PWD, +# then from the script's directory, until .buildkite/test-nightly.yml exists +# +# Optional environment: +# REPO_ROOT - vllm-omni root (working directory for pytest); see above +# YML - path to test-nightly.yml (default: $REPO_ROOT/.buildkite/test-nightly.yml) +# LOG_DIR - logs + generated job scripts (default: $REPO_ROOT/logs/nightly_jobs) +# TEST_TYPE - comma-separated and/or repeated flags (default: all); see above +# MODEL_TYPE - comma-separated and/or repeated flags (default: all); see above +# LABEL_SUBSTR - YAML mode: substring of Buildkite step label; stability: substring of path/key/filename +# DRY_RUN=1 - print extracted commands only; do not write scripts or run pytest +# +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +LABEL_SUBSTR="${LABEL_SUBSTR:-}" +TEST_TYPE_ENV="${TEST_TYPE:-all}" +MODEL_TYPE_ENV="${MODEL_TYPE:-all}" +TEST_TYPE_CLI_PARTS=() +MODEL_TYPE_CLI_PARTS=() +TEST_TYPE_FROM_CLI=0 +MODEL_TYPE_FROM_CLI=0 +DRY_RUN="${DRY_RUN:-0}" +# REPO_ROOT / YML / LOG_DIR resolved after CLI (do not assume script path vs repo root) + +usage() { + sed -n '2,66p' "$0" | sed 's/^# \{0,1\}//' +} + +# Append comma-separated tokens from $2 into array named $1 (lowercased, trimmed). +_split_append_csv_array() { + local -n _arr="${1}" + local _raw="${2}" + local _ifs="${IFS}" + local _part + IFS=',' + for _part in ${_raw}; do + IFS="${_ifs}" + _part="$(printf '%s' "${_part}" | tr '[:upper:]' '[:lower:]' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')" + [[ -n "${_part}" ]] || continue + _arr+=("${_part}") + done +} + +# Build comma-separated TEST_TYPE from CLI parts or env; collapse to a single "all" if present. +_finalize_test_type_csv() { + if [[ "${TEST_TYPE_FROM_CLI}" -eq 1 ]]; then + if ((${#TEST_TYPE_CLI_PARTS[@]} == 0)); then + echo "--test-type requires a non-empty value" >&2 + exit 2 + fi + local _has_all=0 _has_stability=0 _has_local=0 + local _x + for _x in "${TEST_TYPE_CLI_PARTS[@]}"; do + if [[ "${_x}" == all ]]; then + _has_all=1 + fi + if [[ "${_x}" == stability ]]; then + _has_stability=1 + fi + if [[ "${_x}" == local ]]; then + _has_local=1 + fi + done + if [[ "${_has_all}" -eq 1 ]]; then + local _out="all" + if [[ "${_has_stability}" -eq 1 ]]; then + _out="${_out},stability" + fi + if [[ "${_has_local}" -eq 1 ]]; then + _out="${_out},local" + fi + printf '%s' "${_out}" + return 0 + fi + local -A _seen=() + local -a _out=() + for _x in "${TEST_TYPE_CLI_PARTS[@]}"; do + [[ ${_seen["${_x}"]+isset} ]] && continue + _seen["${_x}"]=1 + _out+=("${_x}") + done + (IFS=','; printf '%s' "${_out[*]}") + else + printf '%s' "${TEST_TYPE_ENV}" + fi +} + +_finalize_model_type_csv() { + if [[ "${MODEL_TYPE_FROM_CLI}" -eq 1 ]]; then + if ((${#MODEL_TYPE_CLI_PARTS[@]} == 0)); then + echo "--model-type requires a non-empty value" >&2 + exit 2 + fi + local -A _seen=() + local -a _out=() + local _x + for _x in "${MODEL_TYPE_CLI_PARTS[@]}"; do + if [[ "${_x}" == all ]]; then + printf '%s' "all" + return 0 + fi + [[ ${_seen["${_x}"]+isset} ]] && continue + _seen["${_x}"]=1 + _out+=("${_x}") + done + (IFS=','; printf '%s' "${_out[*]}") + else + printf '%s' "${MODEL_TYPE_ENV}" + fi +} + +while [[ $# -gt 0 ]]; do + case "$1" in + -h | --help) + usage + exit 0 + ;; + --dry-run) + DRY_RUN=1 + shift + ;; + --log-dir) + LOG_DIR="$2" + shift 2 + ;; + --yaml) + YML="$2" + shift 2 + ;; + --repo-root) + REPO_ROOT="$2" + shift 2 + ;; + --label-substr) + LABEL_SUBSTR="$2" + shift 2 + ;; + --model-type) + MODEL_TYPE_FROM_CLI=1 + _split_append_csv_array MODEL_TYPE_CLI_PARTS "$2" + shift 2 + ;; + --test-type) + TEST_TYPE_FROM_CLI=1 + _split_append_csv_array TEST_TYPE_CLI_PARTS "$2" + shift 2 + ;; + *) + echo "Unknown option: $1" >&2 + usage >&2 + exit 2 + ;; + esac +done + +TEST_TYPE="$(_finalize_test_type_csv)" +MODEL_TYPE="$(_finalize_model_type_csv)" + +BUILDKITE_REL=".buildkite/test-nightly.yml" + +_find_repo_containing_nightly() { + local dir="${1:-}" + [[ -n "$dir" ]] || return 1 + dir="$(cd "$dir" && pwd)" || return 1 + while true; do + if [[ -f "${dir}/${BUILDKITE_REL}" ]]; then + printf '%s\n' "$dir" + return 0 + fi + [[ "${dir}" == "/" ]] && return 1 + dir="$(dirname "${dir}")" + done +} + +_derive_repo_root_from_yml() { + local yml="$1" + local d + d="$(cd "$(dirname "${yml}")" && pwd)" + [[ "$(basename "${d}")" == ".buildkite" ]] || return 1 + printf '%s\n' "$(dirname "${d}")" +} + +# Resolve REPO_ROOT, YML (no relative path between script and repo) +if [[ -n "${YML:-}" && -n "${REPO_ROOT:-}" ]]; then + REPO_ROOT="$(cd "${REPO_ROOT}" && pwd)" + YML="$(cd "$(dirname "${YML}")" && pwd)/$(basename "${YML}")" +elif [[ -n "${YML:-}" ]]; then + YML="$(cd "$(dirname "${YML}")" && pwd)/$(basename "${YML}")" + if ! REPO_ROOT="$(_derive_repo_root_from_yml "${YML}")"; then + echo "Could not derive REPO_ROOT from YML=${YML} (expected file at /.buildkite/test-nightly.yml)." >&2 + echo "Set REPO_ROOT explicitly (or pass --repo-root) for pytest working directory." >&2 + exit 2 + fi +elif [[ -n "${REPO_ROOT:-}" ]]; then + REPO_ROOT="$(cd "${REPO_ROOT}" && pwd)" + YML="${REPO_ROOT}/${BUILDKITE_REL}" +else + REPO_ROOT="" + if command -v git >/dev/null 2>&1; then + REPO_ROOT="$(git -C "${PWD}" rev-parse --show-toplevel 2>/dev/null || true)" + fi + if [[ -z "${REPO_ROOT}" ]]; then + REPO_ROOT="$(_find_repo_containing_nightly "${PWD}" || true)" + fi + if [[ -z "${REPO_ROOT}" ]]; then + REPO_ROOT="$(_find_repo_containing_nightly "${SCRIPT_DIR}" || true)" + fi + if [[ -z "${REPO_ROOT}" ]]; then + echo "Could not locate ${BUILDKITE_REL}. Set REPO_ROOT or YML, run from inside the vllm-omni clone," >&2 + echo "or place this script (or run from a cwd) under the repository tree." >&2 + exit 2 + fi + YML="${REPO_ROOT}/${BUILDKITE_REL}" +fi + +LOG_DIR="${LOG_DIR:-${REPO_ROOT}/logs/nightly_jobs}" + +# Require nightly YAML only when at least one non-stability test kind needs it (validated in Python too). +_needs_nightly_yml=0 +IFS=',' read -ra _TTOK <<< "${TEST_TYPE}" +for _t in "${_TTOK[@]}"; do + _t="$(printf '%s' "${_t}" | tr '[:upper:]' '[:lower:]' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')" + case "${_t}" in + perf | acc | function | all) + _needs_nightly_yml=1 + ;; + esac +done +if [[ "${_needs_nightly_yml}" -eq 1 ]] && [[ ! -f "${YML}" ]]; then + echo "YAML not found: ${YML}" >&2 + exit 2 +fi + +if ! command -v python3 >/dev/null 2>&1; then + echo "python3 is required." >&2 + exit 1 +fi + +mkdir -p "${LOG_DIR}/jobs" +LOG_DIR="$(cd "${LOG_DIR}" && pwd)" +export REPO_ROOT LOG_DIR YML LABEL_SUBSTR TEST_TYPE MODEL_TYPE DRY_RUN + +# Drop stale wrapper scripts, old tee logs, and perf manifest so a filtered run does not +# mix with the previous outputs. +if [[ "${DRY_RUN}" != "1" ]]; then + shopt -s nullglob + _stale=( "${LOG_DIR}/jobs"/*.sh ) + _stale_logs=( "${LOG_DIR}"/*.log ) + shopt -u nullglob + if ((${#_stale[@]})); then + rm -f "${_stale[@]}" + fi + if ((${#_stale_logs[@]})); then + rm -f "${_stale_logs[@]}" + fi + rm -f "${LOG_DIR}/jobs/.perf_job_keys" + rm -f "${REPO_ROOT}/logs/nightly_perf_manual.xlsx" + rm -f "${LOG_DIR}/nightly_perf_manual.xlsx" +fi + +# shellcheck disable=SC2016,SC1078,SC1079 +python3 - <<'PY' +from __future__ import annotations + +import os +import re +import stat +import sys +from pathlib import Path + +try: + import yaml +except ImportError: + print("Missing PyYAML. Install with: pip install pyyaml", file=sys.stderr) + sys.exit(1) + +REPO_ROOT = Path(os.environ["REPO_ROOT"]).resolve() +LOG_DIR = Path(os.environ["LOG_DIR"]).resolve() +YML = Path(os.environ["YML"]).resolve() +LABEL_SUBSTR = (os.environ.get("LABEL_SUBSTR") or "").strip() +DRY_RUN = os.environ.get("DRY_RUN", "0") == "1" + +ALLOWED_TEST_TYPES = frozenset({"all", "perf", "acc", "function", "stability", "local"}) +ALLOWED_MODEL_TYPES = frozenset({"all", "omni", "tts", "diffusion"}) + + +def parse_model_types(raw: str) -> list[str]: + parts = [p.strip().lower() for p in (raw or "").split(",") if p.strip()] + if not parts: + parts = ["all"] + bad = [p for p in parts if p not in ALLOWED_MODEL_TYPES] + if bad: + print( + f"Invalid MODEL_TYPE / --model-type value(s): {bad!r} " + f"(allowed: {', '.join(sorted(ALLOWED_MODEL_TYPES))})", + file=sys.stderr, + ) + sys.exit(2) + if "all" in parts: + return ["all"] + out: list[str] = [] + seen: set[str] = set() + for p in parts: + if p not in seen: + seen.add(p) + out.append(p) + return out + + +def parse_test_types(raw: str) -> list[str]: + """If 'all' appears, YAML dimension is unconstrained; extras like stability/local are kept.""" + parts = [p.strip().lower() for p in (raw or "").split(",") if p.strip()] + if not parts: + parts = ["all"] + bad = [p for p in parts if p not in ALLOWED_TEST_TYPES] + if bad: + print( + f"Invalid TEST_TYPE / --test-type value(s): {bad!r} " + f"(allowed: {', '.join(sorted(ALLOWED_TEST_TYPES))})", + file=sys.stderr, + ) + sys.exit(2) + if "all" in parts: + out: list[str] = ["all"] + if "stability" in parts: + out.append("stability") + if "local" in parts: + out.append("local") + return out + out: list[str] = [] + seen: set[str] = set() + for p in parts: + if p not in seen: + seen.add(p) + out.append(p) + return out + + +PYTEST_CMD_RE = re.compile( + r"(?:timeout\s+\S+\s+)?(?:python3? -m\s+)?pytest\s+[^\n&|;]*" +) + + +def iter_leaf_steps(steps, group: str | None = None): + """Flatten Buildkite steps; yield (leaf_step_dict, group_title).""" + for raw in steps or []: + if not isinstance(raw, dict): + continue + nested = raw.get("steps") + if isinstance(nested, list) and nested: + g = raw.get("group") + next_group: str | None + if isinstance(g, str): + next_group = g + elif g is not None: + next_group = str(g) + else: + next_group = group + yield from iter_leaf_steps(nested, next_group) + continue + if raw.get("commands"): + yield raw, group + + +def raw_command_text(step: dict) -> str: + raw = step.get("commands") or [] + if isinstance(raw, str): + raw = [raw] + text = "\n".join((c.strip() if isinstance(c, str) else "") for c in raw if c) + return text.replace("$$", "$") + + +def export_lines(text: str) -> list[str]: + out: list[str] = [] + for line in text.splitlines(): + s = line.strip() + if not s or s.startswith("#"): + continue + if s.startswith("export "): + out.append(s) + return out + + +def pytest_lines(text: str) -> list[str]: + out: list[str] = [] + for m in PYTEST_CMD_RE.finditer(text): + line = m.group(0).strip() + line_start = text.rfind("\n", 0, m.start()) + 1 + before = text[line_start : m.start()] + if before.lstrip().startswith("#"): + continue + if line: + out.append(line) + return out + + +def label_matches_test_type(label: str, test_types: list[str]) -> bool: + """perf / acc / function / all — aligned with nightly step labels (OR across test_types).""" + if "all" in test_types: + return True + has_perf = "Perf Test" in label + has_acc = "Accuracy Test" in label + is_function = not has_perf and not has_acc + for t in test_types: + if t in ("stability", "local"): + continue + if t == "perf" and has_perf: + return True + if t == "acc" and has_acc: + return True + if t == "function" and is_function: + return True + return False + + +def label_matches_model_type(label: str, model_types: list[str]) -> bool: + """Omni · / TTS · / Diffusion — OR across model_types.""" + if "all" in model_types: + return True + if "omni" in model_types and "Omni ·" in label: + return True + if "tts" in model_types and "TTS ·" in label: + return True + if "diffusion" in model_types and "Diffusion" in label: + return True + return False + + +def job_key_from_step(step: dict, label: str) -> str: + k = step.get("key") + if isinstance(k, str) and k.strip(): + return k.strip() + slug = re.sub(r"[^\w\-.]+", "_", label or "job", flags=re.UNICODE) + slug = re.sub(r"_+", "_", slug).strip("_") or "job" + return slug + + +# Fixed stability scripts (when stability is selected); MODEL_TYPES narrows which run. +STABILITY_CASES: list[tuple[str, str, tuple[str, ...]]] = [ + ("stability_qwen3_omni", "tests/dfx/stability/scripts/test_stability_qwen3_omni.py", ("omni",)), + ("stability_qwen3_tts", "tests/dfx/stability/scripts/test_stability_qwen3_tts.py", ("tts",)), + ("stability_qwen_image", "tests/dfx/stability/scripts/test_stability_qwen_image.py", ("diffusion",)), + ("stability_wan22", "tests/dfx/stability/scripts/test_stability_wan22.py", ("diffusion",)), +] + + +def _stability_model_matches(model_types: list[str], families: tuple[str, ...]) -> bool: + if "all" in model_types: + return True + return bool(set(model_types) & set(families)) + + +def _write_job_script(key: str, script_lines: list[str], jobs_dir: Path) -> None: + body = "\n".join(script_lines) + "\n" + if DRY_RUN: + print(f"=== {key} ===") + print(body) + return + job_path = jobs_dir / f"{key}.sh" + job_path.write_text(body, encoding="utf-8") + try: + mode = job_path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH + job_path.chmod(mode) + except OSError: + pass + print(f"generated {job_path}", file=sys.stderr) + + +def local_pytest_marker_expr(model_types: list[str]) -> str: + """Build -m expression: (omni|tts|diffusion) markers combined with local_model.""" + if "all" in model_types: + families: tuple[str, ...] = ("omni", "tts", "diffusion") + else: + families = tuple(model_types) + if len(families) == 1: + return f"{families[0]} and local_model" + return f"({' or '.join(families)}) and local_model" + + +def run_local_mode(jobs_dir: Path, model_types: list[str]) -> int: + """Single job: pytest -sv -m ' and local_model' (testpaths from pytest config).""" + marker_expr = local_pytest_marker_expr(model_types) + key = "local_pytest" + pytest_line = f'pytest -sv -m "{marker_expr}"' + script_lines = [ + "#!/usr/bin/env bash", + f"# Local marker tests — MODEL_TYPE={','.join(model_types)}", + "set -euo pipefail", + f'cd "{REPO_ROOT}"', + pytest_line, + ] + _write_job_script(key, script_lines, jobs_dir) + return 1 + + +def run_stability_mode(jobs_dir: Path, model_types: list[str]) -> int: + matched = 0 + for key, rel_posix, families in STABILITY_CASES: + if not _stability_model_matches(model_types, families): + continue + if LABEL_SUBSTR: + fn = Path(rel_posix).name + if ( + LABEL_SUBSTR not in rel_posix + and LABEL_SUBSTR not in key + and LABEL_SUBSTR not in fn + ): + continue + rel_path = REPO_ROOT / rel_posix + if not rel_path.is_file(): + print(f"# skip (missing file): {rel_path}", file=sys.stderr) + continue + pytest_line = f"pytest -s -v {rel_posix}" + matched += 1 + script_lines = [ + "#!/usr/bin/env bash", + f"# Stability: {key} — {rel_posix}", + "set -euo pipefail", + f'cd "{REPO_ROOT}"', + pytest_line, + ] + _write_job_script(key, script_lines, jobs_dir) + return matched + + +def main() -> None: + test_types = parse_test_types(os.environ.get("TEST_TYPE", "all")) + model_types = parse_model_types(os.environ.get("MODEL_TYPE", "all")) + + want_stability = "stability" in test_types + want_local = "local" in test_types + yaml_test_types = [t for t in test_types if t not in ("stability", "local")] + needs_yaml = bool(yaml_test_types) + + jobs_dir = LOG_DIR / "jobs" + if not DRY_RUN: + jobs_dir.mkdir(parents=True, exist_ok=True) + + matched_stability = 0 + matched_yaml = 0 + perf_job_keys: list[str] = [] + + if want_stability: + matched_stability = run_stability_mode(jobs_dir, model_types) + + if want_local: + run_local_mode(jobs_dir, model_types) + + if needs_yaml: + if not YML.is_file(): + print(f"YAML not found: {YML}", file=sys.stderr) + sys.exit(1) + + data = yaml.safe_load(YML.read_text(encoding="utf-8")) + top_steps = (data or {}).get("steps") or [] + + for step, _grp in iter_leaf_steps(top_steps): + label = step.get("label") or "" + if LABEL_SUBSTR and LABEL_SUBSTR not in label: + continue + if not label_matches_test_type(label, yaml_test_types): + continue + if not label_matches_model_type(label, model_types): + continue + text = raw_command_text(step) + exports = export_lines(text) + pys = pytest_lines(text) + if not pys: + print(f"# skip (no pytest line): {label!r}", file=sys.stderr) + continue + + key = job_key_from_step(step, label) + matched_yaml += 1 + + script_lines = [ + "#!/usr/bin/env bash", + f'# From Buildkite label: {label.replace(chr(10), " ")}', + "set -euo pipefail", + f'cd "{REPO_ROOT}"', + ] + script_lines.extend(exports) + script_lines.extend(pys) + _write_job_script(key, script_lines, jobs_dir) + if "Perf Test" in label and key not in perf_job_keys: + perf_job_keys.append(key) + + if not DRY_RUN: + manifest_path = LOG_DIR / "jobs" / ".perf_job_keys" + if perf_job_keys: + manifest_path.write_text("\n".join(perf_job_keys) + "\n", encoding="utf-8") + elif manifest_path.is_file(): + manifest_path.unlink() + + errs: list[str] = [] + if want_stability and matched_stability == 0: + errs.append( + f"No stability jobs matched TEST_TYPE={test_types!r} MODEL_TYPE={model_types!r} " + f"LABEL_SUBSTR={LABEL_SUBSTR!r} under {REPO_ROOT}" + ) + if needs_yaml and matched_yaml == 0: + errs.append( + f"No YAML steps matched TEST_TYPE={test_types!r} MODEL_TYPE={model_types!r} " + f"LABEL_SUBSTR={LABEL_SUBSTR!r} in {YML}" + ) + if errs: + for e in errs: + print(e, file=sys.stderr) + sys.exit(2) + + +main() +PY + +if [[ "${DRY_RUN}" == "1" ]]; then + exit 0 +fi + +# Run jobs: YAML perf steps first (if manifest exists), then generate Excel (always when +# manifest non-empty, even if some perf jobs failed; Excel uses whatever JSON exists on disk), +# then remaining jobs. Paths align with nightly Buildkite (tests/dfx/perf/results). +run_generated_jobs_with_tee() { + set -o pipefail + local any_fail=0 + local perf_list="${LOG_DIR}/jobs/.perf_job_keys" + local -a perf_keys=() + if [[ -f "${perf_list}" ]] && [[ -s "${perf_list}" ]]; then + mapfile -t perf_keys < "${perf_list}" + fi + + _run_one_job() { + local _job="$1" + local base out job_status + base="$(basename "${_job}" .sh)" + out="${LOG_DIR}/${base}.log" + echo "==> ${_job} (tee ${out})" >&2 + (cd "${REPO_ROOT}" && bash "${_job}") 2>&1 | tee "${out}" + job_status="${PIPESTATUS[0]}" + return "${job_status}" + } + + local k + for k in "${perf_keys[@]}"; do + [[ -n "${k}" ]] || continue + if [[ ! -f "${LOG_DIR}/jobs/${k}.sh" ]]; then + echo "warning: perf job script missing, skip: ${LOG_DIR}/jobs/${k}.sh" >&2 + any_fail=1 + continue + fi + _run_one_job "${LOG_DIR}/jobs/${k}.sh" || any_fail=1 + done + + if ((${#perf_keys[@]})); then + local excel_out excel_log excel_status excel_repo_logs + excel_repo_logs="${REPO_ROOT}/logs" + mkdir -p "${excel_repo_logs}" + excel_out="${excel_repo_logs}/nightly_perf_$(date -u +%Y%m%d-%H%M%S).xlsx" + excel_log="${excel_repo_logs}/generate_nightly_perf_excel.log" + if [[ "${any_fail}" -ne 0 ]]; then + echo "Note: one or more perf jobs failed; still running generate_nightly_perf_excel.py " \ + "(report reflects JSON already under tests/dfx/perf/results)." >&2 + fi + echo "==> python3 tools/nightly/generate_nightly_perf_excel.py -> ${excel_out} (tee ${excel_log})" >&2 + ( + cd "${REPO_ROOT}" && python3 tools/nightly/generate_nightly_perf_excel.py \ + --input-dir "${REPO_ROOT}/tests/dfx/perf/results" \ + --diffusion-input-dir "${REPO_ROOT}/tests/dfx/perf/results" \ + --output-file "${excel_out}" + ) 2>&1 | tee "${excel_log}" + excel_status="${PIPESTATUS[0]}" + if [[ "${excel_status}" -ne 0 ]]; then + any_fail=1 + echo "generate_nightly_perf_excel.py failed (exit ${excel_status}). See ${excel_log}" >&2 + fi + fi + + local -A _perf_done=() + for k in "${perf_keys[@]}"; do + [[ -n "${k}" ]] && _perf_done["${k}"]=1 + done + + local _job base + shopt -s nullglob + for _job in "${LOG_DIR}/jobs"/*.sh; do + base="$(basename "${_job}" .sh)" + [[ ${_perf_done["${base}"]+isset} ]] && continue + _run_one_job "${_job}" || any_fail=1 + done + shopt -u nullglob + + if [[ "${any_fail}" -ne 0 ]]; then + echo "One or more selected jobs failed. See logs under ${LOG_DIR}." >&2 + return 1 + fi + echo "All selected jobs finished OK. Logs: ${LOG_DIR}/*.log" >&2 + return 0 +} + +run_generated_jobs_with_tee || exit 1