From 72998560b2ed18e12b66301a8df1536afee3946f Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:25:58 -0700 Subject: [PATCH 01/13] [TRTLLM-16460][fix] agent_flow: the CLI agent backend launches a CLI the organization accepts The agent SDK launches its bundled CLI binary, which only moves when the SDK is upgraded; an organization that requires a newer CLI then refuses every in-process turn at start-up ("older than the minimum version required by your organization"). A layer now launches the CLI named by the backend's CLI-path environment variable, or else the CLI on PATH when it is newer than the bundled one, and otherwise keeps the SDK's choice. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- agent-flow/agent_flow/backends/claude_code.py | 117 ++++++++++++++++-- agent-flow/tests/test_backends.py | 74 +++++++++++ 2 files changed, 184 insertions(+), 7 deletions(-) diff --git a/agent-flow/agent_flow/backends/claude_code.py b/agent-flow/agent_flow/backends/claude_code.py index 2203a6e6c326..213907574a17 100644 --- a/agent-flow/agent_flow/backends/claude_code.py +++ b/agent-flow/agent_flow/backends/claude_code.py @@ -6,11 +6,12 @@ import subprocess import sys import tempfile -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, contextmanager +from contextvars import ContextVar from importlib.metadata import PackageNotFoundError from importlib.metadata import version as _pkg_version from pathlib import Path -from typing import Any, AsyncIterator +from typing import Any, AsyncIterator, Callable, Iterator import claude_agent_sdk from claude_agent_sdk import ( @@ -370,13 +371,114 @@ def _claude_sdk_version() -> str: return "" +#: A ``claude`` binary every layer launches instead of the SDK's own choice. +CLI_PATH_ENV = "AGENT_FLOW_CLAUDE_CLI" + + +def _version_of(cli_path: str) -> tuple[int, ...]: + """``claude --version`` of ``cli_path`` as a comparable tuple; empty when unreadable.""" + try: + out = subprocess.run( + [cli_path, "--version"], capture_output=True, text=True, timeout=_CLI_VERSION_TIMEOUT_S + ) + except (OSError, subprocess.SubprocessError): + return () + token = out.stdout.strip().split()[0] if out.returncode == 0 and out.stdout.strip() else "" + try: + return tuple(int(part) for part in token.split(".")) + except ValueError: + return () + + +def _cli_override() -> str | None: + """The ``claude`` binary layers launch, when it is not the SDK's bundled one. + + ``$AGENT_FLOW_CLAUDE_CLI`` names one outright. Otherwise the ``claude`` + on ``PATH`` is used when it is newer than the bundled binary: an + organization can require a minimum CLI version, and the bundled one + only moves when the SDK is upgraded, so it is refused while the + user's own CLI runs. ``None`` keeps the SDK's choice. + """ + named = os.environ.get(CLI_PATH_ENV, "").strip() + if named: + return named + global _CLI_OVERRIDE_CACHE + if _CLI_OVERRIDE_CACHE is None: + bundled = Path(claude_agent_sdk.__file__).parent / "_bundled" / "claude" + on_path = shutil.which("claude") + choice = "" + if on_path and bundled.is_file(): + path_version = _version_of(on_path) + if path_version and path_version > _version_of(str(bundled)): + choice = on_path + _CLI_OVERRIDE_CACHE = choice + return _CLI_OVERRIDE_CACHE or None + + +_CLI_OVERRIDE_CACHE: str | None = None + + +def launched_cli() -> tuple[str, str]: + """The ``claude`` binary a layer launches here, and its version; ``""`` for what is unknown.""" + return _find_claude_cli() or "", _claude_cli_version() + + +#: Given each line a CLI started in this context writes to stderr; see :func:`cli_stderr_to`. +_CLI_STDERR_LISTENER: ContextVar[Callable[[str], None] | None] = ContextVar( + "claude_cli_stderr_listener", default=None +) + + +@contextmanager +def cli_stderr_to(listener: Callable[[str], None]) -> Iterator[None]: + """Hand ``listener`` every stderr line of a CLI started inside the block, too. + + The lines still reach this process's stderr, and a listener set by an + enclosing block still gets them. A layer called in the block's thread + starts its CLI in the block's context, so the lines of one turn reach the + listener of that turn and of no other running beside it. A listener that + raises does not stop the CLI's stderr from being read. + """ + outer = _CLI_STDERR_LISTENER.get() + + def both(line: str) -> None: + if outer is not None: + _tell(outer, line) + _tell(listener, line) + + token = _CLI_STDERR_LISTENER.set(both) + try: + yield + finally: + _CLI_STDERR_LISTENER.reset(token) + + +def _tell(listener: Callable[[str], None], line: str) -> None: + try: + listener(line) + except Exception: # noqa: BLE001 - an observer cannot fail the CLI it reads + pass + + +def _forward_stderr(line: str) -> None: + """Print a CLI stderr line, so a start-up failure is diagnosable, and pass it to the listener.""" + print(f"[claude-cli stderr] {line}", file=sys.stderr, flush=True) + listener = _CLI_STDERR_LISTENER.get() + if listener is not None: + _tell(listener, line) + + def _find_claude_cli() -> str | None: - """Locate the ``claude`` CLI binary the SDK would actually launch. + """Locate the ``claude`` CLI binary the SDK will actually launch. - Mirrors the SDK's lookup order — bundled binary first, then ``PATH``. - Returns ``None`` when nothing is found so the caller can simply omit - the CLI version from the rendered string. + The override (:func:`_cli_override`) first, then the SDK's own lookup + order — bundled binary first, then ``PATH``. Returns ``None`` when + nothing is found so the caller can simply omit the CLI version from + the rendered string. """ + override = _cli_override() + if override: + return override bundled = Path(claude_agent_sdk.__file__).parent / "_bundled" / "claude" if bundled.is_file(): return str(bundled) @@ -576,6 +678,7 @@ async def create_client( cwd=work_dir, sandbox={"enabled": False}, permission_mode="bypassPermissions", + cli_path=_cli_override(), hooks=hooks, disallowed_tools=[ *(disallowed_tools or []), @@ -583,7 +686,7 @@ async def create_client( ], # Forward the CLI's stderr so startup failures are diagnosable # instead of vanishing with the subprocess. - stderr=lambda line: print(f"[claude-cli stderr] {line}", file=sys.stderr, flush=True), + stderr=_forward_stderr, ) async with ClaudeSDKClient(options=options) as sdk_client: diff --git a/agent-flow/tests/test_backends.py b/agent-flow/tests/test_backends.py index 3085f65ff244..0c4794f9a6e1 100644 --- a/agent-flow/tests/test_backends.py +++ b/agent-flow/tests/test_backends.py @@ -851,6 +851,13 @@ async def test_create_client_disables_bash_sandbox(self, monkeypatch): options = await self._capture_options(monkeypatch) assert options.sandbox == {"enabled": False} + async def test_create_client_forwards_the_clis_stderr(self, monkeypatch): + # Printed, and handed to the listener of the context the CLI started in. + from agent_flow.backends import claude_code + + options = await self._capture_options(monkeypatch) + assert options.stderr is claude_code._forward_stderr + async def test_create_client_uses_bypass_permission_mode(self, monkeypatch): # ``bypassPermissions`` is the SDK setting that suppresses every # permission prompt the CLI would otherwise raise. @@ -1137,3 +1144,70 @@ def test_spill_files_are_content_addressed(self, tmp_path): cc_mod.spill_system_prompt(prompt, tmp_path, threshold=1024) cc_mod.spill_system_prompt(prompt + "y", tmp_path, threshold=1024) assert len(list((tmp_path / ".agent-flow").glob("system-prompt-*.md"))) == 2 + + +def test_a_named_claude_cli_is_the_one_every_layer_launches(monkeypatch): + from agent_flow.backends import claude_code + + monkeypatch.setenv(claude_code.CLI_PATH_ENV, "/opt/claude/bin/claude") + assert claude_code._cli_override() == "/opt/claude/bin/claude" + assert claude_code._find_claude_cli() == "/opt/claude/bin/claude" + + +def test_the_cli_on_path_wins_only_when_newer_than_the_bundled_one(monkeypatch, tmp_path): + from agent_flow.backends import claude_code + + monkeypatch.delenv(claude_code.CLI_PATH_ENV, raising=False) + monkeypatch.setattr(claude_code.shutil, "which", lambda name: "/usr/local/bin/claude") + versions = {"/usr/local/bin/claude": (2, 1, 296)} + monkeypatch.setattr(claude_code, "_version_of", lambda path: versions.get(path, (2, 1, 259))) + monkeypatch.setattr(claude_code.Path, "is_file", lambda self: True) + monkeypatch.setattr(claude_code, "_CLI_OVERRIDE_CACHE", None) + assert claude_code._cli_override() == "/usr/local/bin/claude" + versions["/usr/local/bin/claude"] = (2, 1, 100) + monkeypatch.setattr(claude_code, "_CLI_OVERRIDE_CACHE", None) + assert claude_code._cli_override() is None + + +def test_cli_stderr_reaches_the_listeners_of_the_context_it_started_in(capsys): + from agent_flow.backends import claude_code + + outer: list[str] = [] + inner: list[str] = [] + with claude_code.cli_stderr_to(outer.append): + claude_code._forward_stderr("one") + with claude_code.cli_stderr_to(inner.append): + claude_code._forward_stderr("two") + claude_code._forward_stderr("three") + assert outer == ["one", "two"] and inner == ["two"] + err = capsys.readouterr().err + assert "[claude-cli stderr] one" in err and "[claude-cli stderr] three" in err + + +def test_a_cli_stderr_listener_is_its_own_threads_and_cannot_fail_the_reader(): + import threading + + from agent_flow.backends import claude_code + + mine: list[str] = [] + + def broken(line: str) -> None: + raise OSError("disk full") + + def elsewhere() -> None: + claude_code._forward_stderr("from another turn") + + with claude_code.cli_stderr_to(mine.append), claude_code.cli_stderr_to(broken): + thread = threading.Thread(target=elsewhere) + thread.start() + thread.join() + claude_code._forward_stderr("mine") + assert mine == ["mine"] + + +def test_the_launched_cli_is_the_override_with_its_version(monkeypatch): + from agent_flow.backends import claude_code + + monkeypatch.setenv(claude_code.CLI_PATH_ENV, "/opt/claude/bin/claude") + monkeypatch.setattr(claude_code, "_claude_cli_version", lambda: "2.1.296") + assert claude_code.launched_cli() == ("/opt/claude/bin/claude", "2.1.296") From f1d3904bbd29d34c4b0e0f57d382b95818a13bf6 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:25:59 -0700 Subject: [PATCH 02/13] [TRTLLM-16460][feat] trtllm_modeling_v2: the task spec, intake and configuration variants The spec a run starts from, validated before any agent is constructed: - spec/task_schema.py: the task.yaml schema with closed keys and batched errors: checkpoint, units and features, parallel members, the serving envelope, accuracy workloads and floors, runtime blocks, sessions, the launch environment, source embargo, agents and the perf block. - spec/intake/: `trtllm-modeling-v2 intake check|confirm|compile`. A user-facing task.yaml is checked against its checkpoint and device, confirmed (and the run refuses one changed since), and compiled with a per-cluster site.yaml into the resolved spec. - spec/variants.py: configuration variants, a unit's target served under another configuration, as members of their own. - task.example.yaml documents every key. - spec/intake/derive.py: the fast path. `--checkpoint --site site.yaml --goal "..."` derives the request from config.json, the site's device and defaults blocks and a one-line goal, compiles it with per-unit options, discriminators and per-member perf configs, and starts the run; its intake checks are the first part of phase 0, and confirmation is a sidecar file. - Task validation refuses a TMPDIR too long for the LLM API's unix socket. - `coverage.accept` names requested cells the run accepts leaving uncovered; `knowledge: production|measurement` chooses what earlier runs' knowledge applies; a storage variant boots at every member of its unit. - accuracy.anchor: pending (per unit or variant): the benchmark is no check, the logit gate is the oracle, and delivery stays incomplete until a written anchor is recorded. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../trtllm_modeling_v2/spec/__init__.py | 4 + .../spec/intake/__init__.py | 11 + .../trtllm_modeling_v2/spec/intake/checks.py | 417 +++ .../trtllm_modeling_v2/spec/intake/cli.py | 92 + .../trtllm_modeling_v2/spec/intake/compile.py | 480 ++++ .../trtllm_modeling_v2/spec/intake/derive.py | 663 +++++ .../trtllm_modeling_v2/spec/intake/freeze.py | 160 ++ .../trtllm_modeling_v2/spec/intake/schema.py | 269 ++ .../trtllm_modeling_v2/spec/schema_keys.py | 41 + .../trtllm_modeling_v2/spec/task_schema.py | 2477 +++++++++++++++++ .../trtllm_modeling_v2/spec/variants.py | 324 +++ .../trtllm_modeling_v2/task.example.yaml | 839 ++++++ .../workflows/trtllm_modeling_v2/__init__.py | 2 + .../workflows/trtllm_modeling_v2/conftest.py | 103 + .../trtllm_modeling_v2/test_intake.py | 370 +++ .../test_intake_fast_path.py | 408 +++ .../trtllm_modeling_v2/test_task_schema.py | 1104 ++++++++ .../trtllm_modeling_v2/test_variants.py | 171 ++ 18 files changed, 7935 insertions(+) create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/__init__.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/__init__.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/checks.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/cli.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/compile.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/derive.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/freeze.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/schema.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/schema_keys.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/task_schema.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/variants.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/task.example.yaml create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/__init__.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/conftest.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_intake.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_intake_fast_path.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_task_schema.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_variants.py diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/__init__.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/__init__.py new file mode 100644 index 000000000000..d41a53a2a781 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Requirements: the user-facing task, its checks and confirmation, and the resolved spec.""" diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/__init__.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/__init__.py new file mode 100644 index 000000000000..e9321e3b9b27 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/__init__.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Intake: the ``task.yaml`` a run starts from (see ``INTAKE.md``). + +- ``schema.py`` the user-facing ``task.yaml`` and the per-cluster ``site.yaml`` +- ``derive.py`` the fast path: a ``task.yaml`` derived from the checkpoint, the site and a goal +- ``checks.py`` feasibility checks against the checkpoint and the request +- ``freeze.py`` confirmation: a digest of what the user agreed to, kept beside the request +- ``compile.py`` ``task.yaml`` + ``site.yaml`` -> the resolved spec the pipeline runs +- ``cli.py`` ``trtllm-modeling-v2 intake derive|check|confirm|compile`` +""" diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/checks.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/checks.py new file mode 100644 index 000000000000..aa9deee54069 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/checks.py @@ -0,0 +1,417 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Feasibility checks: can this checkpoint, on this device, meet this request (INTAKE.md section 6). + +Each check reads the request and the checkpoint's ``config.json`` and says +``pass``, ``warn`` or ``block``. A ``block`` is an open question: the task +cannot be confirmed while one stands. +""" + +from __future__ import annotations + +import json +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping + +from .schema import ( + GOAL_ADD_FEATURES, + GOAL_ONBOARD, + GOAL_TUNE, + STATUS_OUT, + STATUS_REQUIRED, + feature_entry, + feature_status, + validate_task, +) + +#: Speculative algorithms whose draft model is a checkpoint of its own. +DRAFTER_ALGOS: tuple[str, ...] = ("eagle3", "dflash") + +PASS, WARN, BLOCK = "pass", "warn", "block" +_ARCH = re.compile(r"^sm_(\d+)$") + + +@dataclass(frozen=True) +class Finding: + level: str + area: str + detail: str + fix: str = "" + + def render(self) -> str: + fix = f" -- {self.fix}" if self.fix else "" + return f"[{self.level}] {self.area}: {self.detail}{fix}" + + +def read_config(checkpoint: str | Path) -> dict[str, Any] | None: + """The checkpoint's ``config.json``, its ``text_config`` lifted to the top; None if unreadable.""" + path = Path(str(checkpoint)) / "config.json" + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + if not isinstance(data, dict): + return None + # A multimodal checkpoint keeps the language model's shape under text_config. + if isinstance(data.get("text_config"), dict): + data = {**data["text_config"], **{k: v for k, v in data.items() if k != "text_config"}} + return data + + +def _config(task: Mapping[str, Any]) -> dict[str, Any] | None: + return read_config((task.get("model") or {}).get("checkpoint", "")) + + +#: Where a config states its multi-token-prediction layers, as ``onboard/survey.py`` reads them. +DRAFT_LAYER_KEYS: tuple[str, ...] = ( + "num_nextn_predict_layers", + "num_mtp_modules", + "mtp_num_hidden_layers", +) + + +def draft_layers(config: Mapping[str, Any]) -> int: + """How many draft (MTP) layers the checkpoint carries; 0 for none.""" + for key in DRAFT_LAYER_KEYS: + value = config.get(key) + if isinstance(value, int) and not isinstance(value, bool) and value > 0: + return value + return 0 + + +def _goals(task: Mapping[str, Any]) -> list[Finding]: + goals = set(task.get("goals") or ()) + out = [] + if not goals: + out.append(Finding(BLOCK, "goals", "no goal stated", "onboard, add_features and/or tune")) + if {GOAL_ONBOARD, GOAL_ADD_FEATURES} <= goals: + out.append( + Finding( + BLOCK, + "goals", + "onboard and add_features together", + "onboard builds the target; list the features under it", + ) + ) + if goals and GOAL_ONBOARD not in goals and not (task.get("model") or {}).get("base"): + out.append(Finding(BLOCK, "goals", "no onboard goal and no model.base to start from")) + if GOAL_TUNE in goals and not (task.get("perf") or task.get("workload")): + out.append(Finding(BLOCK, "goals", "tune needs a workload or perf.configs")) + return out + + +def _speculative(task: Mapping[str, Any], config: Mapping[str, Any]) -> list[Finding]: + status = feature_status(task, "speculative") + if status in (None, STATUS_OUT): + if draft_layers(config) and status is None: + return [ + Finding( + BLOCK, + "speculative", + "the checkpoint carries draft layers the request does not mention", + "say whether speculative decoding is required or out_of_scope", + ) + ] + return [] + algo = feature_entry(task["features"]["speculative"]).get("params", {}).get("algo") + dspark = any(k.startswith("dspark") for k in config) + if algo == "mtp" and dspark: + return [ + Finding( + BLOCK, + "speculative", + "the checkpoint's draft stages are a DSpark drafter (dspark_* keys)", + "use algo: dspark", + ) + ] + if algo in ("mtp", "dspark") and not draft_layers(config): + return [ + Finding( + BLOCK, + "speculative", + f"{algo} requested; config states no draft layers ({', '.join(DRAFT_LAYER_KEYS)})", + ) + ] + block = config.get("dspark_block_size") + draft_len = feature_entry(task["features"]["speculative"]).get("params", {}).get("draft_len") + if algo == "dspark" and block and draft_len != block: + return [ + Finding( + BLOCK, + "speculative", + f"the DSpark drafter proposes blocks of {block} tokens; draft_len is {draft_len}", + f"set features.speculative.params.draft_len: {block}", + ) + ] + if algo in DRAFTER_ALGOS: + drafter = feature_entry(task["features"]["speculative"]).get("params", {}).get("drafter") + if not drafter: + return [ + Finding( + BLOCK, + "speculative", + f"{algo} drafts with a checkpoint of its own, and none is named", + "set features.speculative.params.drafter (drafter= in --goal)", + ) + ] + if not (Path(str(drafter)) / "config.json").is_file(): + return [ + Finding(BLOCK, "speculative", f"the {algo} drafter {drafter} has no config.json") + ] + return [Finding(PASS, "speculative", f"{algo} drafter {drafter}")] + return [Finding(PASS, "speculative", f"{algo} weights present")] + + +def _multimodal(task: Mapping[str, Any], config: Mapping[str, Any]) -> list[Finding]: + status = feature_status(task, "multimodal") + towers = [k for k in ("vision_config", "audio_config") if k in config] + if status == STATUS_REQUIRED and not towers: + return [ + Finding( + BLOCK, "multimodal", "required, but the checkpoint has no vision or audio tower" + ) + ] + if status is None and towers: + return [ + Finding( + BLOCK, + "multimodal", + f"the checkpoint has {', '.join(towers)} the request does not mention", + "say whether multimodal is required or out_of_scope", + ) + ] + return [] + + +def _divisibility(task: Mapping[str, Any], config: Mapping[str, Any]) -> list[Finding]: + out = [] + heads = config.get("num_attention_heads") + kv_heads = config.get("num_key_value_heads") or heads + experts = ( + config.get("num_experts") + or config.get("n_routed_experts") + or config.get("num_local_experts") + ) + device = task.get("device") or {} + per_node = device.get("gpus_per_node") + capacity = per_node * (device.get("nodes") or 1) if isinstance(per_node, int) else None + for p in task.get("parallelism") or (): + kind, world = p["kind"], p["world"] + name = f"{kind}{world}" + if capacity is not None and world > capacity: + out.append( + Finding(BLOCK, "parallelism", f"{name}: {world} GPUs, the device has {capacity}") + ) + # Attention DP replicates attention; tp and tep split its heads. + if kind in ("tp", "tep") and heads and heads % world: + out.append( + Finding(BLOCK, "parallelism", f"{name}: {heads} heads do not split {world} ways") + ) + if kind in ("tp", "tep") and kv_heads and kv_heads % world and world % kv_heads: + out.append( + Finding( + BLOCK, "parallelism", f"{name}: {kv_heads} KV heads do not split {world} ways" + ) + ) + if kind in ("tep", "dep"): + if not experts: + out.append( + Finding(BLOCK, "parallelism", f"{name}: expert parallelism on a dense model") + ) + elif experts % world: + out.append( + Finding( + BLOCK, "parallelism", f"{name}: {experts} experts do not split {world} ways" + ) + ) + return out + + +def _length(task: Mapping[str, Any], config: Mapping[str, Any]) -> list[Finding]: + limit = config.get("max_position_embeddings") + if not limit: + return [Finding(WARN, "length", "config.json states no max_position_embeddings")] + scaling = config.get("rope_scaling") or {} + factor = float(scaling.get("factor", 1.0)) if isinstance(scaling, Mapping) else 1.0 + reach = int(limit * max(factor, 1.0)) if scaling else int(limit) + out = [] + workload = task.get("workload") or {} + wanted = workload.get("max_seq_len") or (workload.get("isl", 0) + workload.get("osl", 0)) + if wanted and wanted > reach: + out.append(Finding(BLOCK, "length", f"isl + osl = {wanted} exceeds the model's {reach}")) + if not wanted: + out.append( + Finding( + WARN, + "length", + f"no workload: the engine would size its cache for {limit} positions", + "state workload.isl/osl or max_seq_len", + ) + ) + long = (task.get("features") or {}).get("long_context") + if long and feature_entry(long).get("status") != STATUS_OUT: + length = (feature_entry(long).get("params") or {}).get("length", 0) + if length > reach: + out.append(Finding(BLOCK, "long_context", f"{length} exceeds the model's {reach}")) + return out + + +def _feature_members(task: Mapping[str, Any]) -> list[Finding]: + """Every member a feature is narrowed to is one of the run's parallel configurations.""" + configs = {f"{p['kind']}{p['world']}" for p in task.get("parallelism") or ()} or {"tp1"} + out = [] + for name, raw in (task.get("features") or {}).items(): + if name == "custom": + continue + entry = feature_entry(raw) + members = entry.get("members") + if members is None or entry.get("status") == STATUS_OUT: + continue + if not isinstance(members, list) or not members: + out.append(Finding(BLOCK, name, "'members' must be a non-empty list")) + continue + absent = [str(m) for m in members if str(m) not in configs] + if absent: + out.append( + Finding( + BLOCK, + name, + f"members {', '.join(absent)} are not among the run's " + f"{', '.join(sorted(configs))}", + "add them to parallelism, or narrow the feature to members the run has", + ) + ) + return out + + +def _reuse_snapshots(task: Mapping[str, Any], config: Mapping[str, Any]) -> list[Finding]: + """Block reuse on a hybrid checkpoint needs a Mamba state snapshot policy.""" + from ...run.preflight import hybrid_marker + + status = feature_status(task, "prefix_reuse") + if status in (None, STATUS_OUT): + return [] + marker = hybrid_marker(config) + if marker is None: + return [] + params = feature_entry(task["features"]["prefix_reuse"]).get("params") or {} + if params.get("snapshot_interval"): + return [ + Finding( + PASS, + "prefix_reuse", + f"hybrid checkpoint ({marker}): state snapshots every " + f"{params['snapshot_interval']} tokens", + ) + ] + return [ + Finding( + BLOCK, + "prefix_reuse", + f"a hybrid checkpoint ({marker}) with block reuse and no Mamba state snapshot " + f"policy: the engine turns reuse off, and every gate fails on its warning", + "set features.prefix_reuse.params.snapshot_interval (snapshot= in --goal), " + "or make prefix_reuse out_of_scope", + ) + ] + + +def _quant_arch(task: Mapping[str, Any], config: Mapping[str, Any]) -> list[Finding]: + match = _ARCH.match(str((task.get("device") or {}).get("arch", ""))) + algo = str(((config.get("quantization_config") or {}).get("quant_algo") or "")).lower() + if match and "fp4" in algo and int(match.group(1)) < 100: + return [Finding(BLOCK, "quantization", f"{algo} needs sm_100 or later")] + return [] + + +def _perf_configs(task: Mapping[str, Any]) -> list[Finding]: + out = [] + configs = {f"{p['kind']}{p['world']}" for p in task.get("parallelism") or ()} or {"tp1"} + for entry in (task.get("perf") or {}).get("configs") or (): + parallel = entry.get("parallel") + if parallel and parallel not in configs: + out.append( + Finding( + BLOCK, + "perf", + f"config {entry.get('name')!r} runs at {parallel}, which is not among the " + f"run's {', '.join(sorted(configs))}", + ) + ) + for name, on in (entry.get("features") or {}).items(): + if on in (True, "on") and feature_status(task, name) in (None, STATUS_OUT): + out.append( + Finding( + BLOCK, + "perf", + f"config {entry.get('name')!r} turns on {name}, which is not a capability", + ) + ) + return out + + +def run_checks(task: Mapping[str, Any]) -> list[Finding]: + """Every finding for ``task``. + + With shape errors, those and the open questions only: no other check is + meaningful on a malformed request, and an open question often says what + the shape is missing. + """ + errors = validate_task(task) + questions = task.get("open_questions") if isinstance(task, Mapping) else None + asked = ( + [Finding(BLOCK, "open question", str(q)) for q in questions or ()] + if isinstance(questions, list) + else [] + ) + if errors: + return [Finding(BLOCK, "schema", e) for e in errors] + asked + findings = _goals(task) + asked + config = _config(task) + if config is None: + findings.append( + Finding(BLOCK, "checkpoint", "model.checkpoint has no readable config.json") + ) + else: + findings.append( + Finding(PASS, "checkpoint", f"architecture {(config.get('architectures') or ['?'])[0]}") + ) + for check in ( + _speculative, + _multimodal, + _divisibility, + _length, + _quant_arch, + _reuse_snapshots, + ): + findings += check(task, config) + findings += _feature_members(task) + findings += _perf_configs(task) + return findings + + +def blocking(findings: list[Finding]) -> list[Finding]: + return [f for f in findings if f.level == BLOCK] + + +def render_findings(findings: list[Finding]) -> str: + """Every finding on its own line, blocks first.""" + order = {BLOCK: 0, WARN: 1, PASS: 2} + return "\n".join(f.render() for f in sorted(findings, key=lambda f: order.get(f.level, 3))) + + +__all__ = [ + "BLOCK", + "DRAFTER_ALGOS", + "DRAFT_LAYER_KEYS", + "Finding", + "PASS", + "WARN", + "blocking", + "draft_layers", + "read_config", + "render_findings", + "run_checks", +] diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/cli.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/cli.py new file mode 100644 index 000000000000..0c96f2018484 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/cli.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``trtllm-modeling-v2 intake derive|check|confirm|compile``. + +derive --checkpoint C --site S [--goal G] -o OUT derive a request (derive.py); exit 1 on a block +check TASK run the feasibility checks; exit 1 on any block +confirm TASK --by NAME record the confirmation in TASK.confirmed (refused while a block stands) +compile TASK --site S -o OUT write the resolved spec the pipeline runs +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import yaml + +from ..task_schema import TaskSchemaError +from .checks import blocking, run_checks +from .compile import compile_files +from .derive import REPORT_FILENAME, derive_task, parse_goal, render_report +from .freeze import confirm +from .schema import IntakeError + + +def _derive(args: argparse.Namespace) -> int: + try: + site = yaml.safe_load(args.site.read_text(encoding="utf-8")) or {} + derivation = derive_task(args.checkpoint, site, parse_goal(args.goal), site_path=args.site) + except IntakeError as exc: + print(f"error: {exc}") + return 2 + findings = run_checks(derivation.task) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(yaml.safe_dump(derivation.task, sort_keys=False), encoding="utf-8") + report = args.out.with_name(REPORT_FILENAME) + report.write_text(render_report(derivation, args.out, findings), encoding="utf-8") + for finding in findings: + print(finding.render()) + print(f"wrote {args.out} and {report}") + return 1 if blocking(findings) else 0 + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog="trtllm-modeling-v2 intake") + sub = parser.add_subparsers(dest="cmd", required=True) + derive = sub.add_parser("derive") + derive.add_argument("--checkpoint", required=True, type=Path) + derive.add_argument("--site", required=True, type=Path) + derive.add_argument("--goal", default=None) + derive.add_argument("-o", "--out", required=True, type=Path) + check = sub.add_parser("check") + check.add_argument("task", type=Path) + conf = sub.add_parser("confirm") + conf.add_argument("task", type=Path) + conf.add_argument("--by", required=True) + comp = sub.add_parser("compile") + comp.add_argument("task", type=Path) + comp.add_argument("--site", required=True, type=Path) + comp.add_argument("-o", "--out", required=True, type=Path) + args = parser.parse_args(argv) + + if args.cmd == "derive": + return _derive(args) + task = yaml.safe_load(args.task.read_text(encoding="utf-8")) + findings = run_checks(task) + if args.cmd == "check": + for finding in findings: + print(finding.render()) + return 1 if blocking(findings) else 0 + if args.cmd == "confirm": + if blocking(findings): + for finding in blocking(findings): + print(finding.render()) + print("not confirmed: answer the blocks above first") + return 1 + try: + stamp = confirm(args.task, args.by) + except IntakeError as exc: + print(f"error: {exc}") + return 2 + print(f"confirmed by {stamp['by']} at {stamp['at']} ({stamp['sha256'][:12]})") + return 0 + try: + compiled = compile_files(args.task, args.site, args.out) + except (IntakeError, TaskSchemaError) as exc: + print(f"error: {exc}") + return 2 + for warning in compiled.warnings: + print(f"warning: {warning}") + print(f"wrote {args.out}") + return 0 diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/compile.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/compile.py new file mode 100644 index 000000000000..1cca27dae7ba --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/compile.py @@ -0,0 +1,480 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``task.yaml`` + ``site.yaml`` -> the resolved spec the pipeline runs (INTAKE.md section 8). + +Deterministic: the same two files always compile to the same spec. Requirements +the pipeline cannot enforce yet are compiled into nothing and returned as +warnings, so the user sees which of their requirements no check will hold. + +Each assembly unit gets an LLM-API options file (:attr:`Compiled.llm_api_options`, +written by :func:`compile_files` under :data:`LLM_API_DIRNAME`): ``trtllm-eval`` +reads no runtime flags, so without one the accuracy gate measures the engine's +defaults rather than the configuration the other gates run, and a speculative +unit's accuracy is measured with no draft at all. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Mapping + +import yaml + +from ..task_schema import load_and_validate_task_yaml +from .schema import ( + GOAL_ADD_FEATURES, + GOAL_ONBOARD, + GOAL_TUNE, + INCREMENT_FEATURES, + PARALLEL_SEGMENT, + STATUS_NICE, + STATUS_OUT, + UNGATED_FEATURES, + IntakeError, + feature_entry, + validate_site, + validate_task, +) + +#: Core feature -> the runtime switch it sets. +_RUNTIME_SWITCH = { + "cuda_graph": "cuda_graph", + "overlap_scheduler": "overlap_scheduler", + "chunked_prefill": "chunked_prefill", + "prefix_reuse": "kv_cache_reuse", +} +REFERENCES_DIR = "tests/integration/defs/accuracy/references" +#: Where :func:`compile_files` writes each unit's LLM-API options, beside the spec. +LLM_API_DIRNAME = "llm_api" + +#: Speculative algorithm -> its reference rows' ``spec_dec_algo`` and the LLM +#: API's ``speculative_config.decoding_type`` (the two spell them alike). +SPEC_DEC_ALGO = {"mtp": "MTP", "eagle3": "Eagle3", "dflash": "DFlash", "dspark": "DSpark"} +#: Weight formats a reference row states no ``quant_algo`` for. +_UNQUANTIZED = frozenset({"", "none", "auto", "bf16", "bfloat16", "fp16", "float16", "fp32"}) +#: KV-cache dtype -> the reference rows' ``kv_cache_quant_algo``. +_KV_QUANT_ALGO = {"fp8": "FP8", "nvfp4": "NVFP4"} + + +@dataclass +class Compiled: + spec: dict[str, Any] + warnings: list[str] = field(default_factory=list) + #: Assembly unit -> the LLM-API options its gates run under. + llm_api_options: dict[str, dict[str, Any]] = field(default_factory=dict) + + +def unit_name(name: str, raw: Any) -> str | None: + """The assembly unit an intake feature builds, or None for a runtime switch or a note.""" + entry = feature_entry(raw) + if entry.get("status") == STATUS_OUT: + return None + if name == "speculative": + return str((entry.get("params") or {}).get("algo") or "") or None + return INCREMENT_FEATURES.get(name) + + +def _parallel(task: Mapping[str, Any]) -> dict[str, Any]: + """The parallel segment and the member set, per unit where a feature narrows it.""" + configs = [f"{p['kind']}{p['world']}" for p in task.get("parallelism") or []] + if not configs: + return {"parallel": "tp1"} + if len(configs) == 1: + return {"parallel": configs[0]} + narrowed: dict[str, list[str]] = {} + for name, raw in (task.get("features") or {}).items(): + unit = unit_name(name, raw) if name != "custom" else None + members = feature_entry(raw).get("members") if unit else None + if unit and isinstance(members, list) and members: + narrowed[unit] = [str(m) for m in members] + if not narrowed: + return {"parallel": "unified", "members": configs} + return {"parallel": "unified", "members": {"default": configs, **narrowed}} + + +def _params(task: Mapping[str, Any], name: str) -> dict[str, Any]: + raw = (task.get("features") or {}).get(name) + if raw is None or feature_entry(raw).get("status") == STATUS_OUT: + return {} + return dict(feature_entry(raw).get("params") or {}) + + +def _discriminators(task: Mapping[str, Any], built: list[str]) -> dict[str, Any] | None: + """The reference row each unit anchors to: the weight and KV formats, and the speculator. + + A reference row matches only its exact discriminator set, so a run that + names none anchors a quantized checkpoint to the bf16 row. Every unit from + the speculative one on runs with its draft, so they anchor to its row too. + """ + explicit = (task.get("accuracy") or {}).get("discriminators") + if explicit is not None: + return dict(explicit) + base: dict[str, Any] = {} + weights = str(((task.get("model") or {}).get("quantization") or {}).get("weights") or "") + if weights.strip().lower() not in _UNQUANTIZED: + base["quant_algo"] = weights.strip().upper() + kv = str(_params(task, "kv_cache").get("dtype") or "").lower() + if kv in _KV_QUANT_ALGO: + base["kv_cache_quant_algo"] = _KV_QUANT_ALGO[kv] + per_unit: dict[str, dict[str, Any]] = {} + algo = None + for unit in built: + algo = SPEC_DEC_ALGO.get(unit, algo) + if algo: + per_unit[unit] = {**base, "spec_dec_algo": algo} + if per_unit: + return {"default": base, **per_unit} + return base or None + + +def _llm_api_options( + task: Mapping[str, Any], built: list[str], runtime: Mapping[str, Any], max_batch_size: Any +) -> dict[str, dict[str, Any]]: + """Each unit's LLM-API options: the runtime switches, the KV cache, and its speculator. + + The capture list is stated here and nowhere in the runtime flags, so the + quickstart gates derive theirs from it and both launch paths capture one + list. A unit carries the speculator of every feature up to it, as its + target implements them all. + """ + kv_params = _params(task, "kv_cache") + kv: dict[str, Any] = {} + if str(kv_params.get("dtype") or "").lower() in _KV_QUANT_ALGO: + kv["dtype"] = str(kv_params["dtype"]).lower() + if kv_params.get("memory_fraction"): + kv["free_gpu_memory_fraction"] = float(kv_params["memory_fraction"]) + if "kv_cache_reuse" in runtime: + kv["enable_block_reuse"] = bool(runtime["kv_cache_reuse"]) + snapshot = _params(task, "prefix_reuse").get("snapshot_interval") + if snapshot and runtime.get("kv_cache_reuse", True): + kv["mamba_state_config"] = {"periodic_snapshot_interval": int(snapshot)} + base: dict[str, Any] = {} + if kv: + base["kv_cache_config"] = kv + if "chunked_prefill" in runtime: + base["enable_chunked_prefill"] = bool(runtime["chunked_prefill"]) + chunk = _params(task, "chunked_prefill").get("chunk_tokens") + if chunk: + base["max_num_tokens"] = int(chunk) + if runtime.get("overlap_scheduler") is False: + base["disable_overlap_scheduler"] = True + if runtime.get("cuda_graph"): + graphs: dict[str, Any] = {"enable_padding": True} + if max_batch_size: + graphs["max_batch_size"] = int(max_batch_size) + base["cuda_graph_config"] = graphs + elif "cuda_graph" in runtime: + base["cuda_graph_config"] = None + + speculative = _params(task, "speculative") + options = {"core": dict(base)} + draft: dict[str, Any] | None = None + for unit in built: + if unit in SPEC_DEC_ALGO: + draft = {"decoding_type": SPEC_DEC_ALGO[unit]} + draft["max_draft_len"] = int(speculative.get("draft_len") or 1) + if speculative.get("drafter"): + draft["speculative_model"] = str(speculative["drafter"]) + options[unit] = {**base, **({"speculative_config": dict(draft)} if draft else {})} + return options + + +def _features(task: Mapping[str, Any], warnings: list[str]) -> tuple[list[str], dict, str]: + """The pipeline's feature list, the runtime switches, and the scope.""" + features = task.get("features") or {} + built: list[str] = [] + runtime: dict[str, Any] = {} + scope = "full" + for name, raw in features.items(): + if name == "custom": + for item in raw or (): + warnings.append( + f"custom feature {item.get('name', '?')!r}: no gate checks it; it reaches " + "the agents as a note" + ) + continue + entry = feature_entry(raw) + status = entry.get("status") + if status == STATUS_NICE: + warnings.append( + f"{name}: nice_to_have is built as required (the pipeline has no optional units)" + ) + if name in _RUNTIME_SWITCH: + runtime[_RUNTIME_SWITCH[name]] = status != STATUS_OUT + elif name == "speculative" and status != STATUS_OUT: + built.append(entry["params"]["algo"]) + elif name in INCREMENT_FEATURES: + if status != STATUS_OUT: + built.append(INCREMENT_FEATURES[name]) + elif name == "multimodal": + scope = "text_only" + elif name in UNGATED_FEATURES and status != STATUS_OUT: + warnings.append(f"{name}: recorded; no gate checks it yet") + if features.get("kv_cache"): + dtype = (feature_entry(features["kv_cache"]).get("params") or {}).get("dtype") + if dtype: + runtime["kv_cache_dtype"] = dtype + return built, runtime, scope + + +#: Feature parameters the compiled spec carries; any other is recorded and warned. +_COMPILED_PARAMS: dict[str, tuple[str, ...]] = { + "kv_cache": ("dtype", "memory_fraction"), + "chunked_prefill": ("chunk_tokens",), + "prefix_reuse": ("snapshot_interval",), + "speculative": ("algo", "draft_len", "drafter"), +} + + +def _unconsumed(task: Mapping[str, Any], warnings: list[str]) -> None: + """Warn of each requirement the request states and the compiled spec does not carry.""" + for name, raw in (task.get("features") or {}).items(): + if name == "custom" or name in UNGATED_FEATURES: + continue + entry = feature_entry(raw) + if entry.get("status") == STATUS_OUT: + continue + for key in sorted(set(entry.get("params") or {}) - set(_COMPILED_PARAMS.get(name, ()))): + warnings.append(f"features.{name}.params.{key}: recorded; nothing compiles it yet") + if entry.get("accept"): + warnings.append(f"features.{name}.accept: recorded; the pipeline's gates decide") + if entry.get("members") and unit_name(name, raw) is None: + warnings.append( + f"features.{name}.members: {name} is a runtime switch of every member, not " + "a unit of its own, so it cannot be narrowed" + ) + accuracy = task.get("accuracy") or {} + if accuracy.get("reference"): + warnings.append( + "accuracy.reference: recorded; the anchor is the reference_yaml row, which a " + "measurement must be written into first" + ) + for key in sorted(set(accuracy.get("protocol") or {}) - {"chat_template"}): + warnings.append( + f"accuracy.protocol.{key}: recorded; only chat_template compiles " + "(to --apply_chat_template)" + ) + perf = task.get("perf") or {} + if perf.get("metric"): + warnings.append("perf.metric: recorded; the keep rule judges both headline metrics") + if perf.get("baseline") not in (None, "stock"): + warnings.append( + f"perf.baseline {perf['baseline']!r}: the tune stage measures against the stock " + "path or nothing; it compiles to reference: none" + ) + if task.get("deliverable"): + warnings.append("deliverable: recorded; delivery is the branch the run commits to") + model = task.get("model") or {} + for key in ("in_tree", "reference_impl"): + if model.get(key): + warnings.append(f"model.{key}: recorded; say how the agents use it in `notes`") + + +def compile_task(task: Mapping[str, Any], site: Mapping[str, Any]) -> Compiled: + """The resolved spec for ``task`` on ``site``. Raises :class:`IntakeError` on a bad shape.""" + errors = validate_task(task) + validate_site(site) + if errors: + raise IntakeError("\n".join(errors)) + warnings: list[str] = [] + goals = list(task["goals"]) + model = task["model"] + workload = task.get("workload") or {} + if GOAL_ADD_FEATURES in goals: + warnings.append( + "add_features: the pipeline has no entry from an existing target; the run " + "rebuilds core with the new features" + ) + if GOAL_ONBOARD not in goals and GOAL_ADD_FEATURES not in goals and GOAL_TUNE in goals: + warnings.append("tune only: run it with `trtllm-modeling-v2 tune` on model.base") + + built, runtime, scope = _features(task, warnings) + _unconsumed(task, warnings) + extra_flags: list[str] = [] + kv_dtype = runtime.pop("kv_cache_dtype", None) + concurrency = workload.get("concurrency") or [] + max_seq_len = workload.get("max_seq_len") or ( + workload["isl"] + workload["osl"] if "isl" in workload and "osl" in workload else None + ) + max_batch_size = workload.get("max_batch_size") or (max(concurrency) if concurrency else None) + if max_seq_len: + extra_flags.append(f"--max_seq_len={max_seq_len}") + if max_batch_size: + extra_flags.append(f"--max_batch_size={max_batch_size}") + if kv_dtype: + extra_flags.append(f"--kv_cache_dtype={kv_dtype}") + chunk = _params(task, "chunked_prefill").get("chunk_tokens") + if chunk: + extra_flags.append(f"--max_num_tokens={chunk}") + fraction = _params(task, "kv_cache").get("memory_fraction") + if fraction: + extra_flags.append(f"--kv_cache_fraction={fraction}") + + accuracy_in = task.get("accuracy") or {} + benchmark = accuracy_in.get("benchmark", "gsm8k") + accuracy: dict[str, Any] = { + "benchmark": benchmark, + "reference_yaml": accuracy_in.get("reference_yaml") or f"{REFERENCES_DIR}/{benchmark}.yaml", + } + if accuracy_in.get("model_key"): + accuracy["model_key"] = accuracy_in["model_key"] + if (accuracy_in.get("protocol") or {}).get("chat_template"): + accuracy["benchmark_args"] = ["--apply_chat_template"] + if max_seq_len: + accuracy["max_seq_len"] = max_seq_len + if max_batch_size: + accuracy["max_batch_size"] = max_batch_size + discriminators = _discriminators(task, built) + if discriminators: + accuracy["discriminators"] = discriminators + for key in ("boot", "witness"): + if accuracy_in.get(key): + accuracy[key] = dict(accuracy_in[key]) + + spec: dict[str, Any] = { + "checkpoint_path": model["checkpoint"], + "trtllm_repo_path": site["trtllm_repo_path"], + "arch": task["device"]["arch"], + **_parallel(task), + "features": built, + "scope": scope, + "runtime": {**runtime, "extra_flags": extra_flags}, + "environment": dict(site["environment"]), + "accuracy": accuracy, + } + if model.get("target"): + spec["target"] = dict(model["target"]) + constraints = task.get("constraints") or {} + if constraints.get("embargo"): + spec["source_embargo"] = list(constraints["embargo"]) + commit_policy = dict(site.get("commit_policy") or {}) + if constraints.get("base_ref"): + commit_policy["base_ref"] = constraints["base_ref"] + if commit_policy: + spec["commit_policy"] = commit_policy + if constraints.get("budget"): + warnings.append("constraints.budget: recorded; no run-wide budget enforces it yet") + if site.get("agents"): + spec["agents"] = dict(site["agents"]) + if task.get("notes"): + spec["notes"] = str(task["notes"]) + options = _llm_api_options(task, built, spec["runtime"], max_batch_size) + if GOAL_TUNE in goals: + draft = options[built[-1]].get("speculative_config") if built else None + spec["perf"] = _perf(task, site, warnings, draft) + return Compiled(spec=spec, warnings=warnings, llm_api_options=options) + + +def member_options(member: str) -> dict[str, Any]: + """The LLM-API parallelism one member (``tp4``, ``tep4``, ``dep4``) serves at.""" + match = PARALLEL_SEGMENT.match(str(member)) + if match is None: + return {} + kind, world = match.group(1), int(match.group(2)) + options: dict[str, Any] = {"tensor_parallel_size": world} + if kind in ("tep", "dep"): + options["moe_expert_parallel_size"] = world + if kind == "dep": + options["enable_attention_dp"] = True + return options + + +def _perf( + task: Mapping[str, Any], + site: Mapping[str, Any], + warnings: list[str], + draft: Mapping[str, Any] | None, +) -> dict[str, Any]: + """The tune stage's block: each config at its member's parallelism, with the speculator. + + The tune stage serves each config under the options it derives from the + workload, overlaid with the config's ``llm_api_options``; the parallelism + and the draft are stated there, or it measures TP1 without a draft. + ``draft`` is the last speculative unit's ``speculative_config``. Without + ``perf.configs``, one config per member of the run at the run's workload. + """ + perf_in = task.get("perf") or {} + site_perf = site.get("perf") or {} + perf: dict[str, Any] = { + "enabled": True, + "reference": "stock" if perf_in.get("baseline", "stock") == "stock" else "none", + **{k: v for k, v in site_perf.items()}, + } + configs = [] + for entry in perf_in.get("configs") or []: + config: dict[str, Any] = {"name": entry["name"], "workload": dict(entry["workload"])} + if entry.get("constraint"): + config["constraint"] = dict(entry["constraint"]) + options = member_options(entry["parallel"]) if entry.get("parallel") else {} + for name, on in (entry.get("features") or {}).items(): + if name == "speculative": + if on in (True, "on") and draft: + options["speculative_config"] = dict(draft) + else: + warnings.append( + f"perf config {entry['name']!r}: `features.{name}` is recorded; the tune " + f"stage serves {name} as the workload's derived options set it" + ) + options.update(entry.get("llm_api_options") or {}) + if options: + config["llm_api_options"] = options + configs.append(config) + if not configs: + workload = task.get("workload") or {} + shared = { + "isl": workload.get("isl", 1024), + "osl": workload.get("osl", 1024), + "concurrency": workload.get("concurrency") or [1], + } + members = [f"{p['kind']}{p['world']}" for p in task.get("parallelism") or []] or ["tp1"] + for member in members: + options = member_options(member) + if draft: + options["speculative_config"] = dict(draft) + configs.append({"name": member, "workload": dict(shared), "llm_api_options": options}) + perf["configs"] = configs + goal = perf_in.get("goal") + if isinstance(goal, Mapping): + compiled = {} + if goal.get("ratio_to_baseline"): + compiled["ratio_to_reference"] = goal["ratio_to_baseline"] + if goal.get("pct_of_sol"): + compiled["pct_of_sol"] = goal["pct_of_sol"] + if compiled: + perf["goal"] = compiled + return perf + + +def compile_files(task_path: Path, site_path: Path, out: Path) -> Compiled: + """Compile two files into ``out`` and validate it with the pipeline's own schema. + + Each unit's LLM-API options are written under ``/llm_api/`` + and named in the spec's ``accuracy.llm_api_options``. + """ + task = yaml.safe_load(Path(task_path).read_text(encoding="utf-8")) + site = yaml.safe_load(Path(site_path).read_text(encoding="utf-8")) + compiled = compile_task(task, site) + out.parent.mkdir(parents=True, exist_ok=True) + options_dir = out.parent.resolve() / LLM_API_DIRNAME + options_dir.mkdir(parents=True, exist_ok=True) + paths: dict[str, str] = {} + for unit, options in compiled.llm_api_options.items(): + path = options_dir / f"{unit}.yaml" + path.write_text(yaml.safe_dump(options, sort_keys=False), encoding="utf-8") + paths[unit] = str(path) + if paths: + compiled.spec["accuracy"]["llm_api_options"] = paths + out.write_text(yaml.safe_dump(compiled.spec, sort_keys=False), encoding="utf-8") + load_and_validate_task_yaml(out) + return compiled + + +__all__ = [ + "LLM_API_DIRNAME", + "SPEC_DEC_ALGO", + "Compiled", + "compile_files", + "compile_task", + "member_options", + "unit_name", +] diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/derive.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/derive.py new file mode 100644 index 000000000000..842e144e42ce --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/derive.py @@ -0,0 +1,663 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fast path: a ``task.yaml`` derived in seconds from the checkpoint, ``site.yaml`` and a goal. + +``trtllm-modeling-v2 --checkpoint --site site.yaml --workspace W [--goal ...]`` +reads the checkpoint's ``config.json`` (and ``hf_quant_config.json``), takes the +device, the default members and workload, the accuracy anchor, the boot prompts +and the witness declarations from ``site.yaml``, and the features, members and +workload from the goal, and writes the request to ``/task.intake.yaml`` +with every value's provenance, beside ``intake.report.md``. The run starts from it +at once; its intake checks run in phase 0, beside the survey and the substrate +gate (``run/stages/preflight.py``). What cannot be derived is an open question, +which those checks report as a block. + +A feature the checkpoint carries (draft layers, a vision or audio tower) is +built unless the goal puts it out of scope; anything else is built only when +the goal names it. The goal is one line of words, separated by spaces, commas +or semicolons: + +==================================== ============================================= +``mtp`` ``eagle3`` ``dflash`` speculative decoding with that algorithm +``dspark`` +``disagg`` ``multimodal`` ... a feature, by its INTAKE.md section 5 name +``@tp1/tep4`` the feature, gated on those members only +``-`` ``no-`` out of scope; ``text_only`` is ``-multimodal`` +``tp1`` ``tep4`` ``dep4`` ... the members (default: the site's, else tp1) +``isl=`` ``osl=`` ``max_seq_len=`` the workload; ``concurrency=1/16/64`` +``max_batch_size=`` +``tune`` tune the target once it is delivered +``arch=sm_100`` the device (default: the site's) +``benchmark=`` ``model_key=`` the accuracy anchor +``reference=`` +``draft_len=`` ``drafter=`` speculative decoding's depth and drafter +``kv=fp8`` ``memory=0.6`` the KV cache's dtype and memory fraction +``chunk=`` the prefill chunk (``max_num_tokens``) +``snapshot=`` a hybrid checkpoint's state snapshot interval +==================================== ============================================= +""" + +from __future__ import annotations + +import json +import re +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Mapping + +import yaml + +from .checks import draft_layers, read_config +from .compile import unit_name +from .freeze import confirmation_status +from .schema import ( + CORE_FEATURES, + GOAL_ONBOARD, + GOAL_TUNE, + INTAKE_VERSION, + KNOWN_FEATURES, + PARALLEL_SEGMENT, + SPECULATIVE_ALGOS, + STATUS_OUT, + STATUS_REQUIRED, + WORKLOAD_KEYS, + IntakeError, + validate_site, +) + +#: The derived request and its report, in the workspace. +REQUEST_FILENAME = "task.intake.yaml" +REPORT_FILENAME = "intake.report.md" + +#: Provenance of a value: stated in the goal, read from the checkpoint or the +#: site, or a default of this module. +USER, INFERRED, DEFAULT = "user", "inferred", "default" + +#: Members when neither the goal nor the site names any. +DEFAULT_MEMBERS: tuple[str, ...] = ("tp1",) +#: The workload when neither the goal nor the site states one. Stated rather +#: than left to the engine, which would size its cache for the checkpoint's +#: whole ``max_position_embeddings``. +DEFAULT_WORKLOAD: dict[str, int] = {"isl": 4096, "osl": 1024, "max_batch_size": 64} +#: Tokens between Mamba state snapshots on a hybrid checkpoint with block +#: reuse: without a snapshot policy the engine turns reuse off. +DEFAULT_SNAPSHOT_INTERVAL = 2048 +#: The boot gate's prompts when the site declares none: one whose answer any +#: working target states, checked, and one open continuation. +DEFAULT_BOOT: dict[str, Any] = { + "prompts": [ + {"prompt": "The capital of France is", "keywords_any": ["Paris"]}, + "In one sentence, photosynthesis is", + ], + "max_tokens": 32, +} +DEFAULT_BENCHMARK = "gsm8k" + +#: Pipeline feature -> the witness keys it cannot judge without, which nothing +#: can derive from the checkpoint. A replay witness's ``window`` and prompt +#: length are derived from ``sliding_window`` (:func:`_witness`). The +#: disaggregated witness needs none: its leg B is the flow's own +#: (``checks/harness/disagg_leg_b.py``) unless the site names a +#: ``contrast_command``, which then needs its ``contrast_env``. +WITNESS_REQUIRED: dict[str, tuple[str, ...]] = { + "decoder_replay": ("contrast_env",), + "encoder_replay": ("contrast_env",), + "kv_quant": ("expected_bytes_per_token", "contrast_env"), + "host_offload": ("contrast_env",), + "disagg": (), +} +#: A replay witness's long prompt (or shared prefix), in multiples of the window. +REPLAY_PROMPT_WINDOWS = 4 + +#: Goal word -> intake feature, for the words that are not a feature's own name. +_FEATURE_WORDS = {"disagg": "disaggregated", "spec": "speculative"} +#: ``key=value`` goal words -> (feature, parameter, type). +_PARAM_WORDS: dict[str, tuple[str, str, type]] = { + "draft_len": ("speculative", "draft_len", int), + "drafter": ("speculative", "drafter", str), + "kv": ("kv_cache", "dtype", str), + "memory": ("kv_cache", "memory_fraction", float), + "chunk": ("chunked_prefill", "chunk_tokens", int), + "snapshot": ("prefix_reuse", "snapshot_interval", int), +} +_ACCURACY_WORDS = { + "benchmark": "benchmark", + "model_key": "model_key", + "reference": "reference_yaml", +} +_KV_DTYPES = ("auto", "fp8", "nvfp4") +_REPO_NAME = re.compile(r"^[A-Za-z0-9][\w.-]*/[\w.-]+$") +#: The core switches every derived request turns on unless the goal turns one off. +_CORE_SWITCHES = tuple(name for name in CORE_FEATURES if name != "kv_cache") + + +@dataclass +class Goal: + """What a ``--goal`` line asks for.""" + + text: str = "" + #: Intake feature -> status, for every feature the goal names. + features: dict[str, str] = field(default_factory=dict) + #: The speculative algorithm the goal names. + algo: str | None = None + #: Intake feature -> the members the goal narrows it to. + feature_members: dict[str, list[str]] = field(default_factory=dict) + members: list[str] = field(default_factory=list) + workload: dict[str, Any] = field(default_factory=dict) + #: Intake feature -> parameters. + params: dict[str, dict[str, Any]] = field(default_factory=dict) + accuracy: dict[str, str] = field(default_factory=dict) + arch: str | None = None + tune: bool = False + + +def _feature(word: str) -> tuple[str, str | None] | None: + """``(intake feature, speculative algorithm)`` a goal word names, or None.""" + if word in SPECULATIVE_ALGOS: + return "speculative", word + name = _FEATURE_WORDS.get(word, word) + if name in KNOWN_FEATURES and name != "custom": + return name, None + return None + + +def parse_goal(text: str | None) -> Goal: + """Read a goal line (module docstring). Raises :class:`IntakeError` naming every unknown word.""" + goal = Goal(text=" ".join((text or "").split())) + errors: list[str] = [] + for word in re.split(r"[\s,;]+", goal.text): + if not word: + continue + if "=" in word: + _key_value(goal, word, errors) + continue + lowered = word.lower() + if lowered == GOAL_TUNE: + goal.tune = True + continue + if PARALLEL_SEGMENT.match(lowered): + goal.members.append(lowered) + continue + if lowered == "text_only": + goal.features["multimodal"] = STATUS_OUT + continue + negated = lowered.startswith(("-", "no-", "no_")) + name, _, members = ( + lowered.lstrip("-").removeprefix("no-").removeprefix("no_").partition("@") + ) + found = _feature(name) + if found is None: + errors.append(f"goal: {word!r} is not a feature, a member or key=value") + continue + feature, algo = found + goal.features[feature] = STATUS_OUT if negated else STATUS_REQUIRED + if algo and not negated: + goal.algo = algo + if members: + listed = [m for m in members.split("/") if m] + bad = [m for m in listed if not PARALLEL_SEGMENT.match(m)] + if bad or not listed: + errors.append(f"goal: {word!r} names no members of the form tp, tep, dep") + else: + goal.feature_members[feature] = listed + if errors: + raise IntakeError("\n".join(errors)) + return goal + + +def _key_value(goal: Goal, word: str, errors: list[str]) -> None: + key, _, value = word.partition("=") + key = key.lower() + try: + if key == "concurrency": + goal.workload[key] = [int(v) for v in value.split("/") if v] + elif key in WORKLOAD_KEYS: + goal.workload[key] = int(value) + elif key == "arch": + goal.arch = value.lower() + elif key in _ACCURACY_WORDS: + goal.accuracy[_ACCURACY_WORDS[key]] = value + elif key in _PARAM_WORDS: + feature, param, kind = _PARAM_WORDS[key] + parsed = kind(value.lower() if kind is str and key == "kv" else value) + if key == "kv" and parsed not in _KV_DTYPES: + errors.append(f"goal: kv={value} is not one of {list(_KV_DTYPES)}") + return + goal.params.setdefault(feature, {})[param] = parsed + else: + errors.append(f"goal: {key!r} in {word!r} is not a key the goal takes") + except ValueError: + errors.append(f"goal: {word!r} needs a number") + + +@dataclass(frozen=True) +class CheckpointFacts: + """What a derivation reads from a checkpoint directory.""" + + config: dict[str, Any] + architecture: str + #: The weight format and the KV-cache format the checkpoint was quantized to. + quant_algo: str | None + kv_cache_quant_algo: str | None + #: A quantization the checkpoint declares without a readable ``quant_algo``. + unread_quantization: str | None + moe: bool + #: ``mtp`` or ``dspark`` for draft layers in the checkpoint, and their depth. + speculator: str | None + draft_len: int | None + towers: tuple[str, ...] + #: What marks the checkpoint hybrid (``run/preflight.hybrid_marker``), or None. + hybrid: str | None + #: The ``/`` name the config states, the reference rows' key. + name: str | None + + +def inspect_checkpoint(path: Path) -> CheckpointFacts: + """Read ``path``'s config files. Raises :class:`IntakeError` without a ``config.json``.""" + from ...run.preflight import hybrid_marker + + config = read_config(path) + if config is None: + raise IntakeError(f"{path}: no readable config.json; this is not a HuggingFace checkpoint") + quant = config.get("quantization_config") + quant = dict(quant) if isinstance(quant, Mapping) else {} + try: + hf_quant = json.loads((Path(path) / "hf_quant_config.json").read_text(encoding="utf-8")) + quant = {**dict(hf_quant.get("quantization") or {}), **quant} + except (OSError, ValueError, AttributeError): + pass + quant_algo = quant.get("quant_algo") + method = str(quant.get("quant_method") or "").lower() + if not quant_algo and method == "fp8": + quant_algo = "FP8_BLOCK_SCALES" if quant.get("weight_block_size") else "FP8" + unread = None + if quant and not quant_algo: + unread = method or "a quantization_config without quant_algo" + speculator = None + draft_len = None + if any(str(k).startswith("dspark") for k in config): + speculator, draft_len = "dspark", config.get("dspark_block_size") + elif draft_layers(config): + speculator, draft_len = "mtp", draft_layers(config) + name = config.get("_name_or_path") + return CheckpointFacts( + config=config, + architecture=str((config.get("architectures") or ["?"])[0]), + quant_algo=str(quant_algo) if quant_algo else None, + kv_cache_quant_algo=str(quant.get("kv_cache_quant_algo") or "") or None, + unread_quantization=unread, + moe=bool( + config.get("num_experts") + or config.get("n_routed_experts") + or config.get("num_local_experts") + ), + speculator=speculator, + draft_len=int(draft_len) if isinstance(draft_len, int) and draft_len > 0 else None, + towers=tuple(k for k in ("vision_config", "audio_config") if k in config), + hybrid=hybrid_marker(config), + name=str(name) if isinstance(name, str) and _REPO_NAME.match(name) else None, + ) + + +@dataclass +class Derivation: + """A derived request, and what it holds that the goal did not state.""" + + task: dict[str, Any] + #: ``(field, value, source)`` for every value read from the checkpoint or + #: the site, or defaulted; ``source`` is ``checkpoint``, ``site`` or ``default``. + unstated: list[tuple[str, Any, str]] = field(default_factory=list) + + +class _Builder: + """Records each value's provenance as the request is assembled.""" + + def __init__(self) -> None: + self.provenance: dict[str, str] = {} + self.unstated: list[tuple[str, Any, str]] = [] + self.questions: list[str] = [] + + def note(self, where: str, value: Any, source: str) -> Any: + self.provenance[where] = ( + USER if source == USER else (DEFAULT if source == DEFAULT else INFERRED) + ) + if source != USER: + self.unstated.append((where, value, source)) + return value + + +def derive_task( + checkpoint: Path, + site: Mapping[str, Any], + goal: Goal, + *, + site_path: Path | None = None, + now: datetime | None = None, +) -> Derivation: + """The request ``goal`` asks of ``checkpoint`` on ``site``; see the module docstring.""" + errors = validate_site(site) + if errors: + raise IntakeError("\n".join(errors)) + checkpoint = Path(checkpoint).resolve() + facts = inspect_checkpoint(checkpoint) + defaults = dict(site.get("defaults") or {}) + b = _Builder() + + model: dict[str, Any] = {"checkpoint": str(checkpoint)} + model["architecture"] = b.note("model.architecture", facts.architecture, "checkpoint") + model["kind"] = b.note("model.kind", "moe" if facts.moe else "dense", "checkpoint") + if facts.quant_algo: + quantization = {"weights": facts.quant_algo} + if facts.kv_cache_quant_algo: + quantization["kv_cache"] = facts.kv_cache_quant_algo + model["quantization"] = b.note("model.quantization", quantization, "checkpoint") + elif facts.unread_quantization: + b.questions.append( + f"the checkpoint is quantized ({facts.unread_quantization}) and states no " + f"quant_algo: set accuracy.discriminators to the reference row's (quant_algo, " + f"kv_cache_quant_algo), so the anchor is not the unquantized model's" + ) + + task: dict[str, Any] = { + "version": INTAKE_VERSION, + "goals": [GOAL_ONBOARD, *([GOAL_TUNE] if goal.tune else [])], + "model": model, + "device": _device(goal, site, b), + "parallelism": _parallelism(goal, defaults, b), + "workload": _workload(goal, defaults, facts, b), + "features": _features(goal, facts, b), + "accuracy": {}, + } + task["accuracy"] = _accuracy(goal, defaults, facts, task["features"], b) + task["provenance"] = dict(sorted(b.provenance.items())) + task["open_questions"] = list(b.questions) + stamp = (now or datetime.now(timezone.utc)).strftime("%Y-%m-%dT%H:%M:%SZ") + task["derived"] = { + "checkpoint": str(checkpoint), + "site": str(Path(site_path).resolve()) if site_path else "", + "goal": goal.text, + "at": stamp, + } + return Derivation(task=task, unstated=b.unstated) + + +def _device(goal: Goal, site: Mapping[str, Any], b: _Builder) -> dict[str, Any]: + site_device = dict(site.get("device") or {}) + device: dict[str, Any] = {} + if goal.arch: + device["arch"] = b.note("device.arch", goal.arch, USER) + elif site_device.get("arch"): + device["arch"] = b.note("device.arch", site_device["arch"], "site") + else: + b.questions.append( + "device.arch: name the GPU architecture (arch=sm_100 in --goal, or `device: " + "{arch: ...}` in site.yaml); receipts are keyed by it and it is never guessed" + ) + for key in ("gpus_per_node", "nodes"): + if site_device.get(key): + device[key] = b.note(f"device.{key}", site_device[key], "site") + return device + + +def _parallelism(goal: Goal, defaults: Mapping[str, Any], b: _Builder) -> list[dict[str, Any]]: + if goal.members: + members, source = goal.members, USER + elif defaults.get("members"): + members, source = list(defaults["members"]), "site" + else: + members, source = list(DEFAULT_MEMBERS), DEFAULT + b.note("parallelism", list(dict.fromkeys(members)), source) + out = [] + for member in dict.fromkeys(members): + match = PARALLEL_SEGMENT.match(member) + out.append({"kind": match.group(1), "world": int(match.group(2))}) + return out + + +def _workload( + goal: Goal, defaults: Mapping[str, Any], facts: CheckpointFacts, b: _Builder +) -> dict[str, Any]: + site_workload = dict(defaults.get("workload") or {}) + if site_workload: + workload, source = site_workload, "site" + else: + workload, source = dict(DEFAULT_WORKLOAD), DEFAULT + reach = facts.config.get("max_position_embeddings") + if isinstance(reach, int) and workload["isl"] + workload["osl"] > reach: + # A default must fit the model; a stated workload that does not is a block. + workload["osl"] = min(workload["osl"], reach // 4) + workload["isl"] = reach - workload["osl"] + for key, value in workload.items(): + if key not in goal.workload: + b.note(f"workload.{key}", value, source) + for key, value in goal.workload.items(): + workload[key] = b.note(f"workload.{key}", value, USER) + return workload + + +def _features(goal: Goal, facts: CheckpointFacts, b: _Builder) -> dict[str, Any]: + features: dict[str, Any] = {} + + def entry(name: str, status: str, source: str) -> dict[str, Any]: + value: dict[str, Any] = {"status": b.note(f"features.{name}", status, source)} + params = dict(goal.params.get(name) or {}) + for key in params: + b.note(f"features.{name}.params.{key}", params[key], USER) + if params and status != STATUS_OUT: + value["params"] = params + if goal.feature_members.get(name) and status != STATUS_OUT: + value["members"] = b.note(f"features.{name}.members", goal.feature_members[name], USER) + features[name] = value + return value + + for name in _CORE_SWITCHES: + status = goal.features.get(name) + entry(name, status or STATUS_REQUIRED, USER if status else DEFAULT) + reuse = features["prefix_reuse"] + if facts.hybrid and reuse["status"] != STATUS_OUT: + params = reuse.setdefault("params", {}) + if "snapshot_interval" not in params: + params["snapshot_interval"] = b.note( + "features.prefix_reuse.params.snapshot_interval", DEFAULT_SNAPSHOT_INTERVAL, DEFAULT + ) + + kv_dtype = (goal.params.get("kv_cache") or {}).get("dtype") + native_kv = str(facts.kv_cache_quant_algo or "").lower() + if goal.features.get("kv_cache") or kv_dtype or native_kv in _KV_DTYPES: + stated = goal.features.get("kv_cache") or kv_dtype + kv = entry( + "kv_cache", + goal.features.get("kv_cache") or STATUS_REQUIRED, + USER if stated else "checkpoint", + ) + if not kv_dtype and native_kv in _KV_DTYPES and kv["status"] != STATUS_OUT: + kv.setdefault("params", {})["dtype"] = b.note( + "features.kv_cache.params.dtype", native_kv, "checkpoint" + ) + + spec_status = goal.features.get("speculative") + if spec_status == STATUS_OUT: + entry("speculative", STATUS_OUT, USER) + elif goal.algo or facts.speculator: + algo = goal.algo or facts.speculator + source = USER if goal.algo else "checkpoint" + spec = entry("speculative", STATUS_REQUIRED, source) + params = spec.setdefault("params", {}) + params["algo"] = b.note("features.speculative.params.algo", algo, source) + if "draft_len" not in params and facts.draft_len and algo == facts.speculator: + params["draft_len"] = b.note( + "features.speculative.params.draft_len", facts.draft_len, "checkpoint" + ) + elif spec_status: + b.questions.append( + "speculative: the goal asks for speculative decoding and the checkpoint carries " + "no draft layers; name the algorithm (mtp, eagle3, dflash or dspark)" + ) + + if "multimodal" in goal.features: + entry("multimodal", goal.features["multimodal"], USER) + elif facts.towers: + entry("multimodal", STATUS_REQUIRED, "checkpoint") + + handled = {*_CORE_SWITCHES, "kv_cache", "speculative", "multimodal"} + for name, status in goal.features.items(): + if name not in handled: + entry(name, status, USER) + return features + + +def _accuracy( + goal: Goal, + defaults: Mapping[str, Any], + facts: CheckpointFacts, + features: Mapping[str, Any], + b: _Builder, +) -> dict[str, Any]: + site_accuracy = dict(defaults.get("accuracy") or {}) + accuracy: dict[str, Any] = {} + for key, value in site_accuracy.items(): + accuracy[key] = b.note(f"accuracy.{key}", value, "site") + for key, value in goal.accuracy.items(): + accuracy[key] = b.note(f"accuracy.{key}", value, USER) + if "benchmark" not in accuracy: + accuracy["benchmark"] = b.note("accuracy.benchmark", DEFAULT_BENCHMARK, DEFAULT) + if "model_key" not in accuracy and facts.name: + accuracy["model_key"] = b.note("accuracy.model_key", facts.name, "checkpoint") + if defaults.get("boot"): + accuracy["boot"] = b.note("accuracy.boot", dict(defaults["boot"]), "site") + else: + accuracy["boot"] = b.note("accuracy.boot", json.loads(json.dumps(DEFAULT_BOOT)), DEFAULT) + witness = _witness(defaults, facts, features, b) + if witness: + accuracy["witness"] = witness + return accuracy + + +def _witness( + defaults: Mapping[str, Any], facts: CheckpointFacts, features: Mapping[str, Any], b: _Builder +) -> dict[str, Any]: + """Each built feature's witness declaration: the site's, the window, and questions for the rest.""" + site_witness = dict(defaults.get("witness") or {}) + window = facts.config.get("sliding_window") + out: dict[str, Any] = {} + for name, raw in features.items(): + unit = unit_name(name, raw) + if unit not in WITNESS_REQUIRED: + continue + block = dict(site_witness.get(unit) or {}) + if block: + b.note(f"accuracy.witness.{unit}", dict(block), "site") + if unit in ("decoder_replay", "encoder_replay") and isinstance(window, int) and window > 0: + length_key = "prompt_tokens" if unit == "decoder_replay" else "prefix_tokens" + if "window" not in block: + block["window"] = b.note(f"accuracy.witness.{unit}.window", window, "checkpoint") + if length_key not in block: + block[length_key] = b.note( + f"accuracy.witness.{unit}.{length_key}", + REPLAY_PROMPT_WINDOWS * int(block["window"]), + DEFAULT, + ) + missing = [key for key in WITNESS_REQUIRED[unit] if not block.get(key)] + if unit in ("decoder_replay", "encoder_replay") and "window" not in block: + missing.append("window") + if unit == "disagg" and block.get("contrast_command") and not block.get("contrast_env"): + missing.append("contrast_env") + if missing: + b.questions.append( + f"accuracy.witness.{unit}: declare {', '.join(missing)} (site.yaml " + f"`defaults.witness.{unit}`, or in this file): without them the {unit} " + f"witness cannot tell a working feature from an absent one" + ) + if block or unit == "disagg": + out[unit] = block + return out + + +def resume_problem(request: Mapping[str, Any], *, checkpoint: Path, goal: str | None) -> str | None: + """Why the request already in a workspace is not the one this command asks for, or None. + + A request is re-derived only on ``--clean``: the one on disk may carry the + operator's edits, and its run may have started from it. + """ + derived = request.get("derived") if isinstance(request, Mapping) else None + if not isinstance(derived, Mapping): + return "the workspace's request was not derived by --checkpoint; run it with --task" + if Path(str(derived.get("checkpoint") or "")).resolve() != Path(checkpoint).resolve(): + return ( + f"the workspace's request was derived from {derived.get('checkpoint')}, not " + f"{checkpoint}: use another --workspace, or --clean to derive a new one" + ) + if goal is not None and " ".join(goal.split()) != str(derived.get("goal") or ""): + return ( + f"the workspace's request was derived for the goal {derived.get('goal')!r}: edit " + f"the request, or --clean to derive it again for {goal!r}" + ) + return None + + +def render_report(derivation: Derivation, request_path: Path, findings: list | None = None) -> str: + """``intake.report.md``: what was read, what was defaulted, and what is still open.""" + task = derivation.task + derived = task.get("derived") or {} + lines = [ + "# Intake: the derived request", + "", + f"Derived from `{derived.get('checkpoint')}` ({task['model'].get('architecture')}), " + f"the site `{derived.get('site') or '(unnamed)'}` and the goal " + f"`{derived.get('goal') or '(none)'}` at {derived.get('at')}.", + "", + f"The run starts from `{request_path}` at once, and its intake checks run in " + "phase 0, beside the survey and the substrate gate; a block stops the run there " + "with every block listed. Edit the request and re-run the same command to change " + "it; once the run is past preflight, an edit runs only when confirmed.", + "", + f"Confirmation: {confirmation_status(task, None)}. " + f"`trtllm-modeling-v2 intake confirm {request_path} --by NAME` records it, or " + "`--confirm` on the run's command.", + "", + "## What the goal did not state", + "", + "| field | from | value |", + "|---|---|---|", + ] + for where, value, source in derivation.unstated: + shown = json.dumps(value, sort_keys=True) if not isinstance(value, str) else value + lines.append(f"| `{where}` | {source} | `{shown}` |") + questions = task.get("open_questions") or [] + lines += ["", f"## Open questions ({len(questions)})", ""] + lines += [f"- {q}" for q in questions] or ["None."] + if findings: + lines += ["", "## Checks", "", "```"] + [f.render() for f in findings] + ["```"] + return "\n".join(lines) + "\n" + + +def write_request(workspace: Path, derivation: Derivation, findings: list | None = None) -> Path: + """Write the request and its report into ``workspace``; returns the request's path.""" + workspace.mkdir(parents=True, exist_ok=True) + path = workspace / REQUEST_FILENAME + path.write_text(yaml.safe_dump(derivation.task, sort_keys=False), encoding="utf-8") + (workspace / REPORT_FILENAME).write_text( + render_report(derivation, path, findings), encoding="utf-8" + ) + return path + + +__all__ = [ + "DEFAULT_BOOT", + "DEFAULT_MEMBERS", + "DEFAULT_SNAPSHOT_INTERVAL", + "DEFAULT_WORKLOAD", + "REPORT_FILENAME", + "REQUEST_FILENAME", + "WITNESS_REQUIRED", + "CheckpointFacts", + "Derivation", + "Goal", + "derive_task", + "inspect_checkpoint", + "parse_goal", + "render_report", + "resume_problem", + "write_request", +] diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/freeze.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/freeze.py new file mode 100644 index 000000000000..3d332979ed1f --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/freeze.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Confirmation: the user agreed to exactly this ``task.yaml``, and the run checks it still is. + +A confirmation is a digest of the request's content, recorded beside it in +``.confirmed`` so the request itself, comments included, is never +rewritten. A ``confirmed`` block inside the request is read too. + +A request written by hand runs only once confirmed. A request the fast path +derived (``derive.py``) runs at once: its confirmation is recorded when the +operator gives it (``--confirm``, or ``intake confirm`` later) rather than +awaited, and until then ``.started``, the digest of the content its +run started from, stands in for it. Either way an edit made after the run +started is refused until it is confirmed. +""" + +from __future__ import annotations + +import hashlib +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Mapping + +import yaml + +from .schema import IntakeError, validate_task + +#: The blocks a request may carry about itself, left out of its digest. +STAMP_KEYS: tuple[str, ...] = ("confirmed", "derived") +#: Beside the request: who confirmed which content, and what its run started from. +CONFIRMED_SUFFIX = ".confirmed" +STARTED_SUFFIX = ".started" + + +def digest(task: Mapping[str, Any]) -> str: + """sha256 of the task without its stamps (:data:`STAMP_KEYS`), independent of key order.""" + body = {k: v for k, v in task.items() if k not in STAMP_KEYS} + return hashlib.sha256(yaml.safe_dump(body, sort_keys=True).encode("utf-8")).hexdigest() + + +def _now() -> str: + return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def confirmation_path(path: Path) -> Path: + return Path(path).with_name(Path(path).name + CONFIRMED_SUFFIX) + + +def started_path(path: Path) -> Path: + return Path(path).with_name(Path(path).name + STARTED_SUFFIX) + + +def _read_stamp(path: Path) -> dict[str, Any] | None: + try: + data = yaml.safe_load(path.read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError): + return None + return dict(data) if isinstance(data, Mapping) and data.get("sha256") else None + + +def read_confirmation(path: Path) -> dict[str, Any] | None: + """The confirmation recorded for the request at ``path``: its sidecar, else its own block.""" + stamp = _read_stamp(confirmation_path(path)) + if stamp is not None: + return stamp + try: + task = yaml.safe_load(Path(path).read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError): + return None + confirmed = task.get("confirmed") if isinstance(task, Mapping) else None + return dict(confirmed) if isinstance(confirmed, Mapping) and confirmed.get("sha256") else None + + +def confirm(path: Path, by: str) -> dict[str, Any]: + """Record that ``by`` confirmed the task at ``path``; refuse one with open questions. + + The record goes to ``.confirmed``; the request is not rewritten. + """ + task = yaml.safe_load(Path(path).read_text(encoding="utf-8")) + errors = validate_task(task) + if isinstance(task, Mapping) and task.get("open_questions"): + errors.append(f"{len(task['open_questions'])} open question(s) must be answered first") + if errors: + raise IntakeError("\n".join(errors)) + stamp = {"by": by, "at": _now(), "sha256": digest(task)} + confirmation_path(path).write_text(yaml.safe_dump(stamp, sort_keys=False), encoding="utf-8") + return stamp + + +def unconfirmed_reason( + task: Mapping[str, Any], confirmation: Mapping[str, Any] | None = None +) -> str | None: + """Why ``task`` may not run, or None when it is confirmed and unchanged since. + + ``confirmation`` is :func:`read_confirmation`'s answer; without it the + task's own ``confirmed`` block is read. + """ + confirmed = confirmation if confirmation is not None else task.get("confirmed") + if not isinstance(confirmed, Mapping) or not confirmed.get("sha256"): + return "task.yaml is not confirmed: run `trtllm-modeling-v2 intake confirm`" + if confirmed["sha256"] != digest(task): + return "task.yaml changed after it was confirmed: review it and confirm again" + return None + + +def start_reason(path: Path, *, started: bool) -> str | None: + """Why the derived request at ``path`` may not start or resume its run, or None. + + A confirmed request runs while it is unchanged since its confirmation. One + not yet confirmed runs too; but once its run is past preflight + (``started``), it must still be the content that run started from + (:func:`stamp_start`), and an edit is refused until it is confirmed. + """ + task = yaml.safe_load(Path(path).read_text(encoding="utf-8")) + confirmation = read_confirmation(path) + if confirmation is not None: + return unconfirmed_reason(task, confirmation) + if not started: + return None + stamp = _read_stamp(started_path(path)) + if stamp is None or stamp["sha256"] != digest(task): + return ( + "the request changed after its run started from it: confirm the change " + f"(`trtllm-modeling-v2 intake confirm {path} --by NAME`, or `--confirm` on this " + "command), or start over with --clean" + ) + return None + + +def stamp_start(path: Path) -> dict[str, Any]: + """Record the content at ``path`` as what its run starts from, in ``.started``.""" + task = yaml.safe_load(Path(path).read_text(encoding="utf-8")) + stamp = {"at": _now(), "sha256": digest(task)} + started_path(path).write_text(yaml.safe_dump(stamp, sort_keys=False), encoding="utf-8") + return stamp + + +def confirmation_status(task: Mapping[str, Any], confirmation: Mapping[str, Any] | None) -> str: + """One line on whether ``task`` is confirmed, for the journal and the console.""" + if confirmation is None: + return "not confirmed; the run does not wait for it" + if confirmation.get("sha256") != digest(task): + return "changed after it was confirmed" + return f"confirmed by {confirmation.get('by', '?')} at {confirmation.get('at', '?')}" + + +__all__ = [ + "CONFIRMED_SUFFIX", + "STAMP_KEYS", + "STARTED_SUFFIX", + "confirm", + "confirmation_path", + "confirmation_status", + "digest", + "read_confirmation", + "stamp_start", + "start_reason", + "started_path", + "unconfirmed_reason", +] diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/schema.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/schema.py new file mode 100644 index 000000000000..f34be3ecca86 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/intake/schema.py @@ -0,0 +1,269 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The user-facing ``task.yaml`` and the per-cluster ``site.yaml``: closed keys, batched errors.""" + +from __future__ import annotations + +import re +from typing import Any, Mapping + +from ..schema_keys import reject_unknown_keys + +INTAKE_VERSION = 1 + +GOAL_ONBOARD = "onboard" +GOAL_ADD_FEATURES = "add_features" +GOAL_TUNE = "tune" +GOALS: tuple[str, ...] = (GOAL_ONBOARD, GOAL_ADD_FEATURES, GOAL_TUNE) + +STATUS_REQUIRED = "required" +STATUS_NICE = "nice_to_have" +STATUS_OUT = "out_of_scope" +STATUSES: tuple[str, ...] = (STATUS_REQUIRED, STATUS_NICE, STATUS_OUT) + +TASK_KEYS: tuple[str, ...] = ( + "version", + "goals", + "model", + "device", + "parallelism", + "workload", + "features", + "accuracy", + "perf", + "constraints", + "deliverable", + "notes", + "provenance", + "open_questions", + "confirmed", + # Written by the fast path (`derive.py`): the checkpoint, site and goal the + # request was derived from, and when. + "derived", +) +MODEL_KEYS = ( + "checkpoint", + "architecture", + "quantization", + "kind", + "in_tree", + "reference_impl", + "base", + "target", +) +DEVICE_KEYS = ("arch", "gpus_per_node", "nodes") +WORKLOAD_KEYS = ("isl", "osl", "max_seq_len", "max_batch_size", "concurrency") +ACCURACY_KEYS = ( + "benchmark", + "protocol", + "reference", + "reference_yaml", + "model_key", + # The reference row's discriminators, when the checkpoint's own cannot be read. + "discriminators", + # The boot gate's prompts and the witness declarations, carried to the + # resolved spec's `accuracy.boot` and `accuracy.witness` as they are. + "boot", + "witness", +) +PERF_KEYS = ("metric", "baseline", "goal", "configs") +CONSTRAINT_KEYS = ("embargo", "base_ref", "budget") +FEATURE_ENTRY_KEYS = ("status", "params", "members", "accept") +PARALLEL_KINDS = ("tp", "tep", "dep") +#: One parallel configuration, as a member is written: `tp1`, `tep4`, `dep8`. +PARALLEL_SEGMENT = re.compile(r"^(tp|tep|dep)([1-9][0-9]*)$") + +#: Runtime behaviours every target has; their status decides the gate switches. +CORE_FEATURES = ("cuda_graph", "overlap_scheduler", "chunked_prefill", "prefix_reuse", "kv_cache") +#: Features that become a build increment of today's pipeline, by their pipeline name. +INCREMENT_FEATURES = { + "kv_quant": "kv_quant", + "decoder_replay": "decoder_replay", + "encoder_replay": "encoder_replay", + "host_offload": "host_offload", + "disaggregated": "disagg", + "multimodal": "multimodal", +} +SPECULATIVE_ALGOS = ("mtp", "eagle3", "dflash", "dspark") +#: Features recorded for the user but with no Step 2 consumer yet. +UNGATED_FEATURES = ( + "long_context", + "guided_decoding", + "logprobs", + "reasoning_parser", + "tool_parser", + "lora", +) +KNOWN_FEATURES = ( + *CORE_FEATURES, + "speculative", + *INCREMENT_FEATURES, + *UNGATED_FEATURES, + "custom", +) + +SITE_KEYS = ( + "trtllm_repo_path", + "environment", + "agents", + "perf", + "commit_policy", + # The cluster's devices, and what a derived request takes when its goal is silent. + "device", + "defaults", +) +SITE_PERF_KEYS = ("spools", "dev_spools", "env_script") +#: `site.yaml defaults`: the members, workload, accuracy anchor, boot prompts and +#: witness declarations (by pipeline feature name) a derived request starts from. +SITE_DEFAULTS_KEYS = ("members", "workload", "accuracy", "boot", "witness") +SITE_DEFAULT_ACCURACY_KEYS = ("benchmark", "protocol", "reference", "reference_yaml", "model_key") + + +class IntakeError(ValueError): + """A ``task.yaml`` or ``site.yaml`` that cannot be compiled, every problem listed.""" + + +def feature_entry(value: Any) -> dict[str, Any]: + """``required`` or ``{status: required, params: ...}`` as one shape.""" + if isinstance(value, str): + return {"status": value} + return dict(value or {}) + + +def feature_status(task: Mapping[str, Any], name: str) -> str | None: + """The declared status of ``name``, or None when the task does not mention it.""" + features = task.get("features") or {} + if name not in features: + return None + return feature_entry(features[name]).get("status") + + +def _mapping(task: Mapping[str, Any], key: str, keys, errors: list[str]) -> Mapping[str, Any]: + block = task.get(key) + if block is None: + return {} + if not isinstance(block, Mapping): + errors.append(f"'{key}' must be a mapping") + return {} + reject_unknown_keys(block, keys, errors, where=key) + return block + + +def validate_task(task: Any) -> list[str]: + """Every problem with the shape of a user-facing ``task.yaml``.""" + if not isinstance(task, Mapping): + return ["task.yaml must be a mapping"] + errors: list[str] = [] + reject_unknown_keys(task, TASK_KEYS, errors, where="task.yaml") + if task.get("version") != INTAKE_VERSION: + errors.append(f"'version' must be {INTAKE_VERSION}") + goals = task.get("goals") + if not isinstance(goals, list) or not all(g in GOALS for g in goals): + errors.append(f"'goals' must be a list drawn from {list(GOALS)}") + model = _mapping(task, "model", MODEL_KEYS, errors) + if not model.get("checkpoint"): + errors.append("'model.checkpoint' is required") + device = _mapping(task, "device", DEVICE_KEYS, errors) + if not device.get("arch"): + errors.append("'device.arch' is required (sm_100, sm_103, ...)") + parallelism = task.get("parallelism") + if parallelism is not None: + if not isinstance(parallelism, list) or not parallelism: + errors.append("'parallelism' must be a non-empty list of {kind, world}") + else: + for i, entry in enumerate(parallelism): + if not isinstance(entry, Mapping) or entry.get("kind") not in PARALLEL_KINDS: + errors.append(f"'parallelism[{i}].kind' must be one of {list(PARALLEL_KINDS)}") + elif not isinstance(entry.get("world"), int) or entry["world"] < 1: + errors.append(f"'parallelism[{i}].world' must be an integer >= 1") + workload = _mapping(task, "workload", WORKLOAD_KEYS, errors) + for key in ("isl", "osl", "max_seq_len", "max_batch_size"): + if key in workload and (not isinstance(workload[key], int) or workload[key] < 1): + errors.append(f"'workload.{key}' must be an integer >= 1") + features = task.get("features") or {} + if not isinstance(features, Mapping): + errors.append("'features' must be a mapping of feature -> status or entry") + features = {} + for name, value in features.items(): + if name not in KNOWN_FEATURES: + errors.append(f"'features.{name}' is not a known feature: {sorted(KNOWN_FEATURES)}") + continue + if name == "custom": + if not isinstance(value, list): + errors.append("'features.custom' must be a list") + continue + entry = feature_entry(value) + reject_unknown_keys(entry, FEATURE_ENTRY_KEYS, errors, where=f"features.{name}") + if entry.get("status") not in STATUSES: + errors.append(f"'features.{name}.status' must be one of {list(STATUSES)}") + if name == "speculative" and entry.get("status") != STATUS_OUT: + algo = (entry.get("params") or {}).get("algo") + if algo not in SPECULATIVE_ALGOS: + errors.append( + f"'features.speculative.params.algo' must be one of {list(SPECULATIVE_ALGOS)}" + ) + _mapping(task, "accuracy", ACCURACY_KEYS, errors) + _mapping(task, "perf", PERF_KEYS, errors) + _mapping(task, "constraints", CONSTRAINT_KEYS, errors) + questions = task.get("open_questions") + if questions is not None and not isinstance(questions, list): + errors.append("'open_questions' must be a list") + derived = task.get("derived") + if derived is not None and not isinstance(derived, Mapping): + errors.append("'derived' must be a mapping") + return errors + + +def validate_site(site: Any) -> list[str]: + """Every problem with the shape of a ``site.yaml``.""" + if not isinstance(site, Mapping): + return ["site.yaml must be a mapping"] + errors: list[str] = [] + reject_unknown_keys(site, SITE_KEYS, errors, where="site.yaml") + for key in ("trtllm_repo_path", "environment"): + if not site.get(key): + errors.append(f"site.yaml: '{key}' is required") + perf = site.get("perf") or {} + if isinstance(perf, Mapping): + reject_unknown_keys(perf, SITE_PERF_KEYS, errors, where="site.yaml perf") + device = site.get("device") + if device is not None: + if not isinstance(device, Mapping): + errors.append("site.yaml: 'device' must be a mapping") + else: + reject_unknown_keys(device, DEVICE_KEYS, errors, where="site.yaml device") + defaults = site.get("defaults") + if defaults is not None: + if not isinstance(defaults, Mapping): + errors.append("site.yaml: 'defaults' must be a mapping") + else: + reject_unknown_keys(defaults, SITE_DEFAULTS_KEYS, errors, where="site.yaml defaults") + members = defaults.get("members") + if members is not None and ( + not isinstance(members, list) + or not members + or not all(isinstance(m, str) and PARALLEL_SEGMENT.match(m) for m in members) + ): + errors.append( + "site.yaml: 'defaults.members' must be a non-empty list of tp, tep, dep" + ) + for key in ("workload", "accuracy", "boot", "witness"): + if defaults.get(key) is not None and not isinstance(defaults[key], Mapping): + errors.append(f"site.yaml: 'defaults.{key}' must be a mapping") + if isinstance(defaults.get("workload"), Mapping): + reject_unknown_keys( + defaults["workload"], WORKLOAD_KEYS, errors, where="site.yaml defaults.workload" + ) + if isinstance(defaults.get("accuracy"), Mapping): + reject_unknown_keys( + defaults["accuracy"], + SITE_DEFAULT_ACCURACY_KEYS, + errors, + where="site.yaml defaults.accuracy", + ) + return errors + + +def is_intake_task(data: Any) -> bool: + """A user-facing ``task.yaml`` (as opposed to an already resolved spec).""" + return isinstance(data, Mapping) and "goals" in data and "version" in data diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/schema_keys.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/schema_keys.py new file mode 100644 index 000000000000..65ee77956025 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/schema_keys.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Closed-key validation, shared by every block of the task schema.""" + +from __future__ import annotations + +import difflib +from typing import Any, Mapping, Sequence + + +def reject_unknown_keys( + block: Mapping[str, Any], + known: Sequence[str], + errors: list[str], + *, + where: str = "", +) -> None: + """Append an error for any key in ``block`` that ``known`` does not list. + + Every block in this schema is closed: a misspelled key that is silently + ignored produces a run that gates a configuration nobody asked for, and + the report still says it was covered. ``feature:`` instead of + ``features:`` would resolve to ``[]``, so the run would onboard core + only, gate it green, and report success for a feature set it never + built; ``runtme:`` instead of ``runtime:`` would drop ``max_seq_len`` and + fail later with ``Impossible to fit in any sequence in kvCache``. + + The nearest known key is suggested, because the whole class of error is + off-by-one-character and a suggestion turns the message into the fix. + """ + prefix = f"{where}." if where else "" + for key in sorted(set(block) - set(known)): + close = difflib.get_close_matches(str(key), [str(k) for k in known], n=1, cutoff=0.6) + hint = f" — did you mean '{prefix}{close[0]}'?" if close else "" + errors.append( + f"unknown key '{prefix}{key}'{hint} Known keys here: {sorted(str(k) for k in known)}" + ) + + +__all__ = ["reject_unknown_keys"] diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/task_schema.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/task_schema.py new file mode 100644 index 000000000000..0744fa8ce986 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/task_schema.py @@ -0,0 +1,2477 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema validation for trtllm-modeling-v2's ``task.yaml`` input. + +A modeling_v2 target's identity is the ``(checkpoint, GPU arch, parallel +config)`` triple, so those three fields are the spec's core and the run's +slug is derived from them. ``arch`` is **required and never inferred**: +catalog receipts are per-architecture, so a different arch means the whole +certification is redone, and guessing it wrong means gating against +receipts that describe another machine. + +The ``environment`` block is the part that earns its keep. It declares — +once, up front — how GPU work is launched, and the workflow resolves it +into a single command template that every agent prompt carries verbatim. +Agents never construct a launcher command, never probe for one, and never +call a test-runner skill. That is deliberate: the failure mode this +workflow exists to avoid is a run that spends its budget deciding how to +submit a job instead of onboarding a model. + +Validation runs at the CLI boundary so an invalid spec aborts before any +agent is constructed, and every detected problem is batched into one +message so the user fixes them in a single pass rather than one per run. +:func:`load_and_validate_task_yaml` also fills in the workflow's defaults, +so the resolved spec copied into the workspace is fully explicit. +""" + +from __future__ import annotations + +import re +import shlex +from pathlib import Path +from typing import Any, Mapping + +import yaml + +from agent_flow.agent_runtime import AGENTS_FIELD, validate_agents + +from ..agents.roles import ROLES +from ..checks import stats +from ..infra.committer import COMMIT_POLICY_FIELD, validate_commit_policy_block +from ..infra.substrate.sessions import validate_sessions +from ..onboard.embargo import PACKAGE_PREFIX, package_entry_problem, resolve_entry +from ..tune.schema import validate_perf_block +from .schema_keys import reject_unknown_keys +from .variants import VARIANTS_FIELD, validate_variants + +REQUIRED_PATH_FIELDS: tuple[str, ...] = ( + "checkpoint_path", + "trtllm_repo_path", +) + +ARCH_FIELD = "arch" +PARALLEL_FIELD = "parallel" +FEATURES_FIELD = "features" +ENVIRONMENT_FIELD = "environment" +ACCURACY_FIELD = "accuracy" + +# ``sm_100`` (B200), ``sm_103`` (GB300), ... Matched rather than +# enumerated: a new architecture should not need a code change here, and +# preflight asserts the value against the host anyway, which is the check +# that actually protects the receipts. +ARCH_PATTERN = re.compile(r"^sm_\d+$") + +# ``tp4`` / ``tep4`` / ``dep4`` carry their world size in the trailing +# number; ``unified`` means one modeling codebase over a family of +# parallel configurations and takes its world size from the member set. +PARALLEL_PATTERN = re.compile(r"^(tp|tep|dep)(\d+)$") +PARALLEL_UNIFIED = "unified" +DEFAULT_PARALLEL = PARALLEL_UNIFIED + +# Which concrete configurations a ``unified`` segment is actually gated over. +MEMBERS_FIELD = "members" + +# The default member set, and the reason there is one: ``unified`` is +# ``DEFAULT_PARALLEL``, so a spec that omits ``parallel`` gets the form with +# the broadest claim. Gated at ``--tp_size=N`` alone, the expert-parallel and +# attention-DP branches that only a unified codebase carries would never be +# executed by any gate. +# +# TEP and DEP at the world size are the baseline pair +# ``agents/skills/trtllm-modeling-v2-parallel/references/unified.md`` names; TP is +# the degenerate path, and it is not free — it is the only member that folds +# every branch off, so it runs the general path with nothing switched on and +# no other member covers that. +DEFAULT_MEMBER_KINDS: tuple[str, ...] = ("tp", "tep", "dep") + +# Feature name -> the skill the assembler loads for it. This mapping is +# the reason the feature set is closed: a feature with no skill is an +# assembler that has been told to build something nobody wrote down how +# to build, which surfaces as a confused mid-run report rather than as an +# error here. +FEATURE_SKILLS: dict[str, str] = { + "mtp": "trtllm-modeling-v2-speculator", + "eagle3": "trtllm-modeling-v2-speculator", + "dflash": "trtllm-modeling-v2-speculator", + # A block drafter: one 3-stage pass proposes a 5-token block, where MTP + # chains one-token heads. Its own kind rather than a spelling of `mtp` + # because the two read the SAME checkpoint namespace: a checkpoint can ship + # its DSpark stages as `mtp.*` with `num_nextn_predict_layers: 3`, and + # `--spec_decode_algo=MTP` over that builds three chained heads that either + # fail to load or draft garbage the verifier silently rejects. + "dspark": "trtllm-modeling-v2-speculator", + # "SWA Bounded Replay", decoder and encoder halves. NOT + # output-identical -- both deliberately truncate sliding-window state -- so + # accuracy is a real check here; but a unit that implements nothing passes + # it too, which is why each carries a work witness. See + # `witness.run_replay_witness`. + "decoder_replay": "trtllm-modeling-v2-bounded-replay", + "encoder_replay": "trtllm-modeling-v2-bounded-replay", + # KV pools switch from bf16 containers holding the reference's + # fake-quantized values (exact, and what core ships) to the checkpoint's + # native storage formats. Output-identical by construction when attention + # dequantizes exactly, so boot and accuracy cannot see whether it happened; + # its witness counts stored bytes per token against a declared figure. + "kv_quant": "trtllm-modeling-v2-kv-quant", + "multimodal": "trtllm-modeling-v2-multimodal", + # A weight tensor too large to keep resident lives in pinned host memory + # and the forward gathers what it needs. Named for the mechanism, not for + # any one tensor (a tens-of-GB per-layer-embedding n-gram table, say), + # because the mechanism is what generalizes. + "host_offload": "trtllm-modeling-v2-host-offload", + # Prefill and decode on separate engines. A modeling feature rather than a + # serving one: only per-sequence state the transceiver enumerates survives + # the hop, so a target that keeps state elsewhere passes every + # single-engine gate and is wrong across it. + "disagg": "trtllm-modeling-v2-disagg", +} + +# Features that read the checkpoint's `mtp.*` namespace, which holds MTP +# modules OR DSpark stages and never both. Listing two of them is not a richer +# target; it is two readings of one set of weights, one of which is wrong. +MTP_NAMESPACE_FEATURES: tuple[str, ...] = ("mtp", "dspark") + +# The engine batch limit BOTH legs of a token-comparing witness (residency, +# kv_quant, equivalence) run at. It replaces the runtime block's +# `--max_batch_size` in leg A's quickstart argv, and a `contrast_command` leg B +# is handed it as MV2FLOW_WITNESS_MAX_BATCH_SIZE. `1` is the value that +# gives both legs the same schedule. A row's bits depend on the rows beside it +# (fp32 GEMMs, sparse attention at a large verify batch, the all-reduce +# strategy by token count), and batched admission follows arrival timing, so +# two batched legs can pack rows differently and flip a near-tie that neither +# leg's feature caused. Unset keeps the runtime block's limit. The replay +# witnesses count work rather than compare tokens, and do not take it. +WITNESS_BATCH_SIZE_KEY = "max_batch_size" + +# The keys every witness block may carry. `expected_mib` is the residency +# witness's prediction, `contrast_command` the equivalence witness's second +# launcher. `timeout` is seconds PER LEG, where the one-engine boot budget does +# not fit -- a disagg leg boots two engines in sequence. +# The workload a token-comparing witness (residency, kv_quant, equivalence) +# declares, as the acceptance gate does: `prompts` or `prompts_file`, +# `max_tokens`, and `min_generated_tokens` (default 32) -- a witness whose +# continuations are shorter fails as too short to compare instead of passing +# on two-token answers. With none, the declared acceptance prompts, else the +# boot prompts. +WITNESS_WORKLOAD_KEYS: tuple[str, ...] = ( + "prompts", + "prompts_file", + "max_tokens", + "min_generated_tokens", +) +WITNESS_KEYS: tuple[str, ...] = ( + "contrast_env", + "expected_mib", + "contrast_command", + "timeout", + WITNESS_BATCH_SIZE_KEY, + *WITNESS_WORKLOAD_KEYS, +) + +# Features gated by the EQUIVALENCE witness (token identity across a hop). +# +# contrast_command optional. Leg B's launcher: a context engine and a +# generation engine with the state handed between them, +# which no environment variable over the aggregated command +# turns into. Without one the witness runs the flow's own +# (`checks/harness/disagg_leg_b.py`); with one, `contrast_env` +# is required as before. A word of it naming a file beside +# the spec is resolved against the spec's directory +# prompt_tokens optional list: one generated prompt of at least that many +# tokens each, added to the boot prompts on BOTH legs, for the +# lengths a hop has edges at (over a prefill chunk, odd, one +# past a page) that hand-written prompts rarely land on +# legs_env optional mapping, see WITNESS_LEGS_ENV_KEY +EQUIVALENCE_WITNESS_FEATURES: tuple[str, ...] = ("disagg",) + +# The environment BOTH legs of an equivalence witness run under: leg A's +# process, and leg B's launcher, which is handed it as +# MV2FLOW_WITNESS_LEGS_ENV to apply to every engine it starts. A witness +# holds when its legs differ in exactly one variable, and leg B's launcher is a +# different program from leg A's: a code path a disaggregated engine refuses or +# guards off (a replay that cannot run across the hop) stays on in the +# aggregated leg, so the legs differ in it as well as in the hop and token +# identity does not hold by construction. Pinning it here on both legs takes +# it out of the comparison; its own witness covers it. A variable may not be in +# both this and `contrast_env`: there it is the contrast. The residency and +# kv_quant witnesses run the same aggregated command on both legs, so every +# code path other than `contrast_env` is already alike and they do not take it. +WITNESS_LEGS_ENV_KEY = "legs_env" + +# The cache transceiver both engines of the flow's own leg B run, as the LLM +# API's `CacheTransceiverConfig` takes it (`backend`, `transceiver_runtime`, +# ...). Unset, `{backend: DEFAULT, transceiver_runtime: auto}`. A +# `contrast_command` is handed it as MV2FLOW_WITNESS_TRANSCEIVER. +WITNESS_TRANSCEIVER_KEY = "cache_transceiver_config" + +EQUIVALENCE_WITNESS_KEYS: tuple[str, ...] = ( + "contrast_env", + "contrast_command", + "prompt_tokens", + "timeout", + WITNESS_BATCH_SIZE_KEY, + WITNESS_LEGS_ENV_KEY, + WITNESS_TRANSCEIVER_KEY, + *WITNESS_WORKLOAD_KEYS, +) + +# The bounded-replay witness parameters, per feature, each REQUIRED. Declared +# in the spec because a work witness that sizes its workload from the run it is +# judging can only agree with itself -- and because both verdicts are +# comparisons against `window`, which is a property of the checkpoint, not of +# anything the target prints. +# +# window the checkpoint's sliding window (e.g. 128) +# prompt_tokens decoder_replay: the long prompt's length, at least. Must +# exceed `window`, or replay-on and replay-off do the same work +# prefix_tokens encoder_replay: the shared prefix's length. Must exceed +# `window`, or a prefix hit replays everything it reused and +# "re-forwarded the whole prefix" looks like a pass +REPLAY_WITNESS_REQUIRED: dict[str, tuple[str, ...]] = { + "decoder_replay": ("contrast_env", "window", "prompt_tokens"), + "encoder_replay": ("contrast_env", "window", "prefix_tokens"), +} +# `reuse_alignment`: the reuse boundary must be a multiple of this. A model that +# compresses global KV in pairs (compress ratio 2) needs 2, since a boundary that +# splits a pair leaves one half of it nowhere; SGLang refuses an odd boundary for +# that reason. Default 1, i.e. no constraint. +REPLAY_WITNESS_OPTIONAL: dict[str, tuple[str, ...]] = { + "decoder_replay": ("timeout",), + "encoder_replay": ("reuse_alignment", "timeout"), +} + +# The witness builds its long prompt as a quickstart `--prompt` argument, one +# word per token at least. Linux caps ONE argv string at 128 KiB +# (MAX_ARG_STRLEN) and fails the exec with E2BIG, which surfaces as a leg that +# "did not complete" -- so the cap is stated here, where it costs nothing. +# 16384 words of the witness's prose (`checks/witness_text.py`) are ~89 KB. +MAX_WITNESS_PROMPT_WORDS = 16384 + +# The `kv_quant` witness block. +# +# expected_bytes_per_token REQUIRED. Persistent KV bytes one rank stores per +# token, computed from the checkpoint's formats +# before the run (e.g. 890 for an NVFP4 checkpoint) +# exact default true: attention dequantizes the native +# values exactly, so output must be token-identical +# to bf16 storage of the same values. `false` names +# a lossy attention path (an FP8-query kernel) and +# drops the identity leg, leaving accuracy as the +# only output check -- a weaker claim, recorded so +# contrast_env the leg WITH bf16 storage; required when exact +# bytes_tolerance relative, default 0.02. Enough for row padding, +# not for a different layout: padding the 4-byte +# indexer scale to 16 moves 890 to 920 (+3.4%) +KV_QUANT_WITNESS_KEYS: tuple[str, ...] = ( + "contrast_env", + "expected_bytes_per_token", + "exact", + "bytes_tolerance", + "timeout", + WITNESS_BATCH_SIZE_KEY, + *WITNESS_WORKLOAD_KEYS, +) +DEFAULT_BYTES_TOLERANCE = 0.02 + +# Every feature whose witness block is a DECLARATION -- a workload or a figure +# the verdict is judged against -- and is therefore required at spec time. +DECLARED_WITNESS_FEATURES: tuple[str, ...] = (*REPLAY_WITNESS_REQUIRED, "kv_quant") + +# The acceptance gate's WORKLOAD, declared in `accuracy.acceptance` beside the +# per-unit floors: at the block's top level for every unit, or inside one +# unit's entry for that unit alone. With none declared, the gate takes the +# first `num_prompts` (default 32) natural-text questions of the accuracy +# benchmark's dataset -- a local path, or a Hugging Face id read from the local +# cache only. Never the boot prompts: a 2-token answer ends in EOS before a +# drafter is asked anything, and the rate reads 1.00 for every drafter. +# +# prompts a list of prompt strings, passed as-is (no chat +# template is applied: write them the way the boot +# prompts are written) +# prompts_file or a file holding them: a YAML/JSON list of strings, +# or a mapping with `prompts:` (boot.yaml's shape). +# Relative paths resolve against the spec's directory +# dataset or a dataset: a local file or directory (JSON, JSON +# lines, YAML, Parquet, Arrow) or a cached dataset id +# num_prompts questions taken from the dataset (default 32) +# text_field the dataset field holding the question (default: the +# first of question/prompt/problem/input/text/...) +# prompt_template how a question becomes a prompt, `{text}` naming it; +# write the model's chat markers here if it needs them +# split, dataset_config which split and config of a dataset id +# max_tokens tokens each prompt may generate (default 256) +# min_draft_iterations generation iterations the whole workload must reach +# before its rate is read (default 32); fewer fails as +# "workload too short" rather than as a low number +# min_prompts prompts the workload must hold (default 32); fewer +# fails as "workload too small": four prompts read a +# change in generated text as an acceptance regression +ACCEPTANCE_WORKLOAD_KEYS: tuple[str, ...] = ( + "prompts", + "prompts_file", + "dataset", + "num_prompts", + "text_field", + "prompt_template", + "split", + "dataset_config", + "max_tokens", + "min_draft_iterations", + "min_prompts", +) +# The paired comparison against the previous tree's verdict, at the block's top +# level only: `test` (t or wilcoxon, default t), `alpha` (one-sided, default +# 0.05), `min_effect` (accepted tokens per iteration, default 0.1). A drop must +# be both significant and at least `min_effect` to fail. +ACCEPTANCE_REGRESSION_KEYS: tuple[str, ...] = ("test", "alpha", "min_effect") +ACCEPTANCE_BLOCK_KEYS: tuple[str, ...] = (*ACCEPTANCE_WORKLOAD_KEYS, "regression") +# What one unit's entry in `accuracy.acceptance` may carry: +# min_absolute the absolute floor used when the `=off` leg cannot run +# measured_reference what a reference drafter measured on this workload; the +# floor is this x 0.95 when `min_absolute` is not given +# floor_source where that measurement was taken. A floor with neither +# a measured reference nor a source is used, and every +# verdict says it is not tied to a measurement +ACCEPTANCE_UNIT_KEYS: tuple[str, ...] = ( + "min_absolute", + "measured_reference", + "floor_source", + *ACCEPTANCE_WORKLOAD_KEYS, +) + +TARGET_FIELD = "target" +RUNTIME_FIELD = "runtime" +SCOPE_FIELD = "scope" +NOTES_FIELD = "notes" +SOURCE_EMBARGO_FIELD = "source_embargo" + +# Runtime switches the gates must exercise, and which therefore have to be +# declared rather than left to whatever the engine defaults to. +# +# The values are what the *user* wants true of the run, not what the flag +# says: two of these are spelled as negations on the command line +# (``--disable_overlap_scheduler``, ``--disable_kv_cache_reuse``), so a +# reader translating "block reuse on" into a flag has to remember to emit +# nothing. Doing that translation once, in Python, is the point — +# :func:`~.gates.runtime_flags` owns it and no agent ever types it. +RUNTIME_SWITCHES: dict[str, bool] = { + "cuda_graph": False, + "chunked_prefill": False, + "overlap_scheduler": True, + "kv_cache_reuse": True, +} + +# ---- memory margin (checks/gates.py, `MemoryMargin`) ------------------------- +# The free device memory a gate's engine must keep once it is up, below which +# its verdict carries a LOW MARGIN warning: the larger of an absolute floor and +# a fraction of the device. Optional and never filled in: validation leaves a +# runtime block that sets neither as it was. Scalars, so a flat block is still +# told apart from the per-unit form by its mapping values. +RUNTIME_MARGIN_DEFAULTS: dict[str, float] = { + "min_free_memory_gib": 2.0, + "min_free_memory_fraction": 0.02, +} + +# Scopes a target may be built at. ``text_only`` is the useful one: a +# multimodal checkpoint whose vision tower is out of scope still has to +# load, route and gate as a language model, and saying so up front is what +# keeps an assembler from treating the absent vision path as a defect. +SCOPES: tuple[str, ...] = ("full", "text_only") +DEFAULT_SCOPE = "full" + +# Every block below is closed; see :func:`reject_unknown_keys` for why. +# Listed as data rather than derived from the validators so that adding a +# field and forgetting to accept it is a test failure, not a silent drop. +TOP_LEVEL_KEYS: tuple[str, ...] = ( + *REQUIRED_PATH_FIELDS, + ARCH_FIELD, + PARALLEL_FIELD, + MEMBERS_FIELD, + FEATURES_FIELD, + ENVIRONMENT_FIELD, + ACCURACY_FIELD, + TARGET_FIELD, + RUNTIME_FIELD, + SCOPE_FIELD, + NOTES_FIELD, + SOURCE_EMBARGO_FIELD, + # Per-role backend, model and reasoning effort; see `roles.py`. + AGENTS_FIELD, + COMMIT_POLICY_FIELD, + # The tune stage (tune/stage.py). Closed keys, validated there. + "perf", + # Configurations of a built unit the target must also serve; see `variants.py`. + VARIANTS_FIELD, + # Requested cells the run accepts leaving uncovered; see `run/requirements.py`. + "coverage", + # What earlier runs' knowledge the run starts with; see `KNOWLEDGE_MODES`. + "knowledge", +) +#: ``coverage.accept``: the requested feature@member cells the contract may leave uncovered. +COVERAGE_FIELD = "coverage" +COVERAGE_KEYS: tuple[str, ...] = ("accept",) +#: ``knowledge``: ``production`` (the default) starts from everything earlier runs +#: learned, the engine-defect registry included; ``measurement`` leaves out the +#: registry and the lessons of the model's family, so a run measures how the +#: workflow does on a model it has not met. +KNOWLEDGE_FIELD = "knowledge" +KNOWLEDGE_PRODUCTION = "production" +KNOWLEDGE_MEASUREMENT = "measurement" +KNOWLEDGE_MODES: tuple[str, ...] = (KNOWLEDGE_PRODUCTION, KNOWLEDGE_MEASUREMENT) + + +def knowledge_mode(task_data: Mapping[str, Any]) -> str: + """The run's ``knowledge`` mode; ``production`` when the spec names none.""" + return str(task_data.get(KNOWLEDGE_FIELD) or KNOWLEDGE_PRODUCTION) + + +TARGET_KEYS: tuple[str, ...] = ("family", "checkpoint_segment", "keep_local") +#: The ``target`` keys that name a Python package directory. +TARGET_NAME_KEYS: tuple[str, ...] = ("family", "checkpoint_segment") + +ACCURACY_KEYS: tuple[str, ...] = ( + # The absolute `acceptance_length` floor a unit falls back to when the + # `=off` leg cannot exist -- a genuinely new architecture has no built-in + # implementation to compare a speculator against. See + # `gates.resolve_acceptance_floor`. Also the acceptance gate's declared + # workload; see `ACCEPTANCE_WORKLOAD_KEYS`. + "acceptance", + "benchmark", + "reference_yaml", + "model_key", + "tolerance", + "dataset_path", + "llm_api_options", + # Which ROW of reference_yaml this run anchors to. The references key a + # model to a list discriminated by quant_algo / kv_cache_quant_algo / + # spec_dec_algo, and the match is exact -- so a quantized checkpoint with + # no discriminators here anchors to the bf16 row, or to nothing at all. + "discriminators", + # Per-unit configuration for the two-leg WITNESS gates, which exist for + # the features that are output-identical by construction. Supplies the one + # thing the orchestrator cannot derive: the environment that turns the + # feature OFF, and (for host_offload) the memory delta predicted from the + # checkpoint's own shapes. The bounded-replay features also declare their + # workload here -- see `REPLAY_WITNESS_REQUIRED`. + "witness", + # `trtllm-eval` takes these two as its OWN flags rather than through the + # options YAML, so they have to be declarable here -- see + # `gates.eval_command`, which reads exactly these names. + "max_seq_len", + "max_batch_size", + "benchmark_args", + # How the accuracy floor is derived; see `FLOOR_KEYS`. + "floor", + # The boot gate's workload, declared here; see `BOOT_KEYS`. + "boot", + # The target at its declared limits; see `ENVELOPE_KEYS`. + "envelope", + # The target served by `trtllm-serve`; see `SERVE_KEYS`. + "serve", + # The benchmark's sample count, scalar or per unit, for the standard + # error behind the statistical floor. Needed only for a benchmark whose + # count the gate does not know (`stats.KNOWN_BENCHMARK_SAMPLES`) and whose + # eval does not print one. + "samples", + # Which members of a unit run the benchmark; see `BENCHMARK_MEMBERS`. + "benchmark_members", + # Whether the anchor is written down yet, scalar or per unit; see `ANCHOR_PENDING`. + "anchor", +) + +# `accuracy.anchor: pending` (scalar, or per unit): no written anchor exists +# for the unit yet. Its accuracy benchmark is not a check, since a score +# compared with nothing proves nothing, and the run must not measure its own +# reference; the logit gate against the HF reference (`accuracy.logits`) is +# the unit's correctness oracle, and delivery stays incomplete until a row is +# recorded from a written source. +ANCHOR_PENDING = "pending" +ANCHOR_RECORDED = "recorded" +ANCHOR_STATES: tuple[str, ...] = (ANCHOR_RECORDED, ANCHOR_PENDING) + +# `accuracy.benchmark_members`: where each unit's accuracy benchmark runs. +# +# sampled (default) with the logit gate on: at the narrowest member (tp1), +# the widest, and every member whose served engine the logit gate +# does not run whole (attention DP; a speculator, unless the boot +# compares greedy tokens). The logit gate, owed at every member, +# proves the rest equivalent to the reference. Without the logit +# gate, at every member. +# every at every member, whatever else the contract owes. +BENCHMARK_MEMBERS_SAMPLED = "sampled" +BENCHMARK_MEMBERS_EVERY = "every" +BENCHMARK_MEMBERS: tuple[str, ...] = (BENCHMARK_MEMBERS_SAMPLED, BENCHMARK_MEMBERS_EVERY) + +# The boot gate's workload, `accuracy.boot`: run-wide keys, plus optional +# per-unit entries (unit name -> the same keys) layered over them. With no +# block the workspace `boot.yaml` is read. +# +# prompts strings, or mappings {prompt, keywords, keywords_any, images, +# render}: a mapping's keywords are checked against THAT +# prompt's generated text; `images` are files sent with it +# (relative paths resolve against the spec's directory when the +# file is there, else against the tree under test, where the +# repository's test assets are); `render` is `loader` (the +# tree's multimodal input loader applies the chat template) or +# `raw` (the prompt is pre-rendered, placeholders included) +# prompts_file a YAML/JSON list of plain prompts, or boot.yaml's shape +# keywords each must appear in SOME generated text +# max_tokens per prompt (default 32) +# greedy how the HF reference's greedy tokens are compared (`GREEDY_KEYS`); +# a level holding only this declares no workload +# +# Keywords are matched against the generated texts only, never the log. A unit +# whose target implements `multimodal` and declares no image prompt sends the +# repository's own `tests/integration/test_input_files/merlion.png`. +BOOT_KEYS: tuple[str, ...] = ("prompts", "prompts_file", "keywords", "max_tokens", "greedy") +BOOT_PROMPT_KEYS: tuple[str, ...] = ("prompt", "keywords", "keywords_any", "images", "render") +BOOT_RENDERS: tuple[str, ...] = ("loader", "raw") + +# `accuracy.boot.greedy` (run-wide, or a unit's own entry): wherever the run has +# an HF logit reference (`accuracy.logits.reference`), the boot also sends its +# first text prompts as token ids and the engine's first greedy tokens must be +# the reference's, up to a near-tie (checks/gates.py, `GreedySetting`). +# +# enabled default true; false compares no tokens. A configuration variant +# compares none unless its own entry says true. +# prompts how many of the reference's text prompts (default 4) +# tokens greedy tokens compared per prompt (default 24, at most the +# reference's continuation) +GREEDY_KEYS: tuple[str, ...] = ("enabled", "prompts", "tokens") + +# The envelope boot, `accuracy.envelope` (envelope_boot.py). All optional: +# +# enabled default true; it runs once per member at delivery +# max_tokens tokens generated per request (default 16) +# agree_tokens leading token ids each probe must keep under load (default 4) +# long_prompts requests at the sequence limit (default 2) +# probes the short prompts compared against a batch-1 run +# timeout seconds per leg +ENVELOPE_KEYS: tuple[str, ...] = ( + "enabled", + "max_tokens", + "agree_tokens", + "long_prompts", + "probes", + "timeout", +) + +# The serve gate, `accuracy.serve` (checks/serve_gate.py): `trtllm-serve` starts +# each unit's target and answers one completion, and one chat with an image for +# a unit whose boot sends one. All optional: +# +# enabled default: on wherever the server's configuration can be derived, +# that is wherever the spec bounds the sequence length the server +# sizes its cache for (`accuracy.max_seq_len`, or `max_seq_len` in +# the run-wide options file); without one the server would size it +# for the checkpoint's own maximum. The loader writes the value it +# chose (`derive_gate_defaults`); false turns it off +# prompt the completion's prompt (default "The capital of France is") +# max_tokens tokens per answer (default 32) +# timeout seconds the server may take to answer /health (default 1800) +SERVE_KEYS: tuple[str, ...] = ("enabled", "prompt", "max_tokens", "timeout") + +# The `accuracy.floor` block. +# +# rule `statistical` (default): floor = anchor - k * SE, the SE +# binomial at the anchor over the benchmark's sample count. +# `tolerance`: floor = anchor - `accuracy.tolerance`, a +# fixed width; an explicit choice +# k default 2.0. GSM8K's 1319 questions at 92 % give SE 0.74, +# so the floor sits ~1.5 points under the anchor, where a +# fixed 5.0 sits ~7 SE under it +# shift_min_members default 2. When every member of a unit scores more than +# 1 SE under the anchor, that is a FINDING that blocks +# delivery, even though each member passed its own floor -- +# but only for a unit gated over at least this many members +# acknowledge unit -> reason. A finding on that unit is recorded with +# the reason and does not block delivery +FLOOR_KEYS: tuple[str, ...] = ("rule", "k", "shift_min_members", "acknowledge") + +ENVIRONMENT_KEYS: tuple[str, ...] = ( + "kind", + "python", + "env", + "spool", + "partition", + "account", + "container_image", + "mounts", + "time_limit", + "gpus_per_node", + # Pyxis reuses an already-imported container of this name on the node, + # which is the difference between paying 4-8 minutes per job for a 26 GB + # image and paying it once. + "container_name", + # A checkpoint the STOCK modeling path already supports, used by + # `preflight.substrate_gate` when the run's own checkpoint has an + # architecture that path has never heard of. Optional, because a + # checkpoint the stock path supports needs no proxy -- and declared + # rather than guessed, because this block is closed and because which + # model stands in for another is a judgement nothing here can make. + "substrate_checkpoint", + # The compute sessions the run stands on. Validated by + # `substrate.sessions.validate_sessions`. + "sessions", + # How long the orchestrator waits for a gate session that stopped + # heart-beating before it closes the run resumably. Also validated there. + "session_wait", +) + +ENVIRONMENT_KINDS: tuple[str, ...] = ("local", "session_spool") +ENVIRONMENT_REQUIRED_FIELDS: tuple[str, ...] = ("kind", "python") + +# Accuracy tolerance, in benchmark points, under `accuracy.floor.rule: +# tolerance` only. The default rule is statistical (`FLOOR_KEYS`): 5.0 points +# is about seven standard errors on GSM8K, wide enough that a defect moving +# every member by half a point passes every gate. +DEFAULT_ACCURACY_TOLERANCE = 5.0 + + +class TaskSchemaError(ValueError): + """Raised when ``task.yaml`` fails trtllm-modeling-v2 schema validation.""" + + +def parse_parallel_segment(segment: str) -> tuple[str, int | None]: + """Return ``(kind, world_size)`` for a parallel segment. + + ``world_size`` is ``None`` for ``unified``, whose member set — not the + segment string — decides how many devices the run claims. + """ + if segment == PARALLEL_UNIFIED: + return PARALLEL_UNIFIED, None + match = PARALLEL_PATTERN.match(segment) + if match is None: + raise TaskSchemaError( + f"'{PARALLEL_FIELD}' must be 'tp', 'tep', 'dep' or " + f"'{PARALLEL_UNIFIED}'; got {segment!r}" + ) + world_size = int(match.group(2)) + if world_size < 1: + raise TaskSchemaError(f"'{PARALLEL_FIELD}' world size must be >= 1; got {segment!r}") + return match.group(1), world_size + + +def members_for(data: Mapping[str, Any], unit: str | None, world_size: int) -> list[str]: + """The concrete parallel configurations ``unit`` is gated over. + + One function and one code path for both forms on purpose. A concrete + segment is the **degenerate** member set — itself, once — so the gate + loop in :meth:`~.workflow.TrtllmModelingV2Workflow._assemble_one_unit` + never branches on whether the run is unified, and no caller can forget + the unified case. + + ``members`` takes the same two shapes ``runtime`` and + ``accuracy.llm_api_options`` take: a list applies to every unit, and a + mapping from assembly unit to a list is layered over ``default``. The + per-unit form is what a run uses to buy a cheaper answer honestly — the + ``multimodal`` unit gated at one member, named, rather than every unit + gated at one member and the target documented as covering three. + + Validation lives in :func:`_validate_members`; this resolves an + already-validated spec and falls back to the default set. + """ + segment = str(data.get(PARALLEL_FIELD) or DEFAULT_PARALLEL) + if segment != PARALLEL_UNIFIED: + return [segment] + block = data.get(MEMBERS_FIELD) + if isinstance(block, dict): + chosen = block.get(unit) + block = block.get("default") if chosen is None else chosen + if isinstance(block, list) and block: + return [str(member) for member in block] + return [f"{kind}{world_size}" for kind in DEFAULT_MEMBER_KINDS] + + +def widest_world_size(data: Mapping[str, Any], default: int) -> int: + """The largest world size any member of a validated spec declares. + + A concrete segment answers for itself. A ``unified`` target takes the + widest of its members, per-unit lists and configuration variants + included, because that is the most devices any one gate asks for; with + no members declared it is ``default``. + """ + _, size = parse_parallel_segment(str(data.get(PARALLEL_FIELD) or DEFAULT_PARALLEL)) + if size is not None: + return size + block = data.get(MEMBERS_FIELD) + declared: list[Any] = [] + if isinstance(block, dict): + for entries in block.values(): + declared.extend(entries or ()) + elif isinstance(block, list): + declared.extend(block) + for entry in (data.get(VARIANTS_FIELD) or {}).values(): + declared.extend((entry or {}).get("members") or ()) + sizes = [] + for member in declared: + try: + sizes.append(parse_parallel_segment(str(member))[1] or 0) + except TaskSchemaError: + continue + return max(sizes, default=0) or default + + +def derive_slug(data: Mapping[str, Any]) -> str: + """Return ``--`` for a validated spec. + + The checkpoint segment is the checkpoint directory's own name, + lowercased with non-alphanumerics folded to hyphens — it names the + branch, the workspace subdirectory and the run's reports, so it has to + survive being a path component on any filesystem. + """ + checkpoint = Path(str(data["checkpoint_path"]).rstrip("/")).name + checkpoint = re.sub(r"[^a-zA-Z0-9]+", "-", checkpoint).strip("-").lower() + return f"{checkpoint}-{data[ARCH_FIELD]}-{data[PARALLEL_FIELD]}" + + +def derive_family(architecture: str) -> str: + """``GptOssForCausalLM`` -> ``gpt_oss``: the modeling_v2 family directory. + + The family groups every target for one architecture, and in tree it is + snake_case with the ``For`` suffix dropped — ``gpt_oss``, + ``deepseek_v3``. Derived rather than declared because the checkpoint + already names its architecture, and a hand-typed family that disagrees + with the routing index is a target nothing can reach. + """ + name = re.sub(r"For(Causal|Conditional|Seq2Seq|Sequence)?LM$", "", architecture) + name = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "_", name) + name = re.sub(r"(?<=[A-Z])(?=[A-Z][a-z])", "_", name) + return re.sub(r"[^a-zA-Z0-9]+", "_", name).strip("_").lower() + + +def derive_checkpoint_segment(checkpoint_path: str | Path) -> str: + """``DeepSeek-R1-0528-NVFP4`` -> ``deep_seek_r1_0528_nvfp4``. + + **A guess, and the one identity component worth overriding.** In tree the + segment sits *under* the family directory, so the family's own name is + usually redundant in it — ``deepseek_v3/r1_0528_nvfp4__sm_103__dep4`` rather + than ``deepseek_v3/deepseek_r1_0528_nvfp4__sm_103__dep4``. Nothing can derive + which part is redundant, so set ``target.checkpoint_segment`` in + task.yaml when the default reads badly. + """ + name = Path(str(checkpoint_path).rstrip("/")).name + # Split camelCase only at a genuine word boundary — a lowercase letter + # followed by a capital. Allowing a digit on the left would break the + # size suffixes these names are full of: `Qwen3-30B-A3B` would come out + # as `qwen3_30_b_a3_b`. + name = re.sub(r"(?<=[a-z])(?=[A-Z])", "_", name) + return re.sub(r"[^a-zA-Z0-9]+", "_", name).strip("_").lower() + + +def _validate_coverage(block: Any, errors: list[str]) -> None: + """``coverage``: a mapping whose ``accept`` lists ``@`` or ``@*``.""" + if block is None: + return + if not isinstance(block, dict): + errors.append(f"'{COVERAGE_FIELD}' must be a mapping, got {type(block).__name__}") + return + reject_unknown_keys(block, COVERAGE_KEYS, errors, where=COVERAGE_FIELD) + raw = block.get("accept") + if raw is not None and ( + not isinstance(raw, list) or not all(isinstance(x, str) and "@" in x for x in raw) + ): + errors.append( + f"'{COVERAGE_FIELD}.accept' must be a list of '@' or '@*'" + ) + + +def _validate_target(block: Any, errors: list[str]) -> None: + """Validate the optional ``target`` block, which overrides the derivation.""" + if not isinstance(block, dict): + raise TaskSchemaError(f"'{TARGET_FIELD}' must be a mapping, got {type(block).__name__}") + reject_unknown_keys(block, TARGET_KEYS, errors, where=TARGET_FIELD) + keep = block.get("keep_local") + if keep is not None and ( + not isinstance(keep, list) + or not all(isinstance(e, str) and re.match(r"^[a-z_]+/[a-z0-9_]+$", e) for e in keep) + ): + errors.append( + f"'{TARGET_FIELD}.keep_local' lists shared catalog entries as " + f"'/' (lowercase with underscores); got {keep!r}" + ) + for name in TARGET_NAME_KEYS: + value = block.get(name) + if value is None: + continue + if not isinstance(value, str) or not value.strip(): + errors.append(f"'{TARGET_FIELD}.{name}' must be a non-empty string") + elif not re.match(r"^[a-z0-9_]+$", value): + errors.append( + f"'{TARGET_FIELD}.{name}' is a Python package directory name, so it must " + f"be lowercase with underscores; got {value!r}" + ) + + +def _validate_runtime_switches(block: dict, errors: list[str], where: str, *, fill: bool) -> None: + """Validate one runtime block. ``fill`` writes the defaults in.""" + unknown = sorted( + set(block) - set(RUNTIME_SWITCHES) - set(RUNTIME_MARGIN_DEFAULTS) - {"extra_flags"} + ) + if unknown: + errors.append( + f"'{where}' has unknown switch(es) {unknown}; known switches are " + f"{sorted(RUNTIME_SWITCHES)} plus 'extra_flags' and " + f"{sorted(RUNTIME_MARGIN_DEFAULTS)}" + ) + for name in RUNTIME_MARGIN_DEFAULTS: + value = block.get(name) + if value is None: + continue + if isinstance(value, bool) or not isinstance(value, (int, float)) or value < 0: + errors.append(f"'{where}.{name}' must be a non-negative number, got {value!r}") + elif name == "min_free_memory_fraction" and value >= 1: + errors.append(f"'{where}.{name}' is a fraction of the device, below 1; got {value!r}") + + for name, default in RUNTIME_SWITCHES.items(): + value = block.get(name) + if value is None: + if fill: + block[name] = default + elif not isinstance(value, bool): + errors.append(f"'{where}.{name}' must be a boolean, got {type(value).__name__}") + + extra = block.get("extra_flags") + if extra is None: + if fill: + block["extra_flags"] = [] + elif not isinstance(extra, list) or not all(isinstance(f, str) for f in extra): + errors.append(f"'{where}.extra_flags' must be a list of strings") + + +def _validate_runtime(block: Any, errors: list[str]) -> None: + """Validate the optional ``runtime`` block in place, filling defaults. + + Unknown keys are an error rather than a warning. A misspelled switch + that is silently ignored produces a run that gates a feature set + nobody asked for, and the report still says the feature was covered — + which is the expensive kind of wrong. + """ + if not isinstance(block, dict): + raise TaskSchemaError(f"'{RUNTIME_FIELD}' must be a mapping, got {type(block).__name__}") + + # Per-unit form: a mapping whose values are themselves mappings. The units + # of one run can genuinely DISAGREE about a switch -- `core` and `mtp` can + # pass with block reuse ON while the multimodal path requires it OFF -- + # and then no run-wide value is right for all three. A per-unit block + # is layered over `default`, so a unit states only what differs. + if any(isinstance(v, dict) for v in block.values()): + nested = {k: v for k, v in block.items() if isinstance(v, dict)} + if len(nested) != len(block): + errors.append( + f"'{RUNTIME_FIELD}' mixes the flat form (cuda_graph: true) with the " + f"per-unit form (core: {{...}}); pick one" + ) + if "default" not in block: + errors.append( + f"'{RUNTIME_FIELD}' is a per-unit mapping, so it needs a 'default' key " + f"-- a unit nobody listed would otherwise be gated under the engine's " + f"defaults rather than this run's" + ) + for unit, sub in nested.items(): + # Only `default` is filled in: a per-unit block is an OVERLAY, so + # filling its absent switches with the schema defaults would + # silently override `default` with them. + _validate_runtime_switches( + sub, errors, f"{RUNTIME_FIELD}.{unit}", fill=(unit == "default") + ) + return + + _validate_runtime_switches(block, errors, RUNTIME_FIELD, fill=True) + + +def _validate_members( + block: Any, parallel: Any, errors: list[str], *, environment_kind: str = "local" +) -> None: + """Validate the optional ``members`` block: what a unified run is gated over. + + Three rules, each preventing a distinct silent wrong: + + - **Only under ``parallel: unified``.** A concrete segment already names + one configuration, so a ``members`` list beside it is a second and + quieter declaration of the same thing, and the two can disagree. There + is no reading of ``parallel: tep4`` plus ``members: [dep4]`` that is + not a mistake in one of them. + - **No ``unified`` member.** A member is a configuration to *serve at*; + ``unified`` is the family, not a member of it. Accepting it would let + a unified target be gated over one path under another name — the gap + this key exists to close. + - **One world size, on a local run.** Preflight claims one set of + devices for a ``local`` run, so a member asking for a different device + count cannot be run by it at all. The mismatch is reported with both + values rather than quietly re-sized to the claim, because a member + silently served at half its declared width is a measurement of something + nobody asked for. Under a scheduler (``session_spool``) each gate job + asks for its own member's devices, so the members may differ. + """ + if parallel != PARALLEL_UNIFIED: + errors.append( + f"'{MEMBERS_FIELD}' is only legal with '{PARALLEL_FIELD}: {PARALLEL_UNIFIED}'; " + f"this spec says '{PARALLEL_FIELD}: {parallel}', which already names the one " + f"configuration the run is gated at" + ) + return + + # Same two shapes `runtime` and `accuracy.llm_api_options` take: a list + # applies to every unit; a mapping from unit to list is layered over + # `default`. + if isinstance(block, dict): + if "default" not in block: + errors.append( + f"'{MEMBERS_FIELD}' is a per-unit mapping, so it needs a 'default' key " + f"-- a unit nobody listed would otherwise fall back to the built-in " + f"member set and be gated over configurations this spec did not choose" + ) + lists = {f"{MEMBERS_FIELD}.{unit}": value for unit, value in block.items()} + elif isinstance(block, list): + lists = {MEMBERS_FIELD: block} + else: + errors.append( + f"'{MEMBERS_FIELD}' must be a list of parallel segments, or a mapping from " + f"assembly unit to one; got {type(block).__name__}" + ) + return + + # Keyed by the segment string so the mismatch message can name the two + # members that disagree rather than only the two numbers. + world_sizes: dict[str, int] = {} + for where, entries in lists.items(): + if not isinstance(entries, list) or not entries: + errors.append( + f"'{where}' must be a non-empty list of parallel segments; an empty one " + f"means a unit that is built and never gated" + ) + continue + for index, entry in enumerate(entries): + if not isinstance(entry, str): + errors.append(f"'{where}[{index}]' must be a string, got {type(entry).__name__}") + continue + try: + kind, size = parse_parallel_segment(entry) + except TaskSchemaError as exc: + errors.append(f"'{where}[{index}]' is not a parallel segment: {exc}") + continue + if kind == PARALLEL_UNIFIED or size is None: + errors.append( + f"'{where}[{index}]' is '{PARALLEL_UNIFIED}', which is the family " + f"rather than a configuration to serve at. A member is concrete: " + f"'tp4', 'tep4', 'dep4'." + ) + continue + world_sizes[entry] = size + + if len(set(world_sizes.values())) > 1 and environment_kind == "local": + widest = max(world_sizes, key=lambda m: world_sizes[m]) + narrowest = min(world_sizes, key=lambda m: world_sizes[m]) + errors.append( + f"'{MEMBERS_FIELD}' mixes world sizes: {narrowest} wants " + f"{world_sizes[narrowest]} device(s) and {widest} wants {world_sizes[widest]}. " + f"A local run claims one set of devices for every member, so they must all " + f"declare the same world size; under `environment.kind: session_spool` each " + f"gate asks the scheduler for its own member's devices." + ) + + +def _validate_source_embargo(block: Any, errors: list[str], repo: Any) -> None: + """Validate ``source_embargo``: paths agents may not read. + + Two shapes, and the split is on where the path points rather than on how + it is spelled: + + - **repo-relative**, resolved under ``trtllm_repo_path``. The common + case: TensorRT-LLM already ships an implementation of this checkpoint, + so name it and every role prompt forbids opening it. + - **absolute and outside the repo.** The tempting file is often not in + this repository at all — a sibling vLLM or SGLang checkout on the same + filesystem implements the same checkpoint, is one ``Read`` away from + every agent, and is *more* likely to be found because the run + deliberately starts one of those engines as its accuracy oracle. What + makes the embargo necessary is the file being reachable, not which + repository it lives in. + + An absolute path that resolves *inside* the repo is rejected: spell it + repo-relative so the list reads the same from any working directory. + + Either way the entry must **exist**. An embargo naming a path that is not + there reads as protection while protecting nothing, and the typo it + usually is would leave the real file readable. A **directory** entry is + legal and embargoes the tree beneath it — a sibling engine checkout is a + tree, and asking a spec to enumerate its model files is asking it to go + stale. + """ + if not isinstance(block, list) or not all(isinstance(p, str) for p in block): + raise TaskSchemaError( + f"'{SOURCE_EMBARGO_FIELD}' must be a list of paths: repo-relative, absolute for a " + f"tree outside the repo, or `{PACKAGE_PREFIX}` inside the workflow package" + ) + if not isinstance(repo, str): + return + repo_root = Path(repo).resolve() + for entry in block: + if entry.startswith(PACKAGE_PREFIX): + problem = package_entry_problem(entry) + if problem: + errors.append(f"'{SOURCE_EMBARGO_FIELD}' names {entry!r}, which {problem}.") + elif not resolve_entry(entry).exists(): + errors.append( + f"'{SOURCE_EMBARGO_FIELD}' names {entry!r}, which does not exist under " + f"the workflow package. An embargo on a path that is not there protects " + f"nothing." + ) + continue + if entry.startswith("/"): + resolved = Path(entry).resolve() + if resolved == repo_root or repo_root in resolved.parents: + errors.append( + f"'{SOURCE_EMBARGO_FIELD}' names {entry!r}, which is inside " + f"{repo}. Spell an in-repo path repo-relative, so the list means " + f"the same thing from any working directory." + ) + elif not resolved.exists(): + errors.append( + f"'{SOURCE_EMBARGO_FIELD}' names {entry!r}, which does not exist. " + f"An embargo on a path that is not there protects nothing." + ) + elif not (Path(repo) / entry).exists(): + errors.append( + f"'{SOURCE_EMBARGO_FIELD}' names {entry!r}, which does not exist under " + f"{repo}. An embargo on a path that is not there protects nothing." + ) + + +def _validate_environment(block: Any, errors: list[str]) -> None: + """Validate the ``environment`` block in place, appending to ``errors``.""" + if not isinstance(block, dict): + raise TaskSchemaError( + f"'{ENVIRONMENT_FIELD}' must be a mapping, got {type(block).__name__}" + ) + + reject_unknown_keys(block, ENVIRONMENT_KEYS, errors, where=ENVIRONMENT_FIELD) + + validate_sessions(block, errors, where=ENVIRONMENT_FIELD) + + for name in ENVIRONMENT_REQUIRED_FIELDS: + value = block.get(name) + if not isinstance(value, str) or not value.strip(): + errors.append(f"'{ENVIRONMENT_FIELD}.{name}' must be a non-empty string") + + kind = block.get("kind") + if isinstance(kind, str) and kind not in ENVIRONMENT_KINDS: + errors.append( + f"'{ENVIRONMENT_FIELD}.kind' must be one of {ENVIRONMENT_KINDS}; got {kind!r}" + ) + + env = block.get("env") + if env is None: + block["env"] = {} + elif not isinstance(env, dict): + errors.append(f"'{ENVIRONMENT_FIELD}.env' must be a mapping, got {type(env).__name__}") + else: + for key, value in env.items(): + if not isinstance(value, (str, int, float)): + errors.append( + f"'{ENVIRONMENT_FIELD}.env.{key}' must be a scalar the shell can " + f"carry, got {type(value).__name__}" + ) + # Imported here: the preflight module imports this one. + from ..run.preflight import PreflightError, check_tmpdir + + try: + check_tmpdir({str(k): str(v) for k, v in env.items()}, require_set=False) + except PreflightError as exc: + errors.append(f"'{ENVIRONMENT_FIELD}.env': {exc}") + + mounts = block.get("mounts") + if mounts is not None and ( + not isinstance(mounts, list) or not all(isinstance(m, str) for m in mounts) + ): + errors.append(f"'{ENVIRONMENT_FIELD}.mounts' must be a list of strings") + + # Same must-exist rule the embargo list carries, for the same reason: a + # proxy that is not there is discovered by the substrate gate, which is + # the longest step in preflight and the one the proxy exists to rescue. + proxy = block.get("substrate_checkpoint") + if proxy is not None: + if not isinstance(proxy, str) or not proxy.strip(): + errors.append( + f"'{ENVIRONMENT_FIELD}.substrate_checkpoint' must be a path to a " + f"checkpoint directory the stock modeling path supports" + ) + elif not Path(proxy).exists(): + errors.append( + f"'{ENVIRONMENT_FIELD}.substrate_checkpoint' does not exist: {proxy}. " + f"The substrate gate falls back to it only after the run's own " + f"checkpoint has already failed, so a missing proxy turns the " + f"longest step in preflight into two failures instead of one." + ) + + if kind == "session_spool": + # The spool is the whole contract. It must already exist, because a + # spool nobody is watching accepts work forever and completes none + # of it — the run would hang on its first gate with no error. + spool = block.get("spool") + if not isinstance(spool, str) or not spool.strip(): + errors.append(f"'{ENVIRONMENT_FIELD}.spool' is required when kind is 'session_spool'") + elif not Path(spool).is_dir(): + errors.append( + f"'{ENVIRONMENT_FIELD}.spool' does not exist: {spool}. Start the compute " + f"session first — an unwatched spool accepts work and never completes it." + ) + + +def _positive_int(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value >= 1 + + +def _number(value: Any) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) + + +def _validate_contrast_env(entry: Mapping[str, Any], errors: list[str], where: str) -> None: + contrast_env = entry.get("contrast_env") + if not isinstance(contrast_env, dict) or not contrast_env: + errors.append( + f"'{where}.contrast_env' is required and must be a non-empty mapping: " + f"a witness with no changed variable compares a run to itself" + ) + elif not all(isinstance(v, (str, int, float, bool)) for v in contrast_env.values()): + errors.append(f"'{where}.contrast_env' values must be scalars the shell can carry") + + +def _validate_witness_timeout(entry: Mapping[str, Any], errors: list[str], where: str) -> None: + timeout = entry.get("timeout") + if timeout is not None and not _positive_int(timeout): + errors.append(f"'{where}.timeout' must be a whole number of seconds >= 1, got {timeout!r}") + + +def _validate_witness_workload(entry: Mapping[str, Any], errors: list[str], where: str) -> None: + """Validate a token-comparing witness block's declared workload.""" + _validate_acceptance_workload( + {k: v for k, v in entry.items() if k in ("prompts", "prompts_file", "max_tokens")}, + errors, + where, + ) + minimum = entry.get("min_generated_tokens") + if minimum is not None and not _positive_int(minimum): + errors.append(f"'{where}.min_generated_tokens' must be an integer >= 1") + + +def _validate_witness_batch_size(entry: Mapping[str, Any], errors: list[str], where: str) -> None: + value = entry.get(WITNESS_BATCH_SIZE_KEY) + if value is not None and not _positive_int(value): + errors.append( + f"'{where}.{WITNESS_BATCH_SIZE_KEY}' must be an integer >= 1 (1 gives both legs " + f"the same schedule), got {value!r}" + ) + + +_ENV_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") + + +def _validate_legs_env(entry: Mapping[str, Any], errors: list[str], where: str) -> None: + """Validate the optional ``legs_env``: shell variables, none of them the contrast.""" + legs_env = entry.get(WITNESS_LEGS_ENV_KEY) + if legs_env is None: + return + key = f"{where}.{WITNESS_LEGS_ENV_KEY}" + if not isinstance(legs_env, dict) or not legs_env: + errors.append(f"'{key}' must be a non-empty mapping from variable to value") + return + for name, value in legs_env.items(): + if not isinstance(name, str) or not _ENV_NAME.fullmatch(name): + errors.append(f"'{key}' names {name!r}, which is not a shell variable name") + if not isinstance(value, (str, int, float, bool)): + errors.append(f"'{key}.{name}' must be a scalar the shell can carry") + contrast_env = entry.get("contrast_env") + if isinstance(contrast_env, dict): + shared = sorted(str(name) for name in set(legs_env) & set(contrast_env)) + if shared: + errors.append( + f"'{key}' and '{where}.contrast_env' both set {', '.join(shared)}: a variable " + f"leg B sets differently is the contrast, not something the legs share" + ) + + +def _validate_equivalence_witness(entry: Mapping[str, Any], errors: list[str], where: str) -> None: + """Validate an equivalence witness block: leg B is the flow's launcher, or the run's.""" + reject_unknown_keys(entry, EQUIVALENCE_WITNESS_KEYS, errors, where=where) + command = entry.get("contrast_command") + if command is not None and (not isinstance(command, str) or not command.strip()): + errors.append( + f"'{where}.contrast_command' must be a non-empty string, or left out for the " + f"flow's own leg B" + ) + if command is not None or entry.get("contrast_env") is not None: + _validate_contrast_env(entry, errors, where) + _validate_legs_env(entry, errors, where) + transceiver = entry.get(WITNESS_TRANSCEIVER_KEY) + if transceiver is not None and ( + not isinstance(transceiver, dict) + or not transceiver + or not all(isinstance(k, str) for k in transceiver) + ): + errors.append( + f"'{where}.{WITNESS_TRANSCEIVER_KEY}' must be a non-empty mapping of " + f"CacheTransceiverConfig fields (e.g. backend: NIXL)" + ) + lengths = entry.get("prompt_tokens") + if lengths is not None: + if not isinstance(lengths, list) or not lengths: + errors.append(f"'{where}.prompt_tokens' must be a non-empty list of integers") + else: + for value in lengths: + if not _positive_int(value): + errors.append( + f"'{where}.prompt_tokens' entries must be integers >= 1, got {value!r}" + ) + elif value > MAX_WITNESS_PROMPT_WORDS: + errors.append( + f"'{where}.prompt_tokens' entry {value} is over " + f"{MAX_WITNESS_PROMPT_WORDS}; each prompt travels as one argv string, " + f"which Linux caps at 128 KiB" + ) + _validate_witness_timeout(entry, errors, where) + _validate_witness_batch_size(entry, errors, where) + _validate_witness_workload(entry, errors, where) + + +def _validate_replay_witness( + feature: str, entry: Mapping[str, Any], errors: list[str], where: str +) -> None: + """Validate a bounded-replay witness block: the workload is DECLARED here. + + Every size check below exists because the opposite value turns the witness + into one that passes for a target that implements nothing: a prompt no + longer than the window makes replay-on and replay-off do identical work, so + "the decoder saw at most `window` rows" is true of both. + """ + reject_unknown_keys( + entry, + (*REPLAY_WITNESS_REQUIRED[feature], *REPLAY_WITNESS_OPTIONAL[feature]), + errors, + where=where, + ) + for key in REPLAY_WITNESS_REQUIRED[feature]: + if entry.get(key) is None: + errors.append( + f"'{where}.{key}' is required: the {feature} witness sizes its workload " + f"from the spec, never from the run it judges" + ) + _validate_contrast_env(entry, errors, where) + + window = entry.get("window") + if window is not None and not _positive_int(window): + errors.append(f"'{where}.window' must be an integer >= 1, got {window!r}") + window = None + for key in ("prompt_tokens", "prefix_tokens"): + value = entry.get(key) + if value is None or key not in REPLAY_WITNESS_REQUIRED[feature]: + continue + if not _positive_int(value): + errors.append(f"'{where}.{key}' must be an integer >= 1, got {value!r}") + elif window is not None and value <= window: + errors.append( + f"'{where}.{key}' is {value}, which does not exceed window={window}. At or " + f"under the window, replay on and replay off do the same work, and the " + f"witness passes for a target that implements nothing." + ) + elif value > MAX_WITNESS_PROMPT_WORDS: + errors.append( + f"'{where}.{key}' is {value}; the witness prompt travels as one argv " + f"string, which Linux caps at 128 KiB, so the limit is " + f"{MAX_WITNESS_PROMPT_WORDS}" + ) + alignment = entry.get("reuse_alignment") + if alignment is not None and not _positive_int(alignment): + errors.append(f"'{where}.reuse_alignment' must be an integer >= 1, got {alignment!r}") + _validate_witness_timeout(entry, errors, where) + + +def _validate_kv_quant_witness(entry: dict[str, Any], errors: list[str], where: str) -> None: + """Validate the ``kv_quant`` witness block, filling ``exact``'s default. + + ``exact`` is written back explicitly because it changes what the unit + claims, and the resolved spec in the workspace is what a reader checks: a + unit gated without its identity leg should say so there, not by omission. + """ + reject_unknown_keys(entry, KV_QUANT_WITNESS_KEYS, errors, where=where) + expected = entry.get("expected_bytes_per_token") + if not _positive_int(expected): + errors.append( + f"'{where}.expected_bytes_per_token' is required and must be an integer >= 1, " + f"computed from the checkpoint's storage formats before the run; got {expected!r}" + ) + exact = entry.setdefault("exact", True) + if not isinstance(exact, bool): + errors.append(f"'{where}.exact' must be true or false, got {exact!r}") + elif exact or entry.get("contrast_env") is not None: + # The identity leg needs a second storage to compare against; without + # one it compares native storage with itself and passes. + _validate_contrast_env(entry, errors, where) + tolerance = entry.get("bytes_tolerance") + if tolerance is not None and ( + not isinstance(tolerance, (int, float)) + or isinstance(tolerance, bool) + or not 0 <= tolerance < 1 + ): + errors.append(f"'{where}.bytes_tolerance' must be a fraction in [0, 1), got {tolerance!r}") + _validate_witness_timeout(entry, errors, where) + _validate_witness_batch_size(entry, errors, where) + _validate_witness_workload(entry, errors, where) + + +def load_prompts_file(path: str | Path) -> tuple[str, ...]: + """Read a prompt list for the acceptance gate's declared workload. + + Two shapes, both already on disk in a run: a YAML or JSON **list** of + strings (a harness's ``prompts.json``), or a **mapping** with a + ``prompts:`` list (``boot.yaml``'s shape). Raises + :class:`TaskSchemaError` naming the file for anything else, including an + empty list: a workload with no prompts measures nothing. + """ + source = Path(path) + try: + data = yaml.safe_load(source.read_text(encoding="utf-8")) + except OSError as exc: + raise TaskSchemaError(f"cannot read the prompts file {source}: {exc}") from exc + except yaml.YAMLError as exc: + raise TaskSchemaError(f"the prompts file {source} is not valid YAML/JSON: {exc}") from exc + if isinstance(data, dict): + data = data.get("prompts") + if not isinstance(data, list) or not data: + raise TaskSchemaError( + f"the prompts file {source} must hold a non-empty list of strings, or a " + f"mapping with one under 'prompts'" + ) + if not all(isinstance(p, str) and p.strip() for p in data): + raise TaskSchemaError(f"every prompt in {source} must be a non-empty string") + return tuple(data) + + +def _validate_acceptance_workload(entry: Mapping[str, Any], errors: list[str], where: str) -> None: + """Validate the workload keys one level of ``accuracy.acceptance`` declares.""" + prompts = entry.get("prompts") + prompts_file = entry.get("prompts_file") + if prompts is not None and prompts_file is not None: + errors.append( + f"'{where}' declares both 'prompts' and 'prompts_file'; keep one, so the " + f"workload a verdict was measured on has a single source" + ) + if prompts is not None and ( + not isinstance(prompts, list) + or not prompts + or not all(isinstance(p, str) and p.strip() for p in prompts) + ): + errors.append(f"'{where}.prompts' must be a non-empty list of non-empty strings") + if prompts_file is not None: + if not isinstance(prompts_file, str) or not prompts_file.strip(): + errors.append(f"'{where}.prompts_file' must be a non-empty path") + elif not Path(prompts_file).is_file(): + errors.append(f"'{where}.prompts_file' does not exist: {prompts_file}") + else: + try: + load_prompts_file(prompts_file) + except TaskSchemaError as exc: + errors.append(f"'{where}.prompts_file': {exc}") + sources = [k for k in ("prompts", "prompts_file", "dataset") if entry.get(k) is not None] + if "dataset" in sources and len(sources) > 1: + errors.append( + f"'{where}' declares 'dataset' beside {sources[0]!r}; keep one, so the workload " + f"a verdict was measured on has a single source" + ) + dataset = entry.get("dataset") + if dataset is not None and (not isinstance(dataset, str) or not dataset.strip()): + errors.append( + f"'{where}.dataset' must be a local path or a cached dataset id (no download)" + ) + for key in ("text_field", "prompt_template", "split", "dataset_config"): + value = entry.get(key) + if value is not None and (not isinstance(value, str) or not value.strip()): + errors.append(f"'{where}.{key}' must be a non-empty string") + template = entry.get("prompt_template") + if isinstance(template, str) and "{text}" not in template: + errors.append(f"'{where}.prompt_template' must name the question as {{text}}") + for key in ("max_tokens", "min_draft_iterations", "num_prompts", "min_prompts"): + value = entry.get(key) + if value is not None and ( + not isinstance(value, int) or isinstance(value, bool) or value < 1 + ): + errors.append(f"'{where}.{key}' must be an integer >= 1") + + +def _validate_acceptance_regression(block: Any, errors: list[str], where: str) -> None: + """Validate ``accuracy.acceptance.regression``.""" + if not isinstance(block, dict): + errors.append(f"'{where}' must be a mapping with keys {list(ACCEPTANCE_REGRESSION_KEYS)}") + return + reject_unknown_keys(block, ACCEPTANCE_REGRESSION_KEYS, errors, where=where) + test = block.get("test") + if test is not None and test not in stats.PAIRED_TESTS: + errors.append(f"'{where}.test' must be one of {stats.PAIRED_TESTS}, got {test!r}") + alpha = block.get("alpha") + if alpha is not None and (not _number(alpha) or not 0 < alpha < 1): + errors.append(f"'{where}.alpha' must be a probability in (0, 1), got {alpha!r}") + effect = block.get("min_effect") + if effect is not None and (not _number(effect) or effect < 0): + errors.append(f"'{where}.min_effect' must be a number >= 0, got {effect!r}") + + +def _validate_acceptance(block: Any, errors: list[str]) -> None: + """Validate ``accuracy.acceptance``: per-unit floors plus the gate's workload. + + Every key that is not a workload key names a unit, and its value is the + unit's floor -- a number, or a mapping with ``min_absolute`` -- optionally + with that unit's own workload beside it. + """ + where = f"{ACCURACY_FIELD}.acceptance" + if not isinstance(block, dict): + errors.append( + f"'{where}' must be a mapping from unit to floor, plus the optional " + f"workload keys {list(ACCEPTANCE_WORKLOAD_KEYS)}" + ) + return + _validate_acceptance_workload(block, errors, where) + if block.get("regression") is not None: + _validate_acceptance_regression(block["regression"], errors, f"{where}.regression") + for unit, entry in block.items(): + if unit in ACCEPTANCE_BLOCK_KEYS: + continue + unit_where = f"{where}.{unit}" + if isinstance(entry, dict): + reject_unknown_keys(entry, ACCEPTANCE_UNIT_KEYS, errors, where=unit_where) + _validate_acceptance_workload(entry, errors, unit_where) + floor = entry.get("min_absolute") + floor_where = f"{unit_where}.min_absolute" + measured = entry.get("measured_reference") + if measured is not None and (not _number(measured) or measured < 1.0): + errors.append( + f"'{unit_where}.measured_reference' must be a number >= 1.0; got {measured!r}" + ) + source = entry.get("floor_source") + if source is not None and (not isinstance(source, str) or not source.strip()): + errors.append(f"'{unit_where}.floor_source' must be a non-empty string") + else: + floor = entry + floor_where = unit_where + if floor is None: + continue + # Acceptance length counts the target's own token on every step, so + # it is at least 1.0 by construction and a floor below that passes + # everything. + if not isinstance(floor, (int, float)) or isinstance(floor, bool) or floor < 1.0: + errors.append(f"'{floor_where}' must be a number >= 1.0; got {floor!r}") + + +def _validate_boot_level(level: Mapping[str, Any], errors: list[str], where: str) -> None: + """Validate one level of ``accuracy.boot``.""" + prompts = level.get("prompts") + if prompts is not None: + if not isinstance(prompts, list) or not prompts: + errors.append(f"'{where}.prompts' must be a non-empty list") + else: + for n, entry in enumerate(prompts): + at = f"{where}.prompts[{n}]" + if isinstance(entry, str): + if not entry.strip(): + errors.append(f"'{at}' must be a non-empty string") + continue + if not isinstance(entry, dict): + errors.append(f"'{at}' must be a string or a mapping") + continue + reject_unknown_keys(entry, BOOT_PROMPT_KEYS, errors, where=at) + if not isinstance(entry.get("prompt"), str) or not entry["prompt"].strip(): + errors.append(f"'{at}.prompt' is required and must be a non-empty string") + for key in ("keywords", "keywords_any", "images"): + value = entry.get(key) + if value is not None and ( + not isinstance(value, list) + or not all(isinstance(v, str) and v.strip() for v in value) + ): + errors.append(f"'{at}.{key}' must be a list of non-empty strings") + if entry.get("render") is not None and entry["render"] not in BOOT_RENDERS: + errors.append(f"'{at}.render' must be one of {BOOT_RENDERS}") + prompts_file = level.get("prompts_file") + if prompts_file is not None: + if not isinstance(prompts_file, str) or not Path(prompts_file).is_file(): + errors.append(f"'{where}.prompts_file' does not exist: {prompts_file}") + keywords = level.get("keywords") + if keywords is not None and ( + not isinstance(keywords, list) or not all(isinstance(k, str) for k in keywords) + ): + errors.append(f"'{where}.keywords' must be a list of strings") + tokens = level.get("max_tokens") + if tokens is not None and not _positive_int(tokens): + errors.append(f"'{where}.max_tokens' must be an integer >= 1") + greedy = level.get("greedy") + if greedy is not None: + at = f"{where}.greedy" + if not isinstance(greedy, dict): + errors.append(f"'{at}' must be a mapping") + return + reject_unknown_keys(greedy, GREEDY_KEYS, errors, where=at) + if greedy.get("enabled") is not None and not isinstance(greedy["enabled"], bool): + errors.append(f"'{at}.enabled' must be true or false") + for key in ("prompts", "tokens"): + if greedy.get(key) is not None and not _positive_int(greedy[key]): + errors.append(f"'{at}.{key}' must be an integer >= 1") + + +def _validate_boot(block: Any, errors: list[str]) -> None: + """Validate ``accuracy.boot``: the run-wide keys, then each unit's entry.""" + where = f"{ACCURACY_FIELD}.boot" + if not isinstance(block, dict): + errors.append(f"'{where}' must be a mapping") + return + _validate_boot_level(block, errors, where) + for unit, entry in block.items(): + if unit in BOOT_KEYS: + continue + if not isinstance(entry, dict): + errors.append( + f"'{where}.{unit}' is neither a boot key ({list(BOOT_KEYS)}) nor a unit's " + f"mapping of them" + ) + continue + reject_unknown_keys(entry, BOOT_KEYS, errors, where=f"{where}.{unit}") + _validate_boot_level(entry, errors, f"{where}.{unit}") + + +# The sample count a run actually scored, where the eval prints one: the MMLU +# line carries it in parentheses, and lm-eval's partial-score line (printed +# under `TLLM_EVAL_PARTIAL_SCORES_EVERY`) ends at `/`. +_MMLU_SAMPLES = re.compile(r"MMLU weighted average accuracy:\s*[0-9.]+\s*\((\d+)\)") +_PARTIAL_TOTAL = re.compile(r"Partial scores after\s+(\d+)/(\d+)\s+responses") +# `--num_samples=N` or `--num_samples N` among the benchmark arguments. +_NUM_SAMPLES_ARG = re.compile(r"^--num_samples(?:=(\d+))?$") + + +def resolve_per_unit(accuracy: Mapping[str, Any], key: str, unit: str | None) -> Any: + """Read ``accuracy[key]``, honouring the per-unit mapping form. + + Several fields of the accuracy block are legitimately different for a + feature increment than for core, and the mapping form is how a spec says + so. The scalar form means "the whole run". + + ``benchmark`` and ``reference_yaml`` are the load-bearing case: the + ``multimodal`` unit's release criterion is a multimodal benchmark, and + gating it on the text benchmark core was gated on is not a weaker check + — it is a check that cannot see the vision tower at all, so it passes + whatever the vision tower does. + """ + value = accuracy.get(key) + if isinstance(value, dict): + chosen = value.get(unit) + if chosen is None: + chosen = value.get("default") + return chosen + return value + + +def anchor_pending(accuracy: Mapping[str, Any], unit: str | None) -> bool: + """Whether ``unit`` has no written accuracy anchor yet (``accuracy.anchor: pending``).""" + return resolve_per_unit(accuracy, "anchor", unit) == ANCHOR_PENDING + + +def resolve_samples( + accuracy: Mapping[str, Any], + unit: str | None, + benchmark: str, + output: str = "", +) -> tuple[int | None, str]: + """How many samples the accuracy score was (or will be) measured on, and from where. + + Most specific first: the count the eval itself printed (what was actually + scored), then ``--num_samples`` among the unit's benchmark arguments, then + ``accuracy.samples`` (scalar or per unit), then + :data:`.stats.KNOWN_BENCHMARK_SAMPLES`. ``(None, reason)`` when none of + them answers; the statistical floor cannot be computed then, and the + schema refuses such a task before a run starts. + """ + if output: + mmlu = _MMLU_SAMPLES.findall(output) + if mmlu: + return int(mmlu[-1]), "the eval's own output" + partial = _PARTIAL_TOTAL.findall(output) + if partial: + return int(partial[-1][1]), "the eval's own output" + args = [str(a) for a in resolve_per_unit(accuracy, "benchmark_args", unit) or ()] + for i, arg in enumerate(args): + matched = _NUM_SAMPLES_ARG.match(arg) + if not matched: + continue + value = matched.group(1) or (args[i + 1] if i + 1 < len(args) else None) + try: + return int(str(value)), "--num_samples in benchmark_args" + except (TypeError, ValueError): + break + declared = resolve_per_unit(accuracy, "samples", unit) + if declared is not None: + try: + return int(declared), "accuracy.samples" + except (TypeError, ValueError): + pass + known = stats.KNOWN_BENCHMARK_SAMPLES.get(benchmark) + if known is not None: + return known, f"the {benchmark} split size" + return None, f"no sample count for `{benchmark}`: declare `accuracy.samples`" + + +def _validate_floor(block: Mapping[str, Any], errors: list[str]) -> None: + """Validate ``accuracy.floor`` and that its rule can be computed for every unit. + + The statistical rule needs a sample count per benchmark. Asked here + rather than at the gate: a floor that cannot be computed is found at the + end of a unit's first accuracy run otherwise, which is the most expensive + place to learn that a number was never declared. + """ + where = f"{ACCURACY_FIELD}.floor" + floor = block.get("floor") + if floor is not None: + if not isinstance(floor, dict): + errors.append(f"'{where}' must be a mapping with keys {list(FLOOR_KEYS)}") + floor = {} + reject_unknown_keys(floor, FLOOR_KEYS, errors, where=where) + rule = floor.get("rule") + if rule is not None and rule not in stats.FLOOR_RULES: + errors.append(f"'{where}.rule' must be one of {stats.FLOOR_RULES}, got {rule!r}") + k = floor.get("k") + if k is not None and (not _number(k) or k <= 0): + errors.append(f"'{where}.k' must be a number > 0, got {k!r}") + members = floor.get("shift_min_members") + if members is not None and not _positive_int(members): + errors.append(f"'{where}.shift_min_members' must be an integer >= 1, got {members!r}") + acknowledge = floor.get("acknowledge") + if acknowledge is not None and ( + not isinstance(acknowledge, dict) + or not all(isinstance(v, str) and v.strip() for v in acknowledge.values()) + ): + errors.append( + f"'{where}.acknowledge' must map a unit to the reason its systematic shift " + f"is accepted, as a non-empty string" + ) + else: + floor = {} + + samples = block.get("samples") + if samples is not None: + values = samples.values() if isinstance(samples, dict) else [samples] + if not all(_positive_int(v) for v in values): + errors.append( + f"'{ACCURACY_FIELD}.samples' must be an integer >= 1, or a mapping from unit " + f"to one; got {samples!r}" + ) + return + + if (floor.get("rule") or stats.FLOOR_STATISTICAL) != stats.FLOOR_STATISTICAL: + return + + benchmarks = block.get("benchmark") + pairs = ( + list(benchmarks.items()) + if isinstance(benchmarks, dict) + else [(None, benchmarks or "gsm8k")] + ) + for unit, benchmark in pairs: + if not isinstance(benchmark, str): + continue + count, source = resolve_samples(block, None if unit == "default" else unit, benchmark) + if count is None: + errors.append( + f"'{ACCURACY_FIELD}' uses the statistical floor, and there is {source} " + f"(unit {unit or 'every unit'}); the floor is anchor - k * SE and the SE " + f"needs it" + ) + + +def _validate_envelope(block: Any, errors: list[str]) -> None: + """Validate ``accuracy.envelope``.""" + where = f"{ACCURACY_FIELD}.envelope" + if not isinstance(block, dict): + errors.append(f"'{where}' must be a mapping") + return + reject_unknown_keys(block, ENVELOPE_KEYS, errors, where=where) + if block.get("enabled") is not None and not isinstance(block["enabled"], bool): + errors.append(f"'{where}.enabled' must be true or false") + for key in ("max_tokens", "agree_tokens", "timeout"): + if block.get(key) is not None and not _positive_int(block[key]): + errors.append(f"'{where}.{key}' must be an integer >= 1") + long_prompts = block.get("long_prompts") + if long_prompts is not None and ( + not isinstance(long_prompts, int) or isinstance(long_prompts, bool) or long_prompts < 0 + ): + errors.append(f"'{where}.long_prompts' must be an integer >= 0") + probes = block.get("probes") + if probes is not None and ( + not isinstance(probes, list) + or not probes + or not all(isinstance(p, str) and p.strip() for p in probes) + ): + errors.append(f"'{where}.probes' must be a non-empty list of prompts") + + +def _validate_serve(block: Any, errors: list[str]) -> None: + """Validate ``accuracy.serve``.""" + where = f"{ACCURACY_FIELD}.serve" + if not isinstance(block, dict): + errors.append(f"'{where}' must be a mapping") + return + reject_unknown_keys(block, SERVE_KEYS, errors, where=where) + if block.get("enabled") is not None and not isinstance(block["enabled"], bool): + errors.append(f"'{where}.enabled' must be true or false") + for key in ("max_tokens", "timeout"): + if block.get(key) is not None and not _positive_int(block[key]): + errors.append(f"'{where}.{key}' must be an integer >= 1") + prompt = block.get("prompt") + if prompt is not None and (not isinstance(prompt, str) or not prompt.strip()): + errors.append(f"'{where}.prompt' must be a non-empty string") + + +# ------------------------------------------------------------- accuracy.logits +# +# The logit gate, `accuracy.logits` (checks/logit_agreement.py): each unit's +# engine predicts, position by position, what the Hugging Face reference +# recorded in a reference file (onboard/logit_reference.py writes it). The +# contract owes it wherever `reference` names a file. A spec that names none, +# and does not turn the gate off, gets `DEFAULT_LOGIT_REFERENCE` beside it when +# that file is there (`derive_gate_defaults`). +# +# reference the reference file; relative to the spec's directory +# enabled default true; false keeps the block and owes nothing +# prompts the reference's prompt ids to send (default: every one the +# unit can take; an image prompt only to a multimodal unit) +# threshold_factor how many times the reference's own bf16-vs-fp32 spread the +# engine may differ by (default 3.0) +# timeout seconds for the job (default 1800) +LOGITS_KEYS: tuple[str, ...] = ("enabled", "reference", "prompts", "threshold_factor", "timeout") +#: The reference a spec that names none gets, when the file is beside the spec. +DEFAULT_LOGIT_REFERENCE = "logits.reference.json.gz" +#: The block joins the accuracy keys here, beside its validation. +ACCURACY_KEYS = (*ACCURACY_KEYS, "logits") + + +def _validate_logits(block: Any, errors: list[str]) -> None: + """Validate ``accuracy.logits``.""" + where = f"{ACCURACY_FIELD}.logits" + if not isinstance(block, dict): + errors.append(f"'{where}' must be a mapping") + return + reject_unknown_keys(block, LOGITS_KEYS, errors, where=where) + enabled = block.get("enabled") + if enabled is not None and not isinstance(enabled, bool): + errors.append(f"'{where}.enabled' must be true or false") + reference = block.get("reference") + if reference is not None and (not isinstance(reference, str) or not reference.strip()): + errors.append(f"'{where}.reference' must be a path") + elif reference is None and enabled is not False: + errors.append( + f"'{where}.reference' names the reference file " + f"(onboard/logit_reference.py writes it); without one set enabled: false" + ) + elif reference is not None and enabled is not False and not Path(reference).is_file(): + errors.append(f"'{where}.reference' does not exist: {reference}") + prompts = block.get("prompts") + if prompts is not None and ( + not isinstance(prompts, list) + or not prompts + or not all(isinstance(p, str) and p.strip() for p in prompts) + or len(set(prompts)) != len(prompts) + ): + errors.append(f"'{where}.prompts' must be a non-empty list of distinct prompt ids") + factor = block.get("threshold_factor") + if factor is not None and (not _number(factor) or not factor > 0): + errors.append(f"'{where}.threshold_factor' must be a number > 0") + if block.get("timeout") is not None and not _positive_int(block["timeout"]): + errors.append(f"'{where}.timeout' must be an integer >= 1") + + +# --------------------------------------------------------- end accuracy.logits + + +def _bounds_server_length(accuracy: Mapping[str, Any]) -> bool: + """Whether the spec says what sequence length ``trtllm-serve`` sizes its cache for.""" + if accuracy.get("max_seq_len") is not None: + return True + options = accuracy.get("llm_api_options") + if isinstance(options, dict): + options = options.get("default") + if not isinstance(options, str) or not Path(options).is_file(): + return False + try: + loaded = yaml.safe_load(Path(options).read_text(encoding="utf-8")) + except (OSError, yaml.YAMLError): + return False + return isinstance(loaded, dict) and loaded.get("max_seq_len") is not None + + +def derive_gate_defaults(data: dict[str, Any], spec_dir: Path) -> None: + """Turn on, in place, the gates the spec leaves unsaid and the run can derive. + + - ``accuracy.logits``: a block that names no reference and does not turn + the gate off gets :data:`DEFAULT_LOGIT_REFERENCE` in ``spec_dir`` when + that file is there; + - ``accuracy.serve``: a block without ``enabled`` gets it, true wherever + the spec bounds the server's sequence length (:func:`_bounds_server_length`). + + What it chose is written into the resolved spec, so the workspace copy + says which gates the contract owes. A spec that says ``enabled: false`` + keeps either off. + """ + accuracy = data.get(ACCURACY_FIELD) + if not isinstance(accuracy, dict): + return + logits = accuracy.get("logits") + default = spec_dir / DEFAULT_LOGIT_REFERENCE + unnamed = isinstance(logits, dict) and logits.get("reference") is None + if (logits is None or (unnamed and logits.get("enabled") is not False)) and default.is_file(): + accuracy["logits"] = {**(logits or {}), "reference": str(default.resolve())} + serve = accuracy.get("serve") + if serve is None or (isinstance(serve, dict) and "enabled" not in serve): + accuracy["serve"] = {**(serve or {}), "enabled": _bounds_server_length(accuracy)} + + +def _validate_accuracy(block: Any, errors: list[str]) -> None: + """Validate the ``accuracy`` block in place, appending to ``errors``.""" + if not isinstance(block, dict): + raise TaskSchemaError(f"'{ACCURACY_FIELD}' must be a mapping, got {type(block).__name__}") + + reject_unknown_keys(block, ACCURACY_KEYS, errors, where=ACCURACY_FIELD) + + # Both of these accept the per-unit mapping form, for the same reason + # `llm_api_options` does. The `multimodal` increment's release criterion + # is a multimodal benchmark in a different reference file; gating it on + # the text benchmark core was gated on is not a weaker check, it is a + # check that cannot see the vision tower and passes whatever it does. + for name in ("benchmark", "reference_yaml"): + value = block.get(name) + if isinstance(value, str): + if not value.strip(): + errors.append(f"'{ACCURACY_FIELD}.{name}' must be a non-empty string") + continue + if isinstance(value, dict): + if "default" not in value: + errors.append( + f"'{ACCURACY_FIELD}.{name}' is a per-unit mapping, so it needs a " + f"'default' key -- otherwise a unit nobody listed silently has no " + f"{name} and the gate falls back to a hard-coded one" + ) + for unit, entry in value.items(): + if not isinstance(entry, str) or not entry.strip(): + errors.append(f"'{ACCURACY_FIELD}.{name}[{unit}]' must be a non-empty string") + continue + errors.append( + f"'{ACCURACY_FIELD}.{name}' must be a non-empty string, or a mapping " + f"from assembly unit to one" + ) + + tolerance = block.get("tolerance") + if tolerance is None: + block["tolerance"] = DEFAULT_ACCURACY_TOLERANCE + elif not isinstance(tolerance, (int, float)) or isinstance(tolerance, bool): + errors.append(f"'{ACCURACY_FIELD}.tolerance' must be a number") + elif tolerance < 0: + errors.append(f"'{ACCURACY_FIELD}.tolerance' must be >= 0; got {tolerance}") + + # ``dataset_path`` is optional but load-bearing for MMLU: trtllm-eval + # fetches that one dataset itself from a URL that now 404s, writes the + # error page to data.tar, and dies inside tarfile.open — so the failure + # reads as a corrupt archive rather than a dead link. Every other + # benchmark comes through lm-eval and is unaffected. + dataset_path = block.get("dataset_path") + if dataset_path is not None and (not isinstance(dataset_path, str) or not dataset_path.strip()): + errors.append(f"'{ACCURACY_FIELD}.dataset_path' must be a non-empty string when present") + + # `trtllm-eval` takes none of the quickstart runtime flags, so the + # configuration the gate measures under has to arrive as a YAML. It must + # exist: a missing file means the gate silently measures the engine's + # defaults, which is how a run compares a number against an anchor taken + # under a different configuration and calls it a regression. + options = block.get("llm_api_options") + # A mapping unit -> path, because a feature unit is not measured under + # the same configuration as core: MTP needs a `speculative_config`, and + # without one the accuracy gate re-measures the base path while the + # acceptance gate compares 1.0 against 1.0. + if isinstance(options, dict): + for unit, path in options.items(): + if not isinstance(path, str) or not Path(path).is_file(): + errors.append(f"'{ACCURACY_FIELD}.llm_api_options[{unit}]' does not exist: {path}") + options = None + if options is not None: + if not isinstance(options, str) or not options.strip(): + errors.append(f"'{ACCURACY_FIELD}.llm_api_options' must be a non-empty path") + elif not Path(options).is_file(): + errors.append( + f"'{ACCURACY_FIELD}.llm_api_options' does not exist: {options}. Without it " + f"the gate measures the engine's defaults, not the configuration the " + f"anchor was measured under." + ) + + for name in ("max_seq_len", "max_batch_size"): + value = block.get(name) + if value is None: + continue + if not isinstance(value, int) or isinstance(value, bool) or value < 1: + errors.append(f"'{ACCURACY_FIELD}.{name}' must be an integer >= 1") + + if block.get("acceptance") is not None: + _validate_acceptance(block["acceptance"], errors) + + _validate_floor(block, errors) + + if block.get("boot") is not None: + _validate_boot(block["boot"], errors) + + if block.get("envelope") is not None: + _validate_envelope(block["envelope"], errors) + + if block.get("logits") is not None: + _validate_logits(block["logits"], errors) + + if block.get("serve") is not None: + _validate_serve(block["serve"], errors) + + anchor = block.get("anchor") + for where, state in ( + anchor.items() if isinstance(anchor, dict) else [("", anchor)] if anchor is not None else [] + ): + if state not in ANCHOR_STATES: + errors.append( + f"'{ACCURACY_FIELD}.anchor{'.' + str(where) if where else ''}' must be one of " + f"{list(ANCHOR_STATES)}; got {state!r}" + ) + + members = block.get("benchmark_members") + if members is not None and members not in BENCHMARK_MEMBERS: + errors.append( + f"'{ACCURACY_FIELD}.benchmark_members' must be one of {list(BENCHMARK_MEMBERS)}; " + f"got {members!r}" + ) + + witness = block.get("witness") + if witness is not None: + if not isinstance(witness, dict): + errors.append(f"'{ACCURACY_FIELD}.witness' must be a mapping from unit to a block") + else: + for unit, entry in witness.items(): + where = f"{ACCURACY_FIELD}.witness.{unit}" + if not isinstance(entry, dict): + errors.append(f"'{where}' must be a mapping") + continue + # Closed per feature: `window` under `host_offload` is a key + # nothing reads, and `expected_mib` under a replay feature reads + # as a check that is not being made. + if unit in REPLAY_WITNESS_REQUIRED: + _validate_replay_witness(str(unit), entry, errors, where) + continue + if unit == "kv_quant": + _validate_kv_quant_witness(entry, errors, where) + continue + if unit in EQUIVALENCE_WITNESS_FEATURES: + _validate_equivalence_witness(entry, errors, where) + continue + reject_unknown_keys(entry, WITNESS_KEYS, errors, where=where) + _validate_witness_timeout(entry, errors, where) + _validate_witness_batch_size(entry, errors, where) + _validate_witness_workload(entry, errors, where) + command = entry.get("contrast_command") + if command is not None and (not isinstance(command, str) or not command.strip()): + errors.append(f"'{where}.contrast_command' must be a non-empty string") + _validate_contrast_env(entry, errors, where) + # A scalar, or per member with a `default`: the witness runs + # at every member, and the delta can differ between them. + predicted = entry.get("expected_mib") + values = predicted.values() if isinstance(predicted, dict) else [predicted] + if predicted is not None and ( + not all(_number(v) and v > 0 for v in values) + or (isinstance(predicted, dict) and not predicted) + ): + errors.append( + f"'{where}.expected_mib' must be a positive number, or a mapping from " + f"member (and `default`) to one" + ) + + # A mapping of scalars (the whole run) or of unit -> mapping of scalars. + # Validated rather than trusted because a discriminator the reference file + # does not use produces `no entry matching {...}` at the END of an + # assembly, which is the most expensive place to learn about a typo. + discriminators = block.get("discriminators") + if discriminators is not None: + if not isinstance(discriminators, dict): + errors.append(f"'{ACCURACY_FIELD}.discriminators' must be a mapping") + else: + nested = [v for v in discriminators.values() if isinstance(v, dict)] + if nested and len(nested) != len(discriminators): + errors.append( + f"'{ACCURACY_FIELD}.discriminators' mixes the flat form " + f"(quant_algo: ...) with the per-unit form (core: {{...}}); pick one" + ) + for group in nested or [discriminators]: + for key, value in group.items(): + if not isinstance(value, (str, int, float, bool)): + errors.append( + f"'{ACCURACY_FIELD}.discriminators.{key}' must be a scalar " + f"matching a reference-file discriminator, got " + f"{type(value).__name__}" + ) + + # The evaluation *protocol* -- chat template, few-shot handling, seed -- + # is expressed as options of the benchmark subcommand, so it cannot live + # in the llm_api_options YAML. It has to be declared, because a protocol + # mismatch between the anchor and the gate is indistinguishable from an + # accuracy regression, and can be tens of points wide. + benchmark_args = block.get("benchmark_args") + + def _is_arg_list(value: Any) -> bool: + return isinstance(value, list) and all(isinstance(a, str) for a in value) + + if benchmark_args is None: + block["benchmark_args"] = [] + elif isinstance(benchmark_args, dict): + # Per unit for a harder reason than the other fields: these are + # options of the BENCHMARK subcommand, and the subcommands do not + # share an option set. `--apply_chat_template` is a gsm8k option and + # not an mmmu one, so a run-wide list makes the multimodal unit's + # gate fail to start rather than fail to measure. + if "default" not in benchmark_args: + errors.append( + f"'{ACCURACY_FIELD}.benchmark_args' is a per-unit mapping, so it needs " + f"a 'default' key" + ) + for unit, value in benchmark_args.items(): + if not _is_arg_list(value): + errors.append( + f"'{ACCURACY_FIELD}.benchmark_args[{unit}]' must be a list of strings" + ) + elif not _is_arg_list(benchmark_args): + errors.append( + f"'{ACCURACY_FIELD}.benchmark_args' must be a list of strings, or a " + f"mapping from assembly unit to one" + ) + + +def _require_declared_witnesses(data: Mapping[str, Any], errors: list[str]) -> None: + """A feature whose witness judges against a DECLARED value needs it here. + + Here and not at the unit boundary, where `host_offload` and `disagg` are + checked: these witnesses compare the target's own counters against a + workload size or a byte figure, and the one place such a declaration is + certainly made before the run it judges is this one. Found missing at the + gate, it would be written after the assembler's turns -- and chosen, + however honestly, by someone who has seen the target's numbers. + """ + features = data.get(FEATURES_FIELD) + accuracy = data.get(ACCURACY_FIELD) + if not isinstance(features, list) or not isinstance(accuracy, dict): + return + witness = accuracy.get("witness") + witness = witness if isinstance(witness, dict) else {} + for feature in features: + if feature not in DECLARED_WITNESS_FEATURES or feature in witness: + continue + keys = REPLAY_WITNESS_REQUIRED.get(feature) or ("expected_bytes_per_token", "exact") + errors.append( + f"feature {feature!r} needs 'accuracy.witness.{feature}' ({', '.join(keys)}): " + f"a unit that implements nothing passes both boot and accuracy, so the " + f"witness is the only gate that can tell, and what it judges against has to " + f"be declared before the run." + ) + + +def _beside_spec(command: str, spec_dir: Path) -> str: + """``command`` with each word naming a file in ``spec_dir`` made absolute; else unchanged.""" + try: + words = shlex.split(command) + except ValueError: + return command + moved = [ + str((spec_dir / word).resolve()) + if word and not Path(word).is_absolute() and (spec_dir / word).is_file() + else word + for word in words + ] + return shlex.join(moved) if moved != words else command + + +def _resolve_spec_relative_paths(data: dict[str, Any], spec_dir: Path) -> None: + """Make paths the user writes NEXT TO the spec portable, in place. + + ``accuracy.llm_api_options``, ``accuracy.witness.*.prompts_file`` and + ``accuracy.acceptance[.].prompts_file`` point at files the author + writes alongside ``task.yaml`` — the per-unit runtime YAMLs, prompt + sets. So may a witness's ``contrast_command``: each word of it naming a + file the spec's directory holds is made absolute, and any other word -- + an interpreter, a flag, a path of the tree under test, which the command + reaches from inside it -- is left as written. Written absolute, **the spec stops working on any other machine**, + which is the opposite of what a spec is for. + + Resolved against the **spec's own directory**, because that is where the + author put them, and written back as absolute so no consumer downstream + needs to know where the spec lives — the gates take a string and pass it + to ``trtllm-eval``. + + Absolute paths are untouched: a site that wants one keeps it, and a + workspace copy of an already-resolved spec re-resolves to itself, which is + the case that matters because the workflow copies ``task.yaml`` into the + workspace and reads it again. + + Deliberately NOT applied to ``checkpoint_path``, ``trtllm_repo_path`` or + ``environment.spool``. Those genuinely differ per machine and are read by + processes that never see the spec's directory; making them relative would + move the portability problem rather than solve it. + """ + + def absolute(value: Any) -> Any: + if not isinstance(value, str) or not value.strip(): + return value + candidate = Path(value) + return value if candidate.is_absolute() else str((spec_dir / candidate).resolve()) + + accuracy = data.get(ACCURACY_FIELD) + if not isinstance(accuracy, dict): + return + + options = accuracy.get("llm_api_options") + if isinstance(options, dict): + accuracy["llm_api_options"] = {k: absolute(v) for k, v in options.items()} + elif options is not None: + accuracy["llm_api_options"] = absolute(options) + + witness = accuracy.get("witness") + if isinstance(witness, dict): + for block in witness.values(): + if isinstance(block, dict) and block.get("prompts_file") is not None: + block["prompts_file"] = absolute(block["prompts_file"]) + if isinstance(block, dict) and isinstance(block.get("contrast_command"), str): + block["contrast_command"] = _beside_spec(block["contrast_command"], spec_dir) + + # The acceptance workload's prompt file, at the block level and per unit. + # A `dataset` may be a cached dataset id ("openai/gsm8k"), which reads like + # a relative path, so it is resolved only when the spec's directory holds it. + acceptance = accuracy.get("acceptance") + if isinstance(acceptance, dict): + for level in (acceptance, *(v for v in acceptance.values() if isinstance(v, dict))): + if level.get("prompts_file") is not None: + level["prompts_file"] = absolute(level["prompts_file"]) + dataset = level.get("dataset") + if isinstance(dataset, str) and dataset.strip() and not Path(dataset).is_absolute(): + if (spec_dir / dataset).exists(): + level["dataset"] = absolute(dataset) + + # Boot prompt files, and boot images that sit beside the spec. An image + # that does not is left relative: the gate resolves it against the tree + # under test, where the repository's own test assets are. + boot = accuracy.get("boot") + if isinstance(boot, dict): + for level in (boot, *(v for v in boot.values() if isinstance(v, dict))): + if level.get("prompts_file") is not None: + level["prompts_file"] = absolute(level["prompts_file"]) + for entry in level.get("prompts") or (): + if isinstance(entry, dict) and isinstance(entry.get("images"), list): + entry["images"] = [ + absolute(i) + if isinstance(i, str) + and not Path(i).is_absolute() + and (spec_dir / i).exists() + else i + for i in entry["images"] + ] + + # The logit gate's reference file, written beside the spec. + logits = accuracy.get("logits") + if isinstance(logits, dict) and logits.get("reference") is not None: + logits["reference"] = absolute(logits["reference"]) + + +def load_and_validate_task_yaml(path: str | Path) -> dict[str, Any]: + """Parse ``path`` as YAML, validate it, and fill in the workflow's defaults. + + Returns the resolved mapping. Raises :class:`TaskSchemaError` with + **every** detected problem batched into a single message. + """ + task_path = Path(path) + if not task_path.is_file(): + raise TaskSchemaError(f"task file not found: {task_path}") + + text = task_path.read_text(encoding="utf-8") + try: + data = yaml.safe_load(text) + except yaml.YAMLError as exc: + raise TaskSchemaError(f"{task_path} is not valid YAML: {exc}") from exc + + if data is None: + data = {} + if not isinstance(data, dict): + raise TaskSchemaError( + f"{task_path} must be a YAML mapping at the top level, got {type(data).__name__}" + ) + + # BEFORE validation, deliberately. `accuracy.llm_api_options` has an + # existence check, and that check is a *consumer* of the resolved path: + # resolving afterwards makes a perfectly good relative path fail as + # "does not exist". Resolving first also means the error message names + # the absolute path the run would really have used. + _resolve_spec_relative_paths(data, task_path.resolve().parent) + # After the paths, before validation: what it derives is validated like + # what the spec says. + derive_gate_defaults(data, task_path.resolve().parent) + + errors: list[str] = [] + + validate_commit_policy_block(data.get(COMMIT_POLICY_FIELD), errors) + + errors.extend(validate_perf_block(data)) + # Before the required fields are read: a key this schema does not know + # is an error, not a no-op. Placed here so the message a user gets names + # their typo rather than the consequence of it — `feature:` silently + # resolving to no features reads downstream as "the run built core and + # said it was done", which is a much longer trip back to the cause. + reject_unknown_keys(data, TOP_LEVEL_KEYS, errors) + errors.extend(validate_agents(data, ROLES)) + + for name in REQUIRED_PATH_FIELDS: + if name not in data: + errors.append(f"missing required field '{name}'") + continue + value = data[name] + if not isinstance(value, str) or not value.strip(): + errors.append(f"'{name}' must be a non-empty string, got {type(value).__name__}") + continue + if not Path(value).exists(): + errors.append(f"'{name}' points to a non-existent path: {value}") + + arch = data.get(ARCH_FIELD) + if arch is None: + errors.append( + f"missing required field '{ARCH_FIELD}'. It is never inferred: catalog " + f"receipts are per-architecture, so the wrong value gates this target " + f"against another machine's certification." + ) + elif not isinstance(arch, str) or not ARCH_PATTERN.match(arch): + errors.append(f"'{ARCH_FIELD}' must look like 'sm_100' / 'sm_103'; got {arch!r}") + + parallel = data.get(PARALLEL_FIELD) + if parallel is None: + data[PARALLEL_FIELD] = DEFAULT_PARALLEL + elif not isinstance(parallel, str): + errors.append(f"'{PARALLEL_FIELD}' must be a string, got {type(parallel).__name__}") + else: + try: + parse_parallel_segment(parallel) + except TaskSchemaError as exc: + errors.append(str(exc)) + + # After `parallel` has been defaulted, since legality depends on it. + if MEMBERS_FIELD in data: + environment = data.get(ENVIRONMENT_FIELD) + kind = environment.get("kind", "local") if isinstance(environment, dict) else "local" + _validate_members( + data[MEMBERS_FIELD], data.get(PARALLEL_FIELD), errors, environment_kind=str(kind) + ) + + features = data.get(FEATURES_FIELD) + if features is None: + data[FEATURES_FIELD] = [] + elif not isinstance(features, list): + errors.append( + f"'{FEATURES_FIELD}' must be a list of strings, got {type(features).__name__}" + ) + else: + for index, feature in enumerate(features): + if not isinstance(feature, str): + errors.append( + f"'{FEATURES_FIELD}[{index}]' must be a string, got {type(feature).__name__}" + ) + elif feature not in FEATURE_SKILLS: + errors.append( + f"'{FEATURES_FIELD}[{index}]' is {feature!r}, which no assembler skill " + f"covers. Known features: {sorted(FEATURE_SKILLS)}" + ) + readers = [f for f in MTP_NAMESPACE_FEATURES if f in features] + if len(readers) > 1: + errors.append( + f"'{FEATURES_FIELD}' lists {readers}, and all of them read the checkpoint's " + f"`mtp.*` weights, which are MTP modules or DSpark stages, never both. " + f"Keep the one the checkpoint ships. An `mtp.*` that holds DSpark stages " + f"(`dspark_*` keys in its config) is spelled `dspark`, even when the " + f"request says 'MTP'." + ) + + if isinstance(data.get(FEATURES_FIELD), list): + validate_variants( + data.get(VARIANTS_FIELD), + units=assembly_units(data), + speculators={f for f, s in FEATURE_SKILLS.items() if s.endswith("-speculator")}, + unified=data.get(PARALLEL_FIELD) == PARALLEL_UNIFIED, + errors=errors, + ) + + _validate_coverage(data.get(COVERAGE_FIELD), errors) + if data.get(KNOWLEDGE_FIELD) is not None and data[KNOWLEDGE_FIELD] not in KNOWLEDGE_MODES: + errors.append(f"'{KNOWLEDGE_FIELD}' must be one of {list(KNOWLEDGE_MODES)}") + + if ENVIRONMENT_FIELD not in data: + errors.append( + f"missing required field '{ENVIRONMENT_FIELD}'. It declares how GPU work is " + f"launched; without it the agents would have to discover that themselves, " + f"which is exactly what this workflow refuses to do." + ) + else: + try: + _validate_environment(data[ENVIRONMENT_FIELD], errors) + except TaskSchemaError as exc: + errors.append(str(exc)) + + if TARGET_FIELD not in data: + data[TARGET_FIELD] = {} + else: + try: + _validate_target(data[TARGET_FIELD], errors) + except TaskSchemaError as exc: + errors.append(str(exc)) + + if RUNTIME_FIELD not in data: + data[RUNTIME_FIELD] = dict(RUNTIME_SWITCHES) | {"extra_flags": []} + else: + try: + _validate_runtime(data[RUNTIME_FIELD], errors) + except TaskSchemaError as exc: + errors.append(str(exc)) + + notes = data.get(NOTES_FIELD) + if notes is not None and (not isinstance(notes, str) or not notes.strip()): + errors.append(f"'{NOTES_FIELD}' must be a non-empty string") + + scope = data.get(SCOPE_FIELD) + if scope is None: + data[SCOPE_FIELD] = DEFAULT_SCOPE + elif scope not in SCOPES: + errors.append(f"'{SCOPE_FIELD}' must be one of {SCOPES}; got {scope!r}") + + if SOURCE_EMBARGO_FIELD not in data: + data[SOURCE_EMBARGO_FIELD] = [] + else: + try: + _validate_source_embargo( + data[SOURCE_EMBARGO_FIELD], errors, data.get("trtllm_repo_path") + ) + except TaskSchemaError as exc: + errors.append(str(exc)) + + if ACCURACY_FIELD not in data: + errors.append( + f"missing required field '{ACCURACY_FIELD}'. The accuracy gate is the release " + f"criterion, and its anchor comes from a written source — never from the run." + ) + else: + try: + _validate_accuracy(data[ACCURACY_FIELD], errors) + except TaskSchemaError as exc: + errors.append(str(exc)) + _require_declared_witnesses(data, errors) + + if errors: + bullet = "\n - " + raise TaskSchemaError( + f"{task_path} failed trtllm-modeling-v2 schema validation:{bullet}{bullet.join(errors)}" + ) + + return data + + +def features_through(task_data: Mapping[str, Any], unit: str) -> tuple[str, ...]: + """The features a unit's target implements: every feature up to it.""" + features = [str(f) for f in task_data.get(FEATURES_FIELD) or ()] + if unit not in features: + return () + return tuple(features[: features.index(unit) + 1]) + + +def assembly_units(data: Mapping[str, Any]) -> list[str]: + """Return the assembly units in dispatch order: ``core`` then each feature. + + Core is always first and always green before the first feature: the + accuracy gate's tolerance is wide enough that a core defect and a + feature defect are indistinguishable when they land together. It also + means a first assembly with ``features`` follows exactly the same path + as a core assembly plus later increments, so there is only one path to + get right. + """ + return ["core", *[str(f) for f in data.get(FEATURES_FIELD, [])]] + + +__all__ = [ + "ANCHOR_PENDING", + "ANCHOR_RECORDED", + "ANCHOR_STATES", + "anchor_pending", + "COVERAGE_FIELD", + "KNOWLEDGE_FIELD", + "KNOWLEDGE_MEASUREMENT", + "KNOWLEDGE_MODES", + "KNOWLEDGE_PRODUCTION", + "COVERAGE_KEYS", + "resolve_per_unit", + "resolve_samples", + "BENCHMARK_MEMBERS", + "BENCHMARK_MEMBERS_EVERY", + "BENCHMARK_MEMBERS_SAMPLED", + "DEFAULT_LOGIT_REFERENCE", + "GREEDY_KEYS", + "derive_gate_defaults", + "ACCEPTANCE_UNIT_KEYS", + "ACCEPTANCE_WORKLOAD_KEYS", + "ACCURACY_FIELD", + "ACCURACY_KEYS", + "ARCH_FIELD", + "ARCH_PATTERN", + "DEFAULT_ACCURACY_TOLERANCE", + "DEFAULT_MEMBER_KINDS", + "DEFAULT_PARALLEL", + "DEFAULT_SCOPE", + "ENVIRONMENT_FIELD", + "ENVIRONMENT_KEYS", + "ENVIRONMENT_KINDS", + "DECLARED_WITNESS_FEATURES", + "DEFAULT_BYTES_TOLERANCE", + "FEATURES_FIELD", + "FEATURE_SKILLS", + "KV_QUANT_WITNESS_KEYS", + "LOGITS_KEYS", + "MAX_WITNESS_PROMPT_WORDS", + "MEMBERS_FIELD", + "MTP_NAMESPACE_FEATURES", + "REPLAY_WITNESS_OPTIONAL", + "REPLAY_WITNESS_REQUIRED", + "EQUIVALENCE_WITNESS_FEATURES", + "EQUIVALENCE_WITNESS_KEYS", + "WITNESS_BATCH_SIZE_KEY", + "WITNESS_KEYS", + "WITNESS_LEGS_ENV_KEY", + "WITNESS_TRANSCEIVER_KEY", + "PARALLEL_FIELD", + "PARALLEL_UNIFIED", + "REQUIRED_PATH_FIELDS", + "RUNTIME_FIELD", + "RUNTIME_MARGIN_DEFAULTS", + "RUNTIME_SWITCHES", + "SCOPES", + "SCOPE_FIELD", + "SERVE_KEYS", + "reject_unknown_keys", + "SOURCE_EMBARGO_FIELD", + "TARGET_FIELD", + "TARGET_KEYS", + "TOP_LEVEL_KEYS", + "knowledge_mode", + "TaskSchemaError", + "assembly_units", + "features_through", + "derive_checkpoint_segment", + "derive_family", + "derive_slug", + "load_and_validate_task_yaml", + "load_prompts_file", + "members_for", + "parse_parallel_segment", + "widest_world_size", +] diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/variants.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/variants.py new file mode 100644 index 000000000000..1a992558feb8 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/spec/variants.py @@ -0,0 +1,324 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Configuration variants: a unit's target, served under another configuration. + +A variant is one configuration of an already-built unit that the target must +also serve: a KV-cache or SSM-state storage type, a speculation depth. It +builds nothing of its own. It adds checks to the contract, and those checks +run with the unit's spec overlaid with the variant's switches:: + + variants: + kv_fp8: + unit: core # whose target and settings it starts from + members: [tep4] # default: the unit's members + gates: [boot, accuracy] # the default; `acceptance` needs a speculator + runtime: {extra_flags: [--kv_cache_dtype=fp8]} # boot: appended to the unit's + llm_api_options: {kv_cache_config: {dtype: fp8}} # merged into the unit's YAML + env: {} # added to the gate command's environment + accuracy: {discriminators: {}} # accuracy keys that differ from the unit's + note: KV cache stored as FP8 + +Each check is named by the variant (``kv_fp8:accuracy[gsm8k]@tep4``), so its +verdicts, per-prompt baselines and findings are its own, never the unit's. +:func:`variant_spec` is the one place the overlay is applied; the +orchestrator's checks and ``trtllm-modeling-v2 gate --unit `` both +go through it. +""" + +from __future__ import annotations + +import copy +import re +from pathlib import Path +from typing import Any, Mapping + +import yaml + +from .schema_keys import reject_unknown_keys + +VARIANTS_FIELD = "variants" +VARIANT_KEYS: tuple[str, ...] = ( + "unit", + "members", + "gates", + "runtime", + "llm_api_options", + "env", + "accuracy", + "note", +) +VARIANT_GATES: tuple[str, ...] = ("boot", "accuracy", "acceptance") +DEFAULT_VARIANT_GATES: tuple[str, ...] = ("boot", "accuracy") +#: Accuracy keys a variant may set for itself; every other key is the unit's. +VARIANT_ACCURACY_KEYS: tuple[str, ...] = ( + "benchmark", + "reference_yaml", + "model_key", + "discriminators", + "benchmark_args", + "samples", + "acceptance", +) +RUNTIME_OVERLAY_KEYS: tuple[str, ...] = ( + "cuda_graph", + "chunked_prefill", + "overlap_scheduler", + "kv_cache_reuse", + "extra_flags", + "min_free_memory_gib", + "min_free_memory_fraction", +) +#: Where a variant's merged LLM-API options are written, under the workspace. +OPTIONS_DIRNAME = "variants" +_NAME = re.compile(r"^[a-z][a-z0-9_]{0,40}$") + + +def variants(data: Mapping[str, Any]) -> dict[str, dict[str, Any]]: + """The declared variants, by name, in declaration order.""" + block = data.get(VARIANTS_FIELD) or {} + return {str(k): dict(v or {}) for k, v in block.items()} if isinstance(block, Mapping) else {} + + +def is_variant(data: Mapping[str, Any], name: str | None) -> bool: + return bool(name) and str(name) in variants(data) + + +def base_unit(data: Mapping[str, Any], name: str) -> str: + """The unit whose target a variant configures.""" + return str(variants(data)[name].get("unit") or "core") + + +def variant_gates(data: Mapping[str, Any], name: str) -> tuple[str, ...]: + return tuple(variants(data)[name].get("gates") or DEFAULT_VARIANT_GATES) + + +def validate_variants( + block: Any, + *, + units: list[str], + speculators: set[str], + unified: bool, + errors: list[str], +) -> None: + """Every problem with the ``variants`` block, batched into ``errors``.""" + if block is None: + return + if not isinstance(block, Mapping): + errors.append(f"'{VARIANTS_FIELD}' must be a mapping from variant name to its entry") + return + for name, entry in block.items(): + where = f"{VARIANTS_FIELD}.{name}" + if not isinstance(name, str) or not _NAME.match(name): + errors.append(f"'{where}': a variant name is lowercase letters, digits and '_'") + continue + if name in units or name == "default": + errors.append(f"'{where}': {name!r} is a unit name; a variant needs its own") + continue + if not isinstance(entry, Mapping): + errors.append(f"'{where}' must be a mapping") + continue + reject_unknown_keys(entry, VARIANT_KEYS, errors, where=where) + unit = str(entry.get("unit") or "core") + if unit not in units: + errors.append(f"'{where}.unit' is {unit!r}, not one of the units {units}") + continue + gates = entry.get("gates") or list(DEFAULT_VARIANT_GATES) + if not isinstance(gates, list) or not gates or any(g not in VARIANT_GATES for g in gates): + errors.append( + f"'{where}.gates' must be a non-empty list drawn from {list(VARIANT_GATES)}" + ) + elif "acceptance" in gates: + built = set(units[1 : units.index(unit) + 1]) + if not built & speculators: + errors.append( + f"'{where}.gates' asks for acceptance, but unit {unit!r} implements no " + f"speculator ({sorted(speculators)})" + ) + members = entry.get("members") + if members is not None: + if not unified: + errors.append(f"'{where}.members' is only legal with 'parallel: unified'") + elif not isinstance(members, list) or not members: + errors.append(f"'{where}.members' must be a non-empty list of parallel segments") + else: + from .task_schema import TaskSchemaError, parse_parallel_segment + + for member in members: + try: + _, size = parse_parallel_segment(str(member)) + except TaskSchemaError as exc: + errors.append(f"'{where}.members': {exc}") + continue + if size is None: + errors.append(f"'{where}.members': {member!r} is not a concrete member") + runtime = entry.get("runtime") + if runtime is not None: + if not isinstance(runtime, Mapping): + errors.append(f"'{where}.runtime' must be a mapping") + else: + reject_unknown_keys(runtime, RUNTIME_OVERLAY_KEYS, errors, where=f"{where}.runtime") + flags = runtime.get("extra_flags") + if flags is not None and not ( + isinstance(flags, list) and all(isinstance(f, str) for f in flags) + ): + errors.append(f"'{where}.runtime.extra_flags' must be a list of strings") + for key in ("llm_api_options", "env", "accuracy"): + if entry.get(key) is not None and not isinstance(entry[key], Mapping): + errors.append(f"'{where}.{key}' must be a mapping") + if isinstance(entry.get("accuracy"), Mapping): + reject_unknown_keys( + entry["accuracy"], VARIANT_ACCURACY_KEYS, errors, where=f"{where}.accuracy" + ) + if not (entry.get("runtime") or entry.get("llm_api_options") or entry.get("env")): + errors.append( + f"'{where}' changes nothing: give it `runtime`, `llm_api_options` or `env`" + ) + + +def _deep_merge(base: Mapping[str, Any], overlay: Mapping[str, Any]) -> dict[str, Any]: + merged = dict(base) + for key, value in overlay.items(): + if isinstance(value, Mapping) and isinstance(merged.get(key), Mapping): + merged[key] = _deep_merge(merged[key], value) + else: + merged[key] = copy.deepcopy(value) + return merged + + +def _per_unit(key: str, value: Any, name: str, chosen: Any) -> dict[str, Any]: + """Accuracy key ``key`` in its per-unit form, with ``name`` set to ``chosen``. + + ``acceptance`` holds the workload and the per-unit floors side by side; + ``discriminators`` is per-unit when its values are mappings; every other + key is per-unit exactly when it is a mapping. + """ + if key == "acceptance": + return {**dict(value or {}), name: chosen} + if key == "discriminators": + per_unit = isinstance(value, Mapping) and any( + isinstance(v, Mapping) for v in value.values() + ) + else: + per_unit = isinstance(value, Mapping) + out = dict(value) if per_unit else ({} if value is None else {"default": value}) + out[name] = chosen + return out + + +def variant_anchor_pending(data: Mapping[str, Any], spec: Mapping[str, Any], name: str) -> bool: + """Whether variant ``name`` has no written anchor (``accuracy.anchor``). + + Its own say where the anchor mapping names it, else its base unit's: a + variant of a unit whose anchor is pending has none of its own either. + """ + # task_schema validates the variants with this module, so it is read here, when called. + from .task_schema import ANCHOR_PENDING, anchor_pending + + own = (spec.get("accuracy") or {}).get("anchor") + if isinstance(own, Mapping) and name in own: + return own[name] == ANCHOR_PENDING + return anchor_pending(data.get("accuracy") or {}, base_unit(data, name)) + + +def variant_spec( + data: Mapping[str, Any], + name: str, + *, + world_size: int, + options_dir: Path | None = None, +) -> dict[str, Any]: + """The spec a variant's checks run under: the unit's, with the variant's overlay. + + Every per-unit lookup the gates make (``resolve_per_unit``, the runtime + block, ``members``) then answers for the variant's name what it answers + for its unit, except where the variant says otherwise. With + ``options_dir`` the unit's LLM-API YAML merged with the variant's options + is written there; without it (compiling the contract) the options are not + resolved. + """ + from ..checks.gates import resolve_llm_api_options, resolve_runtime, write_atomically + from .task_schema import members_for + + entry = variants(data)[name] + unit = base_unit(data, name) + out = copy.deepcopy(dict(data)) + accuracy = dict(out.get("accuracy") or {}) + for key, value in list(accuracy.items()): + if isinstance(value, dict) and unit in value: + value[name] = copy.deepcopy(value[unit]) + + if options_dir is not None: + base_options = resolve_llm_api_options(accuracy, unit) + loaded: dict[str, Any] = {} + if base_options and Path(base_options).is_file(): + loaded = yaml.safe_load(Path(base_options).read_text(encoding="utf-8")) or {} + merged = _deep_merge(loaded, entry.get("llm_api_options") or {}) + options_dir.mkdir(parents=True, exist_ok=True) + path = options_dir / f"{name}.llm_api.yaml" + # Replaced whole: a gate of another variant may be reading it now. + write_atomically(path, yaml.safe_dump(merged, sort_keys=False)) + accuracy["llm_api_options"] = _per_unit( + "llm_api_options", accuracy.get("llm_api_options"), name, str(path) + ) + for key, value in (entry.get("accuracy") or {}).items(): + accuracy[key] = _per_unit(key, accuracy.get(key), name, copy.deepcopy(value)) + out["accuracy"] = accuracy + + runtime = dict(resolve_runtime(data, unit)) + overlay = dict(entry.get("runtime") or {}) + flags = [*runtime.get("extra_flags", ()), *overlay.pop("extra_flags", ())] + runtime.update(overlay) + runtime["extra_flags"] = flags + block = data.get("runtime") or {} + if any(isinstance(v, dict) for v in block.values()): + out["runtime"] = {**copy.deepcopy(dict(block)), name: runtime} + else: + out["runtime"] = {"default": copy.deepcopy(dict(block)), name: runtime} + + members = entry.get("members") or members_for(data, unit, world_size) + declared = data.get("members") + if isinstance(declared, dict): + out["members"] = {**declared, name: list(members)} + elif str(data.get("parallel")) == "unified": + default = list(declared or members_for(data, None, world_size)) + out["members"] = {"default": default, name: list(members)} + + if entry.get("env"): + environment = dict(out.get("environment") or {}) + environment["env"] = {**(environment.get("env") or {}), **entry["env"]} + out["environment"] = environment + return out + + +def describe(data: Mapping[str, Any], name: str) -> str: + """What a variant asks of the target, for the planner and the assembler.""" + entry = variants(data)[name] + lines = [ + f"Variant **{name}** of unit `{base_unit(data, name)}`: {entry.get('note') or ''}".rstrip() + ] + if entry.get("runtime"): + lines.append(f"- boot-gate switches added: `{dict(entry['runtime'])}`") + if entry.get("llm_api_options"): + lines.append( + f"- LLM-API options merged into the unit's: `{dict(entry['llm_api_options'])}`" + ) + if entry.get("env"): + lines.append(f"- environment added: `{dict(entry['env'])}`") + lines.append(f"- gates: {', '.join(entry.get('gates') or DEFAULT_VARIANT_GATES)}") + return "\n".join(lines) + + +__all__ = [ + "DEFAULT_VARIANT_GATES", + "OPTIONS_DIRNAME", + "VARIANTS_FIELD", + "VARIANT_GATES", + "base_unit", + "describe", + "is_variant", + "validate_variants", + "variant_anchor_pending", + "variant_gates", + "variant_spec", + "variants", +] diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/task.example.yaml b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/task.example.yaml new file mode 100644 index 000000000000..7a079ea36b8b --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/task.example.yaml @@ -0,0 +1,839 @@ +# trtllm-modeling-v2 task specification. +# +# Copy this file, fill it in, and run: +# +# trtllm-modeling-v2 --task task.yaml --workspace workspace/trtllm-modeling-v2 +# +# Start with --dry-run: it runs preflight, resolves the launch environment +# and stops before the first agent turn, which is the cheapest way to find +# out whether this host can host this run. +# +# Or do not write one: `trtllm-modeling-v2 --checkpoint --site site.yaml +# --workspace W --goal "..."` derives the request from the checkpoint and the +# site, compiles it into this schema (W/task.resolved.yaml) and runs it +# (INTAKE.md, the fast path). +# +# Every block here is CLOSED: a key this schema does not know is an error, +# not a no-op, and the message suggests the nearest real key. A misspelled +# key that is silently ignored (`feature:` for `features:`) would resolve to +# no features: the run would build core, gate it green, and report success +# for a feature set it never built. +# +# Commented-out fields below are optional. Uncomment what you need; the +# defaults are stated next to each. + +# ---------------------------------------------------------------- identity +# A target's identity is the (checkpoint, GPU arch, parallel config) triple. +# All three end up in the run's slug, the branch name and the target's path. + +checkpoint_path: /path/to/Qwen3-30B-A3B + +# REQUIRED and never inferred. Catalog receipts are per-architecture, so the +# wrong value gates this target against another machine's certification. +# sm_100 = B200, sm_103 = GB300. +arch: sm_103 + +# tp | tep | dep | unified. The trailing number is the world size +# and therefore how many devices preflight claims. Omit for `unified` — one +# modeling codebase over a family of parallel configurations, whose world +# size comes from the member set instead. +parallel: tp1 + +# Which configurations a `unified` target is actually GATED at. Legal only +# with `parallel: unified`, because a concrete segment already names one. +# +# The failure this prevents: a unified target gated once, at pure tensor +# parallelism, leaves expert parallelism and attention data parallelism (the +# branches a unified codebase exists to carry: for a large MoE, the expert +# sharding and the all-to-all) switched on by no gate. An EP bug is invisible +# to a TP-only gate by construction. +# +# On a `local` run every member must declare the SAME world size: preflight +# claims one set of devices for them all. Under `environment.kind: +# session_spool` each gate job asks the scheduler for its own member's +# devices, so members may differ in width ([tp1, tep2, tep4]). +# `unified` is not a legal member — it is the family, not a configuration. +# +# Default when omitted, at the unified world size W (4 unless a member says +# otherwise): [tpW, tepW, depW]. TEP and DEP are the baseline pair; TP is the +# degenerate path, and it is not free — it is the only member that folds every +# branch off, so no other member covers the code it runs. +# +# COST. Each member pays boot + accuracy per unit, and accuracy is roughly an +# hour. Three members over four units is twelve accuracy gates, which +# preflight prints before the run starts. Naming one member is how a run buys +# a cheaper answer honestly, rather than getting a cheap answer labelled as an +# expensive one. +# +# members: [tp4, tep4, dep4] +# +# Per-unit, layered over `default`, same shape as `runtime:` — for when one +# unit is worth less gate budget than the rest: +# members: +# default: [tp4, tep4, dep4] +# multimodal: [tep4] + +# Feature increments, built one at a time *after* core is green. A concise +# target is the default, not an option. Known values: mtp, eagle3, dflash, +# dspark, multimodal, host_offload, disagg, kv_quant, decoder_replay, +# encoder_replay. +# +# Several are output-identical BY CONSTRUCTION and therefore carry a third +# gate, because boot and accuracy pass whether or not the feature is doing +# anything: the speculators (mtp/eagle3/dflash/dspark) get the acceptance +# gate, host_offload a residency witness, disagg an equivalence witness, +# kv_quant a bytes-per-token witness. decoder_replay and encoder_replay are NOT +# output-identical -- they truncate sliding-window state on purpose -- but a +# unit that implements nothing still passes accuracy, so they carry a work +# witness. All of those but the speculators need `accuracy.witness` below; for +# kv_quant and the two replay features it is required here, at spec time. +# +# `dspark` is not a spelling of `mtp`. A checkpoint whose `mtp.*` stages are a +# DSpark block drafter (its config carries `dspark_*` keys) asks for `dspark` +# even when the request says "MTP"; `mtp` over those weights builds chained +# one-token heads the checkpoint never trained. The two are refused together. +features: [] + +# Configurations a built unit must also serve, each gated on its own: +# a KV-cache or SSM-state dtype, a speculation depth. A variant builds nothing +# of its own; its checks (`:@`) run on its unit's target +# under the unit's settings plus the variant's overlay, and the planner claims +# them with a `feature` task whose `ref` is the variant. +# +# variants: +# kv_fp8: +# unit: core # default core +# members: [tep4] # default: the unit's members +# gates: [boot, accuracy] # the default; `acceptance` needs a speculator unit +# runtime: {extra_flags: [--kv_cache_dtype=fp8]} # boot gate, appended +# llm_api_options: {kv_cache_config: {dtype: fp8}} # merged into the unit's YAML +# env: {} # added to the gate command's environment +# accuracy: {discriminators: {kv_cache_quant_algo: FP8}} # where its anchor differs +# note: KV cache stored as FP8 +# +# A variant with no `acceptance` gate (a cache dtype, a chunk size) also boots at +# every member of its unit: its boot attests the setting took effect, cheaply; +# its accuracy stays at the members it names. The boot fails when a setting the +# overlay names is not observed on every rank (`checks/observables.py`); a +# chunked variant's boot also sends a prompt the engine has to split. A variant +# changing only lossless settings (a chunk size, block reuse, the capture list) +# is held to the logit reference's greedy tokens as its unit is. + +# What the contract proves of the request is a matrix, feature by member +# (`coverage.yaml`, and the delivery report). A requested cell no check covers +# stops the run at preflight unless it is accepted here, so "fully supported" +# says what was proven. `@`, or `@*` for a whole row; +# a pair of units no check reaches together is `+`. +# coverage: +# accept: [mtp7@tp4, disagg+multimodal@*] + +# What earlier runs' knowledge the run starts with: `production` (the default) +# uses everything, the engine-defect registry included; `measurement` leaves the +# registry and the lessons of the model's family out, to measure how the +# workflow does on a model it has not met. A measuring run's improver also turns +# while the run works; a production run's improver turns once, at the close. +# knowledge: production + +trtllm_repo_path: /path/to/TensorRT-LLM + +# Overrides for the identity the workflow derives. The family comes from the +# checkpoint's own `architectures[0]` and is almost never worth setting; the +# checkpoint segment is a GUESS made from the directory name and is the one +# component worth setting by hand. Both are Python package directory names, +# so lowercase-with-underscores. +# target: +# family: qwen3_5 +# checkpoint_segment: qwen3_5_27b_nvfp4 +# +# A kernel task on a shared catalog entry lands a local variant in the +# target's own ops/ (ops//__.py), and a promote task +# merges an entry's variants into the catalog once they have landed and the +# target imports them, beside the feature work; delivery refuses a +# target that still holds one. Name a shared entry here to keep its variants +# local instead: it is never promoted, and delivery accepts them. Default: none. +# keep_local: [norm/flashinfer_rmsnorm] + +# `full` (default) | `text_only`. A multimodal checkpoint whose vision tower +# is out of scope still has to load, route and gate as a language model; +# saying so here is what keeps an assembler from treating the absent vision +# path as a defect. +# scope: text_only + +# Free-text guidance from the user, passed verbatim to every role: code style, +# reference implementations or kernels to follow, infrastructure notes. It +# never relaxes a gate, the embargo or the scope. +# notes: | +# Reuse the NVFP4 MoE kernels in tensorrt_llm/_torch/modules/fused_moe. + +# Paths the agents may not read. Each must EXIST — an embargo naming a path +# that is not there reads as protection while protecting nothing. Use this to +# verify the pipeline derives a target rather than transliterating one that +# already exists. +# +# Two spellings, split on where the path points: +# - repo-relative, resolved under `trtllm_repo_path` — the engine's own +# implementation of this checkpoint, when it ships one; +# - ABSOLUTE, when it resolves outside the repo — a sibling vLLM or SGLang +# checkout that implements this checkpoint. Those are the trees most +# likely to be copied from on a new architecture, because the run +# deliberately starts one of them as its accuracy oracle, and they are one +# `Read` away from every agent. An absolute path INSIDE the repo is +# rejected; spell that one repo-relative. +# +# A directory entry embargoes the tree beneath it, so a checkout is named once +# rather than enumerated file by file and left to go stale. +# +# Running either kind is expected — the in-repo one under +# `TRTLLM_MODELING_V2=off`, the out-of-repo one as a black-box server. Only +# reading is forbidden. The embargo is a rule, not a tool hook: every role's +# prompt states it and lists these paths. +# +# A third spelling, `package:`, is relative to this workflow package +# wherever it is installed — for its own worked example of the model being +# onboarded (`package:examples/.md`). +# source_embargo: +# - tensorrt_llm/_torch/models/modeling_qwen3_5.py +# - /path/to/vllm +# - package:examples/.md + +# ----------------------------------------------------------------- runtime +# What must be true of every gate command. These are declared rather than +# left to the engine's defaults because the gates have to exercise them, and +# because two are spelled as negations on the command line — translating +# "block reuse on" into "emit no flag" is done once here, in Python, so no +# agent ever types it. +# +# runtime: +# cuda_graph: false # default false +# chunked_prefill: false # default false +# overlap_scheduler: true # default true (--disable_overlap_scheduler) +# kv_cache_reuse: true # default true (--disable_kv_cache_reuse) +# +# PER UNIT when the units disagree, layered over `default`: +# +# runtime: +# default: +# kv_cache_reuse: true +# extra_flags: [--max_seq_len=8192] +# multimodal: +# kv_cache_reuse: false # the vision path requires it off +# +# This block and not `llm_api_options`, which is already per-unit, because +# `quickstart_advanced.py` takes no `--extra_llm_api_options` — so the BOOT +# gate, the one that runs first and most often, cannot read that file at all +# and this is its only channel. A setting both launchers take (the KV +# fraction, KV and SSM cache dtypes, tokens per block, max_num_tokens, +# trust_remote_code) is stated once, in either place, and reaches every +# launcher: a flag here is written into the eval, serve and tune options, and +# a key in the YAML becomes the quickstart's flag. Two values for one setting +# are refused at preflight. A checkpoint that ships its own code (`auto_map`) +# is trusted by every launcher without saying so. +# +# # Anything else the gate command needs, verbatim. `max_seq_len` belongs +# # here for most checkpoints: without it the engine inherits the +# # checkpoint's `max_position_embeddings` — 262144 is common — and the +# # run dies with `Impossible to fit in any sequence in kvCache`, which +# # reads as a model defect and is not one. +# extra_flags: +# - --max_seq_len=8192 +# - --max_batch_size=8 +# - --kv_cache_fraction=0.6 + +# The runtime block's limits (max_batch_size, max_num_tokens, max_seq_len and +# the CUDA-graph capture batches) are the serving envelope: what the target is +# claimed to serve. The plan roles judge catalog coverage against them, smiths +# certify their entries' call sites at every captured batch size, and the +# reviewer checks the target's capacity constants against them — never +# against the shapes one gate happens to launch. + +# ------------------------------------------------------------- environment +# How GPU work is launched. Resolved ONCE into a literal command prefix that +# every agent prompt carries verbatim; no agent ever builds a launcher line. + +environment: + # local | session_spool + # + # `local` runs on this host, on the devices preflight claims. On a cluster + # use `session_spool`: one long-lived allocation runs the spool scheduler, + # and work is a file, so anything that can write a file can run GPU work -- + # no container import (26 GB, 4-8 minutes) or queue wait per command, and + # no need for this host to reach the Slurm CLI. + kind: local + + # The interpreter every Python command runs under. Point this at the + # editable TensorRT-LLM install you want measured — resolving it here is + # what keeps a run from silently gating against a different build. + python: /path/to/TensorRT-LLM/.venv-3.12/bin/python + + env: + # Mandatory for anything you will attribute to modeling_v2. Under + # `auto`, a configuration that misses a target's routing criteria + # silently gets the built-in implementation, and a number measured that + # way reads as modeling_v2's. `require` raises instead. + TRTLLM_MODELING_V2: require + # Engine builds, NVRTC scratch and harness output all default to TMPDIR. + # On a shared host another tenant can fill /tmp to zero mid-run, so + # preflight rejects a TMPDIR that points there. + TMPDIR: /path/to/scratch/tmp + CLAUDE_CODE_TMPDIR: /path/to/scratch/tmp + + # A checkpoint the STOCK modeling path already supports, used only as a + # stand-in by the substrate gate — the preflight step that proves this host + # can generate a token at all before any agent exists. + # + # That gate runs on the run's OWN checkpoint, which is right for two of the + # three failures it catches (a kvCache sizing fault and an NVRTC one depend + # on this checkpoint's geometry and this run's flags). But a genuinely new + # architecture has no stock path by definition, so the gate can never pass + # for exactly the runs this workflow exists for: it fails with an + # architecture-not-supported error that reads as a target defect. + # + # When that specific failure happens — and no other; a real kvCache fault is + # still a hard stop — the gate re-runs against this checkpoint under the same + # environment and the same `runtime` block. `MPI_ERR_SPAWN` and the NVRTC + # PATH fault are still proven, because neither depends on the checkpoint. The + # kvCache one is NOT, and preflight's notes say so verbatim: a sizing fault + # will then surface at the first boot gate instead of here. + # + # Prefer this over `--skip-substrate-gate`, which gives up all three. + # substrate_checkpoint: /path/to/Qwen3-30B-A3B + + # --- the session allocations (session_spool) ---------------------------- + # Submission parameters every session below inherits unless it sets its + # own; `trtllm-modeling-v2 session submission` renders the call from them. + # partition: batch + # account: your-account + # container_image: /path/to/image.sqsh + # mounts: + # - /lustre:/lustre + # time_limit: "04:00:00" + # + # GPUs per node of a session allocation. + # gpus_per_node: 4 + # + # Pyxis skips the container import when a container of this name already + # exists on the node. A 26 GB image costs 4-8 minutes per job to import, + # so this is the difference between paying that once and paying it every + # time a session starts. + # container_name: trtllm-modeling-v2 + # + # PREFER A LOCAL SQUASHFS over a registry reference. pyxis takes either, and + # the difference is the single biggest lever on how parallel a run can be: a + # registry reference is fetched and unpacked PER ALLOCATION, measured at 5-7 + # minutes for a 26 GB image and at 25-35 minutes on a loaded cluster, three + # times in a row. A file on shared storage mounts in seconds. + # + # That cost is the reason a run ends up on one shared session, and a shared + # session has ONE FAILURE DOMAIN — a leaked server or a stray process-group + # kill takes down every other job with it. + # + # enroot import -o /lustre//images/trtllm.sqsh \ + # docker://artifactory.example.com#ns/trtllm:tag + # + # then point container_image at that path. Preflight warns if you do not. + + # --- how multi-rank actually launches ------------------------------------ + # Worth knowing, because the failure is loud and the cause is not. + # + # At world size > 1 the workflow launches the command under `mpirun -n + # ` and wraps it in `trtllm-llmapi-launch`, so the ranks EXIST before + # the engine starts. Without that the MPI world size is 1, and the LLM API + # concludes it must create the other ranks itself — + # `need_spawn_mpi_workers` in `llmapi/mpi_session.py` is exactly + # `get_mpi_world_size() == 1 and model_world_size > 1`. Under pyxis that + # dynamic spawn is forbidden and the run dies with `MPI_ERR_SPAWN`. + # + # At world size 1 the opposite applies: no wrapper, and the MPI environment + # is cleared before exec, because mpi4py auto-initializes MPI whenever PMI + # variables are present and inside a PMIx step they always are. + # + # Both rules are TensorRT-LLM's own. Nothing here is this workflow's + # invention, and no agent ever assembles this line. + + # --- kind: session_spool only ------------------------------------------- + # The directory a long-lived compute session watches. Work is a file: + # write `/pending/.sh`, wait for `/done/.rc`, + # read `/logs/.log`. + # spool: /path/to/scratch/spool + # + # The sessions the run stands on. Each one is an allocation whose main + # process is the spool scheduler (`infra/substrate/scheduler.py`): it allocates + # devices per job (`--gpus N|all`), enforces wall-clock and memory limits, + # removes the allocation's MPI context from every job, and never kills a + # process that still carries it. Render a session's entry script with + # trtllm-modeling-v2 session render --task task.yaml --name a + # and submit the allocation with that script as its command; there is no + # hand-edited entry script to go stale. + # + # `time_limit` lets the scheduler plan the hand-over: `handover_lead` before + # the end it writes a successor request (the driver answers it, see + # `trtllm-modeling-v2 driver wait`) and stops admitting jobs that could not + # finish in time, so they wait in the spool for the successor instead of + # dying with the allocation. `nodes: 2` makes a multi-node session for + # `infra/scripts/mn/mn_mpirun.sh`. + # + # sessions: + # - name: a + # spool: /path/to/scratch/spool # the gate spool above + # role: gates # gates | work | any + # gpus_per_node: 4 + # time_limit: "24:00:00" + # handover_lead: "01:00:00" + # scheduler: {default_wall: 2h, pressure_floor_gb: 32} + # - name: b + # spool: /path/to/scratch/spool-b + # role: work # catalog tests, microbenchmarks + # + # How long the orchestrator waits for a gate's session to heart-beat again + # when it stops (a requeued allocation, a successor still starting). A gate + # lost that way measured nothing: it is withdrawn from the spool, re-run once + # a session is back, and neither recorded as a failure nor charged an + # assembler turn. Past this wait the run closes `session-lost:`, which + # re-running the same command resumes. Seconds, 90m, 2h or HH:MM:SS. + # session_wait: 1h + +# ---------------------------------------------------------------- accuracy +# The release criterion. One-sided: measured >= reference - k x SE, the SE +# binomial over the benchmark's sample count (see `floor` below). +# +# The anchor comes from a written source — TensorRT-LLM's own accuracy +# references — and is NEVER invented. If this checkpoint has no entry yet, +# the plan phase must add one from a source it cites, before assembly. +# +# Preflight resolves the row EVERY unit will look up — through the accuracy +# gate's own lookup, so it fails exactly where the gate would — and stops the +# run if one is missing. That is deliberate: a unit whose anchor is absent +# otherwise assembles, boots on every member and dies at its accuracy gate +# hours later, on this. The unit most likely to be the missing one is the +# increment that names a different benchmark and a different file. + +accuracy: + # The benchmark AND its protocol. A reasoning model scored as 5-shot + # completion lands near the floor whatever its ability, so a mismatch here + # reads as an assembly defect. + benchmark: gsm8k + reference_yaml: tests/integration/defs/accuracy/references/gsm8k.yaml + + # How the floor is derived. `statistical` (default): reference - k x SE, + # where SE = sqrt(p(1-p)/n) at the reference over the benchmark's sample + # count -- 0.73 points for GSM8K's 1319 questions at 92 %, so k = 2 puts the + # floor 1.5 points down. A fixed `tolerance` of 5.0 sat ~7 SE down, wide + # enough that a defect in every member passed every gate of the run that + # had it. `rule: tolerance` keeps the fixed width as an explicit choice. + # + # When EVERY member of a unit (at least `shift_min_members`, default 2) + # scores more than 1 SE under the reference, the unit still passes but the + # run records a FINDING, and delivery refuses it until it is explained: + # name the unit under `acknowledge` with the reason, and relaunch. + # floor: + # rule: statistical + # k: 2.0 + # shift_min_members: 2 + # acknowledge: + # dspark: "every member -0.9..-1.1: drafter changes rounding, see 06-gates/..." + + # In benchmark points; read only under `floor.rule: tolerance`. + tolerance: 5.0 + + # The benchmark's sample count, for the SE. Needed only when the gate cannot + # know it: it reads the count the eval prints, then `--num_samples` in + # `benchmark_args`, then this, then the split sizes it ships for gsm8k, + # mmlu, mmmu and gpqa. A scalar, or a mapping from unit to a count. + # samples: {default: 1319, multimodal: 900} + + # Where each unit's benchmark runs. `sampled` (default): with the logit gate + # owed, at the narrowest member (tp1), the widest, and every member whose + # served engine the logit gate does not run whole (attention DP, whose other + # ranks a one-prompt-at-a-time gate never reaches; a speculator, unless the + # boot compares greedy tokens, which it does with the speculator on). The + # logit gate, owed at every member, proves the others equivalent to the + # reference. Without the logit gate the benchmark runs at every member. + # `every` runs it at every member regardless. contract.yaml's `notes` say + # where each unit's runs and why. + # benchmark_members: sampled + + # `pending` where no written anchor exists for a unit yet (scalar, or per + # unit; default `recorded`). Its benchmark is then no check: a score + # compared with nothing proves nothing, and the run never measures its own + # anchor. The logit gate (`logits` below) must be on, as its oracle, and + # delivery stays incomplete until the row is recorded from a written source. + # An anchor of 1 or less is refused: anchors are in percent. + # anchor: {multimodal: pending} + + # Which row of reference_yaml this checkpoint is anchored to. Defaults to + # a value derived from the checkpoint directory name, which is right for a + # stock checkpoint and wrong for a quantized derivative that should anchor + # to its parent. + # model_key: openai/gpt-oss-120b + + # `model_key` picks the model; `discriminators` picks the ROW. The + # references hold a LIST per model, tagged with `quant_algo` / + # `kv_cache_quant_algo` / `spec_dec_algo`, and the match is EXACT — no + # more, no less. Omitting this therefore asks for the untagged row, which + # for a quantized checkpoint is the bf16 number, and for a model that has + # no untagged row is an error raised at the end of an assembly. + # + # Per unit for the same reason `llm_api_options` is: `core` and `mtp` are + # different rows of the same list and different questions. + # discriminators: + # core: {quant_algo: MIXED_PRECISION} + # mtp: {quant_algo: MIXED_PRECISION, spec_dec_algo: MTP} + + # The two-leg WITNESS gates, for the features that are output-identical by + # construction. `host_offload` and `disagg` each pass boot AND accuracy + # whether or not the feature is doing anything — the same blindness the + # acceptance gate exists for, with a different observable — so a unit with + # no witness can be declared green for a feature nobody built. + # + # Supplies the one thing the orchestrator cannot derive: the environment + # for the SECOND leg -- which turns the offload off, and turns disagg on, so + # the key is not named for a direction. `expected_mib` is the delta predicted + # from the CHECKPOINT's own shapes, passed in rather than derived, because + # a witness that computes its own expectation from the run it is judging + # can only ever agree with itself. + # + # witness: + # host_offload: + # contrast_env: {MYMODEL_HOST_OFFLOAD: "0"} # the leg WITHOUT the offload + # expected_mib: 48800 # 20 M rows x 2560 dims, fp8 + # # The target prints `__DEVICE_MEMORY_MIB__ ` once per rank after its + # # weights load and before the KV cache pool is allocated. + # disagg: + # # Leg B is the flow's own (checks/harness/disagg_leg_b.py): leg A's + # # engine split into a context and a generation engine on twice the + # # member's devices, each prompt handed over context-only -> generation- + # # only. Nothing else is owed; the keys below are all optional. + # # The cache transceiver both engines run (CacheTransceiverConfig fields); + # # default {backend: DEFAULT, transceiver_runtime: auto}. + # cache_transceiver_config: {backend: NIXL} + # # A launcher of the run's own replaces it: run as ONE process on every + # # device of the claim, handed leg A's workload as MV2FLOW_WITNESS_* + # # variables (see witness.py), and then `contrast_env` is required. A + # # word naming a file beside this spec is resolved against its directory. + # # contrast_command: scripts/my_leg_b.py + # # contrast_env: {MYMODEL_DISAGG: "1"} # leg B's environment alone + # prompt_tokens: [131, 2050, 2051, 4097] # extra prompts, BOTH legs + # timeout: 2700 # seconds per leg + # max_batch_size: 1 # one schedule on BOTH legs + # # The environment of BOTH legs (`contrast_env` is leg B's alone). + # legs_env: + # MYMODEL_ENCODER_REPLAY: "0" + # MYMODEL_DECODER_REPLAY: "0" + # + # `legs_env` (disagg only): the two legs must differ in exactly ONE thing, + # the hop. Leg B is a different launcher, and a code path its engines refuse + # or guard off under disaggregation (above: both bounded replays) stays ON in + # the aggregated leg unless pinned here -- then the legs differ in that path + # too, and token identity no longer holds by construction. Leg A's process + # gets these variables; the contrast command gets them in its environment + # and, as JSON, in MV2FLOW_WITNESS_LEGS_ENV, and must apply them to every + # engine it starts. The pinned path is witnessed by its own block. A variable + # may not appear in both `legs_env` and `contrast_env`. The host_offload and + # kv_quant witnesses run one aggregated command on both legs, so nothing but + # `contrast_env` differs between them and they do not take the key. + # + # Any witness block may set `timeout` (seconds per leg; default 1800). + # + # A witness that compares TOKENS declares its workload the way the acceptance + # gate does: `prompts` or `prompts_file`, `max_tokens`, and + # `min_generated_tokens` (default 32). With none it uses the declared + # acceptance prompts, else the boot prompts -- and a witness whose + # continuations come back shorter than the minimum fails as too short to + # compare: identity over a 2-token answer says nothing end to end. Both legs + # print token ids under TRTLLM_WITNESS_FULL_OUTPUT=1 (`[i] Witness token ids: + # [...]`, from the tree's quickstart_advanced.py and from a contrast_command); + # text that agrees without them fails rather than passes. + # + # A witness that compares TOKENS (`host_offload`, `kv_quant`, `disagg`) may + # set `max_batch_size`: both legs run at it, in place of the runtime block's + # `--max_batch_size`, and a `contrast_command` gets it as + # MV2FLOW_WITNESS_MAX_BATCH_SIZE for every engine it starts. Set it to 1 + # when the comparison must be exact: a row's bits depend on the rows beside + # it, and batched admission follows arrival timing, so two batched legs can + # pack rows differently and flip a near-tie neither feature caused. Unset, + # the legs run at the runtime block's limit, as before. + # + # `kv_quant`: the target prints `MODELING_V2_WITNESS kv_bytes_per_token=` + # once per rank under TRTLLM_MODELING_V2_WITNESS=1, and it must match the + # figure DECLARED here from the checkpoint's formats. `exact: true` (default) + # also demands token-identical output against bf16 storage; `exact: false` + # names a lossy attention path (an FP8-query kernel) and drops that leg. + # + # kv_quant: + # contrast_env: {MYMODEL_KV_STORAGE: bf16} # the leg WITH bf16 storage + # # One rank's persistent bytes per token, summed over its pools and + # # computed from the checkpoint's storage formats before the run. + # expected_bytes_per_token: 656 + # exact: true + # bytes_tolerance: 0.02 # default; row padding, not a layout + # max_batch_size: 1 # both legs one request at a time + # + # The bounded-replay features: the target prints one + # `MODELING_V2_WITNESS forwarded_tokens= decoder_tokens= ...` line + # per context request, and the witness judges it against a window and a + # workload declared HERE -- sized before the run, never from it. Each size + # must exceed the window (or replay on and off do the same work) and stay + # under 16384, since the prompt travels as one argv string. + # + # decoder_replay: + # contrast_env: {MYMODEL_DECODER_REPLAY: "0"} # replay OFF + # window: 128 # the checkpoint's sliding window + # prompt_tokens: 4096 + # encoder_replay: + # contrast_env: {MYMODEL_ENCODER_REPLAY: "0"} # replay OFF + # window: 128 + # prefix_tokens: 4096 + # reuse_alignment: 2 # e.g. a KV compression ratio the boundary may not split + + # The ACCEPTANCE gate runs in every unit whose target keeps speculation on + # (the unit that adds it and every unit after it). It measures accepted + # tokens per iteration PER PROMPT and judges them three ways: + # - against the same checkpoint under BUILT-IN modeling (`=off`), same + # session: fails when the ratio is under 0.95 AND a paired test over the + # prompts finds the drop significant; + # - against the previous tree's verdict for the same member and workload + # (kept in /gate_results.jsonl): a paired t-test (or Wilcoxon) + # with a minimum effect -- a drop must be significant AND >= `min_effect`; + # - against 1.0: every draft rejected. + # + # A genuinely new architecture has no built-in modeling, so `=off` cannot + # load it. Declare a floor for that case, and tie it to a MEASUREMENT: a + # reference drafter's acceptance on this workload (`measured_reference`; the + # floor is 0.95 x it) or a `min_absolute` with the `floor_source` it was + # derived from. A floor with neither still applies, and every verdict warns + # that it is not tied to a measured reference: one picked below every real + # drafter's acceptance passes a wrong drafter too. The run records the + # weaker claim verbatim either way: the draft path accepts at a real rate, + # NOT that it matches stock. + # + # The WORKLOAD defaults to the first 32 questions of the accuracy dataset + # (natural text, long answers) -- `dataset_path` when the run keeps it + # locally, else the benchmark's dataset id read from the local cache. Nothing + # is downloaded. Fewer than `min_prompts` (32) prompts fails as "workload too + # small": four prompts read a change in generated text as a regression. + # Fewer than `min_draft_iterations` generation iterations fails as "workload + # too short". Prompts are passed as-is; `prompt_template` (with `{text}`) + # adds the model's chat markers to dataset questions. Every workload key + # also works inside one unit's entry, for that unit only. + # + # acceptance: + # # dataset: /data/gsm8k # or a cached id; default: the benchmark's + # num_prompts: 32 + # prompt_template: "{text} Solve it step by step, then give the answer." + # # prompts: [...] # or declare the prompts outright + # # prompts_file: prompts.json # a YAML/JSON list, or boot.yaml's shape + # max_tokens: 256 + # min_draft_iterations: 32 + # min_prompts: 32 + # regression: {test: t, alpha: 0.05, min_effect: 0.1} + # mtp: + # min_absolute: 1.5 + # floor_source: "03-plan/acceptance-spike.md: reference 2.31, wrong controls <= 1.4" + # dspark: + # measured_reference: 3.32 # floor 3.15 + + # The BOOT gate's workload, declared here (else the workspace boot.yaml). It + # runs through the engine driver, which prints one structured record per + # prompt, and every keyword is matched against the GENERATED TEXT only -- + # never the log, where `72` once matched `CPUs 72-143`. A prompt may be a + # mapping with checks of its own and the images it sends; a unit whose target + # implements `multimodal` and declares no image prompt sends the repository's + # tests/integration/test_input_files/merlion.png (a git-lfs file: pulled + # in the tree under test). Per-unit entries layer over the run-wide keys. + # + # Wherever the run has a logit reference (`logits` below), the boot also + # sends the reference's first text prompts as token ids, every request runs + # its whole budget past any end token, and the engine's first greedy tokens + # must be the reference's, token for token. A differing token is allowed only + # where the reference's own top-1 leads it by at most the logit gate's + # largest |delta| limit (a near-tie; the prompt is compared no further). + # `greedy` sets how many prompts and tokens; `enabled: false` compares none. + # A configuration variant compares none unless its own entry turns it on. + # + # boot: + # prompts: + # - "The capital of France is" + # - prompt: "Natalia sold 48 clips in April and half as many in May. How many in all?" + # keywords: ["72"] + # max_tokens: 32 + # greedy: {enabled: true, prompts: 4, tokens: 24} + # multimodal: + # prompts: + # - prompt: "<|User|>What vegetable is this? One word.<|Assistant|>" + # images: [examples/images/carrots.jpeg] # beside this file, or in the tree + # render: raw # pre-rendered; or `loader` + # keywords_any: [carrot] + + # The ENVELOPE BOOT, once per member at delivery: the finished target served + # at the limits the runtime block declares -- max_batch_size requests in + # flight, prompts at max_seq_len, prefill chunks of max_num_tokens, every + # CUDA-graph capture batch replayed -- with a few probe prompts that must + # keep their first `agree_tokens` tokens from a batch-1 run. It needs + # `--max_seq_len` in runtime.extra_flags. Every other gate drives a handful + # of short prompts, so this is the only one that reaches the limits. + # + # envelope: + # enabled: true + # max_tokens: 16 + # agree_tokens: 4 + # long_prompts: 2 + # probes: ["The capital of France is", "Two plus two equals"] + # timeout: 3600 + + # The SERVE GATE: `trtllm-serve` starts each unit's target with the accuracy + # gate's options file, at the unit's widest member, and must answer one + # /v1/completions request -- and, for a unit whose boot sends an image, one + # /v1/chat/completions request carrying it -- with text. Every other gate + # builds the engine in its own process, so a defect only the server reaches + # passes all of them. Owed after every change to a unit, and once more at + # delivery on the last unit at one rank. On by default wherever the spec + # bounds the server's sequence length (`max_seq_len` below, or in the + # run-wide options file); `enabled: false` turns it off. + # + # serve: + # enabled: true + # prompt: "The capital of France is" + # max_tokens: 32 + # timeout: 1800 + + # The LOGIT GATE, on wherever a reference is named: each unit's engine (the + # boot gate's, speculator off) runs the reference's fixed prompts with prompt + # and generation log-probs, and at every position its log-prob of the HF + # reference's top-1 token is held to the reference's own bf16-vs-fp32 spread + # times `threshold_factor`. A benchmark's standard error hides the drift a + # fusion or a rounding change makes; this does not. Make the file once, in + # the anchors phase, on the HF implementation: + # python onboard/logit_reference.py --model \ + # --out logits.reference.json.gz [--trust-remote-code] \ + # [--image tests/integration/test_input_files/merlion.png --image-root ] + # A spec that names none gets `logits.reference.json.gz` beside it when that + # file is there; `enabled: false` turns the gate off. Owed after every change + # to a unit, at every member, and with it the benchmark is sampled + # (`benchmark_members` above) and the boot compares greedy tokens. + # + # logits: + # reference: logits.reference.json.gz + # prompts: [capital, clips, long, image0] # default: every one the unit takes + # threshold_factor: 3.0 + # timeout: 1800 + + # `trtllm-eval`'s own flags, not LLM-API settings, so they cannot live in + # llm_api_options. Without `max_seq_len` the eval inherits the checkpoint's + # `max_position_embeddings` and the run dies in kvCache allocation. + # max_seq_len: 8192 + # max_batch_size: 32 + + # MMLU only. trtllm-eval fetches that one dataset itself from a URL that + # now 404s; it writes the error page to data.tar and dies inside + # tarfile.open, so the failure reads as a corrupt archive rather than a + # dead link. Every other benchmark comes through lm-eval and is unaffected. + # dataset_path: /path/to/datasets/mmlu + + # The LLM-API YAML the accuracy gate runs under — a path, or a mapping + # from assembly unit to a path when the units need different settings. + # + # Per-unit is not an optimization. The `mtp` unit's gate must carry a + # `speculative_config`, and without it the run re-measures the BASE path + # and reports the number as the speculator's. Measured. + # + # It is also where the boot and acceptance gates learn the draft DEPTH + # (`speculative_config.max_draft_len`, or `block_size` for DSpark), since the + # quickstart they drive cannot read this file. For DSpark that is not a + # knob: the depth must equal the drafter's trained block size, or the engine + # rejects the config before it starts. + # + # WRITE THESE RELATIVE. A path that is not absolute is resolved against THIS + # FILE'S OWN DIRECTORY, because that is where these YAMLs are written — so + # `scripts/core_llm_api.yaml` means the same thing on every machine, and an + # absolute spelling means it on exactly one. One run produced a 220-line + # task.yaml with six absolute paths in it, and its most reusable artifact was + # unusable anywhere else. An absolute path is still honoured untouched, so a + # site that wants one keeps it. + # llm_api_options: + # core: scripts/core_llm_api.yaml + # mtp: scripts/mtp_llm_api.yaml + # dspark: scripts/dspark_llm_api.yaml # decoding_type: DSpark, max_draft_len: 5 + + # Options of the benchmark subcommand itself — the PROTOCOL, which cannot + # live in llm_api_options. A protocol mismatch between the anchor and the + # gate is indistinguishable from a real regression: measured 63.50 against + # a 95.53 anchor once, a 27-point phantom, because the anchor was scored + # with a chat template and the gate was not. + # benchmark_args: + # - --apply_chat_template + +# ------------------------------------------------------- agent backends +# Every role runs on the agent that drives the run, at its model and reasoning +# effort. Route a role elsewhere only for a specific need -- a plan reviewer on +# another provider for a cross-model review, say. The roles are plan_drafter, +# plan_reviewer, smith, assembler, reviewer and tuner. +# +# `defaults` applies to every role and `roles.` overrides it per field. +# Under `--agent-transport external` a routed role's agent reaches the driver +# as `agent` in the turn's request file. +# agents: +# roles: +# plan_reviewer: +# backend: codex +# model: gpt-6-astra +# # Portable external MCP servers (STDIO or streamable HTTP) belong under +# # defaults or an individual role: +# # extra_mcp_servers: +# # knowledge-base: {type: http, url: https://example.test/mcp} + +# ----------------------------------------------------------- commit policy +# What the agents' commit helper (`/bin/commit.sh`) accepts, and +# what delivery checks on the branch. Every commit is authored, committed and +# signed off (DCO) as the repository's configured user — `git -C +# config user.name/user.email`, or both of +# MV2FLOW_GIT_NAME / MV2FLOW_GIT_EMAIL — so the sign-off email always +# equals the author email. The run refuses to start without one. +# +# Whatever this block says, the helper refuses directories, globs, absolute +# paths, symlinks git does not track (a worktree's link to the native build), +# sandbox placeholders (`.mcp.json`, anything under `.claude/`), and the run's +# own workspace when it sits inside the repository. +# commit_policy: +# allowed_paths: [tensorrt_llm/, tests/, cpp/, examples/] # the default +# allow_co_authors: false # the default: no Co-authored-by trailer anywhere +# base_ref: origin/main # delivery audits the commits in merge-base(HEAD, base_ref)..HEAD + +# ------------------------------------------------------------- performance +# Tuning (README "Tuning"). With `enabled: true` the planned run holds a +# `tune` task that starts once the core is correct; on an already-onboarded +# target, `trtllm-modeling-v2 tune --task task.yaml --workspace W`. Perf is +# measured, never a release gate; every kept change passes the accuracy gate. +# perf: +# enabled: true +# identity_configs: [] # optional: derived from the workload when absent (tune/serving_config.py) +# workload: {isl: 8192, osl: 1024, concurrency: [1, 128]} +# configs: # instead of `workload`: several named perf configs, all measured +# - name: chat # per candidate; a candidate must lose on none of them +# workload: {isl: 1024, osl: 256, concurrency: [1, 16]} +# constraint: {tpot_ms_max: 30} # breaking it is a loss (also ttft_ms_max) +# - name: batch +# workload: {isl: 8192, osl: 1024, concurrency: [128]} +# llm_api_options: {enable_chunked_prefill: true} # over the derived serving config +# goal: {ratio_to_reference: 1.1} # stop dispatching once every headline is 1.1x the stock reference; +# # or/and pct_of_sol: 0.7, 70% of each headline's SOL ceiling +# sol: {enabled: true, gpu: gb200} # a turn derives the SOL ceilings (default on); `gpu` is the +# # part the toolkit's SOL calculator prices +# reference: stock # the same checkpoint and configs under modeling_v2: "off" +# spools: [] +# dev_spools: [] +# env_script: +# noise_repeats: 2 # identity repeats; the noise floor is the same-host spread +# campaign: +# min_gain: 0.01 # keep threshold = max(min_gain, noise_factor x spread) +# noise_factor: 3.0 +# max_parallel_tuners: 4 +# max_iterations: 12 +# split_variants: true # a LOSS split by member asks for one member-gated variant (default on) +# max_consecutive_misses: 3 +# max_roadmaps: 2 # a spent roadmap with budget left is re-profiled on the moved best tree +# max_hours: 24 # the campaign's wall budget; omit it for none +# smoke: true # one short bench per world size before the baseline (default on) +# smoke_wall: 45m # its wall limit +# accuracy: {unit: core, members: [tp1, dep4]} # the keep gate's unit and members; or {gate: off} +# source_embargo: [] +# sister_engines: [] diff --git a/agent-flow/tests/workflows/trtllm_modeling_v2/__init__.py b/agent-flow/tests/workflows/trtllm_modeling_v2/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/agent-flow/tests/workflows/trtllm_modeling_v2/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/agent-flow/tests/workflows/trtllm_modeling_v2/conftest.py b/agent-flow/tests/workflows/trtllm_modeling_v2/conftest.py new file mode 100644 index 000000000000..eeeae09885a7 --- /dev/null +++ b/agent-flow/tests/workflows/trtllm_modeling_v2/conftest.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fixtures shared across the trtllm_modeling_v2 tests. + +``repo`` is a minimal TensorRT-LLM checkout holding one merged target, +``checkpoint`` a HuggingFace checkpoint directory for it, and ``task_file`` a +task spec that points at both. pytest resolves them by name, so no test +module imports another to get them. +""" + +from __future__ import annotations + +import json + +import pytest +import yaml + +from .test_delivery import _build_repo + +SCRATCH = "/scratch/mv2flow-run-tmp" + + +@pytest.fixture +def repo(tmp_path): + return _build_repo(tmp_path) + + +@pytest.fixture +def checkpoint(tmp_path): + path = tmp_path / "gpt-oss-120b" + path.mkdir() + (path / "config.json").write_text(json.dumps({"architectures": ["GptOssForCausalLM"]})) + return path + + +@pytest.fixture +def task_file(tmp_path, repo, checkpoint): + spec = { + "checkpoint_path": str(checkpoint), + "arch": "sm_103", + "parallel": "tp1", + "trtllm_repo_path": str(repo), + "target": {"family": "gpt_oss", "checkpoint_segment": "gpt_oss_120b"}, + "environment": { + "kind": "local", + "python": "/opt/venv/bin/python", + "env": {"TMPDIR": SCRATCH, "CLAUDE_CODE_TMPDIR": SCRATCH}, + }, + "accuracy": { + "benchmark": "gsm8k", + "reference_yaml": "tests/integration/defs/accuracy/references/gsm8k.yaml", + "model_key": "openai/gpt-oss-120b", + "tolerance": 5.0, + # The fake shells answer the unit gates only; the envelope boot + # has its own tests (test_envelope_boot.py) and one through delivery. + "envelope": {"enabled": False}, + }, + } + path = tmp_path / "task.yaml" + path.write_text(yaml.safe_dump(spec), encoding="utf-8") + return path + + +@pytest.fixture(autouse=True) +def _planned_runs_start_from_the_planners_graph(monkeypatch): + """A scripted planner here writes the whole graph; the survey is not started for it. + + A new run starts with the survey beside the planner's first turn (phase 0, + ``PlanStage.survey_first``), which would add a survey turn every scripted + run had to expect. Its own tests turn it back on. + """ + from agent_flow.workflows.trtllm_modeling_v2.run.stages.plan import PlanStage + + monkeypatch.setattr(PlanStage, "survey_first", False) + + +@pytest.fixture(autouse=True) +def _git_stays_under_the_base_temp(tmp_path_factory, monkeypatch): + """Stop git's repository search at pytest's base temp directory. + + The workflow asks git about the trees these tests build -- delivery, the + commit audit and the target marker all do -- and git searches upward + from where it starts. A fixture tree that is not a checkout has to read + as not one; but when the base temp directory sits inside some other + checkout (an ignored build directory, say), the search would reach that + checkout and answer for it instead. + + git never enters a ceiling directory but still examines the ones below + it, so a repository a test creates at its ``tmp_path`` -- a child of the + base temp directory -- is found from anywhere inside it. + """ + monkeypatch.setenv("GIT_CEILING_DIRECTORIES", str(tmp_path_factory.getbasetemp())) + + +@pytest.fixture(autouse=True) +def _no_agent_probe(monkeypatch): + """Keep preflight's agent probe from starting a real agent turn. + + These tests script every turn through fakes of the role layers, and the + probe builds a layer of its own. Its tests (test_agent_probe.py) unset this. + """ + monkeypatch.setenv("MV2FLOW_SKIP_AGENT_PROBE", "1") diff --git a/agent-flow/tests/workflows/trtllm_modeling_v2/test_intake.py b/agent-flow/tests/workflows/trtllm_modeling_v2/test_intake.py new file mode 100644 index 000000000000..411d6c76f8ed --- /dev/null +++ b/agent-flow/tests/workflows/trtllm_modeling_v2/test_intake.py @@ -0,0 +1,370 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Intake: the user-facing task.yaml, its checks, its confirmation and its compilation.""" + +from __future__ import annotations + +import copy +import json + +import pytest +import yaml + +from agent_flow.workflows.trtllm_modeling_v2.spec.intake import checks +from agent_flow.workflows.trtllm_modeling_v2.spec.intake.cli import main as intake_main +from agent_flow.workflows.trtllm_modeling_v2.spec.intake.compile import compile_files, compile_task +from agent_flow.workflows.trtllm_modeling_v2.spec.intake.freeze import ( + confirm, + digest, + read_confirmation, + unconfirmed_reason, +) +from agent_flow.workflows.trtllm_modeling_v2.spec.intake.schema import IntakeError, validate_task + + +@pytest.fixture +def checkpoint(tmp_path): + path = tmp_path / "ckpt" + path.mkdir() + config = { + "architectures": ["Qwen3MoeForCausalLM"], + "num_attention_heads": 32, + "num_key_value_heads": 4, + "num_experts": 128, + "max_position_embeddings": 40960, + "num_nextn_predict_layers": 1, + } + (path / "config.json").write_text(json.dumps(config)) + return path + + +def _task(checkpoint) -> dict: + return { + "version": 1, + "goals": ["onboard", "tune"], + "model": {"checkpoint": str(checkpoint)}, + "device": {"arch": "sm_100"}, + "parallelism": [{"kind": "tep", "world": 4}, {"kind": "dep", "world": 4}], + "workload": {"isl": 8192, "osl": 1024, "concurrency": [1, 64]}, + "features": { + "cuda_graph": "required", + "overlap_scheduler": "required", + "prefix_reuse": "out_of_scope", + "kv_cache": {"status": "required", "params": {"dtype": "fp8"}}, + "speculative": {"status": "required", "params": {"algo": "mtp"}}, + "multimodal": "out_of_scope", + "long_context": {"status": "nice_to_have", "params": {"length": 32768}}, + }, + "accuracy": {"benchmark": "gsm8k", "protocol": {"chat_template": True}}, + "perf": { + "configs": [ + { + "name": "chat", + "workload": {"isl": 1024, "osl": 256, "concurrency": [1, 16]}, + "constraint": {"tpot_ms_max": 30}, + } + ], + "goal": {"ratio_to_baseline": 1.1}, + }, + "constraints": {"embargo": ["tensorrt_llm/x.py"], "base_ref": "origin/main"}, + "notes": "Reuse the NVFP4 MoE kernels.", + } + + +_SITE = { + "trtllm_repo_path": "/repo", + "environment": {"kind": "local", "python": "/opt/venv/bin/python"}, + "perf": {"spools": ["/spool/a"], "env_script": "/env.sh"}, +} + + +def test_a_well_formed_task_has_no_shape_errors(checkpoint): + assert validate_task(_task(checkpoint)) == [] + + +def test_shape_errors_are_batched(): + errors = validate_task({"version": 2, "goals": ["deploy"], "feature": {}}) + joined = "\n".join(errors) + assert "'version'" in joined and "'goals'" in joined + assert "did you mean 'task.yaml.features'" in joined + assert "'model.checkpoint' is required" in joined + + +def test_the_checks_pass_a_feasible_request(checkpoint): + findings = checks.run_checks(_task(checkpoint)) + assert checks.blocking(findings) == [], [f.render() for f in findings] + + +@pytest.mark.parametrize( + "mutate, expected", + [ + (lambda t: t.update(goals=["onboard", "add_features"]), "onboard and add_features"), + (lambda t: t.update(goals=["tune"]), "no model.base"), + (lambda t: t["parallelism"].append({"kind": "tp", "world": 3}), "32 heads do not split 3"), + (lambda t: t["workload"].update(isl=40000), "exceeds the model's 40960"), + (lambda t: t.update(open_questions=["which benchmark?"]), "which benchmark?"), + (lambda t: t["features"].pop("speculative"), "draft layers the request does not mention"), + ( + lambda t: t["perf"]["configs"][0].update(features={"multimodal": "on"}), + "turns on multimodal", + ), + ], +) +def test_each_infeasible_request_is_a_block(checkpoint, mutate, expected): + task = copy.deepcopy(_task(checkpoint)) + mutate(task) + blocks = [f.render() for f in checks.blocking(checks.run_checks(task))] + assert any(expected in b for b in blocks), blocks + + +def test_mtp_over_a_dspark_drafter_is_refused(checkpoint): + config = json.loads((checkpoint / "config.json").read_text()) + (checkpoint / "config.json").write_text(json.dumps({**config, "dspark_block_size": 5})) + blocks = [f.render() for f in checks.blocking(checks.run_checks(_task(checkpoint)))] + assert any("use algo: dspark" in b for b in blocks) + + +def test_confirmation_holds_until_the_task_changes(tmp_path, checkpoint): + path = tmp_path / "task.yaml" + path.write_text(yaml.safe_dump(_task(checkpoint))) + assert "not confirmed" in unconfirmed_reason(yaml.safe_load(path.read_text())) + confirm(path, "reviewer") + assert unconfirmed_reason(yaml.safe_load(path.read_text()), read_confirmation(path)) is None + edited = yaml.safe_load(path.read_text()) + edited["workload"]["isl"] = 4096 + assert "changed after it was confirmed" in unconfirmed_reason(edited, read_confirmation(path)) + + +def test_confirming_leaves_the_users_file_and_its_comments_as_they_are(tmp_path, checkpoint): + path = tmp_path / "task.yaml" + text = "# reviewed with the model owner\n" + yaml.safe_dump(_task(checkpoint)) + path.write_text(text) + confirm(path, "reviewer") + assert path.read_text() == text + assert read_confirmation(path)["by"] == "reviewer" + + +def test_a_confirmed_block_inside_the_task_is_still_honoured(tmp_path, checkpoint): + task = _task(checkpoint) + task["confirmed"] = {"by": "reviewer", "at": "2026-01-01T00:00:00Z", "sha256": digest(task)} + path = tmp_path / "task.yaml" + path.write_text(yaml.safe_dump(task)) + assert unconfirmed_reason(task, read_confirmation(path)) is None + + +def test_open_questions_block_confirmation(tmp_path, checkpoint): + path = tmp_path / "task.yaml" + path.write_text(yaml.safe_dump({**_task(checkpoint), "open_questions": ["tp or dep?"]})) + with pytest.raises(IntakeError, match="open question"): + confirm(path, "reviewer") + + +def test_experts_named_num_local_experts_make_a_moe_model(checkpoint): + config = json.loads((checkpoint / "config.json").read_text()) + config.pop("num_experts") + (checkpoint / "config.json").write_text(json.dumps({**config, "num_local_experts": 8})) + task = _task(checkpoint) + blocks = [f.render() for f in checks.blocking(checks.run_checks(task))] + assert not any("dense model" in b for b in blocks), blocks + task["parallelism"].append({"kind": "tep", "world": 3}) + blocks = [f.render() for f in checks.blocking(checks.run_checks(task))] + assert any("8 experts do not split 3" in b for b in blocks), blocks + + +def test_draft_layers_under_mtp_num_hidden_layers_are_draft_layers(checkpoint): + config = json.loads((checkpoint / "config.json").read_text()) + config.pop("num_nextn_predict_layers") + (checkpoint / "config.json").write_text(json.dumps({**config, "mtp_num_hidden_layers": 1})) + assert checks.blocking(checks.run_checks(_task(checkpoint))) == [] + task = _task(checkpoint) + task["features"].pop("speculative") + blocks = [f.render() for f in checks.blocking(checks.run_checks(task))] + assert any("draft layers the request does not mention" in b for b in blocks) + + +@pytest.mark.parametrize( + "mutate, expected", + [ + ( + lambda t: t["features"].update( + disaggregated={"status": "required", "members": ["tp2"]} + ), + "members tp2 are not among", + ), + (lambda t: t.update(device={"arch": "sm_100", "gpus_per_node": 2}), "the device has 2"), + ( + lambda t: t["features"].update( + speculative={"status": "required", "params": {"algo": "eagle3"}} + ), + "none is named", + ), + (lambda t: t["perf"]["configs"][0].update(parallel="tp8"), "not among the run's"), + ], +) +def test_each_new_infeasible_request_is_a_block(checkpoint, mutate, expected): + task = copy.deepcopy(_task(checkpoint)) + mutate(task) + blocks = [f.render() for f in checks.blocking(checks.run_checks(task))] + assert any(expected in b for b in blocks), blocks + + +def test_block_reuse_on_a_hybrid_checkpoint_needs_a_snapshot_policy(checkpoint): + config = json.loads((checkpoint / "config.json").read_text()) + config["layer_types"] = ["linear_attention", "full_attention"] + (checkpoint / "config.json").write_text(json.dumps(config)) + task = _task(checkpoint) + task["features"]["prefix_reuse"] = "required" + blocks = [f.render() for f in checks.blocking(checks.run_checks(task))] + assert any("snapshot policy" in b for b in blocks), blocks + task["features"]["prefix_reuse"] = {"status": "required", "params": {"snapshot_interval": 512}} + assert checks.blocking(checks.run_checks(task)) == [] + + +def test_a_dspark_drafter_drafts_its_own_block_length(checkpoint): + config = json.loads((checkpoint / "config.json").read_text()) + (checkpoint / "config.json").write_text(json.dumps({**config, "dspark_block_size": 5})) + task = _task(checkpoint) + task["features"]["speculative"] = {"status": "required", "params": {"algo": "dspark"}} + blocks = [f.render() for f in checks.blocking(checks.run_checks(task))] + assert any("blocks of 5 tokens" in b for b in blocks), blocks + task["features"]["speculative"]["params"]["draft_len"] = 5 + assert checks.blocking(checks.run_checks(task)) == [] + + +def test_compile_maps_requirements_onto_the_pipeline_spec(checkpoint): + compiled = compile_task(_task(checkpoint), _SITE) + spec = compiled.spec + assert spec["parallel"] == "unified" and spec["members"] == ["tep4", "dep4"] + assert spec["features"] == ["mtp"] + assert spec["scope"] == "text_only" + assert spec["runtime"]["cuda_graph"] is True + assert spec["runtime"]["kv_cache_reuse"] is False + assert "--max_seq_len=9216" in spec["runtime"]["extra_flags"] + assert "--max_batch_size=64" in spec["runtime"]["extra_flags"] + assert "--kv_cache_dtype=fp8" in spec["runtime"]["extra_flags"] + assert spec["accuracy"]["benchmark_args"] == ["--apply_chat_template"] + assert spec["source_embargo"] == ["tensorrt_llm/x.py"] + assert spec["commit_policy"]["base_ref"] == "origin/main" + perf = spec["perf"] + assert perf["enabled"] and perf["spools"] == ["/spool/a"] + assert perf["configs"][0]["constraint"] == {"tpot_ms_max": 30} + # The FP8 KV cache and the MTP draft each name the reference row. + assert spec["accuracy"]["discriminators"] == { + "default": {"kv_cache_quant_algo": "FP8"}, + "mtp": {"kv_cache_quant_algo": "FP8", "spec_dec_algo": "MTP"}, + } + options = compiled.llm_api_options + assert "speculative_config" not in options["core"] + assert options["mtp"]["speculative_config"] == {"decoding_type": "MTP", "max_draft_len": 1} + assert options["core"]["kv_cache_config"] == {"dtype": "fp8", "enable_block_reuse": False} + assert options["core"]["cuda_graph_config"] == {"enable_padding": True, "max_batch_size": 64} + assert perf["goal"] == {"ratio_to_reference": 1.1} + assert any("long_context" in w for w in compiled.warnings) + assert spec["notes"] == "Reuse the NVFP4 MoE kernels." + assert not any(w.startswith("notes") for w in compiled.warnings) + + +def test_a_quantized_checkpoint_anchors_to_its_own_row(checkpoint): + task = _task(checkpoint) + task["model"]["quantization"] = {"weights": "nvfp4"} + task["features"].pop("kv_cache") + spec = compile_task(task, _SITE).spec + assert spec["accuracy"]["discriminators"] == { + "default": {"quant_algo": "NVFP4"}, + "mtp": {"quant_algo": "NVFP4", "spec_dec_algo": "MTP"}, + } + task["accuracy"]["discriminators"] = {"quant_algo": "MIXED_PRECISION"} + assert compile_task(task, _SITE).spec["accuracy"]["discriminators"] == { + "quant_algo": "MIXED_PRECISION" + } + + +def test_a_feature_narrowed_to_members_is_gated_on_those_only(checkpoint): + task = _task(checkpoint) + task["features"]["disaggregated"] = {"status": "required", "members": ["tep4"]} + spec = compile_task(task, _SITE).spec + assert spec["members"] == {"default": ["tep4", "dep4"], "disagg": ["tep4"]} + + +def test_every_requirement_compile_does_not_carry_is_warned(checkpoint): + task = _task(checkpoint) + task["features"]["cuda_graph"] = {"status": "required", "params": {"batch_sizes": [1, 2]}} + task["accuracy"]["protocol"]["shots"] = 0 + task["perf"]["metric"] = "output_tps_per_gpu" + task["deliverable"] = {"branch": "auto"} + warnings = "\n".join(compile_task(task, _SITE).warnings) + for expected in ( + "features.cuda_graph.params.batch_sizes", + "accuracy.protocol.shots", + "perf.metric", + "deliverable", + ): + assert expected in warnings, warnings + + +def test_tune_measures_each_member_at_its_parallelism_with_the_draft(checkpoint): + task = _task(checkpoint) + task["perf"].pop("configs") + configs = compile_task(task, _SITE).spec["perf"]["configs"] + assert [c["name"] for c in configs] == ["tep4", "dep4"] + assert configs[1]["llm_api_options"] == { + "tensor_parallel_size": 4, + "moe_expert_parallel_size": 4, + "enable_attention_dp": True, + "speculative_config": {"decoding_type": "MTP", "max_draft_len": 1}, + } + task = _task(checkpoint) + task["perf"]["configs"][0].update(parallel="tep4", features={"speculative": "on"}) + options = compile_task(task, _SITE).spec["perf"]["configs"][0]["llm_api_options"] + assert options["moe_expert_parallel_size"] == 4 and "speculative_config" in options + + +def test_a_sol_goal_compiles_beside_the_ratio(checkpoint): + task = _task(checkpoint) + task["perf"]["goal"] = {"ratio_to_baseline": 1.1, "pct_of_sol": 0.7} + goal = compile_task(task, _SITE).spec["perf"]["goal"] + assert goal == {"ratio_to_reference": 1.1, "pct_of_sol": 0.7} + + +def test_compiled_spec_passes_the_pipeline_schema(tmp_path, checkpoint): + repo = tmp_path / "repo" + repo.mkdir() + task, site = tmp_path / "task.yaml", tmp_path / "site.yaml" + task.write_text(yaml.safe_dump(_task(checkpoint))) + site.write_text(yaml.safe_dump({**_SITE, "trtllm_repo_path": str(repo)})) + (repo / "tensorrt_llm").mkdir() + (repo / "tensorrt_llm" / "x.py").write_text("") + out = tmp_path / "ws" / "task.resolved.yaml" + compile_files(task, site, out) + spec = yaml.safe_load(out.read_text()) + assert spec["arch"] == "sm_100" + options = spec["accuracy"]["llm_api_options"] + assert set(options) == {"core", "mtp"} + mtp = yaml.safe_load((tmp_path / "ws" / "llm_api" / "mtp.yaml").read_text()) + assert options["mtp"] == str(tmp_path / "ws" / "llm_api" / "mtp.yaml") + assert mtp["speculative_config"]["decoding_type"] == "MTP" + + +def test_the_cli_checks_confirms_and_compiles(tmp_path, checkpoint, capsys): + task = tmp_path / "task.yaml" + task.write_text(yaml.safe_dump(_task(checkpoint))) + assert intake_main(["check", str(task)]) == 0 + assert intake_main(["confirm", str(task), "--by", "reviewer"]) == 0 + assert "confirmed by reviewer" in capsys.readouterr().out + bad = tmp_path / "bad.yaml" + bad.write_text(yaml.safe_dump({**_task(checkpoint), "goals": []})) + assert intake_main(["check", str(bad)]) == 1 + + +def test_the_run_refuses_an_unconfirmed_intake_task(tmp_path, checkpoint): + from agent_flow.workflows.trtllm_modeling_v2.cli import _resolve_intake_task + + task, site = tmp_path / "task.yaml", tmp_path / "site.yaml" + task.write_text(yaml.safe_dump(_task(checkpoint))) + site.write_text(yaml.safe_dump(_SITE)) + with pytest.raises(IntakeError, match="not confirmed"): + _resolve_intake_task(task, site, tmp_path / "ws") + resolved = tmp_path / "resolved.yaml" + resolved.write_text(yaml.safe_dump({"checkpoint_path": "/c"})) + assert _resolve_intake_task(resolved, None, tmp_path / "ws") == resolved diff --git a/agent-flow/tests/workflows/trtllm_modeling_v2/test_intake_fast_path.py b/agent-flow/tests/workflows/trtllm_modeling_v2/test_intake_fast_path.py new file mode 100644 index 000000000000..6195bcf5d388 --- /dev/null +++ b/agent-flow/tests/workflows/trtllm_modeling_v2/test_intake_fast_path.py @@ -0,0 +1,408 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The fast path: a request derived from the checkpoint, the site and a goal, checked in phase 0.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import yaml + +from agent_flow.logger import get_logger +from agent_flow.workflows.trtllm_modeling_v2 import cli +from agent_flow.workflows.trtllm_modeling_v2.run.preflight import check_intake_request +from agent_flow.workflows.trtllm_modeling_v2.run.stages.preflight import Phase0Probe +from agent_flow.workflows.trtllm_modeling_v2.run.state import ( + STAGE_PLANNED, + STATE_FILENAME, + WorkflowState, + save_state, +) +from agent_flow.workflows.trtllm_modeling_v2.spec.intake import derive +from agent_flow.workflows.trtllm_modeling_v2.spec.intake.checks import blocking, run_checks +from agent_flow.workflows.trtllm_modeling_v2.spec.intake.cli import main as intake_main +from agent_flow.workflows.trtllm_modeling_v2.spec.intake.compile import compile_files +from agent_flow.workflows.trtllm_modeling_v2.spec.intake.schema import IntakeError +from agent_flow.workflows.trtllm_modeling_v2.spec.task_schema import load_and_validate_task_yaml + +from .conftest import SCRATCH +from .test_workflow import _workflow + +REFERENCES = "tests/integration/defs/accuracy/references/gsm8k.yaml" + + +def _checkpoint(root: Path, **config) -> Path: + path = root / "newmodel" + path.mkdir(parents=True, exist_ok=True) + base = { + "architectures": ["NewModelForCausalLM"], + "_name_or_path": "org/newmodel", + "num_attention_heads": 16, + "num_key_value_heads": 4, + "num_experts": 16, + "max_position_embeddings": 32768, + "num_nextn_predict_layers": 1, + "vision_config": {"model_type": "vit"}, + } + (path / "config.json").write_text(json.dumps({**base, **config})) + return path + + +def _repo(root: Path, rows: dict | None = None) -> Path: + repo = root / "repo" + references = repo / REFERENCES + references.parent.mkdir(parents=True, exist_ok=True) + rows = rows if rows is not None else {"org/newmodel": [{"accuracy": 80.0}]} + rows["org/newmodel"] = [ + *rows.get("org/newmodel", []), + {"spec_dec_algo": "MTP", "accuracy": 80.0}, + ] + references.write_text(yaml.safe_dump(rows)) + return repo + + +def _site(root: Path, repo: Path, **extra) -> dict: + return { + "trtllm_repo_path": str(repo), + "environment": { + "kind": "session_spool", + "python": "/opt/venv/bin/python", + "spool": str(root / "spool"), + "env": {"TMPDIR": SCRATCH, "CLAUDE_CODE_TMPDIR": SCRATCH}, + }, + "device": {"arch": "sm_100", "gpus_per_node": 4}, + "defaults": {"members": ["tp1", "tep4", "dep4"]}, + **extra, + } + + +def _site_file(root: Path, site: dict) -> Path: + (root / "spool").mkdir(exist_ok=True) + path = root / "site.yaml" + path.write_text(yaml.safe_dump(site)) + return path + + +# ------------------------------------------------------------------ the goal + + +def test_the_goal_names_features_members_and_the_workload(): + goal = derive.parse_goal( + "mtp, disagg@tp1/tep4 -multimodal tp1 tep4; isl=8192 concurrency=1/16 tune" + ) + assert goal.features == { + "speculative": "required", + "disaggregated": "required", + "multimodal": "out_of_scope", + } + assert goal.algo == "mtp" + assert goal.feature_members == {"disaggregated": ["tp1", "tep4"]} + assert goal.members == ["tp1", "tep4"] + assert goal.workload == {"isl": 8192, "concurrency": [1, 16]} + assert goal.tune + + +def test_every_word_the_goal_does_not_know_is_named_at_once(): + with pytest.raises(IntakeError) as raised: + derive.parse_goal("mtpp tp1 isl=lots colour=blue kv=int3") + message = str(raised.value) + for word in ("'mtpp'", "isl=lots", "'colour'", "kv=int3"): + assert word in message, message + + +# ------------------------------------------------------------ the derivation + + +def test_features_the_checkpoint_carries_are_built_with_their_provenance(tmp_path): + repo = _repo(tmp_path) + derivation = derive.derive_task( + _checkpoint(tmp_path), _site(tmp_path, repo), derive.parse_goal("") + ) + task = derivation.task + assert task["features"]["speculative"] == { + "status": "required", + "params": {"algo": "mtp", "draft_len": 1}, + } + assert task["features"]["multimodal"] == {"status": "required"} + assert task["features"]["cuda_graph"] == {"status": "required"} + assert [f"{p['kind']}{p['world']}" for p in task["parallelism"]] == ["tp1", "tep4", "dep4"] + assert task["workload"] == derive.DEFAULT_WORKLOAD + assert task["accuracy"]["model_key"] == "org/newmodel" + assert task["accuracy"]["boot"] == derive.DEFAULT_BOOT + assert task["provenance"]["features.speculative.params.algo"] == "inferred" + assert task["provenance"]["workload.isl"] == "default" + assert task["provenance"]["parallelism"] == "inferred" + assert task["open_questions"] == [] + assert task["derived"]["goal"] == "" + assert blocking(run_checks(task)) == [] + + +def test_the_goal_overrides_what_the_checkpoint_and_the_site_say(tmp_path): + repo = _repo(tmp_path) + goal = derive.parse_goal("-mtp text_only tep4 isl=1024 osl=256 arch=sm_103") + task = derive.derive_task(_checkpoint(tmp_path), _site(tmp_path, repo), goal).task + assert task["features"]["speculative"] == {"status": "out_of_scope"} + assert task["features"]["multimodal"] == {"status": "out_of_scope"} + assert task["parallelism"] == [{"kind": "tep", "world": 4}] + assert task["workload"]["isl"] == 1024 and task["provenance"]["workload.isl"] == "user" + assert task["device"]["arch"] == "sm_103" + + +def test_what_cannot_be_derived_is_an_open_question_and_blocks(tmp_path): + repo = _repo(tmp_path) + site = _site(tmp_path, repo) + site.pop("device") + task = derive.derive_task(_checkpoint(tmp_path), site, derive.parse_goal("disagg")).task + questions = "\n".join(task["open_questions"]) + assert "device.arch" in questions + assert "accuracy.witness.disagg" not in questions + assert len(blocking(run_checks(task))) >= 1 + + +def test_a_disagg_witness_without_a_launcher_runs_the_flows_leg_b(tmp_path): + repo = _repo(tmp_path) + site = _site(tmp_path, repo) + task = derive.derive_task(_checkpoint(tmp_path), site, derive.parse_goal("disagg")).task + assert task["accuracy"]["witness"] == {"disagg": {}} + assert task["open_questions"] == [] + site = _site(tmp_path, repo) + site["defaults"]["witness"] = {"disagg": {"contrast_command": "leg_b.py"}} + task = derive.derive_task(_checkpoint(tmp_path), site, derive.parse_goal("disagg")).task + assert any("accuracy.witness.disagg: declare contrast_env" in q for q in task["open_questions"]) + + +def test_the_sites_witness_declaration_is_carried_to_the_run(tmp_path): + repo = _repo(tmp_path) + witness = {"contrast_env": {"LEG": "B"}, "contrast_command": "leg_b.py"} + site = _site(tmp_path, repo) + site["defaults"]["witness"] = {"disagg": witness} + task = derive.derive_task(_checkpoint(tmp_path), site, derive.parse_goal("disagg")).task + assert task["accuracy"]["witness"] == {"disagg": witness} + assert task["open_questions"] == [] + + +def test_a_hybrid_checkpoint_gets_a_snapshot_policy_for_block_reuse(tmp_path): + repo = _repo(tmp_path) + checkpoint = _checkpoint(tmp_path, layer_types=["linear_attention", "full_attention"]) + task = derive.derive_task(checkpoint, _site(tmp_path, repo), derive.parse_goal("")).task + params = task["features"]["prefix_reuse"]["params"] + assert params == {"snapshot_interval": derive.DEFAULT_SNAPSHOT_INTERVAL} + assert blocking(run_checks(task)) == [] + # The check this default answers still fails without it. + task["features"]["prefix_reuse"] = {"status": "required"} + assert any("snapshot policy" in f.render() for f in blocking(run_checks(task))) + + +def test_a_default_workload_fits_the_model(tmp_path): + repo = _repo(tmp_path) + checkpoint = _checkpoint(tmp_path, max_position_embeddings=2048) + task = derive.derive_task(checkpoint, _site(tmp_path, repo), derive.parse_goal("")).task + assert task["workload"]["isl"] + task["workload"]["osl"] <= 2048 + stated = derive.parse_goal("isl=4096 osl=1024") + task = derive.derive_task(checkpoint, _site(tmp_path, repo), stated).task + assert any("exceeds the model" in f.render() for f in blocking(run_checks(task))) + + +# ---------------------------------------------- compiled, and checked in phase 0 + + +def _compiled(tmp_path, task: dict, site: dict) -> tuple[Path, dict]: + request = tmp_path / "ws" / derive.REQUEST_FILENAME + request.parent.mkdir(parents=True, exist_ok=True) + request.write_text(yaml.safe_dump(task)) + resolved = tmp_path / "ws" / "task.resolved.yaml" + compile_files(request, _site_file(tmp_path, site), resolved) + return resolved, load_and_validate_task_yaml(resolved) + + +def test_a_derived_request_compiles_to_a_spec_the_phase_0_checks_pass(tmp_path): + repo = _repo(tmp_path) + site = _site(tmp_path, repo) + task = derive.derive_task(_checkpoint(tmp_path), site, derive.parse_goal("-multimodal")).task + resolved, spec = _compiled(tmp_path, task, site) + assert spec["features"] == ["mtp"] + assert spec["accuracy"]["discriminators"] == {"default": {}, "mtp": {"spec_dec_algo": "MTP"}} + mtp = yaml.safe_load(Path(spec["accuracy"]["llm_api_options"]["mtp"]).read_text()) + assert mtp["speculative_config"] == {"decoding_type": "MTP", "max_draft_len": 1} + assert spec["accuracy"]["boot"]["prompts"][0]["keywords_any"] == ["Paris"] + findings = check_intake_request(task, spec) + assert blocking(findings) == [], [f.render() for f in findings] + + +def test_a_quantized_checkpoint_without_its_row_is_a_phase_0_block(tmp_path): + repo = _repo(tmp_path) + site = _site(tmp_path, repo) + checkpoint = _checkpoint(tmp_path, quantization_config={"quant_algo": "NVFP4"}) + task = derive.derive_task(checkpoint, site, derive.parse_goal("-multimodal")).task + _, spec = _compiled(tmp_path, task, site) + blocks = [f.render() for f in blocking(check_intake_request(task, spec))] + # Both units miss their NVFP4 row; neither falls back to the unquantized one. + assert sum("accuracy anchor" in b for b in blocks) == 2, blocks + + +# --------------------------------------------------------------- the command + + +def _fast(tmp_path, site_path, goal=None, confirm_by=None, clean=False): + return cli._derive_intake_task( + _checkpoint(tmp_path), + site_path, + tmp_path / "ws", + goal=goal, + confirm_by=confirm_by, + clean=clean, + ) + + +def _start(workspace: Path) -> None: + save_state( + workspace / STATE_FILENAME, + WorkflowState(task_path=str(workspace / "task.yaml"), stage=STAGE_PLANNED), + ) + + +def test_the_command_derives_writes_and_compiles_in_one_step(tmp_path): + site_path = _site_file(tmp_path, _site(tmp_path, _repo(tmp_path))) + resolved, request = _fast(tmp_path, site_path, goal="-multimodal") + assert request == tmp_path / "ws" / derive.REQUEST_FILENAME + assert (tmp_path / "ws" / derive.REPORT_FILENAME).is_file() + assert load_and_validate_task_yaml(resolved)["features"] == ["mtp"] + + +def test_an_edit_before_the_run_starts_is_taken_and_after_it_is_refused(tmp_path): + site_path = _site_file(tmp_path, _site(tmp_path, _repo(tmp_path))) + _, request = _fast(tmp_path, site_path, goal="-multimodal") + text = request.read_text() + request.write_text(text.replace("osl: 1024", "osl: 512")) + resolved, _ = _fast(tmp_path, site_path) + assert "--max_seq_len=4608" in load_and_validate_task_yaml(resolved)["runtime"]["extra_flags"] + + _start(tmp_path / "ws") + _fast(tmp_path, site_path) # unchanged since the run started + request.write_text(request.read_text().replace("osl: 512", "osl: 256")) + with pytest.raises(IntakeError, match="changed after its run started"): + _fast(tmp_path, site_path) + _fast(tmp_path, site_path, confirm_by="reviewer") + request.write_text(request.read_text().replace("osl: 256", "osl: 128")) + with pytest.raises(IntakeError, match="changed after it was confirmed"): + _fast(tmp_path, site_path) + + +def test_another_goal_or_checkpoint_needs_clean(tmp_path): + site_path = _site_file(tmp_path, _site(tmp_path, _repo(tmp_path))) + _fast(tmp_path, site_path, goal="-multimodal") + with pytest.raises(IntakeError, match="--clean"): + _fast(tmp_path, site_path, goal="-multimodal -mtp") + _, request = _fast(tmp_path, site_path, goal="-multimodal -mtp", clean=True) + assert yaml.safe_load(request.read_text())["features"]["speculative"]["status"] == ( + "out_of_scope" + ) + + +def test_confirming_a_request_with_a_block_is_refused_with_every_block(tmp_path): + site = _site(tmp_path, _repo(tmp_path)) + site.pop("device") + site_path = _site_file(tmp_path, site) + with pytest.raises(IntakeError) as raised: + _fast(tmp_path, site_path, goal="disagg", confirm_by="reviewer") + assert "device.arch" in str(raised.value) and "contrast_command" not in str(raised.value) + + +def test_the_command_line_takes_a_checkpoint_or_a_task(tmp_path): + with pytest.raises(SystemExit): + cli._parse_args(["--checkpoint", str(tmp_path)]) + with pytest.raises(SystemExit): + cli._parse_args(["--task", "t.yaml", "--goal", "mtp"]) + args = cli._parse_args( + ["--checkpoint", str(tmp_path), "--site", "s.yaml", "--goal", "mtp", "--confirm"] + ) + assert args.confirm == "" and args.goal == "mtp" + + +def test_a_reload_does_not_confirm_again(monkeypatch): + seen = {} + monkeypatch.setattr(cli.os, "execve", lambda path, argv, env: seen.update(argv=argv)) + argv = ["python", "--checkpoint", "c", "--confirm", "--site", "s.yaml", "--confirm", "me"] + cli._relauncher([*argv, "--confirm=me", "--workspace", "ws"], {})() + assert seen["argv"][1:] == ["--checkpoint", "c", "--site", "s.yaml", "--workspace", "ws"] + + +def test_intake_derive_writes_the_request_and_its_report(tmp_path, capsys): + site_path = _site_file(tmp_path, _site(tmp_path, _repo(tmp_path))) + out = tmp_path / "out" / "task.yaml" + argv = ["derive", "--checkpoint", str(_checkpoint(tmp_path)), "--site", str(site_path)] + assert intake_main([*argv, "--goal=-multimodal", "-o", str(out)]) == 0 + assert (out.parent / derive.REPORT_FILENAME).is_file() + assert intake_main(["check", str(out)]) == 0 + + +# --------------------------------------------------------------------- phase 0 + + +def _request(tmp_path, checkpoint: Path, **changes) -> Path: + task = { + "version": 1, + "goals": ["onboard"], + "model": {"checkpoint": str(checkpoint)}, + "device": {"arch": "sm_103"}, + "parallelism": [{"kind": "tp", "world": 1}], + "workload": {"isl": 1024, "osl": 256}, + **changes, + } + path = tmp_path / derive.REQUEST_FILENAME + path.write_text(yaml.safe_dump(task)) + return path + + +def test_phase_0_stops_on_every_intake_block_before_the_host_is_touched( + tmp_path, task_file, checkpoint +): + checkpoint.joinpath("config.json").write_text( + json.dumps({"architectures": ["GptOssForCausalLM"], "num_attention_heads": 8}) + ) + workflow = _workflow(tmp_path) + log = get_logger().console + state = workflow._init_state(task_file, log) + workflow.intake_request = _request( + tmp_path, + checkpoint, + parallelism=[{"kind": "tp", "world": 3}], + open_questions=["which benchmark?"], + ) + assert workflow._run_preflight(state, log) is False + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "intake checks failed: 3 block(s)" in journal + assert "8 heads do not split 3" in journal and "which benchmark?" in journal + assert workflow.environment is None + + +def test_phase_0_records_the_confirmation_and_goes_on(tmp_path, task_file, checkpoint): + workflow = _workflow(tmp_path) + log = get_logger().console + state = workflow._init_state(task_file, log) + workflow.intake_request = _request(tmp_path, checkpoint) + assert workflow._run_preflight(state, log) is True + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "intake: not confirmed; the run does not wait for it" in journal + + +def test_phase_0_probes_run_beside_and_a_failing_one_does_not_stop_the_run(tmp_path, task_file): + seen = [] + + def measure(workflow, result): + seen.append(result.world_size) + + def broken(workflow, result): + raise RuntimeError("no engine") + + workflow = _workflow(tmp_path) + workflow.phase0_probes = (Phase0Probe("measure", measure), Phase0Probe("broken", broken)) + log = get_logger().console + state = workflow._init_state(task_file, log) + assert workflow._run_preflight(state, log) is True + assert workflow._await_phase0_probes(timeout=10) + assert seen == [1] + assert "phase-0 probe `broken` failed" in (tmp_path / "ws" / "journal.md").read_text() diff --git a/agent-flow/tests/workflows/trtllm_modeling_v2/test_task_schema.py b/agent-flow/tests/workflows/trtllm_modeling_v2/test_task_schema.py new file mode 100644 index 000000000000..fafacfb08f07 --- /dev/null +++ b/agent-flow/tests/workflows/trtllm_modeling_v2/test_task_schema.py @@ -0,0 +1,1104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for trtllm-modeling-v2's ``task.yaml`` validation.""" + +from __future__ import annotations + +import shlex + +import pytest +import yaml + +from agent_flow.workflows.trtllm_modeling_v2.spec import task_schema + + +def _valid_spec(tmp_path) -> dict: + checkpoint = tmp_path / "ckpt" + checkpoint.mkdir() + repo = tmp_path / "repo" + repo.mkdir() + return { + "checkpoint_path": str(checkpoint), + "arch": "sm_103", + "parallel": "tp1", + "trtllm_repo_path": str(repo), + "environment": { + "kind": "local", + "python": "/usr/bin/python3", + "env": {"TRTLLM_MODELING_V2": "require", "TMPDIR": "/scratch/tmp"}, + }, + "accuracy": { + "benchmark": "gsm8k", + "reference_yaml": "tests/integration/defs/accuracy/references/gsm8k.yaml", + }, + } + + +def _write(tmp_path, spec) -> str: + path = tmp_path / "task.yaml" + path.write_text(yaml.safe_dump(spec), encoding="utf-8") + return str(path) + + +def test_valid_spec_round_trips_and_fills_defaults(tmp_path): + spec = _valid_spec(tmp_path) + del spec["parallel"] + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert data["parallel"] == task_schema.PARALLEL_UNIFIED + assert data["features"] == [] + assert data["accuracy"]["tolerance"] == task_schema.DEFAULT_ACCURACY_TOLERANCE + assert data["environment"]["env"]["TRTLLM_MODELING_V2"] == "require" + + +def test_missing_arch_is_rejected_with_the_reason(tmp_path): + """``arch`` is never inferred, and the message must say why.""" + spec = _valid_spec(tmp_path) + del spec["arch"] + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert "arch" in str(excinfo.value) + assert "per-architecture" in str(excinfo.value) + + +def test_every_problem_is_batched_into_one_message(tmp_path): + """A user fixing a spec should see all of it at once, not one per run.""" + spec = _valid_spec(tmp_path) + del spec["arch"] + del spec["environment"] + del spec["accuracy"] + spec["checkpoint_path"] = "/definitely/not/here" + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + message = str(excinfo.value) + for expected in ("arch", "environment", "accuracy", "non-existent path"): + assert expected in message + + +@pytest.mark.parametrize("arch", ["sm100", "SM_103", "blackwell", "sm_", ""]) +def test_malformed_arch_is_rejected(tmp_path, arch): + spec = _valid_spec(tmp_path) + spec["arch"] = arch + with pytest.raises(task_schema.TaskSchemaError, match="arch"): + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + + +@pytest.mark.parametrize( + "segment,expected", + [ + ("tp1", ("tp", 1)), + ("tp8", ("tp", 8)), + ("tep4", ("tep", 4)), + ("dep4", ("dep", 4)), + ("unified", ("unified", None)), + ], +) +def test_parallel_segments_parse(segment, expected): + assert task_schema.parse_parallel_segment(segment) == expected + + +@pytest.mark.parametrize("segment", ["single", "tp", "tp0", "ep4", "tp-4", "TP4"]) +def test_bad_parallel_segments_are_rejected(segment): + with pytest.raises(task_schema.TaskSchemaError): + task_schema.parse_parallel_segment(segment) + + +# ------------------------------------------------------------- the member set + + +def _unified(tmp_path, members=None) -> dict: + spec = _valid_spec(tmp_path) + spec["parallel"] = task_schema.PARALLEL_UNIFIED + if members is not None: + spec["members"] = members + return spec + + +def test_a_unified_run_defaults_to_the_baseline_pair_plus_the_degenerate_path(tmp_path): + """`unified` is DEFAULT_PARALLEL. + + So the default member set is what most runs actually get gated over. + """ + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, _unified(tmp_path))) + assert task_schema.members_for(data, "core", 4) == ["tp4", "tep4", "dep4"] + assert task_schema.members_for(data, "core", 8) == ["tp8", "tep8", "dep8"] + + +def test_a_concrete_segment_is_its_own_member_set(tmp_path): + """One code path: the single-configuration run is the degenerate case.""" + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, _valid_spec(tmp_path))) + assert task_schema.members_for(data, "core", 4) == ["tp1"] + + +def test_an_explicit_list_applies_to_every_unit(tmp_path): + spec = _unified(tmp_path, ["tep4", "dep4"]) + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert task_schema.members_for(data, "core", 4) == ["tep4", "dep4"] + assert task_schema.members_for(data, "mtp", 4) == ["tep4", "dep4"] + + +def test_a_per_unit_mapping_is_layered_over_default(tmp_path): + """A per-unit mapping is layered over `default`. + + The shape `runtime:` already uses, and the knob for buying a cheaper + answer honestly: one unit gated at one member, named. + """ + spec = _unified(tmp_path, {"default": ["tp4", "tep4", "dep4"], "multimodal": ["tep4"]}) + spec["features"] = ["multimodal"] + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert task_schema.members_for(data, "core", 4) == ["tp4", "tep4", "dep4"] + assert task_schema.members_for(data, "multimodal", 4) == ["tep4"] + + +def test_a_per_unit_mapping_needs_a_default(tmp_path): + spec = _unified(tmp_path, {"multimodal": ["tep4"]}) + with pytest.raises(task_schema.TaskSchemaError, match="default"): + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + + +def test_an_unknown_segment_is_rejected_as_a_member(tmp_path): + spec = _unified(tmp_path, ["tp4", "ep4"]) + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert "members[1]" in str(excinfo.value) + + +def test_unified_is_not_a_member_of_itself(tmp_path): + """Accepting it would be the TP-only gate under another name.""" + spec = _unified(tmp_path, ["tp4", "unified"]) + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert "family" in str(excinfo.value) + + +def test_mixed_world_sizes_are_rejected_naming_both(tmp_path): + """One allocation serves every member, so a wider member cannot be run.""" + spec = _unified(tmp_path, ["tp4", "dep8"]) + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + message = str(excinfo.value) + assert "tp4" in message and "dep8" in message + assert "4" in message and "8" in message + + +def test_members_beside_a_concrete_segment_is_an_error_that_says_so(tmp_path): + """`parallel: tep4` already names the one configuration. + + Two declarations of the same thing can disagree. + """ + spec = _valid_spec(tmp_path) + spec["parallel"] = "tep4" + spec["members"] = ["dep4"] + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert "only legal with" in str(excinfo.value) + assert "tep4" in str(excinfo.value) + + +def test_an_empty_member_list_is_a_unit_that_is_never_gated(tmp_path): + spec = _unified(tmp_path, []) + with pytest.raises(task_schema.TaskSchemaError, match="non-empty"): + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + + +def test_unknown_feature_is_rejected_by_name(tmp_path): + """A feature with no assembler skill is an error here, not a confused agent later.""" + spec = _valid_spec(tmp_path) + spec["features"] = ["mtp", "telepathy"] + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert "telepathy" in str(excinfo.value) + assert "mtp" in str(excinfo.value) # listed among the known features + + +def test_every_known_feature_maps_to_a_skill(): + assert set(task_schema.FEATURE_SKILLS) >= {"mtp", "host_offload", "multimodal"} + for feature, skill in task_schema.FEATURE_SKILLS.items(): + assert skill.startswith("trtllm-modeling-v2-"), feature + + +@pytest.mark.parametrize("feature", ["eplb", "dwdp"]) +def test_features_with_no_target_behind_them_are_refused(tmp_path, feature): + """No in-tree target exercises them, so there is no skill to build them from.""" + spec = _valid_spec(tmp_path) + spec["features"] = [feature] + with pytest.raises(task_schema.TaskSchemaError, match=feature): + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + + +@pytest.mark.parametrize("kind", ["slurm", "slurm_broker"]) +def test_only_local_and_session_spool_are_environment_kinds(tmp_path, kind): + spec = _valid_spec(tmp_path) + spec["environment"]["kind"] = kind + with pytest.raises(task_schema.TaskSchemaError, match="session_spool"): + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + + +def test_local_kind_does_not_require_slurm_fields(tmp_path): + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, _valid_spec(tmp_path))) + assert data["environment"]["kind"] == "local" + + +def test_negative_tolerance_is_rejected(tmp_path): + spec = _valid_spec(tmp_path) + spec["accuracy"]["tolerance"] = -1 + with pytest.raises(task_schema.TaskSchemaError, match="tolerance"): + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + + +def test_derive_slug_is_a_safe_path_component(tmp_path): + checkpoint = tmp_path / "DeepSeek-R1-0528_NVFP4" + checkpoint.mkdir() + spec = _valid_spec(tmp_path) + spec["checkpoint_path"] = str(checkpoint) + spec["parallel"] = "dep4" + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + slug = task_schema.derive_slug(data) + assert slug == "deepseek-r1-0528-nvfp4-sm_103-dep4" + assert "/" not in slug and " " not in slug + + +def test_assembly_units_put_core_first(tmp_path): + spec = _valid_spec(tmp_path) + spec["features"] = ["mtp", "multimodal"] + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert task_schema.assembly_units(data) == ["core", "mtp", "multimodal"] + + +def test_missing_file_is_reported_as_such(tmp_path): + with pytest.raises(task_schema.TaskSchemaError, match="task file not found"): + task_schema.load_and_validate_task_yaml(tmp_path / "nope.yaml") + + +def test_non_mapping_yaml_is_rejected(tmp_path): + path = tmp_path / "task.yaml" + path.write_text("- just\n- a\n- list\n", encoding="utf-8") + with pytest.raises(task_schema.TaskSchemaError, match="mapping"): + task_schema.load_and_validate_task_yaml(path) + + +def test_example_task_yaml_matches_the_schema_fields(): + """The shipped example must not drift from the fields the schema reads.""" + from pathlib import Path + + import agent_flow.workflows.trtllm_modeling_v2 as pkg + + example = Path(pkg.__file__).parent / "task.example.yaml" + data = yaml.safe_load(example.read_text(encoding="utf-8")) + for field in (*task_schema.REQUIRED_PATH_FIELDS, task_schema.ARCH_FIELD): + assert field in data, field + assert data["environment"]["kind"] in task_schema.ENVIRONMENT_KINDS + assert "python" in data["environment"] + assert set(data["accuracy"]) <= {"benchmark", "reference_yaml", "tolerance", "dataset_path"} + for feature in data.get("features") or []: + assert feature in task_schema.FEATURE_SKILLS + + +# ------------------------------------------- the substrate gate's proxy checkpoint + + +def test_a_substrate_checkpoint_that_exists_is_accepted(tmp_path): + """The escape hatch for an architecture the stock path has never heard of. + + Optional, because a checkpoint the stock path already supports needs no + proxy — and declared rather than guessed, because which model stands in + for another is a judgement nothing in the schema can make. + """ + proxy = tmp_path / "proxy-ckpt" + proxy.mkdir() + spec = _valid_spec(tmp_path) + spec["environment"]["substrate_checkpoint"] = str(proxy) + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert data["environment"]["substrate_checkpoint"] == str(proxy) + + +def test_a_substrate_checkpoint_that_does_not_exist_is_rejected(tmp_path): + """The gate reaches for it only AFTER the run's own checkpoint has failed. + + That is the longest step in preflight, so a missing proxy turns one + failure into two, with a wait in between. + """ + spec = _valid_spec(tmp_path) + spec["environment"]["substrate_checkpoint"] = "/definitely/not/here" + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert "substrate_checkpoint' does not exist" in str(excinfo.value) + + +@pytest.mark.parametrize("value", ["", " ", 4]) +def test_a_substrate_checkpoint_that_is_not_a_path_is_rejected(tmp_path, value): + spec = _valid_spec(tmp_path) + spec["environment"]["substrate_checkpoint"] = value + with pytest.raises(task_schema.TaskSchemaError, match="substrate_checkpoint"): + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + + +# ------------------------------------------------------ closed key sets + + +def test_a_misspelled_feature_key_does_not_silently_build_core_only(tmp_path): + """The most expensive typo available in this file. + + Accepted, `feature:` instead of `features:` would resolve to `[]`: the + run onboards core, gates it green, and reports success for a feature set + it never built — a green run that delivers the wrong thing. + """ + spec = _valid_spec(tmp_path) | {"feature": ["mtp"]} + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert "unknown key 'feature'" in str(excinfo.value) + assert "did you mean 'features'" in str(excinfo.value) + + +def test_a_misspelled_runtime_block_does_not_silently_drop_max_seq_len(tmp_path): + """A dropped `runtme:` surfaces as `Impossible to fit in any sequence in kvCache`.""" + spec = _valid_spec(tmp_path) | {"runtme": {"max_seq_len": 8192}} + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert "did you mean 'runtime'" in str(excinfo.value) + + +@pytest.mark.parametrize( + "block,typo,expected", + [ + # The misspellings are the fixture, so they are spelled wrong on + # purpose and codespell is told to leave them alone. + ("environment", "pyhton", "environment.python"), # codespell:ignore + ("accuracy", "toleranace", "accuracy.tolerance"), + ("target", "familly", "target.family"), # codespell:ignore + ], +) +def test_nested_blocks_are_closed_too(tmp_path, block, typo, expected): + spec = _valid_spec(tmp_path) + spec.setdefault(block, {})[typo] = "x" + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert f"did you mean '{expected}'" in str(excinfo.value) + + +def test_an_unrecognizable_key_is_still_rejected_without_a_suggestion(tmp_path): + """No near match is not a reason to accept it.""" + spec = _valid_spec(tmp_path) | {"zzzzzz": 1} + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert "unknown key 'zzzzzz'" in str(excinfo.value) + assert "did you mean" not in str(excinfo.value) + + +def test_every_key_the_schema_accepts_actually_validates(tmp_path): + """The key lists and the validators must not drift apart. + + Adding a field to a validator and forgetting to list it here would make + the new field an error; listing it here and forgetting the validator + would make it silently ignored. This asserts the first direction — + a spec using *every* declared key validates. + """ + repo = tmp_path / "repo" + spec = _valid_spec(tmp_path) | { + "features": ["mtp"], + "scope": "text_only", + "source_embargo": [], + "target": {"family": "gpt_oss", "checkpoint_segment": "gpt_oss_120b"}, + "runtime": {"cuda_graph": True, "extra_flags": ["--max_seq_len=8192"]}, + } + spec["accuracy"] = { + "benchmark": "gsm8k", + "reference_yaml": "r.yaml", + "model_key": "openai/gpt-oss-120b", + "tolerance": 5.0, + "dataset_path": str(repo), + "llm_api_options": {"core": str(repo / "o.yaml")}, + "benchmark_args": ["--apply_chat_template"], + } + spool = tmp_path / "spool" + spool.mkdir() + spec["environment"] = { + "kind": "session_spool", + "spool": str(spool), + "python": "/usr/bin/python3", + "env": {"TMPDIR": "/scratch/tmp", "CLAUDE_CODE_TMPDIR": "/scratch/tmp"}, + "partition": "batch", + "account": "acct", + "container_image": "/img.sqsh", + "mounts": ["/lustre:/lustre"], + "time_limit": "04:00:00", + "gpus_per_node": 8, + } + (repo / "o.yaml").write_text("k: v\n", encoding="utf-8") + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert data["scope"] == "text_only" + assert data["runtime"]["cuda_graph"] is True + + +def test_the_declared_key_lists_cover_the_documented_example(tmp_path): + """`task.example.yaml` is a user's only map of the config surface. + + If it used a key the schema rejects, every copy of it would fail; if the + schema accepted a key the example never shows, that field is invisible. + This asserts the first direction — the shipped example validates as-is + against the closed key sets. + """ + from pathlib import Path + + example = Path(task_schema.__file__).parents[1] / "task.example.yaml" + spec = yaml.safe_load(example.read_text(encoding="utf-8")) + errors: list[str] = [] + task_schema.reject_unknown_keys(spec, task_schema.TOP_LEVEL_KEYS, errors) + task_schema.reject_unknown_keys( + spec["environment"], task_schema.ENVIRONMENT_KEYS, errors, where="environment" + ) + task_schema.reject_unknown_keys( + spec["accuracy"], task_schema.ACCURACY_KEYS, errors, where="accuracy" + ) + assert errors == [] + + +def test_the_example_documents_every_key_the_schema_accepts(): + """A field absent from the example does not exist, for a user. + + `task.example.yaml` is the only map of the config surface anyone reads, + and the optional fields (`runtime`, `target`, `scope`, `source_embargo`, + the accuracy protocol options) are the ones a run fails without. + """ + import re + from pathlib import Path + + example = Path(task_schema.__file__).parents[1] / "task.example.yaml" + text = example.read_text(encoding="utf-8") + # Commented-out keys count: an optional field is documented by showing + # it commented with its default, not by setting it. + mentioned = set(re.findall(r"^#?\s*#?\s*([a-z_]+):", text, re.M)) + + for label, keys in ( + ("top level", task_schema.TOP_LEVEL_KEYS), + ("accuracy", task_schema.ACCURACY_KEYS), + ("environment", task_schema.ENVIRONMENT_KEYS), + ("target", task_schema.TARGET_KEYS), + ): + undocumented = sorted(k for k in keys if k not in mentioned) + assert not undocumented, f"{label} keys missing from task.example.yaml: {undocumented}" + + +def test_every_environment_kind_is_shown_in_the_example(): + """Picking a kind is the first decision a user makes.""" + from pathlib import Path + + example = Path(task_schema.__file__).parents[1] / "task.example.yaml" + text = example.read_text(encoding="utf-8") + for kind in task_schema.ENVIRONMENT_KINDS: + assert kind in text, f"environment kind {kind!r} is undocumented" + + +# ------------------------------------------------------------------- CLI + + +def test_every_workflow_knob_is_reachable_from_the_command_line(): + """A constructor parameter with no flag is a feature nobody can use. + + An escape hatch such as `skip_substrate_gate` would otherwise be + reachable only by editing Python. + """ + import inspect + + from agent_flow.workflows.trtllm_modeling_v2.cli import _parse_args + from agent_flow.workflows.trtllm_modeling_v2.run.workflow import TrtllmModelingV2Workflow + + args = vars(_parse_args(["--task", "t.yaml"])) + parameters = inspect.signature(TrtllmModelingV2Workflow.__init__).parameters + # Injection seams for the tests, not operator knobs. + injected = {"self", "workspace", "prompts", "command_runner", "shell_runner", "git_runner"} + unreachable = sorted( + name for name in parameters if name not in injected and name.replace("-", "_") not in args + ) + assert not unreachable, f"no CLI flag reaches: {unreachable}" + + +def test_the_substrate_gate_flag_reaches_the_workflow(): + from agent_flow.workflows.trtllm_modeling_v2.cli import _parse_args + + assert _parse_args(["--task", "t.yaml"]).skip_substrate_gate is False + assert _parse_args(["--task", "t.yaml", "--skip-substrate-gate"]).skip_substrate_gate is True + + +# ------------------------------- removed launch keys + + +@pytest.mark.parametrize("key", ["workloads", "orchestrator", "allocator", "mpi"]) +def test_a_removed_environment_key_is_rejected(tmp_path, key): + spec = _valid_spec(tmp_path) + spec["environment"][key] = {} + with pytest.raises(task_schema.TaskSchemaError, match=key): + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + + +def test_the_session_wait_is_an_environment_key(tmp_path): + spec = _valid_spec(tmp_path) + spec["environment"]["session_wait"] = "2h" + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert data["environment"]["session_wait"] == "2h" + + +# ------------------------------- a spec must work on another machine + + +def test_llm_api_options_written_relative_resolve_next_to_the_spec(tmp_path): + """The per-unit runtime YAMLs are written NEXT TO task.yaml. + + Written absolute, the spec stops working on any other machine, and its + most reusable artifact is unusable anywhere else. + """ + (tmp_path / "core.yaml").write_text("kv_cache_config: {}\n", encoding="utf-8") + spec = _valid_spec(tmp_path) + spec.setdefault("accuracy", {})["llm_api_options"] = "core.yaml" + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert data["accuracy"]["llm_api_options"] == str((tmp_path / "core.yaml").resolve()) + + +def test_a_relative_option_path_in_a_subdirectory_resolves_too(tmp_path): + """The usual layout: a `scripts/` directory beside the spec.""" + scripts = tmp_path / "scripts" + scripts.mkdir() + (scripts / "core_llm_api.yaml").write_text("kv_cache_config: {}\n", encoding="utf-8") + spec = _valid_spec(tmp_path) + spec.setdefault("accuracy", {})["llm_api_options"] = "scripts/core_llm_api.yaml" + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert data["accuracy"]["llm_api_options"] == str((scripts / "core_llm_api.yaml").resolve()) + + +def test_the_per_unit_mapping_form_resolves_too(tmp_path): + (tmp_path / "core.yaml").write_text("kv_cache_config: {}\n", encoding="utf-8") + (tmp_path / "mtp.yaml").write_text("kv_cache_config: {}\n", encoding="utf-8") + spec = _valid_spec(tmp_path) + spec.setdefault("accuracy", {})["llm_api_options"] = { + "default": "core.yaml", + "mtp": str(tmp_path / "mtp.yaml"), + } + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + options = data["accuracy"]["llm_api_options"] + assert options["default"] == str((tmp_path / "core.yaml").resolve()) + # Absolute paths are untouched: a site that wants one keeps it. + assert options["mtp"] == str(tmp_path / "mtp.yaml") + + +def test_a_relative_option_path_is_checked_for_existence_after_resolution(tmp_path): + """Resolution runs BEFORE validation, so the message names the real path. + + Resolving afterwards would make a perfectly good relative path fail as + "does not exist", which is the opposite of the point. + """ + spec = _valid_spec(tmp_path) + spec.setdefault("accuracy", {})["llm_api_options"] = "absent.yaml" + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert str((tmp_path / "absent.yaml").resolve()) in str(excinfo.value) + + +def test_resolving_an_already_resolved_spec_is_idempotent(tmp_path): + """The workflow copies task.yaml into the workspace and re-reads it.""" + (tmp_path / "core.yaml").write_text("kv_cache_config: {}\n", encoding="utf-8") + spec = _valid_spec(tmp_path) + spec.setdefault("accuracy", {})["llm_api_options"] = "core.yaml" + once = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + spec["accuracy"]["llm_api_options"] = once["accuracy"]["llm_api_options"] + twice = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert twice["accuracy"]["llm_api_options"] == once["accuracy"]["llm_api_options"] + + +def test_machine_specific_paths_are_deliberately_left_alone(tmp_path): + """checkpoint_path and friends differ per machine and are read elsewhere. + + Making them spec-relative would move the portability problem, not solve it. + """ + import inspect + + spec = _valid_spec(tmp_path) + expected = spec["checkpoint_path"] + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + src = inspect.getsource(task_schema._resolve_spec_relative_paths) + assert "checkpoint_path" in src and "Deliberately NOT applied" in src + assert data["checkpoint_path"] == expected + + +def test_a_spec_with_no_accuracy_block_survives_resolution(tmp_path): + """Resolution runs before validation, so it meets malformed specs first.""" + data: dict = {} + task_schema._resolve_spec_relative_paths(data, tmp_path) + assert data == {} + data = {"accuracy": "not-a-mapping"} + task_schema._resolve_spec_relative_paths(data, tmp_path) + assert data == {"accuracy": "not-a-mapping"} + + +# ---------------------------- the speculator, replay and kv-quant feature kinds + + +def test_the_new_feature_kinds_map_to_their_skills(): + skills = task_schema.FEATURE_SKILLS + assert skills["dspark"] == "trtllm-modeling-v2-speculator" + assert skills["decoder_replay"] == "trtllm-modeling-v2-bounded-replay" + assert skills["encoder_replay"] == "trtllm-modeling-v2-bounded-replay" + assert skills["kv_quant"] == "trtllm-modeling-v2-kv-quant" + + +def test_an_unknown_feature_lists_the_new_kinds(tmp_path): + spec = _valid_spec(tmp_path) + spec["features"] = ["swa_bounded_replay"] + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + message = str(excinfo.value) + for kind in ("dspark", "decoder_replay", "encoder_replay", "kv_quant"): + assert kind in message + + +def test_mtp_and_dspark_are_refused_together(tmp_path): + """Both read `mtp.*`; the namespace holds one drafter, not two.""" + spec = _valid_spec(tmp_path) + spec["features"] = ["mtp", "dspark"] + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert "spelled `dspark`" in str(excinfo.value) + + +def test_dspark_alone_validates(tmp_path): + spec = _valid_spec(tmp_path) + spec["features"] = ["dspark"] + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert task_schema.assembly_units(data) == ["core", "dspark"] + + +_DECODER = { + "contrast_env": {"TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY": "0"}, + "window": 128, + "prompt_tokens": 4096, +} +_ENCODER = { + "contrast_env": {"TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY": "0"}, + "window": 128, + "prefix_tokens": 4096, + "reuse_alignment": 2, +} +_KV = { + "contrast_env": {"TRTLLM_DSV41_KV_STORAGE": "bf16"}, + "expected_bytes_per_token": 890, +} + + +def _with_witness(tmp_path, features, witness): + spec = _valid_spec(tmp_path) + spec["features"] = features + spec["accuracy"]["witness"] = witness + return _write(tmp_path, spec) + + +def test_the_replay_and_kv_quant_witness_blocks_validate(tmp_path): + data = task_schema.load_and_validate_task_yaml( + _with_witness( + tmp_path, + ["kv_quant", "decoder_replay", "encoder_replay"], + {"decoder_replay": _DECODER, "encoder_replay": _ENCODER, "kv_quant": _KV}, + ) + ) + # `exact` changes what the unit claims, so the resolved spec states it. + assert data["accuracy"]["witness"]["kv_quant"]["exact"] is True + + +@pytest.mark.parametrize("feature", ["decoder_replay", "encoder_replay", "kv_quant"]) +def test_a_declared_witness_is_required_at_spec_time(tmp_path, feature): + """Declared before the run it judges, which is only certain here.""" + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_with_witness(tmp_path, [feature], {})) + assert f"accuracy.witness.{feature}" in str(excinfo.value) + + +def test_a_replay_workload_inside_the_window_is_refused(tmp_path): + """At or under the window, replay on and off do the same work.""" + block = _DECODER | {"prompt_tokens": 128} + with pytest.raises(task_schema.TaskSchemaError, match="does not exceed window=128"): + task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["decoder_replay"], {"decoder_replay": block}) + ) + + +def test_a_replay_workload_past_the_argv_cap_is_refused(tmp_path): + block = _ENCODER | {"prefix_tokens": task_schema.MAX_WITNESS_PROMPT_WORDS + 1} + with pytest.raises(task_schema.TaskSchemaError, match="128 KiB"): + task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["encoder_replay"], {"encoder_replay": block}) + ) + + +@pytest.mark.parametrize("missing", ["contrast_env", "window", "prompt_tokens"]) +def test_every_replay_parameter_is_required(tmp_path, missing): + block = {k: v for k, v in _DECODER.items() if k != missing} + with pytest.raises(task_schema.TaskSchemaError, match=missing): + task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["decoder_replay"], {"decoder_replay": block}) + ) + + +def test_witness_blocks_are_closed_per_feature(tmp_path): + """`expected_mib` under a replay feature reads as a check nobody makes.""" + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml( + _with_witness( + tmp_path, + ["decoder_replay", "host_offload"], + { + "decoder_replay": _DECODER | {"expected_mib": 1, "prefix_tokens": 4096}, + "host_offload": {"contrast_env": {"X": "0"}, "window": 128}, + }, + ) + ) + message = str(excinfo.value) + assert "accuracy.witness.decoder_replay.expected_mib" in message + assert "accuracy.witness.decoder_replay.prefix_tokens" in message + assert "accuracy.witness.host_offload.window" in message + + +def test_an_exact_kv_quant_needs_a_contrast(tmp_path): + """Without one, token identity compares native storage with itself.""" + exact = {"expected_bytes_per_token": 890} + with pytest.raises(task_schema.TaskSchemaError, match="contrast_env"): + task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["kv_quant"], {"kv_quant": exact}) + ) + + +def test_a_lossy_kv_quant_does_not_need_a_contrast(tmp_path): + """`exact: false` drops the identity leg, and with it the contrast.""" + lossy = {"expected_bytes_per_token": 890, "exact": False} + data = task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["kv_quant"], {"kv_quant": lossy}) + ) + assert data["accuracy"]["witness"]["kv_quant"]["exact"] is False + + +@pytest.mark.parametrize( + "bad", + [ + {"expected_bytes_per_token": 0}, + {"expected_bytes_per_token": "890"}, + {"bytes_tolerance": 1.5}, + {"exact": "yes"}, + ], +) +def test_kv_quant_values_are_checked(tmp_path, bad): + with pytest.raises(task_schema.TaskSchemaError): + task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["kv_quant"], {"kv_quant": _KV | bad}) + ) + + +# ------------------------------------------------- the equivalence witness + +_DISAGG = { + "contrast_env": {"TRTLLM_DISAGG": "1"}, + "contrast_command": "/onboard/scripts/disagg_leg_b.py", +} + + +def test_a_disagg_witness_without_a_launcher_of_its_own_runs_the_flows(tmp_path): + """Leg B is the package's launcher unless the run names one; nothing else is owed.""" + data = task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["disagg"], {"disagg": {"max_batch_size": 1}}) + ) + assert "contrast_command" not in data["accuracy"]["witness"]["disagg"] + + +@pytest.mark.parametrize( + "block,message", + [ + ({"contrast_command": "/onboard/leg_b.py"}, "contrast_env' is required"), + ({"contrast_command": " ", "contrast_env": {"X": "1"}}, "or left out for the flow's"), + ({"contrast_env": {}}, "contrast_env' is required and must be a non-empty mapping"), + ({"cache_transceiver_config": "NIXL"}, "mapping of CacheTransceiverConfig fields"), + ({"cache_transceiver_config": {}}, "mapping of CacheTransceiverConfig fields"), + ], +) +def test_a_disagg_witness_block_is_checked(tmp_path, block, message): + with pytest.raises(task_schema.TaskSchemaError, match=message): + task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["disagg"], {"disagg": block}) + ) + + +def test_the_flows_leg_b_takes_the_cache_transceiver_it_runs(tmp_path): + block = {task_schema.WITNESS_TRANSCEIVER_KEY: {"backend": "NIXL"}} + data = task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["disagg"], {"disagg": block}) + ) + assert data["accuracy"]["witness"]["disagg"]["cache_transceiver_config"] == {"backend": "NIXL"} + + +def test_a_launcher_beside_the_spec_is_found_from_anywhere(tmp_path): + """A word of the command naming a file beside the spec is made absolute. + + A check runs inside a snapshot of the tree, where the spec's file is not; + a word naming no file beside it (a flag, a path of the tree) is left as written. + """ + (tmp_path / "leg_b.py").write_text("print('leg B')\n", encoding="utf-8") + command = "python leg_b.py --anchors=table.yaml examples/disaggregated/launch.sh" + data = task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["disagg"], {"disagg": _DISAGG | {"contrast_command": command}}) + ) + resolved = data["accuracy"]["witness"]["disagg"]["contrast_command"] + assert shlex.split(resolved) == [ + "python", + str((tmp_path / "leg_b.py").resolve()), + "--anchors=table.yaml", + "examples/disaggregated/launch.sh", + ] + untouched = "bash 'examples/disaggregated/launch.sh' --n=2" + elsewhere = tmp_path / "elsewhere" + elsewhere.mkdir() + data = task_schema.load_and_validate_task_yaml( + _with_witness(elsewhere, ["disagg"], {"disagg": _DISAGG | {"contrast_command": untouched}}) + ) + assert data["accuracy"]["witness"]["disagg"]["contrast_command"] == untouched + + +def test_a_disagg_witness_may_declare_its_prompt_lengths_and_timeout(tmp_path): + block = _DISAGG | {"prompt_tokens": [131, 2050, 2051, 4097], "timeout": 2700} + data = task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["disagg"], {"disagg": block}) + ) + assert data["accuracy"]["witness"]["disagg"]["prompt_tokens"] == [131, 2050, 2051, 4097] + + +@pytest.mark.parametrize( + "bad,message", + [ + ({"prompt_tokens": []}, "non-empty list"), + ({"prompt_tokens": 2051}, "non-empty list"), + ({"prompt_tokens": [0]}, "integers >= 1"), + ({"prompt_tokens": [task_schema.MAX_WITNESS_PROMPT_WORDS + 1]}, "128 KiB"), + ({"timeout": 0}, "timeout"), + ({"timeout": "45m"}, "timeout"), + ({"expected_mib": 1}, "expected_mib"), + ], +) +def test_disagg_witness_values_are_checked(tmp_path, bad, message): + with pytest.raises(task_schema.TaskSchemaError, match=message): + task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["disagg"], {"disagg": _DISAGG | bad}) + ) + + +def test_every_witness_block_takes_a_timeout(tmp_path): + """Seconds per leg; the one-engine boot budget is only a default.""" + data = task_schema.load_and_validate_task_yaml( + _with_witness( + tmp_path, + ["host_offload", "kv_quant", "decoder_replay", "encoder_replay"], + { + "host_offload": {"contrast_env": {"X": "0"}, "expected_mib": 1, "timeout": 900}, + "kv_quant": _KV | {"timeout": 900}, + "decoder_replay": _DECODER | {"timeout": 900}, + "encoder_replay": _ENCODER | {"timeout": 900}, + }, + ) + ) + assert all(block["timeout"] == 900 for block in data["accuracy"]["witness"].values()) + + +def test_prompt_lengths_are_not_a_residency_key(tmp_path): + """Closed per feature: nothing would read them.""" + block = {"contrast_env": {"X": "0"}, "prompt_tokens": [2051]} + with pytest.raises(task_schema.TaskSchemaError, match="host_offload.prompt_tokens"): + task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["host_offload"], {"host_offload": block}) + ) + + +_REPLAYS_OFF = { + "TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY": "0", + "TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY": "0", +} + + +def test_a_disagg_witness_may_pin_an_environment_on_both_legs(tmp_path): + """Paths leg B's engines guard off under disagg are pinned off in leg A too.""" + data = task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["disagg"], {"disagg": _DISAGG | {"legs_env": _REPLAYS_OFF}}) + ) + assert data["accuracy"]["witness"]["disagg"]["legs_env"] == _REPLAYS_OFF + + +@pytest.mark.parametrize( + "legs_env,message", + [ + ({}, "non-empty mapping"), + (["TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY=0"], "non-empty mapping"), + ({"1_REPLAY": "0"}, "not a shell variable name"), + ({"A-B": "0"}, "not a shell variable name"), + ({"REPLAY": ["0"]}, "scalar"), + # Set one way on leg A and another on leg B: that is the contrast. + ({"TRTLLM_DISAGG": "0"}, "is the contrast"), + ], +) +def test_legs_env_values_are_checked(tmp_path, legs_env, message): + with pytest.raises(task_schema.TaskSchemaError, match=message): + task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["disagg"], {"disagg": _DISAGG | {"legs_env": legs_env}}) + ) + + +@pytest.mark.parametrize( + "feature,block", + [ + ("host_offload", {"contrast_env": {"X": "0"}}), + ("kv_quant", _KV), + ("decoder_replay", _DECODER), + ], +) +def test_legs_env_is_an_equivalence_key_only(tmp_path, feature, block): + """The other witnesses run one command on both legs; only `contrast_env` differs.""" + with pytest.raises(task_schema.TaskSchemaError, match=f"{feature}.legs_env"): + task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, [feature], {feature: block | {"legs_env": _REPLAYS_OFF}}) + ) + + +# -------------------------------------------- the acceptance gate's workload + + +def _with_acceptance(tmp_path, acceptance) -> str: + spec = _valid_spec(tmp_path) + spec["accuracy"]["acceptance"] = acceptance + return _write(tmp_path, spec) + + +def test_the_per_unit_floors_alone_still_validate(tmp_path): + """Per-unit floors without a declared workload validate as they are.""" + acceptance = {"dspark": {"min_absolute": 2.0}, "multimodal": {"min_absolute": 2.0}, "mtp": 1.5} + data = task_schema.load_and_validate_task_yaml(_with_acceptance(tmp_path, acceptance)) + assert data["accuracy"]["acceptance"] == acceptance + + +def test_a_declared_workload_validates_run_wide_and_per_unit(tmp_path): + acceptance = { + "prompts": ["Explain, step by step, how a heat pump works."], + "max_tokens": 256, + "min_draft_iterations": 64, + "dspark": {"min_absolute": 2.0}, + "multimodal": {"min_absolute": 2.0, "prompts": ["Describe the image."]}, + } + data = task_schema.load_and_validate_task_yaml(_with_acceptance(tmp_path, acceptance)) + assert data["accuracy"]["acceptance"]["max_tokens"] == 256 + + +def test_a_relative_prompts_file_resolves_next_to_the_spec(tmp_path): + (tmp_path / "scripts").mkdir() + (tmp_path / "scripts" / "long.json").write_text('["a long question"]', encoding="utf-8") + (tmp_path / "boot.yaml").write_text("prompts: [from boot]\n", encoding="utf-8") + acceptance = { + "prompts_file": "scripts/long.json", + "dspark": {"min_absolute": 2.0, "prompts_file": "boot.yaml"}, + } + data = task_schema.load_and_validate_task_yaml(_with_acceptance(tmp_path, acceptance)) + block = data["accuracy"]["acceptance"] + assert block["prompts_file"] == str((tmp_path / "scripts" / "long.json").resolve()) + assert block["dspark"]["prompts_file"] == str((tmp_path / "boot.yaml").resolve()) + + +@pytest.mark.parametrize( + "acceptance,message", + [ + ({"prompts": ["a"], "prompts_file": "p.json"}, "both 'prompts' and 'prompts_file'"), + ({"prompts": []}, "acceptance.prompts' must be a non-empty list"), + ({"prompts": ["ok", 3]}, "acceptance.prompts' must be a non-empty list"), + ({"prompts": "one string"}, "acceptance.prompts' must be a non-empty list"), + ({"prompts_file": "missing.json"}, "prompts_file' does not exist"), + ({"max_tokens": 0}, "acceptance.max_tokens' must be an integer >= 1"), + ({"min_draft_iterations": True}, "min_draft_iterations' must be an integer >= 1"), + ({"dspark": {"max_tokens": "256"}}, "dspark.max_tokens' must be an integer >= 1"), + ( + {"dspark": {"min_absolut": 2.0}}, + "did you mean 'accuracy.acceptance.dspark.min_absolute'", + ), + ({"dspark": {"min_absolute": 0.5}}, "dspark.min_absolute' must be a number >= 1.0"), + ({"mtp": "high"}, "acceptance.mtp' must be a number >= 1.0"), + (["dspark"], "acceptance' must be a mapping"), + ], +) +def test_a_malformed_acceptance_block_is_rejected(tmp_path, acceptance, message): + with pytest.raises(task_schema.TaskSchemaError, match=message): + task_schema.load_and_validate_task_yaml(_with_acceptance(tmp_path, acceptance)) + + +def test_a_prompts_file_that_holds_no_prompts_is_rejected(tmp_path): + (tmp_path / "empty.yaml").write_text("keywords: [x]\n", encoding="utf-8") + with pytest.raises(task_schema.TaskSchemaError, match="non-empty list of strings"): + task_schema.load_and_validate_task_yaml( + _with_acceptance(tmp_path, {"prompts_file": "empty.yaml"}) + ) + + +# ----------------------------------- one schedule for token-comparing legs + + +def test_token_comparing_witnesses_take_a_batch_limit(tmp_path): + """Batch composition moves tokens; `max_batch_size: 1` gives both legs one schedule.""" + data = task_schema.load_and_validate_task_yaml( + _with_witness( + tmp_path, + ["host_offload", "kv_quant", "disagg"], + { + "host_offload": { + "contrast_env": {"X": "0"}, + "expected_mib": 1, + "max_batch_size": 1, + }, + "kv_quant": _KV | {"max_batch_size": 1}, + "disagg": _DISAGG | {"max_batch_size": 1}, + }, + ) + ) + witness = data["accuracy"]["witness"] + assert {unit: block["max_batch_size"] for unit, block in witness.items()} == { + "host_offload": 1, + "kv_quant": 1, + "disagg": 1, + } + + +def test_a_block_without_a_batch_limit_still_validates(tmp_path): + """The key is optional, and an absent one gets no default.""" + data = task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["kv_quant", "disagg"], {"kv_quant": _KV, "disagg": _DISAGG}) + ) + assert all("max_batch_size" not in b for b in data["accuracy"]["witness"].values()) + + +def test_a_replay_witness_takes_no_batch_limit(tmp_path): + """It counts work, not tokens; the key would read as a check nobody makes.""" + with pytest.raises(task_schema.TaskSchemaError, match="decoder_replay.max_batch_size"): + task_schema.load_and_validate_task_yaml( + _with_witness( + tmp_path, ["decoder_replay"], {"decoder_replay": _DECODER | {"max_batch_size": 1}} + ) + ) + + +@pytest.mark.parametrize("bad", [0, -1, "1", 1.0, True]) +@pytest.mark.parametrize("unit", ["kv_quant", "disagg", "host_offload"]) +def test_a_batch_limit_is_a_positive_integer(tmp_path, unit, bad): + base = { + "kv_quant": _KV, + "disagg": _DISAGG, + "host_offload": {"contrast_env": {"X": "0"}, "expected_mib": 1}, + }[unit] + with pytest.raises(task_schema.TaskSchemaError, match=f"{unit}.max_batch_size"): + task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, [unit], {unit: base | {"max_batch_size": bad}}) + ) + + +def test_notes_are_an_optional_non_empty_string(tmp_path): + spec = _valid_spec(tmp_path) + spec["notes"] = "Reuse the in-tree NVFP4 MoE kernels." + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert data["notes"] == "Reuse the in-tree NVFP4 MoE kernels." + spec["notes"] = " " + with pytest.raises(task_schema.TaskSchemaError, match="'notes' must be a non-empty string"): + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) diff --git a/agent-flow/tests/workflows/trtllm_modeling_v2/test_variants.py b/agent-flow/tests/workflows/trtllm_modeling_v2/test_variants.py new file mode 100644 index 000000000000..358ece990033 --- /dev/null +++ b/agent-flow/tests/workflows/trtllm_modeling_v2/test_variants.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Configuration variants, and members of different widths under a scheduler.""" + +from __future__ import annotations + +import pytest +import yaml + +from agent_flow.workflows.trtllm_modeling_v2.checks import gate_cli +from agent_flow.workflows.trtllm_modeling_v2.checks import gates as gates_mod +from agent_flow.workflows.trtllm_modeling_v2.run import contract as contract_mod +from agent_flow.workflows.trtllm_modeling_v2.spec import task_schema +from agent_flow.workflows.trtllm_modeling_v2.spec import variants as V + + +def _spec(tmp_path, **over) -> dict: + checkpoint, repo = tmp_path / "ckpt", tmp_path / "repo" + checkpoint.mkdir(exist_ok=True) + repo.mkdir(exist_ok=True) + (tmp_path / "spool").mkdir(exist_ok=True) + (tmp_path / "mtp.yaml").write_text( + yaml.safe_dump({"speculative_config": {"decoding_type": "MTP", "max_draft_len": 3}}) + ) + spec = { + "checkpoint_path": str(checkpoint), + "arch": "sm_103", + "parallel": "unified", + "members": {"default": ["tp1", "tep2", "tep4"], "mtp": ["tep4"]}, + "trtllm_repo_path": str(repo), + "features": ["mtp"], + "environment": { + "kind": "session_spool", + "spool": str(tmp_path / "spool"), + "python": "/usr/bin/python3", + "env": {"TMPDIR": "/scratch/tmp"}, + }, + "runtime": {"cuda_graph": True, "extra_flags": ["--max_seq_len=9216"]}, + "accuracy": { + "benchmark": "gsm8k", + "reference_yaml": "x.yaml", + "llm_api_options": {"mtp": str(tmp_path / "mtp.yaml")}, + "discriminators": {"default": {}, "mtp": {"spec_dec_algo": "MTP"}}, + }, + "variants": { + "kv_fp8": { + "members": ["tp1"], + "runtime": {"extra_flags": ["--kv_cache_dtype=fp8"]}, + "llm_api_options": {"kv_cache_config": {"dtype": "fp8"}}, + "accuracy": {"discriminators": {"kv_cache_quant_algo": "FP8"}}, + }, + "mtp7": { + "unit": "mtp", + "gates": ["boot", "accuracy", "acceptance"], + "llm_api_options": {"speculative_config": {"max_draft_len": 7}}, + "env": {"X": "1"}, + }, + }, + } + spec.update(over) + return spec + + +def _load(tmp_path, spec) -> dict: + path = tmp_path / "task.yaml" + path.write_text(yaml.safe_dump(spec), encoding="utf-8") + return task_schema.load_and_validate_task_yaml(path) + + +def test_members_of_different_widths_are_legal_under_a_scheduler_only(tmp_path): + data = _load(tmp_path, _spec(tmp_path)) + assert task_schema.widest_world_size(data, 4) == 4 + local = _spec(tmp_path) + local["environment"] = {**local["environment"], "kind": "local"} + local["environment"].pop("spool") + with pytest.raises(task_schema.TaskSchemaError, match="mixes world sizes"): + _load(tmp_path, local) + + +def test_every_variant_problem_is_named(tmp_path): + spec = _spec(tmp_path) + spec["variants"] = { + "core": {"runtime": {"cuda_graph": False}}, + "Bad-Name": {"runtime": {"cuda_graph": False}}, + "nothing": {"note": "changes nothing"}, + "orphan": {"unit": "eagle3", "runtime": {"cuda_graph": False}}, + "spec": {"gates": ["acceptance"], "llm_api_options": {"a": 1}}, + "typo": {"runtme": {}, "env": {"A": "1"}, "members": ["tq4"]}, + } + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + _load(tmp_path, spec) + text = str(excinfo.value) + assert "'core' is a unit name" in text + assert "variants.Bad-Name" in text + assert "'variants.nothing' changes nothing" in text + assert "not one of the units" in text + assert "implements no speculator" in text + assert "runtme" in text and "variants.typo.members" in text + + +def test_variant_checks_are_named_by_the_variant_at_its_own_members(tmp_path): + data = _load(tmp_path, _spec(tmp_path)) + contract = contract_mod.compile_contract(data, world_size=4) + proof = [c.key for c in contract.by_tier(contract_mod.TIER_PROOF)] + assert "kv_fp8:boot@tp1" in proof and "kv_fp8:accuracy[gsm8k]@tp1" in proof + # A storage variant boots at every member of its unit; its accuracy stays where it names. + tep = {k for k in proof if k.startswith("kv_fp8:") and "@tep" in k} + assert tep == {"kv_fp8:boot@tep2", "kv_fp8:boot@tep4"} + # A speculation variant runs every gate only where it names. + assert {"mtp7:boot@tep4", "mtp7:accuracy[gsm8k]@tep4", "mtp7:acceptance@tep4"} <= set(proof) + assert not [k for k in proof if k.startswith("mtp7:") and "@tp1" in k] + + +def test_a_variant_runs_under_its_units_spec_with_its_overlay(tmp_path): + data = _load(tmp_path, _spec(tmp_path)) + spec = V.variant_spec(data, "mtp7", world_size=4, options_dir=tmp_path / "ws" / "variants") + accuracy = spec["accuracy"] + options = yaml.safe_load(open(gates_mod.resolve_llm_api_options(accuracy, "mtp7"))) + assert options["speculative_config"] == {"decoding_type": "MTP", "max_draft_len": 7} + assert gates_mod.resolve_max_draft_len(accuracy, "mtp7") == 7 + assert gates_mod.resolve_max_draft_len(accuracy, "mtp") == 3 # the unit is untouched + assert gates_mod.resolve_discriminators(accuracy, "mtp7") == {"spec_dec_algo": "MTP"} + assert task_schema.members_for(spec, "mtp7", 4) == ["tep4"] + assert spec["environment"]["env"]["X"] == "1" + assert gate_cli.features_for(spec, "mtp7") == ["mtp"] + + +def test_a_variant_appends_its_flags_and_overrides_its_anchor_row(tmp_path): + data = _load(tmp_path, _spec(tmp_path)) + spec = V.variant_spec(data, "kv_fp8", world_size=4, options_dir=tmp_path / "v") + runtime = gates_mod.resolve_runtime(spec, "kv_fp8") + assert runtime["cuda_graph"] is True + assert runtime["extra_flags"] == ["--max_seq_len=9216", "--kv_cache_dtype=fp8"] + assert gates_mod.resolve_runtime(spec, "core")["extra_flags"] == ["--max_seq_len=9216"] + accuracy = spec["accuracy"] + assert gates_mod.resolve_discriminators(accuracy, "kv_fp8") == {"kv_cache_quant_algo": "FP8"} + assert gates_mod.resolve_discriminators(accuracy, "core") == {} + assert yaml.safe_load(open(gates_mod.resolve_llm_api_options(accuracy, "kv_fp8"))) == { + "kv_cache_config": {"dtype": "fp8"} + } + assert task_schema.members_for(spec, "kv_fp8", 4) == ["tp1"] + assert task_schema.members_for(spec, "core", 4) == ["tp1", "tep2", "tep4"] + assert gate_cli.features_for(spec, "kv_fp8") == [] + + +def test_the_widest_member_counts_variants_and_per_unit_lists(tmp_path): + data = {"parallel": "unified", "members": {"default": ["tp1"], "mtp": ["tep2"]}} + assert task_schema.widest_world_size(data, 4) == 2 + data["variants"] = {"v": {"members": ["dep4"]}} + assert task_schema.widest_world_size(data, 8) == 4 + assert task_schema.widest_world_size({"parallel": "tp2"}, 8) == 2 + assert task_schema.widest_world_size({"parallel": "unified"}, 8) == 8 + + +def test_each_member_is_launched_and_scheduled_at_its_own_width(tmp_path): + from agent_flow.workflows.trtllm_modeling_v2.infra.environment import RunEnvironment + from agent_flow.workflows.trtllm_modeling_v2.run.workflow import TrtllmModelingV2Workflow + + workflow = TrtllmModelingV2Workflow(workspace=tmp_path / "ws") + workflow.environment = RunEnvironment( + kind="session_spool", python="python3", world_size=4, spool=str(tmp_path / "spool") + ) + workflow.task_data = {"trtllm_repo_path": str(tmp_path)} + assert workflow._member_env("tp1").world_size == 1 + assert "mpirun -n 2" in workflow._member_env("tep2").command("python x.py") + assert "mpirun" not in workflow._member_env("tp1").command("python x.py") + assert workflow._member_runner("tp1") is workflow._member_runner("tp1") + workflow._member_runner("tep2") + workflow._member_runner("tep2", "all") + assert sorted(workflow._member_runners) == ["1", "2", "all"] From 91c88480abb9adc6bb03bd9273462805bb863306 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Wed, 21 Oct 2026 02:26:01 -0700 Subject: [PATCH 03/13] [TRTLLM-16460][feat] trtllm_modeling_v2: compute sessions, the spool scheduler and command execution Where GPU work runs. A compute session is a Slurm allocation that runs a spool scheduler; jobs are files in its spool: - infra/substrate/scheduler.py: a priority queue over device slots with wall and memory limits and a scrubbed MPI environment; a device is quarantined until its memory is released; a job lost with its allocation runs again or is closed out node_lost; a native rebuild is held until the jobs that map the libraries are done (rebuild.py). - spoolctl.py: submit, launch, wait, status and cancel, also runnable by path from a shell without the CLI. placement.py routes a job to a session that can run it now and lets it borrow idle devices. - sessions.py and cli.py: sessions declared in task.yaml, their entry scripts and their hand-over; watchdog.py returns when one needs a human. - provenance.py and the resolve hook record which checkout every process of a job imported, so a job with mixed checkouts is refused. - infra/execution/: one command or one agent turn as a request and an outcome, run here or through a spool, and how many checks and tasks the sessions' free devices admit. - Within a priority, pending jobs run in the order they were queued (mtime, then a queue stamp), and a native rebuild also holds jobs whose native build is linked into the rebuilt checkout. - infra/execution/agent_gate.py: one cap on the agents working at once, slots granted by priority (planner, critical, required, side, background). Every role turn takes a slot, one stays free for the planner, and background turns take only idle slots. - The planner's reserved slot takes nothing from the tasks; critic and improver turns use only the slots the foreground work leaves idle (or the next one free after 20 minutes). - A queued gate moves to a session that would start it now (an atomic rename, so it never runs twice), and a dev session lends its free devices while nothing of the campaign's own runs there. - A session near its end judges a job by how long such jobs ran (twice the 90th-percentile of its spool's history), not by its worst-case wall limit. - incremental_build.sh builds from a venv of its own that sees the run's packages, and fails if the run's venv changed under it. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../infra/execution/__init__.py | 233 ++ .../infra/execution/agent_gate.py | 164 + .../infra/execution/capacity.py | 259 ++ .../infra/execution/commands.py | 130 + .../infra/execution/spool.py | 403 +++ .../infra/execution/turns.py | 539 ++++ .../infra/scripts/gpu_slot.sh | 51 + .../infra/scripts/incremental_build.sh | 140 + .../infra/scripts/mn/README.md | 66 + .../infra/scripts/mn/mn_local_mpirun.sh | 28 + .../infra/scripts/mn/mn_mpirun.sh | 136 + .../infra/scripts/mn/mn_node_agent.sh | 42 + .../trtllm_modeling_v2/infra/spool_watcher.sh | 42 + .../infra/substrate/__init__.py | 26 + .../trtllm_modeling_v2/infra/substrate/cli.py | 217 ++ .../infra/substrate/placement.py | 482 +++ .../infra/substrate/provenance.py | 119 + .../infra/substrate/rebuild.py | 255 ++ .../substrate/resolvehook/sitecustomize.py | 131 + .../infra/substrate/scheduler.py | 2741 +++++++++++++++++ .../infra/substrate/sessions.py | 799 +++++ .../infra/substrate/spoolctl.py | 696 +++++ .../infra/substrate/watchdog.py | 222 ++ .../trtllm_modeling_v2/test_agent_gate.py | 135 + .../trtllm_modeling_v2/test_check_capacity.py | 318 ++ .../trtllm_modeling_v2/test_commands.py | 238 ++ .../trtllm_modeling_v2/test_handover_state.py | 362 +++ .../test_inprocess_turns.py | 507 +++ .../trtllm_modeling_v2/test_session_ack.py | 228 ++ .../test_spool_device_memory.py | 322 ++ .../trtllm_modeling_v2/test_spool_hangs.py | 578 ++++ .../test_spool_placement.py | 601 ++++ .../trtllm_modeling_v2/test_spool_rebuild.py | 380 +++ .../trtllm_modeling_v2/test_spool_requeue.py | 469 +++ .../test_spool_scheduler.py | 841 +++++ .../trtllm_modeling_v2/test_spool_sessions.py | 424 +++ .../trtllm_modeling_v2/test_spool_tree.py | 203 ++ .../trtllm_modeling_v2/test_turns.py | 391 +++ 38 files changed, 13918 insertions(+) create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/__init__.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/agent_gate.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/capacity.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/commands.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/spool.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/turns.py create mode 100755 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/gpu_slot.sh create mode 100755 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/incremental_build.sh create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/mn/README.md create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/mn/mn_local_mpirun.sh create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/mn/mn_mpirun.sh create mode 100755 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/mn/mn_node_agent.sh create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/spool_watcher.sh create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/substrate/__init__.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/substrate/cli.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/substrate/placement.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/substrate/provenance.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/substrate/rebuild.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/substrate/resolvehook/sitecustomize.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/substrate/scheduler.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/substrate/sessions.py create mode 100644 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/substrate/spoolctl.py create mode 100755 agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/substrate/watchdog.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_agent_gate.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_check_capacity.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_commands.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_handover_state.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_inprocess_turns.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_session_ack.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_spool_device_memory.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_spool_hangs.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_spool_placement.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_spool_rebuild.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_spool_requeue.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_spool_scheduler.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_spool_sessions.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_spool_tree.py create mode 100644 agent-flow/tests/workflows/trtllm_modeling_v2/test_turns.py diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/__init__.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/__init__.py new file mode 100644 index 000000000000..5f22c174e3fa --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/__init__.py @@ -0,0 +1,233 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Where this run's commands and agent turns go, decided once from the resolved environment. + +`environment.py` answers *what the command line is*: the environment, the +launcher wrapper, the MPI scrub. This package answers *who runs it*, and it is +the only place in the workflow that maps an ``environment.kind`` onto a +mechanism: + +- ``local``: a subprocess here; +- ``session_spool``: a file in a spool a long-lived allocation drains + (:mod:`.spool`). + +Agent turns run in process (:class:`InProcessTurnExecutor`) unless the run was +given another transport (`transport.ExternalAgentTransport`). Either bounds a +turn by its role's timeout; one in process leaves its prompt, streamed activity +and answer under ``/turns/`` (:mod:`.turns`). + +How many checks run at once follows the devices those spools have free +(:func:`check_capacity`, :class:`CheckAdmission`). +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Mapping + +from ..environment import RunEnvironment +from ..substrate import placement +from .capacity import ( + DEFAULT_CHECK_CAP, + DEFAULT_CHECK_FALLBACK, + DEFAULT_CHECK_FLOOR, + CheckAdmission, + more_checks, +) +from .commands import ( + NOT_FOUND_RETURNCODE, + TIMEOUT_RETURNCODE, + CommandExecutor, + CommandResult, + CommandSpec, + LocalCommandExecutor, +) +from .spool import SpoolCommandExecutor +from .turns import ( + DEFAULT_ROLE_TURN_TIMEOUTS, + DEFAULT_TURN_TIMEOUT_S, + TURNS_DIRNAME, + InProcessTurnExecutor, + LayerProvider, + TurnExecutionError, + TurnExecutor, + TurnOutcome, + TurnRequest, + TurnTimeout, + turns_in_flight, +) + +KIND_LOCAL = "local" +KIND_SESSION_SPOOL = "session_spool" + + +def select_command_executor( + environment: RunEnvironment, + *, + name_prefix: str = "gate", + tree: Path | str | None = None, + gpus: str = "all", +) -> CommandExecutor: + """The executor this run's commands go through. + + ``tree`` is the checkout under test: on a spool every command then runs with + it as its working directory and first ``PYTHONPATH`` entry. ``gpus`` is what + each spool job asks the scheduler for: a device count, or ``all`` for the + whole node. The spools of every session declared with ``role: gates`` form + a pool beside ``environment.spool``, and each command goes to the least + loaded live one. + """ + kind = environment.kind + if kind == KIND_SESSION_SPOOL and environment.spool: + # Every gate session's spool is a place a gate may run. + pool = [spool for _, spool, role in environment.sessions if role == "gates"] + # Another session's idle devices take a gate no gate session can start now. + overflow = [spool for _, spool, role in environment.sessions if role != "gates"] + return SpoolCommandExecutor( + environment.spool, + name_prefix=name_prefix, + tree=tree, + gpus=gpus, + pool=pool, + overflow=overflow, + # The tune campaign's dev sessions, only while nothing runs or waits there. + last_resort=environment.dev_spools, + ) + return LocalCommandExecutor() + + +def startable_checks( + environment: RunEnvironment, + gpus: str | None = None, + wall_s: float | None = None, + *, + now: float | None = None, +) -> int | None: + """How many more checks would start at once on this run's sessions. + + A check asks for ``gpus`` devices (``all`` for a whole node; by default + ``environment.world_size``, its member's width) with ``wall_s`` to run + (the session's default when ``None``), and goes where the executor + :func:`select_command_executor` builds sends it + (:meth:`SpoolCommandExecutor.startable`). ``None`` when no session has a + live watcher, and under ``local``: nothing there says what is free. + """ + want = str(environment.world_size if gpus is None else gpus) + executor = select_command_executor(environment, gpus=want) + if not isinstance(executor, SpoolCommandExecutor): + return None + job = placement.JobNeeds( + name="check", + gpus=None if want == "all" else max(1, int(want)), + wall_s=wall_s, + holds_gpus=True, + owner=executor.name_prefix, + ) + return executor.startable(job, now) + + +def check_capacity( + environment: RunEnvironment, + now: float | None = None, + *, + gpus: str | None = None, + wall_s: float | None = None, + in_flight: int = 0, + starting: int = 0, + floor: int = DEFAULT_CHECK_FLOOR, + cap: int = DEFAULT_CHECK_CAP, + fallback: int = DEFAULT_CHECK_FALLBACK, +) -> int: + """How many more checks may start now, with ``in_flight`` already running. + + :func:`startable_checks` read through :func:`more_checks`: never fewer + than ``floor`` in flight, never more than ``cap``, and ``fallback`` when + no session can be read. File reads only. + """ + return more_checks( + startable_checks(environment, gpus, wall_s, now=now), + in_flight=in_flight, + starting=starting, + floor=floor, + cap=cap, + fallback=fallback, + ) + + +def wait_for_session(executor: CommandExecutor | None, timeout_s: float) -> bool: + """Wait up to ``timeout_s`` for the session ``executor`` submits to to heart-beat again. + + What a result whose ``infrastructure`` is set asks of its caller before the + command is worth running again. False at once for an executor with no + session behind it: nothing there can come back. + """ + if isinstance(executor, SpoolCommandExecutor): + return executor.wait_for_session(timeout_s) + return False + + +def shell_runner_for(executor: CommandExecutor, *, default_timeout: int = 1800) -> Any: + """``executor`` as the ``(command, env, timeout) -> CommandResult`` callable gates take.""" + + def run( + command: str, + env: Mapping[str, str] | None = None, + timeout: int = default_timeout, + ) -> CommandResult: + return executor.execute(CommandSpec(argv=command, env=env, timeout_s=timeout)) + + return run + + +def select_turn_executor( + agent_transport: TurnExecutor | None, + layer_provider: LayerProvider, + *, + timeout_s: float | None = None, + records: Path | str | None = None, +) -> TurnExecutor: + """The transport this run's role turns go through: in process, unless one was given. + + ``timeout_s`` bounds every in-process turn alike (each role has its own + bound without it), and ``records`` is where each one leaves its files; a + given transport keeps its own. + """ + return agent_transport or InProcessTurnExecutor( + layer_provider, timeout_s=timeout_s, records=records + ) + + +__all__ = [ + "DEFAULT_CHECK_CAP", + "DEFAULT_CHECK_FALLBACK", + "DEFAULT_CHECK_FLOOR", + "DEFAULT_ROLE_TURN_TIMEOUTS", + "DEFAULT_TURN_TIMEOUT_S", + "KIND_LOCAL", + "KIND_SESSION_SPOOL", + "NOT_FOUND_RETURNCODE", + "TIMEOUT_RETURNCODE", + "TURNS_DIRNAME", + "CheckAdmission", + "CommandExecutor", + "CommandResult", + "CommandSpec", + "InProcessTurnExecutor", + "LayerProvider", + "LocalCommandExecutor", + "SpoolCommandExecutor", + "TurnExecutionError", + "TurnExecutor", + "TurnOutcome", + "TurnRequest", + "TurnTimeout", + "check_capacity", + "more_checks", + "select_command_executor", + "select_turn_executor", + "shell_runner_for", + "startable_checks", + "turns_in_flight", + "wait_for_session", +] diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/agent_gate.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/agent_gate.py new file mode 100644 index 000000000000..12adf3d22413 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/agent_gate.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""One cap on the agents a run has working at once, shared by priority. + +Every role turn passes through `Workflow._run_turn` and takes a slot here for +as long as it runs, so the run never has more than ``capacity`` agents alive. +Who gets a slot is decided by the turn's priority, which the code that starts +the turn declares (:func:`turn_priority`), never guessed from where it runs: + +- ``PLANNER``: the planner turns a finished task into the next one, so it owns + the reserved slot and goes first; it never waits hours behind task turns. +- ``CRITICAL``, then ``REQUIRED``: the graph's tasks, those on the critical + path first. +- ``SIDE``: the critic and the improver, advisory work that waits for the tasks: + no more of them run than the open foreground work leaves idle + (``side_limit``), unless one has waited ``side_max_wait_s``. +- ``BACKGROUND``: a tune campaign's turns, filler for capacity nothing else + uses: no more of them run than the open foreground work leaves idle + (``background_limit``). + +A waiting turn of a higher priority goes before any of a lower one; only the +planner may take the reserved slot, and a planner turn in it takes no slot from +the tasks: the loop that starts tasks counts them alone. +""" + +from __future__ import annotations + +import contextlib +import enum +import threading +import time +from typing import Callable, Iterator + + +class Priority(enum.IntEnum): + """A turn's claim on an agent slot; a lower value goes first.""" + + PLANNER = 0 + CRITICAL = 1 + REQUIRED = 2 + SIDE = 3 + BACKGROUND = 4 + + +_local = threading.local() + + +@contextlib.contextmanager +def turn_priority(priority: Priority) -> Iterator[None]: + """Turns started in this thread, inside the block, take slots at ``priority``.""" + previous = getattr(_local, "priority", None) + _local.priority = priority + try: + yield + finally: + _local.priority = previous + + +def current_priority() -> Priority: + """The priority this thread's turns take slots at; ``REQUIRED`` when none was declared.""" + priority = getattr(_local, "priority", None) + return Priority.REQUIRED if priority is None else priority + + +def run_at(priority: Priority, fn: Callable[..., object], *args: object) -> object: + """``fn(*args)`` with its turns at ``priority``: for work submitted to a pool.""" + with turn_priority(priority): + return fn(*args) + + +#: How often a background turn waiting for idle capacity looks again; the +#: foreground's demand changes without anyone telling the gate. +_BACKGROUND_POLL_S = 5.0 +#: How long a side turn waits for idle slots before it takes the next one free. +SIDE_MAX_WAIT_S = 20 * 60.0 + + +class AgentGate: + """At most ``capacity`` turns at once, granted by :class:`Priority`. + + ``reserve`` slots are the planner's alone. ``background_limit`` says how + many background turns may run now: the slots the open foreground work + leaves idle (no limit when not given). ``side_limit`` says the same of + side turns, which take a slot anyway once they have waited + ``side_max_wait_s``. + """ + + def __init__( + self, + capacity: int, + reserve: int = 1, + background_limit: Callable[[], int] | None = None, + side_limit: Callable[[], int] | None = None, + side_max_wait_s: float = SIDE_MAX_WAIT_S, + ) -> None: + self.capacity = max(1, capacity) + self.reserve = min(max(0, reserve), self.capacity - 1) + self.background_limit = background_limit + self.side_limit = side_limit + self.side_max_wait_s = side_max_wait_s + self._active: dict[Priority, int] = {p: 0 for p in Priority} + self._waiting: dict[Priority, int] = {p: 0 for p in Priority} + self._cond = threading.Condition() + + @property + def task_slots(self) -> int: + """How many turns other than the planner's may run at once.""" + return self.capacity - self.reserve + + @property + def active(self) -> int: + with self._cond: + return sum(self._active.values()) + + def active_at(self, priority: Priority) -> int: + with self._cond: + return self._active[priority] + + def _admissible(self, priority: Priority, waited_s: float = 0.0) -> bool: + active = sum(self._active.values()) + if priority == Priority.PLANNER: + return active < self.capacity + # The planner's turns in the reserved slots are not the tasks'. + planner = self._active[Priority.PLANNER] + if active - min(planner, self.reserve) >= self.task_slots: + return False + if any(self._waiting[p] for p in Priority if p < priority): + return False + if priority == Priority.BACKGROUND and self.background_limit is not None: + return self._active[Priority.BACKGROUND] < max(0, self.background_limit()) + if ( + priority == Priority.SIDE + and self.side_limit is not None + and waited_s < self.side_max_wait_s + ): + return self._active[Priority.SIDE] < max(0, self.side_limit()) + return True + + @contextlib.contextmanager + def slot(self, priority: Priority | None = None) -> Iterator[None]: + """Hold a slot for one turn, at ``priority`` or the thread's (:func:`turn_priority`).""" + priority = current_priority() if priority is None else priority + polled = priority in (Priority.BACKGROUND, Priority.SIDE) + poll = _BACKGROUND_POLL_S if polled else None + since = time.monotonic() + with self._cond: + self._waiting[priority] += 1 + try: + while not self._admissible(priority, time.monotonic() - since): + self._cond.wait(timeout=poll) + finally: + self._waiting[priority] -= 1 + self._active[priority] += 1 + self._cond.notify_all() + try: + yield + finally: + with self._cond: + self._active[priority] -= 1 + self._cond.notify_all() + + +__all__ = ["AgentGate", "Priority", "current_priority", "run_at", "turn_priority"] diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/capacity.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/capacity.py new file mode 100644 index 000000000000..d8f69d211438 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/capacity.py @@ -0,0 +1,259 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""How many checks run at once, and what a running task holds of the sessions' devices. + +A fixed number of checks in flight is wrong both ways. While the gate +sessions are full it queues more jobs behind each other in their spools, +where each was placed when it was submitted and cannot move to a session +that frees up later; while the work and dev sessions stand idle it leaves +their devices unused. So a check starts when a session would start its job +at once (:meth:`.spool.SpoolCommandExecutor.startable`, counted over every +spool the job may go to) and waits in this process otherwise, to be placed +when a device frees. + +:func:`more_checks` turns that count into how many more checks may start; +:class:`CheckAdmission` is the gate a run's check threads pass through. + +A running task holds the ``gpus`` it declares of the run's sessions (the gate +and work sessions), except a task whose every job runs on sessions of its own: +the tune campaign once ``perf.spools`` names its measuring sessions +(:func:`offsite_kinds`, :func:`held_gpus`). Its reservation would otherwise keep +devices from the gate and work sessions for the campaign's whole life while none +of its jobs ever runs there. +""" + +from __future__ import annotations + +import itertools +import threading +import time +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Callable, Collection, Iterator, Mapping + +from ..environment import RunEnvironment + +#: Checks in flight however busy the sessions read: their jobs queue in a spool +#: and the run keeps moving even when every record is wrong. +DEFAULT_CHECK_FLOOR = 2 +#: Checks in flight at most, however many devices stand free. +DEFAULT_CHECK_CAP = 32 +#: Checks in flight when no session can be read (``local``, or no live watcher). +DEFAULT_CHECK_FALLBACK = 8 + +#: How many more checks like the one asked about would start at once on the +#: run's sessions: ``(environment, gpus, wall_s) -> int``, ``None`` when none +#: can be read. +Startable = Callable[[RunEnvironment, str | None, float | None], int | None] + + +def more_checks( + startable: int | None, + *, + in_flight: int = 0, + starting: int = 0, + floor: int = DEFAULT_CHECK_FLOOR, + cap: int = DEFAULT_CHECK_CAP, + fallback: int = DEFAULT_CHECK_FALLBACK, +) -> int: + """How many more checks may start now, with ``in_flight`` already running. + + Args: + startable: How many would start at once on the sessions; ``None`` when + none can be read, and then ``fallback`` run at once. + in_flight: Checks running now. + starting: Those of ``in_flight`` started too recently for their jobs to + show in a spool yet; they still take what the sessions read as free. + floor: Checks that may always run at once. + cap: Checks that may run at once at most. + fallback: Checks that may run at once when nothing can be read. + """ + target = fallback if startable is None else in_flight + max(0, startable - starting) + return max(0, min(cap, max(floor, target)) - in_flight) + + +class CheckAdmission: + """Holds each check back until :func:`more_checks` lets it start, and counts it while it runs. + + One per run, shared by every check thread; the environments it is asked + about differ only in width. :meth:`slot` blocks its caller until the check + may start and holds its place while it runs. + + - A check started less than ``settle_s`` ago counts against what the + sessions read as free: its job may not have reached a spool yet. + - A check that has waited ``patience_s`` goes first: none that arrived + after it starts before it, so a check that needs a whole node gets the + next node that drains instead of losing every device to narrower ones. + - The sessions are read at most every ``refresh_s`` per width and wall + limit, and again as soon as a check ends. + + Args: + startable: Reads the sessions (:data:`Startable`; the package's ``startable_checks``). + floor: See :func:`more_checks`. + cap: See :func:`more_checks`; also the size the check pool needs. + fallback: See :func:`more_checks`. + settle_s: How long a started check counts as not yet in a spool. + patience_s: How long a check waits before none may pass it. + poll_s: How often a waiting check reads the sessions again. + refresh_s: How long one reading of the sessions stands. + clock: Monotonic seconds. + """ + + def __init__( + self, + startable: Startable, + *, + floor: int = DEFAULT_CHECK_FLOOR, + cap: int = DEFAULT_CHECK_CAP, + fallback: int = DEFAULT_CHECK_FALLBACK, + settle_s: float = 60.0, + patience_s: float = 600.0, + poll_s: float = 10.0, + refresh_s: float = 2.0, + clock: Callable[[], float] = time.monotonic, + ) -> None: + self._startable = startable + self.floor = max(1, floor) + self.cap = max(self.floor, cap) + self.fallback = fallback + self.settle_s = settle_s + self.patience_s = patience_s + self.poll_s = poll_s + self.refresh_s = refresh_s + self._clock = clock + self._cond = threading.Condition() + self._tickets = itertools.count() + #: Checks running, by ticket: when each started. + self._running: dict[int, float] = {} + #: Checks waiting, by ticket: when each arrived. + self._waiting: dict[int, float] = {} + #: The last reading per width and wall limit: when, and what it said. + self._readings: dict[tuple[str, float | None], tuple[float, int | None]] = {} + + @property + def in_flight(self) -> int: + """Checks running now.""" + with self._cond: + return len(self._running) + + @property + def waiting(self) -> int: + """Checks waiting to start.""" + with self._cond: + return len(self._waiting) + + @contextmanager + def slot( + self, + environment: RunEnvironment, + gpus: str | None = None, + wall_s: float | None = None, + ) -> Iterator[None]: + """Wait until a check may start; hold its place while it runs. + + Args: + environment: The check's environment; its world size is the check's width. + gpus: The devices its job asks for instead (``all``: a whole node). + wall_s: Its job's wall limit; ``None`` takes each session's default. + """ + ticket = next(self._tickets) + with self._cond: + self._waiting[ticket] = self._clock() + try: + while not self._may_start(ticket, environment, gpus, wall_s): + self._cond.wait(self.poll_s) + finally: + del self._waiting[ticket] + self._running[ticket] = self._clock() + # Whoever waits behind it reads what is left now rather than at its next poll. + self._cond.notify_all() + try: + yield + finally: + with self._cond: + del self._running[ticket] + self._readings.clear() + self._cond.notify_all() + + def _may_start( + self, ticket: int, environment: RunEnvironment, gpus: str | None, wall_s: float | None + ) -> bool: + now = self._clock() + in_flight = len(self._running) + if in_flight >= self.cap: + return False + overdue = [t for t, since in self._waiting.items() if now - since >= self.patience_s] + if overdue and ticket != min(overdue): + return False + if in_flight < self.floor: + return True + starting = sum(1 for since in self._running.values() if now - since < self.settle_s) + more = more_checks( + self._read(environment, gpus, wall_s, now), + in_flight=in_flight, + starting=starting, + floor=self.floor, + cap=self.cap, + fallback=self.fallback, + ) + return more > 0 + + def _read( + self, environment: RunEnvironment, gpus: str | None, wall_s: float | None, now: float + ) -> int | None: + key = (str(environment.world_size if gpus is None else gpus), wall_s) + held = self._readings.get(key) + if held is None or now - held[0] >= self.refresh_s: + held = (now, self._startable(environment, gpus, wall_s)) + self._readings[key] = held + return held[1] + + +#: The task kind that runs the tune campaign. +TUNE_KIND = "tune" + + +def offsite_kinds(perf: Mapping[str, Any] | None, environment: RunEnvironment) -> frozenset[str]: + """The task kinds none of whose jobs runs on the run's sessions: ``{"tune"}`` or none. + + A ``tune`` task measures on ``perf.spools``, and gates its keeps and runs its tuners' + jobs on ``perf.dev_spools`` (the measuring spools without them). Once ``perf.spools`` + is set and none of those is one of the run's sessions (``environment.spool`` and + ``environment.sessions``), its ``gpus`` are never the sessions' to hold. Without + ``perf.spools`` the campaign measures on the run's own sessions, and its ``gpus`` + count as any task's. + """ + if not isinstance(perf, Mapping) or not perf.get("spools"): + return frozenset() + measuring = [str(spool) for spool in perf["spools"]] + its_own = {*measuring, *(str(spool) for spool in perf.get("dev_spools") or measuring)} + sessions = {environment.spool, *(spool for _, spool, _ in environment.sessions)} - {None, ""} + if {_resolved(s) for s in its_own} & {_resolved(str(s)) for s in sessions}: + return frozenset() + return frozenset({TUNE_KIND}) + + +def held_gpus(task: Any, offsite: Collection[str] = ()) -> int: + """The devices of the run's sessions ``task`` holds while it runs. + + Its ``gpus``; none for a task of a kind in ``offsite`` (:func:`offsite_kinds`). + """ + return 0 if task.kind in offsite else max(0, int(task.gpus or 0)) + + +def _resolved(spool: str) -> str: + return str(Path(spool).expanduser().resolve()) + + +__all__ = [ + "DEFAULT_CHECK_CAP", + "DEFAULT_CHECK_FALLBACK", + "DEFAULT_CHECK_FLOOR", + "TUNE_KIND", + "CheckAdmission", + "Startable", + "held_gpus", + "more_checks", + "offsite_kinds", +] diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/commands.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/commands.py new file mode 100644 index 000000000000..17e9fdf79f90 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/commands.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""One command, what it may take, and what came back. + +The caller builds a :class:`CommandSpec`; an executor runs it somewhere and +returns a :class:`CommandResult`. **A timeout is not an error here**: +``timed_out`` is separate from the exit code because a gate that timed out has +said nothing about what it was gating, while a gate that exited non-zero has. + +**Neither is a lost session.** ``infrastructure`` names the third outcome: the +substrate, not the command, ended the wait (the allocation that would run it +stopped heart-beating, none ever watched, or it was closed out as orphaned). A +caller that would re-measure on a timeout should instead wait for the substrate +to come back. +""" + +from __future__ import annotations + +import abc +import os +import shlex +import subprocess +from dataclasses import dataclass +from typing import Mapping, Sequence + +#: What a command that never finished reports: ``timeout(1)``'s code. +TIMEOUT_RETURNCODE = 124 + +#: What a command whose executable does not exist reports: the shell's own +#: "command not found". Distinct from a non-zero exit: the probe could not run, +#: rather than ran and disagreed. +NOT_FOUND_RETURNCODE = 127 + + +@dataclass(frozen=True) +class CommandSpec: + """What to run, and the bound it runs under. + + ``argv`` is either a **string**, run as a shell line, or a **sequence**, + executed directly with no shell between. A resolved launch prefix is shell + syntax; a probe like ``nvidia-smi --query-gpu=...`` has no reason to be + exposed to word splitting. + """ + + argv: str | Sequence[str] + env: Mapping[str, str] | None = None + timeout_s: int = 600 + #: Working directory, where the executor supports one. + cwd: str | None = None + + @property + def rendered(self) -> str: + """``argv`` as a single line, for logs and for shell execution.""" + if isinstance(self.argv, str): + return self.argv + return " ".join(shlex.quote(str(part)) for part in self.argv) + + @property + def is_shell_line(self) -> bool: + return isinstance(self.argv, str) + + +@dataclass(frozen=True) +class CommandResult: + """What a command did: its exit code, its output, and whether it ran out of time.""" + + returncode: int + stdout: str = "" + stderr: str = "" + timed_out: bool = False + #: Why the substrate ended this command rather than the command itself; + #: empty when the command ran to its own outcome. A result that carries it + #: also has ``timed_out`` set. + infrastructure: str = "" + + @property + def ok(self) -> bool: + return self.returncode == 0 and not self.timed_out + + +class CommandExecutor(abc.ABC): + """Run one command somewhere and report what happened.""" + + @abc.abstractmethod + def execute(self, spec: CommandSpec) -> CommandResult: + """Run ``spec`` to completion or to its deadline.""" + + +class LocalCommandExecutor(CommandExecutor): + """A subprocess on this host. + + The environment is the parent's with ``spec.env`` layered on top: + everything here runs inside an environment somebody already set up (a + venv's ``PATH``, a container's library paths, ``TMPDIR``), and replacing it + turns a working host into a broken one. + """ + + def execute(self, spec: CommandSpec) -> CommandResult: + try: + completed = subprocess.run( + spec.rendered if spec.is_shell_line else list(spec.argv), + shell=spec.is_shell_line, + capture_output=True, + text=True, + timeout=spec.timeout_s, + env={**os.environ, **(spec.env or {})}, + cwd=spec.cwd, + check=False, + ) + except subprocess.TimeoutExpired: + return CommandResult(returncode=TIMEOUT_RETURNCODE, timed_out=True) + except FileNotFoundError as exc: + # Only reachable on the argv path: a shell reports this itself. + return CommandResult(returncode=NOT_FOUND_RETURNCODE, stderr=str(exc)) + return CommandResult( + returncode=completed.returncode, + stdout=completed.stdout or "", + stderr=completed.stderr or "", + ) + + +__all__ = [ + "NOT_FOUND_RETURNCODE", + "TIMEOUT_RETURNCODE", + "CommandExecutor", + "CommandResult", + "CommandSpec", + "LocalCommandExecutor", +] diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/spool.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/spool.py new file mode 100644 index 000000000000..03d9dcf32cb9 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/spool.py @@ -0,0 +1,403 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run a command through a compute-session spool and block on its receipt. + +The session is an allocation acquired once whose main process drains +``/pending`` (`infra/substrate/scheduler.py`): one queue wait and one +container import for the run instead of one of each per gate, and work any +process able to write a file can submit. The protocol is `infra/substrate/spoolctl.py`'s; +this executor is the gate's policy over it: + +- **The deadline counts from the job's start**, not from its submission: time + queued behind other work says nothing about the command. While it is queued, + the wait is bounded by the session's liveness instead, so a dead session + cannot hang its caller. +- **A command given up on is withdrawn** (``cancel/``), so a session that + comes back does not run work nobody is waiting for. A spool whose heartbeat is + already stale gets nothing queued at all. +- **A lost session is not a failed command**: when the session, not the command, + ended the wait, the result says so in ``infrastructure``, and + :meth:`SpoolCommandExecutor.wait_for_session` is how a caller waits for one to + come back before running the command again. +- **A pool spreads the work**: with several spools (one per gate session), each + command goes to the live spool with the least work queued or running among + those that can take it now (``placement``: not in hand-over, and with the + command's wall limit left before the allocation ends); to the least loaded + live one when none can. +- **Idle devices elsewhere take the overflow**: a command no spool of the pool + would start at once goes to another session's spool that would (its devices + free, nothing queued); failing that, to a last-resort spool (the tune + campaign's dev sessions) while nothing of the campaign's runs or waits there. +- **A queued command moves to where it starts**: while it waits in one queue, + a spool that would start it now takes it (:meth:`SpoolCommandExecutor._relocate`), + so devices freed elsewhere, or a session that died, do not keep it waiting. + :meth:`SpoolCommandExecutor.startable` counts by the same rules how many + more commands would start at once. +""" + +from __future__ import annotations + +import itertools +import os +import threading +import time +import uuid +from pathlib import Path +from typing import Sequence + +from ..substrate import placement, spoolctl +from ..substrate.scheduler import SESSION_LOSS_REASONS, SpoolPaths, job_spec_from_script +from .commands import TIMEOUT_RETURNCODE, CommandExecutor, CommandResult, CommandSpec + +#: Where a command may go, in the order :meth:`SpoolCommandExecutor._pick` tries. +_POOL, _OVERFLOW, _LAST_RESORT = "pool", "overflow", "last_resort" + + +class SpoolCommandExecutor(CommandExecutor): + """Submit through a spool directory and block on the receipt.""" + + def __init__( + self, + spool: str | Path, + *, + poll_s: float = 3.0, + name_prefix: str = "gate", + gpus: str = "all", + tree: str | Path | None = None, + queue_timeout_s: float | None = None, + stale_s: float = 180.0, + grace_s: float = 180.0, + pool: Sequence[str | Path] = (), + overflow: Sequence[str | Path] = (), + last_resort: Sequence[str | Path] = (), + ) -> None: + self.root = Path(spool) + #: Every spool a command may go to, ``spool`` first. + self.pool = [self.root, *(Path(p) for p in pool if Path(p) != self.root)] + #: Spools of other sessions a command goes to only when it starts there + #: at once and no spool of the pool can start it now. + self.overflow = [Path(p) for p in overflow if Path(p) not in self.pool] + #: Spools kept for other work, lent only when neither the pool nor the + #: overflow would start a command at once, and only while nothing waits + #: there and nothing but commands it lent already runs there. + self.last_resort = [ + Path(p) for p in last_resort if Path(p) not in (*self.pool, *self.overflow) + ] + self._pick_lock = threading.Lock() + self.poll_s = poll_s + self.name_prefix = name_prefix + self.gpus = gpus + #: The checkout a job runs against: its working directory and first + #: ``PYTHONPATH`` entry; the session refuses a result that imported the + #: package under test from anywhere else. + self.tree = str(Path(tree).resolve()) if tree else None + #: Bound on time spent queued; ``None`` waits as long as the session is alive. + self.queue_timeout_s = queue_timeout_s + #: A session whose watchers have not heart-beat for this long is gone. + self.stale_s = stale_s + #: Allowance past a job's own deadline for the session to stop it. + self.grace_s = grace_s + self._seq = itertools.count(1) + + def _name(self) -> str: + """Unique across processes and threads: two gates in one millisecond must not collide. + + The pid alone does not separate processes: each sandboxed command runs in + its own pid namespace, where two gates started together share a small pid. + """ + return ( + f"{self.name_prefix}-{int(time.time() * 1000)}-{os.getpid()}-" + f"{uuid.uuid4().hex[:6]}-{next(self._seq)}" + ) + + def _alive(self, root: Path) -> bool: + age = spoolctl.watcher_age(root) + return age is not None and age <= self.stale_s + + def session_alive(self) -> bool: + """Whether a watcher on any spool of the pool has heart-beat within ``stale_s``.""" + return any(self._alive(root) for root in self.pool) + + def _pick(self, job: placement.JobNeeds | None = None, now: float | None = None) -> Path: + """The live spool to send a command to; the first of the pool when none is live. + + With ``job``, only the live spools that can take it now compete + (:attr:`placement.SpoolView.pickable`): a session in hand-over, or + one whose allocation ends before the job's wall limit, would park it + for a successor that may be hours away. One that would start it at + once comes first, then the least loaded. When no spool can take it, + the least loaded live one is still the best place to queue it. + + When no spool of the pool would start ``job`` at once, an + :attr:`overflow` spool that would (its devices free, nothing queued, + out of hand-over) takes it: the gates wait behind each other while + another session's devices stand idle. When none would either, a + :attr:`last_resort` spool that would start it takes it, while nothing + waits there and nothing of its own runs (:meth:`_lendable`). + """ + now = time.time() if now is None else now + if len(self.pool) == 1 and not (job is not None and (self.overflow or self.last_resort)): + return self.root + live: list[tuple[int, Path]] = [] + ready: list[tuple[bool, int, Path]] = [] + starts_now = False + for root in self.pool: + if not self._alive(root): + continue + paths = SpoolPaths(root) + load = sum(1 for _ in paths.pending.glob("*.sh")) + sum( + 1 for _ in paths.running.glob("*.sh") + ) + live.append((load, root)) + if job is not None: + view = placement.assess(root, job, stale_s=self.stale_s, now=now) + if view.pickable: + starts = view.starts_now and _nothing_queued(root) + ready.append((not starts, load, root)) + starts_now = starts_now or starts + if job is not None and not starts_now: + for spools, free_enough in ( + (self.overflow, _nothing_queued), + (self.last_resort, self._lendable), + ): + idle = [ + view + for view in ( + placement.assess(root, job, stale_s=self.stale_s, now=now) + for root in spools + if self._alive(root) + ) + if view.pickable and view.starts_now and free_enough(view.path) + ] + if idle: + return min(idle, key=placement.rank).path + if ready: + # One that starts the job at once first, then the least loaded. + return min(ready, key=lambda choice: choice[:2])[2] + return min(live, key=lambda choice: choice[0])[1] if live else self.root + + def startable(self, job: placement.JobNeeds, now: float | None = None) -> int | None: + """How many more jobs like ``job`` would start at once where :meth:`_pick` sends them. + + Summed over every spool with a live watcher (:meth:`_offer`); ``None`` + when there is none, so there is nothing to count. File reads only: each + spool's watcher records and its queue on disk. + """ + now = time.time() if now is None else now + offers = [ + self._offer(root, job, now, tier) + for tier, spools in ( + (_POOL, self.pool), + (_OVERFLOW, self.overflow), + (_LAST_RESORT, self.last_resort), + ) + for root in spools + ] + live = [offer for offer in offers if offer is not None] + return sum(live) if live else None + + def _offer(self, root: Path, job: placement.JobNeeds, now: float, tier: str) -> int | None: + """How many jobs like ``job`` ``root`` would start at once; ``None`` with no live watcher. + + Each live watcher that can take the job (:attr:`placement.SpoolView.pickable`) + fits as many as its free devices hold. A spool of the pool offers those + fits less what the jobs queued there on disk will take first; an + overflow spool offers them only with nothing queued, and a last-resort + spool only while it is :meth:`_lendable`: what :meth:`_pick` would send + to each. + """ + views = placement.watcher_views(root, job, stale_s=self.stale_s, now=now) + if not views: + return None + fits = sum(_fits(job, view) for view in views if view.pickable) + if tier == _OVERFLOW: + return fits if _nothing_queued(root) else 0 + if tier == _LAST_RESORT: + return fits if self._lendable(root) else 0 + node = max(view.devices for view in views) + taken = sum(_fits_taken(job, queued, node) for queued in placement.queued_jobs(root)) + return max(0, fits - taken) + + def wait_for_session(self, timeout_s: float, *, poll_s: float | None = None) -> bool: + """Block until a watcher heart-beats on this spool again, for at most ``timeout_s``.""" + deadline = time.monotonic() + max(0.0, timeout_s) + interval = self.poll_s if poll_s is None else poll_s + while not self.session_alive(): + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + time.sleep(min(interval, remaining)) + return True + + def _lendable(self, root: Path) -> bool: + """Whether a last-resort spool lends its free devices: nothing waits, and nothing of its own runs. + + Its own work (a campaign's profile, say) must not share the node with + a gate; gates this executor lent it already may. + """ + if not _nothing_queued(root): + return False + for script in SpoolPaths(root).running.glob("*.sh"): + try: + text = script.read_text(encoding="utf-8", errors="replace") + except OSError: + return False + if job_spec_from_script(script.stem, text).owner != self.name_prefix: + return False + return True + + def _elsewhere(self) -> bool: + """Whether a command could go to more than one spool.""" + return len(self.pool) + len(self.overflow) + len(self.last_resort) > 1 + + def _relocate(self, job: placement.JobNeeds, current: Path) -> Path | None: + """Move ``job``, still queued in ``current``, to a spool that would start it now; that spool. + + A job is placed when it is submitted, and a session that frees + devices later, or a queue whose session died, would otherwise keep it + waiting while devices elsewhere stand idle. The script is renamed into + the other spool's queue: a watcher that claimed it first has renamed it + into its own ``running/``, and the rename fails, so a job never runs + twice. ``None`` when it stays. + """ + now = time.time() + with self._pick_lock: + target = self._pick(job, now) + if target == current or not self._alive(target): + return None + view = placement.assess(target, job, stale_s=self.stale_s, now=now) + if not (view.pickable and view.starts_now and _nothing_queued(target)): + return None + source = SpoolPaths(current).pending / f"{job.name}.sh" + destination = SpoolPaths(target).pending / f"{job.name}.sh" + if destination.exists(): + return None + try: + destination.parent.mkdir(parents=True, exist_ok=True) + os.rename(source, destination) + except OSError: + return None + return target + + def _withdraw(self, root: Path, name: str) -> str: + try: + flag = spoolctl.cancel_job(root, name) + except OSError as exc: + return f"; withdrawing it failed ({exc}), so a session draining {root} may still run it" + return f"; withdrawn through {flag}" + + def _log(self, root: Path, name: str) -> str: + log = SpoolPaths(root).log(name) + return log.read_text(encoding="utf-8", errors="replace") if log.exists() else "" + + def execute(self, spec: CommandSpec) -> CommandResult: + name = self._name() + token = uuid.uuid4().hex + submitted = time.time() + header = { + "gpus": self.gpus, + "wall": int(spec.timeout_s) if spec.timeout_s and spec.timeout_s > 0 else None, + "owner": self.name_prefix, + "token": token, + "tree": self.tree if self.tree and not spec.cwd else None, + } + script = spoolctl.render_job(spec.rendered, header=header, cwd=spec.cwd, env=spec.env) + with self._pick_lock: + root = self._pick(placement.JobNeeds.from_script(name, script), submitted) + age = spoolctl.watcher_age(root, submitted) + if age is None or age <= self.stale_s: + root.mkdir(parents=True, exist_ok=True) + spoolctl.queue_job( + root, name, script, owner=self.name_prefix, require_watcher=False + ) + if age is not None and age > self.stale_s: + reason = ( + f"the session watching {root} stopped heart-beating {age:.0f}s ago; " + f"nothing was queued" + ) + return CommandResult( + returncode=TIMEOUT_RETURNCODE, + stdout=f"=== spool executor: {reason} ===\n", + timed_out=True, + infrastructure=reason, + ) + job = placement.JobNeeds.from_script(name, script) + where = {"root": root} + + def relocate(current: Path) -> Path | None: + moved = self._relocate(job, current) + if moved is not None: + where["root"] = moved + return moved + + try: + record = spoolctl.wait_for_job( + root, + name, + token=token, + poll_s=self.poll_s, + timeout_s=spec.timeout_s, + stale_s=self.stale_s, + grace_s=self.grace_s, + queue_timeout_s=self.queue_timeout_s, + unwatched_s=spec.timeout_s, + submitted_at=submitted, + relocate=relocate if self._elsewhere() else None, + ) + except spoolctl.SpoolError as exc: + root = where["root"] + reason = str(exc) + self._withdraw(root, name) + lost = exc.exit_code == spoolctl.EXIT_NO_SESSION + return CommandResult( + returncode=TIMEOUT_RETURNCODE, + stdout=f"{self._log(root, name)}\n=== spool executor: {reason} ===\n", + timed_out=True, + infrastructure=reason if lost else "", + ) + root = where["root"] + lost = "" + if record.get("reason") in SESSION_LOSS_REASONS: + lost = ( + f"the session that ran it stopped heart-beating and the job was closed " + f"out as {record['reason']}: {record.get('detail') or 'no detail recorded'}" + ) + return CommandResult( + returncode=int(record["rc"]), + stdout=self._log(root, name), + # The scheduler stops a job at its wall limit or when it stops making + # progress, and a job whose session died under it measured nothing. + timed_out=bool(record.get("timed_out")) or bool(lost), + infrastructure=lost, + ) + + +def _nothing_queued(root: Path) -> bool: + """No job waits in ``root``'s queue on disk. + + The watcher reports its free devices on its next tick, so a burst of + submissions would all read the same free devices and pile onto one + session; a job already waiting in the queue is seen at once. + """ + return not any(SpoolPaths(root).pending.glob("*.sh")) + + +def _width(gpus: int | None, node: int) -> int: + """Devices a job asking for ``gpus`` (``None``: the whole node) holds on a node of ``node``.""" + return max(1, node if gpus is None else gpus) + + +def _fits(job: placement.JobNeeds, view: placement.SpoolView) -> int: + """How many jobs like ``job`` the free devices of ``view``'s watcher hold.""" + return view.free // _width(job.gpus, view.devices) if view.devices else 0 + + +def _fits_taken(job: placement.JobNeeds, queued: placement.JobNeeds, node: int) -> int: + """How many fits of ``job`` the queued job ``queued`` takes on nodes of ``node`` devices.""" + if not queued.holds_gpus: + return 0 + held, each = _width(queued.gpus, node), _width(job.gpus, node) + return -(-held // each) + + +__all__ = ["SESSION_LOSS_REASONS", "SpoolCommandExecutor"] diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/turns.py b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/turns.py new file mode 100644 index 000000000000..95456e1fd229 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/execution/turns.py @@ -0,0 +1,539 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""One agent turn, as a request and an outcome, and the executor that runs it here. + +A *turn* is one round trip with one role: the orchestrator hands over a system +prompt and a task, the role does its work, and a text answer comes back (its +typed decision arrives separately, through the decision file or tool). The two +transports, in process and through the filesystem (`transport.py`), take the +same request and return the same outcome, so the orchestrator never branches on +which one is in force. Both bound a turn by its role's timeout +(:data:`DEFAULT_ROLE_TURN_TIMEOUTS`) and raise :class:`TurnTimeout` past it. + +An in-process turn leaves the evidence a dispatched one leaves in +``dispatch/``, in the directory the run names +(:meth:`InProcessTurnExecutor.keep_records`, ``/turns/``), each file +written as the turn goes, so an operator can tail one while it runs:: + + .request.json role, task, agent, process, start, bound, and the files below + .system.md the role's system prompt + .prompt.md the task prompt + .activity.log each tool call, text and thinking block as it streams, and the CLI's stderr + .result.md the answer + .done how it ended (answered, raised, timed out, interrupted); written last + +:func:`turns_in_flight` reads the ones that have not ended. +""" + +from __future__ import annotations + +import abc +import copy +import dataclasses +import inspect +import itertools +import json +import os +import socket +import threading +import time +from contextvars import ContextVar +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Callable, Mapping + +import anyio + +from agent_flow.backends.claude_code import cli_stderr_to + +from ..substrate.scheduler import read_json, write_json_atomic, write_text_atomic + + +class TurnExecutionError(RuntimeError): + """A turn could not be delivered or could not be answered.""" + + +class TurnTimeout(TurnExecutionError): + """Nothing answered the turn inside its deadline.""" + + +#: How long a turn may take when neither the caller nor the role says. +DEFAULT_TURN_TIMEOUT_S = 3 * 60 * 60 + +#: Per-role turn timeouts, in seconds. One bound for every role is wrong in +#: both directions: a core assembly for a new architecture (modeling, weights, +#: routing, golden and long-context proofs) runs well past 3 h, while a stuck +#: plan review should not wait as long as an assembly. +DEFAULT_ROLE_TURN_TIMEOUTS: dict[str, float] = { + "plan_drafter": 4 * 3600.0, + "plan_reviewer": 3 * 3600.0, + "smith": 6 * 3600.0, + "assembler": 16 * 3600.0, + "reviewer": 4 * 3600.0, + # One candidate: implement, test on GPU, micro-bench -- hours of spool jobs. + "tuner": 8 * 3600.0, + "planner": 2 * 3600.0, +} + +#: Where in-process turns leave their files, under the workspace. +TURNS_DIRNAME = "turns" + +#: One activity entry's text is cut at this many characters. +ACTIVITY_ENTRY_LIMIT = 4000 + + +@dataclass(frozen=True) +class TurnRequest: + """What one role is being asked to do. + + ``system_prompt`` travels **with** the request: a transport that hands the + turn to another session must deliver it, because a driver given only the + task prompt runs a different agent and still returns a plausible answer. + ``decision_file`` is the exact path the role must write its typed decision + to. ``backend``, ``model`` and ``reasoning_effort`` are set only when + ``task.yaml`` routes the role to a particular agent; a transport that + hands the turn on passes them along, and empty means "run it on the + driving agent's own model and effort". ``work`` names the work when + ``what`` does not. + """ + + role: str + what: str + prompt: str + system_prompt: str = "" + decision_file: str | None = None + backend: str | None = None + model: str | None = None + reasoning_effort: str | None = None + #: The work the turn is for (a graph task's id), when ``what`` alone does + #: not name it: two tasks verifying one entry are both asked ``verify + #: ``. A transport adopts a turn a previous process left only for + #: the same work. Empty: ``what`` names it. + work: str = "" + + +@dataclass(frozen=True) +class TurnOutcome: + """The role's prose answer. **Not a control-plane input**: decisions are typed.""" + + role: str + text: str + #: Which executor ran it, for the journal. + transport: str = "" + #: The transport's own id for this turn, where it has one. + request_id: str = "" + + +class TurnExecutor(abc.ABC): + """Deliver one turn and return its answer.""" + + name = "turn-executor" + + @abc.abstractmethod + def run(self, request: TurnRequest) -> TurnOutcome: + """Deliver ``request`` and block until it is answered.""" + + def bound(self, session: Any) -> "TurnExecutor": + """A view of this executor pinned to one caller-owned session. + + A batch that builds its own agent per item calls this unconditionally; + a transport with no session concept (a mailbox) ignores it. + """ + return self + + +#: Given a request, the callable that runs it: ``layer(prompt) -> str``, in +#: practice an :class:`agent_flow.AgentLayer`. +LayerProvider = Callable[[TurnRequest], Any] + + +def _now_text(epoch: float | None = None) -> str: + moment = datetime.fromtimestamp(time.time() if epoch is None else epoch, timezone.utc) + return moment.isoformat(timespec="seconds") + + +def _process() -> dict: + return {"pid": os.getpid(), "host": socket.gethostname()} + + +def _seq(name: str) -> int: + head = name.split("-", 1)[0] + return int(head) if head.isdigit() else 0 + + +def _describe(event: Any) -> str: + """One streamed event as text: a tool call's name and input, a block's text, or its fields.""" + name, payload = getattr(event, "name", None), getattr(event, "input", None) + if isinstance(name, str) and payload is not None: + return f"{name} {json.dumps(payload, ensure_ascii=False, default=str)}" + text = getattr(event, "text", None) + if isinstance(text, str): + return text + if dataclasses.is_dataclass(event) and not isinstance(event, type): + return json.dumps(dataclasses.asdict(event), ensure_ascii=False, default=str) + return str(event) + + +class TurnRecord: + """The files one in-process turn leaves, written as it goes. + + Bookkeeping only: a write that fails (a full filesystem) never fails the + turn it describes. + """ + + def __init__(self, root: Path, request_id: str) -> None: + self.request_id = request_id + self.request_file = root / f"{request_id}.request.json" + self.system_file = root / f"{request_id}.system.md" + self.prompt_file = root / f"{request_id}.prompt.md" + self.activity_file = root / f"{request_id}.activity.log" + self.result_file = root / f"{request_id}.result.md" + self.done_file = root / f"{request_id}.done" + self.started = time.time() + self._lock = threading.Lock() + + def begin(self, request: TurnRequest, timeout_s: float) -> None: + """The request, its prompts and the empty activity log, before the turn starts.""" + try: + write_text_atomic(self.system_file, request.system_prompt) + write_text_atomic(self.prompt_file, request.prompt) + self.activity_file.touch() + write_json_atomic( + self.request_file, + { + "id": self.request_id, + "role": request.role, + "what": request.what, + "decision_file": request.decision_file, + "agent": { + "backend": request.backend, + "model": request.model, + "reasoning_effort": request.reasoning_effort, + }, + **_process(), + "started_at": self.started, + "timeout_s": timeout_s, + "system_prompt_file": str(self.system_file), + "prompt_file": str(self.prompt_file), + "activity_file": str(self.activity_file), + "result_file": str(self.result_file), + "done_file": str(self.done_file), + }, + ) + except OSError: + pass + + def note(self, kind: str, text: str) -> None: + """Append one entry to the activity log: its time and kind, then its text, indented.""" + if len(text) > ACTIVITY_ENTRY_LIMIT: + cut = len(text) - ACTIVITY_ENTRY_LIMIT + text = f"{text[:ACTIVITY_ENTRY_LIMIT]} ... ({cut} more characters)" + first, *rest = text.splitlines() or [""] + entry = "".join( + [f"{_now_text()[11:19]} {kind}: {first}\n", *(f" {line}\n" for line in rest)] + ) + with self._lock: + try: + with self.activity_file.open("a", encoding="utf-8") as handle: + handle.write(entry) + except OSError: + pass + + def activity(self, kind: str, event: Any) -> None: + """One event the layer streamed; a subagent's carries its label.""" + label = getattr(event, "agent_label", None) + self.note(f"{kind} [{label}]" if label else kind, _describe(event)) + + def finish(self, outcome: str, *, text: str | None = None, error: str = "") -> None: + """The answer, when there is one, then ``.done``: how the turn ended.""" + try: + if text is not None: + write_text_atomic(self.result_file, text) + ended = time.time() + write_json_atomic( + self.done_file, + { + "outcome": outcome, + "error": error, + "ended_at": ended, + "elapsed_s": round(ended - self.started, 1), + }, + ) + except OSError: + pass + + +class _TurnLog: + """Where an executor and its bound views keep their records, and the ids they share.""" + + def __init__(self) -> None: + self.root: Path | None = None + self._lock = threading.Lock() + self._counter: Any = None + + def keep_in(self, root: Path | str) -> None: + # Absolute, so the paths a record names hold from any working directory. + root = Path(root).absolute() + with self._lock: + if self.root == root: + return + root.mkdir(parents=True, exist_ok=True) + seen = [_seq(p.name) for p in root.glob("*.request.json")] + self.root = root + # Continued across a resume, so a later process never overwrites a record. + self._counter = itertools.count(max(seen, default=0) + 1) + + def open(self, request: TurnRequest, timeout_s: float) -> TurnRecord | None: + with self._lock: + if self.root is None: + return None + request_id = f"{next(self._counter):03d}-{request.role}" + root = self.root + record = TurnRecord(root, request_id) + record.begin(request, timeout_s) + return record + + +#: The record of the turn running in this context. A layer called in a turn's +#: thread streams its events in that thread's context, so each event reaches +#: its own turn's record however many turns of one layer run at once. +_ACTIVE: ContextVar[TurnRecord | None] = ContextVar("in_process_turn_record", default=None) + + +class _ActivityRouter: + """A layer's ``on_activity``: each event to the record of the turn it belongs to, then on.""" + + def __init__(self, previous: Callable[[str, Any], None] | None) -> None: + self.previous = previous + + def __call__(self, kind: str, event: Any) -> None: + record = _ACTIVE.get() + if record is not None: + record.activity(kind, event) + if self.previous is not None: + self.previous(kind, event) + + +_ROUTER_LOCK = threading.Lock() + + +def _route_activity(layer: Any) -> None: + """Give ``layer`` the observer that writes its events to the running turn's record, once. + + A layer is configured by whoever built it, and its observer is part of + that configuration; one it already had still sees every event, after the + record. A layer with no dataclass ``config.on_activity`` streams nothing + here. + """ + with _ROUTER_LOCK: + config = getattr(layer, "config", None) + if ( + config is None + or not dataclasses.is_dataclass(config) + or not hasattr(config, "on_activity") + or isinstance(config.on_activity, _ActivityRouter) + ): + return + try: + layer.config = dataclasses.replace( + config, on_activity=_ActivityRouter(config.on_activity) + ) + except (AttributeError, TypeError): + return + + +def _cancellable(layer: Any) -> bool: + """Whether a timeout can stop ``layer``'s turn: a coroutine ``aforward`` on a fresh session. + + Every :class:`agent_flow.AgentLayer` the workflow builds is one. A layer + that keeps a persistent session runs it on a loop of its own, which a + turn's loop must not borrow; it, and a plain callable, run unbounded. + """ + session = getattr(getattr(layer, "config", None), "session", None) + if getattr(session, "mode", None) == "persistent": + return False + return inspect.iscoroutinefunction(getattr(layer, "aforward", None)) + + +#: What a bounded call returns when its deadline cancelled the turn. +_EXPIRED = object() + + +class InProcessTurnExecutor(TurnExecutor): + """Run a turn by calling the role's layer in this process. + + Every request is bounded by its role's timeout, as on the external + transport: without ``timeout_s`` each role has its default + (:data:`DEFAULT_ROLE_TURN_TIMEOUTS`), with it every role has that one + bound. A turn past its bound is cancelled -- the CLI it started is + stopped -- and :class:`TurnTimeout` is raised, so a caller handles it as it + handles an external turn nobody answered. Once :meth:`keep_records` names a + directory, each turn leaves its files there (the module docstring). + """ + + name = "inprocess" + + def __init__( + self, + layer_provider: LayerProvider, + *, + timeout_s: float | None = None, + role_timeouts: Mapping[str, float] | None = None, + records: Path | str | None = None, + ) -> None: + self._provider = layer_provider + self.timeout_s = DEFAULT_TURN_TIMEOUT_S if timeout_s is None else timeout_s + timeouts = dict(DEFAULT_ROLE_TURN_TIMEOUTS) if timeout_s is None else {} + timeouts.update(role_timeouts or {}) + self.role_timeouts = timeouts + self._log = _TurnLog() + if records is not None: + self._log.keep_in(records) + + def timeout_for(self, role: str) -> float: + """How long a turn of ``role`` may take before the run gives up on it.""" + return float(self.role_timeouts.get(role, self.timeout_s)) + + def keep_records(self, root: Path | str) -> None: + """Leave each turn's files in ``root`` from now on, numbered on from those already there. + + Shared with every view :meth:`bound` returns. + """ + self._log.keep_in(root) + + @property + def records(self) -> Path | None: + """Where turns leave their files, or ``None`` before :meth:`keep_records`.""" + return self._log.root + + def bound(self, session: Any) -> "InProcessTurnExecutor": + """Pin every turn to ``session``, whatever the request asks for.""" + view = copy.copy(self) + view._provider = lambda request: session + return view + + def run(self, request: TurnRequest) -> TurnOutcome: + layer = self._provider(request) + if layer is None: + raise TurnExecutionError(f"no layer available for role {request.role!r}") + timeout_s = self.timeout_for(request.role) + record = self._log.open(request, timeout_s) + if record is not None: + _route_activity(layer) + token = _ACTIVE.set(record) + outcome, text, error = "raised", None, "" + try: + with cli_stderr_to(_stderr_sink(record)): + answer = self._call(layer, request, timeout_s, record) + text = answer if isinstance(answer, str) else str(answer) + outcome = "answered" + except TurnTimeout as exc: + outcome, error = "timed out", str(exc) + raise + except BaseException as exc: + outcome = "interrupted" if isinstance(exc, KeyboardInterrupt) else "raised" + error = f"{type(exc).__name__}: {exc}" + raise + finally: + _ACTIVE.reset(token) + if record is not None: + record.finish(outcome, text=text, error=error) + return TurnOutcome( + role=request.role, + text=text, + transport=self.name, + request_id=record.request_id if record is not None else "", + ) + + @staticmethod + def _call(layer: Any, request: TurnRequest, timeout_s: float, record: TurnRecord | None) -> Any: + """``layer(prompt)``, cancelled at ``timeout_s`` where the layer can be.""" + if not _cancellable(layer): + return layer(request.prompt) + + async def bounded() -> Any: + with anyio.move_on_after(timeout_s): + return await layer.aforward(request.prompt) + return _EXPIRED + + answer = anyio.run(bounded) + if answer is _EXPIRED: + name = record.request_id if record is not None else request.role + where = f" Its activity is in {record.activity_file}." if record is not None else "" + raise TurnTimeout( + f"{name}: no answer within {timeout_s:g}s (the {request.role} turn timeout); " + f"the turn was stopped.{where}" + ) + return answer + + +def _stderr_sink(record: TurnRecord | None) -> Callable[[str], None]: + """Where a turn's CLI stderr lines go: its activity log, or nowhere without a record.""" + if record is None: + return lambda line: None + return lambda line: record.note("stderr", line) + + +def turns_in_flight(root: Path | str, now: float | None = None) -> list[dict]: + """The in-process turns recorded under ``root`` that have not ended, oldest first. + + Each is its ``.request.json`` with ``age_s``; ``overdue`` once past its + bound; and ``alive``: whether the process running it still exists, ``None`` + when it ran on another host. A turn whose process is gone never ends. + """ + now = time.time() if now is None else now + host = socket.gethostname() + out = [] + for path in sorted(Path(root).glob("*.request.json"), key=lambda p: _seq(p.name)): + data = read_json(path) + if data is None: + continue + request_id = str(data.get("id") or path.name[: -len(".request.json")]) + if (path.parent / f"{request_id}.done").exists(): + continue + age = now - float(data.get("started_at") or now) + timeout = float(data.get("timeout_s") or 0) + alive = None + if data.get("host") == host: + alive = _pid_alive(int(data.get("pid") or 0)) + out.append( + { + **data, + "id": request_id, + "age_s": round(age, 1), + "overdue": bool(timeout) and age > timeout, + "alive": alive, + } + ) + return out + + +def _pid_alive(pid: int) -> bool: + if pid <= 0: + return False + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +__all__ = [ + "ACTIVITY_ENTRY_LIMIT", + "DEFAULT_ROLE_TURN_TIMEOUTS", + "DEFAULT_TURN_TIMEOUT_S", + "TURNS_DIRNAME", + "InProcessTurnExecutor", + "LayerProvider", + "TurnExecutionError", + "TurnExecutor", + "TurnOutcome", + "TurnRecord", + "TurnRequest", + "TurnTimeout", + "turns_in_flight", +] diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/gpu_slot.sh b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/gpu_slot.sh new file mode 100755 index 000000000000..a5e1e6a44369 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/gpu_slot.sh @@ -0,0 +1,51 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Run a command on N GPUs. +# +# spoolctl.py submit free- 'bash /gpu_slot.sh N ' +# spoolctl.py submit --gpus N '' # the same thing, preferred +# +# Inside a spool job the scheduler has already allocated the devices: it reads +# `gpu_slot.sh N` in the job body as `gpus=N`, exports CUDA_VISIBLE_DEVICES to +# exactly those devices and removes the allocation's MPI context. This script +# then only checks that the job holds enough devices and runs the command. +# +# Outside a spool job it is a convenience wrapper: it removes the MPI context, +# pins CUDA_VISIBLE_DEVICES to the first N devices if unset, and applies a +# deadline (GPU_SLOT_TIMEOUT seconds, default 7200, 0 = none). +ulimit -c 0 +set -u +N=${1:?usage: gpu_slot.sh } +shift + +if [ -n "${MV2FLOW_SPOOL_JOB:-}" ]; then + held=0 + [ -n "${MV2FLOW_SPOOL_DEVICES:-}" ] && held=$(awk -F, '{print NF}' <<< "$MV2FLOW_SPOOL_DEVICES") + if [ "$N" != all ] && [ "$held" -lt "$N" ]; then + echo "gpu_slot: job $MV2FLOW_SPOOL_JOB holds $held device(s) but asks for $N; submit it with --gpus $N" >&2 + exit 64 + fi + echo "gpu_slot: CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-} (allocated by the spool scheduler)" + if [ -n "${GPU_SLOT_TIMEOUT:-}" ] && [ "${GPU_SLOT_TIMEOUT}" != 0 ]; then + exec timeout -k 60 "$GPU_SLOT_TIMEOUT" "$@" + fi + exec "$@" +fi + +for v in $(env | sed -n 's/^\(\(PMI\|PMIX\|OMPI\|SLURM\|PRTE\|PALS\|MPI_LOCALRANKID\|MPI_LOCALNRANKS\)[A-Za-z0-9_]*\)=.*/\1/p'); do + unset "$v" +done +if [ -z "${CUDA_VISIBLE_DEVICES:-}" ] && [ "$N" != all ]; then + CUDA_VISIBLE_DEVICES=$(seq -s, 0 $((N - 1))) + export CUDA_VISIBLE_DEVICES +fi +echo "gpu_slot: not inside a spool job; no slot accounting. CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-}" >&2 +if [ "${GPU_SLOT_TIMEOUT:-7200}" != 0 ]; then + timeout -k 60 "${GPU_SLOT_TIMEOUT:-7200}" "$@" + rc=$? + [ "$rc" = 124 ] && echo "gpu_slot: TIMED OUT after ${GPU_SLOT_TIMEOUT:-7200}s -- the job was stopped, not failed" >&2 + exit $rc +fi +exec "$@" diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/incremental_build.sh b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/incremental_build.sh new file mode 100755 index 000000000000..898ac5bcb191 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/incremental_build.sh @@ -0,0 +1,140 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Incremental native rebuild of a TensorRT-LLM checkout after a C++/CUDA change, +# without crashing processes that have the old libraries mapped. +# +# incremental_build.sh [--repo DIR] [--spool DIR]... [--cuda-archs '100-real;103-real'] +# [--log-dir DIR] [--no-import-check] [-- ] +# +# Run it as a whole-node spool job, so nothing else on that session runs during +# the build: +# +# spoolctl.py submit --gpus all --wall 3h build- \ +# 'bash /incremental_build.sh --repo --spool --spool ' +# +# What it does, in order: +# 1. takes a per-checkout lock, so two builds never run over one cpp/build; +# 2. announces the rebuild on every listed spool: each session's scheduler +# holds new GPU jobs against this checkout until the build ends; +# 3. gives every installed shared library a fresh inode (copy + rename), so a +# running process keeps the old file and build_wheel.py's in-place rewrite +# lands on a file nobody has mapped; +# 4. runs build_wheel.py --skip_building_wheel with ccache, from a venv of +# the build's own (/.dev-build/build-venv) that sees the run's +# packages through a .pth: build_wheel.py pip-installs its development +# requirements into the venv it runs in, and the run's venv is the one +# every gate imports from, so a build must never change it (it fails if +# the run's venv changed under it); +# 5. checks the checkout still imports, and retracts the announcement. +# +# Defaults, all overridable by flag: +# --repo $MV2FLOW_REPO, else the git checkout containing the current directory +# --spool $MV2FLOW_SESSION_SPOOLS (':'-separated), plus $MV2FLOW_SPOOL (this job's own) +# --cuda-archs $MV2FLOW_CUDA_ARCHS, else build_wheel.py's own default +# --log-dir $MV2FLOW_BUILD_LOG_DIR, else /.dev-build/logs +# build python $MV2FLOW_BUILD_PYTHON, else python3 (the venv on PATH): the +# venv the build's own is made beside; MV2FLOW_BUILD_ISOLATED=0 +# builds in it directly +# +# The build needs the user site (the CUTLASS configure step imports +# cutlass_library from it), so PYTHONNOUSERSITE and PYTHONPATH are unset for +# the build and only the import check runs with PYTHONNOUSERSITE=1. +set -u + +HERE=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +REBUILD="$HERE/../substrate/rebuild.py" +PY=${MV2FLOW_PYTHON:-python3} + +REPO=${MV2FLOW_REPO:-} +ARCHS=${MV2FLOW_CUDA_ARCHS:-} +LOG_DIR=${MV2FLOW_BUILD_LOG_DIR:-} +IMPORT_CHECK=1 +SPOOLS=() +EXTRA=() +while [ $# -gt 0 ]; do + case "$1" in + --repo) REPO=$2; shift 2 ;; + --spool) SPOOLS+=("$2"); shift 2 ;; + --cuda-archs) ARCHS=$2; shift 2 ;; + --log-dir) LOG_DIR=$2; shift 2 ;; + --no-import-check) IMPORT_CHECK=0; shift ;; + --) shift; EXTRA=("$@"); break ;; + *) echo "incremental_build: unknown argument $1" >&2; exit 64 ;; + esac +done + +if [ -z "$REPO" ]; then + REPO=$(git rev-parse --show-toplevel 2>/dev/null) || { echo "incremental_build: no --repo and not inside a git checkout" >&2; exit 64; } +fi +REPO=$(cd "$REPO" && pwd) +[ -f "$REPO/scripts/build_wheel.py" ] || { echo "incremental_build: $REPO has no scripts/build_wheel.py" >&2; exit 64; } + +if [ ${#SPOOLS[@]} -eq 0 ] && [ -n "${MV2FLOW_SESSION_SPOOLS:-}" ]; then + IFS=: read -r -a SPOOLS <<< "$MV2FLOW_SESSION_SPOOLS" +fi +if [ -n "${MV2FLOW_SPOOL:-}" ]; then + SPOOLS+=("$MV2FLOW_SPOOL") +fi +SPOOL_ARGS=() +for s in "${SPOOLS[@]}"; do + [ -n "$s" ] && SPOOL_ARGS+=(--spool "$s") +done + +LOG_DIR=${LOG_DIR:-$REPO/.dev-build/logs} +LOCK=${MV2FLOW_BUILD_LOCK:-$REPO/.dev-build/native-build.lock} +mkdir -p "$LOG_DIR" "$(dirname "$LOCK")" +LOG=$LOG_DIR/incremental-$(date +%Y%m%d-%H%M%S).log + +( + flock 8 + echo "=== incremental build start $(date -Is) on $(hostname) repo=$REPO ===" + TOKEN=$("$PY" "$REBUILD" begin --tree "$REPO" "${SPOOL_ARGS[@]}" --note "incremental build on $(hostname)") + trap '"$PY" "$REBUILD" end --token "$TOKEN" "${SPOOL_ARGS[@]}"; echo "=== rebuild announcement $TOKEN retracted ==="' EXIT + echo "=== rebuild announced as $TOKEN on ${#SPOOLS[@]} spool(s): ${SPOOLS[*]:-none} ===" + git -C "$REPO" log --oneline -1 + git -C "$REPO" status --porcelain -uno -- cpp | head -40 + "$PY" "$REBUILD" refresh-inodes --root "$REPO/tensorrt_llm" || exit 1 + ( + unset PYTHONNOUSERSITE PYTHONPATH + export CCACHE_DIR=${CCACHE_DIR:-$REPO/.dev-build/ccache} + export PIP_CACHE_DIR=${PIP_CACHE_DIR:-$REPO/.dev-build/pip-cache} + cd "$REPO" || exit 1 + start=$(date +%s) + ARCH_ARGS=() + [ -n "$ARCHS" ] && ARCH_ARGS=(--cuda_architectures "$ARCHS") + RUN_PY=${MV2FLOW_BUILD_PYTHON:-python3} + BUILD_PY=$RUN_PY + run_lib=$("$RUN_PY" -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])') || exit 1 + installed() { ls -1 "$run_lib" 2>/dev/null | grep -E '\.(dist|egg)-info$' | sort; } + before=$(installed) + if [ "${MV2FLOW_BUILD_ISOLATED:-1}" = 1 ]; then + BUILD_VENV=$REPO/.dev-build/build-venv + if [ ! -x "$BUILD_VENV/bin/python" ]; then + "$RUN_PY" -m venv --system-site-packages "$BUILD_VENV" || exit 1 + fi + own_lib=$("$BUILD_VENV/bin/python" -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])') || exit 1 + echo "$run_lib" > "$own_lib/mv2flow-run-venv.pth" + BUILD_PY=$BUILD_VENV/bin/python + echo "=== building from $BUILD_VENV (the run's packages through a .pth) ===" + fi + "$BUILD_PY" ./scripts/build_wheel.py --use_ccache --skip_building_wheel \ + "${ARCH_ARGS[@]}" --yes "${EXTRA[@]}" + rc=$? + echo "INCREMENTAL_BUILD_WALL_SECONDS=$(( $(date +%s) - start )) rc=$rc" + if [ "$(installed)" != "$before" ]; then + echo "incremental_build: the run's venv ($run_lib) changed during the build:" >&2 + diff <(echo "$before") <(installed) >&2 + [ "$rc" = 0 ] && rc=65 + fi + exit $rc + ) + rc=$? + if [ "$rc" = 0 ] && [ "$IMPORT_CHECK" = 1 ]; then + (cd "$REPO" && PYTHONNOUSERSITE=1 "${MV2FLOW_BUILD_PYTHON:-python3}" -c \ + "import tensorrt_llm; print('import ok', tensorrt_llm.__version__, tensorrt_llm.__file__)") || rc=$? + fi + exit $rc +) 8>"$LOCK" 2>&1 | tee "$LOG" +exit "${PIPESTATUS[0]}" diff --git a/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/mn/README.md b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/mn/README.md new file mode 100644 index 000000000000..84efd1046733 --- /dev/null +++ b/agent-flow/agent_flow/workflows/trtllm_modeling_v2/infra/scripts/mn/README.md @@ -0,0 +1,66 @@ +# Multi-node launches on a multi-node compute session (optional) + +A run that needs ranks on more than one node declares a session with `nodes: 2` (or more): + +```yaml +environment: + sessions: + - name: mn + spool: /shared/run/spool-mn + role: gates + nodes: 2 + gpus_per_node: 4 + time_limit: "12:00:00" +``` + +`trtllm-modeling-v2 session render --task task.yaml --name mn` writes `/session_entry.sh`. +Submit the allocation with `--nodes=2 --gpus-per-node=4` (and `--segment=2` where the cluster uses +block topology) and **no** `--ntasks-per-node`, so the batch step runs the entry once per node, each +copy in its own container. `trtllm-modeling-v2 session submission --task task.yaml --name mn` prints +the broker call. Each copy becomes a spool scheduler: + +| node | spool | used for | +|---|---|---| +| 0 | `` | the session's spool: submit here | +| k | `-n` | node k's spool: hold jobs and the remote MPI daemons | + +Each node registers `/hosts/`, which is how the launcher finds the other nodes. + +## Launching across the nodes + +```bash +spoolctl.py submit --spool --gpus all --tree \ + 'bash /mn/mn_mpirun.sh 4 python