diff --git a/.gitignore b/.gitignore index dfa6897a04..2742301898 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ .infx-tooling/ dist/ experimental/multiturn/vllm_benchmark/results/ +*.resolved.yaml diff --git a/benchmarks/multi_node/srt-slurm-recipe-identities.yaml b/benchmarks/multi_node/srt-slurm-recipe-identities.yaml new file mode 100644 index 0000000000..6b4ccaa1c8 --- /dev/null +++ b/benchmarks/multi_node/srt-slurm-recipe-identities.yaml @@ -0,0 +1,3 @@ +# CONFIG_FILE selectors that replaced flat recipes, mapped to the replaced +# CONFIG_FILE. Matrix fingerprints and curve identity use the replaced path. +{} diff --git a/docs/configuration-procedures.md b/docs/configuration-procedures.md index c6941009a7..dc1f2147e7 100644 --- a/docs/configuration-procedures.md +++ b/docs/configuration-procedures.md @@ -183,6 +183,15 @@ Mapping source: [`benchmarks/multi_node/srt-slurm-recipes/RECIPES.md`](../benchm Do not ship one side alone. `srtctl` reads the recipe, while matrix generation reads the master config. Recipe-only changes can mislabel results. Master-only changes do not alter the deployed recipe. +### Variants of one recipe + +Related topologies can share one override-format recipe (`base` plus `zip_override_` sections; see srt-slurm's [overrides guide](https://github.com/NVIDIA/srt-slurm/blob/main/docs/overrides.md)). Each master entry then selects exactly one variant: `CONFIG_FILE=recipes//.yaml:zip_override_[]`, `:override_`, or `:base`. Group or glob selectors are rejected because one matrix row is one job. + +- Launchers call `materialize_srt_configs` ([`runners/slurm_utils.sh`](../runners/slurm_utils.sh)) first. It writes the selected variant beside its source as `..resolved.yaml` and points `CONFIG_FILE` at it, so power detection, launcher patches and `srtctl` all see a standalone recipe. Expansion follows srtctl ([`infx/srt_slurm/recipe_selector.py`](../infx/srt_slurm/recipe_selector.py)). +- Matrix generation takes the node count from the selected variant. +- Give every zip group a `name` list; otherwise srtctl renames variants to `__`. Inside a zip section every list is a zip dimension: wrap list-valued settings in an outer list, and note that a `null` element deletes the key. +- When moving an existing flat recipe into a variant, add `: ` to [`benchmarks/multi_node/srt-slurm-recipe-identities.yaml`](../benchmarks/multi_node/srt-slurm-recipe-identities.yaml). Recipe fingerprints and curve identity then keep the old path, so history and Klaud baselines still match. Verify that the resolved variant equals the old file before deleting it. + ## Register an llm-d recipe Sources: [`benchmarks/llm-d/README.md`](../benchmarks/llm-d/README.md), [`benchmarks/multi_node/llm-d/README.md`](../benchmarks/multi_node/llm-d/README.md), [`llm-d-recipes/`](../benchmarks/multi_node/llm-d-recipes/), and the current [`llmd-vllm` benchmark wrapper](../benchmarks/multi_node/dsv4_fp4_gb200_llmd-vllm-disagg.sh). diff --git a/docs/configuration-procedures_zh.md b/docs/configuration-procedures_zh.md index 54986db75c..aa12d9bbd8 100644 --- a/docs/configuration-procedures_zh.md +++ b/docs/configuration-procedures_zh.md @@ -160,6 +160,15 @@ B200 Nscale 的 GLM-5.1 可用 `MODEL_PATH` 指定已有共享权重，覆盖默 不得只提交一侧：`srtctl` 读取配方，而矩阵生成读取主配置。仅改配方可能给结果贴错标签；仅改主配置不会改变实际部署的配方。 +### 同一配方的多个变体 + +相关拓扑可以共用一个 override 格式的配方（`base` 加 `zip_override_` 段；见 srt-slurm 的 [overrides 指南](https://github.com/NVIDIA/srt-slurm/blob/main/docs/overrides.md)）。每个主配置条目只选择一个变体：`CONFIG_FILE=recipes//.yaml:zip_override_[]`、`:override_` 或 `:base`。一个矩阵行对应一个作业，因此不接受整组或通配符选择器。 + +- 启动脚本首先调用 `materialize_srt_configs`（[`runners/slurm_utils.sh`](../runners/slurm_utils.sh)）。它把选中的变体写到源文件旁的 `..resolved.yaml`，并让 `CONFIG_FILE` 指向该文件，使功耗检测、启动脚本补丁和 `srtctl` 都看到独立配方。展开规则与 srtctl 一致（[`infx/srt_slurm/recipe_selector.py`](../infx/srt_slurm/recipe_selector.py)）。 +- 矩阵生成从选中的变体读取节点数。 +- 每个 zip 组都要提供 `name` 列表，否则 srtctl 会把变体重命名为 `__`。zip 段中的每个列表都是 zip 维度：列表值设置需再包一层列表；`null` 元素会删除该键。 +- 把现有平铺配方迁入变体时，在 [`benchmarks/multi_node/srt-slurm-recipe-identities.yaml`](../benchmarks/multi_node/srt-slurm-recipe-identities.yaml) 中添加 `<新 CONFIG_FILE>: <旧 CONFIG_FILE>`。配方指纹和曲线身份因此沿用旧路径，历史数据和 Klaud 基线仍可匹配。删除旧文件前，先确认解析后的变体与旧文件一致。 + ## 注册 llm-d 配方 来源：[`benchmarks/llm-d/README.md`](../benchmarks/llm-d/README.md)、[`benchmarks/multi_node/llm-d/README.md`](../benchmarks/multi_node/llm-d/README.md)、[`llm-d-recipes/`](../benchmarks/multi_node/llm-d-recipes/) 和当前 [`llmd-vllm` 基准 wrapper](../benchmarks/multi_node/dsv4_fp4_gb200_llmd-vllm-disagg.sh)。 diff --git a/infx/matrix/generate.py b/infx/matrix/generate.py index 5bbad46d3c..eb21f31ecc 100644 --- a/infx/matrix/generate.py +++ b/infx/matrix/generate.py @@ -11,6 +11,7 @@ import yaml from infx.config import repository_root +from infx.srt_slurm.recipe_selector import load_recipe, recipe_source, split_config_file from .validation import ( DEFAULT_AGENTIC_DURATION_SECONDS, @@ -238,23 +239,21 @@ def recipe_node_count(prefill: dict, decode: dict) -> int | None: config_file = config_files.pop() repo_root = repository_root() - recipe_root = repo_root / "benchmarks" / "multi_node" / "srt-slurm-recipes" - if config_file.startswith("benchmarks/multi_node/srt-slurm-recipes/"): - recipe_path = repo_root / config_file - else: - recipe_path = recipe_root / config_file.removeprefix("recipes/") + recipe_path = recipe_source(config_file, repo_root) if not recipe_path.exists(): # Some srt-slurm recipes live only in the runtime image. Their master # config topology remains the best available scheduling estimate. return None - recipe = yaml.safe_load(recipe_path.read_text()) + if split_config_file(config_file)[1] is None and "base" in yaml.safe_load( + recipe_path.read_text() + ): + # Without a selector srtctl submits every variant, which has no single + # node count. The master topology supplies the estimate. + return None + recipe = load_recipe(config_file, repo_root) if recipe.get("schema") != 2: raise ValueError(f"srt-slurm recipes must declare schema: 2: {recipe_path}") - if "base" in recipe: - # A file with several override variants has no single authoritative - # node count. The selected master topology supplies the estimate. - return None roles = recipe.get("roles") if roles: # Schema 2 groups node allocations by role. A colocated decode role diff --git a/infx/matrix/plan.py b/infx/matrix/plan.py index a964e1984f..5e3fe2a46d 100644 --- a/infx/matrix/plan.py +++ b/infx/matrix/plan.py @@ -18,7 +18,8 @@ import yaml -from infx.config import GENERATE_SWEEPS_PY_SCRIPT, MASTER_CONFIGS, RUNNER_CONFIG +from infx.config import GENERATE_SWEEPS_PY_SCRIPT, MASTER_CONFIGS, RUNNER_CONFIG, repository_root +from infx.srt_slurm.recipe_selector import recipe_identities from .generate import ( EvalMode, @@ -169,12 +170,36 @@ def generation_inputs_at_ref(ref: str) -> Iterator[GenerationInputs]: ) +def _recipe_identity(entry: dict) -> dict: + """Name a consolidated recipe variant by the flat recipe path it replaced. + + Moving a recipe into an override file only changes its CONFIG_FILE spelling; + keeping the old spelling here preserves fingerprints and curve identity. + """ + identities = recipe_identities(repository_root()) + if not identities: + return entry + view = dict(entry) + for role in ("prefill", "decode"): + worker = entry.get(role) + if not isinstance(worker, dict): + continue + settings = worker.get("additional-settings") or [] + mapped = [] + for setting in settings: + name, sep, value = setting.partition("=") + mapped.append(f"{name}={identities[value]}" if sep and value in identities else setting) + if mapped != settings: + view[role] = {**worker, "additional-settings": mapped} + return view + + def _matrix_curve_key(entry: dict) -> tuple: """Identify one curve while deliberately excluding point-level fields.""" return tuple( sorted( (key, freeze_config_value(value)) - for key, value in entry.items() + for key, value in _recipe_identity(entry).items() if key not in {"conc", "exp-name", "recipe-fingerprint"} ) ) @@ -184,7 +209,7 @@ def recipe_fingerprint(entry: dict) -> str: """Hash the generated recipe independently of point-level concurrency/name.""" recipe = { key: value - for key, value in entry.items() + for key, value in _recipe_identity(entry).items() if key not in {"conc", "exp-name", "recipe-fingerprint"} } canonical = json.dumps( diff --git a/infx/srt_slurm/recipe_selector.py b/infx/srt_slurm/recipe_selector.py new file mode 100644 index 0000000000..078210b016 --- /dev/null +++ b/infx/srt_slurm/recipe_selector.py @@ -0,0 +1,169 @@ +"""Resolve srt-slurm ``CONFIG_FILE`` selectors into standalone recipes. + +A master config may point one benchmark at a variant of an override-format +recipe, e.g. ``recipes/x/disagg.yaml:zip_override_lowlat[2]``. Launchers read +recipes as text (power telemetry, name and health-check patches) before srtctl +ever runs, so they materialize the selected variant as a flat recipe first. The +expansion mirrors srtctl's ``generate_override_configs`` and is checked against +it in ``utils/test_recipe_selector.py``. +""" + +from __future__ import annotations + +import argparse +import copy +import functools +import re +import sys +from pathlib import Path +from typing import Any + +import yaml + +RECIPE_ROOT = Path("benchmarks/multi_node/srt-slurm-recipes") +RUNTIME_PREFIX = "recipes/" +IDENTITIES_FILE = Path("benchmarks/multi_node/srt-slurm-recipe-identities.yaml") +_ZIP_SELECTOR = re.compile(r"(zip_override_[\w-]+)\[(\d+)\]") + + +def split_config_file(config_file: str) -> tuple[str, str | None]: + """Split ``path[:selector]``; a missing selector means a flat recipe.""" + path, sep, selector = config_file.partition(":") + return path, (selector if sep else None) + + +def recipe_source(config_file: str, repo_root: Path) -> Path: + """Map either CONFIG_FILE spelling to the checked-in recipe file.""" + path, _ = split_config_file(config_file) + if path.startswith(f"{RECIPE_ROOT}/"): + return repo_root / path + return repo_root / RECIPE_ROOT / path.removeprefix(RUNTIME_PREFIX) + + +def deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]: + """Merge like srtctl: dicts recurse, lists and scalars replace, ``None`` deletes.""" + result = copy.deepcopy(base) + for key, value in override.items(): + if value is None: + result.pop(key, None) + elif isinstance(value, dict) and isinstance(result.get(key), dict): + result[key] = deep_merge(result[key], value) + else: + result[key] = copy.deepcopy(value) + return result + + +def _list_lengths(section: dict[str, Any]) -> list[int]: + lengths = [] + for value in section.values(): + if isinstance(value, list): + lengths.append(len(value)) + elif isinstance(value, dict): + lengths.extend(_list_lengths(value)) + return lengths + + +def _zip_length(section: dict[str, Any]) -> int: + lengths = _list_lengths(section) + if not lengths or 0 in lengths: + raise ValueError("zip_override section needs non-empty list values") + widths = {n for n in lengths if n != 1} + if len(widths) > 1: + raise ValueError(f"Incompatible zip lengths {sorted(widths)}") + return widths.pop() if widths else 1 + + +def _zip_slice(section: dict[str, Any], index: int) -> dict[str, Any]: + result = {} + for key, value in section.items(): + if isinstance(value, list): + result[key] = value[0 if len(value) == 1 else index] + elif isinstance(value, dict): + result[key] = _zip_slice(value, index) + else: + result[key] = value + return result + + +def resolve_variant(raw: dict[str, Any], selector: str) -> dict[str, Any]: + """Expand exactly one variant; one CONFIG_FILE must describe one job.""" + if "base" not in raw: + raise ValueError(f"Selector {selector!r} requires an override-format recipe") + base = raw["base"] + if selector == "base": + variant = copy.deepcopy(base) + elif match := _ZIP_SELECTOR.fullmatch(selector): + key, index = match.group(1), int(match.group(2)) + if key not in raw: + raise ValueError(f"{key!r} not found in recipe") + section = raw[key] + width = _zip_length(section) + if index >= width: + raise ValueError(f"Index [{index}] out of range for {key!r} ({width} variants)") + variant = deep_merge(base, _zip_slice(section, index)) + if not isinstance(section.get("name"), list): + variant["name"] = ( + f"{base.get('name', 'unnamed')}_{key.removeprefix('zip_override_')}_{index}" + ) + elif selector.startswith("override_") and selector in raw: + variant = deep_merge(base, raw[selector]) + if "name" not in raw[selector]: + variant["name"] = f"{base.get('name', 'unnamed')}_{selector.removeprefix('override_')}" + else: + raise ValueError(f"Unsupported or unknown recipe selector {selector!r}") + if "schema" in raw: + variant.setdefault("schema", raw["schema"]) + return variant + + +def load_recipe(config_file: str, repo_root: Path) -> dict[str, Any]: + """Return the single recipe a CONFIG_FILE selects.""" + raw = yaml.safe_load(recipe_source(config_file, repo_root).read_text()) + if not isinstance(raw, dict): + raise ValueError(f"Recipe must be a mapping: {config_file}") + _, selector = split_config_file(config_file) + return raw if selector is None else resolve_variant(raw, selector) + + +def materialize(config_file: str, repo_root: Path) -> str: + """Write the selected variant beside its source and return its CONFIG_FILE.""" + path, selector = split_config_file(config_file) + if selector is None: + return config_file + recipe = load_recipe(config_file, repo_root) + source = recipe_source(config_file, repo_root) + label = re.sub(r"[^\w-]+", "-", selector).strip("-") + target = source.with_name(f"{source.stem}.{label}.resolved.yaml") + target.write_text(yaml.safe_dump(recipe, sort_keys=False, default_flow_style=False)) + return str(Path(path).with_name(target.name)) + + +@functools.cache +def recipe_identities(repo_root: Path) -> dict[str, str]: + """Selector CONFIG_FILEs that replaced flat recipes, keyed to their old path.""" + path = repo_root / IDENTITIES_FILE + if not path.exists(): + return {} + identities = yaml.safe_load(path.read_text()) or {} + if not isinstance(identities, dict): + raise ValueError(f"{IDENTITIES_FILE} must be a mapping") + return identities + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="command", required=True) + mat = sub.add_parser("materialize", help="print a flat CONFIG_FILE for a selector") + mat.add_argument("config_file") + mat.add_argument("--repo-root", type=Path, default=Path.cwd()) + args = parser.parse_args(argv) + try: + print(materialize(args.config_file, args.repo_root)) + except (OSError, ValueError, yaml.YAMLError) as error: + print(f"ERROR: recipe selector: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/runners/launch_b200-nscale-slurm.sh b/runners/launch_b200-nscale-slurm.sh index cd238f252e..16e781323c 100755 --- a/runners/launch_b200-nscale-slurm.sh +++ b/runners/launch_b200-nscale-slurm.sh @@ -18,6 +18,7 @@ check_env_vars SLURM_PARTITION SLURM_ACCOUNT # shellcheck source=runners/slurm_utils.sh source "$(dirname "${BASH_SOURCE[0]}")/slurm_utils.sh" || exit 1 +materialize_srt_configs || exit 1 set -x diff --git a/runners/launch_b300-dsxe.sh b/runners/launch_b300-dsxe.sh index a9c21d65f9..608a5642ba 100755 --- a/runners/launch_b300-dsxe.sh +++ b/runners/launch_b300-dsxe.sh @@ -5,6 +5,7 @@ check_env_vars ENROOT_IMPORT_TIME_LIMIT EVAL_ONLY IS_AGENTIC IS_MULTINODE RUN_EV # shellcheck source=runners/slurm_utils.sh source "$(dirname "${BASH_SOURCE[0]}")/slurm_utils.sh" || exit 1 +materialize_srt_configs || exit 1 # B300 DSXE Slurm cluster (dsxe-sa-b300-prd0); runners run as sa-gha-runner. # Cluster-specific facts live in this block. Multi-node jobs go through diff --git a/runners/launch_gb200-nv.sh b/runners/launch_gb200-nv.sh index 60dd034fcd..a4dec04f2f 100755 --- a/runners/launch_gb200-nv.sh +++ b/runners/launch_gb200-nv.sh @@ -7,6 +7,7 @@ check_env_vars EVAL_ONLY IS_AGENTIC IS_MULTINODE RUN_EVAL SALLOC_TIME_LIMIT set -x source "$(dirname "${BASH_SOURCE[0]}")/slurm_utils.sh" || exit 1 +materialize_srt_configs || exit 1 export SLURM_PARTITION="batch" export SLURM_ACCOUNT="benchmark" diff --git a/runners/launch_gb300-nv.sh b/runners/launch_gb300-nv.sh index 83ecb73eb6..81d976a960 100644 --- a/runners/launch_gb300-nv.sh +++ b/runners/launch_gb300-nv.sh @@ -8,6 +8,7 @@ set -exo pipefail # shellcheck source=runners/slurm_utils.sh source "$(dirname "${BASH_SOURCE[0]}")/slurm_utils.sh" || exit 1 +materialize_srt_configs || exit 1 check_env_vars SLURM_PARTITION export SBATCH_PARTITION="$SLURM_PARTITION" diff --git a/runners/launch_h100-dgxc-slurm.sh b/runners/launch_h100-dgxc-slurm.sh index 3eccf3eba5..4391ac6b55 100644 --- a/runners/launch_h100-dgxc-slurm.sh +++ b/runners/launch_h100-dgxc-slurm.sh @@ -6,6 +6,7 @@ set -e # shellcheck source=runners/slurm_utils.sh source "$(dirname "${BASH_SOURCE[0]}")/slurm_utils.sh" || exit 1 +materialize_srt_configs || exit 1 SLURM_PARTITION="hpc-gpu-1" SLURM_ACCOUNT="customer" diff --git a/runners/launch_h200-dgxc-slurm.sh b/runners/launch_h200-dgxc-slurm.sh index 36405594c2..4f69ad6755 100755 --- a/runners/launch_h200-dgxc-slurm.sh +++ b/runners/launch_h200-dgxc-slurm.sh @@ -14,6 +14,7 @@ DSV4_MODEL_REPO="deepseek-ai/DeepSeek-V4-Pro-0813" set -x source "$(dirname "${BASH_SOURCE[0]}")/slurm_utils.sh" || exit 1 +materialize_srt_configs || exit 1 if [[ "$IS_MULTINODE" == "true" ]]; then diff --git a/runners/slurm_utils.sh b/runners/slurm_utils.sh index 87e7767747..025cf4b3c8 100644 --- a/runners/slurm_utils.sh +++ b/runners/slurm_utils.sh @@ -8,6 +8,27 @@ SRTCTL_EVAL_ARGS=( --set 'post_eval.command=["bash", "{infmax_workspace}/benchmarks/multi_node/srt_eval.sh", "{endpoint}", "{infmax_workspace}"]' ) +# Replace a "recipe.yaml:" CONFIG_FILE or EVAL_CONFIG_FILE with a flat +# recipe written beside its source. Launchers read recipes as text (power +# telemetry, name and health-check patches) before srtctl runs, so each variant +# must look exactly like a standalone recipe. Flat recipes pass through untouched. +materialize_srt_configs() { + local name value + for name in CONFIG_FILE EVAL_CONFIG_FILE; do + value="${!name}" + [[ "$value" == *:* ]] || continue + check_env_vars GITHUB_WORKSPACE + # Login nodes do not all ship PyYAML; the selector path is the only user. + value=$(PYTHONPATH="$INFERENCEX_SLURM_UTILS_DIR/..${PYTHONPATH:+:$PYTHONPATH}" \ + uv run --no-project --quiet --with pyyaml \ + python3 -m infx.srt_slurm.recipe_selector materialize "$value" \ + --repo-root "$GITHUB_WORKSPACE") || return 1 + echo "Resolved $name ${!name} -> $value" + printf -v "$name" '%s' "$value" + export "${name?}" + done +} + # Write a job-local cluster config; profiles contain only native srt-slurm settings. write_srt_cluster_config() { if [[ $# -lt 3 || -z "$1" || -z "$2" || ( "$3" != 0 && "$3" != 1 ) ]]; then diff --git a/utils/test_recipe_selector.py b/utils/test_recipe_selector.py new file mode 100644 index 0000000000..8726ebd77f --- /dev/null +++ b/utils/test_recipe_selector.py @@ -0,0 +1,181 @@ +"""Resolve srt-slurm CONFIG_FILE selectors into the recipe a job actually runs.""" + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import yaml + +import infx.config +from infx.matrix import generate, plan +from infx.srt_slurm import recipe_selector +from infx.srt_slurm.recipe_selector import materialize, resolve_variant + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "utils/srt-slurm/src")) + +# The pinned upstream implementation is the contract every selector must match. +from srtctl.core.config import generate_override_configs # noqa: E402 + +RECIPE = { + "schema": 2, + "base": { + "name": "fam", + "roles": {"prefill": {"nodes": 1}, "decode": {"nodes": 2, "env": {"A": "1"}}}, + "benchmark": {"concurrencies": [4, 8]}, + "health_check": {"max_attempts": 180}, + }, + "zip_override_grid": { + "name": ["fam-a", "fam-b"], + "roles": {"decode": {"nodes": [2, 4], "env": [{"B": "2"}, None]}}, + "benchmark": {"concurrencies": [[4], [8, 16]]}, + "health_check": {"max_attempts": [None]}, + }, + "zip_override_auto": {"roles": {"prefill": {"nodes": [3, 5]}}}, + "override_solo": {"roles": {"decode": {"nodes": 7}}}, +} + + +@pytest.fixture +def repo(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + recipe = tmp_path / "benchmarks/multi_node/srt-slurm-recipes/fam/disagg.yaml" + recipe.parent.mkdir(parents=True) + recipe.write_text(yaml.safe_dump(RECIPE, sort_keys=False)) + (tmp_path / "configs").mkdir() + monkeypatch.setattr(infx.config, "__file__", str(tmp_path / "infx/config.py")) + recipe_selector.recipe_identities.cache_clear() + yield tmp_path + recipe_selector.recipe_identities.cache_clear() + + +@pytest.mark.parametrize( + "selector", + ["base", "zip_override_grid[0]", "zip_override_grid[1]", "zip_override_auto[1]", "override_solo"], +) +def test_resolve_variant_matches_srtctl(selector): + assert resolve_variant(RECIPE, selector) == generate_override_configs(RECIPE, selector)[0][1] + + +def test_resolve_variant_slices_deletes_and_names(): + assert resolve_variant(RECIPE, "zip_override_grid[1]") == { + "name": "fam-b", + "roles": {"prefill": {"nodes": 1}, "decode": {"nodes": 4}}, + "benchmark": {"concurrencies": [8, 16]}, + "health_check": {}, + "schema": 2, + } + assert resolve_variant(RECIPE, "zip_override_auto[0]")["name"] == "fam_auto_0" + + +@pytest.mark.parametrize( + ("selector", "message"), + [ + ("zip_override_grid[2]", "out of range"), + ("zip_override_grid", "Unsupported"), + ("override_*", "Unsupported"), + ], +) +def test_resolve_variant_rejects_selectors_that_are_not_one_job(selector, message): + with pytest.raises(ValueError, match=message): + resolve_variant(RECIPE, selector) + + +def test_materialize_cli_writes_a_flat_recipe_beside_its_source(repo): + result = subprocess.run( + [ + sys.executable, "-m", "infx.srt_slurm.recipe_selector", "materialize", + "recipes/fam/disagg.yaml:zip_override_grid[0]", "--repo-root", str(repo), + ], + capture_output=True, text=True, check=True, cwd=ROOT, + ) + + assert result.stdout.strip() == "recipes/fam/disagg.zip_override_grid-0.resolved.yaml" + written = repo / "benchmarks/multi_node/srt-slurm-recipes/fam/disagg.zip_override_grid-0.resolved.yaml" + assert yaml.safe_load(written.read_text()) == { + "name": "fam-a", + "roles": {"prefill": {"nodes": 1}, "decode": {"nodes": 2, "env": {"A": "1", "B": "2"}}}, + "benchmark": {"concurrencies": [4]}, + "health_check": {}, + "schema": 2, + } + # Launchers patch top-level and nested keys by indentation. + assert written.read_text().startswith("name: fam-a\n") + assert materialize("recipes/fam/disagg.yaml", repo) == "recipes/fam/disagg.yaml" + + +def test_materialize_cli_reports_bad_selectors(repo): + result = subprocess.run( + [ + sys.executable, "-m", "infx.srt_slurm.recipe_selector", "materialize", + "recipes/fam/disagg.yaml:zip_override_missing[0]", "--repo-root", str(repo), + ], + capture_output=True, text=True, cwd=ROOT, + ) + + assert result.returncode == 1 + assert "zip_override_missing" in result.stderr + + +def test_launcher_helper_replaces_selector_config_files(repo, tmp_path): + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + # Stub uv: drop its own arguments and run the requested command locally. + (bin_dir / "uv").write_text('#!/usr/bin/env bash\nshift 5\nexec "$@"\n') + (bin_dir / "uv").chmod(0o755) + (bin_dir / "python3").write_text(f'#!/usr/bin/env bash\nexec {sys.executable} "$@"\n') + (bin_dir / "python3").chmod(0o755) + script = ( + f'source {ROOT}/runners/slurm_utils.sh && materialize_srt_configs && ' + 'bash -c \'echo "$CONFIG_FILE|$EVAL_CONFIG_FILE"\'' + ) + env = { + **os.environ, + "PATH": f"{bin_dir}:{os.environ['PATH']}", + "GITHUB_WORKSPACE": str(repo), + "CONFIG_FILE": "recipes/fam/disagg.yaml:override_solo", + "EVAL_CONFIG_FILE": "recipes/fam/flat.yaml", + } + + result = subprocess.run( + ["bash", "-c", script], capture_output=True, text=True, env=env, check=True + ) + + assert result.stdout.strip().splitlines()[-1] == ( + "recipes/fam/disagg.override_solo.resolved.yaml|recipes/fam/flat.yaml" + ) + + +def test_node_count_reads_the_selected_variant(repo): + def settings(config_file): + return {"additional-settings": [f"CONFIG_FILE={config_file}"]} + + assert generate.recipe_node_count(settings("recipes/fam/disagg.yaml:zip_override_grid[1]"), {}) == 5 + assert generate.recipe_node_count(settings("recipes/fam/disagg.yaml:override_solo"), {}) == 8 + # Without a selector srtctl submits every variant; the master estimate applies. + assert generate.recipe_node_count(settings("recipes/fam/disagg.yaml"), {}) is None + + +def test_fingerprint_keeps_the_identity_of_a_consolidated_recipe(repo): + def entry(config_file): + return { + "model": "m", + "conc": [4], + "prefill": {"tp": 4, "additional-settings": [f"CONFIG_FILE={config_file}", "X=1"]}, + "decode": {"tp": 8, "additional-settings": [f"EVAL_CONFIG_FILE={config_file}"]}, + } + + old = entry("recipes/fam/disagg-1p2d.yaml") + new = entry("recipes/fam/disagg.yaml:zip_override_grid[0]") + before = plan.recipe_fingerprint(new) + (repo / recipe_selector.IDENTITIES_FILE).write_text( + yaml.safe_dump({"recipes/fam/disagg.yaml:zip_override_grid[0]": "recipes/fam/disagg-1p2d.yaml"}) + ) + recipe_selector.recipe_identities.cache_clear() + + assert before != plan.recipe_fingerprint(old) + assert plan.recipe_fingerprint(new) == plan.recipe_fingerprint(old) + assert plan.recipe_fingerprint(entry("recipes/fam/disagg.yaml:zip_override_grid[1]")) != ( + plan.recipe_fingerprint(old) + )