diff --git a/agent-flow/agent_flow/workflows/perf_optimize/bench_cli.py b/agent-flow/agent_flow/workflows/perf_optimize/bench_cli.py new file mode 100644 index 000000000000..e785ac624065 --- /dev/null +++ b/agent-flow/agent_flow/workflows/perf_optimize/bench_cli.py @@ -0,0 +1,422 @@ +"""Reading the benchmark suite's configs and results, without driving it. + +A SOL-track campaign measures through `ibc-bench`, the CLI the +`ibc-trtllm-harness` package ships. The agent drives it — `submit sweep`, +`jobs check`, `process ctx`, `process frontier` are Bash commands in the +track's prompt section, run by the role that needs them, the same way an +aggregate campaign runs `trtllm-serve` and `benchmark_serving`. + +Python's part is smaller than it looks, and deliberately so: it reads two +kinds of file the harness owns, and runs nothing. + +- **The sweep** — at schema validation, before any agent exists, so + `task.yaml` can be reconciled against the run that will happen. The + expansion is computed here rather than asked of the CLI because + `ibc-bench` has no read-only "what would this submit" command: + `submit sweep --dry-run` answers, but by writing a case directory per + case, which is not something schema validation may do. The rule it + applies is the harness's own and is narrow — a `gen_configs` row is + `[ctx, gen, tp, batch, max_num_tokens, attention_dp, gpu_mem_frac, mtp, + eplb, concurrency_list]` and expands over that last field; a ctx + `benchmarks` entry expands over `max_batch` × `tp_size` × `mtp_range`. + The agent's own `--dry-run` step re-derives it against the harness + before an allocation is spent, so a drift shows up as a refusal rather + than as a campaign quoting points it never measured. + +- **The frontier CSV** — what `ibc-bench process frontier` writes. Its + columns are already the names this workflow scores on: + `throughput_per_user` needs no translation, and `output_tput_per_gpu` + and `ctx_gen_inst_ratio_round_float` carry the frontier the gate's + metric is not. + +- **The gen-only CSV** — what `get_gen_only_perf` writes, for a gen + campaign that declares no ctx anchor. The gate's metric is + ``accept_rate / avg_iteration_time``, which has no context term in it; + the frontier is the only thing that needs the anchor. Reading the score + out of the frontier therefore invented a dependency the measurement + does not have — a gen campaign had to wait for somebody's ctx run + before it could score a decode change that ctx cannot affect. Both + readers are kept because they answer different questions, and which one + applies is declared in `task.yaml` rather than guessed from the + directory. + +What this module does **not** do is reconstruct the identity model the +previous CLI carried (`config_id` / `code_id` / per-measurement records). +Nothing here can tell an attempt from a repeat of its baseline, so the +prompts make that an obligation the evaluator discharges by reading the +case's materialized config — stated as such, rather than implied. +""" + +from __future__ import annotations + +import csv +from pathlib import Path +from typing import Any, Iterable, Mapping + +#: The console script `ibc-trtllm-harness` installs. Named for the prompt +#: sections and error messages; nothing here executes it. +IBC_BENCH = "ibc-bench" + +#: Where `process frontier` leaves the scored points. One per (mtp tag, +#: variant), so a run dir can hold several; the reader takes them all. +FRONTIER_CSV_GLOB = "*_frontier_*.csv" + +#: The columns this workflow reads. The rest of the row is carried +#: through untouched, because a number a reader can trace beats a number +#: that arrived alone. +CASE_KEY = ("name", "concurrency") +GEN_METRIC = "throughput_per_user" + +#: Where `get_gen_only_perf` leaves the anchor-free scores — one file per +#: run dir, written beside the cases it read. +GEN_ONLY_CSV = "gen_only_perf.csv" + +#: What that file calls :data:`GEN_METRIC`. The same quantity by the same +#: formula (`accept_rate / avg_iteration_time`); only the column name +#: differs, so the rename happens here and nowhere else. +GEN_ONLY_METRIC = "tps_per_user" + +#: How the harness exposes the gen-only extractor. It has no `ibc-bench` +#: subcommand as of v0.5.7 — `process frontier` is the only scored gen +#: path on the CLI — so the prompt runs the module directly. Named here so +#: the error messages and the prompt cannot drift apart. +GEN_ONLY_MODULE = "ibc_trtllm_harness.process_data.get_gen_only_perf" + + +class BenchCliError(RuntimeError): + """A sweep or a result set could not be read as the harness writes it.""" + + +def _int(value: Any) -> int | None: + try: + parsed = int(str(value).strip()) + except (TypeError, ValueError): + return None + return parsed + + +def _float(value: Any) -> float | None: + try: + parsed = float(str(value).strip()) + except (TypeError, ValueError): + return None + return parsed if parsed == parsed else None # NaN is not a measurement + + +def _bool(value: Any) -> bool: + """A pandas-written boolean cell. Anything unrecognised is ``False``. + + ``get_gen_only_perf`` writes Python bools through ``DataFrame.to_csv``, + so the cell is the literal ``True``/``False``. Only this column feeds + the ``dep``/``tep`` half of a case name, and a name that silently + changed shape would be worse than one that is visibly wrong, so the + accepted spellings are enumerated rather than inferred from + truthiness — ``"False"`` is a non-empty string. + """ + return str(value).strip().lower() in {"true", "1", "yes"} + + +# ------------------------------------------------------------------ the sweep + + +def _one(values: Iterable[Any]) -> Any: + """The single distinct value, or ``None`` if there is not exactly one. + + Disagreement is not resolved to a first or a maximum: two rows at two + input lengths are two workloads, and a campaign that quoted one of + them would describe half its own cases wrongly with nothing raising. + """ + seen = [value for value in values if value is not None] + unique = {str(value): value for value in seen} + return next(iter(unique.values())) if len(unique) == 1 else None + + +def workload(sweep: Mapping[str, Any]) -> dict[str, Any]: + """What the run will serve, as the sweep states it. + + The two sweep kinds spell this differently and the difference is not + cosmetic. A **gen** sweep is one deployment, so the model and the + corpus are top-level and ``isl``/``osl`` with them. A **ctx** sweep is + a list of prefill benchmarks: the model sits under ``model:`` and each + ``benchmarks`` entry carries its own lengths, because sweeping the + input length is the normal thing to do there. Read only the gen shape + and a ctx campaign resolves to a workload of ``None``s -- which does + not fail, it simply reconciles nothing, and `task.yaml` keeps the + defaults block's ``random_input_len: 1024`` as this campaign's stated + input length whatever the sweep measures. + + So the ctx lengths are taken from the ``benchmarks`` entries, and only + when they agree. A campaign is frozen to one operating point, so they + do; a sweep spanning two input lengths has no single workload to + state, and saying nothing is the honest answer there. + + ``isl`` is deliberately not read as the request length on the gen + side. The harness spends it on ``max_seq_len`` and the client's + ``input_length``; the requests come from ``dataset_file``, a corpus + with its own distribution. The checked-in 8k sweep pairs ``isl: 8192`` + with a ``...-8192-1024-200000-...`` corpus and both are right. On a + ctx sweep the two coincide -- a prefill-only run at ``osl: 1`` reads a + corpus generated for that length -- but they are still read from where + each sweep puts them, not assumed equal. + """ + model = sweep.get("model") + model = model if isinstance(model, Mapping) else {} + entries = [entry for entry in sweep.get("benchmarks") or [] if isinstance(entry, Mapping)] + return { + "model": sweep.get("model_id") or model.get("model_card"), + "model_path": sweep.get("model_path") or model.get("model_path"), + "dataset": sweep.get("dataset_file") or model.get("dataset_file"), + "precision": sweep.get("precision"), + "isl": sweep.get("isl") + if sweep.get("isl") is not None + else _one(e.get("isl") for e in entries), + "osl": sweep.get("osl") + if sweep.get("osl") is not None + else _one(e.get("osl") for e in entries), + "benchmark_client": sweep.get("benchmark_client"), + } + + +def gen_cases(sweep: Mapping[str, Any]) -> list[dict[str, Any]]: + """Every gen case a `gen_configs` sweep expands to. + + The row order is the harness': `[ctx_num, gen_num, tp_size, batch, + max_num_tokens, attention_dp, gpu_memory_fraction, mtp, eplb, + concurrency_list]`. The case name mirrors what the postprocessor's + `name` column carries — `{dep|tep}_{tp}_eplb{N}_mtp{M}` — so a point + read back out of the CSV addresses the row that produced it. + """ + cases: list[dict[str, Any]] = [] + for row in sweep.get("gen_configs") or []: + if isinstance(row, Mapping): # the dict form the harness also accepts + row = [ + row.get(k) + for k in ( + "ctx_num", + "gen_num", + "gen_tp_size", + "gen_batch_size", + "gen_max_num_tokens", + "gen_enable_attention_dp", + "gen_gpu_memory_fraction", + "gen_mtp_size", + "gen_eplb_num_slots", + "gen_concurrency_list", + ) + ] + if not isinstance(row, (list, tuple)) or len(row) < 10: + continue + ctx_num, gen_num, tp, batch, mnt, adp, gmf, mtp, eplb, concurrencies = row[:10] + shape = f"{'dep' if adp else 'tep'}_{tp}_eplb{eplb}_mtp{mtp}" + for token in str(concurrencies).split(","): + concurrency = _int(token) + if concurrency is None: + continue + cases.append( + { + "case": f"{shape}_conc{concurrency}", + "name": shape, + "stage": "gen", + "config": { + "ctx_num": _int(ctx_num), + "gen_num": _int(gen_num), + "tp_size": _int(tp), + "batch_size": _int(batch), + "max_num_tokens": _int(mnt), + "attention_dp": bool(adp), + "gpu_memory_fraction": _float(gmf), + "mtp_size": _int(mtp), + "eplb_num_slots": _int(eplb), + "concurrency": concurrency, + }, + } + ) + return cases + + +def ctx_cases(sweep: Mapping[str, Any]) -> list[dict[str, Any]]: + """Every ctx case a `benchmarks` block expands to. + + A ctx entry has no concurrency: `max_batch` is the in-flight request + count for a prefill-only run, which is what this workflow means by + concurrency everywhere else. + """ + cases: list[dict[str, Any]] = [] + for entry in sweep.get("benchmarks") or []: + if not isinstance(entry, Mapping): + continue + isl, osl = entry.get("isl"), entry.get("osl", 1) + for batch in entry.get("max_batch") or []: + for tp in entry.get("tp_size") or []: + for ratio in entry.get("ratio") or [None]: + for mtp in entry.get("mtp_range") or [0]: + cases.append( + { + "case": f"ctx-isl{isl}_osl{osl}_b{batch}_tp{tp}_mtp{mtp}", + "stage": "ctx", + "config": { + "isl": _int(isl), + "osl": _int(osl), + "max_batch": _int(batch), + "tp_size": _int(tp), + "ratio": ratio, + "mtp": _int(mtp), + }, + } + ) + return cases + + +def plan(sweep: Mapping[str, Any]) -> dict[str, Any]: + """The sweep as a workload plus the cases it expands to.""" + return {"workload": workload(sweep), "cases": gen_cases(sweep) + ctx_cases(sweep)} + + +def operating_point(config: Mapping[str, Any]) -> int | None: + """Total requests in flight at one case, or ``None`` if unstated. + + Not ``concurrency`` verbatim on a gen case: the sweep row's value is + **per generation server**, the client is driven at ``concurrency * + gen_num``, and the harness names the result directory after that + product. A ctx case has no concurrency at all — ``max_batch`` is the + in-flight count for a prefill-only run. + """ + listed = config.get("concurrency") + if not isinstance(listed, int) or isinstance(listed, bool): + batch = config.get("max_batch") + return batch if isinstance(batch, int) and not isinstance(batch, bool) else None + gen_num = config.get("gen_num", 1) + if not isinstance(gen_num, int) or isinstance(gen_num, bool) or gen_num < 1: + gen_num = 1 + return listed * gen_num + + +def operating_points(plan_data: Mapping[str, Any]) -> list[int]: + """The concurrency axis a `task.yaml` should carry, from a plan.""" + points = (operating_point(case.get("config") or {}) for case in plan_data.get("cases") or []) + return sorted({point for point in points if point is not None}) + + +# ---------------------------------------------------------------- the results + + +def frontier_points(run_dir: Path) -> list[dict[str, Any]]: + """Every scored point `process frontier` wrote under ``run_dir``. + + Keyed by ``(name, concurrency)`` — the shape the row was measured at + and the point on its curve — because that pair is what survives a code + change, and therefore what an attempt and its baseline have in common. + """ + found: dict[tuple[str, int], dict[str, Any]] = {} + csvs = sorted(Path(run_dir).rglob(FRONTIER_CSV_GLOB)) + if not csvs: + raise BenchCliError( + f"no {FRONTIER_CSV_GLOB} under {run_dir}: `{IBC_BENCH} process frontier` " + f"has not run, or ran without producing a curve. `{IBC_BENCH} jobs check " + f"-f {run_dir}/job_status.csv --summary` says whether the cases behind it " + f"measured." + ) + for path in csvs: + try: + rows: Iterable[dict[str, str]] = list(csv.DictReader(path.open(encoding="utf-8"))) + except OSError as exc: # pragma: no cover - message path + raise BenchCliError(f"could not read {path}: {exc}") from exc + for row in rows: + name = (row.get("name") or "").strip() + concurrency = _int(row.get("concurrency")) + value = _float(row.get(GEN_METRIC)) + if not name or concurrency is None or value is None: + continue + found[(name, concurrency)] = { + "case": f"{name}_conc{concurrency}", + "name": name, + "concurrency": concurrency, + "metrics": { + GEN_METRIC: value, + "output_tput_per_gpu": _float(row.get("output_tput_per_gpu")), + "ctx_gen_inst_ratio": _float(row.get("ctx_gen_inst_ratio_round_float")), + "ctx_request_rate": _float(row.get("ctx_request_rate")), + }, + "gpus": { + "ctx": _float(row.get("ctx_gpus_round")), + "gen": _float(row.get("gen_num_round")), + "total": _float(row.get("total_gpus_round")), + }, + "source_csv": str(path), + } + if not found: + raise BenchCliError( + f"the frontier CSVs under {run_dir} carry no row with a name, a " + f"concurrency and a {GEN_METRIC}. Check that the postprocessor scored " + f"the cases rather than only listing them." + ) + return [found[key] for key in sorted(found)] + + +def gen_only_points(run_dir: Path) -> list[dict[str, Any]]: + """Every scored point `get_gen_only_perf` wrote under ``run_dir``. + + The same shape :func:`frontier_points` returns, keyed the same way, so + the rest of the workflow cannot tell which reader produced a point — + except by what is missing, which is the whole e2e half: + ``output_tput_per_gpu``, the ctx:gen ratio and the GPU counts all + divide by a context rate this campaign never measured. + + They are **absent rather than zero**, and the caller is expected to + say so in the result it writes. A frontier column filled with a + plausible default is the failure this track is built to refuse: the + curve still plots, the deployment it describes never existed. + + ``output_tps_per_gen_gpu`` is carried but deliberately not renamed to + ``output_tput_per_gpu``. They differ in the denominator — this one + divides by the generation GPUs alone, the frontier's by the whole + rate-matched deployment — so the two are never the same number, and + the smaller denominator makes this the flattering one. + """ + path = Path(run_dir) / GEN_ONLY_CSV + if not path.is_file(): + raise BenchCliError( + f"no {GEN_ONLY_CSV} in {run_dir}: this campaign declares no ctx anchor, " + f"so its score comes from the anchor-free extractor. Run " + f"`python -m {GEN_ONLY_MODULE} -i {run_dir}` after the sweep, then " + f"collect again. `{IBC_BENCH} jobs check -f {run_dir}/job_status.csv " + f"--summary` says whether the cases behind it measured." + ) + try: + rows: Iterable[dict[str, str]] = list(csv.DictReader(path.open(encoding="utf-8"))) + except OSError as exc: # pragma: no cover - message path + raise BenchCliError(f"could not read {path}: {exc}") from exc + + found: dict[tuple[str, int], dict[str, Any]] = {} + for row in rows: + concurrency = _int(row.get("concurrency")) + value = _float(row.get(GEN_ONLY_METRIC)) + tp, mtp, eplb = _int(row.get("tp")), _int(row.get("mtp")), _int(row.get("eplb")) + if concurrency is None or value is None or tp is None: + continue + # Rebuilt from the columns rather than parsed out of `config`, so + # it matches `gen_cases()`'s shape and the frontier CSV's `name` + # exactly. A point read through either reader then addresses the + # same planned row. + name = f"{'dep' if _bool(row.get('adp')) else 'tep'}_{tp}_eplb{eplb}_mtp{mtp}" + found[(name, concurrency)] = { + "case": f"{name}_conc{concurrency}", + "name": name, + "concurrency": concurrency, + "metrics": { + GEN_METRIC: value, + "output_tput": _float(row.get("output_tput")), + "output_tps_per_gen_gpu": _float(row.get("output_tps_per_gen_gpu")), + "avg_itertime_ms": _float(row.get("avg_itertime_ms")), + "num_iters": _int(row.get("num_iters")), + }, + "source_csv": str(path), + "harness_config": (row.get("config") or "").strip() or None, + } + if not found: + raise BenchCliError( + f"{path} carries no row with a concurrency, a tp and a {GEN_ONLY_METRIC}. " + f"The extractor drops a case whose iteration log never reached steady " + f"state, so an empty file means the cases ran but did not settle." + ) + return [found[key] for key in sorted(found)] diff --git a/agent-flow/agent_flow/workflows/perf_optimize/cli.py b/agent-flow/agent_flow/workflows/perf_optimize/cli.py index 5fd791e7578f..201cb5b0c969 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/cli.py +++ b/agent-flow/agent_flow/workflows/perf_optimize/cli.py @@ -4,10 +4,13 @@ import sys from pathlib import Path +import yaml + from agent_flow.workflows.perf_analyze.sol_methodology import resolve_sol_methodology from .disagg import has_disagg from .prompts import build_perf_optimize_prompts +from .sol_track import track_name from .state import STATE_FILENAME from .task_schema import ( TaskSchemaError, @@ -104,11 +107,97 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: "`optimize.item_execution` mode). " "Ignored on resume — the checkpointed budget wins.", ) + _add_dry_run(parser) return parser.parse_args(argv) +def _run_disagg_sol(args) -> None: + """The staged two-track path: fix the operating point, then optimize at it. + + Dispatched on the shape of `task.yaml` rather than on a flag, so one file + describes the whole campaign and there is one place to look for what it + will do. The single-track path below is untouched: a spec without a + `disagg_sol` block reaches it exactly as before. + """ + import json + + import yaml + + from . import spawn + from .disagg_sol import DisaggSolError, supervise + + raw = yaml.safe_load(Path(args.task).read_text(encoding="utf-8")) or {} + block = raw.get("disagg_sol") or {} + root = Path(args.workspace) + try: + record = supervise( + raw, + sweeps={k: Path(v) for k, v in (block.get("sweeps") or {}).items()}, + repos={k: Path(v) for k, v in (block.get("repos") or {}).items()}, + workspace_root=root, + label=block.get("label") or root.name, + incumbent=block.get("incumbent"), + dry_run=bool(getattr(args, "dry_run", False)), + # Only when the spec says where the skill lives. Without it the + # design stays an input and an unestablished one is refused -- + # which is the safer default for a step that costs an order of + # magnitude more than the campaigns it enables. + designer=( + ( + lambda instruction: spawn.design( + instruction, + cwd=Path(block["config_repo"]), + log=root / "design.log", + ) + ) + if block.get("config_repo") + else None + ), + ) + except DisaggSolError as exc: + print(f"error: {exc}", file=sys.stderr) + sys.exit(2) + print(json.dumps(record, indent=2, default=str)) + + +def _add_dry_run(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--dry-run", + action="store_true", + help=( + "staged disagg only: select the operating point and write every " + "campaign's task.yaml, then stop without starting them" + ), + ) + + def main(argv: list[str] | None = None) -> None: args = _parse_args(argv) + # Before schema validation: a staged spec is a different shape, and the + # single-track schema would reject it for the fields it deliberately lacks. + try: + _raw = yaml.safe_load(Path(args.task).read_text(encoding="utf-8")) or {} + except (OSError, yaml.YAMLError) as exc: + print(f"error: could not read {args.task}: {exc}", file=sys.stderr) + sys.exit(2) + if isinstance(_raw, dict) and "disagg_sol" in _raw: + _run_disagg_sol(args) + return + # Refused rather than ignored. The flag is on the shared parser because + # argparse cannot know which path the spec takes until the file is read, + # and only the staged path implements it -- so on a single-track spec it + # used to be accepted and then silently dropped, which turns "show me + # what this would do" into a full multi-hour campaign on real hardware. + # The one failure mode a dry run must not have. + if getattr(args, "dry_run", False): + print( + "error: --dry-run is implemented for the staged disagg path only " + f"(a spec with a 'disagg_sol' block); {args.task} is a single-track " + "campaign, and running it without the flag would start it for real. " + "Re-run without --dry-run when you mean to.", + file=sys.stderr, + ) + sys.exit(2) try: task_data = load_and_validate_task_yaml( args.task, @@ -139,6 +228,7 @@ def main(argv: list[str] | None = None) -> None: kernel_coverage=kernel_coverage(task_data), sol_methodology=methodology.name, include_disagg=has_disagg(task_data), + sol_track=track_name(task_data), ) with PerfOptimizeWorkflow( workspace=args.workspace, diff --git a/agent-flow/agent_flow/workflows/perf_optimize/disagg_sol.py b/agent-flow/agent_flow/workflows/perf_optimize/disagg_sol.py new file mode 100644 index 000000000000..c5d53c475742 --- /dev/null +++ b/agent-flow/agent_flow/workflows/perf_optimize/disagg_sol.py @@ -0,0 +1,1262 @@ +"""Staged disagg optimization: fix the operating point, then optimize at it. + +A SOL-track campaign (see :mod:`.sol_track`) optimizes at an operating point +it is *handed*. That is the right division of labour, and it has one failure +mode: nothing checks where the point came from. The two campaigns this module +was written for inherited row 8 of a checked-in sweep — a row of unknown +provenance — and spent seven hours and thirteen submissions improving a point +nobody had measured. A gain at the wrong point is worth less than moving to +the right one, and neither campaign could have told the difference. + +So this module is the layer above a campaign: it **establishes the point**, +and only then starts the campaigns that optimize at it. + +**It is not a merge layer.** The two tracks stay what they are — two +independent campaigns, two workspaces, two checkouts, two branches, each +scored on its own metric, neither able to see the other. Nothing here +computes an end-to-end number, and that is a scope decision rather than an +oversight: see :data:`NO_JOIN`. + +The staging, in dependency order: + +1. **Model facts** — weight bytes, expert count, layers, read from the + checkpoint. Free, and the input to every shape decision below. +2. **Shape space** — which (parallelism, world size) combinations are even + legal for this model on this GPU. Free. +3. **Max-batch probe** — for each shape, the largest batch that fits. This + cannot be computed, only measured: it is where the memory wall is. +4. **Concurrency sweep** — each shape at its *measured* batch, across its + concurrency ladder, scored **anchor-free** (:func:`.bench_cli.gen_only_points`). +5. **Selection** — pick the point the campaigns will freeze on. + +Steps 1–4 are `create-sweep`'s Phases 0–3, minus the frontier build. This +module does not reimplement them; it reads what they wrote. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any, Callable, Iterable, Mapping, Sequence + +import yaml + +from agent_flow.workflows.perf_optimize.bench_cli import ( + GEN_ONLY_CSV, + BenchCliError, + ctx_cases, + gen_cases, + gen_only_points, + operating_point, +) +from agent_flow.workflows.perf_optimize.sweep_design import designer_instruction, load_yaml + +DISAGG_SOL_FIELD = "disagg_sol" +TRACKS_KEY = "tracks" +DESIGN_KEY = "design" +DESIGN_DIR_KEY = "design_dir" +PREFER_KEY = "prefer" +DESIGN_SWEEP_KEY = "design_sweep" + +#: What `create-sweep` calls its state file. Read, never written: this module +#: is a consumer of that skill's output, not a reimplementation of it. +DESIGN_STATE = "state.json" + +#: The two halves. Ordered ctx-first because that is the order the harness +#: itself uses when both are present, not because one gates the other — under +#: this module's scope neither does. +CTX_TRACK = "ctx" +GEN_TRACK = "gen" +TRACKS: tuple[str, ...] = (CTX_TRACK, GEN_TRACK) + +#: What each half inherits from the run's spec. +#: +#: The rule is: **how and where to execute passes through; what to measure does +#: not.** A half's operating point comes from the design and travels in its +#: derived sweep, so `benchmark` is deliberately absent — inheriting one would +#: put a second, unmeasured description of the workload beside the selected row, +#: which is the failure `sol_track` exists to remove. +#: +#: `slurm-environment` is the half of that rule easiest to forget, and it was: +#: it carries `cluster_ssh` and `remote_run_root`, which ARE off-cluster +#: execution. Dropping it did not fail — both halves simply ran as though the +#: flow were on the cluster, with the spec's remote settings absent from the +#: only two files that could act on them. A setting that vanishes silently is +#: worse than one that is refused. +INHERITED: tuple[str, ...] = ( + "checkpoint_path", + "optimize", + "profile", + "slurm-environment", +) + +#: Why the selection frontier is not the deployment frontier. +#: +#: Stated in the selection result, and required to survive into any report +#: quoting it. `output_tps_per_gen_gpu` divides throughput by the GENERATION +#: GPUs alone. The deployment number divides by the whole rate-matched pair:: +#: +#: output_tput_per_gpu = output_throughput / (ctx_gpus * ctx_per_gen + gen_gpus) +#: +#: so a large-expert-parallel shape can lead this ranking while trailing on +#: deployment cost, because the context GPUs it drags behind it are not in the +#: denominator. Recording the reason rather than the omission, because "no +#: end-to-end view was taken" and "the end-to-end view was flat" are different +#: findings that a missing field cannot distinguish. +NO_JOIN = ( + "selection ranked on output_tps_per_gen_gpu, which divides by the GENERATION " + "GPUs alone. This is NOT the deployment frontier: output_tput_per_gpu divides " + "by ctx_gpus * ctx_per_gen + gen_gpus, and no context measurement was rate-" + "matched against these points. A shape that drags more context GPUs behind it " + "therefore ranks better here than it would deploy. Absent, not flat -- do not " + "quote any number selected this way as an end-to-end result." +) + +#: What a one-point campaign cannot see, recorded with the point it froze. +#: +#: The campaign gates on a single operating point, so a change is judged only +#: where it was measured. Nothing here is wrong with that — it is the cheapest +#: honest gate — but it does mean a change that helps at the frozen point and +#: hurts elsewhere on the curve is indistinguishable from one that helps +#: everywhere. That is not hypothetical: the campaign this module supersedes +#: measured `opt-006` at **+1.52 % on one point and -1.85 % on another**, and +#: was only saved from accepting it by having frozen both. With one point, the +#: same attempt is an accept. +#: +#: Stated rather than guarded, because which points a deployment cares about +#: is the same exogenous question :data:`PREFERENCES` answers, and a guard +#: here would be this module inventing an answer to it. +ONE_POINT = ( + "one operating point was frozen, so every gain and regression this campaign " + "reports is measured only there. A change that helps at this point and hurts " + "elsewhere on the curve is indistinguishable here from one that helps " + "everywhere -- the rest of the curve was not re-measured after any attempt. " + "Unobserved, not unchanged." +) + +#: What the selection is allowed to optimize for. Exogenous on purpose: a +#: frontier states the trade-off and cannot state which end of it the +#: deployment is bought for, so the campaign's owner says. +PREFER_INTERACTIVE = "interactive" # max tokens/s/user +PREFER_THROUGHPUT = "throughput" # max tokens/s/gen-GPU +PREFERENCES: tuple[str, ...] = (PREFER_INTERACTIVE, PREFER_THROUGHPUT) + +#: The axes, in the names `gen_only_perf.csv` carries them through +#: :func:`.bench_cli.gen_only_points`. +X_AXIS = "throughput_per_user" +Y_AXIS = "output_tps_per_gen_gpu" + + +class DisaggSolError(ValueError): + """The staged-disagg block, or the design it points at, is unusable.""" + + +# ------------------------------------------------------------------ the block + + +def has_disagg_sol(data: Mapping[str, Any]) -> bool: + """Whether this spec is a staged two-track campaign.""" + return DISAGG_SOL_FIELD in data + + +def _block(data: Mapping[str, Any]) -> Mapping[str, Any]: + block = data.get(DISAGG_SOL_FIELD) + if not isinstance(block, Mapping): + raise DisaggSolError(f"'{DISAGG_SOL_FIELD}' must be a mapping") + return block + + +def tracks(data: Mapping[str, Any]) -> list[str]: + """The halves this spec will optimize, validated. + + Defaults to both. A single-track list is legal and means "run one + campaign, but still fix the point first" — the staging is the point of + this module, not the plurality. + """ + listed = _block(data).get(TRACKS_KEY, list(TRACKS)) + if isinstance(listed, str): + listed = [listed] + if not isinstance(listed, Sequence) or not listed: + raise DisaggSolError(f"'{DISAGG_SOL_FIELD}.{TRACKS_KEY}' must be a non-empty list") + unknown = [t for t in listed if t not in TRACKS] + if unknown: + raise DisaggSolError( + f"'{DISAGG_SOL_FIELD}.{TRACKS_KEY}' names {unknown}, expected any of {list(TRACKS)}" + ) + seen: list[str] = [] + for track in listed: + if track not in seen: + seen.append(track) + return seen + + +def design_dir(data: Mapping[str, Any]) -> Path: + """Where `create-sweep` left the design this campaign freezes on.""" + value = _block(data).get(DESIGN_KEY) + value = value if isinstance(value, Mapping) else {} + stated = value.get(DESIGN_DIR_KEY) + if not isinstance(stated, str) or not stated.strip(): + raise DisaggSolError( + f"'{DISAGG_SOL_FIELD}.{DESIGN_KEY}.{DESIGN_DIR_KEY}' is required: the " + f"`sweep_design/` directory `create-sweep` wrote, which is where the " + f"measured max batch and the measured concurrency ladders live. A " + f"campaign that cannot name one is inheriting its operating point." + ) + return Path(stated.strip()) + + +def preference(data: Mapping[str, Any]) -> str: + """Which end of the frontier this deployment is bought for. + + No default. A frontier is a trade-off between tokens/s/user and + tokens/s/GPU; which end matters is a property of the service being run, + not of the measurement, and a campaign that guesses it optimizes toward + an operating point nobody asked for. The two campaigns this module + replaces froze concurrency 1 *and* 32 — opposite ends of one curve — + without ever saying which one the deployment was for. + """ + value = _block(data).get(DESIGN_KEY) + value = value if isinstance(value, Mapping) else {} + stated = value.get(PREFER_KEY) + if stated not in PREFERENCES: + raise DisaggSolError( + f"'{DISAGG_SOL_FIELD}.{DESIGN_KEY}.{PREFER_KEY}' must be one of " + f"{list(PREFERENCES)}. The measured curve states the trade-off between " + f"{X_AXIS} and {Y_AXIS}; it cannot state which end this deployment is " + f"bought for, so nothing here will infer it." + ) + return str(stated) + + +# ---------------------------------------------------------------- the design + + +def design_state(directory: Path) -> dict[str, Any]: + """`create-sweep`'s own record of how far it got, read not written.""" + path = Path(directory) / DESIGN_STATE + try: + data = json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise DisaggSolError( + f"no {DESIGN_STATE} under {directory}: `create-sweep` records its phase " + f"and artefact paths there after every phase, so its absence means the " + f"design was never run — or was run somewhere else." + ) from exc + except (OSError, json.JSONDecodeError) as exc: + raise DisaggSolError(f"could not read {path}: {exc}") from exc + if not isinstance(data, Mapping): + raise DisaggSolError(f"{path} must be a JSON object, got {type(data).__name__}") + return dict(data) + + +def sweep_points(directory: Path) -> list[dict[str, Any]]: + """Every measured point the design left, across every shape. + + One `gen_only_perf.csv` per run directory, and one run directory per + shape — so the union is the measured space, at each shape's own measured + batch rather than at a batch anybody typed in. + + Anchor-free by construction: this reads what `get_gen_only_perf` wrote, + never a frontier. See :data:`NO_JOIN` for what that costs. + """ + root = Path(directory) + csvs = sorted(root.rglob(GEN_ONLY_CSV)) + if not csvs: + raise DisaggSolError( + f"no {GEN_ONLY_CSV} under {root}: the design's concurrency sweep has " + f"not been scored. Run `python -m " + f"ibc_trtllm_harness.process_data.get_gen_only_perf -i ` for " + f"each shape's run directory, then select again." + ) + points: list[dict[str, Any]] = [] + for path in csvs: + try: + found = gen_only_points(path.parent) + except BenchCliError as exc: # pragma: no cover - message path + raise DisaggSolError(str(exc)) from exc + for point in found: + metrics = dict(point.get("metrics") or {}) + points.append( + { + "shape": point.get("name"), + "concurrency": point.get("concurrency"), + X_AXIS: metrics.get(X_AXIS), + Y_AXIS: metrics.get(Y_AXIS), + "source_csv": point.get("source_csv"), + } + ) + usable = [p for p in points if _usable(p)] + if not usable: + raise DisaggSolError( + f"the {len(csvs)} {GEN_ONLY_CSV} under {root} carry no point with both " + f"{X_AXIS} and {Y_AXIS}. The extractor drops a case whose iteration log " + f"never reached steady state, so this means the cases ran but did not settle." + ) + return usable + + +def _usable(point: Mapping[str, Any]) -> bool: + if not isinstance(point.get("concurrency"), int): + return False + return all( + isinstance(point.get(axis), (int, float)) and not isinstance(point.get(axis), bool) + for axis in (X_AXIS, Y_AXIS) + ) + + +def pareto_front(points: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]: + """The non-dominated points, both axes higher-is-better. + + A point is dropped only when another is at least as good on *both* axes + and strictly better on one. Ties are kept: two shapes that measure the + same pair are two real options, and picking between them is a decision + this function is not entitled to make. + """ + listed = [dict(p) for p in points] + front: list[dict[str, Any]] = [] + for candidate in listed: + dominated = any( + other is not candidate + and other[X_AXIS] >= candidate[X_AXIS] + and other[Y_AXIS] >= candidate[Y_AXIS] + and (other[X_AXIS] > candidate[X_AXIS] or other[Y_AXIS] > candidate[Y_AXIS]) + for other in listed + ) + if not dominated: + front.append(candidate) + return sorted(front, key=lambda p: (-p[X_AXIS], -p[Y_AXIS])) + + +def select_point( + points: Iterable[Mapping[str, Any]], *, prefer: str, incumbent: Mapping[str, Any] | None = None +) -> dict[str, Any]: + """The point the campaigns will freeze on, and why. + + ``incumbent`` is the point a campaign would otherwise have inherited — + the checked-in row. It is not used to choose; it is used to *report*, + because "the selection agrees with what we were already running" and + "the selection moved us" are the two answers this whole staging exists + to distinguish, and only one of them makes the previous campaigns' + measurements still meaningful. + """ + if prefer not in PREFERENCES: + raise DisaggSolError(f"unknown preference {prefer!r}, expected one of {list(PREFERENCES)}") + front = pareto_front(points) + if not front: + raise DisaggSolError("no measured point survived the Pareto filter") + axis = X_AXIS if prefer == PREFER_INTERACTIVE else Y_AXIS + chosen = max(front, key=lambda p: p[axis]) + result = { + "shape": chosen["shape"], + "concurrency": chosen["concurrency"], + X_AXIS: chosen[X_AXIS], + Y_AXIS: chosen[Y_AXIS], + "prefer": prefer, + "ranked_on": axis, + "source_csv": chosen.get("source_csv"), + "pareto_size": len(front), + "measured_points": len(list(points)), + # Not a caveat in prose somewhere: the reason travels with the number. + "e2e_view_absent": NO_JOIN, + "off_point_effects_unobserved": ONE_POINT, + } + if incumbent is not None: + result["incumbent"] = dict(incumbent) + result["moved"] = not _same_point(chosen, incumbent) + result["incumbent_on_pareto"] = any(_same_point(p, incumbent) for p in front) + return result + + +def _same_point(a: Mapping[str, Any], b: Mapping[str, Any]) -> bool: + return (a.get("shape"), a.get("concurrency")) == (b.get("shape"), b.get("concurrency")) + + +# ------------------------------------------------------------- the ctx half + + +#: How a ctx candidate is ranked, and the one place this module's two halves +#: are asymmetric on purpose. +#: +#: The gen side is ranked on a curve, because a decode point trades +#: tokens/s/user against tokens/s/GPU and neither dominates. The ctx side has +#: no such trade: a prefill worker either serves more requests per GPU or it +#: does not, so the objective is scalar and there is nothing to prefer between. +#: The harness' own anchor picker uses exactly this, for exactly this reason — +#: `select_agentx_ctx_anchor.py` ranks on ``req_s_per_gpu`` because the ctx +#: term of the deployment denominator is ``ctx_gpus * ctx_per_gen``, which a +#: higher request rate per GPU shrinks on both factors at once. +#: +#: Note this needs no rate match: it is entirely inside the ctx measurement. +#: Choosing the ctx point is therefore possible without an end-to-end view, +#: which is why this half can be *measured* rather than computed even under +#: this module's no-join scope. +CTX_METRIC = "avg_request_throughput_req_s" +CTX_RANK = "req_s_per_ctx_gpu" + +#: Where a validated ctx measurement keeps its number — the field the harness +#: itself requires before it will call a ctx case successful. Duplicated from +#: :mod:`.sol_track` rather than imported, because the two read it for +#: different purposes and a shared constant would make one module's change +#: silently the other's. +CTX_RESULT_PATH = ("performance", "request_throughput_req_s") + + +def _ctx_case_facts(case_dir: str) -> dict[str, Any] | None: + """``ctx_{isl}_{osl}_ratio{r}_{batch}_{mnt}_{dep|tep}{tp}_MTP{n}_test{k}``. + + Read positionally, like :func:`.sol_track._ctx_concurrency`: a name this + workflow cannot parse must stop the selection rather than let it pick + whichever number happened to match. The GPU count is ``tp`` — a ctx-only + run is a single aggregate worker, so its world size *is* its GPU count. + """ + parts = case_dir.split("_") + if len(parts) < 7 or parts[0] != "ctx": + return None + if not (parts[1].isdigit() and parts[4].isdigit()): + return None + world = parts[6] + kind = world[:3] + if kind not in ("dep", "tep") or not world[3:].isdigit(): + return None + return { + "case": case_dir, + "isl": int(parts[1]), + "max_batch": int(parts[4]), + "adp": kind == "dep", + "ctx_gpus": int(world[3:]), + } + + +def ctx_points(directory: Path) -> list[dict[str, Any]]: + """Every scored ctx candidate the design left. + + One ``run_*.json`` per case, whose + ``performance.request_throughput_req_s`` is the field the harness + validated the case on — so a case that appears here measured, and one + that failed does not appear at all. + """ + found: list[dict[str, Any]] = [] + for result in sorted(Path(directory).rglob("run_*.json")): + if result.name.endswith("_timing.json"): + continue + facts = _ctx_case_facts(result.parent.name) + if facts is None: + continue + value = _read_ctx_json(result) + if value is None or facts["ctx_gpus"] <= 0: + continue + found.append( + { + **facts, + CTX_METRIC: value, + CTX_RANK: value / facts["ctx_gpus"], + "source_run_json": str(result), + } + ) + if not found: + raise DisaggSolError( + f"no scored ctx case under {directory}. A ctx candidate is a " + f"`run_*.json` whose {'.'.join(CTX_RESULT_PATH)} the harness " + f"validated; none was found, so there is nothing to choose between " + f"and the ctx half would fall back to a computed point." + ) + return found + + +def _read_ctx_json(path: Path) -> float | None: + try: + payload: Any = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return None + for key in CTX_RESULT_PATH: + payload = payload.get(key) if isinstance(payload, Mapping) else None + if isinstance(payload, (int, float)) and not isinstance(payload, bool): + return float(payload) + return None + + +#: What makes two ctx measurements the same candidate. +#: +#: Deliberately not the case name. A ctx sweep repeats each case (``rounds``) +#: and expands over ``mtp_range``, so one configuration arrives as several +#: differently-named directories — twelve of them, on the run this was +#: checked against. MTP is not part of a ctx operating point at all: the +#: generation sweep's ``ctx_config`` block has no mtp field, and the measured +#: spread across mtp variants (2.4 %) sat inside the spread across repeats of +#: one variant (4.0 %). +#: +#: Ranking those twelve as twelve candidates does two wrong things: it picks +#: the luckiest repeat rather than the best configuration, and it reports +#: having *moved* when it only moved between repeats of the incumbent. +CTX_CONFIG_KEYS = ("ctx_gpus", "max_batch", "adp") + + +def _config_key(point: Mapping[str, Any]) -> tuple: + facts = point + if not all(key in point for key in CTX_CONFIG_KEYS): + parsed = _ctx_case_facts(str(point.get("case") or "")) + if parsed is None: + return ("?", point.get("case")) + facts = parsed + return tuple(facts[key] for key in CTX_CONFIG_KEYS) + + +def select_ctx_point( + points: Iterable[Mapping[str, Any]], *, incumbent: Mapping[str, Any] | None = None +) -> dict[str, Any]: + """The most GPU-efficient measured prefill *configuration*. + + Scalar, so there is no ``prefer`` here and no Pareto front: unlike the + decode curve, a ctx candidate that serves more requests per GPU is better + in every way a deployment cares about. Ties break toward the smaller GPU + count, because two configurations at one efficiency are not equal — the + smaller one leaves the rest of the node to the generation side. + + Repeats of one configuration are averaged rather than competed. Taking + the maximum over repeats would rank configurations by which one drew the + kindest sample, and the measured spread here — 4.0 % across twelve + repeats of a single configuration — is wider than the difference between + configurations that this selection is meant to resolve. + """ + listed = [dict(p) for p in points] + if not listed: + raise DisaggSolError("no measured ctx candidate to choose between") + + grouped: dict[tuple, list[dict[str, Any]]] = {} + for point in listed: + grouped.setdefault(_config_key(point), []).append(point) + + candidates = [] + for key, repeats in grouped.items(): + ranks = [r[CTX_RANK] for r in repeats] + values = [r[CTX_METRIC] for r in repeats] + first = repeats[0] + candidates.append( + { + "config": key, + "ctx_gpus": first["ctx_gpus"], + "max_batch": first["max_batch"], + "adp": first["adp"], + CTX_METRIC: sum(values) / len(values), + CTX_RANK: sum(ranks) / len(ranks), + "repeats": len(repeats), + "spread_pct": (max(ranks) - min(ranks)) / (sum(ranks) / len(ranks)) * 100.0, + "cases": [r.get("case") for r in repeats], + } + ) + + chosen = max(candidates, key=lambda c: (c[CTX_RANK], -c["ctx_gpus"])) + result = { + "ctx_gpus": chosen["ctx_gpus"], + "max_batch": chosen["max_batch"], + "adp": chosen["adp"], + CTX_METRIC: chosen[CTX_METRIC], + CTX_RANK: chosen[CTX_RANK], + "ranked_on": f"mean {CTX_RANK} over {chosen['repeats']} repeat(s)", + "repeats": chosen["repeats"], + "spread_pct": chosen["spread_pct"], + "cases": chosen["cases"], + "candidates": len(candidates), + "measurements": len(listed), + } + if incumbent is not None: + result["incumbent"] = dict(incumbent) + # By configuration, never by case name: a different repeat of the + # incumbent is not a move, and reporting it as one gives the wrong + # answer to the question this whole staging exists to ask. + result["moved"] = chosen["config"] != _config_key(incumbent) + return result + + +# --------------------------------------------------------------- the campaigns + + +#: A design is reusable once its concurrency sweep has been scored. Checked +#: on the artefacts rather than on ``state.json``'s ``phase`` string: the +#: phase is what the design *says* it reached, the CSVs are what it left, and +#: a resumed or interrupted design can have the first without the second. +def established(directory: Path, track: str = GEN_TRACK) -> bool: + """Whether this design can be selected from for ``track``. + + Per track, because the two halves are established by different artefacts + and — under this module's no-join scope — neither waits on the other. A + design whose prefill candidates have been measured can start the ctx + campaign while its generation sweep is still running; requiring both + would idle a half that is ready on a half that is not, for a dependency + the scope does not have. + + The reason this is a question at all: fixing the operating point costs + roughly an order of magnitude more than the campaigns it enables — six + shapes probed and swept, against one baseline each. Paying that per + campaign would be absurd; paying it once per (model, cluster, workload) + and reusing it is the only shape in which the staging is affordable. + """ + root = Path(directory) + if track == CTX_TRACK: + return any( + path.name.startswith("run_") and not path.name.endswith("_timing.json") + for path in root.rglob("run_*.json") + ) + return bool(sorted(root.rglob(GEN_ONLY_CSV))) + + +#: Why a design may be somewhere other than where it was asked for, and what +#: is done about it. +#: +#: The design agent is given a directory and told to resume from its +#: ``state.json``. When the named directory does not exist and a sibling one +#: does — carrying a state file whose recorded scope matches the instruction — +#: resuming that sibling is the right engineering call: it saved roughly +#: twenty node-hours of re-probing on the run this was written after. Refusing +#: it would forbid a correct decision. +#: +#: But the supervisor checks the directory it *asked* for, so the correct +#: decision surfaced as "the design was never established" and the run +#: stopped. Following the agent silently would be worse: every later artefact +#: would cite a design directory the spec never named. +#: +#: So it is followed and recorded. The redirection travels in the run record +#: and in every campaign's provenance, for the same reason the missing +#: end-to-end view does: a reader must not have to already know. +DESIGN_REDIRECTED = ( + "the design agent was given {requested} and established the design in {actual} " + "instead. That is legitimate -- a design is resumed from its state.json, and " + "resuming an existing one avoids re-measuring what it already measured -- but " + "it means every artefact below was selected from a directory this spec did not " + "name. Point the spec at {actual} to make the next run's request and result " + "agree." +) + + +def resolve_design_dir(requested: Path, wanted: Sequence[str]) -> tuple[Path, str | None]: + """Where the design actually is, and a note if that is not where it was asked. + + Looked for among the requested directory's siblings, which is where an + agent resuming an existing design would find one: the model directory + holds ``sweep_design`` and whatever variants a spec has named. The search + is deliberately narrow -- one level, and only directories that both carry + a ``state.json`` and are established for every wanted track -- because a + supervisor that hunts the filesystem for something that looks like a + design will eventually find one that is not. + + Ambiguity is refused rather than resolved. Two established designs beside + each other is a question about which measurement the campaigns should + freeze on, and this is not the layer that answers it. + """ + requested = Path(requested) + if all(established(requested, track) for track in wanted): + return requested, None + + parent = requested.parent + if not parent.is_dir(): + return requested, None + found = [ + candidate + for candidate in sorted(parent.iterdir()) + if candidate.is_dir() + and candidate != requested + and (candidate / DESIGN_STATE).is_file() + and all(established(candidate, track) for track in wanted) + ] + if not found: + return requested, None + if len(found) > 1: + raise DisaggSolError( + f"{requested} is not established, and {len(found)} sibling directories " + f"are: {[str(f) for f in found]}. Which measurement the campaigns " + f"freeze on is not something this layer may pick -- name one in " + f"'{DISAGG_SOL_FIELD}.{DESIGN_KEY}.{DESIGN_DIR_KEY}'." + ) + return found[0], DESIGN_REDIRECTED.format(requested=requested, actual=found[0]) + + +def rebase_design_sweep(stated: Path, requested: Path, actual: Path) -> Path: + """Carry a redirected design's sweep along with its directory. + + ``design_dir`` and ``design_sweep`` are two fields naming one thing: the + sweep is a file the design writes inside its own directory. When the + design turns out to be somewhere other than where it was asked for, a + sweep path still rooted at the requested directory names a file that was + never written -- and the failure arrives late, after the points have been + read and the operating point chosen, which is the most expensive moment + to discover a path problem. + + Only a sweep that sat under the requested directory is moved, and only + when the file is not where it was stated. A path pointing somewhere else + entirely was a deliberate choice and is left alone. + """ + stated, requested, actual = Path(stated), Path(requested), Path(actual) + if stated.is_file() or requested == actual: + return stated + try: + return actual / stated.relative_to(requested) + except ValueError: + return stated + + +def campaign_workspace(root: Path, track: str, label: str) -> Path: + """Each half gets its own, and the name says which half it is. + + Not cosmetic: the workspace is what a campaign's branch is named after, + and the branch is what claims the checkout. Two campaigns sharing either + would have each reset the other's worktree mid-flight. + """ + return Path(root) / f"ws-{label}-{track}" + + +def campaign_spec( + base: Mapping[str, Any], + *, + track: str, + sweep: Path, + repo: Path, + point: Mapping[str, Any], + design: Path, +) -> dict[str, Any]: + """One half's `task.yaml`, as a single-track SOL campaign. + + The output is an ordinary :mod:`.sol_track` spec — the same shape the two + campaigns this module supersedes were written by hand. That is the point: + nothing below this layer changes, so everything already true of a + single-track campaign stays true, and this layer is additive rather than a + rewrite. + + What it adds is **provenance**. A campaign started from here records the + design it froze on and the point it was given, so a report can answer + "where did this operating point come from" with a path instead of a + shrug. The two campaigns that motivated this module could not. + + No ``ctx_json``: this module does not build an end-to-end view (see + :data:`NO_JOIN`), so the gen half is scored anchor-free and says so. + + What a half inherits is :data:`INHERITED` — see the rule stated there. + """ + if track not in TRACKS: + raise DisaggSolError(f"unknown track {track!r}, expected one of {list(TRACKS)}") + spec: dict[str, Any] = {key: value for key, value in base.items() if key in INHERITED} + spec["trtllm_repo_path"] = str(repo) + spec["sol_track"] = { + "track": track, + "sweep": str(sweep), + # Recorded, not acted on: this layer chose the point, and the + # campaign must be able to say so without reading this module. + "point_provenance": { + "design_dir": str(design), + # Both vocabularies, because the two halves are selected on + # different axes and a campaign has to be able to state its own + # point: a ctx point is (tp_size, max_batch), a gen point is + # (shape, concurrency). Recording only one leaves the other half's + # provenance describing nothing. + "selected": { + key: point[key] + for key in ( + "shape", + "concurrency", + "ctx_gpus", + "max_batch", + "prefer", + "ranked_on", + ) + if key in point + }, + "e2e_view_absent": point.get("e2e_view_absent", NO_JOIN), + }, + } + return spec + + +def verify_sweep_matches_point(sweep: Path, track: str, point: Mapping[str, Any]) -> None: + """Refuse a sweep that does not contain the point this layer selected. + + Without this the selection is a **report**, not a decision. The chosen + point is written into every campaign's ``point_provenance``, and the + campaign runs whatever its sweep says — so a sweep that names a different + row produces a record claiming an operating point the run never used. + That is worse than not selecting at all: the whole reason this layer + exists is to make a campaign's starting point traceable, and an + untraceable point is at least honest about being untraceable. + + Checked rather than generated. Writing the sweep here would put this + module back in the business of authoring configs it does not own, which + is the mistake :mod:`.sweep_design` was just unwound for. The sweep stays + something a person or the design skill wrote; this refuses the pairing. + """ + try: + config = load_yaml(Path(sweep)) + except Exception as exc: # noqa: BLE001 - re-raised with the pairing named + raise DisaggSolError(f"could not read the {track} sweep {sweep}: {exc}") from exc + + if track == GEN_TRACK: + wanted = (point.get("shape"), point.get("concurrency")) + if wanted[0] is None: + return # a ctx-only selection: nothing to check the gen sweep against + try: + # `operating_point`, not the row's own `concurrency`: the selected + # point's number is the deployment total the measured case was + # named after, and the row's is per generation server. They agree + # only at `gen_num == 1`, so comparing the raw fields refuses + # every correct multi-server sweep and says the sweep diverged. + available = { + (case["name"], operating_point(case["config"] or {})) for case in gen_cases(config) + } + except BenchCliError as exc: # pragma: no cover - message path + raise DisaggSolError(str(exc)) from exc + if wanted not in available: + raise DisaggSolError( + f"the gen sweep {sweep} does not contain the selected point " + f"{wanted[0]} @ concurrency {wanted[1]} (requests in flight across " + f"the deployment, which is the row's own list times its gen_num). " + f"It expands to {sorted(available)}. The campaign would run one of " + f"those while its point_provenance claimed the selected one -- a " + f"record that says the run used an operating point it did not." + ) + return + + wanted_ctx = (point.get("ctx_gpus"), point.get("max_batch")) + if wanted_ctx[0] is None: + return + available_ctx = { + ((case["config"] or {}).get("tp_size"), (case["config"] or {}).get("max_batch")) + for case in ctx_cases(config) + } + if wanted_ctx not in available_ctx: + raise DisaggSolError( + f"the ctx sweep {sweep} does not contain the selected point " + f"tp_size {wanted_ctx[0]} @ max_batch {wanted_ctx[1]}. It expands to " + f"{sorted(available_ctx)}. The campaign would run one of those while " + f"its point_provenance claimed the selected one." + ) + + +#: Where a derived sweep may be written, and the whole of the rule. +#: +#: Inside the campaign's own workspace, never anywhere else. The check is +#: blunt on purpose: an earlier module in this package took an output path as +#: a free parameter and would have overwritten a curated config that other +#: people maintain, which is unrecoverable in a way that a wrong measurement +#: is not. A derivation that cannot escape the workspace cannot make that +#: mistake, whatever it is asked to write. +DERIVED_SWEEP_NAME = "sweep-at-selected-point.yaml" + + +def derive_sweep_at_point( + design_sweep: Path, + track: str, + point: Mapping[str, Any], + *, + into: Path, + repo: Path | None = None, +) -> Path: + """Cut the design's sweep down to the one row the campaign will freeze on. + + The last link in the chain, and it exists because of a shape problem + rather than a preference. The design sweep is the whole measured space — + six shapes, each with its own concurrency ladder — and a campaign cannot + run it: two shapes measured at one concurrency both land in + ``concurrency_``, which :func:`.sol_track._place` refuses, because a + campaign's operating points have to be addressable by concurrency alone. + So the campaign needs exactly one row, and which row is not known until + the design has been measured and selected from. + + Derived rather than authored: every field except the row filter comes + from the design's own sweep, so the campaign measures the configuration + the selection was made on and not a hand-typed approximation of it. The + two campaigns this layer replaces differed from their own recorded point + in three fields at once, and nothing noticed. + + Written into the campaign's workspace and nowhere else — see + :data:`DERIVED_SWEEP_NAME`. + """ + into = Path(into).resolve() + out = into / DERIVED_SWEEP_NAME + config = load_yaml(Path(design_sweep)) + + if track == GEN_TRACK: + wanted = (point.get("shape"), point.get("concurrency")) + kept = [row for row in (config.get("gen_configs") or []) if _gen_row_matches(row, wanted)] + if len(kept) != 1: + raise DisaggSolError( + f"{design_sweep} has {len(kept)} rows matching the selected point " + f"{wanted[0]} @ concurrency {wanted[1]}; a campaign needs exactly " + f"one. The design sweep and the measured space have diverged." + ) + row = list(kept[0]) + # Written back in the row's OWN unit. `wanted[1]` is the deployment + # total the measured point was named after; this field is per + # generation server. Writing the total here would restate the point + # as a `gen_num`-times larger one and still look like the row it was + # cut from. + row[9] = str(_gen_row_concurrency(kept[0], wanted)) + config["gen_configs"] = [row] + else: + wanted_ctx = (point.get("ctx_gpus"), point.get("max_batch")) + kept_ctx = [ + {**entry, "tp_size": [wanted_ctx[0]], "max_batch": [wanted_ctx[1]]} + for entry in (config.get("benchmarks") or []) + if isinstance(entry, Mapping) + and wanted_ctx[0] in (entry.get("tp_size") or []) + and wanted_ctx[1] in (entry.get("max_batch") or []) + ] + if len(kept_ctx) != 1: + raise DisaggSolError( + f"{design_sweep} has {len(kept_ctx)} benchmark entries matching the " + f"selected ctx point tp_size {wanted_ctx[0]} @ max_batch " + f"{wanted_ctx[1]}; a campaign needs exactly one." + ) + # `mtp_range` is not part of a ctx point and must not survive as a + # range. `select_ctx_point` groups on (ctx_gpus, max_batch, adp) and + # means the rest, because a ctx case runs at output_length 1: there + # is no decode, so the speculation depth changes nothing it measures. + # The design's four values are four REPEATS, which is why the point + # it produced is a mean over twelve of them and not a reading of one. + # + # Left as a range, the campaign would re-run all four every time it + # measured -- and the run this was found on did not: the agent + # narrowed to MTP0 at submit time, outside the sweep, so the file + # said twelve cases and the campaign booked one draw. Narrowed here + # instead, so what is measured is what is written down. + entry = dict(kept_ctx[0]) + depths = [d for d in (entry.get("mtp_range") or []) if isinstance(d, int)] + if depths: + entry["mtp_range"] = [min(depths)] + config["benchmarks"] = [entry] + config.pop("gpu_overrides", None) + + if repo is not None: + # The design measured the image on purpose -- `sweep_design` refuses a + # build source there, because choosing an operating point on code no + # campaign starts from would pick the point for a different program. + # A campaign is the opposite case: it exists to measure its own edits, + # and a sweep without a rung would run the image and report every + # change as no-gain. So the rung is added here, at the boundary where + # the purpose changes, pointing at this campaign's own checkout. + config["trtllm_install"] = {"trtllm_repo": str(repo)} + config["_derived_from"] = { + "design_sweep": str(design_sweep), + "selected": {k: point.get(k) for k in ("shape", "concurrency", "ctx_gpus", "max_batch")}, + "why": ( + "the design sweep is the whole measured space; a campaign freezes one " + "row of it, because two shapes at one concurrency cannot both be " + "'concurrency_'" + ), + # The one number a reader needs to not over-read the campaign's + # percentages. The point was chosen from a mean; the campaign + # measures its own mean over a different, smaller n, so the two are + # the same KIND of quantity at different precision -- comparable, but + # not interchangeable, and a gain quoted against the design's number + # rather than the campaign's own baseline inherits the difference. + "repeats": { + "design": point.get("repeats"), + "campaign": config.get("rounds"), + "note": ( + "the selected value is a mean over the design's repeats; this " + "campaign re-measures its own baseline over 'campaign' repeats and " + "scores every attempt against that. Quote the campaign's baseline, " + "not the design's point, as what a gain is relative to." + ), + }, + } + into.mkdir(parents=True, exist_ok=True) + out.write_text(yaml.safe_dump(config, sort_keys=False, allow_unicode=True), encoding="utf-8") + return out + + +def _gen_num(row: Any) -> int: + """How many generation servers a sweep row stands up, defaulting to 1. + + The row order is the harness': ``[ctx_num, gen_num, tp_size, batch, mnt, + adp, gmf, mtp, eplb, concurrency_list]``. Defensive about the value for + the same reason :func:`.bench_cli.operating_point` is -- a malformed row must + not silently become a different operating point. + """ + if not isinstance(row, (list, tuple)) or len(row) < 10: + return 1 + value = row[1] + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + return 1 + return value + + +def _gen_row_concurrency(row: Any, wanted: tuple) -> int | None: + """The row's PER-SERVER concurrency at the selected point, or ``None``. + + The two sides of this comparison are in different units, which is the + whole reason this is a function. A measured point's ``concurrency`` comes + from ``gen_only_perf.csv``, and the harness names its result directories + after the requests in flight across the DEPLOYMENT -- ``listed x + gen_num``. The sweep row's own list is PER GENERATION SERVER. Comparing + them directly is correct exactly when ``gen_num == 1`` and silently wrong + otherwise: every multi-generation-server shape would have its own correct + sweep refused with "0 rows matching the selected point". + """ + if not isinstance(row, (list, tuple)) or len(row) < 10: + return None + tp, adp, mtp, eplb = row[2], row[5], row[7], row[8] + name = f"{'dep' if adp else 'tep'}_{tp}_eplb{eplb}_mtp{mtp}" + if name != wanted[0]: + return None + servers = _gen_num(row) + for token in str(row[9]).split(","): + token = token.strip() + if token.isdigit() and int(token) * servers == wanted[1]: + return int(token) + return None + + +def _gen_row_matches(row: Any, wanted: tuple) -> bool: + return _gen_row_concurrency(row, wanted) is not None + + +class CampaignLaunch: + """One campaign, described completely enough to be checked before it runs. + + Separated from the spawning so the *decision* — which track, which + sweep, which checkout, which workspace — is testable without starting a + process. The spawn itself is four lines and has no judgement in it; this + is where the judgement is. + """ + + def __init__(self, track: str, spec: Mapping[str, Any], workspace: Path, task_path: Path): + self.track = track + self.spec = dict(spec) + self.workspace = Path(workspace) + self.task_path = Path(task_path) + + @property + def argv(self) -> list[str]: + return [ + "perf-optimize", + "--task", + str(self.task_path), + "--workspace", + str(self.workspace), + ] + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return f"CampaignLaunch(track={self.track!r}, workspace={self.workspace})" + + +def launch_plan( + base: Mapping[str, Any], + *, + sweeps: Mapping[str, Path], + repos: Mapping[str, Path], + workspace_root: Path, + label: str, + points: Mapping[str, Mapping[str, Any]], + design: Path, + design_sweeps: Mapping[str, Path] | None = None, +) -> list[CampaignLaunch]: + """What this spec would start, in full, before anything starts. + + Every campaign is checked for the two things that made the hand-launched + pair safe and which nothing else enforces: **its own checkout** and **its + own workspace**. Two campaigns pointed at one checkout review each + other's worktree; two pointed at one workspace overwrite each other's + results. Both failures produce plausible numbers, which is why they are + refused here rather than left to the launcher's care. + """ + wanted = tracks(base) + design_sweeps = dict(design_sweeps or {}) + missing = [t for t in wanted if t not in sweeps and t not in design_sweeps] + if missing: + raise DisaggSolError( + f"track(s) {missing} have neither a sweep of their own nor a " + f"'{DESIGN_KEY}.{DESIGN_SWEEP_KEY}' entry to cut one from. A campaign " + f"has to freeze exactly one row, and which row is not known until the " + f"design has been measured -- so either name the design's sweep for " + f"that half and let it be derived, or name a sweep that already " + f"contains its selected point. The two halves are measured by " + f"different sweeps, so `{DESIGN_SWEEP_KEY}` is per track." + ) + missing = [t for t in wanted if t not in repos] + if missing: + raise DisaggSolError(f"no trtllm_repo_path given for track(s) {missing}") + + seen_repos: dict[str, str] = {} + launches: list[CampaignLaunch] = [] + for track in wanted: + repo = Path(repos[track]).resolve() + if str(repo) in seen_repos: + raise DisaggSolError( + f"tracks {seen_repos[str(repo)]!r} and {track!r} both name the checkout " + f"{repo}. A campaign resets the checkout it is given, so two of them " + f"sharing one would each revert the other's work mid-flight — and the " + f"measurement that followed would be of neither's code." + ) + seen_repos[str(repo)] = track + workspace = campaign_workspace(workspace_root, track, label) + point = points.get(track) or {} + if not point: + raise DisaggSolError(f"no selected point for track {track!r}") + if track in sweeps: + # A sweep the caller chose: check it, never rewrite it. + sweep = Path(sweeps[track]) + verify_sweep_matches_point(sweep, track, point) + else: + sweep = derive_sweep_at_point( + Path(design_sweeps[track]), track, point, into=workspace, repo=repo + ) + launches.append( + CampaignLaunch( + track=track, + spec=campaign_spec( + base, + track=track, + sweep=sweep, + repo=repo, + point=point, + design=design, + ), + workspace=workspace, + task_path=workspace / "task.yaml", + ) + ) + return launches + + +# --------------------------------------------------------------- the run + + +#: Where the supervisor records what it selected and started. +RUN_RECORD = "disagg_sol_run.json" + + +def supervise( + base: Mapping[str, Any], + *, + sweeps: Mapping[str, Path], + repos: Mapping[str, Path], + workspace_root: Path, + label: str, + incumbent: Mapping[str, Any] | None = None, + dry_run: bool = False, + designer: Callable[[str], Any] | None = None, +) -> dict[str, Any]: + """Fix the point from an established design, then start each half at it. + + The design is an **input**, not something this produces: establishing it + costs roughly an order of magnitude more than the campaigns it enables, + so it is run once per (model, cluster, workload) and reused. A spec whose + design has not been scored is refused here rather than silently falling + back to whatever operating point the sweeps happen to carry — falling + back is exactly the behaviour this layer exists to remove. + + Returns the record it writes: what was selected, against what incumbent, + and what was started. ``dry_run`` stops after the record, which is the + same code path minus the processes. + """ + from agent_flow.workflows.perf_optimize import spawn + + design = design_dir(base) + prefer = preference(base) + wanted = tracks(base) + + # Resolved BEFORE the designer is considered, not after. A design costs + # about an order of magnitude more than the campaigns it enables, so + # paying for one that already exists next door is the expensive half of + # this mistake -- and the reason the agent resumed a neighbour in the + # first place was to avoid exactly that. + requested_design = design + design, redirect = resolve_design_dir(design, wanted) + unready = [t for t in wanted if not established(design, t)] + if unready and designer is not None: + # The design is normally an input -- it is reused across campaigns and + # costs an order of magnitude more than any of them. But "an input" + # degenerated into "nobody ran it" once already, which is how the two + # campaigns this layer replaces came to inherit an unmeasured row. So + # when a designer is available the run establishes what it needs and + # then re-checks: not a fallback, a first step. + designer( + designer_instruction( + model_dir=Path(design).parent.name, design_dir=design, tracks=unready + ) + ) + # ...and read back again after it runs, because the agent resumes + # from a state file and may have established the design somewhere + # other than where it was sent. + design, after = resolve_design_dir(design, wanted) + redirect = after or redirect + unready = [t for t in wanted if not established(design, t)] + if unready: + what = { + CTX_TRACK: "no scored ctx case (a `run_*.json` the harness validated)", + GEN_TRACK: f"no scored concurrency sweep (a `{GEN_ONLY_CSV}`)", + } + raise DisaggSolError( + f"{design} has " + + "; ".join(f"{what[t]} for track '{t}'" for t in unready) + + f". There is no measured space to choose those halves' operating " + f"points from. Establish the design first — it is reused across " + f"campaigns, so this is paid once — or narrow " + f"'{DISAGG_SOL_FIELD}.{TRACKS_KEY}' to the halves it already covers." + ) + + record: dict[str, Any] = { + "design_dir": str(design), + **({"design_dir_redirected": redirect} if redirect else {}), + "design_state": design_state(design).get("phase"), + "tracks": wanted, + "label": label, + } + + if GEN_TRACK in wanted: + record["gen_point"] = select_point(sweep_points(design), prefer=prefer, incumbent=incumbent) + if CTX_TRACK in wanted: + # No `prefer`: the ctx objective is scalar. See `select_ctx_point`. + record["ctx_point"] = select_ctx_point(ctx_points(design), incumbent=incumbent) + + # One point per half. They are different measurements on different + # axes -- a ctx point is (tp_size, max_batch), a gen point is (shape, + # concurrency) -- so handing one to both halves asks the ctx sweep for a + # row described in the generation half's vocabulary. + points = {track: record[f"{track}_point"] for track in wanted if f"{track}_point" in record} + block = _block(base).get(DESIGN_KEY) + block = block if isinstance(block, Mapping) else {} + stated = block.get(DESIGN_SWEEP_KEY) + stated = {GEN_TRACK: stated} if isinstance(stated, str) else dict(stated or {}) + design_sweeps = { + track: rebase_design_sweep(Path(path), requested_design, design) + for track, path in stated.items() + } + record["design_sweeps"] = {k: str(v) for k, v in design_sweeps.items()} + launches = launch_plan( + base, + sweeps=sweeps, + repos=repos, + workspace_root=workspace_root, + label=label, + points=points, + design=design, + design_sweeps=design_sweeps, + ) + record["campaigns"] = [ + {"track": run.track, "workspace": str(run.workspace), "argv": run.argv} for run in launches + ] + + Path(workspace_root).mkdir(parents=True, exist_ok=True) + (Path(workspace_root) / RUN_RECORD).write_text( + json.dumps(record, indent=2, default=str) + "\n", encoding="utf-8" + ) + if dry_run: + # Write the specs, then stop -- which is what `--dry-run` has always + # advertised ("write every campaign's task.yaml, then stop without + # starting them") and did not do: `materialize` lives inside + # `start_all`, and returning above it meant the one artefact a reader + # would check the plan against was the one thing the dry run skipped. + # `spawn` separates the two calls precisely so this is the same code + # path minus the spawning, rather than a second implementation of it. + record["task_paths"] = {run.track: str(spawn.materialize(run)) for run in launches} + record["started"] = False + (Path(workspace_root) / RUN_RECORD).write_text( + json.dumps(record, indent=2, default=str) + "\n", encoding="utf-8" + ) + return record + + started = spawn.start_all(launches) + record["started"] = True + record["exit_status"] = spawn.wait_all(started) + (Path(workspace_root) / RUN_RECORD).write_text( + json.dumps(record, indent=2, default=str) + "\n", encoding="utf-8" + ) + return record diff --git a/agent-flow/agent_flow/workflows/perf_optimize/prompts/__init__.py b/agent-flow/agent_flow/workflows/perf_optimize/prompts/__init__.py index 6491eea8be70..2c0e5add76c7 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/prompts/__init__.py +++ b/agent-flow/agent_flow/workflows/perf_optimize/prompts/__init__.py @@ -12,6 +12,8 @@ SOL_ANALYZER_CONTEXT, SOL_OPTIMIZE_REPORTER_GUIDANCE, SOL_OPTIMIZER_CONTEXT, + SOL_TRACK_CTX, + SOL_TRACK_GEN, approach_restriction_note, kernel_coverage_analyzer_note, ) @@ -92,6 +94,12 @@ def _append(base: str, extra: str) -> str: ) +#: The section each SOL track composes. Keyed by +#: ``sol_track.track_name()``'s values so the caller passes the track +#: through rather than translating it here. +SOL_TRACK_SECTIONS: dict[str, str] = {"ctx": SOL_TRACK_CTX, "gen": SOL_TRACK_GEN} + + def build_perf_optimize_prompts( include_slurm_environment: bool = False, remote_execution: Mapping[str, Any] | None = None, @@ -101,6 +109,7 @@ def build_perf_optimize_prompts( kernel_coverage: Mapping[str, Any] | None = None, sol_methodology: str = "full", include_disagg: bool = False, + sol_track: str | None = None, ) -> PromptBundle: """Return the workflow's prompt bundle, augmented per the task spec. @@ -167,6 +176,14 @@ def build_perf_optimize_prompts( projector are left alone: neither stands up a server, and the reporter reads the artifacts the others produced either way. + ``sol_track`` is ``"ctx"`` or ``"gen"`` when the task spec carries a + ``sol_track`` block: the campaign optimizes one half of a + disaggregated deployment in isolation. The track's section is + appended to the same five roles and for the same reason as + ``include_disagg``, and the two are mutually exclusive — the schema + refuses a spec carrying both, since each reconciles ``benchmark`` + from a different file. + Composing it here rather than carrying it unconditionally is what keeps the override unambiguous — a role either has the section and it applies, or it does not have it at all. The alternative (always @@ -209,6 +226,15 @@ def build_perf_optimize_prompts( integrator=DISAGG_CAMPAIGN, qa=DISAGG_CAMPAIGN, ) + if sol_track is not None: + section = SOL_TRACK_SECTIONS[sol_track] + bundle = bundle.with_extensions( + benchmarker=section, + analyzer=section, + optimizer=section, + evaluator=section, + qa=section, + ) if kernel_coverage is not None: bundle = bundle.with_extensions( analyzer=kernel_coverage_analyzer_note( diff --git a/agent-flow/agent_flow/workflows/perf_optimize/prompts/_common.py b/agent-flow/agent_flow/workflows/perf_optimize/prompts/_common.py index e86849af039e..f58afa9b238b 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/prompts/_common.py +++ b/agent-flow/agent_flow/workflows/perf_optimize/prompts/_common.py @@ -1542,3 +1542,229 @@ def approach_restriction_note(allowed: Sequence[str]) -> str: allowed approach(es), or REJECT when the item cannot be realized through an allowed approach at all. """ + + +# The two SOL tracks. Each replaces the same guidance DISAGG_CAMPAIGN +# does -- server lifecycle, tuning file, profiling -- but for half a +# deployment measured in isolation, which is the whole point: the knobs +# an optimizer reaches live inside one role, and an end-to-end +# allocation prices in both. +# +# Short by construction. The mechanics live in `bench-disagg`, which the +# benchmark repo ships as its only supported agent-facing interface, and +# in one `sol_track` helper for the single thing that CLI has no notion +# of: this campaign's live tuning file. + +#: Steps 5-6 for the gen track, which has two scorers. Both compute the +#: same column by the same formula; they differ in what ELSE they report, +#: and therefore in whether a CTX anchor is needed at all. +_SOL_SCORING_GEN = """\ +# 5. Score it. WHICH command depends on whether `task.yaml` set +# `sol_track.ctx_json`. Run the one that matches; the other is not a +# fallback. +# +# (a) ctx_json IS set -- the rate-matched frontier, the end-to-end view: +ibc-bench process frontier -i $R --ctx_json --multi_round 1 +# +# (b) ctx_json is NOT set -- the anchor-free extractor: +python -m ibc_trtllm_harness.process_data.get_gen_only_perf -i $R +# +# Never run (a) with an anchor `task.yaml` did not declare. The one +# thing checked about an anchor is that its input length matches this +# sweep's; an undeclared anchor was checked against nothing, and a +# frontier built on a mismatched one still plots a clean curve -- +# describing a deployment whose two halves never served one traffic. + +# 6. Land the score where every later stage reads results from. +python -m agent_flow.workflows.perf_optimize.sol_track --workspace \\ + --collect + +# The **sol-postprocess** skill is the runbook for 5(a) and worth reading +# before your first one -- notably that `accept_rate` comes from the run's +# archived `sweep_config.yaml`, never from a default. The built-in table it +# replaced was measured 15 % high on one model, and it multiplies the metric +# linearly with no symptom. 5(b) reads it from that same archived file and +# refuses without it, for the same reason. +# +# What the two do NOT differ in is the gate. `throughput_per_user` is +# `accept_rate / avg_iteration_time` either way: decode iterations only, no +# context term, so no ctx measurement can move it. Path (b) is missing the +# END-TO-END half -- `output_tput_per_gpu`, the ctx:gen ratio, and the +# `frontier_elasticity` that prices a 1 % gain here against the deployment. +# On (b) those are ABSENT, not zero and not unchanged, and the result file +# says so in `e2e_view_absent`. Never quote a number from a (b) campaign as +# an end-to-end result. In particular `output_tps_per_gen_gpu`, which (b) +# DOES report, is not that number: it divides by the generation GPUs alone, +# which makes it always the flattering one. +# +# Step 6 reads whichever CSV step 5 wrote and lands it under the name +# `task.yaml` scores, at the total in flight rather than the row's listed +# concurrency. Do not hand-write that JSON.""" + +#: Steps 5-6 for the ctx track. No frontier, and running one is an error: +#: `process frontier` selects GEN measurements and takes ctx as its anchor. +_SOL_SCORING_CTX = """\ +# 5. Nothing to run. There is no `process frontier` on this track -- it +# builds the rate-matched GEN curve and takes ctx only as its anchor, +# so it answers with an error however well the ctx jobs ran. The score +# is already on disk: `performance.request_throughput_req_s` in the +# case's `run_*.json`, the field the harness itself validated it on. + +# 6. Land the score where every later stage reads results from. +python -m agent_flow.workflows.perf_optimize.sol_track --workspace \\ + --collect + +# It reads that `run_*.json` for you and files it at the case's `max_batch`, +# which for a prefill-only run IS the in-flight request count. +# Do not hand-write that JSON.""" + +_SOL_TRACK_SHARED = """\ +### Per attempt + +```bash +S=/sweep/ # this campaign's copy, not the original +R=/run + +# 0. approach: code ONLY -- declare how the change reaches the workers, in +# the sweep. The harness runs whatever the image ships unless you do, +# so a source edit alone measures the image and comes back at baseline. +# Take the cheapest rung that fits: +# trtllm_patch: single files, applied in-container, seconds, no rebuild +# trtllm_install.trtllm_repo: python-only editable install, minutes +# trtllm_install.build_wheel: full C++ rebuild +# Validation refuses `code` against a sweep that names none of them. + +# 1. Put the live tuning file into the sweep. Never skip: forgetting it does +# not fail, it measures the PREVIOUS attempt's configuration. +# It edits /sweep/ -- this campaign's own copy. The file the +# task.yaml was written against is never modified. +python -m agent_flow.workflows.perf_optimize.sol_track --workspace + +# 2. See what you are actually overriding, BEFORE spending an allocation. +# --dry-run materializes every case's worker config and queues nothing. +ibc-bench submit sweep -m all -c $S -w $R --dry-run +# Then READ what it wrote: $R/bm_*//gen_config.yaml (and ctx_config.yaml). +# Your tuning file is a PARTIAL OVERLAY deep-merged onto those, so +# "absent from the overlay" does NOT mean "unset" -- the sweep row's +# generator writes most of it. Proposing a value a key already holds +# costs a whole attempt and measures nothing. +# The **check-gen-config** skill lints the sweep's own invariants. + +# 3. Measure. +ibc-bench submit sweep -m all -c $S -w $R + +# 4. Poll until nothing is inflight -- in the FOREGROUND, never a background +# poll. The **check-job** skill diagnoses a hang from the logs; squeue +# alone will not tell you a job is stuck at the fill gate. +ibc-bench jobs check -f $R/bm_*/job_status.csv --summary +{scoring} +``` + +`-w $R` is not a free choice. The run directory the harness creates is named +for the workload and the date and **carries no code identity**, so two +attempts of one campaign submitted into one work dir compute the same name +and the second overwrites the first. Deriving it from the result directory +is what keeps them apart. + +**A measurement nobody wrote down did not happen.** Every later stage — the +baseline gate first — judges this attempt by whether a JSON under the result +directory it named carries `{metric}`. Getting all the way to a good number +and not landing it there is indistinguishable from a campaign that never +measured, and it stops the run. + +**Attribution is yours, and it is not in the numbers.** Nothing in this +stack records which code produced which measurement — there is no +configuration fingerprint to compare. So before you call anything an +improvement, confirm the change reached the workers by reading what was +served: `$R/bm_*//gen_config.yaml` for a tuning change, and the +per-node `trtllm_patch` sha256 manifest beside the case logs for a source +change. **Never infer "it took effect" from "the number moved."** Say in +your report which artifact you read; a claim you cannot point at is a claim +this campaign cannot make. + +Report sample counts alongside any delta, and if this workspace holds a +second measurement of the same configuration, report that repeatability +first — a delta smaller than it is not a result in either direction. + +`CANCELLED` is the **normal** end state for a GEN job: the harness cancels +its own allocation once the client finishes. Judge a case by whether it has +a validated measurement, never by its Slurm state. + +### Frozen for this campaign + +{frozen}. Those live in the sweep row and appear in the case name; the +tuning file is an overlay deep-merged *onto* the config that row generates. +Putting one of them in the overlay fights the row silently -- it is a +REJECT whatever it measured, and step 1 refuses it outright. + +### Profiling + +nsys only, on the {role} workers. torch profiler and ncu have **no path +through this harness**: record `not available in a SOL track campaign` and +plan from nsys -- never fabricate a trace. +""" + +SOL_TRACK_CTX = """\ +## CTX track (supersedes the server-lifecycle, tuning, and profiling guidance above) + +**This campaign optimizes the context half of a disaggregated deployment, +in isolation.** You do not launch `trtllm-serve`, poll `:8000`, or tear a +server down. The measurement is `trtllm-bench throughput` at +`output_length: 1`: one server, prefill only, no disaggregation and no KV +handoff at all. + +The metric is `avg_request_throughput_req_s` -- prefill requests per +second, which is what the end-to-end frontier consumes this half as. Not +tokens/s, not latency. + +- **`/tuning/extra_llm_api_options.yaml`** -- the ctx role's + overlay, and the only file you edit. +- The **gen** role belongs to a separate campaign against the same + deployment. Not yours to edit and not yours to reason about. + +""" + _SOL_TRACK_SHARED.format( + only="", + role="context", + metric="avg_request_throughput_req_s", + scoring=_SOL_SCORING_CTX, + frozen="`max_batch` and `tp_size` -- the anchor the frontier was computed against", +) + +SOL_TRACK_GEN = """\ +## GEN track (supersedes the server-lifecycle, tuning, and profiling guidance above) + +**This campaign optimizes the generation half of a disaggregated +deployment, in isolation.** You do not launch `trtllm-serve`, poll +`:8000`, or tear a server down. The measurement is a real 1-ctx-1-gen +deployment, but only the generation worker's decode iterations are read. + +The metric is `throughput_per_user` -- `tps_per_user` in the CLI's +output. It is the inverse of the steady-state decode step time, and two +consequences follow: + +- It is measured **only on iterations where the decode batch is exactly + full**; every perturbed iteration is filtered out. A change that + improves steady-state decode while worsening KV-handoff or admission + behaviour reads here as a **pure win** and is not one. If you have + reason to think a change moves handoff behaviour, say so in your + report -- this campaign cannot measure it. +- `accept_rate` is a frozen constant in the sweep's `options`, not a + measurement. Anything touching **speculative decoding** is out of + scope: its true effect would land in a number nobody re-measured. + Dismiss such items with that reason. + +- **`/tuning/extra_llm_api_options.yaml`** -- the gen role's + overlay, and the only file you edit. +- The **ctx** role belongs to a separate campaign against the same + deployment. Not yours to edit and not yours to reason about. + +""" + _SOL_TRACK_SHARED.format( + only=" [--only '']", + role="generation", + metric="throughput_per_user", + scoring=_SOL_SCORING_GEN, + frozen=( + "the sweep row -- instance counts, `tp_size`, `attention_dp`, `mtp`, " + "`eplb`, and each point's concurrency" + ), +) diff --git a/agent-flow/agent_flow/workflows/perf_optimize/sol_track.py b/agent-flow/agent_flow/workflows/perf_optimize/sol_track.py new file mode 100644 index 000000000000..f5f3057573a0 --- /dev/null +++ b/agent-flow/agent_flow/workflows/perf_optimize/sol_track.py @@ -0,0 +1,1124 @@ +"""SOL tracks: optimizing the two halves of a disaggregated deployment apart. + +An end-to-end disagg campaign (see :mod:`.disagg`) measures the whole +cluster at once. That is the most expensive measurement the flow can +make, and most of it is wasted: the knobs an optimizer can reach live +inside one role at a time, and a full e2e allocation prices in both. + +The decomposition ``bench-trtllm-disagg`` is built around splits the +measurement in two, and each half is an *aggregate-shaped* campaign: + +- **ctx** — ``trtllm-bench throughput`` at ``output_length: 1``. No + disaggregation at all: one server, prefill only. Its metric is + ``avg_request_throughput_req_s``. +- **gen** — a real 1-ctx-1-gen deployment, read only on the generation + worker's decode iterations. Its metric is ``throughput_per_user``. + +So neither track needs new tuning machinery, no stage changes, and no +new gate: N frontier points become ``benchmark.concurrency``, and +omitting ``optimize.max_regression_pct`` already means *any* point can +veto an attempt. + +**What this module is, and is not.** It does not drive the benchmark +repo's scripts, and it does not re-implement their sweep expansion. That +repo ships ``bench-disagg`` and states plainly that the CLI is the only +supported agent-facing interface, with the scripts as internal backends +— so a campaign points at one orchestration ``sweep.yaml`` and +:mod:`.bench_cli` asks ``sweep plan`` what that expands to. Everything +this file used to compute by parsing YAML — the operating points, the +sequence lengths, the case names — now arrives already expanded, from +the same code that will run them. An expansion computed twice is one +that eventually disagrees with itself, and the disagreement would +surface as a campaign quoting operating points it never measured. + +What remains is the problem :mod:`.disagg` names: **two files describe +one run.** The sweep owns the measurement conditions; ``task.yaml`` is +what the orchestrator reads to build every agent's prompt. When they +disagree nothing raises — the prompts simply quote numbers the run never +measured. So the reconciliation rule is unchanged: a condition the user +did not write is **filled**, one they wrote that **disagrees** is an +**error** naming both values. + +One trap survives the move, because it is a property of the harness +rather than of the file it was read from: a sweep row's concurrency is +**per generation server**. The client is driven at ``concurrency * +gen_num`` and the harness names its result directory after that product, +while the case name keeps the listed value because that is its address. +``task.yaml``'s ``concurrency`` means what an aggregate campaign means +by it, so the product is what belongs there. See +:func:`.bench_cli.operating_point`. + +Two pieces of glue remain, one per direction, and both are code rather +than prompt for the same reason: skipping either does not fail, it +produces a plausible wrong answer. :func:`apply_overlay` carries the +campaign's tuning file *into* the sweep before a submit — forget it and +the run measures the previous attempt's configuration under the new +attempt's name. :func:`collect` carries the score back *out* — forget it +and a measurement that succeeded is indistinguishable, to every later +stage, from one that never ran. +""" + +from __future__ import annotations + +import json +import shutil +from pathlib import Path +from typing import Any, Mapping + +import yaml + +# The full dotted name, not `from ... import bench_cli`: the dashboard +# loads these modules by path with `agent_flow` itself unimportable, and +# only the submodule form resolves from the sys.modules cache. +from agent_flow.workflows.perf_optimize.bench_cli import ( + IBC_BENCH, + BenchCliError, + frontier_points, + gen_only_points, + operating_points, +) + +SOL_TRACK_FIELD = "sol_track" +TRACK_KEY = "track" +SWEEP_KEY = "sweep" +WORKSPACE_KEY = "workspace" +CTX_JSON_KEY = "ctx_json" + +CTX_TRACK = "ctx" +GEN_TRACK = "gen" +TRACKS: tuple[str, ...] = (CTX_TRACK, GEN_TRACK) + +#: The metric each track is scored on. Both are higher-is-better, so +#: neither needs the ``_ms`` suffix rule in ``_normalized_gain_pct``. +#: These are the names ``task.yaml`` carries; ``frontier show`` spells +#: the gen one ``tps_per_user``, and :mod:`.bench_cli` is where the two +#: vocabularies meet. +TRACK_METRICS: dict[str, str] = { + CTX_TRACK: "avg_request_throughput_req_s", + GEN_TRACK: "throughput_per_user", +} + +#: The gen metric, as the frontier CSV already spells it. Kept as a +#: mapping because the ctx track reads a different file for a different +#: key, and the dispatch in :func:`collect` is on the track. +#: Only the gen track has one: ``frontier build`` selects ``stage == GEN`` +#: and refuses outright when a workspace holds none, because a snapshot +#: *is* the rate-matched generation curve. The ctx side enters it as the +#: anchor, never as a point — so there is no snapshot key to read a ctx +#: campaign's score out of. See :func:`collect`. +SNAPSHOT_METRICS: dict[str, str] = {GEN_TRACK: "throughput_per_user"} + +#: Written per operating point under the stage's result directory, so a +#: measurement made by ``bench-disagg`` is discoverable by every part of +#: perf-optimize that globs for result JSONs — the baseline gate first. +SOL_RESULT_NAME = "sol_result.json" + +#: Worker keys the sweep row fixes, and which the tuning overlay +#: therefore may not name. Each is a coordinate of the operating point: +#: a gen row is ``[ctx_num, gen_num, tp_size, batch, max_num_tokens, +#: attention_dp, gpu_mem_frac, mtp, eplb, concurrency]``, and a ctx +#: benchmark entry fixes ``max_batch`` and ``tp_size`` the same way. +#: +#: These are the *generated config's* spellings, because that is what the +#: overlay is deep-merged onto — not the sweep row's. Ported from the +#: script-driven implementation, where a dry run caught an override +#: taking ``tensor_parallel_size`` from 4 to 8 and the job from two nodes +#: to three: the run would have succeeded and the number would have +#: looked plausible, with the node count in a log line as the only trace, +#: while the comparison against a baseline taken at the old point was +#: void rather than merely weaker. +FROZEN_WORKER_KEYS = frozenset( + { + "tensor_parallel_size", + "pipeline_parallel_size", + "context_parallel_size", + "moe_expert_parallel_size", + "enable_attention_dp", + "max_batch_size", + "max_num_tokens", + } +) + +#: Same reason as :data:`.disagg.DISAGG_PROFILE_METHODS` — the harness +#: only knows how to wrap workers in nsys. +SOL_PROFILE_METHODS: tuple[str, ...] = ("nsys",) + +#: The key a sweep stage uses for the overlay deep-merged onto every +#: generated worker config. This is the campaign's tuning surface: it +#: changes what the workers run without changing which points are +#: measured, so case names stay put and ``frontier compare`` can align an +#: attempt against the baseline it is judged against. +TRACK_OVERLAY_KEYS: dict[str, str] = { + CTX_TRACK: "ctx_extra_llm_api", + GEN_TRACK: "gen_extra_llm_api", +} + + +class SolTrackError(ValueError): + """The track block, or the sweep it names, is unusable.""" + + +# --------------------------------------------------------------- the block + + +def has_sol_track(data: Mapping[str, Any]) -> bool: + """Whether the spec enables a SOL track campaign.""" + return SOL_TRACK_FIELD in data + + +def sol_track_block(data: Mapping[str, Any]) -> Mapping[str, Any] | None: + block = data.get(SOL_TRACK_FIELD) + return block if isinstance(block, Mapping) else None + + +def _text(data: Mapping[str, Any], key: str) -> str | None: + block = sol_track_block(data) + if block is None: + return None + value = block.get(key) + return value.strip() if isinstance(value, str) and value.strip() else None + + +def track_name(data: Mapping[str, Any]) -> str | None: + """Which half this campaign optimizes, or ``None`` if unstated.""" + return _text(data, TRACK_KEY) + + +def workspace_name(data: Mapping[str, Any]) -> str | None: + """The ``bench-disagg`` workspace this campaign measures into.""" + return _text(data, WORKSPACE_KEY) + + +def sweep_path(data: Mapping[str, Any]) -> Path | None: + """The orchestration ``sweep.yaml``: stages, server config, options.""" + value = _text(data, SWEEP_KEY) + return None if value is None else Path(value) + + +def ctx_json_path(data: Mapping[str, Any]) -> Path | None: + """An existing CTX anchor this campaign builds its frontier against. + + A **gen** campaign cannot be scored without one. ``tps_per_user`` is a + purely generation-side quantity, but the only thing that reports it is + ``frontier build``, which rate-matches the whole curve and therefore + needs the context request rate; with neither a measured ctx stage nor + this, it raises ``ANCHOR_MISSING`` — *after* the gen jobs have run. + + Pointing at an anchor somebody else measured is legitimate and often + right: a gen campaign freezes the ctx side by construction, so the + anchor is a constant. It is not free of consequence, though — the + build folds the anchor's digest into the snapshot's ``view_id``, so + ``frontier compare`` will report a curve built against a different + anchor as not ``comparable``. + """ + value = _text(data, CTX_JSON_KEY) + return None if value is None else Path(value) + + +def anchor_isl(anchor: Path) -> int | None: + """The input length the ctx anchor was measured at, if it says. + + ``get_ctx_throughput.py`` writes a list of rows, each carrying the + ``isl`` its measurement ran at alongside the throughput. One anchor + file can hold several rows (one per MTP size); they come from one + sweep, so the first that states an ``isl`` is the file's. + """ + try: + rows = json.loads(anchor.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return None + if isinstance(rows, Mapping): + rows = [rows] + if not isinstance(rows, list): + return None + for row in rows: + value = row.get("isl") if isinstance(row, Mapping) else None + if isinstance(value, int) and not isinstance(value, bool): + return value + return None + + +def require_matching_anchor(anchor: Path, plan: Mapping[str, Any]) -> str | None: + """Refuse a borrowed anchor that measured a different request shape. + + The frontier metric divides a decode rate by a prefill rate. Feed it a + ctx anchor taken at one input length and a gen curve taken at another + and it still returns a number, on a curve that still looks like a + frontier, describing a deployment whose two roles were never serving + the same traffic. Nothing about the output looks wrong — which is what + makes it the one failure the rest of this module's care would not + survive. + + Only the input length is compared, and only when the anchor states + one. A ctx measurement runs at ``osl: 1`` by construction, so an + output length would never match and comparing it would reject every + anchor; the dataset path differs legitimately too, since the tracks + read the ``_for_bench`` and ``_for_serve`` variants of one corpus. + Returns a note when the anchor is silent, because "not checked" and + "checked and matched" must not read alike. + """ + workload = plan.get("workload") + expected = (workload or {}).get("isl") if isinstance(workload, Mapping) else None + measured = anchor_isl(anchor) + if measured is None: + return ( + f"{anchor} states no 'isl', so nothing verified that it measured this " + f"campaign's requests — the frontier would rate-match against it either way" + ) + if isinstance(expected, int) and not isinstance(expected, bool) and measured != expected: + raise SolTrackError( + f"the ctx anchor {anchor} was measured at isl {measured}, but this " + f"campaign's sweep runs at isl {expected}. The frontier divides a decode " + f"rate by a prefill rate, so it would still produce a number and a curve " + f"— describing a deployment whose two roles never served the same traffic. " + f"Point at an anchor measured on this workload, or enable the sweep's ctx " + f"stage so the campaign measures its own." + ) + return None + + +def load_sweep(path: Path) -> dict[str, Any]: + """Read the orchestration file, or raise with a legible message.""" + try: + data = yaml.safe_load(path.read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError) as exc: # pragma: no cover - message path + raise SolTrackError(f"could not read sol_track sweep {path}: {exc}") from exc + if not isinstance(data, Mapping): + raise SolTrackError( + f"sol_track sweep {path} must be a YAML mapping, got {type(data).__name__}" + ) + return dict(data) + + +#: The rungs of the harness' escalation ladder, cheapest first. A sweep +#: that names none of them runs whatever the image already ships. +BUILD_SOURCE_KEYS = ("trtllm_patch", "trtllm_install") + + +def build_source(sweep: Mapping[str, Any]) -> dict[str, Any] | None: + """How a source change reaches the workers, if the sweep says. + + ``trtllm_patch`` is single-file overlays applied inside the container + before the workers start -- seconds, no rebuild, sha256-manifested per + node. ``trtllm_install`` is the two rungs above it: an editable + install from a repo (minutes) or a full wheel build. + """ + for key in BUILD_SOURCE_KEYS: + block = sweep.get(key) + if key == "trtllm_patch" and block: + return {key: block} + if isinstance(block, Mapping) and any( + str(block.get(k) or "").strip() for k in ("trtllm_repo", "trtllm_wheel_path") + ): + return {key: dict(block)} + return None + + +def require_build_source(task_data: Mapping[str, Any], approaches: Any) -> None: + """Refuse a source-changing campaign whose sweep installs nothing. + + The harness runs whatever the image already has unless the sweep says + otherwise. So with ``approach: code`` and no such block, the optimizer + edits ``trtllm_repo_path``, the job measures the image, the number + comes back at the baseline, and the evaluator rejects an untested + change as "no gain" -- a full allocation spent learning nothing about + it, with nothing anywhere reporting an error. + """ + if "code" not in (approaches or []): + return + track = track_name(task_data) + sweep_file = sweep_path(task_data) + if track not in TRACKS or sweep_file is None or not sweep_file.is_file(): + return + if build_source(load_sweep(sweep_file)) is None: + raise SolTrackError( + f"'optimize.approaches' includes 'code', but {sweep_file} names neither " + f"'trtllm_patch' nor 'trtllm_install'. The harness runs whatever the " + f"image ships unless one is given, so a source change would never reach " + f"the workers: the measurement would come back at the baseline and the " + f"change would be rejected as 'no gain' without having been tested. " + f"Take the cheapest rung that fits the change:\n" + f" trtllm_patch: # single files, seconds, no rebuild\n" + f" files:\n" + f" - src: patches/.py\n" + f" dst: tensorrt_llm/_torch/.../.py\n" + f" trtllm_install: # python-only editable install, minutes\n" + f" trtllm_repo: \n" + f" trtllm_install: # full C++ rebuild\n" + f" trtllm_repo: ...\n" + f" build_wheel: true\n" + f"or drop 'code' from approaches." + ) + + +def stage_config_path(sweep: Mapping[str, Any], track: str, sweep_file: Path) -> Path | None: + """Where this track's overlay goes: the sweep file itself. + + There is no stage indirection any more. An `ibc-bench` sweep is one + file per (model, workload, mode) -- a gen sweep carries `gen_configs` + and the ctx worker it deploys beside them, a ctx config carries + `benchmarks`. So the file the campaign points at IS the file its + overlay belongs in, and the signature is kept only because callers + read better naming the concept than the path. + """ + return sweep_file if sweep_file.is_file() else None + + +def sweep_accept_rate(sweep: Mapping[str, Any]) -> str | None: + """The frozen acceptance length the sweep carries. + + Top-level ``accept_rate: {rate, source}`` since harness v0.4.7, and + the harness archives the sweep verbatim at submit so every + post-processing entry point reads the rate back out of the run rather + than being handed one. It is required here for the same reason it is + required there: with `mtp > 0` cases and no rate, post-processing + stops rather than falling back on a built-in table -- a table that + was measured, worst case, 15 % high, and which multiplies the metric + linearly with no symptom. + """ + block = sweep.get("accept_rate") + if isinstance(block, Mapping): + value = block.get("rate") + else: + value = block + return value.strip() if isinstance(value, str) and value.strip() else None + + +# --------------------------------------------------------------- the copy + + +#: Where a campaign keeps the sweep it actually measured. +WORKSPACE_SWEEP_DIR = "sweep" +#: Where the adopted copy came from. Written once, on the first adopt, and +#: preserved across resumes -- it is what `--clean` restores from. +ADOPTED_FROM_KEY = "adopted_from" +#: Set only when the spec's sweep was gone and `ADOPTED_FROM_KEY` supplied a +#: live original. Carries why, because the re-adopted file may have changed +#: since the copy was taken. +READOPTED_KEY = "readopted_from_origin" + + +def adopt_sweep(task_data: dict[str, Any], workspace: Path) -> Path | None: + """Copy the sweep into the workspace and point the spec at the copy. + + ``apply_overlay`` writes this campaign's tuning into a sweep stage + config before every submit. Writing that into the file the user named + is a mistake with a long tail: nothing restores it, so the *next* + campaign seeds its tuning file from the *previous* one's accepted + overlay and calls the result a baseline -- measured with an + optimization applied, reported as the sweep's own. The rewrite also + drops every comment in that file, and any key the user had under the + overlay key. + + The disagg campaign already answered this: it synthesizes a fresh + config into its own artifact directory rather than editing the harness + config it was handed. The only difference here is that a sweep is a + *directory* -- a stage config resolves its siblings relative to + itself, and the harness imports the model's ``gen_worker_config.py`` + from beside it -- so the unit of copy is the directory, not the file. + + Copying also makes the campaign self-describing: ``/sweep/`` + is what this run measured, not whatever the original has become since. + + Returns the copied sweep file, or ``None`` when there is nothing to + adopt. Idempotent by design: an existing copy is left alone, so a + resumed run keeps measuring what it started with, and ``--clean`` + (which removes that copy along with the rest of the managed workspace) + is what starts over from the original. + + ``--clean`` is also the one case where the spec's own ``sweep`` can + dangle: a run resumed with ``--task /task.yaml`` names the + copy, and ``--clean`` has just deleted it. That is recoverable rather + than fatal -- ``adopted_from`` still records where the copy came from -- + so the original is re-adopted and the redirection is written into the + block, for the same reason a redirected design directory is: a reader + must not have to already know. + """ + block = dict(task_data.get(SOL_TRACK_FIELD) or {}) + origin = block.get(ADOPTED_FROM_KEY) + origin = origin.strip() if isinstance(origin, str) else "" + + source = sweep_path(task_data) + readopted = "" + if source is not None and not source.is_file() and origin: + candidate = Path(origin).expanduser() + if candidate.is_file(): + readopted = ( + f"{source} is gone -- the workspace copy, removed by --clean -- so the " + f"campaign was re-adopted from {candidate}, which '{ADOPTED_FROM_KEY}' " + f"recorded. What this run measures therefore comes from the original " + f"sweep as it is NOW, not as the deleted copy had it." + ) + source = candidate + if source is None or not source.is_file(): + return None + + destination = workspace / WORKSPACE_SWEEP_DIR + target = destination / source.name + if not destination.exists(): + shutil.copytree( + source.parent, + destination, + # Build artifacts of the harness' own import of the model's + # config generator, and of previous runs. Copying them wastes + # space and, for `__pycache__`, risks a stale module. + ignore=shutil.ignore_patterns("__pycache__", "*.pyc", ".git"), + ) + if not target.is_file(): + raise SolTrackError( + f"the sweep copy at {destination} has no {source.name}; remove it and " + f"re-run, or pass --clean to rebuild the workspace from {source}" + ) + block[SWEEP_KEY] = str(target) + # Only when there is nothing to keep. On a resume the spec's sweep is + # already the workspace copy, and overwriting the origin with that path + # would erase the only pointer back to the file this campaign was cut + # from -- after which --clean has nothing left to restore it from. + if not origin or readopted: + block[ADOPTED_FROM_KEY] = str(source) + if readopted: + block[READOPTED_KEY] = readopted + task_data[SOL_TRACK_FIELD] = block + return target + + +# --------------------------------------------------------------- the tuning seed + + +def tuning_seed_yaml(data: Mapping[str, Any]) -> str: + """The live tuning file's seed: this track's worker overlay. + + Deliberately **one role's** overlay, and deliberately a *partial* one. + ``bench-disagg`` deep-merges it onto the worker config the sweep row + generated, so the row keeps the knobs that define the operating point + — parallel sizes, batch, token limits — and the optimizer edits only + what it is actually tuning. The previous shape, a whole copy of the + role's worker config, put the frozen topology inside the file the + optimizer edits, which is an invitation to change it by accident. + + Empty is the normal starting point: most sweeps carry no overlay. + """ + track = track_name(data) + if track not in TRACKS: + raise SolTrackError(f"unknown sol_track track {track!r}, expected one of {list(TRACKS)}") + path = sweep_path(data) + if path is None: + raise SolTrackError(f"'{SOL_TRACK_FIELD}.{SWEEP_KEY}' is required") + sweep = load_sweep(path) + stage_config = stage_config_path(sweep, track, path) + overlay: Any = None + if stage_config is not None and stage_config.is_file(): + overlay = load_sweep(stage_config).get(TRACK_OVERLAY_KEYS[track]) + if overlay is None: + overlay = {} + if not isinstance(overlay, Mapping): + raise SolTrackError( + f"'{TRACK_OVERLAY_KEYS[track]}' in {stage_config} must be a mapping, " + f"got {type(overlay).__name__}" + ) + return yaml.safe_dump(dict(overlay), sort_keys=False, default_flow_style=False) + + +# --------------------------------------------------------------- reconciliation + + +def _reconcile( + task_data: dict[str, Any], + expected: Mapping[str, Any], + why: Mapping[str, str], + user_set: set[str], +) -> list[str]: + """Fill what the user left out, error on what they wrote differently. + + The same rule :mod:`.disagg` states and for the same reason: a + ``task.yaml`` whose stated operating point is quietly replaced is a + file you cannot read, and "my setting did nothing" is the failure + this avoids. + """ + benchmark = dict(task_data.get("benchmark") or {}) + notes: list[str] = [] + conflicts: list[str] = [] + for key, value in expected.items(): + if key in user_set and benchmark.get(key) != value: + conflicts.append( + f"'benchmark.{key}' is {benchmark.get(key)!r} but the sweep gives " + f"{value!r} ({why[key]})" + ) + elif key not in user_set: + benchmark[key] = value + notes.append(f"benchmark.{key}={value!r} from the sweep ({why[key]})") + if conflicts: + bullet = "\n - " + raise SolTrackError( + f"task.yaml contradicts the sweep, which owns the measurement " + f"conditions:{bullet}{bullet.join(conflicts)}{bullet}" + f"remove these keys to take the sweep's values." + ) + task_data["benchmark"] = benchmark + return notes + + +def apply_plan(task_data: dict[str, Any], plan: Mapping[str, Any], user_set: set[str]) -> list[str]: + """Reconcile ``task_data`` against a ``sweep plan`` envelope, in place. + + The plan is the authority because it is the same expansion that will + run: its ``workload`` block names the corpus the client will read, and + its ``cases`` are the points that will be measured, already multiplied + out from whatever the sweep rows listed. + + **The plan's ``isl`` is not the workload's input length.** It is a + sizing bound and a client argument — ``submit.py`` spends it on + ``ctx_max_seq_len = isl + offset``, ``gen_max_seq_len = isl + osl + + offset``, ``benchmark.input_length`` and the job name — while the + requests themselves come from ``workload.dataset``, a corpus with its + own length distribution. The checked-in reference sweep for this + campaign pairs ``isl: 200000`` with a ``...-coding-c190000-...`` + dataset, and both numbers are right: one bounds the KV allocation, the + other describes the traffic. + + So a dataset run is reconciled as a dataset run. Copying ``isl`` into + ``benchmark.random_input_len`` asserted two things that are not true — + that the dataset is synthetic (``dataset_name`` defaults to + ``random``) and that every request is exactly that long — and every + prompt, the analyzer's reasoning and the report then quoted them. + """ + track = track_name(task_data) + if track not in TRACKS: + raise SolTrackError(f"unknown sol_track track {track!r}, expected one of {list(TRACKS)}") + + stage_cases = [case for case in plan.get("cases") or [] if case.get("stage") == track] + if not stage_cases: + raise SolTrackError( + f"the sweep plans no '{track}' cases. Enable that stage in the sweep's " + f"`stages:` block, or point this campaign at the track the sweep measures." + ) + points = operating_points({"cases": stage_cases}) + if not points: + raise SolTrackError(f"no '{track}' case in the plan carries a usable concurrency") + + workload = plan.get("workload") + workload = workload if isinstance(workload, Mapping) else {} + expected: dict[str, Any] = {"concurrency": points[0] if len(points) == 1 else points} + why = { + "concurrency": ( + "each planned case's max_batch, the in-flight request count of a prefill-only run" + if track == CTX_TRACK + else "each planned case's concurrency x its gen_num" + ) + } + dataset = workload.get("dataset") + dataset = dataset.strip() if isinstance(dataset, str) and dataset.strip() else None + isl, osl = workload.get("isl"), workload.get("osl") + if dataset is not None: + expected["dataset_name"] = Path(dataset).name + expected["dataset_path"] = dataset + why["dataset_name"] = "the plan's workload.dataset — the corpus the client reads" + why["dataset_path"] = "the plan's workload.dataset, resolved by the harness" + else: + # No corpus named, so the lengths really are the request shape. + # Every sweep this has been run against names one; the branch is + # here because a template that leaves the dataset per-row exists. + if isinstance(isl, int) and not isinstance(isl, bool): + expected["random_input_len"] = isl + why["random_input_len"] = "the plan's workload.isl, with no dataset to read from" + if isinstance(osl, int) and not isinstance(osl, bool): + expected["random_output_len"] = osl + why["random_output_len"] = "the plan's workload.osl, with no dataset to read from" + + notes = _reconcile(task_data, expected, why, user_set) + + if dataset is not None: + # The defaults block seeds `random_input_len` / `random_output_len` + # for every campaign, and on a dataset run they describe nothing — + # keeping them is how "1024" or an isl ends up quoted as this + # campaign's input length. Dropped rather than corrected, because + # the honest value is a distribution the sweep does not carry. + # A user who wrote one keeps it: `_reconcile` already refuses a + # value that contradicts the sweep, and this is not that. + benchmark = task_data["benchmark"] + for key in ("random_input_len", "random_output_len"): + if key not in user_set and benchmark.pop(key, None) is not None: + notes.append(f"benchmark.{key} dropped: this campaign reads {dataset}") + notes.append( + f"workload.isl={isl!r} / osl={osl!r} are the harness' sequence-length " + f"bounds (ctx_max_seq_len = isl + offset, gen_max_seq_len = isl + osl + " + f"offset) and the client's input_length — NOT the corpus' measured " + f"lengths. Quote the dataset, not these." + ) + + # Checked here rather than beside the other anchor rules because it + # needs the plan: the sweep's own isl is what the anchor has to match, + # and the plan is where that arrives already resolved. + anchor = ctx_json_path(task_data) + if anchor is not None and anchor.is_file(): + unverified = require_matching_anchor(anchor, plan) + notes.append(unverified or f"ctx anchor {anchor} measured this campaign's isl") + elif anchor is None and track == GEN_TRACK: + # Said once, here, where the resolved spec records it -- so a + # reader who never opens a result still learns which half of the + # deployment this campaign can see. + notes.append(NO_E2E_VIEW) + + # TODO: derive the nsys window from the baseline instead of taking it + # from the spec. ``profile.nsys_iter_range`` is a dead field on both + # SOL tracks -- two real campaigns each asked for 100-150 and each + # captured something else (the gen harness fired 200-250, the ctx one + # 30-50), and on the ctx run the whole benchmark was only 100 + # iterations, so the requested window could not have caught anything + # even if it had been honoured. A window is only meaningful relative + # to how many iterations the run actually has, which the benchmarker + # measures; a number typed in before the baseline exists is a guess + # that nothing reconciles. Until then the field is left alone rather + # than quietly rewritten, so at least the spec and the trace disagree + # visibly. (Same shape as the disagg block's open A3.) + profile = dict(task_data.get("profile") or {}) + dropped = [m for m in (profile.get("methods") or []) if m not in SOL_PROFILE_METHODS] + profile["methods"] = list(SOL_PROFILE_METHODS) + profile.pop("kernel_coverage", None) + task_data["profile"] = profile + if dropped: + notes.append( + f"profile.methods {dropped} dropped: the benchmark harness only wraps " + f"workers in nsys (no torch-profiler env var, no ncu path)" + ) + + if task_data.pop("accuracy", None) is not None: + # The same call disagg makes (`disagg.py:257`). QA is told to run the + # accuracy `command` verbatim against the live server -- and this + # campaign's own prompt section opens by saying it never launches + # one. Left in place, the agent either invents a server or invents a + # score. + notes.append( + "accuracy block ignored: a SOL track never stands up a server for you to " + "query, and the harness has no accuracy pass of its own on either track" + ) + + optimize = task_data.get("optimize") + optimize = dict(optimize) if isinstance(optimize, Mapping) else {} + if optimize.get("target_metric") is None: + optimize["target_metric"] = TRACK_METRICS[track] + notes.append( + f"optimize.target_metric={TRACK_METRICS[track]!r}, what the {track} track is scored on" + ) + task_data["optimize"] = optimize + + # No `code_id` note. The previous CLI digested the config and the + # checkout into one fingerprint and reported it on the plan; this one + # does not, so `plan.get("code_id")` was `None` on every spec this has + # ever written -- a provenance line that answered its own question with + # a shrug, in the one place a reader goes to find out what was measured. + # An absent note is better than a note that is always absent. + return notes + + +# --------------------------------------------------------------- the overlay + + +def apply_overlay(task_data: Mapping[str, Any], tuning: Path) -> Path: + """Put the live tuning file where ``bench-disagg`` will read it. + + The one piece of glue that stays ours. ``bench-disagg`` has no notion + of "this campaign's live tuning file" — it reads the stage config it + is handed — while perf-optimize's whole diff / revert / accepted- + snapshot machinery is built around exactly one file the optimizer + edits. So the overlay is copied from that file into the stage config's + role key before every submit. + + Doing it here rather than in a prompt is the same argument as + everything else this campaign moved into code: forgetting it does not + fail, it measures the previous attempt's configuration and books the + result against the new one. + + Returns the stage config it wrote. + """ + track = track_name(task_data) + if track not in TRACKS: + raise SolTrackError(f"unknown sol_track track {track!r}, expected one of {list(TRACKS)}") + sweep_file = sweep_path(task_data) + if sweep_file is None: + raise SolTrackError(f"'{SOL_TRACK_FIELD}.{SWEEP_KEY}' is required") + stage_config = stage_config_path(load_sweep(sweep_file), track, sweep_file) + if stage_config is None or not stage_config.is_file(): + raise SolTrackError(f"the sweep names no readable '{track}' stage config: {stage_config}") + overlay = yaml.safe_load(tuning.read_text(encoding="utf-8")) + if overlay is None: + overlay = {} + if not isinstance(overlay, Mapping): + raise SolTrackError( + f"{tuning} must be a mapping — it is deep-merged onto the worker config " + f"the sweep row generates — got {type(overlay).__name__}" + ) + overreach = sorted(FROZEN_WORKER_KEYS.intersection(overlay)) + if overreach: + raise SolTrackError( + f"{tuning} names {overreach}, which the sweep row fixes for this " + f"campaign. Those keys ARE the operating point, so changing one measures " + f"a different point — a different world size, a different batch — against " + f"a baseline taken at the old one. The comparison would be void rather " + f"than merely weaker, and nothing downstream would notice: the run " + f"succeeds and the number is plausible. Move the point by editing the " + f"sweep row if that is really what you mean, which starts a new campaign." + ) + stage = load_sweep(stage_config) + key = TRACK_OVERLAY_KEYS[track] + if overlay: + stage[key] = dict(overlay) + else: + stage.pop(key, None) + stage_config.write_text( + yaml.safe_dump(stage, sort_keys=False, default_flow_style=False), encoding="utf-8" + ) + return stage_config + + +# --------------------------------------------------------------- the score + + +def target_metric(task_data: Mapping[str, Any]) -> str: + """What this campaign is scored on, and therefore what to write. + + Not ``TRACK_METRICS[track]``. ``apply_plan`` only *defaults* the target + metric, so an owner who names another one keeps it — and every later + stage, the baseline gate first, then looks up that name. Writing the + track's own spelling instead produces a file the gate cannot see, and + the gate's message says the stage measured nothing, which is the one + thing it certainly did. + """ + optimize = task_data.get("optimize") + named = optimize.get("target_metric") if isinstance(optimize, Mapping) else None + if isinstance(named, str) and named.strip(): + return named.strip() + track = track_name(task_data) + if track not in TRACKS: + raise SolTrackError(f"unknown sol_track track {track!r}, expected one of {list(TRACKS)}") + return TRACK_METRICS[track] + + +def elasticity(metrics: Mapping[str, Any], gpus: Mapping[str, Any]) -> float | None: + """How much of a gen improvement survives into the e2e frontier. + + The two numbers a gen campaign holds are not the same objective. The + gate scores ``throughput_per_user``, which is anchor-free; the + deployment is judged on ``output_tput_per_gpu``, which is not:: + + otpg = output_throughput / (ctx_gpus * ctx_per_gen + gen_gpus) + + and ``ctx_per_gen`` rises with the generation side's own speed -- a + faster decode consumes prefill faster, so it needs proportionally more + context GPUs behind it. Differentiating, a 1 % gain in the gate's + metric is worth ``gen_gpus / denominator`` per cent at the frontier, + approaching zero as the context side becomes the wall. + + On the campaign this one is compared against, that ratio was **0.97 at + concurrency 1** and **0.70 at 32** -- the same measured +1 % meaning + materially different things at two ends of one curve, with nothing + saying so. Recorded rather than applied: the gate stays the track's + own metric, and the evaluator is handed the exchange rate. + + Nothing is assumed: the postprocessor reports both GPU counts and the + ratio itself, so all three terms are read rather than re-derived. + """ + ctx_gpus = gpus.get("ctx") + gen_gpus = gpus.get("gen") + ratio = metrics.get("ctx_gen_inst_ratio") + for value in (ctx_gpus, gen_gpus, ratio): + if not isinstance(value, (int, float)) or isinstance(value, bool): + return None + denominator = ctx_gpus * ratio + gen_gpus + if denominator <= 0 or gen_gpus <= 0: + return None + return gen_gpus / denominator + + +def _write_result(path: Path, payload: dict[str, Any]) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + return path + + +def _place(written: dict[int, Path], cases: dict[int, str], concurrency: int, case: str) -> None: + """Refuse two cases at one operating point rather than overwrite one.""" + if concurrency in written: + raise SolTrackError( + f"cases {cases[concurrency]!r} and {case!r} both measure {concurrency} " + f"requests in flight, so they cannot both be " + f"'concurrency_{concurrency}'. A SOL-track campaign's operating points " + f"have to be addressable by concurrency alone, the way an aggregate " + f"campaign's are. Split the sweep so each point sits at its own " + f"concurrency -- for a ctx sweep that usually means one campaign per " + f"`mtp_range` entry, since those share a `max_batch`." + ) + + +#: Where a stage submits, relative to where it collects. One path, so +#: the two cannot be given different answers. +RUN_SUBDIR = "run" + + +def run_dir(into: Path) -> Path: + """The harness run directory whose score belongs in ``into``. + + Submission and collection are derived from one path on purpose. The + run directory the harness creates is named + ``bm_-----`` and **carries no code + identity**, so two attempts of one campaign on one day compute the + same name: submit them into the same work dir and the second + overwrites the first's case directories, with the post-processor then + averaging across whatever survived. Nothing reports that. + + Tying the work dir to the result dir removes the choice. A stage + submits with ``-w /run`` and collects with ``--collect + ``; there is no second path to get wrong, and every + attempt is isolated because its result directory already was. + """ + root = Path(into) / RUN_SUBDIR + found = sorted(p for p in root.glob("bm_*") if p.is_dir()) + if not found: + raise SolTrackError( + f"no bm_* run directory under {root}. A stage submits with " + f"`{IBC_BENCH} submit sweep -c -w {root}` and collects from the " + f"same place; if the submit ran, check that it was given this -w." + ) + if len(found) > 1: + raise SolTrackError( + f"{root} holds {len(found)} run directories " + f"({', '.join(p.name for p in found)}). One per attempt: two of them " + f"means the score cannot say which run it came from." + ) + return found[0] + + +def collect(task_data: Mapping[str, Any], into: Path, *, snapshot: str = "latest") -> list[Path]: + """Write this attempt's score where perf-optimize looks for a result. + + The other half of :func:`apply_overlay`. The overlay is how a + campaign's tuning reaches the harness; this is how the harness' + answer comes back, in the shape every later stage already reads: + ``optimize.target_metric`` in a JSON under ``concurrency_/``. + Without it the flow's own baseline gate rejects a measurement that + succeeded, which looks exactly like the failure it was built to catch. + + The two tracks are read from different files, because the harness + writes them to different files: + + - **gen** — with a declared ctx anchor, the CSV `ibc-bench process + frontier` scores, whose ``throughput_per_user`` column is already + the name `task.yaml` carries; without one, the ``gen_only_perf.csv`` + the anchor-free extractor writes; + - **ctx** — the ``run_*.json`` a ctx case leaves in its work dir, + whose ``performance.request_throughput_req_s`` is the field the + harness itself validates the case on. + + Which gen reader applies is decided by what `task.yaml` **declares**, + not by which file happens to be on disk. A frontier found in a + campaign that named no anchor is refused rather than read: the + anchor's input length is the one thing :func:`require_matching_anchor` + checks, and an undeclared anchor was checked against nothing. + + ``snapshot`` is accepted and ignored; it belonged to a CLI that kept + frontier snapshots, and the signature is kept so a caller that still + passes it does not break. + """ + track = track_name(task_data) + if track not in TRACKS: + raise SolTrackError(f"unknown sol_track track {track!r}, expected one of {list(TRACKS)}") + metric = target_metric(task_data) + directory = run_dir(into) + if track == GEN_TRACK: + return _collect_gen(directory, into, metric, anchored=ctx_json_path(task_data) is not None) + return _collect_ctx(directory, into, metric) + + +#: Why a gen result carries no e2e columns. Written into the result +#: itself, because "this campaign cannot see the frontier" and "this +#: campaign saw a flat frontier" must not read alike to whoever opens the +#: file next -- including the reporter, which quotes it. +NO_E2E_VIEW = ( + "no sol_track.ctx_json: this campaign declares no ctx measurement, so the " + "rate match that turns a decode rate into a deployment number has no other " + "half. output_tput_per_gpu, ctx_gen_inst_ratio and frontier_elasticity are " + "ABSENT, not zero and not unchanged -- nothing here says what a gain at this " + "point is worth end to end. The gate's own metric is unaffected: it is " + "accept_rate / avg_iteration_time, which has no context term." +) + + +def _collect_gen(directory: Path, into: Path, metric: str, *, anchored: bool) -> list[Path]: + written: dict[int, Path] = {} + cases: dict[int, str] = {} + skipped: list[str] = [] + + try: + points = frontier_points(directory) if anchored else gen_only_points(directory) + except BenchCliError as exc: + raise SolTrackError(str(exc)) from exc + + for point in points: + case = str(point.get("case") or "?") + metrics = dict(point.get("metrics") or {}) + value = metrics.get(SNAPSHOT_METRICS[GEN_TRACK]) + # The CSV row carries the point, not the sweep row, so the + # per-generation-server product is already folded in. + concurrency = point.get("concurrency") + if not isinstance(concurrency, int) or not isinstance(value, (int, float)): + skipped.append(case) + continue + _place(written, cases, concurrency, case) + payload: dict[str, Any] = { + # First, and under the name `optimize.target_metric` + # carries: this key is the whole point of the file. + metric: float(value), + "concurrency": concurrency, + "case": case, + "shape": point.get("name"), + "run_dir": str(directory), + "source_csv": point.get("source_csv"), + } + if anchored: + payload["frontier_elasticity"] = elasticity(metrics, point.get("gpus") or {}) + payload["frontier_metrics"] = metrics + else: + payload["e2e_view"] = None + payload["e2e_view_absent"] = NO_E2E_VIEW + payload["gen_only_metrics"] = metrics + written[concurrency] = _write_result( + into / f"concurrency_{concurrency}" / SOL_RESULT_NAME, payload + ) + cases[concurrency] = case + + if not written: + source = "the frontier" if anchored else "the gen-only extractor" + raise SolTrackError( + f"{source} under {directory} scored no usable point" + + (f" (skipped: {', '.join(skipped)})" if skipped else "") + + f". `{IBC_BENCH} jobs check -f {directory}/job_status.csv --summary` " + f"says whether the cases behind it measured." + ) + return [written[key] for key in sorted(written)] + + +#: Where a validated ctx measurement keeps its number. This is the field +#: the harness itself requires before it will call a ctx case successful. +CTX_RESULT_PATH = ("performance", "request_throughput_req_s") + + +def _collect_ctx(directory: Path, into: Path, metric: str) -> list[Path]: + written: dict[int, Path] = {} + cases: dict[int, str] = {} + skipped: list[str] = [] + + for result in sorted(directory.rglob("run_*.json")): + if result.name.endswith("_timing.json"): + continue + value = _read_ctx_result(str(result)) + if value is None: + skipped.append(result.name) + continue + # `/ctx___ratio____MTP_testN/` + # -- the batch is the in-flight request count of a prefill-only run. + concurrency = _ctx_concurrency(result.parent.name) + if concurrency is None: + skipped.append(f"{result.parent.name} (no batch in the dir name)") + continue + _place(written, cases, concurrency, result.parent.name) + written[concurrency] = _write_result( + into / f"concurrency_{concurrency}" / SOL_RESULT_NAME, + { + metric: value, + "concurrency": concurrency, + "case": result.parent.name, + "run_dir": str(directory), + "source_run_json": str(result), + "source_field": ".".join(CTX_RESULT_PATH), + }, + ) + cases[concurrency] = result.parent.name + + if not written: + raise SolTrackError( + f"no validated ctx measurement under {directory}" + + (f" (skipped: {', '.join(skipped)})" if skipped else "") + + f". `{IBC_BENCH} jobs check -f {directory}/job_status.csv --summary` " + f"reports whether the jobs are queued, failed, or produced artifacts " + f"that did not validate." + ) + return [written[key] for key in sorted(written)] + + +def _ctx_concurrency(case_dir: str) -> int | None: + """The in-flight request count from a ctx case directory name. + + `ctx_1024_1_ratio1_16_16640_dep4_MTP3_test2` -- isl, osl, ratio, then + **max_batch**, then max_num_tokens. The batch is the fourth field, and + reading it positionally rather than by pattern is deliberate: a name + this workflow cannot parse should stop the collect rather than pick + whichever number matched. + """ + parts = case_dir.split("_") + if len(parts) < 5 or parts[0] != "ctx": + return None + batch = parts[4] + return int(batch) if batch.isdigit() else None + + +def _read_ctx_result(path: str) -> float | None: + try: + payload = json.loads(Path(path).read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return None + for key in CTX_RESULT_PATH: + payload = payload.get(key) if isinstance(payload, Mapping) else None + if isinstance(payload, (int, float)) and not isinstance(payload, bool): + return float(payload) + return None + + +def _main(argv: list[str] | None = None) -> int: # pragma: no cover - thin entry + import argparse + + parser = argparse.ArgumentParser(description=apply_overlay.__doc__.splitlines()[0]) + parser.add_argument("--workspace", required=True, help="the perf-optimize workspace") + parser.add_argument( + "--collect", + metavar="RESULT_DIR", + default=None, + help="instead of applying the overlay, write the latest frontier " + "snapshot's score into RESULT_DIR/concurrency_/" + SOL_RESULT_NAME, + ) + parser.add_argument( + "--snapshot", + default="latest", + help="with --collect: a snapshot id, 'latest', or 'baseline' (default: latest)", + ) + args = parser.parse_args(argv) + workspace = Path(args.workspace) + task_data = yaml.safe_load((workspace / "task.yaml").read_text(encoding="utf-8")) + if args.collect: + results = collect(task_data, Path(args.collect), snapshot=args.snapshot) + print(json.dumps({"ok": True, "results": [str(path) for path in results]})) + return 0 + written = apply_overlay(task_data, workspace / "tuning" / "extra_llm_api_options.yaml") + print(json.dumps({"ok": True, "stage_config": str(written)})) + return 0 + + +if __name__ == "__main__": # pragma: no cover - entry point + import sys + + try: + sys.exit(_main()) + except (SolTrackError, BenchCliError) as exc: + import sys as _sys + + # `--collect` shells out, so the CLI's own failures surface here + # too. Carrying `error.code` through means the caller can still + # branch on the taxonomy rather than on message text — NO_DATA + # ("nothing measured yet") and ANCHOR_MISSING ("build it against a + # ctx anchor") ask for different next moves, and a traceback that + # collapsed them would send a reader to the wrong one. + report: dict[str, Any] = {"ok": False, "error": str(exc)} + code = getattr(exc, "code", None) + if code: + report["code"] = code + print(json.dumps(report), file=_sys.stderr) + _sys.exit(1) diff --git a/agent-flow/agent_flow/workflows/perf_optimize/spawn.py b/agent-flow/agent_flow/workflows/perf_optimize/spawn.py new file mode 100644 index 000000000000..5888ee2b13c2 --- /dev/null +++ b/agent-flow/agent_flow/workflows/perf_optimize/spawn.py @@ -0,0 +1,161 @@ +"""Starting a sibling campaign, and the only place in this package that may. + +:mod:`.gitops` opens by stating that it holds the package's only +``subprocess`` call, so that nothing can reach a shell without passing one +reviewed door. That invariant is why a campaign cannot quietly run a +benchmark of its own instead of going through ``ibc-bench``, and it is worth +keeping exactly as strict as it is. + +This module is the second door, and it is narrow on purpose: + +- It starts **two kinds of process**, both of which this workflow owns: + ``perf-optimize`` itself (:func:`start`), and the agent that runs the + design skill (:func:`design`). Not a benchmark, not a harness command, + not a shell. For a campaign the argv is built from a + :class:`.disagg_sol.CampaignLaunch`, which was already validated — its + checkout is unshared, its workspace is its own — so by the time anything + gets here there is no decision left to make. +- It never interprets the child's output. A campaign reports through its + workspace, the same way it does when a person starts it; reading a pipe + would make this module a second, worse channel for results. + +**Why a process rather than a call.** The two campaigns could be run in +this process by importing the workflow twice, and that would need no new +door at all. It would also give up the property that made the hand-started +pair safe: separate processes cannot corrupt each other's state, one dying +does not take the other with it, and each is independently attachable and +killable. Those were not incidental in the run this module is modelled on — +the two campaigns landed on nine and one cluster nodes respectively, ran for +three and a half hours, and neither could have observed the other if it had +tried. A shared interpreter would have made that a claim rather than a fact. + +**What is deliberately not here.** No retry, no backoff, no supervision +loop. A campaign that dies leaves a workspace that says how far it got and +resumes from it; a supervisor that restarted it automatically would be +guessing that the failure was transient, and the failures actually observed +in this stack — a corrupted wheel download, a hung fill gate — are not. +""" + +from __future__ import annotations + +import subprocess +from pathlib import Path +from typing import Iterable, Sequence + +import yaml + +from agent_flow.workflows.perf_optimize.disagg_sol import CampaignLaunch, DisaggSolError + + +def materialize(launch: CampaignLaunch) -> Path: + """Write the campaign's ``task.yaml`` and return where it landed. + + Separate from :func:`start` so the spec a campaign will run under is on + disk, and reviewable, before any process exists — and so a dry run is + the same code path minus one call. + """ + launch.workspace.mkdir(parents=True, exist_ok=True) + launch.task_path.write_text( + yaml.safe_dump(launch.spec, sort_keys=False, allow_unicode=True), encoding="utf-8" + ) + return launch.task_path + + +def start(launch: CampaignLaunch, *, env: dict[str, str] | None = None) -> subprocess.Popen: + """Start one campaign and return without waiting for it. + + Not waiting is the point: the two halves are independent under this + module's scope, so serializing them would add hours of wall clock and buy + nothing. The caller holds the handles and decides when to join. + """ + if not launch.task_path.is_file(): + raise DisaggSolError( + f"{launch.task_path} does not exist — call materialize() before start(), " + f"so that what a campaign runs under is on disk before it runs." + ) + return subprocess.Popen( # noqa: S603 - argv is built, never a shell string + launch.argv, + stdout=subprocess.DEVNULL, + stderr=subprocess.STDOUT, + env=env, + start_new_session=True, + ) + + +def start_all( + launches: Iterable[CampaignLaunch], *, env: dict[str, str] | None = None +) -> list[tuple[CampaignLaunch, subprocess.Popen]]: + """Materialize every campaign, then start every campaign. + + In two passes rather than one, so a spec that cannot be written stops the + run before any process exists. Half-started is the one state a supervisor + must not produce: the campaign that did start would claim a checkout the + other was going to need, and the failure would surface as a repo-claim + refusal minutes later, pointing at the wrong cause. + """ + listed = list(launches) + for launch in listed: + materialize(launch) + return [(launch, start(launch, env=env)) for launch in listed] + + +def wait_all(started: Sequence[tuple[CampaignLaunch, subprocess.Popen]]) -> dict[str, int]: + """Join every campaign, returning each track's exit status. + + Every one is waited on even after one fails: the other is still running + on a cluster, and abandoning it would leave an allocation and a claimed + checkout behind with nothing tracking either. + """ + return {launch.track: process.wait() for launch, process in started} + + +#: The agent binary that runs skills. Named here rather than inlined so a +#: site that installs it elsewhere has one place to look. +AGENT = "claude" + + +def design( + instruction: str, *, cwd: Path, timeout: int | None = None, log: Path | None = None +) -> int: + """Run the design skill to completion, in the checkout that ships it. + + ``cwd`` is not incidental: a skill is discovered from the agent's working + directory, so the design agent has to start inside the config repo or it + cannot invoke `create-sweep` at all — the failure being an agent that + improvises the phases from their names. + + Blocking, unlike :func:`start`. The campaigns are independent of each + other and so are started and joined separately; the design is what every + campaign's operating point comes from, so there is nothing to overlap it + with. + + Permissions are bypassed because this runs unattended for hours and the + skill's own submission gates are the review step -- they are what the + instruction tells it to honour. + + **Its output is kept, unlike a campaign's.** A campaign reports through + its workspace, so discarding its pipe loses nothing; a design that fails + to start leaves no workspace at all, and the caller's only symptom is a + later refusal saying the design was never established -- which points at + the design rather than at whatever stopped it. The first real run of this + function returned 1 because the agent's credentials had expired, and the + message saying so went to /dev/null. + """ + completed = subprocess.run( # noqa: S603 - argv is built, never a shell string + [AGENT, "--dangerously-skip-permissions", "--print", instruction], + cwd=str(cwd), + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + if log is not None: + Path(log).parent.mkdir(parents=True, exist_ok=True) + Path(log).write_text((completed.stdout or "") + (completed.stderr or ""), encoding="utf-8") + if completed.returncode != 0: + tail = ((completed.stderr or "") + (completed.stdout or "")).strip().splitlines() + raise DisaggSolError( + f"the design agent exited {completed.returncode} in {cwd}. Its last " + f"output was: " + (" | ".join(tail[-3:]) or "(nothing)") + ) + return completed.returncode diff --git a/agent-flow/agent_flow/workflows/perf_optimize/sweep_design.py b/agent-flow/agent_flow/workflows/perf_optimize/sweep_design.py new file mode 100644 index 000000000000..4929a89ae688 --- /dev/null +++ b/agent-flow/agent_flow/workflows/perf_optimize/sweep_design.py @@ -0,0 +1,144 @@ +"""Asking the design skill for an operating point, on this workflow's terms. + +Establishing *where* to optimize is `create-sweep`'s job — a skill the +benchmark-config repo ships, carrying rules this workflow has no business +restating: which parallelism tiers a mixture-of-experts model admits, where +the tensor-parallel batch ceiling is, what a concurrency ladder looks like on +each, and where its generated configs may be written. Those rules live in +someone else's repo and will move; a copy here would be a second, quietly +diverging answer. + +So this module does not drive the skill. It hands an agent one instruction — +invoke it, follow it, with two named departures — and reads back what the +skill wrote. + +**Two departures, and why only these.** The skill ends its sweep phase with +``ibc-bench process frontier --ctx_json``: the end-to-end join, which this +staged scope excludes. And its later phases prune irreversibly against a +measured MTP accept rate this model does not have. Both are changes of +*meaning*, which the skill cannot know about, so they are stated. The scored +command is given verbatim rather than described because a described one is +the kind an agent honours on the first turn and restates on the fourth — the +session this module was written after did exactly that. + +**What this module is careful not to be.** Three earlier versions of it +overreached, and the shape of each mistake is worth keeping because the same +trade will look attractive again: + +- Four wrappers generated the skill's own commands, on the reasoning that a + generated command cannot be drifted from. But ``SKILL.md`` also carries a + rule they had no place for — for an existing model directory the generated + YAMLs go into ``sweep_design/``, never on top of the curated configs — and + they took the output path as an unguarded parameter. Trading a recomputable + mistake (a frontier scored the wrong way) for an unrecoverable one + (overwriting a config other people maintain) is not a safety improvement. +- One function rewrote the context sweep to widen its candidate list. Never + called, and it would have been this module editing a config it does not own. +- Two checks read the skill's artefacts back and refused inconsistent ones. + Sound in principle and never wired to anything, so they asserted nothing + while looking like they did. + +What is left is an instruction, the command that instruction overrides, and a +YAML reader. **Nothing here writes a file, edits a config, or starts a +process.** +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Mapping, Sequence + +import yaml + +from agent_flow.workflows.perf_optimize.bench_cli import GEN_ONLY_MODULE + + +class SweepDesignError(ValueError): + """The design skill, or something it produced, is not usable.""" + + +# ------------------------------------------------------------ the one override + + +def postprocess_command(run_dir: Path) -> str: + """How a generation run is scored here, and the one command that changed. + + The skill ends Phase 3 with ``ibc-bench process frontier --ctx_json``, + which rate-matches the curve against a context anchor. That is the join, + and the staged scope has none — so the anchor-free extractor is used + instead. It computes the same ``tps_per_user`` from the same iteration + logs; what it does not compute is the deployment view, which is exactly + the part that is out of scope. + """ + return f"python -m {GEN_ONLY_MODULE} -i {run_dir}" + + +def designer_instruction(*, model_dir: str, design_dir: Path, tracks: Sequence[str]) -> str: + """What the design agent is asked to do: run the skill, with one override. + + Deliberately short. The skill's own ``SKILL.md`` carries the phases, the + probe protocol, the submission gates, the resume spine and — the reason + an earlier version of this function was wrong — the rule about where + generated YAMLs may be written. Restating any of that here would produce + a second copy to drift from; the agent reads the skill. + + What this adds is the one thing the skill cannot know: that this + workflow's scope excludes the end-to-end join, so the sweep is scored + anchor-free. That is a change of meaning rather than of path, and it is + given as the exact command because a described one is the kind an agent + honours on the first turn and restates differently on the fourth. + """ + return ( + f"Fix this model's operating point: invoke the **create-sweep** skill and " + f"follow it. It carries the phases, the probe protocol, the submission " + f"gates and the rules about where its generated YAMLs may be written -- " + f"follow those as written, including its rule that for an EXISTING model " + f"directory the generated configs go into `sweep_design/` and never on " + f"top of the curated ones.\n\n" + f"**Run END-TO-END. This instruction is the confirmation the skill's " + f"submission gates ask for.** Do not stop to ask before submitting: you " + f"are running non-interactively and there is no second turn in which an " + f"answer could reach you, so a gate that waits is a gate that ends the " + f"design. Report the case count and node estimate as the skill asks -- " + f"and then submit.\n\n" + f"- model_dir: `{model_dir}` (existing)\n" + f"- design directory: `{design_dir}` -- read its `{DESIGN_STATE_HINT}` " + f"first and resume from whatever phase it records\n" + f"- halves this campaign will optimize afterwards: {list(tracks)}\n\n" + f"**Two departures from the skill, and only these two.**\n\n" + f"**1. Stop after the concurrency sweep.** Do not run the predict/prune " + f"phase or the final frontier. Both require a measured MTP accept rate " + f"this model does not have, and the pruning is irreversible -- it would " + f"throw away points on the strength of a multiplier nobody measured.\n\n" + f"**2. Score the generation runs with this command and no other:**\n" + f"```bash\n{postprocess_command(Path(''))}\n```\n" + f"Do **not** run `ibc-bench process frontier`, and do not pass " + f"`--ctx_json` to anything. This campaign builds no end-to-end view: that " + f"command rate-matches the generation curve against a context anchor, a " + f"measurement this scope does not take. It would still emit a curve, and " + f"the curve would still look correct -- which is why the replacement is " + f"given as a command rather than described.\n\n" + f"**Do not select the operating point.** Your output is the measured " + f"space; choosing from it happens after you, against a stated preference " + f"you have not been given.\n\n" + f"**A shape whose probe or sweep failed is reported as failed** -- never " + f"omitted, because a silently missing shape becomes a frontier it was " + f"never on. Say which failed and why: a job that never started and a " + f"genuine memory wall are different findings." + ) + + +#: Named separately so the prompt and :mod:`.disagg_sol` cannot drift about +#: which file carries the design's resume state. +DESIGN_STATE_HINT = "state.json" + + +def load_yaml(path: Path) -> dict[str, Any]: + """Read a generated file, or say which one could not be read.""" + try: + data = yaml.safe_load(Path(path).read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError) as exc: + raise SweepDesignError(f"could not read {path}: {exc}") from exc + if not isinstance(data, Mapping): + raise SweepDesignError(f"{path} must be a YAML mapping, got {type(data).__name__}") + return dict(data) diff --git a/agent-flow/agent_flow/workflows/perf_optimize/task.example.yaml b/agent-flow/agent_flow/workflows/perf_optimize/task.example.yaml index 593464a5d000..63e46ba50157 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/task.example.yaml +++ b/agent-flow/agent_flow/workflows/perf_optimize/task.example.yaml @@ -207,3 +207,116 @@ profile: # then evolves. # disagg: # config: /path/to/disagg/config.yaml + +# Optional: optimize ONE HALF of a disaggregated deployment in isolation — a +# SOL track. Mutually exclusive with `disagg` above: that measures the whole +# cluster end-to-end, this measures one role, over the same tuning knobs. +# +# The campaign drives `ibc-bench`, the CLI the harness repo ships as its only +# supported agent-facing interface. Everything about *how* a measurement +# happens — submitting, re-measuring rather than skipping, naming a run, +# recording which image and config produced it — belongs to that tool. This +# block only says which sweep, which half, and which workspace. +# +# track ctx | gen. `ctx` is prefill-only aggregate (trtllm-bench at +# output_length 1 — no disaggregation at all), scored on +# `avg_request_throughput_req_s`. `gen` is a 1-ctx-1-gen +# deployment read only on the generation worker's decode +# iterations, scored on `throughput_per_user`. +# sweep the harness sweep yaml. Paths inside it resolve relative to +# itself, and the harness imports the model's +# `gen_worker_config.py` from beside it — so the unit is the +# DIRECTORY, not the file. It is copied into `/sweep/` +# at startup and the campaign measures the copy; the original is +# never rewritten by the per-attempt tuning overlay. +# ctx_json OPTIONAL, and what it decides is which QUESTION the campaign can +# answer — not whether it can be scored. A gen campaign is scored +# anchor-free, from the `gen_only_perf.csv` that +# `ibc_trtllm_harness.process_data.get_gen_only_perf` writes: +# `throughput_per_user = accept_rate / avg_iteration_time`, which +# has no context term at all. Naming a ctx.json additionally buys +# the rate-matched end-to-end view. WITHOUT it, +# `output_tput_per_gpu`, the ctx:gen ratio and the frontier +# elasticity are ABSENT — not zero, not unchanged — and nothing +# the campaign reports says what a gain is worth end to end. +# workspace the harness workspace name. It fixes one workload on one +# cluster; a new image, build or tuning overlay is not a change. +# +# `accept_rate` lives in the SWEEP, at its TOP LEVEL (`accept_rate: {rate, +# source}`, where the harness has read it since v0.4.7) — and it is required +# only when the sweep actually has `mtp > 0` cases. At mtp 0 the acceptance +# length is 1.0 by definition, so requiring one would ask for a measurement of +# a constant; the harness's own post-processor draws the line in the same +# place. Where it IS required, it is because the length multiplies the metric +# linearly and a wrong one is invisible. +# +# `benchmark` fields you omit are filled from the sweep's own plan — the same +# expansion that will run — and any you set that disagree are an error. Note +# that a sweep row's concurrency is per generation server, so +# `benchmark.concurrency` is `concurrency x gen_num`, while the row keeps the +# listed value. Profiling is nsys-only. `optimize.target_metric` defaults to +# the track's. Leave `optimize.max_regression_pct` unset and every point can +# veto an attempt, which is the gate a frozen-point campaign wants. +# +# The optimizer edits `/tuning/extra_llm_api_options.yaml`, which is +# a PARTIAL overlay deep-merged onto the worker config the sweep row generates +# — so the row keeps the topology and the optimizer only carries what it tunes. +# sol_track: +# track: gen +# sweep: /path/to/campaign/sweep.yaml +# workspace: dsv4pro-8k1k-gen +# # ctx_json: /path/to/a/measured/ctx.json # omit for an anchor-free gen +# # # campaign; see above for what +# # # that makes ABSENT +# +# ── BOTH halves: `disagg_sol` ──────────────────────────────────────────────── +# +# One file starts two campaigns — a ctx one and a gen one — at an operating +# point SELECTED from a design that was measured once. Use this instead of +# writing two `sol_track` specs by hand: those inherit their operating point +# from whoever wrote them, and this is the layer that exists to remove that. +# +# There is deliberately NO `benchmark:` block and no operating point here. The +# point is not knowable until the design has been measured, so each half's +# sweep is cut from the design's own sweep afterwards. +# +# The two halves are NEVER combined into one number: they are measured on +# different metrics over different denominators and no measurement joins them. +# Every selection result carries `e2e_view_absent` saying so. +# +# `slurm-environment` above (if present) is inherited by BOTH halves, which is +# what makes off-cluster execution work: `cluster_ssh` and `remote_run_root` +# ARE remote execution, and a half that did not receive them would silently run +# as though the flow were on the cluster. +# +# perf-optimize --task task.yaml --workspace ws # start both halves +# perf-optimize --task task.yaml --workspace ws --dry-run # select, start nothing +# +# disagg_sol: +# tracks: [ctx, gen] +# design: +# design_dir: /path/to/model/sweep_design # measured once per +# # (model, cluster, workload) +# design_sweep: # two, because the halves are +# ctx: /path/to/model/sweep_design/ctx_config.yaml # measured by +# gen: /path/to/model/sweep_design/8k1k_sol_mtp0.yaml # different sweeps +# prefer: interactive # or `throughput`. No default: a +# # curve states the trade-off and +# # cannot state what it is bought for +# repos: # ONE CHECKOUT PER HALF, and they +# ctx: /path/to/trtllm-ctx # may not be shared -- the flow +# gen: /path/to/trtllm-gen # resets a checkout hard and the +# # evaluator reviews `git diff` on +# # it, so two campaigns on one tree +# # measure fine and read wrong +# # config_repo: /path/to/bench-trtllm-disagg # where the create-sweep skill +# # # lives. Given, and with no design +# # # established, the run measures one +# # # first. Omitted, an unestablished +# # # design is refused rather than +# # # silently fallen back from. +# # incumbent: {shape: tep_4_eplb0_mtp3, concurrency: 32} +# # # what a previous deployment ran at, +# # # so the record says whether the +# # # selection confirmed it or moved off +# # label: mycampaign # names the two child workspaces diff --git a/agent-flow/agent_flow/workflows/perf_optimize/task_schema.py b/agent-flow/agent_flow/workflows/perf_optimize/task_schema.py index 0c51c4aea75e..40adae228ad8 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/task_schema.py +++ b/agent-flow/agent_flow/workflows/perf_optimize/task_schema.py @@ -76,6 +76,8 @@ from agent_flow.workflows.perf_analyze.task_schema import ( load_and_validate_task_yaml as _base_load_and_validate, ) +from agent_flow.workflows.perf_optimize.bench_cli import BenchCliError +from agent_flow.workflows.perf_optimize.bench_cli import plan as sweep_plan from agent_flow.workflows.perf_optimize.disagg import ( DISAGG_CONFIG_KEY, DISAGG_FIELD, @@ -87,6 +89,26 @@ user_set_benchmark_keys, ) from agent_flow.workflows.perf_optimize.roadmap_schema import APPROACHES +from agent_flow.workflows.perf_optimize.sol_track import ( + CTX_JSON_KEY, + GEN_TRACK, + SOL_TRACK_FIELD, + SWEEP_KEY, + TRACK_KEY, + TRACK_METRICS, + TRACKS, + WORKSPACE_KEY, + SolTrackError, + apply_plan, + ctx_json_path, + has_sol_track, + load_sweep, + require_build_source, + sol_track_block, + sweep_accept_rate, + sweep_path, + track_name, +) # Defaults merged under the user's values. ``target_improvement_pct`` is # deliberately absent: when the user does not set it, there is no @@ -138,6 +160,14 @@ for stat in ("mean", "median", "p90", "p99") for kind in ("ttft", "tpot", "itl", "e2el") } + # A SOL track's metric is as real a result key as any above: it is what + # `sol_track.collect` writes into the result JSON every later stage + # reads. Absent from here the two importers refuse it -- the service + # adapter raises and the lint errors -- so a campaign that runs + # perfectly from the CLI cannot be submitted through the dashboard, + # and the message it gets says the key is not a benchmark result key, + # which is the one thing it certainly is. + | set(TRACK_METRICS.values()) ) # The perf-optimize half of the key census the base schema documents. Same @@ -409,6 +439,153 @@ def _validate_disagg_block(data: dict[str, Any], errors: list[str]) -> dict[str, return None +def _validate_sol_track_block(data: dict[str, Any], errors: list[str]) -> dict[str, Any] | None: + """Validate the ``sol_track`` block and ask the CLI what its sweep expands to. + + Same boundary and same reason as :func:`_validate_disagg_block`: a + sweep that cannot be read, or that plans none of the cases this + campaign says it optimizes, must abort before any agent is + constructed rather than an hour into an allocation. + + Returns the ``sweep plan`` envelope's data block — the authority for + the operating points and sequence lengths — or ``None`` when the + block is absent or unusable. ``sweep plan`` is read-only and queues + nothing, which is what makes it usable from here. + """ + if not has_sol_track(data): + return None + block = sol_track_block(data) + if block is None: + errors.append( + f"'{SOL_TRACK_FIELD}' must be a mapping carrying '{TRACK_KEY}' " + f"({' | '.join(TRACKS)}), '{SWEEP_KEY}' and '{WORKSPACE_KEY}', got " + f"{type(data.get(SOL_TRACK_FIELD)).__name__}" + ) + return None + if has_disagg(data): + # One campaign measures one thing. A sol_track campaign optimizes + # one half in isolation; a disagg campaign measures the whole + # cluster. Both reconcile `benchmark` from a different file, so + # combining them means one of the two silently loses. + errors.append( + f"'{SOL_TRACK_FIELD}' cannot be combined with '{DISAGG_FIELD}': a sol_track " + f"campaign optimizes one role in isolation against its own sweep, while a " + f"disagg campaign measures the end-to-end deployment. Run them as separate " + f"campaigns." + ) + return None + if data.get(EXTRA_LLM_API_OPTIONS_FIELD) is not None: + # Same shape as the disagg refusal, and for the same reason: two + # seeds for one live tuning file. A SOL track seeds it from the + # sweep stage's own overlay key, so a named extra_llm_api_options + # would be dropped on the floor -- `workflow.py` reaches the + # `has_sol_track` branch before the `elif extra` one. "My setting + # did nothing" is the failure this codebase refuses to ship. + errors.append( + f"'{EXTRA_LLM_API_OPTIONS_FIELD}' cannot be combined with " + f"'{SOL_TRACK_FIELD}': the live tuning config is seeded from the sweep " + f"stage's own '{{ctx,gen}}_extra_llm_api' overlay, so this file would " + f"never be read. Put its contents in the sweep stage config instead." + ) + return None + track = track_name(data) + if track not in TRACKS: + errors.append( + f"'{SOL_TRACK_FIELD}.{TRACK_KEY}' must be one of {list(TRACKS)}, got " + f"{block.get(TRACK_KEY)!r}" + ) + return None + # `workspace` is optional now, and no longer where results are read + # from: a stage submits into `/run` so submission and + # collection cannot be given different answers. Kept because task.yaml + # files carry it, and a harness work dir root is still a useful thing + # to state. + path = sweep_path(data) + if path is None: + errors.append( + f"'{SOL_TRACK_FIELD}.{SWEEP_KEY}' is required and must be a non-empty " + f"string: the orchestration sweep.yaml naming the stage configs and the " + f"cluster server.config" + ) + return None + if not path.is_file(): + errors.append(f"'{SOL_TRACK_FIELD}.{SWEEP_KEY}' is not a file: {path}") + return None + try: + sweep = load_sweep(path) + except SolTrackError as exc: + errors.append(str(exc)) + return None + # `frontier build` requires it on every build and refuses to infer + # one, so a sweep without it produces a campaign that measures fine + # and cannot be turned into a curve. Cheaper to say so now. + # Only when a row actually uses speculation. The acceptance length is a + # multiplier on the metric, so a wrong one is invisible -- but at mtp 0 it + # is 1.0 by definition and there is nothing to measure. The harness draws + # the same line: its post-processor exits 1 for a missing rate only when + # the sweep has mtp>0 cases. Requiring one here for an mtp0 design sweep + # asks for a measurement of a constant. + if track == GEN_TRACK and _uses_speculation(sweep) and sweep_accept_rate(sweep) is None: + errors.append( + f"{path} sets no top-level 'accept_rate'. Every `process frontier` requires " + f"it and none is inferred: the acceptance length scales both the numerator " + f"and the ctx term of the frontier metric, so a wrong one tilts the whole " + f"curve with no symptom. Freeze the measured value at the sweep's TOP " + f"LEVEL -- `accept_rate: {{rate: , source: }}` -- which is " + f"where the harness has read it since v0.4.7, and where " + f"`sweep_accept_rate` looks." + ) + return None + if track == GEN_TRACK: + # Optional, and the choice is which QUESTION the campaign can + # answer -- not whether it can be scored. + # + # The gate's metric is `accept_rate / avg_iteration_time`: decode + # iterations only, no context term, so a ctx measurement cannot + # move it. What the anchor buys is the e2e half -- the rate match + # that turns a decode rate into `output_tput_per_gpu` and hence + # into `frontier_elasticity`, the exchange rate saying what a gain + # at this point is worth to the deployment. + # + # It was required here once, because the score was read out of the + # frontier CSV, which `process frontier` will not write without an + # anchor. That made a decode campaign wait on somebody's prefill + # run to score a change prefill cannot affect. `get_gen_only_perf` + # computes the same column from the same iteration logs with no + # anchor at all, so the dependency is gone and the anchor is back + # to meaning what it always meant. + anchor = ctx_json_path(data) + if anchor is not None and not anchor.is_file(): + errors.append(f"'{SOL_TRACK_FIELD}.{CTX_JSON_KEY}' is not a file: {anchor}") + return None + try: + # The expansion, from the sweep this campaign will actually + # submit. Read rather than asked of the harness: `ibc-bench` has + # no read-only "what would this submit" command, and the one that + # answers -- `submit sweep --dry-run` -- answers by writing a + # directory per case, which schema validation may not do. The + # agent's own dry-run step re-derives it against the harness + # before an allocation is spent. + return sweep_plan(sweep) + except BenchCliError as exc: + errors.append(f"could not expand {path}: {exc}") + return None + + +def _uses_speculation(sweep: dict) -> bool: + """Whether any row of this sweep runs with a draft length above zero.""" + for row in sweep.get("gen_configs") or []: + if isinstance(row, Mapping): + mtp = row.get("gen_mtp_size") + elif isinstance(row, (list, tuple)) and len(row) > 7: + mtp = row[7] + else: + continue + if isinstance(mtp, int) and not isinstance(mtp, bool) and mtp > 0: + return True + return False + + def load_and_validate_task_yaml( path: str | Path, *, max_rounds_override: int | None = None ) -> dict[str, Any]: @@ -439,8 +616,34 @@ def load_and_validate_task_yaml( } except DisaggConfigError as exc: errors.append(str(exc)) + # Same slot and the same reason for a sol_track campaign: its sweep + # config owns the operating points, and `optimize.target_metric` + # defaults to what the track's post-processor emits — both have to + # land before the blocks validated against them, and before the + # OPTIMIZE_DEFAULTS merge below. + sol_track_cfg = _validate_sol_track_block(data, errors) + if sol_track_cfg is not None: + try: + data[SOL_TRACK_FIELD] = { + **data[SOL_TRACK_FIELD], + "filled_from_sweep_plan": apply_plan( + data, sol_track_cfg, user_set_benchmark_keys(path) + ), + } + except SolTrackError as exc: + errors.append(str(exc)) optimize = _mapping_block(data, "optimize", errors) _validate_optimize_block(optimize, errors) + if sol_track_cfg is not None and isinstance(optimize, Mapping): + # After the optimize block, because it is what names the + # approaches -- and before any agent, because the failure it + # prevents costs a full allocation and reads as a real result. + try: + require_build_source( + data, optimize.get("approaches") or OPTIMIZE_DEFAULTS["approaches"] + ) + except SolTrackError as exc: + errors.append(str(exc)) # An explicitly-null value means "not set" everywhere in this # validator (every check above skips ``None``), so drop those keys # before the defaults merge too — otherwise a bare ``max_rounds:`` diff --git a/agent-flow/agent_flow/workflows/perf_optimize/workflow.py b/agent-flow/agent_flow/workflows/perf_optimize/workflow.py index f5e12459f916..a6bac56888db 100644 --- a/agent-flow/agent_flow/workflows/perf_optimize/workflow.py +++ b/agent-flow/agent_flow/workflows/perf_optimize/workflow.py @@ -45,6 +45,16 @@ ) from .prompts import DEFAULT_PROMPTS, PromptBundle from .roadmap_schema import RoadmapError +from .sol_track import ( + GEN_TRACK, + WORKSPACE_SWEEP_DIR, + adopt_sweep, + ctx_json_path, + has_sol_track, + sweep_path, + track_name, + tuning_seed_yaml, +) from .state import ( ROUND_STAGES, STAGE_ANALYZER, @@ -285,6 +295,15 @@ def __init__( self.final_verification_dir, self.reuse_dir, self.sol_work_dir, + # The adopted sweep copy. `adopt_sweep` leaves an existing + # copy alone -- which is what a RESUME needs, and what makes + # `--clean` the only thing that can restore the original. + # Left here, `apply_overlay` has already written the previous + # campaign's accepted tuning into it, so the "fresh" run would + # seed its baseline from the last accepted optimization and + # book the result as the sweep's own: exactly the failure + # copying into the workspace was introduced to prevent. + self.workspace / WORKSPACE_SWEEP_DIR, ): shutil.rmtree(directory, ignore_errors=True) @@ -573,6 +592,7 @@ def run(self, task: str) -> None: state.reporter_done = True state.done = True self._checkpoint(state) + self._release_repo(state) print_message( f"[bold green]✔ optimization report written to {self.report_path}[/bold green]", log, @@ -645,12 +665,18 @@ def _init_state(self, task: str, log) -> WorkflowState | None: task, max_rounds_override=self.max_rounds_override, ) + # Before the spec is materialized, so the task.yaml every agent + # reads already names the copy this campaign will edit rather than + # the file the user wrote. + if has_sol_track(task_data): + adopt_sweep(task_data, self.workspace) self.task_path.write_text(dump_task_yaml(task_data), encoding="utf-8") # Materialize the live tuning config (the single # --extra_llm_api_options every serve in this workflow uses) and # its last-accepted snapshot. In a disagg campaign the same file - # holds the harness config's ctx / gen worker_config instead, so + # holds the harness config's ctx / gen worker_config instead, and + # in a SOL-track campaign it holds that track's single role — so # the optimizer still edits exactly one file and the diff / # revert / accepted-snapshot machinery applies unchanged. self.tuning_dir.mkdir(parents=True, exist_ok=True) @@ -660,6 +686,13 @@ def _init_state(self, task: str, log) -> WorkflowState | None: self.tuning_config_path.write_text( worker_config_yaml(load_disagg_config(disagg_config)), encoding="utf-8" ) + elif has_sol_track(task_data): + # A SOL track's tuning file is an *overlay*, not a whole role + # config: `bench-disagg` deep-merges it onto the worker config + # its sweep row generated. So it starts as whatever override + # the sweep already carried — usually nothing — and the + # topology the row owns stays out of the optimizer's reach. + self.tuning_config_path.write_text(tuning_seed_yaml(task_data), encoding="utf-8") elif extra: shutil.copyfile(extra, self.tuning_config_path) else: @@ -772,6 +805,7 @@ def _ensure_optimization_branch(self, state: WorkflowState, log) -> None: f"needs git to commit accepted optimizations and revert rejected " f"ones — clone the checkout with git and retry." ) + self._require_unclaimed_repo(repo) timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") state.campaign_git_branch = f"perf-optimize/{self.workspace.name}-{timestamp}" state.campaign_git_base_commit = gitops.rev_parse_head(repo) @@ -1277,7 +1311,7 @@ def _integrate_batch(self, state: WorkflowState, log) -> None: clear_stale_benchmark_results(integration_dir) self._stamp_progress(state, round_no=round_no) self.integrator( - self._disagg_directive() + f"Workspace: {self.workspace}\n" + self._campaign_directive() + f"Workspace: {self.workspace}\n" f"Round: {round_no}\n" f"Integration worktree: {state.integration_worktree_path}\n" f"Integration branch: {state.integration_branch}\n" @@ -1537,32 +1571,150 @@ def _is_nonempty(path: Path) -> bool: """True iff ``path`` exists and holds non-whitespace content.""" return path.is_file() and bool(path.read_text(encoding="utf-8").strip()) - def _disagg_directive(self) -> str: - """The disagg override every stage prompt opens with, or ``""``. + def _campaign_directive(self) -> str: + """The campaign-mode override every stage prompt opens with, or ``""``. The role prompts are composed of two layers: a system prompt built from shared fragments, and the per-stage instruction this - orchestrator writes. ``DISAGG_CAMPAIGN`` supersedes the - single-server guidance in the *first* layer — but the second layer - also names ``trtllm-serve``, ``--extra_llm_api_options`` and a - readiness poll, and it arrives last and reads as the more specific - of the two. Without this the agent is handed a contradiction and - the disagg section can lose on specificity. + orchestrator writes. ``DISAGG_CAMPAIGN`` and the SOL-track + sections supersede the single-server guidance in the *first* + layer — but the second layer also names ``trtllm-serve``, + ``--extra_llm_api_options`` and a readiness poll, and it arrives + last and reads as the more specific of the two. Without this the + agent is handed a contradiction and the overriding section can + lose on specificity. So the stage prompt states the mode up front and points at the section that governs it, rather than every stage's instruction - growing a disagg variant of its own. + growing a variant per campaign mode. """ - if not has_disagg(self._task_data()): - return "" - config = disagg_config_path(self._task_data()) - return ( - f"⚠️ **This campaign is DISAGGREGATED** (harness config: `{config}`). " - f"Nothing below that mentions `trtllm-serve`, `--extra_llm_api_options` " - f"or polling a server applies — your system prompt's " - f"*Disaggregated serving* section replaces all of it.\n\n" + task_data = self._task_data() + if has_disagg(task_data): + config = disagg_config_path(task_data) + return ( + f"⚠️ **This campaign is DISAGGREGATED** (harness config: `{config}`). " + f"Nothing below that mentions `trtllm-serve`, `--extra_llm_api_options` " + f"or polling a server applies — your system prompt's " + f"*Disaggregated serving* section replaces all of it.\n\n" + ) + if has_sol_track(task_data): + track = str(track_name(task_data)) + config = sweep_path(task_data) + directive = ( + f"⚠️ **This campaign is a {track.upper()} TRACK** — one half of a " + f"disaggregated deployment, measured in isolation (sweep: " + f"`{config}`). Nothing below that mentions `trtllm-serve`, " + f"`--extra_llm_api_options` or polling a server applies — your system " + f"prompt's *{track.upper()} track* section replaces all of it.\n\n" + ) + anchor = ctx_json_path(task_data) + if anchor is not None: + # Not discoverable: the campaign measures no ctx stage, so + # the postprocessor would refuse without being told where + # the rate-match's other half comes from. + directive += ( + f"This campaign has no ctx stage, so every " + f"`ibc-bench process frontier` must carry " + f"`--ctx_json {anchor}`.\n\n" + ) + elif track == GEN_TRACK: + # Said here too, not only in the resolved spec: this is the + # stage that will be tempted to reach for a frontier, and + # borrowing an undeclared anchor is the one mistake whose + # output looks exactly like a correct one. + directive += ( + "This campaign declares **no ctx anchor**, so it scores through " + "`get_gen_only_perf`, not `process frontier`. The gate's metric is " + "unaffected — it has no context term — but there is no end-to-end " + "view here: `output_tput_per_gpu` and `frontier_elasticity` are " + "absent, and no number this campaign produces may be quoted as an " + "end-to-end result. Do not supply an anchor of your own to get " + "one.\n\n" + ) + return directive + return "" + + #: A campaign's branch is named for the workspace that owns it, which + #: makes the branch a claim on the checkout without any new bookkeeping + #: -- and bookkeeping is what this cannot have, since `gitops` may be + #: running every command over ssh. + BRANCH_PREFIX = "perf-optimize/" + + def _require_unclaimed_repo(self, repo: str) -> None: + """Refuse a checkout another campaign is already optimizing on. + + The flow resets the checkout hard and branches from it. Two + campaigns pointed at one checkout therefore stomp each other: the + second resets the first's worktree mid-attempt, and the first's + evaluator -- which reads `git status` / `git diff` / `git log` to + review what the optimizer changed -- reviews the wrong tree. With + `approach: code` it is worse than a bad review, because the wheel + built from that tree is what gets measured. + + This matters most for a disaggregated deployment optimized in + halves: the ctx and gen campaigns are independent by design and are + *meant* to run at the same time, so a shared checkout is exactly + the arrangement someone will reach for. It measures fine and reads + wrong, which is this codebase's least favourite shape of bug. + + The branch is the claim. It already carries the owning workspace's + name, so no lock file is needed -- which is the point, because + `gitops` may be talking to a remote host where this process cannot + write files. + """ + try: + branch = gitops.current_branch(repo) + except gitops.GitOpsError: + return + mine = f"{self.BRANCH_PREFIX}{self.workspace.name}-" + if not branch.startswith(self.BRANCH_PREFIX) or branch.startswith(mine): + return + owner = branch[len(self.BRANCH_PREFIX) :].rsplit("-", 2)[0] + raise RuntimeError( + f"{repo} is claimed by another campaign: it is on branch {branch!r}, " + f"which belongs to workspace {owner!r}, not {self.workspace.name!r}. " + f"Give each campaign its own checkout -- they are cheap, and a " + f"disaggregated deployment optimized in halves is meant to run two at " + f"once:\n" + f" git -C {repo} worktree add ../{self.workspace.name}-trtllm \n" + f"then point this campaign's `trtllm_repo_path` at that path. If " + f"{owner!r} has finished and you mean to reuse this one, release it " + f"with `git -C {repo} checkout ` and re-run." ) + def _release_repo(self, state: Any) -> None: + """Hand the checkout back when the campaign is over. + + The branch is this campaign's claim on the checkout, which is what + lets `_require_unclaimed_repo` refuse a second campaign without a + lock file. But a claim nothing releases is indistinguishable from a + live one, so a *finished* campaign would block the next -- which it + did, on the very first run after the guard landed. + + Only when the campaign committed nothing, though. A campaign that + accepted something leaves its work on that branch and is expected + to still be sitting on it -- that is where a reader goes to see + what was accepted, and detaching would hide it. A campaign with + zero accepts has nothing to show and no reason to keep the + checkout, which is the case that was blocking. + + A campaign that crashed also keeps its claim, and that is right: + its checkout is in an unknown state and should not be silently + reused. + """ + repo = self._trtllm_repo_path() + if not repo or not state.campaign_git_base_commit: + return + try: + if gitops.rev_parse_head(repo) != state.campaign_git_base_commit: + return # it committed something; leave it on display + gitops.checkout(repo, state.campaign_git_base_commit) + except gitops.GitOpsError: + # Uncommitted work, a missing commit, anything: the campaign is + # done and its results are written. Failing here would turn a + # finished run into a failed one over bookkeeping. + pass + def _require_baseline_measurement(self) -> None: """Fail loudly when the baseline stage produced no measurement. @@ -2133,7 +2285,7 @@ def _run_benchmarker(self, state: WorkflowState) -> None: "the roadmap's `baseline.value`. " ) self.benchmarker( - self._disagg_directive() + f"Workspace: {self.workspace}\n\n" + self._campaign_directive() + f"Workspace: {self.workspace}\n\n" f"Read `{self.task_path}` for the spec — resolve `checkpoint_path`, " f"`trtllm_repo_path`, and the `benchmark` / `optimize` blocks.\n\n" f"Then **load the `perf-optimization-casebook` skill** (via the " @@ -2255,7 +2407,7 @@ def _run_reused_analyzer(self, state: WorkflowState) -> None: f"`{analysis_dir}` — do not re-derive it.\n\n" ) self.analyzer( - self._disagg_directive() + f"Workspace: {self.workspace}\n" + self._campaign_directive() + f"Workspace: {self.workspace}\n" f"Round: 1 (**reused analysis** — no profiling this round)\n" f"Analysis directory (already populated): {analysis_dir}\n\n" f"This campaign was launched with " @@ -2417,7 +2569,7 @@ def _run_analyzer(self, state: WorkflowState) -> None: f"cannot be closed in this campaign.\n\n" ) self.analyzer( - self._disagg_directive() + f"Workspace: {self.workspace}\n" + self._campaign_directive() + f"Workspace: {self.workspace}\n" f"Round: {round_no}\n" f"Analysis directory (write your artifacts here): {analysis_dir}\n\n" f"Read `{self.task_path}` and `{self.baseline_results_path}` to " @@ -2537,7 +2689,7 @@ def _run_replan_analyzer(self, state: WorkflowState) -> None: f"campaign.\n\n" ) self.analyzer( - self._disagg_directive() + f"Workspace: {self.workspace}\n" + self._campaign_directive() + f"Workspace: {self.workspace}\n" f"Round: {round_no} (**replan only** — no profiling this round)\n" f"Analysis directory (write your artifacts here): {analysis_dir}\n\n" f"Round {state.round_index} accepted **nothing**. " @@ -2664,7 +2816,7 @@ def _run_optimizer( f"your summary, not a claim to re-assert.\n\n" ) (agent or self.optimizer)( - self._disagg_directive() + f"Workspace: {self.workspace}\n" + self._campaign_directive() + f"Workspace: {self.workspace}\n" f"Round: {round_no} — item {state.item_index + 1} of at most " f"{state.max_items_per_round} this round — attempt {attempt_no} " f"of {state.max_attempts_per_item}\n" @@ -2862,7 +3014,7 @@ def _run_evaluator( else: attempt_note = "" (agent or self.evaluator)( - self._disagg_directive() + f"Workspace: {self.workspace}\n" + self._campaign_directive() + f"Workspace: {self.workspace}\n" f"Round: {round_no} — item {state.item_index + 1} of at most " f"{state.max_items_per_round} this round — attempt {attempt_no} " f"of {state.max_attempts_per_item}\n" @@ -2973,7 +3125,7 @@ def _run_qa(self, state: WorkflowState) -> None: "`cumulative_improvement_pct` — from your own measurement" ) self.qa( - self._disagg_directive() + f"Workspace: {self.workspace}\n" + self._campaign_directive() + f"Workspace: {self.workspace}\n" f"Campaign: the optimization loop is over ({state.round_index} " f"round(s) ran); the system under test is the final accepted " f"state.\n" diff --git a/agent-flow/tests/workflows/perf_optimize/test_bench_cli.py b/agent-flow/tests/workflows/perf_optimize/test_bench_cli.py new file mode 100644 index 000000000000..3f4136f413ef --- /dev/null +++ b/agent-flow/tests/workflows/perf_optimize/test_bench_cli.py @@ -0,0 +1,228 @@ +"""Tests for reading the benchmark suite's configs and results. + +This module runs nothing, so there is nothing to stub: every test writes +the file the harness would write and checks what is made of it. That is +also the only way the shape can be caught drifting — a stub agrees with +whatever it was taught, and the sweep row order and the CSV column names +are exactly the things that move upstream. +""" + +from __future__ import annotations + +import pytest + +from agent_flow.workflows.perf_optimize import bench_cli + +#: A gen sweep row, in the harness' order: +#: [ctx, gen, tp, batch, max_num_tokens, attention_dp, gpu_mem_frac, mtp, +#: eplb, concurrency_list] +SWEEP = { + "model_id": "deepseek-ai/DeepSeek-V4-Pro", + "model_path": "/models/DeepSeek-V4-Pro", + "dataset_file": "/data/DeepSeek-V4-8192-1024-200000-ratio-08_for_serve.json", + "precision": "fp4", + "benchmark_client": "trtllm", + "isl": 8192, + "osl": 1024, + "gen_configs": [[1, 1, 4, 64, 256, False, "0.9", 3, 0, "1,32"]], +} + + +def test_a_row_expands_over_its_concurrency_list(): + cases = bench_cli.gen_cases(SWEEP) + assert [c["config"]["concurrency"] for c in cases] == [1, 32] + assert all(c["stage"] == "gen" for c in cases) + + +def test_the_shape_name_matches_the_column_the_postprocessor_writes(): + """`{dep|tep}_{tp}_eplb{N}_mtp{M}`. + + A scored row is addressed by that name plus its concurrency, so the + name this module derives has to be the one the CSV carries or an + attempt cannot be aligned with its baseline. + """ + assert bench_cli.gen_cases(SWEEP)[0]["name"] == "tep_4_eplb0_mtp3" + dep = {**SWEEP, "gen_configs": [[1, 1, 32, 128, 512, True, "0.7", 1, 384, "512"]]} + assert bench_cli.gen_cases(dep)[0]["name"] == "dep_32_eplb384_mtp1" + + +def test_the_dict_row_form_is_accepted_too(): + """The harness takes both; a sweep written either way must expand.""" + rows = { + **SWEEP, + "gen_configs": [ + { + "ctx_num": 1, + "gen_num": 1, + "gen_tp_size": 8, + "gen_batch_size": 128, + "gen_max_num_tokens": 512, + "gen_enable_attention_dp": False, + "gen_gpu_memory_fraction": "0.9", + "gen_mtp_size": 3, + "gen_eplb_num_slots": 0, + "gen_concurrency_list": "4,8", + } + ], + } + cases = bench_cli.gen_cases(rows) + assert [c["name"] for c in cases] == ["tep_8_eplb0_mtp3"] * 2 + assert [c["config"]["concurrency"] for c in cases] == [4, 8] + + +def test_the_operating_point_is_the_product_not_the_listed_value(): + """A row's concurrency is per generation server. + + The client is driven at `concurrency * gen_num` and the harness names + the result directory after that product, so that is the point. + """ + assert bench_cli.operating_point({"concurrency": 64, "gen_num": 2}) == 128 + assert bench_cli.operating_point({"concurrency": 64}) == 64 + + +def test_a_ctx_case_is_addressed_by_max_batch(): + """It has no concurrency; `max_batch` is the in-flight count.""" + assert bench_cli.operating_point({"isl": 1024, "max_batch": 16, "tp_size": 4}) == 16 + assert bench_cli.operating_point({"isl": 1024}) is None + + +def test_a_ctx_block_expands_over_batch_tp_and_mtp(): + cases = bench_cli.ctx_cases( + { + "benchmarks": [ + { + "isl": 1024, + "osl": 1, + "max_batch": [8, 16], + "tp_size": [4], + "ratio": [1], + "mtp_range": [0, 3], + }, + ] + } + ) + assert len(cases) == 4 + assert {c["config"]["max_batch"] for c in cases} == {8, 16} + assert all(c["stage"] == "ctx" for c in cases) + + +def test_isl_is_reported_as_a_bound_beside_the_corpus_not_instead_of_it(): + """The checked-in 8k sweep pairs isl 8192 with a 200000-sample corpus. + + Both numbers are right: one bounds the KV allocation and the client's + request, the other is what actually gets served. + """ + load = bench_cli.workload(SWEEP) + assert load["isl"] == 8192 + assert load["dataset"].endswith("_for_serve.json") + + +#: A ctx sweep as the config repo really writes one. It is not a gen +#: sweep with fewer keys: the model moves under `model:` and the lengths +#: move into the benchmark entries, because sweeping the input length is +#: the normal thing to do on a prefill-only run. +CTX_SWEEP = { + "model": { + "model_card": "deepseek-ai/DeepSeek-V4-Pro", + "model_path": {"aga-gb300": "/models/dsv4-pro"}, + "dataset_file": "${home_dir}/dataset/DeepSeek-V4-8192-1-20000-ratio-08_for_bench.json", + }, + "benchmarks": [ + {"isl": 8192, "osl": 1, "max_batch": [2], "tp_size": [4], "ratio": [0.8], "mtp_range": [3]} + ], +} + + +def test_a_ctx_sweep_states_its_workload_somewhere_else_and_is_read_there(): + """Reading only the gen shape resolves a ctx campaign to all `None`s. + + Which does not fail -- it reconciles nothing, and `task.yaml` keeps + the defaults block's `random_input_len: 1024` as this campaign's + stated input length however long the requests really are. + """ + load = bench_cli.workload(CTX_SWEEP) + assert load["model"] == "deepseek-ai/DeepSeek-V4-Pro" + assert load["isl"] == 8192 + assert load["osl"] == 1 + assert load["dataset"].endswith("_for_bench.json") + assert load["model_path"] == {"aga-gb300": "/models/dsv4-pro"} + + +def test_a_ctx_sweep_spanning_two_input_lengths_states_no_single_one(): + """Two rows at two lengths are two workloads. + + Resolving to the first, or to the largest, would let a campaign quote + one of them and describe half its own cases wrongly, with nothing + raising. + """ + spanning = { + **CTX_SWEEP, + "benchmarks": [ + {"isl": 1024, "osl": 1, "max_batch": [16], "tp_size": [4]}, + {"isl": 8192, "osl": 1, "max_batch": [2], "tp_size": [4]}, + ], + } + load = bench_cli.workload(spanning) + assert load["isl"] is None + assert load["osl"] == 1 # they do agree on this one, so it is stated + + +def test_a_top_level_length_still_wins_over_the_benchmark_entries(): + """The gen shape is not overridden by a sweep that carries both.""" + both = {**CTX_SWEEP, "isl": 4096, "osl": 512} + load = bench_cli.workload(both) + assert (load["isl"], load["osl"]) == (4096, 512) + + +FRONTIER = ( + "name,concurrency,throughput_per_user,output_tput_per_gpu," + "ctx_gen_inst_ratio_round_float,ctx_request_rate,ctx_gpus_round," + "gen_num_round,total_gpus_round\n" + "tep_4_eplb0_mtp3,1,165.3594,39.809,0.016735,10.67,8,4,12\n" + "tep_4_eplb0_mtp3,32,66.1748,368.7572,0.21532,10.67,8,4,12\n" +) + + +def test_the_csv_column_is_already_the_name_this_workflow_scores(tmp_path): + (tmp_path / "sol_frontier_mtp.csv").write_text(FRONTIER) + points = bench_cli.frontier_points(tmp_path) + assert [p["concurrency"] for p in points] == [1, 32] + assert points[0]["metrics"]["throughput_per_user"] == 165.3594 + # And the frontier the gate's metric is not, carried alongside. + assert points[0]["metrics"]["output_tput_per_gpu"] == 39.809 + assert points[0]["gpus"] == {"ctx": 8.0, "gen": 4.0, "total": 12.0} + + +def test_a_row_without_a_usable_metric_is_dropped_not_defaulted(tmp_path): + (tmp_path / "sol_frontier_mtp.csv").write_text( + FRONTIER + "tep_4_eplb0_mtp3,64,,10.0,0.3,10.67,8,4,12\n" + ) + assert [p["concurrency"] for p in bench_cli.frontier_points(tmp_path)] == [1, 32] + + +def test_a_nan_is_not_a_measurement(tmp_path): + """The postprocessor writes NaN into columns it could not compute.""" + (tmp_path / "sol_frontier_mtp.csv").write_text( + "name,concurrency,throughput_per_user\ntep_4_eplb0_mtp3,1,NaN\n" + ) + with pytest.raises(bench_cli.BenchCliError, match="no row with"): + bench_cli.frontier_points(tmp_path) + + +def test_a_run_dir_with_no_curve_names_the_command_that_says_why(tmp_path): + with pytest.raises(bench_cli.BenchCliError, match="process frontier"): + bench_cli.frontier_points(tmp_path) + + +def test_the_later_csv_wins_for_a_repeated_point(tmp_path): + """A run dir can hold several CSVs -- one per mtp tag and variant. + + Sorting makes which one wins deterministic rather than filesystem + order, so two reads of one directory agree. + """ + (tmp_path / "a_frontier_mtp0.csv").write_text(FRONTIER) + (tmp_path / "b_frontier_mtp.csv").write_text( + "name,concurrency,throughput_per_user\ntep_4_eplb0_mtp3,1,200.0\n" + ) + points = {p["concurrency"]: p for p in bench_cli.frontier_points(tmp_path)} + assert points[1]["metrics"]["throughput_per_user"] == 200.0 diff --git a/agent-flow/tests/workflows/perf_optimize/test_disagg.py b/agent-flow/tests/workflows/perf_optimize/test_disagg.py index ac862af45431..79a705c8cb77 100644 --- a/agent-flow/tests/workflows/perf_optimize/test_disagg.py +++ b/agent-flow/tests/workflows/perf_optimize/test_disagg.py @@ -287,7 +287,7 @@ def test_stage_prompts_state_the_mode_so_the_system_prompt_section_wins(tmp_path wf = PerfOptimizeWorkflow.__new__(PerfOptimizeWorkflow) wf._task_data = lambda: task_schema.load_and_validate_task_yaml(task) - directive = wf._disagg_directive() + directive = wf._campaign_directive() assert "DISAGGREGATED" in directive assert "replaces all of it" in directive assert str(harness) in directive @@ -295,4 +295,26 @@ def test_stage_prompts_state_the_mode_so_the_system_prompt_section_wins(tmp_path assert "trtllm-serve" in directive +def test_an_ordinary_campaign_is_handed_no_directive_at_all(tmp_path): + """The regression guard for every campaign that is neither mode. + + ``_campaign_directive`` replaced ``_disagg_directive`` and grew a second + branch for the SOL tracks. Both branches are gated, and a campaign with + neither block must come out exactly as it did before: the stage prompt + opens with ``Workspace:`` and nothing is prepended. + + The empty string is the whole contract. Anything else silently rewrites + the opening line of every aggregate run's prompt -- which is the one + change this refactor must not make, and the one a test of the two + populated branches cannot catch. + """ + from agent_flow.workflows.perf_optimize.workflow import PerfOptimizeWorkflow + + task = _write_task(tmp_path, {}) + wf = PerfOptimizeWorkflow.__new__(PerfOptimizeWorkflow) + wf._task_data = lambda: task_schema.load_and_validate_task_yaml(task) + + assert wf._campaign_directive() == "" + + # ------------------------------------------------- profiling wording diff --git a/agent-flow/tests/workflows/perf_optimize/test_disagg_sol.py b/agent-flow/tests/workflows/perf_optimize/test_disagg_sol.py new file mode 100644 index 000000000000..c4231f9d8cc2 --- /dev/null +++ b/agent-flow/tests/workflows/perf_optimize/test_disagg_sol.py @@ -0,0 +1,1390 @@ +"""The layer above a SOL campaign: fixing the operating point before optimizing at it.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from agent_flow.workflows.perf_optimize import bench_cli, disagg_sol + +FIELD = disagg_sol.DISAGG_SOL_FIELD + +HEADER = ( + "config,concurrency,mtp,tp,adp,eplb,tps_per_user," + "output_tps_per_gen_gpu,output_tput,avg_itertime_ms,num_iters\n" +) + + +def _shape_run(root: Path, name: str, rows: list[tuple]) -> Path: + """One shape's run dir, scored by the anchor-free extractor. + + Each row is (concurrency, tp, adp, eplb, mtp, tps_per_user, per_gen_gpu). + """ + run = Path(root) / name + run.mkdir(parents=True, exist_ok=True) + body = "".join( + f"cfg_{c},{c},{mtp},{tp},{adp},{eplb},{tpu},{ppg},{tpu * c},10.0,400\n" + for c, tp, adp, eplb, mtp, tpu, ppg in rows + ) + (run / bench_cli.GEN_ONLY_CSV).write_text(HEADER + body, encoding="utf-8") + return run + + +def _design(tmp_path: Path, *, state: dict | None = None) -> Path: + d = tmp_path / "sweep_design" + d.mkdir(parents=True, exist_ok=True) + payload = state if state is not None else {"phase": "PHASE3_DONE", "model_dir": "ds-v4"} + (d / disagg_sol.DESIGN_STATE).write_text(json.dumps(payload), encoding="utf-8") + return d + + +def _task(**design) -> dict: + return {FIELD: {"tracks": ["ctx", "gen"], "design": design}} + + +# ------------------------------------------------------------------ the block + + +def test_both_halves_are_the_default_and_the_order_is_stable(): + assert disagg_sol.tracks({FIELD: {}}) == ["ctx", "gen"] + assert disagg_sol.tracks({FIELD: {"tracks": "gen"}}) == ["gen"] + assert disagg_sol.tracks({FIELD: {"tracks": ["gen", "ctx", "gen"]}}) == ["gen", "ctx"] + + +def test_an_unknown_half_is_refused_by_name(): + with pytest.raises(disagg_sol.DisaggSolError, match=r"\['e2e'\]"): + disagg_sol.tracks({FIELD: {"tracks": ["ctx", "e2e"]}}) + + +def test_a_campaign_that_cannot_name_a_design_is_inheriting_its_point(tmp_path): + """The whole point of this layer is that the point was established. + + A spec with no design dir is a spec that took its operating point from + somewhere it cannot cite -- which is exactly the failure this module + exists to make impossible. + """ + with pytest.raises(disagg_sol.DisaggSolError, match="inheriting its operating point"): + disagg_sol.design_dir({FIELD: {"design": {}}}) + + +def test_which_end_of_the_frontier_matters_is_never_inferred(): + """The curve states the trade-off; it cannot state what it is bought for.""" + for bad in (None, "fastest", "balanced", ""): + with pytest.raises(disagg_sol.DisaggSolError, match="cannot state which end"): + disagg_sol.preference({FIELD: {"design": {"prefer": bad}}}) + assert disagg_sol.preference(_task(prefer="interactive")) == "interactive" + assert disagg_sol.preference(_task(prefer="throughput")) == "throughput" + + +# ----------------------------------------------------------------- the design + + +def test_a_design_that_never_ran_is_named_as_such(tmp_path): + with pytest.raises(disagg_sol.DisaggSolError, match="never run"): + disagg_sol.design_state(tmp_path / "nowhere") + + +def test_the_design_state_is_read_not_reimplemented(tmp_path): + d = _design(tmp_path, state={"phase": "PHASE2_PROBES_SUBMITTED", "gpu_name": "GB300"}) + assert disagg_sol.design_state(d)["phase"] == "PHASE2_PROBES_SUBMITTED" + + +def test_every_shape_contributes_its_own_measured_ladder(tmp_path): + """One run dir per shape, each at ITS measured batch -- not a common one.""" + d = _design(tmp_path) + _shape_run( + d, "bm_tep4", [(1, 4, "False", 0, 0, 200.0, 50.0), (64, 4, "False", 0, 0, 60.0, 960.0)] + ) + _shape_run(d, "bm_dep8", [(8, 8, "True", 0, 0, 150.0, 150.0)]) + points = disagg_sol.sweep_points(d) + assert {p["shape"] for p in points} == {"tep_4_eplb0_mtp0", "dep_8_eplb0_mtp0"} + assert {p["concurrency"] for p in points} == {1, 64, 8} + + +def test_an_unscored_design_names_the_command_that_scores_it(tmp_path): + with pytest.raises(disagg_sol.DisaggSolError, match="get_gen_only_perf"): + disagg_sol.sweep_points(_design(tmp_path)) + + +def test_a_case_that_never_settled_is_dropped_not_defaulted(tmp_path): + d = _design(tmp_path) + run = _shape_run(d, "bm_tep4", [(1, 4, "False", 0, 0, 200.0, 50.0)]) + (run / bench_cli.GEN_ONLY_CSV).write_text( + HEADER + "cfg_32,32,0,4,False,0,,,,10.0,400\n" + "cfg_1,1,0,4,False,0,200.0,50.0,200.0,10.0,400\n", + encoding="utf-8", + ) + points = disagg_sol.sweep_points(d) + assert [p["concurrency"] for p in points] == [1] + + +# ----------------------------------------------------------------- the front + + +def test_a_point_beaten_on_both_axes_is_dropped(): + points = [ + { + "shape": "a", + "concurrency": 1, + "throughput_per_user": 200.0, + "output_tps_per_gen_gpu": 50.0, + }, + { + "shape": "b", + "concurrency": 1, + "throughput_per_user": 190.0, + "output_tps_per_gen_gpu": 40.0, + }, + { + "shape": "c", + "concurrency": 64, + "throughput_per_user": 60.0, + "output_tps_per_gen_gpu": 960.0, + }, + ] + front = disagg_sol.pareto_front(points) + assert {p["shape"] for p in front} == {"a", "c"} + + +def test_a_tie_is_kept_because_choosing_between_them_is_not_arithmetic(): + points = [ + { + "shape": "a", + "concurrency": 1, + "throughput_per_user": 200.0, + "output_tps_per_gen_gpu": 50.0, + }, + { + "shape": "b", + "concurrency": 2, + "throughput_per_user": 200.0, + "output_tps_per_gen_gpu": 50.0, + }, + ] + assert len(disagg_sol.pareto_front(points)) == 2 + + +# -------------------------------------------------------------- the selection + +FRONT = [ + { + "shape": "tep_4", + "concurrency": 1, + "throughput_per_user": 214.0, + "output_tps_per_gen_gpu": 53.0, + }, + { + "shape": "tep_4", + "concurrency": 32, + "throughput_per_user": 97.0, + "output_tps_per_gen_gpu": 782.0, + }, + { + "shape": "dep_32", + "concurrency": 256, + "throughput_per_user": 40.0, + "output_tps_per_gen_gpu": 1200.0, + }, +] + + +def test_the_two_ends_of_one_curve_select_different_points(): + interactive = disagg_sol.select_point(FRONT, prefer="interactive") + throughput = disagg_sol.select_point(FRONT, prefer="throughput") + assert (interactive["shape"], interactive["concurrency"]) == ("tep_4", 1) + assert (throughput["shape"], throughput["concurrency"]) == ("dep_32", 256) + + +def test_the_missing_deployment_view_travels_with_the_chosen_number(): + """A shape that drags more context GPUs ranks better here than it deploys. + + `output_tps_per_gen_gpu` has no context term, so the ranking this + selection uses is not the deployment ranking -- and the reason has to be + attached to the result, not left in a docstring, because the result is + what a report quotes. + """ + chosen = disagg_sol.select_point(FRONT, prefer="throughput") + assert "output_tput_per_gpu" in chosen["e2e_view_absent"] + assert "GENERATION" in chosen["e2e_view_absent"] + assert "Absent, not flat" in chosen["e2e_view_absent"] + + +def test_the_incumbent_is_reported_against_but_never_ranked_with(): + """The incumbent is reported against, never ranked with. + + "The selection agrees with what we ran" and "it moved us" are the two + answers this staging exists to distinguish, and only one of them leaves + the previous campaigns' measurements meaningful. + """ + incumbent = {"shape": "tep_4", "concurrency": 32} + chosen = disagg_sol.select_point(FRONT, prefer="interactive", incumbent=incumbent) + assert chosen["moved"] is True + assert chosen["incumbent_on_pareto"] is True + assert chosen["incumbent"] == incumbent + # ...and the incumbent did not change what was picked. + assert disagg_sol.select_point(FRONT, prefer="interactive")["concurrency"] == 1 + + +def test_an_incumbent_that_is_dominated_is_said_to_be(): + dominated = {"shape": "tep_4", "concurrency": 999} + chosen = disagg_sol.select_point(FRONT, prefer="throughput", incumbent=dominated) + assert chosen["incumbent_on_pareto"] is False + assert chosen["moved"] is True + + +def test_selecting_from_nothing_is_refused_rather_than_defaulted(): + with pytest.raises(disagg_sol.DisaggSolError, match="no measured point"): + disagg_sol.select_point([], prefer="interactive") + + +def test_an_unknown_preference_is_refused_at_selection_too(): + with pytest.raises(disagg_sol.DisaggSolError, match="unknown preference"): + disagg_sol.select_point(FRONT, prefer="latency") + + +# --------------------------------------------------------------- the campaigns + + +def test_a_design_is_reusable_once_its_sweep_is_scored(tmp_path): + """Fixing the point costs ~10x the campaigns it enables. + + Paying that per campaign would be absurd, so "is it already established" + has to be answerable -- and answerable from the artefacts, not from what + `state.json` claims, because an interrupted design has the claim without + the CSVs. + """ + d = _design(tmp_path, state={"phase": "PHASE3_DONE"}) + assert disagg_sol.established(d) is False + _shape_run(d, "bm_tep4", [(1, 4, "False", 0, 0, 200.0, 50.0)]) + assert disagg_sol.established(d) is True + + +BASE = { + "checkpoint_path": "/ckpt", + "optimize": {"max_rounds": 3, "approaches": ["code"]}, + "disagg_sol": {"tracks": ["ctx", "gen"], "design": {"prefer": "interactive"}}, +} +CTX_POINT = {"ctx_gpus": 4, "max_batch": 2, "adp": True} +POINT = { + "shape": "tep_4_eplb0_mtp0", + "concurrency": 1, + "prefer": "interactive", + "ranked_on": "throughput_per_user", + "e2e_view_absent": disagg_sol.NO_JOIN, +} + + +def test_each_campaign_is_an_ordinary_single_track_spec(tmp_path): + """This layer is additive: nothing below it changes shape.""" + spec = disagg_sol.campaign_spec( + BASE, + track="gen", + sweep=tmp_path / "s.yaml", + repo=tmp_path / "r", + point=POINT, + design=tmp_path / "d", + ) + assert spec["sol_track"]["track"] == "gen" + assert spec["optimize"]["max_rounds"] == 3 + assert "disagg_sol" not in spec + + +def test_both_halves_inherit_where_to_execute_and_neither_inherits_what(tmp_path): + """`slurm-environment` reaches both halves; `benchmark` reaches neither. + + Off-cluster execution IS `cluster_ssh` + `remote_run_root`. Dropping that + block did not fail -- both halves ran as though the flow were on the + cluster, with the spec's remote settings absent from the only two files + that could act on them. + + `benchmark` is the other half of the same rule and must stay out: the + operating point comes from the design and travels in the derived sweep, so + an inherited workload would be a second, unmeasured description of the run + sitting beside the selected row. + """ + remote = { + "cluster_ssh": "aga-300", + "remote_run_root": "/scratch/run-root", + "docker_image": "/scratch/img.sqsh", + "slurm_partition": "batch", + } + base = dict(BASE, **{"slurm-environment": remote, "benchmark": {"concurrency": 64}}) + + specs = { + track: disagg_sol.campaign_spec( + base, + track=track, + sweep=tmp_path / f"{track}.yaml", + repo=tmp_path / track, + point=POINT if track == disagg_sol.GEN_TRACK else CTX_POINT, + design=tmp_path / "d", + ) + for track in disagg_sol.TRACKS + } + + for track, spec in specs.items(): + assert spec["slurm-environment"] == remote, f"{track} half lost its remote settings" + assert "benchmark" not in spec, f"{track} half inherited a workload it did not measure" + # The checkout is the one thing that must differ, and it is set per track. + assert specs["ctx"]["trtllm_repo_path"] != specs["gen"]["trtllm_repo_path"] + + +def test_the_campaign_can_say_where_its_operating_point_came_from(tmp_path): + """A path instead of a shrug -- the whole reason this layer exists.""" + spec = disagg_sol.campaign_spec( + BASE, + track="gen", + sweep=tmp_path / "s.yaml", + repo=tmp_path / "r", + point=POINT, + design=tmp_path / "design", + ) + prov = spec["sol_track"]["point_provenance"] + assert prov["design_dir"].endswith("design") + assert prov["selected"]["shape"] == "tep_4_eplb0_mtp0" + assert "Absent, not flat" in prov["e2e_view_absent"] + + +def test_no_anchor_is_handed_to_the_gen_half(tmp_path): + """There is no join in this scope, so there is nothing to anchor against.""" + spec = disagg_sol.campaign_spec( + BASE, + track="gen", + sweep=tmp_path / "s.yaml", + repo=tmp_path / "r", + point=POINT, + design=tmp_path / "d", + ) + assert "ctx_json" not in spec["sol_track"] + + +def _sweeps(tmp_path, *, gen_shape="tep", gen_tp=4, gen_conc="1", ctx_tp=4, ctx_batch=2): + """Real sweep files that contain (or deliberately miss) the selected point.""" + import yaml as _yaml + + g = tmp_path / "gen.yaml" + g.write_text( + _yaml.safe_dump( + {"gen_configs": [[1, 1, gen_tp, 64, 64, gen_shape == "dep", "0.9", 0, 0, gen_conc]]} + ) + ) + c = tmp_path / "ctx.yaml" + c.write_text( + _yaml.safe_dump( + {"benchmarks": [{"isl": 8192, "osl": 1, "max_batch": [ctx_batch], "tp_size": [ctx_tp]}]} + ) + ) + return {"ctx": c, "gen": g} + + +def _plan(tmp_path, repos=None): + return disagg_sol.launch_plan( + BASE, + sweeps=_sweeps(tmp_path, gen_shape="tep", gen_tp=4, gen_conc="1"), + repos=repos or {"ctx": tmp_path / "trtllm-ctx", "gen": tmp_path / "trtllm-gen"}, + workspace_root=tmp_path, + label="run1", + points={"ctx": CTX_POINT, "gen": POINT}, + design=tmp_path / "d", + ) + + +def test_the_plan_is_complete_before_anything_starts(tmp_path): + launches = _plan(tmp_path) + assert [run.track for run in launches] == ["ctx", "gen"] + assert launches[0].workspace.name == "ws-run1-ctx" + assert launches[1].workspace.name == "ws-run1-gen" + assert launches[0].argv[:2] == ["perf-optimize", "--task"] + + +def test_two_campaigns_may_not_share_one_checkout(tmp_path): + """A campaign resets the checkout it is given. + + Two sharing one would each revert the other's work mid-flight, and the + measurement that followed would be of neither's code -- while looking + entirely normal. + """ + one = tmp_path / "shared" + with pytest.raises(disagg_sol.DisaggSolError, match="each revert the other"): + _plan(tmp_path, repos={"ctx": one, "gen": one}) + + +def test_a_track_with_no_sweep_or_no_checkout_is_refused_by_name(tmp_path): + with pytest.raises(disagg_sol.DisaggSolError, match=r"neither a sweep of their own"): + disagg_sol.launch_plan( + BASE, + sweeps={"ctx": _sweeps(tmp_path)["ctx"]}, + repos={"ctx": tmp_path / "a", "gen": tmp_path / "b"}, + workspace_root=tmp_path, + label="x", + points={"ctx": CTX_POINT, "gen": POINT}, + design=tmp_path, + ) + with pytest.raises(disagg_sol.DisaggSolError, match=r"trtllm_repo_path.*\['gen'\]"): + disagg_sol.launch_plan( + BASE, + sweeps=_sweeps(tmp_path, gen_tp=4, gen_conc="1"), + repos={"ctx": tmp_path / "a"}, + workspace_root=tmp_path, + label="x", + points={"ctx": CTX_POINT, "gen": POINT}, + design=tmp_path, + ) + + +# ------------------------------------------------------------- the ctx half + + +def _ctx_case(root: Path, name: str, req_s: float) -> Path: + d = Path(root) / name + d.mkdir(parents=True, exist_ok=True) + (d / "run_dep4_MTP3.json").write_text( + json.dumps({"performance": {"request_throughput_req_s": req_s}}), encoding="utf-8" + ) + return d + + +def test_the_ctx_half_is_measured_and_selected_not_computed(tmp_path): + """Ranking prefill needs no rate match, so it is possible without a join. + + That is why this half can be measured even under a no-join scope -- the + objective lives entirely inside the ctx measurement. + """ + d = _design(tmp_path) + _ctx_case(d, "ctx_8192_1_ratio08_2_16416_dep4_MTP3_test1", 9.09) + _ctx_case(d, "ctx_8192_1_ratio08_4_16416_dep8_MTP3_test1", 14.0) + points = disagg_sol.ctx_points(d) + assert {p["ctx_gpus"] for p in points} == {4, 8} + # 9.09/4 = 2.27 beats 14.0/8 = 1.75 -- more total throughput, less per GPU. + chosen = disagg_sol.select_ctx_point(points) + assert chosen["ctx_gpus"] == 4 + assert disagg_sol.CTX_RANK in chosen["ranked_on"] + + +def test_a_tie_on_efficiency_breaks_toward_the_smaller_worker(tmp_path): + """Two configurations at one efficiency are not equal. + + The smaller one leaves the rest of the node to the generation side. + """ + d = _design(tmp_path) + _ctx_case(d, "ctx_8192_1_ratio08_2_16416_dep4_MTP3_test1", 8.0) + _ctx_case(d, "ctx_8192_1_ratio08_4_16416_dep8_MTP3_test1", 16.0) + chosen = disagg_sol.select_ctx_point(disagg_sol.ctx_points(d)) + assert chosen["ctx_gpus"] == 4 + + +def test_an_unparseable_ctx_case_name_is_skipped_not_guessed(tmp_path): + d = _design(tmp_path) + _ctx_case(d, "ctx_8192_1_ratio08_2_16416_dep4_MTP3_test1", 9.09) + _ctx_case(d, "something_else_entirely", 99.0) + assert [p["ctx_gpus"] for p in disagg_sol.ctx_points(d)] == [4] + + +def test_a_design_with_no_scored_ctx_case_says_the_half_would_be_computed(tmp_path): + with pytest.raises(disagg_sol.DisaggSolError, match="fall back to a computed point"): + disagg_sol.ctx_points(_design(tmp_path)) + + +# ------------------------------------------------- what one frozen point hides + + +def test_freezing_one_point_records_what_it_cannot_see(): + """opt-006 measured +1.52 % at one point and -1.85 % at another. + + With both frozen it was rejected; with one it is an accept. The gate is + the cheapest honest one, but the blind spot has to travel with it. + """ + chosen = disagg_sol.select_point(FRONT, prefer="interactive") + assert "Unobserved, not unchanged" in chosen["off_point_effects_unobserved"] + assert "indistinguishable" in chosen["off_point_effects_unobserved"] + + +def test_repeats_of_one_configuration_are_averaged_not_competed(tmp_path): + """Twelve directories, one configuration -- this is the real shape. + + A ctx sweep repeats each case and expands over mtp_range, and mtp is not + part of a ctx operating point: the generation sweep's `ctx_config` block + has no mtp field. Ranking the repeats as candidates picks the luckiest + sample, and the measured spread across repeats (4.0 %) is wider than the + difference between configurations this selection has to resolve. + """ + d = _design(tmp_path) + for mtp in (0, 3): + for test, req_s in ((1, 8.697), (2, 8.701), (3, 8.788)): + _ctx_case(d, f"ctx_8192_1_ratio08_2_16416_dep4_MTP{mtp}_test{test}", req_s) + chosen = disagg_sol.select_ctx_point(disagg_sol.ctx_points(d)) + assert chosen["candidates"] == 1 # one configuration... + assert chosen["measurements"] == 6 # ...measured six times + assert chosen["repeats"] == 6 + # the mean, not the 8.788 maximum + assert chosen[disagg_sol.CTX_METRIC] == pytest.approx(8.7286667, rel=1e-6) + + +def test_a_different_repeat_of_the_incumbent_is_not_a_move(tmp_path): + """`moved` is the answer this whole staging exists to give. + + Reporting True because a different directory won would say the operating + point should change when the selection in fact confirmed it. + """ + d = _design(tmp_path) + _ctx_case(d, "ctx_8192_1_ratio08_2_16416_dep4_MTP0_test3", 8.788) + _ctx_case(d, "ctx_8192_1_ratio08_2_16416_dep4_MTP3_test1", 8.484) + chosen = disagg_sol.select_ctx_point( + disagg_sol.ctx_points(d), + incumbent={"case": "ctx_8192_1_ratio08_2_16416_dep4_MTP3_test1"}, + ) + assert chosen["moved"] is False + + +def test_a_genuinely_different_configuration_is_a_move(tmp_path): + d = _design(tmp_path) + _ctx_case(d, "ctx_8192_1_ratio08_2_16416_dep4_MTP3_test1", 12.0) + _ctx_case(d, "ctx_8192_1_ratio08_4_16416_dep8_MTP3_test1", 8.0) + chosen = disagg_sol.select_ctx_point( + disagg_sol.ctx_points(d), + incumbent={"ctx_gpus": 8, "max_batch": 4, "adp": True}, + ) + assert chosen["ctx_gpus"] == 4 + assert chosen["moved"] is True + + +def test_the_configuration_with_the_better_mean_wins_a_lucky_repeat(tmp_path): + """The inversion the old rule produced, at the measured spread.""" + d = _design(tmp_path) + for test, req_s in ((1, 10.0), (2, 10.0), (3, 10.4)): # mean 10.13, max 10.4 + _ctx_case(d, f"ctx_8192_1_ratio08_2_16416_dep4_MTP0_test{test}", req_s) + for test, req_s in ((1, 10.2), (2, 10.3), (3, 10.2)): # mean 10.23, max 10.3 + _ctx_case(d, f"ctx_8192_1_ratio08_4_16416_dep4_MTP0_test{test}", req_s) + chosen = disagg_sol.select_ctx_point(disagg_sol.ctx_points(d)) + assert chosen["max_batch"] == 4 # the better mean, not the luckier max + assert chosen["spread_pct"] < 2.0 + + +# ---------------------------------------------------------------- the run + + +def _supervisable(tmp_path): + d = _design(tmp_path) + _ctx_case(d, "ctx_8192_1_ratio08_2_16416_dep4_MTP0_test1", 8.7) + _shape_run(d, "bm_tep4", [(1, 4, "False", 0, 0, 214.0, 53.0)]) + spec = { + "checkpoint_path": "/ckpt", + "optimize": {"approaches": ["code"]}, + FIELD: { + "tracks": ["ctx", "gen"], + "design": {"design_dir": str(d), "prefer": "interactive"}, + }, + } + return d, spec + + +def test_a_design_that_was_never_scored_is_refused_not_fallen_back_from(tmp_path): + """Falling back is exactly the behaviour this layer exists to remove.""" + d = _design(tmp_path) + spec = {FIELD: {"tracks": ["gen"], "design": {"design_dir": str(d), "prefer": "interactive"}}} + with pytest.raises(disagg_sol.DisaggSolError, match="no measured space to choose|no scored"): + disagg_sol.supervise( + spec, + sweeps={"gen": tmp_path / "g.yaml"}, + repos={"gen": tmp_path / "r"}, + workspace_root=tmp_path / "ws", + label="t", + dry_run=True, + ) + + +def test_a_dry_run_selects_and_writes_every_spec_without_starting_anything(tmp_path): + """The same code path minus the processes.""" + d, spec = _supervisable(tmp_path) + record = disagg_sol.supervise( + spec, + sweeps=_sweeps(tmp_path, gen_tp=4, gen_conc="1"), + repos={"ctx": tmp_path / "rc", "gen": tmp_path / "rg"}, + workspace_root=tmp_path / "ws", + label="t", + dry_run=True, + ) + assert record["started"] is False + assert record["ctx_point"]["ctx_gpus"] == 4 + assert record["gen_point"]["concurrency"] == 1 + assert [c["track"] for c in record["campaigns"]] == ["ctx", "gen"] + assert (tmp_path / "ws" / disagg_sol.RUN_RECORD).is_file() + # "Writes every spec" was the promise and not the behaviour: `materialize` + # lives inside `start_all`, and returning above it skipped the one + # artefact a reader would check the plan against. + for campaign in record["campaigns"]: + assert (Path(campaign["workspace"]) / "task.yaml").is_file() + assert set(record["task_paths"]) == {"ctx", "gen"} + + +def test_the_record_says_what_was_selected_and_against_what(tmp_path): + """Answerable afterwards from a file, not from a shrug.""" + d, spec = _supervisable(tmp_path) + record = disagg_sol.supervise( + spec, + sweeps=_sweeps(tmp_path, gen_tp=4, gen_conc="1"), + repos={"ctx": tmp_path / "rc", "gen": tmp_path / "rg"}, + workspace_root=tmp_path / "ws", + label="t", + dry_run=True, + incumbent={"shape": "tep_4_eplb0_mtp0", "concurrency": 1}, + ) + assert record["gen_point"]["moved"] is False + assert record["design_dir"] == str(d) + assert "e2e_view_absent" in record["gen_point"] + + +def test_a_half_that_is_ready_does_not_wait_on_one_that_is_not(tmp_path): + """The two halves are established by different artefacts. + + Under a no-join scope neither waits on the other, so requiring both would + idle a ready half on a dependency the scope does not have. Found by + launching the flow: a ctx-only spec was refused for a generation sweep it + had not asked for. + """ + d = _design(tmp_path) + _ctx_case(d, "ctx_8192_1_ratio08_2_16416_dep4_MTP0_test1", 8.7) + assert disagg_sol.established(d, disagg_sol.CTX_TRACK) is True + assert disagg_sol.established(d, disagg_sol.GEN_TRACK) is False + + spec = { + "checkpoint_path": "/ckpt", + "optimize": {"approaches": ["code"]}, + FIELD: {"tracks": ["ctx"], "design": {"design_dir": str(d), "prefer": "interactive"}}, + } + record = disagg_sol.supervise( + spec, + sweeps={"ctx": _sweeps(tmp_path)["ctx"]}, + repos={"ctx": tmp_path / "rc"}, + workspace_root=tmp_path / "ws", + label="t", + dry_run=True, + ) + assert record["ctx_point"]["ctx_gpus"] == 4 + assert "gen_point" not in record + + +def test_the_refusal_names_which_half_is_missing_what(tmp_path): + d = _design(tmp_path) + _ctx_case(d, "ctx_8192_1_ratio08_2_16416_dep4_MTP0_test1", 8.7) + spec = { + FIELD: {"tracks": ["ctx", "gen"], "design": {"design_dir": str(d), "prefer": "interactive"}} + } + with pytest.raises(disagg_sol.DisaggSolError, match=r"concurrency sweep.*for track 'gen'"): + disagg_sol.supervise( + spec, + sweeps=_sweeps(tmp_path), + repos={"ctx": tmp_path / "a", "gen": tmp_path / "b"}, + workspace_root=tmp_path / "ws", + label="t", + dry_run=True, + ) + + +def test_a_sweep_that_does_not_contain_the_selected_point_is_refused(tmp_path): + """Without this the selection is a report, not a decision. + + The chosen point goes into every campaign's `point_provenance` while the + campaign runs whatever its sweep says -- so a mismatched pair produces a + record claiming an operating point the run never used. That is worse than + not selecting at all: an untraceable point is at least honest about being + untraceable. + """ + d, spec = _supervisable(tmp_path) + # the design measured tep_4 @ c=1; hand it a sweep for dep_16 @ c=64 + with pytest.raises(disagg_sol.DisaggSolError, match="does not contain the selected point"): + disagg_sol.supervise( + spec, + sweeps=_sweeps(tmp_path, gen_shape="dep", gen_tp=16, gen_conc="64"), + repos={"ctx": tmp_path / "rc", "gen": tmp_path / "rg"}, + workspace_root=tmp_path / "ws2", + label="t", + dry_run=True, + ) + + +def test_the_refusal_lists_what_the_sweep_does_contain(tmp_path): + """So the reader can see whether the sweep or the selection is wrong.""" + d, spec = _supervisable(tmp_path) + with pytest.raises(disagg_sol.DisaggSolError, match=r"It expands to \[\('dep_16"): + disagg_sol.supervise( + spec, + sweeps=_sweeps(tmp_path, gen_shape="dep", gen_tp=16, gen_conc="64"), + repos={"ctx": tmp_path / "rc", "gen": tmp_path / "rg"}, + workspace_root=tmp_path / "ws3", + label="t", + dry_run=True, + ) + + +def test_a_ctx_sweep_at_the_wrong_worker_shape_is_refused(tmp_path): + """The ctx point is (tp_size, max_batch); a sweep at another is a mismatch.""" + d, spec = _supervisable(tmp_path) + spec = {**spec, FIELD: {**spec[FIELD], "tracks": ["ctx"]}} + with pytest.raises(disagg_sol.DisaggSolError, match="does not contain the selected point"): + disagg_sol.supervise( + spec, + sweeps={"ctx": _sweeps(tmp_path, ctx_tp=8, ctx_batch=16)["ctx"]}, + repos={"ctx": tmp_path / "rc"}, + workspace_root=tmp_path / "ws4", + label="t", + dry_run=True, + ) + + +def test_an_unestablished_design_is_run_rather_than_refused_when_one_can_be(tmp_path): + """An unestablished design is established rather than refused. + + Treating it purely as an input degenerated into nobody running it once + already -- that is how the two campaigns this layer replaces came to + inherit an unmeasured row. When a designer is available the run + establishes what it needs and re-checks: not a fallback, a first step. + """ + d = _design(tmp_path) + _ctx_case(d, "ctx_8192_1_ratio08_2_16416_dep4_MTP0_test1", 8.7) + spec = { + "checkpoint_path": "/ckpt", + "optimize": {"approaches": ["code"]}, + FIELD: {"tracks": ["gen"], "design": {"design_dir": str(d), "prefer": "interactive"}}, + } + seen: list[str] = [] + + def designer(instruction: str) -> None: + seen.append(instruction) + _shape_run(d, "bm_tep4", [(1, 4, "False", 0, 0, 214.0, 53.0)]) # now established + + record = disagg_sol.supervise( + spec, + sweeps={"gen": _sweeps(tmp_path)["gen"]}, + repos={"gen": tmp_path / "rg"}, + workspace_root=tmp_path / "wsd", + label="t", + dry_run=True, + designer=designer, + ) + assert len(seen) == 1 + assert "create-sweep" in seen[0] + assert record["gen_point"]["shape"] == "tep_4_eplb0_mtp0" + + +def test_a_designer_that_did_not_establish_the_design_still_refuses(tmp_path): + """Running it is not the same as it having worked.""" + d = _design(tmp_path) + spec = {FIELD: {"tracks": ["gen"], "design": {"design_dir": str(d), "prefer": "interactive"}}} + with pytest.raises(disagg_sol.DisaggSolError, match="concurrency sweep"): + disagg_sol.supervise( + spec, + sweeps={"gen": _sweeps(tmp_path)["gen"]}, + repos={"gen": tmp_path / "rg"}, + workspace_root=tmp_path / "wse", + label="t", + dry_run=True, + designer=lambda instruction: None, + ) + + +def test_only_the_unready_halves_are_named_to_the_designer(tmp_path): + """No point re-establishing a half that is already measured.""" + d = _design(tmp_path) + _ctx_case(d, "ctx_8192_1_ratio08_2_16416_dep4_MTP0_test1", 8.7) + spec = { + "checkpoint_path": "/ckpt", + FIELD: { + "tracks": ["ctx", "gen"], + "design": {"design_dir": str(d), "prefer": "interactive"}, + }, + } + seen: list[str] = [] + + def designer(instruction: str) -> None: + seen.append(instruction) + _shape_run(d, "bm_tep4", [(1, 4, "False", 0, 0, 214.0, 53.0)]) + + disagg_sol.supervise( + spec, + sweeps=_sweeps(tmp_path), + repos={"ctx": tmp_path / "a", "gen": tmp_path / "b"}, + workspace_root=tmp_path / "wsf", + label="t", + dry_run=True, + designer=designer, + ) + assert "['gen']" in seen[0] + assert "'ctx'" not in seen[0] + + +# ------------------------------------------------------- deriving the sweep + + +DESIGN_SWEEP = { + "model_id": "deepseek-ai/DeepSeek-V4-Pro", + "isl": 8192, + "osl": 1024, + "gen_configs": [ + [1, 1, 4, 64, 64, False, "0.9", 0, 0, "1,2,4,8,16,32,64"], + [1, 1, 8, 64, 64, False, "0.9", 0, 0, "1,2,4,8,16,32,64"], + [1, 1, 16, 32, 32, True, "0.9", 0, 0, "16,64,512,1024,2048"], + ], + "benchmarks": [ + {"isl": 8192, "osl": 1, "max_batch": [1, 2, 4], "tp_size": [4, 8], "ratio": [0.8]} + ], + "gpu_overrides": {"GB300": {"benchmarks": [{"max_batch": [2]}]}}, +} + + +def _design_sweep(tmp_path): + import yaml as _yaml + + p = tmp_path / "sweep_design" / "8k1k_sol_mtp0.yaml" + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(_yaml.safe_dump(DESIGN_SWEEP)) + return p + + +def test_a_campaign_cannot_run_the_whole_measured_space(tmp_path): + """Two shapes at one concurrency both land in `concurrency_`. + + That is what `_place` refuses, so the design sweep -- which is the whole + space -- is not something a campaign can freeze on. Hence the cut. + """ + import yaml as _yaml + + from agent_flow.workflows.perf_optimize import bench_cli + + cases = bench_cli.gen_cases(_yaml.safe_load(_design_sweep(tmp_path).read_text())) + at_64 = [c for c in cases if (c["config"] or {}).get("concurrency") == 64] + assert len({c["name"] for c in at_64}) > 1 # tep_4 and tep_8 collide + + +def test_the_derived_sweep_keeps_one_row_at_the_selected_point(tmp_path): + import yaml as _yaml + + out = disagg_sol.derive_sweep_at_point( + _design_sweep(tmp_path), + "gen", + {"shape": "tep_8_eplb0_mtp0", "concurrency": 64}, + into=tmp_path / "ws", + ) + got = _yaml.safe_load(out.read_text()) + assert len(got["gen_configs"]) == 1 + assert got["gen_configs"][0][2] == 8 # tp_size + assert got["gen_configs"][0][9] == "64" # this point only, not the ladder + + +def test_everything_but_the_row_filter_comes_from_the_design_sweep(tmp_path): + """So the campaign measures what the selection was made on. + + The two campaigns this layer replaces differed from their own recorded + point in three fields at once, and nothing noticed. + """ + import yaml as _yaml + + out = disagg_sol.derive_sweep_at_point( + _design_sweep(tmp_path), + "gen", + {"shape": "tep_4_eplb0_mtp0", "concurrency": 32}, + into=tmp_path / "ws", + ) + got = _yaml.safe_load(out.read_text()) + assert got["model_id"] == DESIGN_SWEEP["model_id"] + assert (got["isl"], got["osl"]) == (8192, 1024) + assert got["_derived_from"]["design_sweep"].endswith("8k1k_sol_mtp0.yaml") + + +def test_the_ctx_cut_narrows_both_axes_and_drops_the_override_table(tmp_path): + import yaml as _yaml + + out = disagg_sol.derive_sweep_at_point( + _design_sweep(tmp_path), "ctx", {"ctx_gpus": 8, "max_batch": 4}, into=tmp_path / "ws" + ) + got = _yaml.safe_load(out.read_text()) + assert got["benchmarks"] == [ + {"isl": 8192, "osl": 1, "max_batch": [4], "tp_size": [8], "ratio": [0.8]} + ] + assert "gpu_overrides" not in got + + +def test_a_derived_sweep_is_written_only_inside_the_campaign_workspace(tmp_path): + """The blunt rule an earlier module in this package lacked. + + It took an output path as a free parameter and would have overwritten a + curated config other people maintain -- unrecoverable in a way a wrong + measurement is not. + """ + ws = tmp_path / "ws" + out = disagg_sol.derive_sweep_at_point( + _design_sweep(tmp_path), "gen", {"shape": "tep_8_eplb0_mtp0", "concurrency": 1}, into=ws + ) + assert out.parent == ws.resolve() + assert out.name == disagg_sol.DERIVED_SWEEP_NAME + # the design's own sweep is untouched + assert "1,2,4,8,16,32,64" in _design_sweep(tmp_path).read_text() + + +def test_a_point_the_design_sweep_cannot_produce_is_refused(tmp_path): + with pytest.raises(disagg_sol.DisaggSolError, match="0 rows matching"): + disagg_sol.derive_sweep_at_point( + _design_sweep(tmp_path), + "gen", + {"shape": "dep_32_eplb0_mtp0", "concurrency": 4096}, + into=tmp_path / "ws", + ) + + +# ------------------------------------------------- where the design landed + + +def _established_at(root: Path, name: str) -> Path: + d = root / name + d.mkdir(parents=True, exist_ok=True) + (d / disagg_sol.DESIGN_STATE).write_text(json.dumps({"phase": "PHASE3_DONE"})) + _ctx_case(d, "ctx_8192_1_ratio08_2_16416_dep4_MTP0_test1", 8.7) + _shape_run(d, "bm_tep4", [(1, 4, "False", 0, 0, 214.0, 53.0)]) + return d + + +def test_the_requested_design_is_used_when_it_is_the_established_one(tmp_path): + d = _established_at(tmp_path / "model", "sweep_design2") + got, note = disagg_sol.resolve_design_dir(d, ["ctx", "gen"]) + assert got == d + assert note is None + + +def test_a_design_established_next_door_is_followed_and_recorded(tmp_path): + """Resuming an existing design instead of re-probing is the right call. + + It saved roughly twenty node-hours on the run this was written after. + Refusing it would forbid a correct decision; following it silently would + leave every later artefact citing a directory the spec never named. + """ + model = tmp_path / "model" + actual = _established_at(model, "sweep_design") + requested = model / "sweep_design2" + requested.mkdir() + + got, note = disagg_sol.resolve_design_dir(requested, ["ctx", "gen"]) + assert got == actual + assert "sweep_design2" in note and "sweep_design" in note + assert "Point the spec at" in note + + +def test_two_established_designs_side_by_side_are_refused(tmp_path): + """Which measurement to freeze on is not this layer's question.""" + model = tmp_path / "model" + _established_at(model, "sweep_design") + _established_at(model, "sweep_design_b") + with pytest.raises(disagg_sol.DisaggSolError, match="is not something this layer may pick"): + disagg_sol.resolve_design_dir(model / "sweep_design2", ["ctx", "gen"]) + + +def test_an_unestablished_neighbour_is_not_followed(tmp_path): + """The search admits only directories that measured what is wanted. + + A supervisor that hunts the filesystem for something that looks like a + design will eventually find one that is not. + """ + model = tmp_path / "model" + empty = model / "sweep_design" + empty.mkdir(parents=True) + (empty / disagg_sol.DESIGN_STATE).write_text(json.dumps({"phase": "PHASE2"})) + requested = model / "sweep_design2" + got, note = disagg_sol.resolve_design_dir(requested, ["gen"]) + assert got == requested and note is None + + +def test_the_redirection_travels_in_the_run_record(tmp_path): + """A reader must not have to already know.""" + model = tmp_path / "model" + _established_at(model, "sweep_design") + requested = model / "sweep_design2" + requested.mkdir() + spec = { + "checkpoint_path": "/ckpt", + FIELD: { + "tracks": ["gen"], + "design": {"design_dir": str(requested), "prefer": "interactive"}, + }, + } + record = disagg_sol.supervise( + spec, + sweeps={"gen": _sweeps(tmp_path)["gen"]}, + repos={"gen": tmp_path / "rg"}, + workspace_root=tmp_path / "wsr", + label="t", + dry_run=True, + designer=lambda instruction: None, # the design is already next door + ) + assert record["design_dir"].endswith("sweep_design") + assert "design_dir_redirected" in record + + +def test_a_design_that_already_exists_is_not_paid_for_again(tmp_path): + """The resolution happens before the designer is considered. + + A design costs about an order of magnitude more than the campaigns it + enables, so running one that already exists next door is the expensive + half of this mistake -- and avoiding exactly that is why the agent + resumed a neighbour to begin with. + """ + model = tmp_path / "model" + _established_at(model, "sweep_design") + requested = model / "sweep_design2" + requested.mkdir() + spec = { + "checkpoint_path": "/ckpt", + FIELD: { + "tracks": ["gen"], + "design": {"design_dir": str(requested), "prefer": "interactive"}, + }, + } + ran: list[str] = [] + record = disagg_sol.supervise( + spec, + sweeps={"gen": _sweeps(tmp_path)["gen"]}, + repos={"gen": tmp_path / "rg"}, + workspace_root=tmp_path / "wsn", + label="t", + dry_run=True, + designer=lambda instruction: ran.append(instruction), + ) + assert ran == [] # the designer was never started + assert record["design_dir"].endswith("sweep_design") + assert "design_dir_redirected" in record + + +def test_a_redirected_design_carries_its_sweep_with_it(tmp_path): + """`design_dir` and `design_sweep` are two fields naming one thing. + + A sweep path still rooted at the requested directory names a file that + was never written, and the failure arrives after the points have been + read and the point chosen -- the most expensive moment to find a path + problem. + """ + requested = tmp_path / "model" / "sweep_design2" + actual = tmp_path / "model" / "sweep_design" + actual.mkdir(parents=True) + (actual / "8k1k_sol_mtp0.yaml").write_text("gen_configs: []\n") + got = disagg_sol.rebase_design_sweep(requested / "8k1k_sol_mtp0.yaml", requested, actual) + assert got == actual / "8k1k_sol_mtp0.yaml" + + +def test_a_sweep_that_exists_where_stated_is_left_alone(tmp_path): + here = tmp_path / "elsewhere.yaml" + here.write_text("gen_configs: []\n") + assert disagg_sol.rebase_design_sweep(here, tmp_path / "a", tmp_path / "b") == here + + +def test_a_sweep_pointing_outside_the_design_is_a_deliberate_choice(tmp_path): + """Not under the requested directory, so not this redirection's business.""" + outside = tmp_path / "curated" / "sweep.yaml" + got = disagg_sol.rebase_design_sweep(outside, tmp_path / "a", tmp_path / "b") + assert got == outside + + +def test_each_half_is_given_its_own_point_not_the_other_s(tmp_path): + """A ctx point is (tp_size, max_batch); a gen point is (shape, concurrency). + + Handing one to both asks the ctx sweep for a row described in the + generation half's vocabulary -- observed as "0 benchmark entries matching + the selected ctx point tp_size None @ max_batch None". + """ + launches = disagg_sol.launch_plan( + BASE, + sweeps=_sweeps(tmp_path, gen_tp=4, gen_conc="1"), + repos={"ctx": tmp_path / "rc", "gen": tmp_path / "rg"}, + workspace_root=tmp_path / "wsp", + label="t", + points={"ctx": CTX_POINT, "gen": POINT}, + design=tmp_path / "d", + ) + by_track = {run.track: run.spec["sol_track"]["point_provenance"] for run in launches} + assert by_track["ctx"]["selected"]["ctx_gpus"] == 4 + assert by_track["ctx"]["selected"]["max_batch"] == 2 + assert by_track["gen"]["selected"]["shape"] == "tep_4_eplb0_mtp0" + assert by_track["gen"]["selected"]["concurrency"] == 1 + + +def test_the_two_halves_are_cut_from_different_design_sweeps(tmp_path): + """They are measured by different sweeps, so `design_sweep` is per track.""" + d, spec = _supervisable(tmp_path) + spec = {**spec, FIELD: {**spec[FIELD], "tracks": ["ctx", "gen"]}} + with pytest.raises(disagg_sol.DisaggSolError, match=r"neither a sweep of their own"): + disagg_sol.supervise( + spec, + sweeps={"gen": _sweeps(tmp_path)["gen"]}, # ctx has neither + repos={"ctx": tmp_path / "rc", "gen": tmp_path / "rg"}, + workspace_root=tmp_path / "wsq", + label="t", + dry_run=True, + ) + + +def test_a_derived_sweep_is_the_one_the_campaign_is_given(tmp_path): + """The resolution and the spec must name the same file. + + Reading the caller's `sweeps` again after deriving one silently hands the + campaign a path that was never derived -- and for a track with no sweep of + its own, no path at all. + """ + d = _design(tmp_path) + _shape_run(d, "bm_tep4", [(1, 4, "False", 0, 0, 214.0, 53.0)]) + design_sweep = tmp_path / "8k1k_sol_mtp0.yaml" + design_sweep.write_text("gen_configs:\n- [1, 1, 4, 64, 64, false, '0.9', 0, 0, '1,32']\n") + launches = disagg_sol.launch_plan( + {**BASE, FIELD: {**BASE[FIELD], "tracks": ["gen"]}}, + sweeps={}, + repos={"gen": tmp_path / "rg"}, + workspace_root=tmp_path / "wsd", + label="t", + points={"gen": POINT}, + design=d, + design_sweeps={"gen": design_sweep}, + ) + named = Path(launches[0].spec["sol_track"]["sweep"]) + assert named.name == disagg_sol.DERIVED_SWEEP_NAME + assert named.is_file() + assert named.parent == launches[0].workspace + + +def test_the_derived_sweep_carries_the_campaign_s_own_build_source(tmp_path): + """The design measures the image; a campaign measures its own edits. + + `sweep_design` refuses a build source in a design sweep -- choosing an + operating point on code no campaign starts from picks it for a different + program. A campaign is the opposite case: without a rung it would run the + image and report every change as no-gain. The rung is added at the + boundary where the purpose changes. + """ + import yaml as _yaml + + out = disagg_sol.derive_sweep_at_point( + _design_sweep(tmp_path), + "gen", + {"shape": "tep_8_eplb0_mtp0", "concurrency": 64}, + into=tmp_path / "ws", + repo=tmp_path / "trtllm-gen", + ) + got = _yaml.safe_load(out.read_text()) + assert got["trtllm_install"]["trtllm_repo"].endswith("trtllm-gen") + # ...and the design's own sweep still has none + assert "trtllm_install" not in _yaml.safe_load(_design_sweep(tmp_path).read_text()) + + +def test_a_derivation_without_a_checkout_adds_no_rung(tmp_path): + import yaml as _yaml + + out = disagg_sol.derive_sweep_at_point( + _design_sweep(tmp_path), + "gen", + {"shape": "tep_8_eplb0_mtp0", "concurrency": 64}, + into=tmp_path / "ws2", + ) + assert "trtllm_install" not in _yaml.safe_load(out.read_text()) + + +def test_the_plan_gives_the_derivation_the_campaign_s_checkout(tmp_path): + """End to end through launch_plan, not just the derivation in isolation. + + The rung reached `derive_sweep_at_point`'s signature and not its call + site, so the unit test passed while every real campaign was refused for + having no build source. + """ + import yaml as _yaml + + d = _design(tmp_path) + _shape_run(d, "bm_tep4", [(1, 4, "False", 0, 0, 214.0, 53.0)]) + design_sweep = tmp_path / "8k1k_sol_mtp0.yaml" + design_sweep.write_text("gen_configs:\n- [1, 1, 4, 64, 64, false, '0.9', 0, 0, '1,32']\n") + launches = disagg_sol.launch_plan( + {**BASE, FIELD: {**BASE[FIELD], "tracks": ["gen"]}}, + sweeps={}, + repos={"gen": tmp_path / "trtllm-gen"}, + workspace_root=tmp_path / "wsk", + label="t", + points={"gen": POINT}, + design=d, + design_sweeps={"gen": design_sweep}, + ) + got = _yaml.safe_load(Path(launches[0].spec["sol_track"]["sweep"]).read_text()) + assert got["trtllm_install"]["trtllm_repo"].endswith("trtllm-gen") + + +# ------------------------------------------- the units the two sides count in + + +GEN_NUM_SWEEP = { + "model_id": "deepseek-ai/DeepSeek-V4-Pro", + "isl": 8192, + "osl": 1024, + # Two generation servers: the row's list is PER SERVER, so this shape was + # measured at 2, 8 and 32 requests in flight across the deployment. + "gen_configs": [[1, 2, 8, 64, 64, False, "0.9", 0, 0, "1,4,16"]], +} + + +def _gen_num_sweep(tmp_path): + import yaml as _yaml + + p = tmp_path / "design2" / "sweep.yaml" + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(_yaml.safe_dump(GEN_NUM_SWEEP)) + return p + + +def test_a_multi_server_shape_is_matched_in_the_deployment_s_own_unit(tmp_path): + """The two sides of this comparison were in different units. + + A measured point's concurrency comes from `gen_only_perf.csv`, and the + harness names its result directories after the requests in flight across + the DEPLOYMENT -- the row's list times `gen_num`. The row's own list is + per generation server. Comparing them directly is right exactly when + `gen_num == 1`, and every multi-server shape had its own correct sweep + refused with "0 rows matching the selected point". + """ + import yaml as _yaml + + out = disagg_sol.derive_sweep_at_point( + _gen_num_sweep(tmp_path), + "gen", + {"shape": "tep_8_eplb0_mtp0", "concurrency": 32}, # 16 per server x 2 + into=tmp_path / "ws", + ) + row = _yaml.safe_load(out.read_text())["gen_configs"][0] + # Written back in the ROW's unit, not the deployment's: the total here + # would restate the point as a twice-as-large one and still look like + # the row it was cut from. + assert row[9] == "16" + assert row[1] == 2 + + +def test_a_point_that_is_not_in_the_deployment_ladder_is_still_refused(tmp_path): + """The unit fix must not turn the check into a rubber stamp.""" + with pytest.raises(disagg_sol.DisaggSolError, match="0 rows matching"): + disagg_sol.derive_sweep_at_point( + _gen_num_sweep(tmp_path), + "gen", + {"shape": "tep_8_eplb0_mtp0", "concurrency": 16}, # a per-server value + into=tmp_path / "ws", + ) + + +def test_the_pairing_check_counts_in_the_deployment_s_unit_too(tmp_path): + """`verify_sweep_matches_point` had the same mismatch, one layer over.""" + sweep = _gen_num_sweep(tmp_path) + disagg_sol.verify_sweep_matches_point( + sweep, "gen", {"shape": "tep_8_eplb0_mtp0", "concurrency": 8} + ) + with pytest.raises(disagg_sol.DisaggSolError, match="requests in flight"): + disagg_sol.verify_sweep_matches_point( + sweep, "gen", {"shape": "tep_8_eplb0_mtp0", "concurrency": 3} + ) + + +# ------------------------------------ what a ctx point is, and is not, made of + + +CTX_REPEATS_SWEEP = { + "rounds": 3, + "benchmarks": [ + { + "isl": 8192, + "osl": 1, + "max_batch": [2], + "tp_size": [4], + "ratio": [0.8], + "mtp_range": [0, 1, 2, 3], + } + ], +} + + +def _ctx_repeats_sweep(tmp_path): + import yaml as _yaml + + p = tmp_path / "design3" / "ctx_config.yaml" + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(_yaml.safe_dump(CTX_REPEATS_SWEEP)) + return p + + +def test_the_ctx_cut_collapses_the_speculation_range_to_one_depth(tmp_path): + """`mtp_range` is not part of a ctx point -- it is how the design repeated. + + A ctx case runs at output_length 1, so there is no decode and the + speculation depth changes nothing it measures. `select_ctx_point` groups + on (ctx_gpus, max_batch, adp) and means the rest for that reason, which + makes the design's four depths four REPEATS. Left as a range, the + campaign would re-run all four every time it measured -- and the run + this was found on did not: the agent narrowed to MTP0 at submit time, + outside the sweep, so the file said twelve cases and the campaign booked + one draw. + """ + import yaml as _yaml + + out = disagg_sol.derive_sweep_at_point( + _ctx_repeats_sweep(tmp_path), + "ctx", + {"ctx_gpus": 4, "max_batch": 2, "repeats": 12}, + into=tmp_path / "ws", + ) + got = _yaml.safe_load(out.read_text()) + assert got["benchmarks"][0]["mtp_range"] == [0] + # `rounds` stays: it is the campaign's OWN repetition, and it is what + # turns a single draw into a mean the noise floor can be read against. + assert got["rounds"] == 3 + + +def test_the_derived_sweep_says_how_many_repeats_each_side_counted(tmp_path): + """Because the selected value and the campaign's baseline are two means. + + The point was chosen from a mean over the design's repeats; the campaign + re-measures its own over a different, smaller n. Comparable, not + interchangeable -- and a gain quoted against the design's number rather + than the campaign's own baseline silently inherits the difference. + """ + import yaml as _yaml + + out = disagg_sol.derive_sweep_at_point( + _ctx_repeats_sweep(tmp_path), + "ctx", + {"ctx_gpus": 4, "max_batch": 2, "repeats": 12}, + into=tmp_path / "ws", + ) + repeats = _yaml.safe_load(out.read_text())["_derived_from"]["repeats"] + assert repeats["design"] == 12 + assert repeats["campaign"] == 3 + assert "baseline" in repeats["note"] + + +# ---------------------------------------------- the flag that meant two things + + +def test_a_dry_run_on_a_single_track_spec_is_refused_not_ignored(tmp_path, capsys): + """Silently ignoring it turns "show me" into a real campaign. + + The flag lives on the shared parser because argparse cannot know which + path a spec takes until the file is read, and only the staged path + implements it. Accepted and dropped, `--dry-run` on an ordinary campaign + started the full multi-hour run on real hardware -- the one failure mode + a dry run must not have. + """ + from agent_flow.workflows.perf_optimize import cli + + # Otherwise valid, so the refusal can only be about the flag: a spec + # that fails validation exits 2 as well, and a test that cannot tell + # the two apart passes on either. + (tmp_path / "ckpt").mkdir() + (tmp_path / "repo").mkdir() + task = tmp_path / "task.yaml" + task.write_text( + f"checkpoint_path: {tmp_path / 'ckpt'}\ntrtllm_repo_path: {tmp_path / 'repo'}\n", + encoding="utf-8", + ) + with pytest.raises(SystemExit) as caught: + cli.main(["--task", str(task), "--workspace", str(tmp_path / "ws"), "--dry-run"]) + assert caught.value.code == 2 + error = capsys.readouterr().err + assert "staged disagg path only" in error + assert "would start it for real" in error diff --git a/agent-flow/tests/workflows/perf_optimize/test_sol_track.py b/agent-flow/tests/workflows/perf_optimize/test_sol_track.py new file mode 100644 index 000000000000..6f42c6de00ca --- /dev/null +++ b/agent-flow/tests/workflows/perf_optimize/test_sol_track.py @@ -0,0 +1,887 @@ +"""Tests for the ``sol_track`` block. + +The authority rule is the one :mod:`.disagg` states — the sweep owns the +measurement conditions, ``task.yaml`` owns only the campaign knobs. + +Nothing here is stubbed. The harness communicates through files: a sweep +YAML going in, a frontier CSV and a ``run_*.json`` coming out. So the +tests write those files and read what this module makes of them, which is +also the only way to catch the shape drifting — a stub agrees with +whatever it was taught. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.perf_optimize import bench_cli, sol_track, task_schema +from agent_flow.workflows.perf_optimize.sol_track import SOL_TRACK_FIELD + +#: A gen sweep as the config repo writes one. The row order is the +#: harness': [ctx, gen, tp, batch, max_num_tokens, attention_dp, +#: gpu_mem_frac, mtp, eplb, concurrency_list]. +GEN_SWEEP = { + "benchmark_mode": "gen_only", + "model_id": "deepseek-ai/DeepSeek-V4-Pro", + "model_path": "/models/DeepSeek-V4-Pro", + "precision": "fp4", + "benchmark_client": "trtllm", + # isl is a sizing bound, not the corpus' length: the checked-in 8k + # sweep pairs isl 8192 with a ...-8192-1024-200000-... dataset. + "isl": 8192, + "osl": 1024, + "dataset_file": "/data/DeepSeek-V4-8192-1024-200000-ratio-08_for_serve.json", + "accept_rate": {"rate": "1:1.93,2:2.54,3:2.82", "source": "upstream"}, + "gen_configs": [[1, 1, 4, 64, 256, False, "0.9", 3, 0, "1,32"]], +} + +CTX_SWEEP = { + "model": {"model_card": "deepseek-ai/DeepSeek-V4-Pro"}, + "benchmarks": [ + {"isl": 1024, "osl": 1, "max_batch": [16], "tp_size": [4], "ratio": [1], "mtp_range": [3]}, + ], +} + + +def _write(path: Path, payload) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(yaml.safe_dump(payload, sort_keys=False), encoding="utf-8") + return path + + +def _sweep_dir(tmp_path, sweep=None, track="gen") -> Path: + """A sweep as it really is: a directory whose files reference siblings.""" + root = tmp_path / "campaign" + root.mkdir(exist_ok=True) + (root / "gen_worker_config.py").write_text("# imported from beside the sweep\n") + return _write( + root / "sweep.yaml", + sweep if sweep is not None else (GEN_SWEEP if track == "gen" else CTX_SWEEP), + ) + + +def _write_task(tmp_path, extra: dict | None = None) -> Path: + for sub in ("ckpt", "repo"): + (tmp_path / sub).mkdir(exist_ok=True) + data = { + "checkpoint_path": str(tmp_path / "ckpt"), + "trtllm_repo_path": str(tmp_path / "repo"), + # An upstream sweep carries no build source, so `config` is what a + # campaign against one can do. `code` needs a rung of the ladder + # and has its own tests. + "optimize": {"approaches": ["config"]}, + } + data.update(extra or {}) + return _write(tmp_path / "task.yaml", data) + + +def _block(tmp_path, track="gen", sweep=None, **extra) -> dict: + return { + "track": track, + "sweep": str(_sweep_dir(tmp_path, sweep, track)), + "workspace": str(tmp_path / "work"), + **extra, + } + + +def _gen(tmp_path, **extra) -> dict: + """A gen block with an anchor -- required, so most tests want it.""" + anchor = tmp_path / "anchor.json" + anchor.write_text(json.dumps([{"isl": 8192, "avg_request_throughput_req_s": 91.2}])) + return _block(tmp_path, ctx_json=str(anchor), **extra) + + +# ------------------------------------------------------------- the expansion + + +def test_the_sweep_is_the_authority_for_points_and_corpus(tmp_path): + """Nothing in task.yaml states the operating points; the sweep does.""" + data = task_schema.load_and_validate_task_yaml( + _write_task(tmp_path, {SOL_TRACK_FIELD: _gen(tmp_path)}) + ) + assert data["benchmark"]["concurrency"] == [1, 32] + assert data["benchmark"]["dataset_name"].endswith("_for_serve.json") + assert data["optimize"]["target_metric"] == "throughput_per_user" + assert data["profile"]["methods"] == ["nsys"] + assert task_schema.is_curve_mode(data) + + +def test_the_operating_point_is_the_product_of_concurrency_and_gen_num(tmp_path): + """The one trap the harness keeps: a row's concurrency is PER GEN SERVER. + + The client is driven at `concurrency * gen_num` and the result + directory is named for that product, so that is what `task.yaml` + means by concurrency. + """ + sweep = {**GEN_SWEEP, "gen_configs": [[1, 2, 4, 64, 256, False, "0.9", 3, 0, "64"]]} + data = task_schema.load_and_validate_task_yaml( + _write_task(tmp_path, {SOL_TRACK_FIELD: _gen(tmp_path, sweep=sweep)}) + ) + assert data["benchmark"]["concurrency"] == 128 # 64 x 2 -> a single point + + +def test_a_ctx_case_is_addressed_by_max_batch(tmp_path): + """A ctx entry has no concurrency at all. + + `max_batch` is the in-flight request count for a prefill-only run, + which is what this workflow means by concurrency everywhere else. + """ + data = task_schema.load_and_validate_task_yaml( + _write_task(tmp_path, {SOL_TRACK_FIELD: _block(tmp_path, track="ctx")}) + ) + assert data["benchmark"]["concurrency"] == 16 + assert data["optimize"]["target_metric"] == "avg_request_throughput_req_s" + + +def test_the_case_name_mirrors_the_column_the_postprocessor_writes(tmp_path): + """`{dep|tep}_{tp}_eplb{N}_mtp{M}` -- so a scored row addresses its row.""" + cases = bench_cli.gen_cases(GEN_SWEEP) + assert [c["name"] for c in cases] == ["tep_4_eplb0_mtp3"] * 2 + assert [c["config"]["concurrency"] for c in cases] == [1, 32] + dep = bench_cli.gen_cases( + {**GEN_SWEEP, "gen_configs": [[1, 1, 8, 64, 256, True, "0.8", 0, 384, "16"]]} + ) + assert dep[0]["name"] == "dep_8_eplb384_mtp0" + + +def test_a_track_the_sweep_does_not_measure_is_refused(tmp_path): + """A ctx campaign against a gen-only sweep would measure nothing.""" + with pytest.raises(task_schema.TaskSchemaError, match="plans no 'ctx' cases"): + task_schema.load_and_validate_task_yaml( + _write_task(tmp_path, {SOL_TRACK_FIELD: _block(tmp_path, track="ctx", sweep=GEN_SWEEP)}) + ) + + +# ------------------------------------------------------------------ the block + + +def test_a_campaign_needs_no_work_dir_of_its_own(tmp_path): + """Submission and collection are one path, so there is nothing to state. + + A stage submits into `/run` and collects from the same + place. `sol_track.workspace` is kept because task.yaml files carry it, + but a spec without one validates: there is no second path left for it + to disagree with. + """ + anchor = tmp_path / "anchor.json" + anchor.write_text(json.dumps([{"isl": 8192}])) + data = task_schema.load_and_validate_task_yaml( + _write_task( + tmp_path, + { + SOL_TRACK_FIELD: { + "track": "gen", + "sweep": str(_sweep_dir(tmp_path)), + "ctx_json": str(anchor), + } + }, + ) + ) + assert data["benchmark"]["concurrency"] == [1, 32] + + +def test_a_missing_sweep_file_is_refused(tmp_path): + with pytest.raises(task_schema.TaskSchemaError, match="is not a file"): + task_schema.load_and_validate_task_yaml( + _write_task( + tmp_path, + {SOL_TRACK_FIELD: {"track": "gen", "sweep": "/nope.yaml", "workspace": "/w"}}, + ) + ) + + +def test_it_cannot_be_combined_with_a_disagg_campaign(tmp_path): + """One campaign measures one thing, and each reconciles from its own file.""" + harness = _write(tmp_path / "disagg.yaml", {"worker_config": {"ctx": {}, "gen": {}}}) + with pytest.raises(task_schema.TaskSchemaError, match="cannot be combined"): + task_schema.load_and_validate_task_yaml( + _write_task( + tmp_path, {SOL_TRACK_FIELD: _gen(tmp_path), "disagg": {"config": str(harness)}} + ) + ) + + +def test_extra_llm_api_options_beside_a_track_is_refused(tmp_path): + """Two seeds for one tuning file: the named one would be discarded.""" + other = _write(tmp_path / "extra.yaml", {"knob": 1}) + with pytest.raises(task_schema.TaskSchemaError, match="cannot be combined"): + task_schema.load_and_validate_task_yaml( + _write_task( + tmp_path, {SOL_TRACK_FIELD: _gen(tmp_path), "extra_llm_api_options": str(other)} + ) + ) + + +def test_a_condition_the_user_wrote_that_disagrees_is_an_error(tmp_path): + """Stated on `concurrency`: it is what each point is scored at.""" + with pytest.raises(task_schema.TaskSchemaError, match="contradicts the sweep"): + task_schema.load_and_validate_task_yaml( + _write_task( + tmp_path, {SOL_TRACK_FIELD: _gen(tmp_path), "benchmark": {"concurrency": [1, 999]}} + ) + ) + + +def test_the_corpus_is_the_workload_not_the_sizing_bound(tmp_path): + """`isl` bounds the KV allocation; the dataset is what gets served. + + The checked-in sweep pairs `isl: 8192` with a `...-200000-...` corpus + and both numbers are right. Copying the bound into `random_input_len` + asserted a synthetic dataset of uniformly 8192-token requests. + """ + data = task_schema.load_and_validate_task_yaml( + _write_task(tmp_path, {SOL_TRACK_FIELD: _gen(tmp_path)}) + ) + bench = data["benchmark"] + assert bench["dataset_path"] == GEN_SWEEP["dataset_file"] + assert "random_input_len" not in bench + assert "random_output_len" not in bench + notes = " ".join(data[SOL_TRACK_FIELD]["filled_from_sweep_plan"]) + assert "sequence-length bounds" in notes + + +def test_no_note_answers_its_own_question_with_none(tmp_path): + """The resolved spec carried ``code_id at plan time: None``, always. + + ``plan()`` emits no such key -- the previous CLI digested the config and + the checkout into one fingerprint and this one does not -- so the line + was structurally absent, in the one place a reader goes to work out what + a number was measured against. + """ + data = task_schema.load_and_validate_task_yaml( + _write_task(tmp_path, {SOL_TRACK_FIELD: _gen(tmp_path)}) + ) + assert not any("code_id" in note for note in data[SOL_TRACK_FIELD]["filled_from_sweep_plan"]) + + +def test_an_owner_who_names_another_metric_still_wins(tmp_path): + data = task_schema.load_and_validate_task_yaml( + _write_task( + tmp_path, + { + SOL_TRACK_FIELD: _gen(tmp_path), + "optimize": {"target_metric": "mine", "approaches": ["config"]}, + }, + ) + ) + assert data["optimize"]["target_metric"] == "mine" + + +# ------------------------------------------------------------------ anchors + + +def _anchor(tmp_path, rows) -> Path: + path = tmp_path / "ctx_anchor.json" + path.write_text(json.dumps(rows), encoding="utf-8") + return path + + +def test_a_gen_track_without_a_ctx_anchor_runs_and_says_what_it_cannot_see(tmp_path): + """The anchor buys the end-to-end view, not the right to be scored. + + The gate's metric is `accept_rate / avg_iteration_time` -- decode + iterations only, no context term -- so no ctx measurement can move it. + Requiring an anchor here made a decode campaign wait on somebody's + prefill run to score a change prefill cannot affect. What must not + happen is that the missing end-to-end half goes unmentioned. + """ + data = task_schema.load_and_validate_task_yaml( + _write_task(tmp_path, {SOL_TRACK_FIELD: _block(tmp_path)}) + ) + assert data["optimize"]["target_metric"] == "throughput_per_user" + notes = " ".join(data[SOL_TRACK_FIELD]["filled_from_sweep_plan"]) + assert "output_tput_per_gpu" in notes + assert "ABSENT" in notes + + +def test_an_anchor_measured_on_other_work_is_refused(tmp_path): + """A frontier is a decode rate over a prefill rate. + + Feed it halves measured at different input lengths and it still + returns a number, on a curve that still looks like a frontier. + """ + anchor = _anchor(tmp_path, [{"isl": 1024, "avg_request_throughput_req_s": 91.2}]) + with pytest.raises(sol_track.SolTrackError, match="measured at isl 1024"): + sol_track.require_matching_anchor(anchor, bench_cli.plan(GEN_SWEEP)) + + +def test_a_matching_anchor_passes_and_an_unstated_one_says_so(tmp_path): + """'not checked' and 'checked and matched' must not read alike.""" + good = _anchor(tmp_path, [{"isl": 8192}]) + assert sol_track.require_matching_anchor(good, bench_cli.plan(GEN_SWEEP)) is None + silent = _anchor(tmp_path, [{"avg_request_throughput_req_s": 10.7}]) + note = sol_track.require_matching_anchor(silent, bench_cli.plan(GEN_SWEEP)) + assert note and "no 'isl'" in note + + +# ------------------------------------------------------------------ the score + + +FRONTIER_HEADER = ( + "name,concurrency,throughput_per_user,output_tput_per_gpu," + "ctx_gen_inst_ratio_round_float,ctx_request_rate,ctx_gpus_round," + "gen_num_round,total_gpus_round\n" +) + + +def _run_dir(into, rows=None, name="bm_deepseek-v4-pro-sol-8192-1024-20260907-GB300") -> Path: + """A harness run directory under the result dir a stage collects into.""" + run = Path(into) / sol_track.RUN_SUBDIR / name + run.mkdir(parents=True, exist_ok=True) + body = ( + rows + if rows is not None + else ( + "tep_4_eplb0_mtp3,1,165.3594,39.809,0.016735,10.67,8,4,12\n" + "tep_4_eplb0_mtp3,32,66.1748,368.7572,0.21532,10.67,8,4,12\n" + ) + ) + (run / "sol_frontier_mtp.csv").write_text(FRONTIER_HEADER + body, encoding="utf-8") + return run + + +GEN_ONLY_HEADER = ( + "config,concurrency,mtp,tp,adp,eplb,tps_per_user," + "output_tps_per_gen_gpu,output_tput,avg_itertime_ms,num_iters\n" +) + + +def _gen_only_dir(into, rows=None, name="bm_deepseek-v4-pro-sol-8192-1024-20260907-GB300") -> Path: + """A run dir scored by the anchor-free extractor instead of a frontier.""" + run = Path(into) / sol_track.RUN_SUBDIR / name + run.mkdir(parents=True, exist_ok=True) + body = ( + rows + if rows is not None + else ( + "ctx1_gen1_tep4_c1_eplb0_mtp3,1,3,4,False,0,165.36,41.34,165.36,11.67,412\n" + "ctx1_gen1_tep4_c32_eplb0_mtp3,32,3,4,False,0,66.17,529.39,2117.57,29.16,388\n" + ) + ) + (run / bench_cli.GEN_ONLY_CSV).write_text(GEN_ONLY_HEADER + body, encoding="utf-8") + return run + + +def _task(track="gen", **extra) -> dict: + """A resolved sol_track block. + + A gen campaign carries an anchor by default, because that is the + end-to-end shape and it is the one the frontier reader applies to. + Pass ``ctx_json=None`` for the anchor-free half; the reader dispatches + on what the campaign DECLARES, never on which file is on disk. + """ + block = {"track": track, **extra} + if track == sol_track.GEN_TRACK: + block.setdefault("ctx_json", "/anchors/ctx_anchor.json") + return {SOL_TRACK_FIELD: block} + + +def test_the_metric_lands_under_the_name_the_campaign_scores(tmp_path): + """The CSV column is already `throughput_per_user`; no rename happens. + + What matters is that it lands under `optimize.target_metric`, because + the baseline gate looks that key up and reports its absence as a stage + that measured nothing. + """ + _run_dir(tmp_path / "baseline") + written = sol_track.collect(_task(), tmp_path / "baseline") + assert [p.parent.name for p in written] == ["concurrency_1", "concurrency_32"] + payload = json.loads(written[0].read_text()) + assert payload["throughput_per_user"] == 165.3594 + assert payload["shape"] == "tep_4_eplb0_mtp3" + assert payload["source_csv"].endswith("sol_frontier_mtp.csv") + + +def test_it_writes_wherever_the_stage_was_told_to(tmp_path): + """Not just the baseline: every attempt has its own result directory.""" + _run_dir(tmp_path / "baseline") + _run_dir(tmp_path / "r1/attempt_1") + written = sol_track.collect(_task(), tmp_path / "r1/attempt_1") + assert written[0] == tmp_path / "r1/attempt_1/concurrency_1" / sol_track.SOL_RESULT_NAME + + +def test_what_a_gen_gain_is_worth_at_the_frontier_travels_with_it(tmp_path): + """The gate's metric is not the deployment's objective. + + `throughput_per_user` is anchor-free; `output_tput_per_gpu` divides by + a denominator the context side owns, so the same measured +1 % is + worth different amounts at two ends of one curve. + """ + _run_dir(tmp_path / "baseline") + sol_track.collect(_task(), tmp_path / "baseline") + at_1 = json.loads((tmp_path / "baseline/concurrency_1" / sol_track.SOL_RESULT_NAME).read_text()) + at_32 = json.loads( + (tmp_path / "baseline/concurrency_32" / sol_track.SOL_RESULT_NAME).read_text() + ) + # gen_gpus / (ctx_gpus * ctx_per_gen + gen_gpus), all three read off the row. + assert at_1["frontier_elasticity"] == pytest.approx(4 / (8 * 0.016735 + 4), rel=1e-6) + assert at_1["frontier_elasticity"] == pytest.approx(0.968, abs=0.005) + assert at_32["frontier_elasticity"] < at_1["frontier_elasticity"] + + +def test_two_points_at_one_concurrency_are_refused_rather_than_overwritten(tmp_path): + """Keeping the last writer scores the campaign on whichever row came second.""" + _run_dir( + tmp_path / "baseline", + rows=( + "tep_4_eplb0_mtp3,8,165.0,39.8,0.0167,10.67,8,4,12\n" + "dep_8_eplb0_mtp3,8,120.0,50.0,0.0200,10.67,8,4,12\n" + ), + ) + with pytest.raises(sol_track.SolTrackError, match="both measure 8 requests in flight"): + sol_track.collect(_task(), tmp_path / "baseline") + + +def test_a_run_dir_with_no_scored_curve_names_the_command_that_says_why(tmp_path): + (tmp_path / "baseline" / sol_track.RUN_SUBDIR / "bm_x").mkdir(parents=True) + with pytest.raises(sol_track.SolTrackError, match="process frontier"): + sol_track.collect(_task(), tmp_path / "baseline") + + +def test_two_run_dirs_under_one_result_dir_are_refused(tmp_path): + """One work dir per campaign per workload. + + Two of them means the score cannot say which run it came from -- and + the run directory name carries no code identity, so a second attempt + submitted into the same place would silently join the first. + """ + _run_dir(tmp_path / "baseline", name="bm_a") + _run_dir(tmp_path / "baseline", name="bm_b") + with pytest.raises(sol_track.SolTrackError, match="2 run directories"): + sol_track.collect(_task(), tmp_path / "baseline") + + +def test_a_row_missing_its_metric_is_skipped_not_defaulted(tmp_path): + _run_dir( + tmp_path / "baseline", + rows=( + "tep_4_eplb0_mtp3,1,165.3594,39.809,0.016735,10.67,8,4,12\n" + "tep_4_eplb0_mtp3,32,,368.7572,0.21532,10.67,8,4,12\n" + ), + ) + written = sol_track.collect(_task(), tmp_path / "baseline") + assert [p.parent.name for p in written] == ["concurrency_1"] + + +# ------------------------------------------------- the anchor-free gen reader + + +def test_a_gen_campaign_with_no_anchor_is_scored_from_the_iteration_logs(tmp_path): + """Same column, same formula, no prefill measurement anywhere in it. + + `tps_per_user` is what `get_gen_only_perf` calls the quantity the + frontier CSV calls `throughput_per_user`; the rename is the only + difference in the number. + """ + _gen_only_dir(tmp_path / "baseline") + written = sol_track.collect(_task(ctx_json=None), tmp_path / "baseline") + assert [p.parent.name for p in written] == ["concurrency_1", "concurrency_32"] + payload = json.loads(written[0].read_text()) + assert payload["throughput_per_user"] == 165.36 + assert payload["source_csv"].endswith(bench_cli.GEN_ONLY_CSV) + + +def test_the_two_gen_readers_name_a_point_the_same_way(tmp_path): + """A point must address the planned row through either reader. + + The extractor writes its own `config` string; the shape is rebuilt + from the columns instead, so it matches both the frontier CSV's `name` + and the sweep expansion's. + """ + _gen_only_dir(tmp_path / "unanchored") + _run_dir(tmp_path / "anchored") + unanchored = json.loads( + sol_track.collect(_task(ctx_json=None), tmp_path / "unanchored")[0].read_text() + ) + anchored = json.loads(sol_track.collect(_task(), tmp_path / "anchored")[0].read_text()) + assert unanchored["shape"] == anchored["shape"] == "tep_4_eplb0_mtp3" + assert {case["name"] for case in bench_cli.gen_cases(GEN_SWEEP)} == {unanchored["shape"]} + + +def test_the_missing_end_to_end_half_is_stated_rather_than_left_blank(tmp_path): + """Absent, not zero and not unchanged. + + A reader who opens one of these files must not be able to mistake + "this campaign could not see the frontier" for "the frontier did not + move" -- so the reason is written into the result, not merely implied + by a missing key. + """ + _gen_only_dir(tmp_path / "baseline") + payload = json.loads( + sol_track.collect(_task(ctx_json=None), tmp_path / "baseline")[0].read_text() + ) + assert payload["e2e_view"] is None + assert "ABSENT" in payload["e2e_view_absent"] + assert "frontier_elasticity" in payload["e2e_view_absent"] + assert "frontier_elasticity" not in payload + assert "frontier_metrics" not in payload + + +def test_the_flattering_per_gen_gpu_number_is_not_renamed_to_the_frontier_s(tmp_path): + """`output_tps_per_gen_gpu` divides by the generation GPUs alone. + + The frontier's `output_tput_per_gpu` divides by the whole rate-matched + deployment, so the two are never the same number and this one is + always the larger. Carried under its own name, or a reader quotes an + end-to-end result the campaign never measured. + """ + _gen_only_dir(tmp_path / "baseline") + payload = json.loads( + sol_track.collect(_task(ctx_json=None), tmp_path / "baseline")[0].read_text() + ) + assert payload["gen_only_metrics"]["output_tps_per_gen_gpu"] == 41.34 + assert "output_tput_per_gpu" not in json.dumps(payload["gen_only_metrics"]) + + +def test_a_frontier_the_campaign_never_declared_an_anchor_for_is_not_read(tmp_path): + """Dispatch is on the declaration, never on what is on disk. + + The only thing checked about an anchor is that its input length is + this sweep's. An anchor nobody declared was checked against nothing, + and the frontier it produces still plots a clean curve. + """ + _run_dir(tmp_path / "baseline") # a frontier CSV, and only that + with pytest.raises(sol_track.SolTrackError, match="get_gen_only_perf"): + sol_track.collect(_task(ctx_json=None), tmp_path / "baseline") + + +def test_a_case_that_never_reached_steady_state_is_skipped_not_defaulted(tmp_path): + """The extractor drops such a case, so its row arrives without a score.""" + _gen_only_dir( + tmp_path / "baseline", + rows=( + "ctx1_gen1_tep4_c1_eplb0_mtp3,1,3,4,False,0,165.36,41.34,165.36,11.67,412\n" + "ctx1_gen1_tep4_c32_eplb0_mtp3,32,3,4,False,0,,529.39,2117.57,29.16,388\n" + ), + ) + written = sol_track.collect(_task(ctx_json=None), tmp_path / "baseline") + assert [p.parent.name for p in written] == ["concurrency_1"] + + +def test_an_attention_dp_row_is_named_dep_rather_than_tep(tmp_path): + """`adp` arrives as the literal `True`/`False` pandas writes.""" + _gen_only_dir( + tmp_path / "baseline", + rows="ctx1_gen1_dep8_c32_eplb256_mtp3,32,3,8,True,256,66.17,529.39,2117.57,29.16,388\n", + ) + payload = json.loads( + sol_track.collect(_task(ctx_json=None), tmp_path / "baseline")[0].read_text() + ) + assert payload["shape"] == "dep_8_eplb256_mtp3" + + +def test_a_ctx_campaign_is_scored_from_the_field_the_harness_validates_on(tmp_path): + """There is no frontier on this track and running one is an error. + + `process frontier` rate-matches the GEN curve and takes ctx only as + its anchor, so a ctx campaign reads the `run_*.json` its own case + left -- whose `performance.request_throughput_req_s` is the field the + harness itself requires before calling the case successful. + """ + case = ( + tmp_path + / "baseline" + / sol_track.RUN_SUBDIR + / "bm_ctx" + / "ctx_1024_1_ratio1_16_16640_dep4_MTP3_test2" + ) + case.mkdir(parents=True) + (case / "run_dep4_MTP3.json").write_text( + json.dumps({"performance": {"request_throughput_req_s": 65.405}}) + ) + written = sol_track.collect(_task("ctx"), tmp_path / "baseline") + assert [p.parent.name for p in written] == ["concurrency_16"] + payload = json.loads(written[0].read_text()) + assert payload["avg_request_throughput_req_s"] == 65.405 + assert payload["source_field"] == "performance.request_throughput_req_s" + + +def test_a_ctx_run_with_nothing_validated_says_which_command_reports_why(tmp_path): + (tmp_path / "baseline" / sol_track.RUN_SUBDIR / "bm_ctx").mkdir(parents=True) + with pytest.raises(sol_track.SolTrackError, match="jobs check"): + sol_track.collect(_task("ctx"), tmp_path / "baseline") + + +# ------------------------------------------------------------------ guards + + +def _overlay(tmp_path, body: dict): + stage = _sweep_dir(tmp_path) + tuning = tmp_path / "tuning.yaml" + tuning.write_text(yaml.safe_dump(body), encoding="utf-8") + return {SOL_TRACK_FIELD: {"track": "gen", "workspace": "/w", "sweep": str(stage)}}, tuning + + +def test_an_overlay_may_not_move_the_operating_point(tmp_path): + """The sweep row's knobs ARE the point; changing one voids the comparison. + + The run succeeds, the number is plausible, and the only trace is a + node count in a log line. + """ + task, tuning = _overlay(tmp_path, {"tensor_parallel_size": 8, "moe_config": {"backend": "X"}}) + with pytest.raises(sol_track.SolTrackError, match="tensor_parallel_size"): + sol_track.apply_overlay(task, tuning) + + +def test_a_knob_that_is_not_the_operating_point_is_still_tunable(tmp_path): + task, tuning = _overlay(tmp_path, {"moe_config": {"backend": "TRTLLM"}}) + written = sol_track.apply_overlay(task, tuning) + assert yaml.safe_load(written.read_text())["gen_extra_llm_api"] == { + "moe_config": {"backend": "TRTLLM"} + } + + +def test_the_campaign_measures_a_copy_and_never_writes_the_original(tmp_path): + """Nothing restores an edited sweep, so the next campaign inherits it.""" + original = _sweep_dir(tmp_path) + task = {SOL_TRACK_FIELD: {"track": "gen", "workspace": "/w", "sweep": str(original)}} + ws = tmp_path / "ws" + ws.mkdir() + + adopted = sol_track.adopt_sweep(task, ws) + assert adopted == ws / "sweep" / "sweep.yaml" + assert task[SOL_TRACK_FIELD]["adopted_from"] == str(original) + # The generator beside the sweep travels with it: the harness imports + # the model plugin from that directory. + assert (ws / "sweep" / "gen_worker_config.py").is_file() + + tuning = tmp_path / "t.yaml" + tuning.write_text(yaml.safe_dump({"moe_config": {"backend": "TRTLLM"}}), encoding="utf-8") + sol_track.apply_overlay(task, tuning) + assert "gen_extra_llm_api" not in yaml.safe_load(original.read_text()) + assert yaml.safe_load((ws / "sweep" / "sweep.yaml").read_text())["gen_extra_llm_api"] + + +def test_a_resumed_run_keeps_measuring_what_it_started_with(tmp_path): + task = { + SOL_TRACK_FIELD: {"track": "gen", "workspace": "/w", "sweep": str(_sweep_dir(tmp_path))} + } + ws = tmp_path / "ws" + ws.mkdir() + sol_track.adopt_sweep(task, ws) + (ws / "sweep" / "sweep.yaml").write_text( + yaml.safe_dump({**GEN_SWEEP, "gen_extra_llm_api": {"mid": "flight"}}), encoding="utf-8" + ) + sol_track.adopt_sweep(dict(task), ws) + assert yaml.safe_load((ws / "sweep" / "sweep.yaml").read_text())["gen_extra_llm_api"] == { + "mid": "flight" + } + + +def test_a_resume_does_not_overwrite_where_the_copy_came_from(tmp_path): + """``adopted_from`` is what ``--clean`` restores from, so it must survive. + + A run resumed with ``--task /task.yaml`` reads a spec whose + ``sweep`` is already the copy. Recording that path as the origin would + leave the workspace pointing only at itself, and nothing able to say + which file the campaign was actually cut from. + """ + original = _sweep_dir(tmp_path) + task = {SOL_TRACK_FIELD: {"track": "gen", "workspace": "/w", "sweep": str(original)}} + ws = tmp_path / "ws" + ws.mkdir() + sol_track.adopt_sweep(task, ws) + + resumed = {SOL_TRACK_FIELD: dict(task[SOL_TRACK_FIELD])} # as the workspace spec reads + sol_track.adopt_sweep(resumed, ws) + assert resumed[SOL_TRACK_FIELD][sol_track.ADOPTED_FROM_KEY] == str(original) + + +def test_clean_leaves_the_spec_pointing_at_a_copy_that_is_gone(tmp_path): + """And the origin is what makes that recoverable rather than fatal. + + ``--clean`` removes the adopted copy -- it has to, or the "fresh" run + re-adopts one ``apply_overlay`` has already written the previous + campaign's accepted tuning into. A spec that named the copy then + dangles, and ``adopted_from`` is the only thing that knows where to + look instead. + """ + import shutil + + original = _sweep_dir(tmp_path) + task = {SOL_TRACK_FIELD: {"track": "gen", "workspace": "/w", "sweep": str(original)}} + ws = tmp_path / "ws" + ws.mkdir() + sol_track.adopt_sweep(task, ws) + + shutil.rmtree(ws / "sweep") # what --clean does + resumed = {SOL_TRACK_FIELD: dict(task[SOL_TRACK_FIELD])} + adopted = sol_track.adopt_sweep(resumed, ws) + + assert adopted == ws / "sweep" / "sweep.yaml" + assert resumed[SOL_TRACK_FIELD][sol_track.ADOPTED_FROM_KEY] == str(original) + # Said out loud: the re-adopted original may have moved on since the + # copy was taken, and a reader must not have to already know. + assert "--clean" in resumed[SOL_TRACK_FIELD][sol_track.READOPTED_KEY] + + +def test_a_sweep_that_is_simply_missing_is_not_readopted(tmp_path): + """The fallback is for a deleted COPY, not for any unreadable path.""" + task = { + SOL_TRACK_FIELD: { + "track": "gen", + "workspace": "/w", + "sweep": str(tmp_path / "nowhere" / "sweep.yaml"), + sol_track.ADOPTED_FROM_KEY: str(tmp_path / "also-nowhere" / "sweep.yaml"), + } + } + ws = tmp_path / "ws" + ws.mkdir() + assert sol_track.adopt_sweep(task, ws) is None + + +# ------------------------------------------------------------------ the prompt + + +def _sections(): + from agent_flow.workflows.perf_optimize.prompts._common import SOL_TRACK_CTX, SOL_TRACK_GEN + + return SOL_TRACK_CTX, SOL_TRACK_GEN + + +def test_neither_track_is_told_to_hand_write_its_result(): + """The code that lands a score and the prompt that asks for it must agree. + + `collect` grew a ctx path once while the ctx prompt still said "write + it down yourself", and the first real ctx campaign duly hand-wrote the + JSON -- correctly, as it happened, by the route the gen track had + already stopped using. + """ + for section in _sections(): + assert "--collect" in section + assert "Do not hand-write that JSON" in section + + +def test_only_the_track_that_has_a_frontier_is_told_to_build_one(): + """`process frontier` is GEN-only and answers with an error otherwise. + + A ctx agent told to run it would spend a step learning that, on a run + that measured perfectly. + """ + ctx, gen = _sections() + assert "process frontier" in gen + assert "ctx_json" in gen + assert "There is no `process frontier` on this track" in ctx + + +def test_both_gen_scorers_are_given_with_the_rule_for_choosing(): + """One command per declaration, and neither is the other's fallback. + + The gen track can be scored with or without an anchor. Naming only the + frontier stranded an unanchored campaign; naming both without the rule + would invite an agent to reach for a frontier by supplying an anchor + nobody declared, which is the one mistake whose output looks right. + """ + _, gen = _sections() + assert "ibc-bench process frontier" in gen + assert bench_cli.GEN_ONLY_MODULE in gen + assert "Never run (a) with an anchor `task.yaml` did not declare" in gen + + +def test_a_missing_gen_only_csv_names_the_command_that_writes_it(): + """The step the agent skipped, not the file the reader wanted. + + The extractor has no `ibc-bench` subcommand, so an agent that only + knows the CLI has no way to guess it -- the error has to carry it. + """ + with pytest.raises(bench_cli.BenchCliError, match=bench_cli.GEN_ONLY_MODULE): + bench_cli.gen_only_points(Path("/nonexistent-run-dir")) + + +def test_the_gate_is_stated_to_be_the_same_number_either_way(): + """Otherwise the two paths read as two different metrics. + + They are one formula -- `accept_rate / avg_iteration_time` -- and what + the anchor adds is only the end-to-end view on top. + """ + _, gen = _sections() + assert "accept_rate / avg_iteration_time" in gen + assert "ABSENT" in gen + assert "flattering" in gen + + +def test_the_work_dir_is_derived_and_the_prompt_says_why(): + """`-w` is not a free choice, and the reason is not obvious.""" + for section in _sections(): + assert "carries no code identity" in section + assert "the second overwrites the first" in section + + +def test_attribution_is_stated_as_the_reader_s_job(): + """Nothing in this stack fingerprints a configuration any more. + + So the obligation the tool used to discharge is named, with the + artifact that discharges it -- rather than left implied and lost. + """ + for section in _sections(): + assert 'Never infer "it took effect" from "the number moved."' in section + assert "gen_config.yaml" in section + + +def test_the_skills_the_config_repo_ships_are_named_where_they_apply(): + """They carry knowledge that is in no `--help`. + + `sol-postprocess`'s first step is not an action but a check: that the + run's archived sweep carries a measured `accept_rate`, because the + built-in table it replaced was 15 % high on one model and multiplies + the metric linearly with no symptom. + """ + ctx, gen = _sections() + assert "sol-postprocess" in gen + assert "15 %" in gen + for section in (ctx, gen): + assert "check-job" in section + + +def test_an_mtp0_sweep_needs_no_measured_accept_rate(tmp_path): + """At mtp 0 the acceptance length is 1.0 by definition. + + Requiring a measured one asks for a measurement of a constant. The + harness draws the same line: its post-processor exits 1 for a missing + rate only when the sweep has mtp>0 cases. A design sweep is mtp0-only, so + every campaign cut from one hit this. + """ + sweep = dict(GEN_SWEEP) + sweep.pop("accept_rate", None) + sweep["gen_configs"] = [[1, 1, 4, 64, 64, False, "0.9", 0, 0, "1,32"]] + data = task_schema.load_and_validate_task_yaml( + _write_task(tmp_path, {SOL_TRACK_FIELD: _block(tmp_path, sweep=sweep)}) + ) + assert data["optimize"]["target_metric"] == "throughput_per_user" + + +def test_a_speculating_sweep_still_needs_one(tmp_path): + """Where it is a multiplier, a wrong value is invisible.""" + sweep = dict(GEN_SWEEP) + sweep.pop("accept_rate", None) + sweep["gen_configs"] = [[1, 1, 4, 64, 256, False, "0.9", 3, 0, "1,32"]] + with pytest.raises(task_schema.TaskSchemaError, match="accept_rate"): + task_schema.load_and_validate_task_yaml( + _write_task(tmp_path, {SOL_TRACK_FIELD: _block(tmp_path, sweep=sweep)}) + ) + + +def test_the_refusal_names_the_key_that_actually_clears_it(tmp_path): + """It said ``options.accept_rate``; nothing reads there. + + ``sweep_accept_rate`` reads the sweep's TOP LEVEL, where the harness has + read it since v0.4.7. A message that names a key the reader does not + read sends someone to add one and re-run into the same refusal, which is + worse than a vague message: it is a confident wrong instruction. + """ + sweep = dict(GEN_SWEEP) + sweep.pop("accept_rate", None) + sweep["gen_configs"] = [[1, 1, 4, 64, 256, False, "0.9", 3, 0, "1,32"]] + with pytest.raises(task_schema.TaskSchemaError) as caught: + task_schema.load_and_validate_task_yaml( + _write_task(tmp_path, {SOL_TRACK_FIELD: _block(tmp_path, sweep=sweep)}) + ) + message = str(caught.value) + assert "top-level 'accept_rate'" in message + assert "options.accept_rate" not in message + # And the fix, spelled: the shape the harness expects. + assert "rate:" in message and "source:" in message diff --git a/agent-flow/tests/workflows/perf_optimize/test_spawn.py b/agent-flow/tests/workflows/perf_optimize/test_spawn.py new file mode 100644 index 000000000000..c716a4331488 --- /dev/null +++ b/agent-flow/tests/workflows/perf_optimize/test_spawn.py @@ -0,0 +1,127 @@ +"""The package's second subprocess door, and how narrow it is.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.perf_optimize import disagg_sol, spawn + +POINT = {"shape": "tep_4_eplb0_mtp3", "concurrency": 1, "prefer": "interactive"} + + +def _launch(tmp_path, track="gen") -> disagg_sol.CampaignLaunch: + return disagg_sol.CampaignLaunch( + track=track, + spec={"checkpoint_path": "/ckpt", "sol_track": {"track": track, "sweep": "/s.yaml"}}, + workspace=tmp_path / f"ws-{track}", + task_path=tmp_path / f"ws-{track}" / "task.yaml", + ) + + +def test_what_a_campaign_runs_under_is_on_disk_before_it_runs(tmp_path): + """Reviewable before any process exists -- and a dry run is the same path.""" + launch = _launch(tmp_path) + path = spawn.materialize(launch) + assert path.is_file() + assert yaml.safe_load(path.read_text())["sol_track"]["track"] == "gen" + + +def test_starting_without_materializing_is_refused(tmp_path): + """Otherwise the child reads a task.yaml nobody wrote.""" + with pytest.raises(disagg_sol.DisaggSolError, match="materialize"): + spawn.start(_launch(tmp_path)) + + +def test_only_perf_optimize_is_ever_started(tmp_path): + """The door admits one kind of process, and argv is built not interpolated.""" + launch = _launch(tmp_path) + assert launch.argv[0] == "perf-optimize" + assert all(isinstance(a, str) for a in launch.argv) + assert not any(";" in a or "|" in a or "&" in a for a in launch.argv) + + +def test_a_spec_that_cannot_be_written_stops_the_run_before_anything_starts(tmp_path, monkeypatch): + """Half-started is the one state a supervisor must not produce. + + The campaign that did start would claim a checkout the other needed, and + the failure would surface as a repo-claim refusal minutes later, pointing + at the wrong cause. + """ + good, bad = _launch(tmp_path, "gen"), _launch(tmp_path, "ctx") + started: list[str] = [] + monkeypatch.setattr(spawn, "start", lambda launch, env=None: started.append(launch.track)) + + real_materialize = spawn.materialize + + def explode(launch): + if launch.track == "ctx": + raise OSError("read-only filesystem") + return real_materialize(launch) + + monkeypatch.setattr(spawn, "materialize", explode) + with pytest.raises(OSError): + spawn.start_all([good, bad]) + assert started == [] + + +def test_every_campaign_is_waited_on_even_after_one_fails(): + """Every campaign is joined even after one has already failed. + + The other is still running on a cluster; abandoning it strands an + allocation and a claimed checkout with nothing tracking either. + """ + + class Fake: + def __init__(self, code): + self.code = code + self.waited = False + + def wait(self): + self.waited = True + return self.code + + a, b = Fake(1), Fake(0) + launches = [ + disagg_sol.CampaignLaunch("ctx", {}, Path("/w1"), Path("/w1/t.yaml")), + disagg_sol.CampaignLaunch("gen", {}, Path("/w2"), Path("/w2/t.yaml")), + ] + assert spawn.wait_all([(launches[0], a), (launches[1], b)]) == {"ctx": 1, "gen": 0} + assert a.waited and b.waited + + +def test_a_design_that_could_not_start_says_why(tmp_path, monkeypatch): + """A campaign reports through its workspace; a failed design has none. + + Discarding the pipe left the caller with a later refusal saying the + design was never established -- pointing at the design rather than at + whatever stopped it. The first real run of this function returned 1 + because the agent's credentials had expired, and the message saying so + went to /dev/null. + """ + import subprocess as sp + + class Done: + returncode = 1 + stdout = "" + stderr = "Anthropic profile login expired - Run /login to use your account\n" + + monkeypatch.setattr(sp, "run", lambda *a, **k: Done()) + with pytest.raises(disagg_sol.DisaggSolError, match="login expired"): + spawn.design("do the design", cwd=tmp_path) + + +def test_the_design_agent_s_output_is_kept_when_a_log_is_named(tmp_path, monkeypatch): + import subprocess as sp + + class Done: + returncode = 0 + stdout = "phase 0 done\n" + stderr = "" + + monkeypatch.setattr(sp, "run", lambda *a, **k: Done()) + log = tmp_path / "logs" / "design.log" + assert spawn.design("go", cwd=tmp_path, log=log) == 0 + assert "phase 0 done" in log.read_text() diff --git a/agent-flow/tests/workflows/perf_optimize/test_sweep_design.py b/agent-flow/tests/workflows/perf_optimize/test_sweep_design.py new file mode 100644 index 000000000000..9ac823f8fe20 --- /dev/null +++ b/agent-flow/tests/workflows/perf_optimize/test_sweep_design.py @@ -0,0 +1,117 @@ +"""Deciding and checking what the design agent runs, without running it.""" + +from __future__ import annotations + +from agent_flow.workflows.perf_optimize import sweep_design + +# ------------------------------------------------------------ the commands + + +def test_the_scored_command_is_the_anchor_free_one(tmp_path): + """The skill ends Phase 3 with the join; the staged scope has none. + + Generated rather than instructed, because "use the other post-processor" + is the kind of instruction honoured on the first run and forgotten on the + second -- and the failure is a frontier that looks like a frontier. + """ + command = sweep_design.postprocess_command(tmp_path / "run") + assert "get_gen_only_perf" in command + assert "process frontier" not in command + assert "ctx_json" not in command + + +# ---------------------------------------------------------- the instruction + + +def _instruction(tmp_path): + return sweep_design.designer_instruction( + model_dir="deepseek-V4-Pro", + design_dir=tmp_path / "sweep_design", + tracks=["ctx", "gen"], + ) + + +def test_the_scored_command_is_handed_over_verbatim_not_described(tmp_path): + """A described command is one an agent can restate differently later. + + The session this was written after did exactly that: the prompt stated + the scope in prose and the agent ran the joined post-processor anyway. + """ + text = _instruction(tmp_path) + assert "get_gen_only_perf" in text + assert "Do **not** run `ibc-bench process frontier`" in text + assert "--ctx_json" in text # ...named, in the prohibition + + +def test_the_reason_the_forbidden_command_is_dangerous_is_given(tmp_path): + """It would still emit a curve, and the curve would still look correct.""" + text = _instruction(tmp_path) + assert "still look" in text and "correct" in text + + +def test_the_agent_is_told_not_to_select_the_point(tmp_path): + """Selection needs a stated preference the design agent was never given.""" + text = _instruction(tmp_path) + assert "Do not select the operating point" in text + + +def test_a_failed_shape_must_be_reported_not_omitted(tmp_path): + text = _instruction(tmp_path) + assert "never omitted" in text + assert "a frontier it was never on" in text + + +def test_infrastructure_failure_and_a_memory_wall_are_kept_apart(tmp_path): + """One is retried, the other is a finding -- conflating them loses a shape.""" + text = _instruction(tmp_path) + assert "never started" in text + assert "different findings" in text + + +def test_the_skill_s_own_output_placement_rule_is_deferred_to(tmp_path): + """The rule an earlier generate-the-commands version had no place for. + + `SKILL.md` writes generated YAMLs into `sweep_design/` for an existing + model dir, never on top of the curated ones. Wrapping its scripts took + the output path as a free parameter with no guard, trading a + recomputable mistake for an unrecoverable one. + """ + text = _instruction(tmp_path) + assert "never on" in text and "curated" in text + assert "sweep_design/" in text + + +def test_the_agent_is_told_to_invoke_the_skill_not_handed_its_commands(tmp_path): + """The agent invokes the skill rather than being handed its commands. + + The phases are not four commands: they are also gates, a resume spine + and an output-placement rule. + """ + text = _instruction(tmp_path) + assert "invoke the **create-sweep** skill and" in text + assert "design_sweep.py" not in text + assert "model_facts.py" not in text + + +def test_the_irreversible_phases_are_excluded_with_the_reason(tmp_path): + text = _instruction(tmp_path) + assert "irreversible" in text + assert "accept rate" in text + + +def test_the_skill_s_submission_gates_are_answered_in_advance(tmp_path): + """A non-interactive design has no turn in which an answer could arrive. + + `SKILL.md` gates every submission on user confirmation "unless the user + said to run end-to-end". Without that clause the agent does exactly what + it is told: it prints the case count and node estimate, asks, and the + `--print` session ends -- so the design terminates having measured + nothing, and the caller sees only "the design was never established". + Observed on the first real run: 38 jobs / 145 nodes reported, then exit 0. + """ + text = _instruction(tmp_path) + assert "Run END-TO-END" in text + assert "this instruction is the confirmation" in text.lower() + assert "no second turn" in text + # ...and it still asks for the estimate, which is the useful half of a gate + assert "node estimate" in text diff --git a/agent-flow/tests/workflows/perf_optimize/test_workflow.py b/agent-flow/tests/workflows/perf_optimize/test_workflow.py index dd1ba5c17178..f92a06f1d6fa 100644 --- a/agent-flow/tests/workflows/perf_optimize/test_workflow.py +++ b/agent-flow/tests/workflows/perf_optimize/test_workflow.py @@ -2710,6 +2710,29 @@ def test_clean_wipes_managed_files_and_dirs(tmp_path): workflow.close() +def test_clean_removes_the_adopted_sweep_copy(tmp_path): + """Or the "fresh" run measures the previous campaign's accepted tuning. + + ``adopt_sweep`` copies the sweep DIRECTORY into ``/sweep/`` + and then leaves an existing copy alone, which is what a resume needs. + ``apply_overlay`` writes each attempt's tuning into that copy. So a copy + that survives ``--clean`` carries the last accepted optimization into + the next campaign's baseline -- measured with a change applied, reported + as the sweep's own. Copying into the workspace exists to prevent exactly + that, and ``--clean`` is named as the thing that starts over. + """ + adopted = tmp_path / "sweep" + adopted.mkdir() + (adopted / "sweep.yaml").write_text("gen_extra_llm_api: {carried: over}\n", encoding="utf-8") + (tmp_path / "rounds").mkdir() + + workflow = Workflow(workspace=tmp_path, clean=True) + try: + assert not adopted.exists() + finally: + workflow.close() + + # --------------------------------------------------------------- agent wiring