diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index aafbfcb1a..de3057a80 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -50,6 +50,8 @@ jobs: run: pip install pre-commit - name: Run pre-commit (all hooks, all files) run: pre-commit run --all-files --show-diff-on-failure + - name: Check version/commit pin consistency + run: python tools/ci/check_version_consistency.py # Flag PRs that introduce known-vulnerable dependencies. Warn-only during # burn-in; drop warn-only to turn it into a hard gate once the team is ready. @@ -237,6 +239,7 @@ jobs: - run: echo "Begin AITER + Primus-Turbo Install." - name: Install AITER run: | + : > "$RUNNER_TEMP/runtime.tsv" # reset the CI runtime log for this job echo "✅ [Uninstall old aiter] started at: $(date)" pip3 uninstall aiter amd-aiter -y || true rm -rf /tmp/aiter || true @@ -253,6 +256,7 @@ jobs: elapsed=$((end_time - start_time)) echo "✅ [Build aiter] ended at: $(date)" echo "⏱️ [Build aiter] Total elapsed time: ${elapsed} seconds" + echo -e "Build aiter\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv" - name: Install Primus-Turbo run: | rm -rf /tmp/Primus-Turbo || true @@ -268,6 +272,7 @@ jobs: elapsed=$((end_time - start_time)) echo "✅ [Pip install requirements] ended at: $(date)" echo "⏱️ [Pip install requirements] Total elapsed time: ${elapsed} seconds" + echo -e "primus-turbo: pip install requirements\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv" start_time=$(date +%s) echo "✅ [build primus-turbo] started at: $(date)" pip3 install --no-build-isolation -e . -v @@ -275,6 +280,7 @@ jobs: elapsed=$((end_time - start_time)) echo "✅ [build primus-turbo] ended at: $(date)" echo "⏱️ [build primus-turbo] Total elapsed time: ${elapsed} seconds" + echo -e "primus-turbo: build/install\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv" - run: echo "🎉 Begin Primus Unit Test." - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 with: @@ -331,7 +337,21 @@ jobs: # Note HSA_NO_SCRATCH_RECLAIM=1 must be set to avoid RCCL perf hit (TAS-8N Node), rocm ver:70125424 export HSA_NO_SCRATCH_RECLAIM=1 mkdir -p "${GITHUB_WORKSPACE}/test-reports" - pytest --maxfail=1 -s ./tests/unit_tests/ \ + # Component-aware selection on PRs; full suite on push/release/dispatch. + # Fail-safe: any failure to compute the diff falls back to the full suite. + TARGETS="./tests/unit_tests/" + if [[ "${{ github.event_name }}" == "pull_request" ]]; then + base_sha="${{ github.event.pull_request.base.sha }}" + git fetch --no-tags --depth=200 origin "${{ github.base_ref }}" 2>/dev/null || true + changed="$(git diff --name-only "${base_sha}" HEAD 2>/dev/null || true)" + if [[ -n "${changed}" ]]; then + sel="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py)" + [[ -n "${sel}" ]] && TARGETS="${sel}" + fi + fi + echo "Selected unit-test targets: ${TARGETS}" + # shellcheck disable=SC2086 # intentional word-splitting of multiple paths + pytest --maxfail=1 -s ${TARGETS} \ --cov=primus --cov-report=term-missing:skip-covered \ --junitxml="${GITHUB_WORKSPACE}/test-reports/core-unit.xml" \ --deselect=tests/unit_tests/megatron/cco/test_tp_overlap.py::TPOverlapTestCase::test_fp8_te_linear \ @@ -361,10 +381,31 @@ jobs: printf '[run]\nparallel = true\nsource = primus\nsigterm = true\n' > "$GITHUB_WORKSPACE/.coveragerc_e2e" echo "COVERAGE_PROCESS_START=$GITHUB_WORKSPACE/.coveragerc_e2e" >> "$GITHUB_ENV" echo "COVERAGE_FILE=$GITHUB_WORKSPACE/.coverage_e2e" >> "$GITHUB_ENV" + - name: Decide E2E scope (torch) + run: | + # Fail-safe: default to running all E2E (push/release/dispatch, or any + # diff failure). On PRs, narrow to the suites the changed files affect. + M=1; T=1 + if [[ "${{ github.event_name }}" == "pull_request" ]]; then + base="${{ github.event.pull_request.base.sha }}" + git fetch --no-tags --depth=200 origin "${{ github.base_ref }}" 2>/dev/null || true + changed="$(git diff --name-only "${base}" HEAD 2>/dev/null || true)" + if [[ -n "${changed}" ]]; then + e2e="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py --e2e)" + echo "Selected torch E2E scope: ${e2e:-}" + if [[ "${e2e}" != "all" ]]; then + echo "${e2e}" | grep -qw megatron || M=0 + echo "${e2e}" | grep -qw torchtitan || T=0 + fi + fi + fi + echo "RUN_MEGATRON_E2E=${M}" >> "$GITHUB_ENV" + echo "RUN_TORCHTITAN_E2E=${T}" >> "$GITHUB_ENV" - name: Run Primus Model Tests -- Megatron-LM env: HF_TOKEN: ${{secrets.HF_TOKEN}} run: | + if [[ "${RUN_MEGATRON_E2E:-1}" != "1" ]]; then echo "Skipping Megatron-LM E2E: no relevant changes."; exit 0; fi echo "Set UT_LOG_PATH: ${{ env.UT_LOG_PATH }}" rm -rf "${{ env.UT_LOG_PATH }}" mkdir -p "${{ env.UT_LOG_PATH }}" @@ -377,6 +418,7 @@ jobs: env: HF_TOKEN: ${{secrets.HF_TOKEN}} run: | + if [[ "${RUN_TORCHTITAN_E2E:-1}" != "1" ]]; then echo "Skipping TorchTitan E2E: no relevant changes."; exit 0; fi echo "Set UT_LOG_PATH: ${{ env.UT_LOG_PATH }}" rm -rf "${{ env.UT_LOG_PATH }}" mkdir -p "${{ env.UT_LOG_PATH }}" @@ -416,6 +458,10 @@ jobs: if: always() run: | python tools/ci/junit_summary.py --title torch test-reports/*.xml >> "$GITHUB_STEP_SUMMARY" || true + - name: Write runtime summary + if: always() + run: | + python tools/ci/runtime_summary.py --title torch "$RUNNER_TEMP/runtime.tsv" >> "$GITHUB_STEP_SUMMARY" || true - name: Build torch coverage json (unit + E2E) if: always() continue-on-error: true @@ -480,6 +526,7 @@ jobs: echo "Primus-Turbo dir: /tmp/Primus-Turbo" git config --global --add safe.directory /tmp/Primus-Turbo cd /tmp/Primus-Turbo + : > "$RUNNER_TEMP/runtime.tsv" # reset the CI runtime log for this job start_time=$(date +%s) echo "✅ [Pip install requirements] started at: $(date)" mkdir -p ${PRIMUS_WORKDIR}/primus-cache @@ -488,6 +535,7 @@ jobs: elapsed=$((end_time - start_time)) echo "✅ [Pip install requirements] ended at: $(date)" echo "⏱️ [Pip install requirements] Total elapsed time: ${elapsed} seconds" + echo -e "pip install/upgrade\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv" start_time=$(date +%s) echo "✅ [build primus-turbo] started at: $(date)" end_time=$(date +%s) @@ -541,10 +589,29 @@ jobs: printf '[run]\nparallel = true\nsource = primus\nsigterm = true\n' > "$GITHUB_WORKSPACE/.coveragerc_e2e" echo "COVERAGE_PROCESS_START=$GITHUB_WORKSPACE/.coveragerc_e2e" >> "$GITHUB_ENV" echo "COVERAGE_FILE=$GITHUB_WORKSPACE/.coverage_e2e" >> "$GITHUB_ENV" + - name: Decide E2E scope (jax) + run: | + # Fail-safe: default to running MaxText E2E; on PRs skip it when no + # relevant (maxtext / shared) changes are detected. + X=1 + if [[ "${{ github.event_name }}" == "pull_request" ]]; then + base="${{ github.event.pull_request.base.sha }}" + git fetch --no-tags --depth=200 origin "${{ github.base_ref }}" 2>/dev/null || true + changed="$(git diff --name-only "${base}" HEAD 2>/dev/null || true)" + if [[ -n "${changed}" ]]; then + e2e="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py --e2e)" + echo "Selected jax E2E scope: ${e2e:-}" + if [[ "${e2e}" != "all" ]]; then + echo "${e2e}" | grep -qw maxtext || X=0 + fi + fi + fi + echo "RUN_MAXTEXT_E2E=${X}" >> "$GITHUB_ENV" - name: Run Unit Tests env: HF_TOKEN: ${{secrets.HF_TOKEN}} run: | + if [[ "${RUN_MAXTEXT_E2E:-1}" != "1" ]]; then echo "Skipping MaxText E2E: no relevant changes."; exit 0; fi echo "Set UT_LOG_PATH: ${{ env.UT_LOG_PATH }}" rm -rf "${{ env.UT_LOG_PATH }}" mkdir -p "${{ env.UT_LOG_PATH }}" @@ -567,6 +634,10 @@ jobs: if: always() run: | python tools/ci/junit_summary.py --title jax test-reports/*.xml >> "$GITHUB_STEP_SUMMARY" || true + - name: Write runtime summary + if: always() + run: | + python tools/ci/runtime_summary.py --title jax "$RUNNER_TEMP/runtime.tsv" >> "$GITHUB_STEP_SUMMARY" || true - name: Build jax coverage json (MaxText E2E) if: always() continue-on-error: true diff --git a/tests/unit_tests/ci/test_select_tests.py b/tests/unit_tests/ci/test_select_tests.py new file mode 100644 index 000000000..91031635f --- /dev/null +++ b/tests/unit_tests/ci/test_select_tests.py @@ -0,0 +1,104 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for tools/ci/select_tests.py (classify-based PR test selection).""" + +import importlib.util +from pathlib import Path + +_ROOT = Path(__file__).resolve().parents[3] +_SPEC = importlib.util.spec_from_file_location("select_tests", _ROOT / "tools/ci/select_tests.py") +_MOD = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_MOD) +select = _MOD.select_targets + +FULL = ["tests/unit_tests/"] + + +# --- unit-test selection --------------------------------------------------- +def test_empty_runs_full(): + assert select([]) == FULL + + +def test_global_change_runs_full(): + assert select([".github/workflows/ci.yaml"]) == FULL + assert select(["tools/ci/select_tests.py"]) == FULL + assert select(["runner/helpers/x.sh"]) == FULL + assert select(["requirements.txt"]) == FULL + assert select(["primus/core/launcher/initialize.py"]) == FULL + + +def test_non_py_under_primus_runs_full(): + # configs / fixtures can't be localized to a unit dir -> fail-safe. + assert select(["primus/configs/x.yaml"]) == FULL + + +def test_backend_maps_to_its_unit_dir(): + out = select(["primus/backends/megatron/training/global_vars.py"]) + assert "tests/unit_tests/backends/megatron/" in out + assert "tests/unit_tests/megatron/" in out # megatron's extra GPU-operator tests + + +def test_backend_without_unit_dir_runs_full(): + # transformer_engine has no tests/unit_tests/backends/transformer_engine/. + assert select(["primus/backends/transformer_engine/x.py"]) == FULL + + +def test_component_maps_to_isomorphic_dir(): + assert select(["primus/core/projection/engine.py"]) == ["tests/unit_tests/core/projection/"] + assert select(["primus/agents/a.py"]) == ["tests/unit_tests/agents/"] + + +def test_changed_unit_test_runs_its_dir(): + assert select(["tests/unit_tests/agents/test_tools.py"]) == ["tests/unit_tests/agents/"] + + +def test_docs_only_runs_full(): + assert select(["README.md", "docs/guide.md"]) == FULL + + +# --- E2E selection (pass a fixed suite set for determinism) ---------------- +SUITES = {"megatron", "torchtitan", "maxtext"} + + +def e2e(files): + return _MOD.select_e2e(files, SUITES) + + +def test_e2e_empty_runs_all(): + assert set(e2e([])) == SUITES + + +def test_e2e_global_runs_all(): + assert set(e2e([".github/workflows/ci.yaml"])) == SUITES + assert set(e2e(["runner/helpers/x.sh"])) == SUITES + + +def test_e2e_backend_with_trainer_runs_its_suite(): + assert e2e(["primus/backends/megatron/x.py"]) == ["megatron"] + assert e2e(["examples/maxtext/configs/x.yaml"]) == ["maxtext"] + assert e2e(["tests/trainer/test_torchtitan_trainer.py"]) == ["torchtitan"] + + +def test_e2e_backend_without_trainer_runs_all(): + # No trainer suite (bridge / hummingbirdxt / transformer_engine) -> fail-safe all. + assert set(e2e(["primus/backends/megatron_bridge/x.py"])) == SUITES + assert set(e2e(["primus/backends/hummingbirdxt/x.py"])) == SUITES + assert set(e2e(["primus/backends/transformer_engine/x.py"])) == SUITES + + +def test_e2e_component_change_runs_all(): + assert set(e2e(["primus/core/trainer/base.py"])) == SUITES + + +def test_e2e_docs_only_runs_none(): + assert e2e(["README.md"]) == [] + + +def test_e2e_added_trainer_is_auto_picked_up(): + # Adding a bridge trainer makes a bridge change run it -- no code change here. + suites = SUITES | {"megatron_bridge"} + assert _MOD.select_e2e(["primus/backends/megatron_bridge/x.py"], suites) == ["megatron_bridge"] diff --git a/tests/unit_tests/tools/test_utils.py b/tests/unit_tests/tools/test_tools_utils.py similarity index 100% rename from tests/unit_tests/tools/test_utils.py rename to tests/unit_tests/tools/test_tools_utils.py diff --git a/tools/ci/check_version_consistency.py b/tools/ci/check_version_consistency.py new file mode 100644 index 000000000..183837011 --- /dev/null +++ b/tools/ci/check_version_consistency.py @@ -0,0 +1,123 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Fail CI on cross-file config drift: a value duplicated across files that fell +out of sync. Runs in the lint job. Covers BASE_IMAGE (ci.yaml vs Dockerfile +ARG), primus.__version__ (PEP 440), pyproject deps vs requirements.txt, action +SHA-pinning, and workflow python-version. Stdlib-only. +""" + +import re +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +CI_YAML = ROOT / ".github/workflows/ci.yaml" +DOCKERFILE = ROOT / ".github/workflows/docker/Dockerfile" +INIT_PY = ROOT / "primus/__init__.py" +PYPROJECT = ROOT / "pyproject.toml" +REQUIREMENTS = ROOT / "requirements.txt" +WORKFLOWS = sorted((ROOT / ".github/workflows").glob("*.y*ml")) + +PEP440 = re.compile(r"^\d+(\.\d+)*((a|b|rc)\d+)?(\.post\d+)?(\.dev\d+)?$") +SHA40 = re.compile(r"^[0-9a-f]{40}$") +NAME_SPEC = re.compile(r"^([A-Za-z0-9._-]+)\s*(.*)$") + + +def _canon(name): + return re.sub(r"[-_.]+", "-", name).lower() + + +def _find(text, pattern, what, where): + match = re.search(pattern, text, re.MULTILINE) + if match is None: + sys.exit(f"ERROR: could not find {what} in {where}") + return match.group(1).strip() + + +def _parse_dep(line): + line = line.split("#", 1)[0].strip() + match = NAME_SPEC.match(line) if line else None + return (_canon(match.group(1)), match.group(2).replace(" ", "")) if match else None + + +def check_base_image(errors): + ci = _find(CI_YAML.read_text(), r"^\s*BASE_IMAGE:\s*(\S+)", "BASE_IMAGE env", CI_YAML) + docker = _find(DOCKERFILE.read_text(), r"^ARG\s+BASE_IMAGE=(\S+)", "ARG BASE_IMAGE default", DOCKERFILE) + if ci != docker: + errors.append(f"BASE_IMAGE drift: ci.yaml={ci!r} vs Dockerfile default={docker!r}.") + + +def check_version(errors): + version = _find(INIT_PY.read_text(), r'__version__\s*=\s*["\']([^"\']+)["\']', "__version__", INIT_PY) + if not PEP440.match(version): + errors.append(f"primus.__version__={version!r} is not a valid PEP 440 version.") + + +def check_deps(errors): + block = re.search(r"^dependencies\s*=\s*\[(.*?)\]", PYPROJECT.read_text(), re.S | re.M) + if block is None: + errors.append("could not find [project].dependencies in pyproject.toml") + return + pyproject_deps = dict(filter(None, (_parse_dep(d) for d in re.findall(r'"([^"]+)"', block.group(1))))) + req = dict(filter(None, (_parse_dep(line) for line in REQUIREMENTS.read_text().splitlines()))) + # pyproject deps must be a subset of requirements with matching specifiers; + # requirements may carry extra dev/CI-only entries (e.g. hip-python). + for name, spec in sorted(pyproject_deps.items()): + if name not in req: + errors.append(f"{name!r} is in pyproject.toml but missing from requirements.txt.") + elif req[name] != spec: + errors.append(f"spec drift for {name!r}: pyproject={spec!r} vs requirements={req[name]!r}.") + + +def check_pinned_actions(errors): + for wf in WORKFLOWS: + for raw in re.findall(r"uses:\s*(\S+)", wf.read_text()): + uses = raw.strip().strip('"').strip("'") + if uses.startswith("./"): + continue + if "@" not in uses: + errors.append(f"{wf.name}: action {uses!r} is not pinned (no @ref).") + elif not SHA40.match(uses.rsplit("@", 1)[1]): + errors.append(f"{wf.name}: action {uses!r} is not pinned to a 40-hex commit SHA.") + + +def check_python_versions(errors): + requires = _find( + PYPROJECT.read_text(), r'requires-python\s*=\s*["\']([^"\']+)["\']', "requires-python", PYPROJECT + ) + floor = re.search(r">=\s*(\d+)\.(\d+)", requires) + min_ver = (int(floor.group(1)), int(floor.group(2))) if floor else None + versions = set() + for wf in WORKFLOWS: + versions.update(re.findall(r'python-version:\s*\[?\s*["\'](\d+\.\d+)["\']', wf.read_text())) + if len(versions) > 1: + errors.append(f"workflow python-version values disagree: {sorted(versions)}.") + for v in versions if min_ver else []: + if tuple(int(x) for x in v.split(".")) < min_ver: + errors.append( + f"workflow python-version {v} is below requires-python >={min_ver[0]}.{min_ver[1]}." + ) + + +def main(): + errors = [] + check_base_image(errors) + check_version(errors) + check_deps(errors) + check_pinned_actions(errors) + check_python_versions(errors) + if errors: + print("CI consistency check FAILED:") + for err in errors: + print(f" - {err}") + return 1 + print("CI consistency OK (base image, version, deps, action pinning, python-version).") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/ci/runtime_summary.py b/tools/ci/runtime_summary.py new file mode 100644 index 000000000..dd55ff08c --- /dev/null +++ b/tools/ci/runtime_summary.py @@ -0,0 +1,70 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Render CI stage wall-clock times (a stageseconds TSV) as a Markdown table. + +Complements junit_summary (test time) by surfacing the heavy build/install +stages. A missing/empty TSV renders nothing and exits 0 so the step never fails +the job; stage order (i.e. execution order) is preserved. +""" + +import argparse + + +def fmt(seconds): + seconds = int(seconds) + hours, rem = divmod(seconds, 3600) + minutes, secs = divmod(rem, 60) + if hours: + return f"{hours}h{minutes:02d}m{secs:02d}s" + if minutes: + return f"{minutes}m{secs:02d}s" + return f"{secs}s" + + +def parse(path): + rows = [] + try: + with open(path) as handle: + for line in handle: + parts = line.rstrip("\n").split("\t") + if len(parts) != 2 or not parts[0].strip(): + continue + try: + rows.append((parts[0].strip(), float(parts[1].strip()))) + except ValueError: + continue + except OSError: + return [] + return rows + + +def render(rows, title=None): + suffix = f" - {title}" if title else "" + lines = [f"## CI runtime{suffix}\n", "| Stage | Time |", "|---|--:|"] + total = 0.0 + for stage, secs in rows: + lines.append(f"| {stage} | {fmt(secs)} |") + total += secs + lines.append(f"| **TOTAL (timed stages)** | **{fmt(total)}** |") + return "\n".join(lines) + + +def main(): + ap = argparse.ArgumentParser(description="Render a stageseconds TSV as a Markdown runtime table.") + ap.add_argument("tsv", help="Path to the runtime TSV (stageseconds per line).") + ap.add_argument("--title", default=None, help="Optional section title (e.g. torch).") + args = ap.parse_args() + + rows = parse(args.tsv) + if not rows: + return 0 # nothing timed; don't emit an empty table or fail the step + print(render(rows, args.title)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/ci/select_tests.py b/tools/ci/select_tests.py new file mode 100644 index 000000000..cfff17855 --- /dev/null +++ b/tools/ci/select_tests.py @@ -0,0 +1,159 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Map a PR's changed files (stdin, one per line) to the tests to run. + +Default prints the minimal unit-test paths; --e2e prints the E2E trainer suites +(or "all"). A single classify() decides each path's blast radius and both +selections build on it. Conventions over hard-coded tables: + - unit dirs mirror the source tree (primus/ -> tests/unit_tests/), + resolved by walking up to the nearest existing dir; + - E2E suites are auto-discovered from tests/trainer/test__trainer.py; + - a backend is named by its dir (primus/backends/ or examples/). +Fail-safe is the only invariant: anything global, unlocatable, or a backend +without a trainer expands to everything -- over-select, never under-select. + + git diff --name-only "$BASE" HEAD | python tools/ci/select_tests.py [--e2e] +""" + +import re +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +FULL = "tests/unit_tests/" + +# The only hard-coded list: changes whose blast radius is the whole repo. Being +# absent here only ever falls back to other fail-safe paths, never to "skip". +GLOBAL_TRIGGERS = ( + ".github/", + "tools/", + "runner/", # the launcher drives all training + "pyproject.toml", + "primus/__init__.py", + "primus/core/launcher/", + "primus/core/utils/", + "primus/core/config/", + "tests/utils.py", + "tests/conftest.py", + "tests/unit_tests/conftest.py", + "tests/run_unit_tests.py", +) + +# megatron's GPU-operator tests aren't path-isomorphic to the backend source. +_BACKEND_EXTRA_UNIT = {"megatron": ("tests/unit_tests/megatron/",)} + +_TRAINER_RE = re.compile(r"test_(.+)_trainer") +_BACKEND_RE = re.compile(r"(?:primus/backends|examples)/([^/]+)/") + + +def discover_e2e_suites(): + return { + m.group(1) + for p in (ROOT / "tests/trainer").glob("test_*_trainer.py") + if (m := _TRAINER_RE.fullmatch(p.stem)) + } + + +def _is_global(path): + if "/" not in path and path.startswith("requirements") and path.endswith(".txt"): + return True + return any(path == t or path.startswith(t) for t in GLOBAL_TRIGGERS) + + +def _nearest_unit_dir(rel): + # rel is source-relative (under primus/ or tests/unit_tests/); walk up to the + # nearest existing tests/unit_tests/<...> dir, or None if none exists. + parts = rel.split("/")[:-1] + while parts: + cand = "tests/unit_tests/" + "/".join(parts) + "/" + if (ROOT / cand).is_dir(): + return cand + parts.pop() + return None + + +def classify(path): + """('global', None) | ('backend', name) | ('component', unit_dir|None) | ('ignore', None).""" + if _is_global(path): + return ("global", None) + backend = _BACKEND_RE.match(path) # primus/backends// or examples// + if backend: + return ("backend", backend.group(1)) + if path.startswith("tests/trainer/"): + m = _TRAINER_RE.search(path) + return ("backend", m.group(1)) if m else ("ignore", None) + for root in ("primus/", "tests/unit_tests/"): + if path.startswith(root): + if not path.endswith(".py"): + return ("global", None) # non-.py here (configs, fixtures) -> fail-safe + return ("component", _nearest_unit_dir(path[len(root) :])) + return ("ignore", None) # docs, README, ... outside the source/test trees + + +def select_targets(files): + files = [f.strip() for f in files if f.strip()] + if not files: + return [FULL] + targets = [] + + def add(d): + if d and d not in targets: + targets.append(d) + + for path in files: + kind, val = classify(path) + if kind == "global": + return [FULL] + if kind == "backend": + base = f"tests/unit_tests/backends/{val}/" + if not (ROOT / base).is_dir(): + return [FULL] # backend without a unit dir (e.g. transformer_engine) -> safe + add(base) + for extra in _BACKEND_EXTRA_UNIT.get(val, ()): + add(extra) + elif kind == "component": + if val is None: + return [FULL] # couldn't localize a unit dir -> safe + add(val) + # ignore -> skip + return targets or [FULL] + + +def select_e2e(files, suites=None): + suites = discover_e2e_suites() if suites is None else set(suites) + files = [f.strip() for f in files if f.strip()] + if not files: + return sorted(suites) + selected = set() + for path in files: + kind, val = classify(path) + if kind == "global": + return sorted(suites) + if kind == "backend": + if val in suites: + selected.add(val) # backend has a trainer -> run its suite + else: + return sorted(suites) # no trainer (bridge/hummingbirdxt/TE) -> all + elif kind == "component": + return sorted(suites) # non-backend source change -> all training + # ignore -> skip + return sorted(selected) + + +def main(): + files = sys.stdin.read().splitlines() + if "--e2e" in sys.argv[1:]: + suites = discover_e2e_suites() + e2e = select_e2e(files, suites) + print("all" if suites and set(e2e) == suites else " ".join(e2e)) + else: + print(" ".join(select_targets(files))) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())