From d4ab0ad733dbba41c241557b5c051c0803514249 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 12 Oct 2026 00:55:45 -0700 Subject: [PATCH 01/70] [TRTLLM-16460][feat] staircase: staircase-onboard workflow: plan, assemble, gate and deliver a model onboarding The staircase-onboard agent-flow workflow: a plan phase that surveys the checkpoint and drafts a staircase of units, per-unit assembly by an assembler agent with kernel and modeling smiths, boot/accuracy/acceptance/witness gates at every parallel member, a bounded remediation cycle, and delivery checks. Includes the agent-flow runtime package (turn and command executors: in-process, file dispatch, local, slurm, spool, delegated) that the workflow routes its turns and gate commands through, the role prompts, the materialized skills, and the package's tests. Squashed from user/williamj/staircase-polish @ c8109d6f90 (42 commits): - [feat] agent-flow: staircase-onboard workflow - [refactor] agent-flow: runtime package for turns and commands - [refactor] agent-flow: route staircase turns and gate commands through the runtime - [feat] staircase: close UNCOVERED with a bounded remediation cycle - [feat] agent-flow: run gates on a cluster only an agent can reach - [feat] staircase: gate every member, survey the demands, admit a new architecture - [fix] staircase: a checkpoint carrying its own implementation is an unsupported architecture - [fix] staircase: the accuracy gate cannot spell the topology the boot gate spells - [fix] staircase: the acceptance gate measured a speculator the unit was not built for - [fix] staircase: make the multi-rank launch and the supply survey actually work - [fix] staircase: the acceptance gate needs a stock path, which a new architecture does not have - [fix] staircase: see the registry's own unsupported-architecture message - [perf] staircase: stop spending assembler turns on restarts, and make a gate say what it measured - [feat] staircase: recover six improvements stranded in a sibling checkout - [fix] staircase: a member overlay must not be written into the source tree - [doc] staircase: carry the findings a real run produced - [chore] staircase: classify the in-tree Triton custom op - [feat] staircase: host offload names its capturable precedent - [feat] staircase: dspark, bounded-replay and kv_quant feature kinds - [fix] staircase: a gate test must not find the checkout that holds TMPDIR - [fix] staircase: the survey must read the in-tree catalog index - [fix] staircase: an unknown model_type is a new architecture too - [feat] staircase: defer feature-only smith failures to their unit - [fix] staircase: a resume collects the answers the dead process never read - [fix] staircase-onboard: commit.sh signs off (DCO) under git's own identity - [fix] staircase-onboard: kernel-source registry claims the run's Triton quantization ops - [fix] staircase-onboard: disagg equivalence witness runs a real, synchronised leg B (F021a-g) - [fix] staircase-onboard: a registration shim without a target counts as an unsupported architecture - [fix] staircase-onboard: an assembler blocked on knowledge_gap or environment can be re-entered - [fix] staircase-onboard: the acceptance gate measures a declared long-answer workload - [fix] staircase-onboard: token-comparing witness legs can run at batch 1 - [fix] staircase-onboard: the equivalence witness hands leg B an absolute prompts path - [feat] staircase-onboard: witness blocks can set an environment for both legs - [feat] staircase: ask targets for their runtime-value dump on every gate and witness run - [chore] agent-flow: reword two codespell findings in staircase_onboard - [chore] agent-flow: one-line docstring summaries in the staircase_onboard tests - [chore] agent-flow: NVIDIA copyright headers on the staircase_onboard files - [fix] staircase-onboard: kernel-source registry claims the four new Triton custom ops - [chore] Apply pre-commit formatting to the DeepSeek-V4.1 changes - [feat] staircase: place targets at models//____/ - [chore] staircase: stop tracking the sandbox's session placeholders - [fix] staircase: no committed file names one operator's home or scratch path Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- agent-flow/agent_flow/runtime/__init__.py | 70 + .../agent_flow/runtime/commands/__init__.py | 24 + .../agent_flow/runtime/commands/base.py | 140 + .../agent_flow/runtime/commands/delegated.py | 204 ++ .../agent_flow/runtime/commands/local.py | 115 + .../agent_flow/runtime/commands/slurm.py | 45 + .../agent_flow/runtime/commands/spool.py | 111 + .../{runtime.py => runtime/portal.py} | 12 +- .../agent_flow/runtime/turns/__init__.py | 29 + agent-flow/agent_flow/runtime/turns/base.py | 126 + .../agent_flow/runtime/turns/file_dispatch.py | 332 ++ .../agent_flow/runtime/turns/inprocess.py | 72 + .../workflows/staircase_onboard/.gitignore | 5 + .../staircase_onboard/BUILD_AND_TEST.md | 245 ++ .../workflows/staircase_onboard/README.md | 278 ++ .../staircase_onboard/TROUBLESHOOTING.md | 543 +++ .../workflows/staircase_onboard/__init__.py | 90 + .../workflows/staircase_onboard/agent_home.py | 204 ++ .../workflows/staircase_onboard/cli.py | 285 ++ .../workflows/staircase_onboard/committer.py | 243 ++ .../workflows/staircase_onboard/delivery.py | 463 +++ .../staircase_onboard/environment.py | 779 ++++ .../workflows/staircase_onboard/execution.py | 219 ++ .../20260922-F001-F003-F004-fix-report.md | 325 ++ ...-F001-substrate-gate-needs-a-stock-path.md | 85 + .../findings/20260922-F002-fix-report.md | 160 + ...02-unified-targets-are-gated-at-tp-only.md | 89 + ...e-survey-misses-four-structural-demands.md | 102 + ...o-cannot-name-a-sibling-engine-checkout.md | 87 + .../findings/20260922-F005-fix-report.md | 80 + ...tory-probe-takes-the-multirank-launcher.md | 123 + .../findings/20260922-F006-F007-fix-report.md | 173 + ...60922-F006-two-multirank-launch-defects.md | 118 + ...007-slurm-vars-defeat-the-mpirun-launch.md | 130 + ...plan-reviewer-cross-model-not-available.md | 88 + ...9-a-turn-can-commit-and-never-hand-back.md | 85 + ...read-once-and-a-gate-error-is-truncated.md | 115 + ...F011-the-gate-measures-the-working-tree.md | 90 + .../staircase_onboard/findings/README.md | 44 + .../RETROSPECTIVE-2026-09-mimo-v2.6-flash.md | 320 ++ .../workflows/staircase_onboard/gates.py | 1874 ++++++++++ .../workflows/staircase_onboard/journal.py | 155 + .../staircase_onboard/kernel_sources.yaml | 233 ++ .../workflows/staircase_onboard/preflight.py | 1014 ++++++ .../workflows/staircase_onboard/progress.py | 769 ++++ .../staircase_onboard/prompts/__init__.py | 201 ++ .../staircase_onboard/prompts/_common.py | 415 +++ .../staircase_onboard/prompts/assembler.py | 197 + .../staircase_onboard/prompts/auditor.py | 135 + .../staircase_onboard/prompts/plan_drafter.py | 228 ++ .../prompts/plan_reviewer.py | 105 + .../staircase_onboard/prompts/reviewer.py | 102 + .../staircase_onboard/prompts/runner.py | 68 + .../staircase_onboard/prompts/smith.py | 161 + .../staircase_onboard/remediation.py | 687 ++++ .../workflows/staircase_onboard/runtime.py | 496 +++ .../staircase_onboard/runtime_index.yaml | 677 ++++ .../staircase-kernel-feasibility/SKILL.md | 214 ++ .../references/derivation-checklist.md | 99 + .../references/mapping-and-triage.md | 83 + .../references/trtllm-ops.md | 423 +++ .../skills/staircase-kernel-onboard/SKILL.md | 125 + .../references/contract-evidence.md | 81 + .../references/decomposition.md | 64 + .../decompositions/fuse-input-embeds.md | 268 ++ .../decompositions/gdn-custom-op.md | 434 +++ .../decompositions/mla-custom-op.md | 255 ++ .../multi-head-ngram-embedding.md | 319 ++ .../references/entry-spec.md | 118 + .../skills/staircase-kernel-verify/SKILL.md | 90 + .../SKILL.md | 245 ++ .../skills/staircase-modeling-core/SKILL.md | 229 ++ .../references/adapter-rules.md | 84 + .../references/attention-sinks.md | 56 + .../references/expert-weight-packing.md | 574 +++ .../references/graph-and-sync.md | 88 + .../references/hybrid-recurrent-state.md | 128 + .../references/latent-kv-cache.md | 306 ++ .../references/routing-identity.md | 118 + .../references/runtime-interface.md | 309 ++ .../references/sliding-window-attention.md | 79 + .../references/step-args.md | 178 + .../skills/staircase-modeling-disagg/SKILL.md | 129 + .../skills/staircase-modeling-dwdp/SKILL.md | 77 + .../references/adapter-interface.md | 39 + .../references/runtime-boundary.md | 35 + .../skills/staircase-modeling-eplb/SKILL.md | 98 + .../references/adapter-interface.md | 45 + .../references/placement.md | 64 + .../references/runtime-boundary.md | 34 + .../references/weight-migration.md | 57 + .../staircase-modeling-host-offload/SKILL.md | 196 + .../staircase-modeling-kv-quant/SKILL.md | 138 + .../staircase-modeling-multimodal/SKILL.md | 244 ++ .../references/vision-path.md | 192 + .../staircase-modeling-parallel/SKILL.md | 95 + .../references/multi-rank.md | 267 ++ .../references/unified.md | 110 + .../staircase-modeling-speculator/SKILL.md | 119 + .../references/acceptance-gate.md | 193 + .../references/algorithms.md | 184 + .../references/multi-token-prediction.md | 300 ++ .../references/runtime-boundary.md | 488 +++ .../staircase_onboard/spool_watcher.sh | 144 + .../workflows/staircase_onboard/state.py | 324 ++ .../workflows/staircase_onboard/survey.py | 1311 +++++++ .../staircase_onboard/task.example.yaml | 564 +++ .../staircase_onboard/task_schema.py | 1740 +++++++++ .../workflows/staircase_onboard/transport.py | 109 + .../workflows/staircase_onboard/vocabulary.py | 598 +++ .../workflows/staircase_onboard/witness.py | 1438 ++++++++ .../workflows/staircase_onboard/workflow.py | 3191 +++++++++++++++++ agent-flow/pyproject.toml | 23 + agent-flow/tests/runtime/__init__.py | 2 + .../tests/runtime/test_command_executors.py | 306 ++ .../tests/runtime/test_delegated_executor.py | 207 ++ .../tests/runtime/test_turn_executors.py | 353 ++ .../workflows/staircase_onboard/__init__.py | 2 + .../workflows/staircase_onboard/conftest.py | 20 + .../staircase_onboard/test_agent_home.py | 154 + .../test_build_and_test_doc.py | 129 + .../staircase_onboard/test_committer.py | 331 ++ .../staircase_onboard/test_decision_files.py | 308 ++ .../staircase_onboard/test_delivery.py | 298 ++ .../staircase_onboard/test_dspark.py | 181 + .../staircase_onboard/test_environment.py | 654 ++++ .../staircase_onboard/test_execution.py | 421 +++ .../staircase_onboard/test_gate_wiring.py | 818 +++++ .../workflows/staircase_onboard/test_gates.py | 1032 ++++++ .../staircase_onboard/test_kernel_sources.py | 205 ++ .../staircase_onboard/test_macro_expansion.py | 231 ++ .../staircase_onboard/test_package_hygiene.py | 66 + .../staircase_onboard/test_preflight.py | 578 +++ .../staircase_onboard/test_prompts.py | 253 ++ .../staircase_onboard/test_remediation.py | 550 +++ .../test_runtime_and_scope.py | 669 ++++ .../staircase_onboard/test_runtime_index.py | 575 +++ .../staircase_onboard/test_skills.py | 326 ++ .../staircase_onboard/test_smith_deferral.py | 487 +++ .../test_smith_resume_adoption.py | 65 + .../workflows/staircase_onboard/test_state.py | 220 ++ .../staircase_onboard/test_substrate_gate.py | 377 ++ .../staircase_onboard/test_survey.py | 702 ++++ .../staircase_onboard/test_survey_probes.py | 504 +++ .../staircase_onboard/test_task_schema.py | 1116 ++++++ .../staircase_onboard/test_transport.py | 376 ++ .../staircase_onboard/test_troubleshooting.py | 149 + .../staircase_onboard/test_vocabulary.py | 420 +++ .../staircase_onboard/test_witness.py | 1140 ++++++ .../staircase_onboard/test_workflow.py | 1390 +++++++ 150 files changed, 47492 insertions(+), 1 deletion(-) create mode 100644 agent-flow/agent_flow/runtime/__init__.py create mode 100644 agent-flow/agent_flow/runtime/commands/__init__.py create mode 100644 agent-flow/agent_flow/runtime/commands/base.py create mode 100644 agent-flow/agent_flow/runtime/commands/delegated.py create mode 100644 agent-flow/agent_flow/runtime/commands/local.py create mode 100644 agent-flow/agent_flow/runtime/commands/slurm.py create mode 100644 agent-flow/agent_flow/runtime/commands/spool.py rename agent-flow/agent_flow/{runtime.py => runtime/portal.py} (86%) create mode 100644 agent-flow/agent_flow/runtime/turns/__init__.py create mode 100644 agent-flow/agent_flow/runtime/turns/base.py create mode 100644 agent-flow/agent_flow/runtime/turns/file_dispatch.py create mode 100644 agent-flow/agent_flow/runtime/turns/inprocess.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/.gitignore create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/BUILD_AND_TEST.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/README.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/TROUBLESHOOTING.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/__init__.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/agent_home.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/cli.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/committer.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/delivery.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/environment.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/execution.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F001-F003-F004-fix-report.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F001-substrate-gate-needs-a-stock-path.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F002-fix-report.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F002-unified-targets-are-gated-at-tp-only.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F003-surface-survey-misses-four-structural-demands.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F004-embargo-cannot-name-a-sibling-engine-checkout.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F005-fix-report.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F005-inventory-probe-takes-the-multirank-launcher.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F006-F007-fix-report.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F006-two-multirank-launch-defects.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F007-slurm-vars-defeat-the-mpirun-launch.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F008-plan-reviewer-cross-model-not-available.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F009-a-turn-can-commit-and-never-hand-back.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F010-task-yaml-is-read-once-and-a-gate-error-is-truncated.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/20260923-F011-the-gate-measures-the-working-tree.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/README.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/findings/RETROSPECTIVE-2026-09-mimo-v2.6-flash.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/gates.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/journal.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/kernel_sources.yaml create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/preflight.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/progress.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/prompts/__init__.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/prompts/_common.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/prompts/assembler.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/prompts/auditor.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_drafter.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_reviewer.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/prompts/reviewer.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/prompts/runner.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/prompts/smith.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/remediation.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/runtime.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/runtime_index.yaml create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/SKILL.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/references/derivation-checklist.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/references/mapping-and-triage.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/references/trtllm-ops.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/SKILL.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/contract-evidence.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decomposition.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/fuse-input-embeds.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/gdn-custom-op.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/mla-custom-op.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/multi-head-ngram-embedding.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/entry-spec.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/SKILL.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-bounded-replay/SKILL.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/SKILL.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/adapter-rules.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/attention-sinks.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/expert-weight-packing.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/graph-and-sync.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/hybrid-recurrent-state.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/latent-kv-cache.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/routing-identity.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/runtime-interface.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/sliding-window-attention.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/step-args.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-disagg/SKILL.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-dwdp/SKILL.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-dwdp/references/adapter-interface.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-dwdp/references/runtime-boundary.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/SKILL.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/adapter-interface.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/placement.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/runtime-boundary.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/weight-migration.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-host-offload/SKILL.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-kv-quant/SKILL.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-multimodal/SKILL.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-multimodal/references/vision-path.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-parallel/SKILL.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-parallel/references/multi-rank.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-parallel/references/unified.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/SKILL.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/acceptance-gate.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/algorithms.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/multi-token-prediction.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/runtime-boundary.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/spool_watcher.sh create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/state.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/survey.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/transport.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/vocabulary.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/witness.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/workflow.py create mode 100644 agent-flow/tests/runtime/__init__.py create mode 100644 agent-flow/tests/runtime/test_command_executors.py create mode 100644 agent-flow/tests/runtime/test_delegated_executor.py create mode 100644 agent-flow/tests/runtime/test_turn_executors.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/__init__.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/conftest.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_agent_home.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_build_and_test_doc.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_committer.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_decision_files.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_delivery.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_dspark.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_environment.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_execution.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_gate_wiring.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_gates.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_kernel_sources.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_macro_expansion.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_package_hygiene.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_preflight.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_prompts.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_remediation.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_runtime_and_scope.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_runtime_index.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_skills.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_smith_deferral.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_smith_resume_adoption.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_state.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_substrate_gate.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_survey.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_survey_probes.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_task_schema.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_transport.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_troubleshooting.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_vocabulary.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_witness.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_workflow.py diff --git a/agent-flow/agent_flow/runtime/__init__.py b/agent-flow/agent_flow/runtime/__init__.py new file mode 100644 index 000000000000..b30442272ad2 --- /dev/null +++ b/agent-flow/agent_flow/runtime/__init__.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Execution substrate shared by every workflow. + +Two questions live here, and neither of them is a workflow's business: + +* **how does one agent turn get delivered and answered** — in this process + through an ``AgentLayer``, or through a filesystem mailbox a separate + authenticated session drains; +* **where and how does one command run** — a local subprocess, an ``srun`` + step, or a file dropped into a spool a long-lived allocation watches. + +A workflow states *what* it needs run and *what* it needs proven. It should +not carry a branch per cluster protocol, and it should not have to change +when a credential store turns read-only. Everything in this package is +therefore free of workflow vocabulary: nothing here knows what a catalog +entry, a smith or a gate is. + +A third concern, older than the other two, already lived at this import +path: running async work on a background loop for a synchronous caller. It +is the same kind of thing — substrate a workflow uses without describing — +so it stayed, as :mod:`~agent_flow.runtime.portal`, and +``from agent_flow.runtime import PortalRunner`` still resolves. +""" + +from .commands import ( + CommandExecutor, + CommandSpec, + DelegatedCommandExecutor, + ExecutionCapabilities, + ExecutionRecord, + LocalCommandExecutor, + SlurmStepExecutor, + SpoolCommandExecutor, + SubmissionRenderer, +) +from .portal import PortalRunner, build_request +from .turns import ( + DEFAULT_TURN_TIMEOUT_S, + FileDispatchTurnExecutor, + InProcessTurnExecutor, + TurnExecutionError, + TurnExecutor, + TurnOutcome, + TurnRequest, + TurnTimeout, +) + +__all__ = [ + "CommandExecutor", + "CommandSpec", + "DEFAULT_TURN_TIMEOUT_S", + "DelegatedCommandExecutor", + "SubmissionRenderer", + "ExecutionCapabilities", + "ExecutionRecord", + "FileDispatchTurnExecutor", + "InProcessTurnExecutor", + "LocalCommandExecutor", + "PortalRunner", + "SlurmStepExecutor", + "SpoolCommandExecutor", + "TurnExecutionError", + "TurnExecutor", + "TurnOutcome", + "TurnRequest", + "TurnTimeout", + "build_request", +] diff --git a/agent-flow/agent_flow/runtime/commands/__init__.py b/agent-flow/agent_flow/runtime/commands/__init__.py new file mode 100644 index 000000000000..670da8665474 --- /dev/null +++ b/agent-flow/agent_flow/runtime/commands/__init__.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Command execution: one interface, one per place a command can land.""" + +from .base import CommandExecutor, CommandSpec, ExecutionCapabilities, ExecutionRecord +from .delegated import DelegatedCommandExecutor, SubmissionRenderer +from .local import NOT_FOUND_RETURNCODE, TIMEOUT_RETURNCODE, LocalCommandExecutor +from .slurm import SlurmStepExecutor +from .spool import SpoolCommandExecutor + +__all__ = [ + "CommandExecutor", + "CommandSpec", + "DelegatedCommandExecutor", + "ExecutionCapabilities", + "ExecutionRecord", + "LocalCommandExecutor", + "NOT_FOUND_RETURNCODE", + "SlurmStepExecutor", + "SubmissionRenderer", + "SpoolCommandExecutor", + "TIMEOUT_RETURNCODE", +] diff --git a/agent-flow/agent_flow/runtime/commands/base.py b/agent-flow/agent_flow/runtime/commands/base.py new file mode 100644 index 000000000000..fb7ac58019d6 --- /dev/null +++ b/agent-flow/agent_flow/runtime/commands/base.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""One command, and where it ran. + +A workflow that measures something needs to say *what to run* and *how long +it may take*. Where that lands — this host, an ``srun`` step, a file in a +spool a long-lived allocation drains — is a property of the deployment, and +every time it leaked into the workflow it arrived as a new branch beside a +stage transition. + +So: the caller builds a :class:`CommandSpec`, an executor returns an +:class:`ExecutionRecord`, and the record carries enough provenance that the +verdict drawn from it can be defended later without knowing which executor +produced it. + +**A timeout is not an error here.** ``timed_out`` is a separate field from +the exit code because the two need different handling — a gate that timed out +has said nothing about the thing it was gating, while a gate that exited +non-zero has. +""" + +from __future__ import annotations + +import abc +import shlex +from dataclasses import dataclass, field +from typing import ClassVar, Mapping, Sequence + + +@dataclass(frozen=True) +class ExecutionCapabilities: + """What an executor can actually run, in machine-readable form. + + Stated rather than inferred, so a caller can refuse an impossible + request before paying for it instead of discovering the limit from a + failure three hours in. + """ + + gpu: bool = False + distributed: bool = False + #: True when the executor runs inside an allocation acquired once and + #: reused, rather than queueing per command. That is the difference + #: between paying a scheduler wait (and a container import) per gate and + #: paying it once for the run. + persistent_allocation: bool = False + max_world_size: int | None = None + architectures: tuple[str, ...] = () + + +@dataclass(frozen=True) +class CommandSpec: + """What to run, and the bounds it runs under. + + ``argv`` is either a **string**, which is a shell line and is run as + one, or a **sequence**, which is executed directly with no shell between. + Both forms are real and neither is a convenience: a resolved launch + prefix (environment assignments, or an ``srun`` line) *is* shell syntax + and re-parsing it into an argv would mean reimplementing the shell, while + a probe like ``nvidia-smi --query-gpu=...`` has no reason to be exposed + to word splitting at all. + """ + + argv: str | Sequence[str] + env: Mapping[str, str] | None = None + timeout_s: int = 600 + #: A short name for this command in logs and artifacts. Not an id. + label: str = "" + #: How many ranks the command expects. Carried so an executor can refuse + #: a request its allocation cannot satisfy; never used to rewrite the + #: command, which is the caller's business. + world_size: int = 1 + gpu_arch: str | None = None + #: Where to run it, when the executor supports a working directory. + 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 ExecutionRecord: + """What happened, with the provenance a verdict rests on.""" + + command: str + returncode: int + stdout: str = "" + stderr: str = "" + timed_out: bool = False + duration_s: float = 0.0 + #: Which executor ran it — ``local``, ``slurm``, ``spool``. + executor: str = "" + label: str = "" + #: Where the full output lives, when the executor wrote one. The tail a + #: record carries is for reading; this is for auditing. + log_path: str | None = None + started_at: float = 0.0 + artifacts: tuple[str, ...] = field(default_factory=tuple) + + @property + def ok(self) -> bool: + return self.returncode == 0 and not self.timed_out + + @property + def combined(self) -> str: + return self.stdout + self.stderr + + def tail(self, lines: int = 40) -> str: + return "\n".join(self.combined.splitlines()[-lines:]) + + +class CommandExecutor(abc.ABC): + """Run one command somewhere and report what happened.""" + + #: Recorded in :attr:`ExecutionRecord.executor`. + name: ClassVar[str] = "command-executor" + + #: Overridden per executor. A class attribute rather than a property so + #: a caller can ask without constructing one. + capabilities: ExecutionCapabilities = ExecutionCapabilities() + + @abc.abstractmethod + def execute(self, spec: CommandSpec) -> ExecutionRecord: + """Run ``spec`` to completion or to its deadline.""" + + +__all__ = [ + "CommandExecutor", + "CommandSpec", + "ExecutionCapabilities", + "ExecutionRecord", +] diff --git a/agent-flow/agent_flow/runtime/commands/delegated.py b/agent-flow/agent_flow/runtime/commands/delegated.py new file mode 100644 index 000000000000..41aae58c1a28 --- /dev/null +++ b/agent-flow/agent_flow/runtime/commands/delegated.py @@ -0,0 +1,204 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run a command somewhere only an agent can reach. + +THE SITUATION THIS EXISTS FOR + +Some clusters put the scheduler behind a tool surface rather than a CLI: +``sinfo`` and ``sbatch`` cannot contact the controller from the process that +wants them, and the only path through is a broker the *agent* holds and the +orchestrator does not. The orchestrator is a plain Python process; it is not +an MCP client and cannot become one just because a gate needs to run. + +THE OBVIOUS ANSWER IS THE WRONG ONE + +"Let the agent run the gate and tell us the number." That hands the verdict +to the thing being audited. Every guarantee in this workflow rests on the +orchestrator reading evidence itself, and an agent that reports `94.7` is +not evidence -- it is a claim about evidence, produced by a process with an +interest in the answer. + +WHAT THIS DOES INSTEAD + +Split the job in two, and only delegate the half that has no verdict in it: + + orchestrator agent (holds the broker tool) + ------------ ----------------------------- + write /.sh + ask for ONE submission --> make exactly the call it was given + (no image to pick, no partition to + discover, no flags to invent) + wait for /.rc + read the log and the exit code + apply its own thresholds + +The submitted script writes its own log and its own exit code to the shared +filesystem, and renames the receipt into place **last** -- so a reader that +sees it can trust the log beside it is complete. That is the same ordering +rule the spool and the dispatch mailbox use, for the same reason. + +The consequence worth stating: **the agent is a trigger, not a witness.** If +it lies about what happened, nothing downstream changes, because nothing +downstream reads what it said. If it never answers, the run times out on the +receipt rather than on the turn. + +WHY THE COMMAND IS A FILE + +The submission templates these brokers take embed the command inside a +quoted string. Real gate commands carry ``--model_dir="/path"``, +``--chat_template_kwargs={"enable_thinking": false}`` and worse, so +interpolating one into a quoted template produces a call that is either +mis-parsed or silently truncated. A path has no quoting hazard. +""" + +from __future__ import annotations + +import time +from pathlib import Path +from typing import Callable, ClassVar + +from ..turns import TurnExecutor, TurnRequest +from .base import CommandExecutor, CommandSpec, ExecutionCapabilities, ExecutionRecord +from .local import TIMEOUT_RETURNCODE + +#: Renders the instruction that makes the far side run ``script_path``. +#: Supplied by whatever knows the cluster's submission surface; this module +#: deliberately does not. +SubmissionRenderer = Callable[[str], str] + + +class DelegatedCommandExecutor(CommandExecutor): + """Ask an agent to launch a command; read the result off the filesystem.""" + + name: ClassVar[str] = "delegated" + capabilities = ExecutionCapabilities(gpu=True, distributed=True) + + def __init__( + self, + turns: TurnExecutor, + *, + root: Path | str, + render_submission: SubmissionRenderer, + role: str = "runner", + poll_s: float = 5.0, + name_prefix: str = "gate", + ) -> None: + self.turns = turns + self.root = Path(root) + self.render_submission = render_submission + self.role = role + self.poll_s = poll_s + self.name_prefix = name_prefix + + def _paths(self, name: str) -> tuple[Path, Path, Path]: + return ( + self.root / f"{name}.sh", + self.root / f"{name}.log", + self.root / f"{name}.rc", + ) + + @staticmethod + def _write_atomic(path: Path, text: str) -> None: + staging = path.with_suffix(path.suffix + ".tmp") + staging.write_text(text, encoding="utf-8") + staging.rename(path) + + def _script(self, spec: CommandSpec, log: Path, receipt: Path) -> str: + """The script the far side runs. It publishes its own verdict. + + Two details are load-bearing, both for the same reason -- the script + has to outlive whatever the command does, or there is no receipt and + the run waits out its full timeout on a job that already finished: + + * ``set -u`` but deliberately **not** ``set -e``. The exit code is + the measurement; aborting on it destroys the thing being measured. + * the command runs in a **subshell**. A payload that calls ``exit`` + -- directly, or through a wrapper that does -- would otherwise + terminate this script before it could record anything. + """ + exports = "".join( + f"export {key}={_shell_quote(str(value))}\n" + for key, value in sorted((spec.env or {}).items()) + ) + return ( + "#!/bin/bash\n" + "set -u\n" + f"{exports}" + f"( {spec.rendered} ) > {log} 2>&1\n" + "rc=$?\n" + # The receipt is written last and renamed into place, so its + # existence means the log beside it is complete. + f"echo $rc > {receipt}.tmp\n" + f"mv {receipt}.tmp {receipt}\n" + ) + + def execute(self, spec: CommandSpec) -> ExecutionRecord: + name = f"{self.name_prefix}-{int(time.time() * 1000)}" + script, log, receipt = self._paths(name) + self.root.mkdir(parents=True, exist_ok=True) + self._write_atomic(script, self._script(spec, log, receipt)) + + started_at = time.time() + started = time.monotonic() + + prompt = ( + f"Run one command on the cluster. **Make exactly this call and change " + f"nothing in it** — the environment was resolved before you started and " + f"is not yours to rediscover:\n\n" + f"```python\n{self.render_submission(str(script))}\n```\n\n" + f"Then poll the job's status until it leaves PENDING/RUNNING, and end " + f"your turn. Do **not** use a blocking submit call: the tool layer aborts " + f"after a few minutes of silence and every job here is longer than that, " + f"so blocking reports a failure for a job that is running fine.\n\n" + f"You do not need to report the result and nothing reads your answer. " + f"The job writes its own log to `{log}` and its own exit code to " + f"`{receipt}`, and the orchestrator reads both itself. Your only job is " + f"to get it submitted.\n\n" + f"If the submission itself fails, say so plainly — that is a real " + f"failure and it is the one thing this turn can tell us." + ) + self.turns.run(TurnRequest(role=self.role, what=f"submit {name}", prompt=prompt)) + + deadline = started + spec.timeout_s + while not receipt.exists(): + if time.monotonic() > deadline: + tail = log.read_text(encoding="utf-8", errors="replace") if log.exists() else "" + return ExecutionRecord( + command=spec.rendered, + returncode=TIMEOUT_RETURNCODE, + stdout=tail, + timed_out=True, + duration_s=time.monotonic() - started, + executor=self.name, + label=spec.label, + log_path=str(log), + started_at=started_at, + ) + time.sleep(self.poll_s) + + try: + returncode = int((receipt.read_text(encoding="utf-8").strip() or "1")) + except ValueError: + returncode = 1 + output = log.read_text(encoding="utf-8", errors="replace") if log.exists() else "" + return ExecutionRecord( + command=spec.rendered, + returncode=returncode, + stdout=output, + duration_s=time.monotonic() - started, + executor=self.name, + label=spec.label, + log_path=str(log), + started_at=started_at, + artifacts=(str(script), str(log), str(receipt)), + ) + + +def _shell_quote(value: str) -> str: + import shlex + + return shlex.quote(value) + + +__all__ = ["DelegatedCommandExecutor", "SubmissionRenderer"] diff --git a/agent-flow/agent_flow/runtime/commands/local.py b/agent-flow/agent_flow/runtime/commands/local.py new file mode 100644 index 000000000000..ad7ad05c710f --- /dev/null +++ b/agent-flow/agent_flow/runtime/commands/local.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run a command in a subprocess of this process.""" + +from __future__ import annotations + +import os +import subprocess +import time +from typing import ClassVar + +from .base import CommandExecutor, CommandSpec, ExecutionCapabilities, ExecutionRecord + +#: What a command that never finished reports. 124 is what ``timeout(1)`` +#: uses, so a caller reading exit codes by hand sees a familiar number. +TIMEOUT_RETURNCODE = 124 + +#: What a command whose executable does not exist reports. 127 is the +#: shell's own "command not found", and it matters that it is distinct from +#: a non-zero exit: the first says the probe could not run, the second says +#: it ran and disagreed. +NOT_FOUND_RETURNCODE = 127 + + +class LocalCommandExecutor(CommandExecutor): + """A subprocess on this host. + + Covers both shapes a caller needs. A shell line is run through the shell + because a resolved launch prefix *is* shell syntax; an argv sequence is + executed directly, because a probe has no reason to be exposed to word + splitting. :attr:`CommandSpec.argv`'s type decides which, so neither + caller has to pass a flag. + + The inherited environment is the parent's, with ``spec.env`` layered on + top. Replacing it outright would be the stricter choice and the wrong + one: everything here runs inside an environment somebody else already + set up — a venv's ``PATH``, a container's library paths, ``TMPDIR`` — + and dropping it turns a working host into a broken one. + """ + + name: ClassVar[str] = "local" + capabilities = ExecutionCapabilities(gpu=True, distributed=False, max_world_size=None) + + def execute(self, spec: CommandSpec) -> ExecutionRecord: + merged = {**os.environ, **(spec.env or {})} + command = spec.rendered + started_at = time.time() + started = time.monotonic() + try: + completed = subprocess.run( + command if spec.is_shell_line else list(spec.argv), + shell=spec.is_shell_line, + capture_output=True, + text=True, + timeout=spec.timeout_s, + env=merged, + cwd=spec.cwd, + check=False, + ) + except subprocess.TimeoutExpired: + return self._record( + spec, + command, + started, + started_at, + returncode=TIMEOUT_RETURNCODE, + timed_out=True, + ) + except FileNotFoundError as exc: + # Only reachable on the argv path -- a shell reports this itself. + return self._record( + spec, + command, + started, + started_at, + returncode=NOT_FOUND_RETURNCODE, + stderr=str(exc), + ) + return self._record( + spec, + command, + started, + started_at, + returncode=completed.returncode, + stdout=completed.stdout or "", + stderr=completed.stderr or "", + ) + + def _record( + self, + spec: CommandSpec, + command: str, + started: float, + started_at: float, + *, + returncode: int, + stdout: str = "", + stderr: str = "", + timed_out: bool = False, + ) -> ExecutionRecord: + return ExecutionRecord( + command=command, + returncode=returncode, + stdout=stdout, + stderr=stderr, + timed_out=timed_out, + duration_s=time.monotonic() - started, + executor=self.name, + label=spec.label, + started_at=started_at, + ) + + +__all__ = ["NOT_FOUND_RETURNCODE", "TIMEOUT_RETURNCODE", "LocalCommandExecutor"] diff --git a/agent-flow/agent_flow/runtime/commands/slurm.py b/agent-flow/agent_flow/runtime/commands/slurm.py new file mode 100644 index 000000000000..a4a06772a0a5 --- /dev/null +++ b/agent-flow/agent_flow/runtime/commands/slurm.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run a command as a Slurm step. + +**The step is still a local subprocess**, and saying so is the point. An +``srun`` line is a shell line: this host runs ``srun``, ``srun`` talks to the +controller, and the exit code and output come back down the same pipe as any +other child process. A separate executor class that secretly did the same +thing as the local one would be decoration. + +What it is *not* decoration for is the capability declaration and the +provenance tag. ``distributed=True`` and ``executor="slurm"`` are the two +facts a caller and an auditor respectively need, and neither is derivable +from the command string without parsing it. + +The ``srun`` line itself — partition, ntasks, ``--gres``, the MPI plugin, the +launcher wrapper — is **not** built here. Those rules belong to whatever +knows the engine being launched, and getting them wrong produces +``MPI_ERR_SPAWN`` rather than an error this layer could diagnose. This +executor receives a line that is already correct. +""" + +from __future__ import annotations + +from typing import ClassVar + +from .base import ExecutionCapabilities +from .local import LocalCommandExecutor + + +class SlurmStepExecutor(LocalCommandExecutor): + """A rendered ``srun`` line, executed here, reported as a Slurm step.""" + + name: ClassVar[str] = "slurm" + capabilities = ExecutionCapabilities(gpu=True, distributed=True) + + def __init__(self, *, max_world_size: int | None = None) -> None: + if max_world_size is not None: + self.capabilities = ExecutionCapabilities( + gpu=True, distributed=True, max_world_size=max_world_size + ) + + +__all__ = ["SlurmStepExecutor"] diff --git a/agent-flow/agent_flow/runtime/commands/spool.py b/agent-flow/agent_flow/runtime/commands/spool.py new file mode 100644 index 000000000000..2ecf82f61692 --- /dev/null +++ b/agent-flow/agent_flow/runtime/commands/spool.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run a command by writing a file a long-lived allocation is watching. + +The session is an allocation acquired once, whose own main process drains +``/pending``. Work is a **file**, which is the property that matters: +anything able to write a file can run GPU work, including a plain Python +process on a host with no GPU and no reachable scheduler. + +That is what hands command execution back to the orchestrator. The +alternative on such a host is to route every command through an agent +holding a broker tool, and then the verdict rests on that agent's account of +what a number was. Here it rests on a file. + +It also removes the per-command scheduler tax: one queue wait and one +container import for the run, instead of one of each per gate. + +Completion is the appearance of ``done/.rc``, which the watcher writes +last and renames into place, so a reader that sees it can trust the log +beside it is complete. Same ordering rule as the dispatch mailbox's +``.done``, for the same reason. +""" + +from __future__ import annotations + +import shlex +import time +from pathlib import Path +from typing import ClassVar + +from .base import CommandExecutor, CommandSpec, ExecutionCapabilities, ExecutionRecord +from .local import TIMEOUT_RETURNCODE + + +class SpoolCommandExecutor(CommandExecutor): + """Submit through a spool directory and block on the receipt.""" + + name: ClassVar[str] = "spool" + capabilities = ExecutionCapabilities( + gpu=True, + distributed=True, + persistent_allocation=True, + ) + + def __init__( + self, + spool: str | Path, + *, + poll_s: float = 3.0, + name_prefix: str = "gate", + ) -> None: + self.root = Path(spool) + self.poll_s = poll_s + self.name_prefix = name_prefix + + def _name(self) -> str: + return f"{self.name_prefix}-{int(time.time() * 1000)}" + + def execute(self, spec: CommandSpec) -> ExecutionRecord: + name = self._name() + command = spec.rendered + exports = "".join( + f"export {k}={shlex.quote(str(v))}\n" for k, v in sorted((spec.env or {}).items()) + ) + script = f"#!/bin/bash\nset -u\n{exports}{command}\n" + + (self.root / "pending").mkdir(parents=True, exist_ok=True) + # Write beside the spool and rename in, so the watcher cannot claim + # a half-written script. + staging = self.root / f".{name}.sh.tmp" + staging.write_text(script, encoding="utf-8") + staging.rename(self.root / "pending" / f"{name}.sh") + + receipt = self.root / "done" / f"{name}.rc" + log = self.root / "logs" / f"{name}.log" + started_at = time.time() + started = time.monotonic() + deadline = started + spec.timeout_s + while not receipt.exists(): + if time.monotonic() > deadline: + tail = log.read_text(encoding="utf-8", errors="replace") if log.exists() else "" + return ExecutionRecord( + command=command, + returncode=TIMEOUT_RETURNCODE, + stdout=tail, + timed_out=True, + duration_s=time.monotonic() - started, + executor=self.name, + label=spec.label, + log_path=str(log), + started_at=started_at, + ) + time.sleep(self.poll_s) + + returncode = int((receipt.read_text(encoding="utf-8").strip() or "1")) + output = log.read_text(encoding="utf-8", errors="replace") if log.exists() else "" + return ExecutionRecord( + command=command, + returncode=returncode, + stdout=output, + duration_s=time.monotonic() - started, + executor=self.name, + label=spec.label, + log_path=str(log), + started_at=started_at, + artifacts=(str(log), str(receipt)), + ) + + +__all__ = ["SpoolCommandExecutor"] diff --git a/agent-flow/agent_flow/runtime.py b/agent-flow/agent_flow/runtime/portal.py similarity index 86% rename from agent-flow/agent_flow/runtime.py rename to agent-flow/agent_flow/runtime/portal.py index 0b29d95e068e..3796ac5766ad 100644 --- a/agent-flow/agent_flow/runtime.py +++ b/agent-flow/agent_flow/runtime/portal.py @@ -1,3 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run async work on a dedicated background loop, from synchronous callers. + +Was ``agent_flow/runtime.py``; moved here when the runtime package grew +turns and commands beside it. ``from agent_flow.runtime import +PortalRunner`` still resolves, so nothing importing it had to change. +""" + from __future__ import annotations import asyncio @@ -7,7 +17,7 @@ import anyio -from .types import AgentRequest +from ..types import AgentRequest class PortalRunner: diff --git a/agent-flow/agent_flow/runtime/turns/__init__.py b/agent-flow/agent_flow/runtime/turns/__init__.py new file mode 100644 index 000000000000..2f48376e3a97 --- /dev/null +++ b/agent-flow/agent_flow/runtime/turns/__init__.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Agent turn transports: one interface, two deliveries.""" + +from .base import TurnExecutionError, TurnExecutor, TurnOutcome, TurnRequest, TurnTimeout +from .file_dispatch import ( + DEFAULT_TURN_TIMEOUT_S, + DISPATCH_DIRNAME, + HEARTBEAT_STALE_S, + FileDispatchTurnExecutor, + TransportNote, +) +from .inprocess import InProcessTurnExecutor, LayerProvider + +__all__ = [ + "DEFAULT_TURN_TIMEOUT_S", + "DISPATCH_DIRNAME", + "FileDispatchTurnExecutor", + "HEARTBEAT_STALE_S", + "InProcessTurnExecutor", + "LayerProvider", + "TransportNote", + "TurnExecutionError", + "TurnExecutor", + "TurnOutcome", + "TurnRequest", + "TurnTimeout", +] diff --git a/agent-flow/agent_flow/runtime/turns/base.py b/agent-flow/agent_flow/runtime/turns/base.py new file mode 100644 index 000000000000..220bc383fddf --- /dev/null +++ b/agent-flow/agent_flow/runtime/turns/base.py @@ -0,0 +1,126 @@ +# 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. + +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 plus — +where the caller asked for one — a typed decision comes back. That is the +entire unit of work an orchestrator cannot do in Python, and it is the only +thing a transport moves. + +Everything else stays with the caller: which stage runs, which item is next, +whether a decision is valid, whether a gate passed. So the two transports +that exist are interchangeable by construction, and the orchestrator should +not be able to tell them apart from its own source. That is what these types +are for. +""" + +from __future__ import annotations + +import abc +from dataclasses import dataclass, field +from typing import Any, ClassVar + + +class TurnExecutionError(RuntimeError): + """A turn could not be delivered or could not be answered.""" + + +class TurnTimeout(TurnExecutionError): + """Nothing answered the turn inside its deadline.""" + + +@dataclass(frozen=True) +class TurnRequest: + """What one role is being asked to do. + + ``system_prompt`` travels **with** the request rather than being set once + on a session. In process it could be either; through a mailbox it has to + be the former, because a driver handed only the task prompt runs a + different agent — no role, no skills, none of the invariants the writing + roles carry — and returns something that still looks like a plausible + answer. + + ``decision_file`` is the path the caller will read the typed decision + back from. It is part of the request because the role has to be told + exactly where to write; a pattern it has to instantiate is one it can + instantiate differently. + """ + + role: str + what: str + prompt: str + system_prompt: str = "" + decision_file: str | None = None + #: The caller's own numbering for this turn (a plan round, a queue item, + #: an assembler turn). Carried so a transport can name its artifacts + #: after it; never interpreted here. + iteration: int | None = None + #: Paths the role should read as input. Named rather than inlined, so a + #: transport that cannot afford the prompt budget can hand over pointers. + context_refs: tuple[str, ...] = () + + +@dataclass(frozen=True) +class TurnOutcome: + """What came back, and enough provenance to defend it later. + + ``text`` is the role's prose. **It is not a control-plane input**: a + caller that branches on it has given the role the ability to talk its way + past a gate. The typed decision arrives separately, validated against a + closed enum, and that is what a state machine may read. + """ + + 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 = "" + #: Where the typed decision was expected, echoed back for the collector. + decision_ref: str | None = None + #: Files this turn produced, oldest first. + artifacts: tuple[str, ...] = field(default_factory=tuple) + #: The conversation this turn belonged to, where the transport has one. + session_ref: str | None = None + duration_s: float = 0.0 + + +class TurnExecutor(abc.ABC): + """Deliver one turn and return its answer. + + Deliberately an abstract base rather than a bare protocol, because + :meth:`bound` has a default that every transport but one wants. + """ + + #: Recorded in :attr:`TurnOutcome.transport`. + name: ClassVar[str] = "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. + + Concurrent work often builds its own session per item — a batch of + catalog dispatches is ten independent agents, not ten turns of one — + and the caller, not the transport, is the thing that knows that. + + The default ignores it, which is correct for every transport with no + session concept: a mailbox has one directory whatever the caller + thought it was handing over. Only the in-process executor overrides + it. Callers can therefore write ``executor.bound(x).run(...)`` + unconditionally and stay free of transport branches. + """ + return self + + +__all__ = [ + "TurnExecutionError", + "TurnExecutor", + "TurnOutcome", + "TurnRequest", + "TurnTimeout", +] diff --git a/agent-flow/agent_flow/runtime/turns/file_dispatch.py b/agent-flow/agent_flow/runtime/turns/file_dispatch.py new file mode 100644 index 000000000000..16afc9e815e4 --- /dev/null +++ b/agent-flow/agent_flow/runtime/turns/file_dispatch.py @@ -0,0 +1,332 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The turn as a file, for a host that cannot authenticate an agent. + +THE PROBLEM, AND WHY IT IS NOT AN AUTHENTICATION PROBLEM + +An orchestrator that spawns a fresh CLI per turn needs that CLI to +authenticate itself, from a credential store the sandbox may make read-only. +A staged, writable home solves the *write* half. It does not solve the other +half, and cannot: + + the staged store holds a copy of a refresh token the parent also holds, + and a refresh rotates it. One of the two copies then loses. + +Measured twice on the staircase project. One run recorded "works for about +forty minutes"; the next saw a staged home authenticate cleanly at 12:12 and +fail with `Not logged in · Please run /login` at 12:25, with byte-identical +credentials. Copying harder does not fix a single-use token, and doing it +under ten concurrent agents makes it worse rather than better. + +THE ANSWER: MOVE THE TRANSPORT, AND NOTHING ELSE + +A role turn is the ONLY thing the orchestrator cannot do here. The stage +machine, the queue walk, the dependency batching, the prompt construction, +the decision-file validation and every gate are ordinary Python and work +fine. So the turn becomes a file: + + orchestrator driver (an authenticated session) + ------------ --------------------------------- + write dispatch/.request.json + write dispatch/.system.md + write dispatch/.prompt.md + wait for dispatch/.done <--- run the turn, however it can + write dispatch/.result.md + write dispatch/.done + read dispatch/.result.md + read the decision file the request named + +The driver can be a human, another agent session, or a service. It is not +trusted with anything: it delivers text and returns text. Every *decision* +still arrives through the caller's own validator, against the same closed +enum an in-process tool schema would have carried, and every *verdict* still +comes from a check the caller ran itself. + +NAMING + +Request ids are ``-`` with a monotonic counter, so a directory +listing sorts into dispatch order and a reader can see what the run did +without parsing anything. Concurrency is fine: ids are unique and each file +is written to a temporary name and renamed, so a driver that sees a file can +trust it is complete. +""" + +from __future__ import annotations + +import itertools +import json +import threading +import time +from pathlib import Path +from typing import Callable, ClassVar + +from .base import TurnExecutor, TurnOutcome, TurnRequest, TurnTimeout + +#: Subdirectory of the workspace the two sides meet in. +DISPATCH_DIRNAME = "dispatch" + +#: How long a single role turn may take before the orchestrator gives up on +#: it. Generous: an assembler turn on a large checkpoint is tens of minutes, +#: and a driver that is a human session may be slower still. The failure mode +#: this bounds is a driver that has stopped, not a turn that is slow. +DEFAULT_TURN_TIMEOUT_S = 3 * 60 * 60 + +#: A driver that has not touched its heartbeat in this long is gone. Checked +#: only if it ever wrote one, so a driver that does not bother is not +#: penalized --- but one that does gets a fast, named failure instead of a +#: three-hour wait. +HEARTBEAT_STALE_S = 300.0 + +#: Builds the paragraph appended to a role's system prompt saying what this +#: transport cannot honour. Supplied by the caller, because what a role +#: prompt promises is the caller's vocabulary, not this module's. +TransportNote = Callable[[str], str] + + +class FileDispatchTurnExecutor(TurnExecutor): + """Deliver a turn through the filesystem and wait for the answer. + + Thread-safe: concurrent turns are ordinary here, and the only shared + state is the id counter. + """ + + name: ClassVar[str] = "file-dispatch" + + def __init__( + self, + workspace: Path | str, + *, + poll_s: float = 5.0, + timeout_s: float = DEFAULT_TURN_TIMEOUT_S, + dispatch_dirname: str = DISPATCH_DIRNAME, + transport_note: TransportNote | None = None, + ) -> None: + self.root = Path(workspace) / dispatch_dirname + self.root.mkdir(parents=True, exist_ok=True) + self.poll_s = poll_s + self.timeout_s = timeout_s + self._note = transport_note + self._counter = itertools.count(self._next_seq()) + self._lock = threading.Lock() + + # ------------------------------------------------------------ helpers + + def _next_seq(self) -> int: + """Continue the numbering across a resume rather than restarting it. + + A resume that renumbers from 1 overwrites the previous attempt's + request files, which is the evidence for what the run asked for. + """ + seen = [ + int(p.name.split("-", 1)[0]) + for p in self.root.glob("*.request.json") + if p.name.split("-", 1)[0].isdigit() + ] + return (max(seen) + 1) if seen else 1 + + @staticmethod + def _write_atomic(path: Path, text: str) -> None: + staging = path.with_suffix(path.suffix + ".tmp") + staging.write_text(text, encoding="utf-8") + staging.rename(path) + + def _driver_is_gone(self) -> bool: + beat = self.root / "driver.heartbeat" + if not beat.is_file(): + return False + try: + return (time.time() - beat.stat().st_mtime) > HEARTBEAT_STALE_S + except OSError: + return False + + def _transport_note(self, role: str) -> str: + """What this transport cannot honour, said once, by the transport. + + A role prompt written for the in-process path ends by telling the + agent to call a typed-decision tool. That tool is real there — the + session it is handed to exists — and absent here, because this + transport has no session to host it in. The prompt therefore arrives + carrying two endings and only the second one works. + + Correcting the role prompt would be the wrong repair: the + instruction is right for the transport it was written for. The + difference belongs to the transport, so the transport states it. + """ + return self._note(role) if self._note is not None else "" + + # -------------------------------------------------------------- public + + def prior_turns(self, role: str) -> list[dict]: + """Every earlier answered turn for ``role``, oldest first. + + This transport starts a fresh session per turn, which silently + breaks any role the caller declares **persistent** — one asked to + remember what it decided so that a later turn revises rather than + re-derives. Their prompts often say so in as many words ("You are in + the same session as your first draft"), and here that sentence is + false. + + Restoring it exactly would mean replaying a conversation. Restoring + it USEFULLY only needs the artifacts, which are already on disk: + what the role was asked, and what it said back. A driver hands those + to the new session as its own prior reasoning. + """ + out = [] + for request in sorted(self.root.glob(f"*-{role}.request.json")): + name = request.name[: -len(".request.json")] + if not (self.root / f"{name}.done").exists(): + continue + out.append( + { + "id": name, + "prompt_file": str(self.root / f"{name}.prompt.md"), + "result_file": str(self.root / f"{name}.result.md"), + } + ) + return out + + def pending(self) -> list[Path]: + """Request files with no answer yet, in dispatch order. For a driver.""" + return sorted( + p + for p in self.root.glob("*.request.json") + if not (self.root / f"{p.name[: -len('.request.json')]}.done").exists() + ) + + def answer(self, request_id: str, result: str) -> None: + """Record a turn's output. For a driver; the orchestrator never calls it.""" + self._write_atomic(self.root / f"{request_id}.result.md", result) + # `.done` last, so a reader that sees it can trust the result beside + # it is complete. Same rule as a compute spool's `.rc`. + self._write_atomic(self.root / f"{request_id}.done", "1\n") + + def _unanswered_twin(self, request: TurnRequest) -> str | None: + """An earlier, still-unanswered request asking exactly this turn. + + A process that dies with turns in flight leaves their request files + behind, and a driver may be working on them. The resumed process asks + the same turns again (same role, same ``what``, same decision file); + writing a second request would duplicate the work -- or, worse, give + two agents the same decision path. Waiting on the first one is the + same answer to the same question. + """ + for path in sorted(self.root.glob(f"*-{request.role}.request.json")): + name = path.name[: -len(".request.json")] + if (self.root / f"{name}.done").exists(): + continue + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + continue + if data.get("what") == request.what and data.get("decision_file") == ( + request.decision_file + ): + return name + return None + + def run(self, request: TurnRequest) -> TurnOutcome: + with self._lock: + adopted = self._unanswered_twin(request) + request_id = adopted or f"{next(self._counter):03d}-{request.role}" + + started = time.monotonic() + # The role's system prompt travels WITH the turn. In-process it is + # set on the session and is easy to forget here -- and a driver that + # delivers only the user prompt is running a different agent: no + # role, no skills, none of the invariants the writing roles carry. + # The result would still look like a plausible answer. + system = self.root / f"{request_id}.system.md" + prompt = self.root / f"{request_id}.prompt.md" + result = self.root / f"{request_id}.result.md" + done = self.root / f"{request_id}.done" + + if adopted is None: + self._write_request(request, request_id, system, prompt, result, done) + + deadline = time.monotonic() + self.timeout_s + while not done.exists(): + if self._driver_is_gone(): + raise TurnTimeout( + f"{request_id}: the dispatch driver's heartbeat is stale. " + f"Nothing is reading {self.root}; the run would otherwise wait " + f"{self.timeout_s:.0f}s for a turn nobody is taking." + ) + if time.monotonic() > deadline: + raise TurnTimeout( + f"{request_id}: no answer within {self.timeout_s:.0f}s. " + f"The prompt is at {prompt}." + ) + time.sleep(self.poll_s) + + return TurnOutcome( + role=request.role, + text=result.read_text(encoding="utf-8") if result.is_file() else "", + transport=self.name, + request_id=request_id, + decision_ref=request.decision_file, + artifacts=(str(system), str(prompt), str(result)), + duration_s=time.monotonic() - started, + ) + + def _write_request(self, request, request_id, system, prompt, result, done) -> None: + self._write_atomic(system, request.system_prompt + self._transport_note(request.role)) + self._write_atomic(prompt, request.prompt) + self._write_atomic( + self.root / f"{request_id}.request.json", + json.dumps( + { + "id": request_id, + "role": request.role, + "what": request.what, + "system_prompt_file": str(system), + "prompt_file": str(prompt), + "result_file": str(result), + "done_file": str(done), + "decision_file": request.decision_file, + "context_refs": list(request.context_refs), + # Named rather than inlined: a driver that wants the + # role's own history reads these, and one that does not + # is not forced to carry them in the prompt. + "prior_turns": self.prior_turns(request.role), + "requested_at": time.time(), + }, + indent=2, + ) + + "\n", + ) + + def __call__( + self, + *, + role: str, + what: str, + prompt: str, + system_prompt: str = "", + decision_file: str | None = None, + ) -> str: + """Call it exactly where a layer would be called. + + Kept because a transport that can be dropped in where an + ``AgentLayer`` was is the cheapest possible migration, and because + the staircase orchestrator was wired this way before + :class:`TurnRequest` existed. + """ + return self.run( + TurnRequest( + role=role, + what=what, + prompt=prompt, + system_prompt=system_prompt, + decision_file=decision_file, + ) + ).text + + +__all__ = [ + "DEFAULT_TURN_TIMEOUT_S", + "DISPATCH_DIRNAME", + "FileDispatchTurnExecutor", + "HEARTBEAT_STALE_S", + "TransportNote", +] diff --git a/agent-flow/agent_flow/runtime/turns/inprocess.py b/agent-flow/agent_flow/runtime/turns/inprocess.py new file mode 100644 index 000000000000..250d6e986a79 --- /dev/null +++ b/agent-flow/agent_flow/runtime/turns/inprocess.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The turn this process runs itself. + +"In process" names where the *call* happens, not where the model does. The +backend still spawns a CLI; what makes this path different from the mailbox +is that the orchestrator calls a layer and gets the answer back on the same +stack, with no file queue and no second party in between. + +It is the shorter path and the default. It stops working when the host +cannot authenticate the CLI it spawns — a read-only credential store, or a +refresh token rotated out from under a child — which is a host failure that +presents as a model that returned nothing. +""" + +from __future__ import annotations + +import time +from typing import Any, Callable, ClassVar + +from .base import TurnExecutionError, TurnExecutor, TurnOutcome, TurnRequest + +#: Resolves the callable that runs a turn for one request. Given a request, +#: return something invocable as ``layer(prompt) -> str`` — in practice an +#: :class:`agent_flow.AgentLayer`, but nothing here requires that. +LayerProvider = Callable[[TurnRequest], Any] + + +class InProcessTurnExecutor(TurnExecutor): + """Run a turn by calling a layer directly. + + The provider indirection is what keeps this package free of workflow + vocabulary: role-to-layer construction (which backend, which model, + which tools, stateless or persistent) is a property of the workflow, so + the workflow supplies it and this executor only decides *when* to call. + """ + + name: ClassVar[str] = "inprocess" + + def __init__(self, layer_provider: LayerProvider) -> None: + self._provider = layer_provider + + def bound(self, session: Any) -> "InProcessTurnExecutor": + """Pin every turn to ``session``, whatever the request asks for. + + Used where the caller already owns the layer — a batch that builds + one agent per item and closes it afterwards. Without this the caller + would have to ask whether its transport has a session concept, which + is exactly the branch this package exists to remove. + """ + return InProcessTurnExecutor(lambda request: session) + + 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}; the provider returned None" + ) + started = time.monotonic() + text = layer(request.prompt) + return TurnOutcome( + role=request.role, + text=text if isinstance(text, str) else str(text), + transport=self.name, + decision_ref=request.decision_file, + session_ref=getattr(layer, "layer_name", None) or None, + duration_s=time.monotonic() - started, + ) + + +__all__ = ["InProcessTurnExecutor", "LayerProvider"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/.gitignore b/agent-flow/agent_flow/workflows/staircase_onboard/.gitignore new file mode 100644 index 000000000000..fdaa0d14ee36 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/.gitignore @@ -0,0 +1,5 @@ +# Placeholders a sandboxed agent session leaves in its working directory. +# Nothing in this package is configuration for a session that runs here: a run +# materializes its skills into its own workspace. +.mcp.json +.claude/ diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/BUILD_AND_TEST.md b/agent-flow/agent_flow/workflows/staircase_onboard/BUILD_AND_TEST.md new file mode 100644 index 000000000000..6c3dae244d37 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/BUILD_AND_TEST.md @@ -0,0 +1,245 @@ +# Build, rebuild, and test at each scale + +A companion to [TROUBLESHOOTING.md](TROUBLESHOOTING.md). That file is keyed on +symptoms — *something broke, what is it*. This one is keyed on tasks — *I need +to build / rebuild / run something at N GPUs, how*. + +Most of it is not this workflow's knowledge. It is collected from the places +in TensorRT-LLM that already do these things correctly, each named so you can +read the current version rather than trust this page: + +| where | what it owns | +|---|---| +| `scripts/build_wheel.py` | the build itself, and every flag below | +| `.claude/skills/exec-slurm-compile/` | building on a Slurm cluster, and `enroot import` | +| `.claude/skills/trtllm-test-script-builder/` | generating Slurm test scripts — **the launch rules** | +| `.claude/skills/exec-local-slurm/` | persistent allocations, container reuse, step cancellation | +| `tests/integration/test_lists/test-db/*.yml` | which tests run at which GPU count | +| `scripts/test_to_stage_mapping.py` | which CI stage a given test lands in | + +--- + +## 0. First ask: do you need to build at all? + +**For a modeling_v2 onboard, almost certainly not.** This is worth checking +before spending an hour, because the answer is usually no: + +| what you changed | rebuild? | +|---|---| +| a target's `modeling.py` / `weights.py` | **no** — pure Python, editable install picks it up | +| a catalog wrapper `.py` or contract `.md` | **no** | +| a Triton kernel (sourcing tier 4) | **no** — Triton is JIT | +| routing, `_router_index.py` | **no** | +| anything under `cpp/` | yes | +| a new `torch.ops.trtllm.*` op | yes | +| pulling a trunk that touched `cpp/` | yes | + +The four sourcing tiers this workflow uses — an exported op, the kernel plane, +native torch, Triton — **none of them require a native rebuild**. That is why +the orchestrator never builds anything and `environment.python` just points at +an install someone else made. + +If you are rebuilding during an onboard, stop and confirm why. + +## 1. Build + +```bash +python3 ./scripts/build_wheel.py \ + --use_ccache \ + --skip_building_wheel \ + --cuda_architectures "100-real;103-real" +``` + +| flag | when | +|---|---| +| `--cuda_architectures` | **always set it.** `100-real` Blackwell B200/GB200, `103-real` GB300, `90-real` Hopper, `89-real` L40S, `80-real` A100. Building all of them is the slowest thing you can do | +| `--use_ccache` | always. A no-op change rebuilds in minutes instead of an hour | +| `--skip_building_wheel` | development — builds in place, no wheel artifact | +| `--fast_build` / `-f` | skips some kernels. Development only — **do not gate on a fast build** | +| `--clean` / `-c` | wipes the build directory. The answer to "weird compile issues" | +| `--configure_cmake` | force a CMake re-configure. Needed when a CMake input changed but the build did not notice | +| `-j ` | lower it when the build OOMs | + +`build_wheel.py` creates its own venv and runs `pip install -e .[devel]` into +it, so the interpreter it produces is the one to put in `environment.python`. + +### ccache does not carry across checkouts + +Measured: pointing `CCACHE_DIR` at another checkout's cache produced no +useful hit rate. Absolute paths are baked into the cache keys. Give each +checkout its own, and expect the first build there to be cold. + +### The one that wastes an afternoon + +A rebuild is *not* the fix for `NVRTC ... could not open source file "cuda.h"`. +Measured: rebuilding with `--nvrtc_dynamic_linking` (98% ccache hits) did not +fix it; the cause was the venv's `bin` missing from `PATH`. See +TROUBLESHOOTING.md before you rebuild for that symptom. + +## 2. Build on Slurm + +`.claude/skills/exec-slurm-compile/` owns this end to end — `submit_compile.sh` +submits, `compile.slurm` allocates, `compile.sh` runs the build inside the +container, `enroot-import` fetches the image. + +Two failure modes from its own troubleshooting table that are not obvious: + +- `enroot import` **auth** failure → `~/.config/enroot/.credentials`; +- `enroot import` producing an **empty or corrupt** `.sqsh` → the image URL + probably has an `https://` prefix it should not have. + +### The environment hazard on a shared cluster + +`/root/.local` and `/root/.cache` may be symlinked into a shared home, so +another checkout's develop-install is on the path of every container run. + +- **At runtime: `PYTHONNOUSERSITE=1`.** Without it an import can resolve to a + different tree than the one you are measuring. +- **During the native build: leave the user site enabled.** The CUTLASS kernel + CMake runs `setup_library.py develop --user` and then imports + `cutlass_library`. Isolate the `pip install -e` step and everything after + it, not the compile itself. + +## 3. Test — the scale ladder + +Scale is declared by `system_gpu_count` in each list's `condition` block, and +the file name follows it: + +| scale | list | `system_gpu_count` | what belongs here | +|---|---|---|---| +| **CPU** | `l0_cpu.yml` | `0–0` | anything that needs no device. `backend: generic` | +| **1–4 GPU** | `l0_gb300.yml`, `l0_b200.yml`, `l0_h100.yml`, … | `1–4` | the ordinary pre-merge matrix | +| **multi-GPU** | `l0_*_multi_gpus.yml` | `4–4`, `8–8` | collectives, TP/EP topologies | +| **multi-node** | `l0_*_multi_nodes*.yml` | `8–8`+ | anything spanning nodes | + +To find where a test you care about runs: + +```bash +python scripts/test_to_stage_mapping.py --tests "test_modeling_v2_allgather" +``` + +### CPU + +`agent-flow`'s own suite is the clearest example: pure CPU, its own stage +(`CPU-AgentFlow-UnitTest`), and `AgentFlowRule` narrows CI so an +agent-flow-only change runs exactly that stage and no GPU stage at all. + +For the engine, CPU tests are the schema / list / prose audits — including the +modeling_v2 contract audits, which are text scans and need no device. + +### Single GPU + +```bash +pytest tests/unittest/_torch/modeling_v2/norm/test_modeling_v2_flashinfer_rmsnorm.py +LLM_MODELS_ROOT=/path/to/models pytest tests/integration/defs/accuracy/test_modeling_v2_gpt_oss.py +``` + +Unit tests need no models. Integration tests need `LLM_MODELS_ROOT` **and** a +GPU; forgetting the variable produces a failure that reads like a missing +model rather than a missing setting. + +### Multi-GPU: two mechanisms, and they are not interchangeable + +**`mpi_pool_executor`** — a fixture, worker count by indirect parametrization: + +```python +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_something(mpi_pool_executor): ... +``` + +**A launcher in a fresh interpreter** — what the modeling_v2 collective +entries use, and the reason is stated in +`tests/unittest/_torch/modeling_v2/comm/_rank_job.py`: + +> A collective that breaks \*hangs\* rather than raising, so the +> deadline is load-bearing and the ``mpi_pool_executor`` fixture has none. + +Two consequences worth carrying into any new multi-rank test: + +- **A hang is the normal failure.** Whatever you write needs a deadline that + kills the process group, or one wedged rank takes the whole run with it. +- **You cannot start an MPI launcher from a pytest process that has already + imported `tensorrt_llm`** — that process has initialized MPI. The launcher + must go in a fresh interpreter, started *by file path*. + +### Multi-node + +The launch rules are the whole content here, and they come from +`.claude/skills/trtllm-test-script-builder/`. **The rule is a function of the +task count and nothing else:** + +| | `ntasks == 1` | `ntasks > 1` | +|---|---|---| +| `--mpi=` | `none` | `pmix` (or the cluster's variant — ask `srun --mpi=list`) | +| payload | clear the MPI environment, then exec | wrap in `trtllm-llmapi-launch` | + +```bash +srun -N 2 --ntasks 8 --ntasks-per-node 4 --mpi=pmix --gres=gpu:4 \ + --container-image=... --container-name=trtllm-run \ + bash -c "trtllm-llmapi-launch python3 -m tensorrt_llm.commands.serve $MODEL ..." +``` + +Four things that are each a whole afternoon if you get them wrong: + +1. **`--ntasks` must match the world size.** Without it `srun` launches one + task, 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 → `MPI_ERR_SPAWN`. +2. **`--gres` is per node**, not per world. They diverge the moment you span + nodes. +3. **At one task, clear the MPI environment.** mpi4py auto-initializes MPI + whenever PMI variables are present, and inside a PMIx-managed step they + always are — so a single-task job initializes MPI for no reason and then + dies spawning a worker. Keep `SLURM_PROCID` and `SLURM_LOCALID`; they + identify the rank rather than bootstrap MPI. +4. **`TLLM_WORKER_USE_SINGLE_PROCESS=1` is not the general fix.** The function + behind it is named `enable_worker_single_process_for_tp1` and the source + calls it a WAR for TP1. It does not scale past one rank. + +This workflow emits all of this from `environment.py`; nothing above is an +agent's decision. It is written out here because a human debugging a gate +needs to recognize a correct line when they see one. + +## 4. Iterating without paying full price + +| cost | how to stop paying it | +|---|---| +| container import, 26 GB / 4–8 min **per job** | `--container-name=` — pyxis skips the import when that name already exists on the node | +| queue wait per command | a persistent allocation: `salloc` once, then `srun --jobid=` steps into it (see `exec-local-slurm`) | +| both, plus a transport that cannot reach Slurm at all | the spool pattern — one long-lived job whose **own main process** watches a directory; work is a file | +| cold native build | `--use_ccache`, and one cache per checkout | + +On the spool: work started inside a `slurm_exec_in_job` step does **not** +survive the step's teardown, `nohup` included — measured. The long-running +process has to be the job's own main process. If your cluster's Slurm CLI is +reachable, prefer the persistent-allocation pattern; the spool exists for the +case where it is not. + +To abort one run without losing the allocation: + +```bash +squeue -s -j # list steps +scancel .0 # cancel one step +``` + +## 5. Where a new test belongs + +1. **No device?** → a CPU list, and check whether a CBTS rule already narrows + CI for your subtree (`jenkins/scripts/cbts/rules/`). +2. **One device?** → the ordinary `l0_.yml` for the architecture you + certified on. Receipts are per-architecture; a pass elsewhere is not a pass + here. +3. **Collectives or a TP/EP topology?** → `l0_*_multi_gpus.yml`, and give it a + deadline. +4. **Spans nodes?** → `l0_*_multi_nodes*.yml`, and it needs the launch rules + in §3. + +Then confirm with `scripts/test_to_stage_mapping.py` that it actually lands in +a stage. A test on no list is a test nobody runs. + +**One trap for documentation-shaped changes:** `ModelingV2Rule` deliberately +excludes `*.md` so a contract-only edit does not pull in multi-GPU GB300 +stages. That is the right trade, but it means a check that *scans* `.md` files +gets no CI trigger from the PRs most likely to break it. Put text audits in a +pre-commit hook, not in a GPU test stage. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/README.md b/agent-flow/agent_flow/workflows/staircase_onboard/README.md new file mode 100644 index 000000000000..6dfdbb671412 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/README.md @@ -0,0 +1,278 @@ +# staircase-onboard + +Brings one HuggingFace checkpoint to a gated **ModelingV2** target: a flat, +self-contained forward assembled only from catalog entries, trusted through +accuracy gates rather than through shared abstractions. + +This workflow is the agent half of [staircase](https://github.com/NVIDIA/TensorRT-LLM). +The code half is already in tree at +`tensorrt_llm/_torch/_experimental/modeling_v2/` — the catalog, the targets, +the router and their tests. What lives here is what builds them: the +orchestrator, the six agent roles, and the ~7,000 lines of skills those agents +load. + +```bash +staircase-onboard --task task.yaml --workspace workspace/staircase-onboard +``` + +Start with `--dry-run`. It runs preflight, resolves the launch environment and +stops before the first agent turn — the cheapest way to find out whether this +host can host this run. + +Two companions to this file: **[TROUBLESHOOTING.md](TROUBLESHOOTING.md)** is keyed on +symptoms — the verbatim error you saw. **[BUILD_AND_TEST.md](BUILD_AND_TEST.md)** is +keyed on tasks — building, rebuilding, and running at CPU / 1 GPU / multi-GPU / +multi-node scale, collected from the places in TensorRT-LLM that already do each +of them correctly. + +**When something breaks, read [TROUBLESHOOTING.md](TROUBLESHOOTING.md) before +you debug the target.** It is keyed on verbatim error text, and the entries +share one property worth knowing in advance: every one of them read as a model +or target defect while it was happening. `MPI_ERR_SPAWN`, a kvCache that +cannot fit a sequence, NVRTC missing `cuda.h`, a 27-point accuracy regression +that was a protocol mismatch — none of those sentences mention the +environment, and all of them were the environment. + +## The pipeline + +``` +preflight -> plan -> smith x N (concurrent) -> assemble x M (serial) -> deliver + Python capped one per one per unit Python + only 1 round catalog entry core, then features + 1 turn +``` + +| stage | who | what it produces | +|---|---|---| +| `preflight` | **nobody** — pure Python | the device claim, the slug, the branch, the resolved launch prefix, the substrate gate, the surface survey | +| `plan` | PlanDrafter ⇄ PlanReviewer | `plan.md`, `vocabulary.yaml` | +| `smith` | one smith per queue item | catalog entries + their GPU tests, committed per entry | +| `assemble` | one assembler per unit | `modeling.py`, `weights.py`, gated | +| `deliver` | Python + one Auditor turn | verification across six trees, then the merge | + +Every stage checkpoints before advancing, so a crash resumes at the phase that +was running. Re-running the same command resumes; `--clean` starts over. + +## Why it is shaped this way + +The predecessor workflow for this job reliably failed the same way: runs spent +their budget deciding *how to launch a job* and re-litigating a plan, and never +reached model onboarding. Four things in this workflow exist to prevent that, +and they are the parts worth not "simplifying" later. + +**Preflight is Python, not an agent turn.** The checkpoint parses, the arch +matches the host, the device set is claimed by `memory.free` and a +lone-process probe, `TMPDIR` is sane, and a multi-rank run's NCCL/MPI +substrate answers — all in seconds, before any model is invoked. A broken host +costs seconds instead of a plan phase and a first assembler turn. + +**The launch environment is declared once and handed to every agent.** +`task.yaml`'s `environment` block resolves into a single command prefix that +every role's prompt carries verbatim. Agents never construct a launcher +command, never probe for one, and never call a test-runner skill. A prefix +that does not fit is a *symptom to report*, not a problem to solve. + +**The plan phase is capped and non-blocking.** `--max-plan-rounds` defaults to +1. Reaching the cap is not a failure: the reviewer's unresolved items are +appended to `plan.md` as `## Open risks` — which the assembler reads — and the +build starts. A plan review never blocks the build. + +**Gates are exit codes, not opinions.** There is no `--min-score` and no agent +scoring an artifact out of ten. `trtllm-eval` prints a number, the anchor comes +from TensorRT-LLM's own accuracy references, and the verdict is +`measured >= reference - tolerance`. + +## The gates + +| gate | when | cost | +|---|---|---| +| **boot** | after every assembler turn | minutes. `quickstart_advanced.py` with the target's topology flags | +| **accuracy** | at each unit boundary | the release criterion. `trtllm-eval` against the anchor. One-sided | +| **acceptance** | only for a distribution-preserving feature | `acceptance_length` under `=require` vs `=off` | + +All three run with **`TRTLLM_MODELING_V2=require`**. 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 — the single most expensive mistake available in this package. + +The **acceptance** gate exists because the first two are blind to speculative +decoding: rejection sampling holds the emitted distribution to the target +model's, so a *miscomputed* draft path produces correct text more slowly. Boot +passes, accuracy passes, only speed moves. + +Under `parallel: unified` — the default — boot and accuracy run **once per +member**, where the members are the concrete configurations the one codebase +serves (`members:` in `task.yaml`, default `[tp, tep, dep]`). A +unified target gated at a single topology leaves the expert-parallel and +attention-DP branches unexecuted while its Coverage claims them, so the +orchestrator enforces the member set rather than trusting the prose. Every +verdict names its member (`accuracy[gsm8k]@dep4`) and the unit advances only +when all of them pass. The witness gates run at the first member only: a +witness is a *contrast between two legs of one configuration*, not a number +per configuration. `--dry-run` prints the member set and the resulting gate +count before anything is spent. + +## The six roles + +| role | session | backend | one run does | +|---|---|---|---| +| `plan_drafter` | persistent | Claude Code | the checkpoint→catalog mapping and the plan | +| `plan_reviewer` | persistent | **Codex** | a paper review of both | +| `smith` | stateless | Claude Code | exactly one catalog entry | +| `assembler` | persistent *within a unit* | Claude Code | one unit: the core, or one feature | +| `reviewer` | stateless | Claude Code | reads one unit's diff against the gate evidence | +| `auditor` | stateless | Claude Code | the runtime-value dump vs the certified columns | + +The PlanReviewer runs on a different model from the PlanDrafter on purpose: a +reviewer that shares the author's blind spots reviews its own reasoning back at +itself. The plan roles are stateful because they are asked to remember what +they decided; the work roles are not, because staircase's unit of work is one +entry or one unit, each in its own context. + +## Workspace layout + +``` +/ +├── task.yaml the resolved spec, defaults filled in +├── plan.md PlanDrafter (+ `## Open risks` if the cap was reached) +├── vocabulary.yaml the typed step→catalog mapping and dispatch queue +├── surface.md the surface survey, unabridged (preflight) +├── progress.yaml every agent's typed decision, by stage +├── journal.md the run's timeline, written as events happen +├── reports/NN--.md each dispatch's prompt and result +├── decisions/-N.yaml typed decisions, when MCP tools are unavailable +├── boot.yaml the boot gate's prompts and keywords (optional) +├── bin/commit.sh the locking, self-verifying commit helper +├── .commit.lock its lock file +├── .claude/skills/ this run's copy of the packaged skills +└── .staircase_onboard_state.json +``` + +The journal and the reports stay here rather than in the repository. They are +real evidence and worth keeping, but they are evidence *about a run*, and a +product tree is the wrong home for it — only a distilled record travels into +the repo at delivery. + +`vocabulary.yaml` is the one artifact with a schema, because the orchestrator +walks it. `vocabulary.py` validates it and will reject a malformed one: +`OK`/`CERT` rows quote the contract line that decides coverage, `GAP` rows name +a candidate, `verify` dispatches name the surface, the two chapters do not +overlap, and a `ceiling` verdict carries its evidence. + +A queue row may also name the `unit` that first needs its entry (`core` when +absent). Every row is dispatched in the smith phase regardless, but only a +failure core needs — by its tag, or as a transitive `depends_on` of a core row — +closes the run there (`smith-blocked`). Any other failure is recorded +`deferred` and retried once before its unit's first assembler turn; if the +retry fails too, the run closes `smith-blocked:`. + +## Flags + +| flag | default | | +|---|---|---| +| `--task` | required | the spec. Validated before any agent is constructed | +| `--workspace` | `workspace/staircase-onboard` | | +| `--dry-run` | off | preflight, then stop. Leaves the checkpoint at `preflight`, so the real run re-claims the devices | +| `--max-plan-rounds` | `1` | drafter, reviewer, at most one re-draft, then proceed regardless | +| `--max-assembler-turns-per-unit` | `6` | bounds the self-heal loop | +| `--max-parallel-smiths` | `10` | where extra compute buys wall-clock. Each dispatch is a full agent session and the upstream rate limit is reported without backoff, so this is bounded by default; `0` opens the whole batch and on a typical queue that is every entry at once | +| `--reviewer` | `light` | `light` reads the evidence and runs nothing; `full` re-runs the boot gate; `off` skips the turn | +| `--plan` / `--vocabulary` | — | skip the plan phase with pre-written artifacts | +| `--skip-substrate-gate` | off | skip the stock-path cold start in preflight. An escape hatch, not a speed-up — see the flag's own help | +| `--agent-transport` | `inprocess` | how a role turn is delivered. `external` writes it to `/dispatch/` and waits — see below | +| `--agent-turn-timeout` | `10800` | seconds one externally-dispatched turn may take | +| `--clean` | off | wipe the checkpoint and managed files | + +## Running where the orchestrator cannot authenticate an agent + +On a sandboxed host the spawned CLI cannot write its credential lock, and the +turn fails with `EROFS ... credentials.lock` reported as +`stop_reason='stop_sequence'` and **zero tokens** — shaped exactly like a model +refusal when in fact no agent started. Two things happen, in order: + +1. **`agent_home.py` runs automatically.** If `$HOME` is not writable it stages + one under the workspace, seeds it from the real credential store and exports + `HOME` / `XDG_CONFIG_HOME` / `CLAUDE_CONFIG_DIR` before any agent exists. + Nothing to configure, and an ordinary workstation is untouched. +2. **If that authenticates and then stops working**, the cause is rotation, not + permissions: the staged store holds a copy of a refresh token the parent + also holds. Copying harder does not fix a single-use token. Use + `--agent-transport external`. + +External transport moves the **turn** and nothing else. The orchestrator writes +`dispatch/.system.md`, `dispatch/.prompt.md` and +`dispatch/.request.json`, then blocks until a driver writes +`dispatch/.result.md` followed by `dispatch/.done`. The driver can be a +human, another session or a service; it is trusted with nothing, because every +*decision* still arrives through `collect_decision_file` — validated against the +same closed enum the MCP tool schema would have — and every *verdict* still +comes from a gate the orchestrator ran itself. + +The one thing a driver must get right: **the role's system prompt travels with +the turn**. A driver that delivers only the user prompt is running a different +agent, and its answer still looks plausible. + +## What this workflow does not do + +- **Performance.** Perf is measured, never gated, and it is not part of + onboarding. A target whose two gates are green is deliverable; the tuning + campaign is separate work. +- **Isolation.** Staircase denied every agent read access to TensorRT-LLM's + per-model definitions, so that a target could not be a transliteration of + the implementation it is compared against. That hook is **not** ported: the + zoo is an ordinary read. What carries the duty instead is the + `modeling-core` invariants — a copied forward cannot pass the + closed-vocabulary audit, is not flat, and leans on shared helpers a target + may not have — plus the rule kept verbatim in all three writing roles: + *read to understand an interface, never to copy a structure*. If a future + run needs the stronger claim, a per-layer `PreToolUse` hook through + `BackendConfig.hooks` is a small, self-contained change. + +## Tests + +```bash +pytest tests/workflows/staircase_onboard/ +``` + +Every test runs on CPU with no model, no GPU and no cluster: the probes and +the shell are injected, and the agent layers are replaced by a fake that +writes the typed decision a real role would have written. Two exceptions earn +their keep by reading the real checkout — `test_delivery.py` verifies the +merged `gpt_oss / sm_103 / tp1` target passes every delivery check, and +`test_skills.py` resolves every repository path the skills name. + +## Feature increments + +`task.yaml`'s `features:` is a closed list, because a feature with no skill is +an assembler told to build something nobody wrote down how to build. + +| feature | skill | gated by | +|---|---|---| +| `mtp` / `eagle3` / `dflash` / `dspark` | `staircase-modeling-speculator` | boot, accuracy, **acceptance** | +| `eplb` | `staircase-modeling-eplb` | boot, accuracy | +| `dwdp` | `staircase-modeling-dwdp` | boot, accuracy | +| `multimodal` | `staircase-modeling-multimodal` | boot, accuracy on a multimodal benchmark | +| `host_offload` | `staircase-modeling-host-offload` | boot, accuracy, **residency witness** | +| `disagg` | `staircase-modeling-disagg` | boot, accuracy, **equivalence witness** | +| `kv_quant` | `staircase-modeling-kv-quant` | boot, accuracy, **kv_quant witness** | +| `decoder_replay` / `encoder_replay` | `staircase-modeling-bounded-replay` | boot, accuracy, **replay witness** | + +The speculator, `host_offload`, `disagg` and `kv_quant` rows share a property +worth stating plainly: they are **output-identical by construction**, so boot +and accuracy pass whether or not the feature is doing anything. Each therefore +carries a third measurement that can only come out one way if the feature is +real — `acceptance_length` under `=require` vs `=off`, device memory with and +without the offload, token-identical output across the disaggregation hop, or +stored KV bytes per token against a declared figure (plus token identity with +bf16 storage). + +The bounded-replay row is the opposite case with the same need. DeepSeek-V4.1's +SWA Bounded Replay truncates sliding-window state on purpose, so accuracy is a +real check of it — but the truncation is designed to stay inside the accuracy +tolerance, and a unit that implements nothing stays inside it too. Its witness +reads per-request row counters the target prints (`MODELING_V2_WITNESS +key= ...` under `TRTLLM_MODELING_V2_WITNESS=1`) against a window and a +workload the spec declares before assembly. + +`dspark` is its own kind and never a spelling of `mtp`: DeepSeek-V4.1 stores +its DSpark block drafter as `mtp.*`, and `--spec_decode_algo=MTP` over those +weights builds the wrong drafter. The schema refuses the two together. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/TROUBLESHOOTING.md b/agent-flow/agent_flow/workflows/staircase_onboard/TROUBLESHOOTING.md new file mode 100644 index 000000000000..b6d08716e345 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/TROUBLESHOOTING.md @@ -0,0 +1,543 @@ +# Troubleshooting — infrastructure, not modeling + +Everything here cost a real run real hours, and every one of them **read as a +model or target defect while it was happening**. That is the pattern worth +internalizing: an engine that will not boot, a gate that reports a 27-point +regression, a test that says nine skills are malformed — none of those +sentences mention the environment, and all of them were the environment. + +So the first question when something fails is not "what is wrong with the +target" but **"has this host ever done this successfully?"** + +## How to use this + +Entries are keyed on **verbatim output** — the string you actually saw, which +is what you can grep. That is deliberate and it is the same rule the catalog +follows: a quoted diagnostic survives a refactor, a line number does not. +Search this file for the error text, not for a description of it. + +Each entry says whether the workflow **already catches it**. If it does, you +should never have reached the symptom, and reaching it anyway is itself worth +reporting — a check that has stopped working is worse than no check. + +## Preflight catches most of this before an agent exists + +Run `--dry-run` first. It resolves the environment and stops before the first +agent turn, and it is the cheapest possible answer to "can this host host this +run". + +| check | what it proves | +|---|---| +| `check_checkpoint` | the checkpoint parses and declares an architecture | +| `check_tmpdir` | `TMPDIR`/`CLAUDE_CODE_TMPDIR` are set and not under `/tmp` | +| `claim_devices` | enough idle devices of the declared arch, by free memory and a lone-process probe | +| `check_multirank_substrate` | NCCL collective + MPI spawn answer under a deadline, for a multi-rank run | +| `substrate_gate` | **this host generates a token** on the stock path under this run's exact runtime block | +| `agent_plane_gate` | one agent turn can complete at all | + +`substrate_gate` is the one that earns its minutes. It *generates* rather than +constructs, because the NVRTC failure below appeared on the first generated +token — after the engine had loaded 22 GB of weights, allocated the KV cache +and finished warm-up. A gate that stopped at "the engine constructed" passes +it straight through. + +--- + +## Launch substrate + +### `MPI_ERR_SPAWN: could not spawn n processes` + +**Not a model problem.** The container runtime forbids `MPI_Comm_spawn`, and +the default LLM API path uses it. Under pyxis/enroot this is the normal state +of affairs, not a misconfiguration. + +The decision is one line, in `llmapi/mpi_session.py`: + +```python +def need_spawn_mpi_workers(model_world_size: int) -> bool: + return get_mpi_world_size() == 1 and model_world_size > 1 +``` + +**So the fix is to make the ranks exist before the engine starts**, which is +what `trtllm-llmapi-launch` is for: one task per rank, rank 0 runs the leader +and the real task, the rest run workers and wait. The workflow emits this +automatically at world size > 1 — an allocation sized without `--ntasks` +sized to match is the whole bug. + +**`TLLM_WORKER_USE_SINGLE_PROCESS=1` is not the general fix.** It works, and +this run used it, but the function behind it is named +`enable_worker_single_process_for_tp1` and the source comments call it a WAR +for TP1. It does not scale past one rank. If you find it in a spec at world +size > 1, that spec is wrong. + +At world size 1 the right answer is the other branch: no wrapper, and clear +the MPI environment before exec. mpi4py auto-initializes MPI whenever PMI +variables are present, and inside a PMIx-managed step they always are — so a +single-task job initializes MPI for no reason and then dies spawning a worker. + +*Caught by:* `substrate_gate`. Note that `check_multirank_substrate` does +**not** catch it: that probe uses `torchrun`, a different launcher, so it +proves NCCL works between local processes and says nothing about whether the +gate's own `srun` line can start a multi-rank engine. It goes green first, +which is misleading. + +### The engine hangs at startup and never returns + +A wrong launch substrate **hangs rather than raises** — which is why the +multi-rank probe has a deadline at all. The usual cause is an MPI connect-back +landing on an interface that does not carry it. Without the deadline this +surfaces as "the engine never booted", hours in. + +*Caught by:* `check_multirank_substrate`, which fails on the timeout rather +than waiting. + +### `There are not enough slots available in the system` + +In full: *"There are not enough slots available in the system to satisfy the 4 +slots that were requested by the application: trtllm-llmapi-launch"*, before +any Python runs. + +Open MPI takes its slot count from the resource manager, and a `session_spool` +job is submitted as **one task** that owns the node (`NumTasks=1` against +`NumCPUs=144` and `gres/gpu:4`). So `-n 4` asks for four slots against one. +The allocation really does hold the whole node, so the slot count is a +bookkeeping guard rather than a real constraint. + +Do **not** fix it by raising the session's `ntasks_per_node`: that gives MPI +its slots and simultaneously starts four copies of the spool watcher, which +then race for every piece of work. + +*Caught by:* `_mpirun_launch`, which adds `--oversubscribe` on the `mpirun` +branch only. Under `slurm`, `srun --ntasks` genuinely sizes the world, so the +same message there is a mis-sized allocation reporting itself and the flag +would hide it. + +### Four `mgmn_leader_node` processes, and a log that stops growing + +No error, no traceback, no exit — the shape above, with a specific cause. A +correct 4-rank launch has one leader and three workers; four leaders means +every rank thinks it is rank 0. + +`trtllm-llmapi-launch` derives rank as +`mpi_rank=${SLURM_PROCID:-${OMPI_COMM_WORLD_RANK:-...}}` and reads +`SLURM_NTASKS` before `OMPI_COMM_WORLD_SIZE`. Slurm wins over Open MPI in +both. That is right under `srun` and exactly wrong under `mpirun` *inside* a +Slurm allocation, where `SLURM_PROCID` is 0 for every rank and `SLURM_NTASKS` +is 1. Measured, four ranks: `SLURM_PROCID` read `0, 0, 0, 0` where +`OMPI_COMM_WORLD_RANK` read `0, 1, 2, 3`. + +MPI itself is not at fault — a bare `mpi4py` probe under the identical +`mpirun` line reports `rank 0..3 of 4`. + +*Caught by:* `_mpirun_launch`, which prefixes `env -u SLURM_PROCID -u +SLURM_NTASKS ...` so the launcher falls through to the Open MPI variables. +The `slurm` branch keeps today's behaviour, where those variables are correct. + +### `UCX ERROR uct_iface_open(tcp/rdma_vf_rail0) failed` + +Noise, on every rank, and a documented red herring. UCX fails to open one IPv6 +RDMA virtual-function interface, falls back, and the run completes with exit +code 0 and correct output. `UCX_TLS=sm,self` silences it and fixes nothing. + +*Not caught, and deliberately not.* If the run is also hanging, the cause is +the entry above, not this. + +### Two jobs die together, both at `MPI_Init` + +Concurrent `MPI_Init` in the same allocation killed two compute sessions in +this run. If you are dropping several scripts into a spool at once, either +serialize the ones that initialize MPI or give them distinct allocations. + +*Not caught.* Observed; the mechanism was not isolated further. + +### Setting `TLLM_DISABLE_MPI=1` makes it worse + +It routes the executor to Ray, which was not installed, so the failure moves +rather than going away. Do not reach for this as a workaround for the spawn +error above — `TLLM_WORKER_USE_SINGLE_PROCESS=1` is the one that works. + +--- + +## The engine boots but will not run + +### `Impossible to fit in any sequence in kvCache` + +The engine inherited the checkpoint's `max_position_embeddings` — 262144 is +common — and sized the KV cache for it. Nothing is wrong with the checkpoint +or the target. + +Declare the runtime explicitly in `task.yaml`'s `runtime.extra_flags`: + +```yaml +runtime: + extra_flags: + - --max_seq_len=8192 + - --max_batch_size=8 + - --kv_cache_fraction=0.6 +``` + +*Caught by:* `substrate_gate`, which runs under the same `runtime` block the +gates will use — so a bad sizing fails in preflight rather than mid-assembly. + +### `NVRTC ... could not open source file "cuda.h"` + +**Two independent things produce this, and one of them is the boring one.** + +1. **The venv's `bin` is missing from `PATH`.** NVRTC needs the toolkit + headers and finds them relative to the interpreter's environment. This was + the root cause in this run, and it was found late because the symptom + points at CUDA rather than at `PATH`. +2. **The generation FMHA kernel takes the NVRTC path.** In + `kernels/trtllmGenKernels/fmha/fmhaKernels.h`, the kernel is selected for + NVRTC unless the KV dtype is `E2m1` — grep for + `options.mDtypeKv != tg::Dtype::E2m1`. On a build whose NVRTC has no header + search path, that path dies compiling. Setting `--kv_cache_dtype fp8` takes + the non-NVRTC branch and the error goes away. + +**A wrong diagnosis worth recording:** rebuilding with +`--nvrtc_dynamic_linking` (NVRTC_DYNAMIC_LINKING OFF→ON, 98% ccache hits) did +**not** fix it. And the `fp8` workaround, while it genuinely works, works for +the second reason above, not the first — so "fp8 fixed it" is a true sentence +with a false explanation attached. If you hit this, check `PATH` before you +rebuild anything. + +*Caught by:* `substrate_gate` — but only because it generates a token. The +fault is on the first generated token, not at construction. + +### The recurrent state runs at the wrong precision, silently + +A checkpoint declaring `mamba_ssm_dtype: float32` still gets the engine flag's +default of bfloat16. Nothing errors; the model is just half as precise as it +was calibrated for. Pass `--mamba_ssm_cache_dtype=float32` in +`runtime.extra_flags`. + +*Not caught.* This is the class the **surface survey** exists to surface: the +checkpoint declares it, so it appears in the survey's demands at plan time. + +--- + +## Cluster access + +### `sinfo` / `sbatch` / `squeue` cannot contact the controller + +On some clusters the Slurm CLI is simply unreachable from the session, and an +MCP broker is the only path. **Do not conclude the cluster is down.** Set +`environment.kind: slurm_broker`. + +One consequence is worth knowing before you choose it: the orchestrator is not +an MCP client, so gate *execution* moves to an agent while gate *adjudication* +stays in Python — the orchestrator reads the log off the shared filesystem and +applies its own thresholds, so no verdict rests on an agent's report of what a +number was. + +### A broker job "fails" after exactly 300 seconds but is still running + +The MCP layer aborts a tool call after 300 s of silence. Every real job here is +longer than that, so the blocking `slurm_run` reliably reports failure for a +healthy job and leaves it running with its output unreachable. + +Use `slurm_submit`, poll `slurm_job_status`, then read the log file. Never the +blocking call. + +### Work started in a job step dies when the step ends + +`nohup` does not save it. Work started inside a `slurm_exec_in_job` step does +not survive the step's teardown — measured. + +This, plus the 300 s cap, plus a 26 GB container image costing 4–8 minutes to +import **per job**, is why `session_spool` exists and why it is the kind to +prefer wherever available: + +- one allocation, acquired once, running a watcher as its own main process; +- **work is a file** — drop `pending/.sh`, wait for `done/.rc`, + read `logs/.log`; +- so anything that can write a file can run GPU work: the orchestrator, an + agent, or a human. + +### `invalid_mount ... world-writable governing directory` + +``` +{"error":"invalid_mount","status":400, + "detail":"mount source has a world-writable governing directory + /.../artifacts (post-validation component-swap vector)"} +``` + +Nothing is wrong with the path you asked for; something **above** it is +world-writable. CI artifact trees usually are (`drwxrwsrwx`), so a checkpoint +published there cannot be mounted where it lives and has to be staged into a +directory you own. + +Staging has its own trap, and it reads as a quota problem: the orchestrator's +sandbox grants write access only to its **declared working directories**, so a +copy into `/models` fails with `Read-only file system` even though +the permission bits are fine. Stage inside the checkout and exclude it via +`.git/info/exclude`. + +### A job submits fine by hand but the workflow's submission is rejected + +Check the GPU request form. NVL72 policy rejects job-scoped `--gpus`; +`gpus_per_node` is the only form that submits there. This is a cluster policy, +not a workflow bug — `environment.gpus_per_node` exists for it. + +--- + +## Python environment + +### The import resolves to a different tree than the one you are measuring + +`/root/.local` and `/root/.cache` are symlinked into a shared, relocated home +on this cluster, so another checkout's develop-install is on the path of every +container run. Export `PYTHONNOUSERSITE=1` at runtime. + +**But not during the native build:** the CUTLASS kernel CMake runs +`setup_library.py develop --user` and then imports `cutlass_library`, so the +user site must be *enabled* for that step. Isolate the `pip install -e` step +and everything after it, not the build itself. + +### The gate measures the wrong build entirely + +`environment.python` is the whole answer: point it at the editable install you +intend to measure. Resolving the interpreter once, in the spec, is what keeps a +run from silently gating a different build — and it is why no agent ever +constructs a launcher command. + +--- + +## The agent plane + +### `You cannot dynamically configure MCP servers when an enterprise MCP config is present` + +The presence of `/etc/claude-code/managed-mcp.json` — not its contents — makes +the SDK refuse every dynamically-registered server, including the in-process +one the typed-decision tools live in. The first symptom is otherwise an opaque +`ProcessError: exit code 1` on the very first agent turn. + +The workflow **detects this and switches to file-based decisions** +automatically (`ENTERPRISE_MCP_CONFIG` in `workflow.py`); each role writes +`decisions/-N.yaml` instead of calling a tool. Nothing to configure. + +### `API Error: EROFS: read-only file system, mkdir '.../credentials.lock'` + +Reported as `stop_reason='stop_sequence'` with **zero tokens**, so it reads as +a model refusal. It is not: no agent ever started. The spawned CLI's first act +is to take a lock beside its credential store, and a sandboxed `$HOME` is +read-only. + +*Caught by:* `agent_home.py`, which stages a writable home under the workspace +and exports `HOME` + `XDG_CONFIG_HOME` + `CLAUDE_CONFIG_DIR` before any agent +is constructed. Four details each cost a cycle to find, so do not "simplify" +them away: + +- `HOME` alone is not enough — the credential store resolves through **XDG**; +- the config directory must **not** be named `.claude`, because a sandbox that + denies `**/.claude` writes cripples the child with a *warning*, not an error; +- the path must be **absolute** — a relative `HOME` resolves against the + *child's* cwd and silently builds `//agent-home/`; +- `$HOME/.gitconfig` has to be written, or every commit dies with *"Please + tell me who you are"* — a new infrastructure failure traded for the old one. + +- a **symlinked** `$HOME` does not work — the lock follows the resolved path. + +### `Not logged in · Please run /login` on a turn that used to work + +This is the *second* half of the problem above, and staging a home does not +fix it. The staged credential store holds a **copy** of a refresh token the +parent session also holds, and a refresh rotates it — so one of the two copies +loses. Measured on this cluster: authenticated cleanly at 12:12, failed at +12:25, byte-identical credentials. A previous run measured "about forty +minutes". Ten concurrent smiths make it worse, not better. + +Copying harder does not fix a single-use token. + +*The answer is `--agent-transport external`.* The turn becomes a file: the +orchestrator writes `/dispatch/.{system,prompt}.md` and blocks +until a driver writes `.result.md` and then `.done`. Any +already-authenticated session can be the driver. Everything the orchestrator +*decides* stays in Python — the stage machine, the queue walk, the dependency +batching, decision-file validation and all three gates. + +One thing to check if you write your own driver: **the role's system prompt +travels with the turn** (`.system.md`). A driver that delivers only the +user prompt is running a different agent — no role, no invariants, no embargo +— and its answer still looks plausible. + +*Caught by:* `agent_plane_gate`, which spends one trivial turn in preflight so +this costs seconds rather than a 24-minute planning turn. + +--- + +## Filesystem and git + +### Everything stops at once with unrelated-looking errors + +Check `TMPDIR`. Engine builds, NVRTC scratch and harness output all default +there, and on a shared host another tenant can fill `/tmp` to zero mid-run. + +*Caught by:* `check_tmpdir`, which rejects a `TMPDIR` under `/tmp` — reading +the value **declared in `task.yaml`**, deliberately not the orchestrator's own +`os.environ`, because a run whose scratch location is implicit is one nobody +can reproduce. + +### `ipc path "..." is longer than 107 characters` + +Followed, from the same path, by `zmq.error.ZMQError: File name too long`. + +107 is `sizeof(sockaddr_un.sun_path)` — a kernel cap on unix domain socket +paths, not a ZeroMQ or TensorRT-LLM choice. The LLM API's RPC binds an IPC +socket under `TMPDIR` named `rpc_test_`, a fixed 45 characters plus a +separator, so **`TMPDIR` must be at most 61**. + +This is about the length of the *name*, not the filesystem. Unix sockets bind +fine on Lustre here — tested on `/tmp`, `/dev/shm` and both Lustre spellings, +all four succeed. Anyone copying a `TMPDIR: /tmp` workaround should know it +works because `/tmp` is *short*, not because it is local; point `TMPDIR` at a +shorter path on the same storage instead. + +On a cluster whose storage paths are deep this is the property that breaks +first: a project scratch directory of the form +`/scratch//portfolios//projects//users//` was 78 +characters before the run appended anything. + +*Caught by:* `check_tmpdir`, which derives the bound from +`UNIX_SOCKET_PATH_MAX` and `RPC_SOCKET_BASENAME_LEN` and puts the arithmetic +in the message. Seconds in preflight instead of a four-rank engine load and a +ZeroMQ traceback that names neither the spec nor the workflow. + +### `Unable to create '.git/index.lock': File exists` + +Two agents committing in the same instant, or a stale lock from a crashed +process. Before removing anything: confirm no git process is running +(`pgrep -a git`) and check the lock's age and size — a zero-byte lock with no +owning process is stale. + +The workflow removes the *cause* rather than the lock: smiths commit through +`bin/commit.sh`, which takes an `flock`, so concurrency becomes a short queue. +Exit code **75** from that script means "waited for the lock and did not get +it" — a queue, not a broken repo. + +### `could not lock config file .git/config: File exists` and the lock will not delete + +Seen on Lustre with `Device or resource busy`. Do not fight it: pass the +identity on the command instead of writing it to the config — + +```bash +git -c user.name="..." -c user.email="..." commit -s ... +``` + +### A suite reports every packaged skill as malformed + +A Claude Code session writes an empty `.claude/.cc-writes/` into whatever tree +it touches, and a filesystem walk treats it as a skill directory. git does not +track it, so `git status` shows nothing — which is the most confusing possible +way for a suite to go red. + +*Caught by:* `_skill_dirs()` in `test_skills.py`, which skips dot-directories. +Apply the same rule to any new walk over a packaged tree. + +--- + +## Measurement traps + +These produce **numbers**, not errors, which makes them the most expensive +category here. A wrong number is believed. + +### The accuracy gate reports a large regression that is not there + +Measured once: **63.50 against a 95.53 anchor, a 27-point phantom.** The anchor +had been scored with a chat template and the gate had not. A benchmark's +*protocol* is part of its identity — a reasoning model scored as 5-shot +completion lands near the floor whatever its ability. + +Declare it: `accuracy.benchmark_args: [--apply_chat_template]`. + +### A feature's gate re-measures the base path + +The `mtp` unit's accuracy gate must carry a `speculative_config`. Without one +it measures the base model and reports the number as the speculator's. Use +per-unit `accuracy.llm_api_options`: + +```yaml +accuracy: + llm_api_options: + core: /path/to/oracle_llm_api.yaml + mtp: /path/to/mtp_llm_api.yaml +``` + +### A gate measured under `auto` is not measuring modeling_v2 + +Under `TRTLLM_MODELING_V2=auto`, a configuration that misses a target's routing +criteria silently gets the built-in implementation — and the number reads as +modeling_v2's. Anything you will attribute to modeling_v2 runs under +`require`, which raises instead of falling back. + +It must be exported **before the ranks start**, not merely before `LLM(...)`: +workers receive the environment as it stood at MPI init, so a driver resolving +a modeling_v2 target while its workers resolve the built-in is exactly the +silent split the switch exists to prevent. + +### `llm.get_stats()` returns nothing + +The feature witnesses for overlap scheduler, chunked prefill and block reuse +are **unverified for this reason**, not because the features are absent. A +self-diagnosing witness reported `[stats] 0 iterations`. Earlier attempts also +failed for unrelated reasons — concurrent `MPI_Init`, `TLLM_DISABLE_MPI=1` +routing to a missing Ray, and wrong stats field names returning all zeros — +so do not assume the stats API is the only thing between you and a witness. + +*Open.* If you need these features proven, this is the thing to fix first. + +### The survey says good catalog entries are "Rotted — do not spend these" + +``` +**Rotted** — the entry names an `impl` that does not resolve in this build. +- `comm/allgather` → `torch.ops.trtllm.allgather` +- `quantization/fp4_quantize` → `torch.ops.trtllm.fp4_quantize` +- `attention/fused_qk_norm_rope` → ... +``` + +Check one by hand before believing it. **`torch.ops.` is a lazy proxy**: +`dir()` lists only the names something has already resolved, not the names that +are registered. Measured on a GB300 build: + +```python +len([n for n in dir(torch.ops.trtllm) if not n.startswith("_")]) # 91 +len({s.name for s in torch._C._jit_get_all_schemas() + if s.name.startswith("trtllm::")}) # 273 +``` + +With the short answer, 13 of 15 "rotted" entries resolved fine — and they were +exactly the vocabulary a quantized hybrid MoE target needs, so the plan would +have routed around all of them and dispatched a smith to re-onboard each one. + +*Caught by:* `survey.py`, which now enumerates from the schema registry and +records `op_enumeration: schema_registry | dir_fallback` so a reader can tell +which measurement they are looking at. + +### A feature unit's gate fails to start with `no such option` + +The evaluation **protocol** is options of the *benchmark subcommand*, and the +subcommands do not share an option set — `gsm8k` takes `--apply_chat_template` +and `--fewshot_as_multiturn`, `mmmu` takes neither. A run-wide +`accuracy.benchmark_args` therefore does not make a multimodal unit's gate +imprecise; it makes it fail to start, with a click usage error that reads as a +spec mistake. + +`benchmark`, `reference_yaml`, `llm_api_options`, `discriminators` and +`benchmark_args` all take the per-unit mapping form. Use it. + +### Two `task.yaml` copies disagree + +The spec is copied into the workspace, and on a resume the two can drift. +`_refresh_workspace_spec` re-copies it, but if you edit a spec mid-run, edit +the source and resume — do not edit the workspace copy. + +--- + +## When none of this fits + +Say so rather than working around it. The agent prompts carry this instruction +verbatim — *"If this prefix does not fit what you need to run, that is a +**symptom to report in your final message** — not a problem to solve"* — and +the same applies to a human operator. A run that spends its budget +rediscovering how to submit a job is the failure mode this whole workflow was +shaped to avoid. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/__init__.py b/agent-flow/agent_flow/workflows/staircase_onboard/__init__.py new file mode 100644 index 000000000000..d2638f1e6a4d --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/__init__.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Staircase onboarding workflow built on ``agent_flow.AgentLayer``. + +Brings one deployment target — a ``(checkpoint, GPU arch, parallel +config)`` triple — from a HuggingFace checkpoint to a gated +``tensorrt_llm/_torch/_experimental/modeling_v2`` target: a flat, +self-contained forward assembled only from catalog entries, trusted +through accuracy gates rather than through shared abstractions. + +The pipeline is five phases, and the first and last are mostly not agents +at all:: + + preflight -> plan -> smith x N (concurrent) -> assemble x M (serial) -> deliver + +- **preflight** is pure Python and runs in three widening rings. The cheap + ring is seconds: the checkpoint parses, the arch matches the host, the + parallel segment is well formed, the device set is claimed by + ``memory.free`` and a lone-process probe, ``TMPDIR`` is sane, and a + multi-rank run's NCCL/MPI substrate answers. Then the **substrate gate** + generates a token on the *stock* modeling path under the run's exact + runtime block — minutes, because it loads the real weights, and worth + every one of them: the three failures it catches + (``MPI_ERR_SPAWN``, ``Impossible to fit ... in kvCache``, NVRTC's + ``cannot open "cuda.h"``) each used to surface hours in, inside an + agent's turn, reading as though the target were at fault. Last is the + **surface survey**, which measures what the checkpoint structurally + demands and what this build supplies, so the plan phase reads an + inventory instead of guessing one. No agent runs in any of it, so a + broken environment costs minutes instead of a whole run. +- **plan** is PlanDrafter <-> PlanReviewer, capped at ``max_plan_rounds`` + (default 1) and **non-blocking**: when the cap is reached the reviewer's + unresolved items are written into ``plan.md`` as open risks and the run + proceeds. Products are ``plan.md`` and ``vocabulary.yaml``, the typed + step-to-catalog mapping that becomes the dispatch queue. +- **smith** walks that queue, one agent per catalog entry, concurrently + within a dependency batch. Each smith onboards or verifies exactly one + entry and commits it through the workspace's commit helper, which takes a + lock and refuses a commit whose contents are not exactly what was asked + for — concurrency into one worktree, made safe in code rather than in + prose. +- **assemble** runs one assembler per unit, serially — core first, then one + increment per feature — with the boot gate after every turn and the + accuracy gate at each unit boundary. Gates are commands with exit codes, + not agent judgements. +- **deliver** verifies the agents' claims against disk across the in-tree + trees, spends one stateless auditor turn on the single check that is + irreducibly prose-reading (the runtime-value dump against the catalog's + certified columns), lints once, and merges. + +Public surface: + +- :class:`StaircaseOnboardWorkflow` — the orchestrator. +- :class:`PromptBundle`, :data:`DEFAULT_PROMPTS`, and + :func:`build_staircase_onboard_prompts` — the five agents' system + prompts and the builder that scopes them to one task. +- ``STAGE_*`` constants — stage identifiers used by the checkpoint schema + (``/.staircase_onboard_state.json``). +""" + +from typing import Any + +from .prompts import DEFAULT_PROMPTS, PromptBundle, build_staircase_onboard_prompts +from .state import STAGE_ASSEMBLE, STAGE_DELIVER, STAGE_PLAN, STAGE_PREFLIGHT, STAGE_SMITH + +__all__ = [ + "DEFAULT_PROMPTS", + "PromptBundle", + "STAGE_ASSEMBLE", + "STAGE_DELIVER", + "STAGE_PLAN", + "STAGE_PREFLIGHT", + "STAGE_SMITH", + "StaircaseOnboardWorkflow", + "build_staircase_onboard_prompts", + "main", +] + + +def __getattr__(name: str) -> Any: + if name == "StaircaseOnboardWorkflow": + from .workflow import StaircaseOnboardWorkflow + + return StaircaseOnboardWorkflow + if name == "main": + from .cli import main + + return main + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/agent_home.py b/agent-flow/agent_flow/workflows/staircase_onboard/agent_home.py new file mode 100644 index 000000000000..af8c9778ff34 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/agent_home.py @@ -0,0 +1,204 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Give the orchestrator a writable ``$HOME`` so it can spawn agents at all. + +THE PROBLEM + +A sandboxed session's ``$HOME`` is read-only. The orchestrator spawns a fresh +``claude`` CLI per role, that CLI authenticates on its own, and its first act +is to take a lock beside the credential store:: + + API Error: EROFS: read-only file system, mkdir + '/home//.config/anthropic/credentials.lock' + +The turn fails with ``stop_reason='stop_sequence'`` and zero tokens, so the +message reads as a model failure. It is not: no agent ever started. + +This cost the previous run its entire agent transport. It worked around it by +emitting prompts for a human-driven session to dispatch, which verified the +prompts and the gates but left the orchestrator's own dispatch path +unexercised. This module is the fix, so that the next run does not have to. + +WHAT IT DOES + +Stage a private home under the workspace, seed it with a *copy* of the real +credential store, and return the environment that points a child process at +it. Applied to ``os.environ`` at startup, every backend inherits it --- Claude +Code and Codex alike --- with no per-backend plumbing. + +WHAT IT DOES NOT DO + +It does not solve credential *rotation*. The copy holds a refresh token the +parent session also holds; a child that refreshes may invalidate the other's. +Measured on the previous run: a copied credential set works for roughly forty +minutes. :func:`refresh_credentials` re-seeds from the real store, which is +cheap enough to call before every launch, and the honest operating model is: +the run checkpoints per stage, so an auth failure is a resume rather than a +loss. Do not read this module as making the problem go away. +""" + +from __future__ import annotations + +import os +import shutil +from dataclasses import dataclass +from pathlib import Path + +#: Where the real credential store lives, relative to the real ``$HOME``. +CREDENTIAL_SUBPATH = Path(".config/anthropic") + +#: Subdirectories of the real home worth copying, beyond the credentials. +#: ``.codex`` is here because the PlanReviewer runs on Codex and keeps its own +#: state there. +OPTIONAL_SUBPATHS = (Path(".codex"),) + +#: The child's Claude Code config directory. Deliberately NOT named +#: ``.claude``: a sandbox that denies writes under ``**/.claude`` --- which is +#: the common shape, since that is where settings and skills live --- would +#: make the child unable to write its own config, and the failure is a warning +#: rather than an error, so it is easy to miss. +CLAUDE_CONFIG_DIRNAME = "claude-config" + + +@dataclass(frozen=True) +class AgentHome: + """A staged home directory and the environment that selects it.""" + + root: Path + env: dict[str, str] + + def apply(self) -> None: + """Export this home into the current process's environment. + + Children inherit it, which is the whole point: no backend needs to + know it exists. + """ + os.environ.update(self.env) + + +def home_is_writable(home: Path | None = None) -> bool: + """Whether the real ``$HOME`` accepts the write the CLI needs. + + Tested by creating and removing a directory beside the credential store, + because that is literally what fails --- ``os.access`` reports the + permission bits and a sandbox is not a permission bit. + """ + root = Path(home or os.path.expanduser("~")) + probe = root / CREDENTIAL_SUBPATH / ".agent_home_write_probe" + try: + probe.mkdir(parents=True, exist_ok=True) + probe.rmdir() + return True + except OSError: + return False + + +def refresh_credentials(root: Path, home: Path | None = None) -> bool: + """Re-seed ``root`` from the real credential store. Returns success. + + Cheap --- two small JSON files --- so call it before every launch rather + than once at setup. A stale copy fails as an authentication error several + minutes into a turn, which is a much more expensive way to find out. + """ + source = Path(home or os.path.expanduser("~")) / CREDENTIAL_SUBPATH + if not source.is_dir(): + return False + destination = root / CREDENTIAL_SUBPATH + try: + destination.parent.mkdir(parents=True, exist_ok=True) + if destination.exists(): + shutil.rmtree(destination) + shutil.copytree(source, destination) + return True + except OSError: + return False + + +def stage_agent_home( + root: Path, + *, + home: Path | None = None, + git_identity: tuple[str, str] | None = None, +) -> AgentHome: + """Create ``root`` as a usable home and return the selecting environment. + + ``git_identity`` is ``(name, email)``. It is written to the staged + ``.gitconfig`` because the committer's ``git commit`` needs an identity + and, once ``$HOME`` is redirected, the real one is no longer on the + search path --- a run that gets this far and then fails every commit with + *"Please tell me who you are"* has traded one infrastructure failure for + another. + """ + # ABSOLUTE, always. The env this returns is read by a CHILD process whose + # cwd is not ours -- the workspace, usually -- so a relative HOME resolves + # somewhere else and the child silently builds a second home underneath + # the first. Observed exactly once, as + # `//agent-home/`, and it is invisible until + # somebody lists the directory. + root = Path(root).resolve() + root.mkdir(parents=True, exist_ok=True) + refresh_credentials(root, home=home) + + real = Path(home or os.path.expanduser("~")) + for sub in OPTIONAL_SUBPATHS: + source, destination = real / sub, root / sub + if source.is_dir() and not destination.exists(): + try: + shutil.copytree(source, destination) + except OSError: + # Best effort. A missing Codex state directory costs a cold + # start, not a run. + pass + + config_dir = root / CLAUDE_CONFIG_DIRNAME + config_dir.mkdir(parents=True, exist_ok=True) + marker = config_dir / ".claude.json" + if not marker.exists(): + marker.write_text("{}\n", encoding="utf-8") + + if git_identity: + name, email = git_identity + (root / ".gitconfig").write_text( + f"[user]\n\tname = {name}\n\temail = {email}\n", encoding="utf-8" + ) + + return AgentHome( + root=root, + env={ + "HOME": str(root), + # The credential store is resolved through XDG, not through HOME, + # on the platforms this runs on. Setting only HOME leaves the + # child reading the real (read-only) store and the failure is + # identical to having done nothing. + "XDG_CONFIG_HOME": str(root / ".config"), + "CLAUDE_CONFIG_DIR": str(config_dir), + }, + ) + + +def ensure_agent_home( + workspace: Path, + *, + home: Path | None = None, + git_identity: tuple[str, str] | None = None, +) -> AgentHome | None: + """Stage a home under ``workspace`` iff the real one cannot be written. + + Returns ``None`` when the real home is fine, so that an ordinary + workstation run is untouched and nothing has to be configured for it. + """ + if home_is_writable(home): + return None + return stage_agent_home(Path(workspace) / "agent-home", home=home, git_identity=git_identity) + + +__all__ = [ + "AgentHome", + "CLAUDE_CONFIG_DIRNAME", + "CREDENTIAL_SUBPATH", + "ensure_agent_home", + "home_is_writable", + "refresh_credentials", + "stage_agent_home", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/cli.py b/agent-flow/agent_flow/workflows/staircase_onboard/cli.py new file mode 100644 index 000000000000..af97b1ac1ae2 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/cli.py @@ -0,0 +1,285 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +from pathlib import Path + +from .agent_home import ensure_agent_home +from .state import STATE_FILENAME +from .task_schema import TaskSchemaError, load_and_validate_task_yaml +from .transport import DEFAULT_TURN_TIMEOUT_S, ExternalAgentTransport +from .workflow import DEFAULT_MAX_PARALLEL_SMITHS + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + prog="staircase-onboard", + description=( + "Bring one checkpoint to a gated modeling_v2 target: " + "preflight -> plan -> smith -> assemble -> deliver." + ), + ) + parser.add_argument( + "--task", + type=Path, + required=True, + help="Path to the task YAML file. It is validated before any agent " + "is constructed and copied verbatim into `/task.yaml`, " + "which every agent reads.", + ) + parser.add_argument( + "--workspace", + type=Path, + default=Path("workspace/staircase-onboard"), + help="Workspace directory for the run's shared state (task.yaml, " + "plan.md, vocabulary.yaml, progress.yaml, journal.md, reports/). " + "Created on demand.", + ) + parser.add_argument( + "--max-plan-rounds", + type=int, + default=1, + help="How many PlanDrafter<->PlanReviewer rounds to spend before " + "proceeding regardless. Default 1: draft, review, and at most one " + "re-draft. Reaching the cap is not a failure — the reviewer's " + "unresolved items are written into plan.md as open risks and the " + "build starts. A plan review never blocks the build.", + ) + parser.add_argument( + "--max-assembler-turns-per-unit", + type=int, + default=6, + help="Bound on the assembler's self-heal loop for one unit (the " + "core assembly, or one feature increment). Exhausting it ends the " + "run with the audit trail rather than looping indefinitely.", + ) + parser.add_argument( + "--max-parallel-smiths", + type=int, + default=DEFAULT_MAX_PARALLEL_SMITHS, + help="Cap on concurrent smith dispatches within one dependency " + f"batch. Default {DEFAULT_MAX_PARALLEL_SMITHS}. Catalog entries are " + "independent by construction, so this is where a run spends extra " + "compute to shorten wall-clock — but each dispatch is a full agent " + "session, and the upstream rate limit is reported without any " + "backoff, so a cap that is hit fails the whole batch together. " + "0 means 'as many as the batch holds' and is not recommended: on " + "a typical queue that is every entry at once.", + ) + parser.add_argument( + "--reviewer", + choices=("off", "light", "full"), + default="light", + help="How much reviewing each assembled unit gets. 'light' " + "(default) reads the diff, the boot-gate output and the vocabulary " + "audit and runs nothing — the expensive evidence already exists " + "because the gates ran. 'full' additionally re-runs the gates " + "itself. 'off' skips the reviewer turn entirely.", + ) + parser.add_argument( + "--plan", + type=Path, + default=None, + help="Use a pre-written plan instead of running the PlanDrafter. " + "Combined with --vocabulary the plan phase is skipped entirely and " + "the run starts at the smith queue. Ignored when resuming from a " + "checkpoint; pass --clean to start fresh.", + ) + parser.add_argument( + "--vocabulary", + type=Path, + default=None, + help="Use a pre-written vocabulary.yaml (the typed step-to-catalog " + "mapping and dispatch queue) instead of deriving one. See --plan.", + ) + parser.add_argument( + "--clean", + action="store_true", + help="Wipe the workspace checkpoint and managed files " + f"({STATE_FILENAME}, plan.md, vocabulary.yaml, progress.yaml, " + "journal.md) and start fresh. Without this flag the workflow " + "resumes from the checkpoint when one is present, and starts fresh " + "otherwise.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Run preflight, resolve the environment, print the plan of " + "record, and stop before the first agent turn. The fastest way to " + "find out whether this host can host this run.", + ) + parser.add_argument( + "--skip-substrate-gate", + action="store_true", + help="Skip preflight's substrate gate — the cold start that proves " + "this host generates a token on the STOCK modeling path under this " + "run's runtime block. It is the single highest-value preflight " + "check and skipping it means the three failures it catches " + "(MPI_ERR_SPAWN, kvCache sizing, NVRTC's missing cuda.h) resurface " + "hours later inside an agent's turn, reading as target defects. " + "Use it only when you have already proven the substrate in this " + "same environment, or when the gate itself is what is broken. It is " + "NOT the answer to an architecture the stock path does not know — " + "declare `environment.substrate_checkpoint` for that, which keeps " + "the two checks that do not depend on the checkpoint.", + ) + parser.add_argument( + "--agent-transport", + choices=("inprocess", "external"), + default="inprocess", + help="How a role turn is delivered. 'inprocess' (default) spawns a " + "CLI per turn, which requires this host to be able to authenticate " + "one. 'external' writes the prompt to `/dispatch/` and " + "waits for an answer, so an already-authenticated session can take " + "the turn instead — use it when turns fail with `EROFS ... " + "credentials.lock` or `Not logged in`, which are host failures that " + "present as zero-token model failures. Everything else — the stage " + "machine, the queue, decision validation, all three gates — stays " + "here either way.", + ) + parser.add_argument( + "--agent-turn-timeout", + type=int, + default=int(DEFAULT_TURN_TIMEOUT_S), + help="Seconds a single externally-dispatched turn may take before " + "the run fails. Only meaningful with --agent-transport external.", + ) + return parser.parse_args(argv) + + +# The identity commits are made under. Needed because redirecting $HOME takes +# the real ~/.gitconfig off git's search path, and this checkout has no +# identity of its own -- so without it every smith commit dies with "Please +# tell me who you are", which is a new infrastructure failure traded for the +# one just fixed. Not an AI attribution: the sign-off is git's own. +def _git_config(key: str) -> str: + """The operator's own git setting, read before $HOME is redirected.""" + try: + out = subprocess.run(["git", "config", key], capture_output=True, text=True, timeout=10) + except (OSError, subprocess.SubprocessError): + return "" + return out.stdout.strip() + + +GIT_IDENTITY_NAME = ( + os.environ.get("STAIRCASE_GIT_NAME") + or _git_config("user.name") + or os.environ.get("USER", "staircase") +) +GIT_IDENTITY_EMAIL = ( + os.environ.get("STAIRCASE_GIT_EMAIL") + or _git_config("user.email") + or f"{GIT_IDENTITY_NAME}@nvidia.com" +) + + +#: Open file objects whose flock this process holds, kept alive deliberately. +_HELD_LOCKS: list = [] + + +def _claim_workspace(workspace: Path): + """Take an exclusive lock on the workspace, or refuse to start. + + Two orchestrators on one workspace is not a race you notice: they share + the checkpoint file, the dispatch directory and the id counter, so they + mint the SAME request id, overwrite each other's request JSON, and both + unblock on one answer. Measured on this run -- a stale process that a + `pkill` had failed to kill raced a fresh one for three seconds and the + survivor's request was the older code's, missing the system-prompt + pointer the driver needed. + + The handle is kept in a module-level list rather than returned-and-hoped, + because `flock` is held by an OPEN FILE DESCRIPTION: let the file object + be collected and the lock silently disappears. A caller that writes + `_claim_workspace(ws)` without binding the result would take a lock that + releases before the first agent turn -- which is worse than no lock, + since it reads as protection. + + Released when the process exits, which is what a crash should do. + """ + import fcntl + + workspace.mkdir(parents=True, exist_ok=True) + path = workspace / ".orchestrator.lock" + handle = path.open("w") + try: + fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError: + handle.close() + print( + f"error: another staircase-onboard process already holds {path}.\n" + f" Two orchestrators on one workspace share the checkpoint, the\n" + f" dispatch directory and the request-id counter, so they overwrite\n" + f" each other and both unblock on one answer. Stop the other one\n" + f" first, or use a different --workspace.", + file=sys.stderr, + ) + sys.exit(3) + handle.write(f"{os.getpid()}\n") + handle.flush() + _HELD_LOCKS.append(handle) + return handle + + +def main(argv: list[str] | None = None) -> None: + """Run the staircase-onboard workflow as a CLI.""" + args = _parse_args(argv) + + # Validate at the CLI boundary: an invalid spec must abort before any + # agent is constructed, and the message must list every problem at + # once so the user fixes them in a single pass. + try: + load_and_validate_task_yaml(args.task) + except TaskSchemaError as exc: + print(f"error: {exc}", file=sys.stderr) + sys.exit(2) + + # A sandboxed session's $HOME is read-only, and the first thing a spawned + # `claude` CLI does is take a lock beside its credential store. Without + # this every agent turn fails with `EROFS ... credentials.lock` and zero + # tokens, which reads as a model failure rather than as "no agent ever + # started". Returns None on an ordinary workstation, where nothing + # changes. + lock = _claim_workspace(Path(args.workspace)) # noqa: F841 -- held for the run + + staged = ensure_agent_home( + Path(args.workspace), + git_identity=(GIT_IDENTITY_NAME, GIT_IDENTITY_EMAIL), + ) + if staged is not None: + staged.apply() + print(f"→ staged a writable agent home at {staged.root} ($HOME is not writable)") + + transport = None + if args.agent_transport == "external": + transport = ExternalAgentTransport(args.workspace, timeout_s=args.agent_turn_timeout) + print( + f"→ role turns go through {transport.root}; a driver session must " + f"answer each `*.request.json` with a `*.result.md` and a `*.done`" + ) + + from .workflow import StaircaseOnboardWorkflow + + with StaircaseOnboardWorkflow( + workspace=args.workspace, + max_plan_rounds=args.max_plan_rounds, + max_assembler_turns_per_unit=args.max_assembler_turns_per_unit, + max_parallel_smiths=args.max_parallel_smiths, + reviewer_mode=args.reviewer, + plan=args.plan, + vocabulary=args.vocabulary, + clean=args.clean, + dry_run=args.dry_run, + skip_substrate_gate=args.skip_substrate_gate, + agent_transport=transport, + ) as workflow: + workflow.run(args.task) + + +if __name__ == "__main__": + main() diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/committer.py b/agent-flow/agent_flow/workflows/staircase_onboard/committer.py new file mode 100644 index 000000000000..bed1898dc8f8 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/committer.py @@ -0,0 +1,243 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""One safe way for concurrent agents to commit into one worktree. + +Smiths run concurrently inside a dependency batch, in the same checkout, and +git's index is a single shared file. The prompts have always told agents to +stage by explicit path and never `git add -A`, and that rule is right but it +is not sufficient. Two failures survived it: + +- **`.git/index.lock` contention.** Two agents committing in the same second + and one gets ``Unable to create '.git/index.lock': File exists``. An agent + that reads that as "the repo is broken" then does something creative. +- **Cross-attribution.** Staircase measured this twice in one run: 576 lines + of another agent's in-progress test file and a mid-edit contract landed in + an unrelated commit, splitting one entry's attribution across two commits + and leaving dangling cross-references. + +The fix is a critical section plus a verification, and neither is judgement, +so neither belongs in a prompt. This module generates a small shell script +into the workspace; the agent prompts hand out its path and nothing else. +What the script guarantees: + +1. **Serialized.** An ``flock`` around the whole stage-and-commit, with a + bounded wait, so concurrency becomes a short queue rather than an error. +2. **Scoped.** ``git commit -- `` commits the working-tree content of + exactly those paths on top of ``HEAD``, whatever else sits in the index. +3. **Verified.** After committing it diffs the new commit against its parent + and **fails** if the file list is not exactly what was asked for. This is + the part that matters: misattribution stops being something a reviewer + notices later and becomes something the tool refuses to do. + +A private ``GIT_INDEX_FILE`` per dispatch was the other candidate and is a +trap. It does stop one agent's ``add`` from seeing another's staged files, +but an index seeded from ``HEAD`` before a concurrent commit lands produces a +tree that silently *reverts* that commit — trading a visible attribution bug +for an invisible data-loss one. +""" + +from __future__ import annotations + +import os +import shlex +import stat +from pathlib import Path + +#: How long an agent waits for the commit lock before giving up. Generous +#: because the section is milliseconds of git and the alternative — an agent +#: that decides the lock is stale — is much worse than a slow queue. +LOCK_TIMEOUT_S = 300 + +SCRIPT_NAME = "commit.sh" + +_SCRIPT = """#!/usr/bin/env bash +# Generated by the staircase-onboard orchestrator. Do not edit. +# +# usage: {name} [path...] +# +# Commits exactly the named paths, under a lock, and verifies afterwards that +# the commit contains exactly those paths and nothing else. +set -uo pipefail + +REPO={repo} +LOCK={lock} +TIMEOUT={timeout} + +if [ "$#" -lt 2 ]; then + echo "usage: $0 [path...]" >&2 + exit 2 +fi +MESSAGE="$1"; shift + +for path in "$@"; do + if [ ! -e "$REPO/$path" ] && [ ! -e "$path" ]; then + echo "commit.sh: $path does not exist; commit by path relative to the repo root" >&2 + exit 2 + fi +done + +exec {fd}>"$LOCK" +if ! flock -w "$TIMEOUT" {fd}; then + echo "commit.sh: waited ${{TIMEOUT}}s for the commit lock and did not get it." >&2 + echo "Another agent is committing. This is a symptom to report, not to work around." >&2 + exit 75 +fi + +cd "$REPO" || exit 1 + +# The identity git would use by itself comes before $USER: a login name is +# not a name a DCO sign-off can carry. +COMMIT_NAME=${{STAIRCASE_GIT_NAME:-$(git -C "$REPO" config user.name 2>/dev/null)}} +COMMIT_NAME=${{COMMIT_NAME:-${{USER:-staircase}}}} +COMMIT_EMAIL=${{STAIRCASE_GIT_EMAIL:-$(git -C "$REPO" config user.email 2>/dev/null)}} +COMMIT_EMAIL=${{COMMIT_EMAIL:-${{COMMIT_NAME}}@nvidia.com}} + +# `git add` so a new file is known to the index; `git commit -- ` so +# the commit carries the working-tree content of exactly these paths on top +# of HEAD, whatever else another agent has staged. +if ! git add -- "$@"; then + echo "commit.sh: git add failed" >&2 + exit 1 +fi +# Identity passed on the command line, never written to config. `-s` adds +# the DCO sign-off the target repository requires of every commit. +# +# Two independent reasons, both measured. A redirected $HOME (see +# agent_home.py) takes the real ~/.gitconfig off git's search path; and this +# checkout cannot accept `git config --local` at all, because a stale +# `.git/config.lock` on Lustre refuses both `rm` ("Device or resource busy") +# and `chmod` ("Read-only file system"). Without this every commit dies with +# +# fatal: unable to auto-detect email address (got 'user@host.(none)') +# +# which commit.sh reported as "nothing to commit is the usual cause" --- so +# each of seven concurrent smiths rediscovered it separately and worked +# around it by hand. +if ! git -c user.name="$COMMIT_NAME" -c user.email="$COMMIT_EMAIL" \ + commit -s -q -m "$MESSAGE" -- "$@"; then + echo "commit.sh: git commit failed (nothing to commit is the usual cause)" >&2 + exit 1 +fi + +SHA=$(git rev-parse HEAD) +ACTUAL=$(git diff-tree --no-commit-id --name-only -r "$SHA" | LC_ALL=C sort) +WANTED=$(printf '%s\\n' "$@" | LC_ALL=C sort) +if [ "$ACTUAL" != "$WANTED" ]; then + echo "commit.sh: REFUSING TO LEAVE THIS COMMIT AS IS." >&2 + echo " asked for:" >&2 + printf ' %s\\n' $WANTED >&2 + echo " commit $SHA actually contains:" >&2 + printf ' %s\\n' $ACTUAL >&2 + # Single quotes, not double: a backtick inside a double-quoted echo is a + # command substitution, so writing the reset as an example inside "..." + # runs it — and then the explicit reset below runs it a second time, + # rewinding a commit that had nothing to do with this one. + echo ' Undoing it with: git reset --soft HEAD^' >&2 + echo " Your work stays in the tree. Report this: it means a path you named" >&2 + echo " matched more than you meant, or another agent's file was already" >&2 + echo " staged for one of them." >&2 + git reset --soft HEAD^ + exit 1 +fi + +echo "commit.sh: $SHA" +printf ' %s\\n' $ACTUAL +""" + + +def render_script(repo: str | Path, lock: str | Path, *, timeout: int = LOCK_TIMEOUT_S) -> str: + """The script body, with this run's repo and lock baked in. + + Baked rather than passed so an agent has no argument it can get wrong, + and so the same call works from any working directory. + """ + return _SCRIPT.format( + name=SCRIPT_NAME, + repo=shlex.quote(str(repo)), + lock=shlex.quote(str(lock)), + timeout=int(timeout), + # An explicit high fd, so the redirect cannot collide with a + # descriptor the surrounding tool already holds open. + fd=9, + ) + + +def install_committer( + workspace: str | Path, repo: str | Path, *, timeout: int = LOCK_TIMEOUT_S +) -> Path: + """Write the commit helper into the workspace and return its path.""" + workspace = Path(workspace) + bin_dir = workspace / "bin" + bin_dir.mkdir(parents=True, exist_ok=True) + lock = workspace / ".commit.lock" + path = bin_dir / SCRIPT_NAME + path.write_text(render_script(repo, lock, timeout=timeout), encoding="utf-8") + path.chmod(path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) + lock.touch(exist_ok=True) + return path + + +def commit_instruction(script: str | Path) -> str: + """The block every committing role's prompt carries.""" + return ( + "## Committing (this replaces the Committing section above)\n" + "\n" + "**Ignore the `git add` / `git commit` recipe given earlier in this prompt.** " + "It states the right rule and leaves you to follow it; this run enforces the " + "rule instead, and running the two recipes against each other is how a commit " + "ends up carrying someone else's file.\n" + "\n" + "Commit through this script and through nothing else:\n" + "\n" + "```bash\n" + f'{script} "" [path...]\n' + "```\n" + "\n" + "Paths are **relative to the repo root**, explicit, and complete — every " + "file of your unit, no directories, no globs that could match a file you " + "did not write.\n" + "\n" + "It takes a lock, so you may wait behind another agent; that is normal and " + "is not an error. It then checks that the commit contains exactly the paths " + "you named, and **undoes the commit** if it does not, leaving your work in " + "the tree. If that happens, report it rather than retrying: it means a path " + "you named swept in more than you meant.\n" + "\n" + "One commit per unit of work, so attribution and revert stay per-unit. No " + "other git command that writes — never `push`, never a branch operation, " + "never a rebase, never `git add -A`. No AI attribution lines in the message.\n" + ) + + +def lock_is_held(workspace: str | Path) -> bool: + """Whether some process currently holds the commit lock. + + Used by the orchestrator to tell "the batch is still committing" from + "the batch is wedged" when a dispatch overruns. + """ + import fcntl + + lock = Path(workspace) / ".commit.lock" + if not lock.exists(): + return False + fd = os.open(lock, os.O_RDWR) + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError: + return True + else: + fcntl.flock(fd, fcntl.LOCK_UN) + return False + finally: + os.close(fd) + + +__all__ = [ + "LOCK_TIMEOUT_S", + "SCRIPT_NAME", + "commit_instruction", + "install_committer", + "lock_is_held", + "render_script", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/delivery.py b/agent-flow/agent_flow/workflows/staircase_onboard/delivery.py new file mode 100644 index 000000000000..adb977543b05 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/delivery.py @@ -0,0 +1,463 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Phase D — verify the agents' claims against disk, before anything merges. + +Staircase's rule: **agent reports are claims.** Verify against disk before +every transition — the gate numbers against the runs, the products against +the directory. Most of that verification is mechanical, so most of it is +here rather than in an agent turn. What is left over — comparing the +runtime-value dump against contracts written in prose — is the auditor's, +and it is the only part of delivery that needs a language model. + +In tree a finished onboarding touches **six** trees, three of them outside +the package the CI selector claims. Every one has a way of being silently +half-done, so every one gets a check: + +1. the catalog entries the target imports +2. their GPU tests, in the tests tree +3. the target itself +4. the routing entry that makes it reachable +5. the accuracy anchor +6. the test-list registration that makes CI run the gate + +Everything here reads files. Nothing imports the target: importing it needs +torch, a built extension and — for a multi-rank target — a live world, none +of which the orchestrator has or should need in order to notice that a file +is missing. The checks that genuinely require a running engine already ran: +they are the gates. +""" + +from __future__ import annotations + +import ast +import re +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, Sequence + +import yaml + +MODELING_V2_PACKAGE = "tensorrt_llm/_torch/_experimental/modeling_v2" +CATALOG_DIR = f"{MODELING_V2_PACKAGE}/catalog" +CATALOG_TESTS_DIR = "tests/unittest/_torch/modeling_v2" +ACCURACY_REFERENCES_DIR = "tests/integration/defs/accuracy/references" +ACCURACY_TESTS_DIR = "tests/integration/defs/accuracy" +TEST_DB_DIR = "tests/integration/test_lists/test-db" + +# ``catalog/torch/`` holds thin mirrors of torch callables. Upstream owns +# their correctness, so by design they carry no contract, no test and no +# receipts — absence there is not a gap. +MIRROR_CATEGORY = "torch" + +_CATALOG_IMPORT = re.compile( + r"tensorrt_llm\._torch\._experimental\.modeling_v2\.catalog\.([a-z_]+)\.([a-z0-9_]+)" +) + + +@dataclass(frozen=True) +class Check: + """One verification, and what it found.""" + + name: str + passed: bool + detail: str + + def render(self) -> str: + return f"[{'ok' if self.passed else 'FAIL'}] {self.name}: {self.detail}" + + +@dataclass(frozen=True) +class DeliveryReport: + checks: tuple[Check, ...] + + @property + def passed(self) -> bool: + return all(check.passed for check in self.checks) + + @property + def failures(self) -> tuple[Check, ...]: + return tuple(check for check in self.checks if not check.passed) + + def render(self) -> str: + return "\n".join(check.render() for check in self.checks) + + +def imported_catalog_entries(modeling_py: Path) -> list[tuple[str, str]]: + """Return ``(category, name)`` for every catalog entry a target imports. + + Parsed from the AST rather than by importing: a target's import pulls in + torch and a built extension, and this has to work on a login node. + """ + tree = ast.parse(modeling_py.read_text(encoding="utf-8")) + found: list[tuple[str, str]] = [] + for node in ast.walk(tree): + module = None + if isinstance(node, ast.ImportFrom) and node.module: + module = node.module + elif isinstance(node, ast.Import): + for alias in node.names: + match = _CATALOG_IMPORT.search(alias.name) + if match: + found.append((match.group(1), match.group(2))) + continue + if module: + match = _CATALOG_IMPORT.search(module) + if match: + found.append((match.group(1), match.group(2))) + # Stable order, de-duplicated: a target may import two symbols from one + # entry, which is one entry to verify, not two. + return sorted(set(found)) + + +def declared_required_ops(modeling_py: Path) -> list[str]: + """Read ``REQUIRED_TRTLLM_OPS`` statically.""" + tree = ast.parse(modeling_py.read_text(encoding="utf-8")) + for node in ast.walk(tree): + if not isinstance(node, ast.Assign): + continue + for target in node.targets: + if isinstance(target, ast.Name) and target.id == "REQUIRED_TRTLLM_OPS": + try: + value = ast.literal_eval(node.value) + except ValueError: + return [] + if isinstance(value, (list, tuple)): + return [str(v) for v in value] + return [] + + +def catalog_index_paths(repo_root: Path) -> set[str]: + """Every entry path recorded in ``catalog/index.yaml``.""" + index = repo_root / CATALOG_DIR / "index.yaml" + if not index.is_file(): + return set() + data = yaml.safe_load(index.read_text(encoding="utf-8")) or {} + return { + str(entry.get("path", "")).strip() + for entry in (data.get("entries") or []) + if isinstance(entry, dict) + } + + +def check_catalog_entries(repo_root: Path, entries: Iterable[tuple[str, str]]) -> list[Check]: + """Every imported entry exists, is indexed, and (unless a mirror) is contracted.""" + checks: list[Check] = [] + indexed = catalog_index_paths(repo_root) + for category, name in entries: + rel = f"{category}/{name}.py" + wrapper = repo_root / CATALOG_DIR / rel + contract = wrapper.with_suffix(".md") + problems: list[str] = [] + if not wrapper.is_file(): + problems.append(f"wrapper missing: {wrapper}") + if category != MIRROR_CATEGORY and not contract.is_file(): + problems.append(f"contract missing: {contract}") + if rel not in indexed: + problems.append(f"not listed in catalog/index.yaml (looked for {rel!r})") + checks.append( + Check( + name=f"catalog:{rel}", + passed=not problems, + detail="; ".join(problems) or "wrapper, contract and index entry present", + ) + ) + return checks + + +def check_catalog_tests(repo_root: Path, entries: Iterable[tuple[str, str]]) -> list[Check]: + """Every non-mirror entry has its GPU test, under the name CI matches. + + The name is not a preference: every file in the subtree is called + ``test_modeling_v2_*`` precisely so it cannot collide with upstream's + test of the same op, and so the CI selector's pattern matches it exactly. + A correctly written test under any other name is a test on no list. + """ + checks: list[Check] = [] + for category, name in entries: + if category == MIRROR_CATEGORY: + continue + rel = f"{CATALOG_TESTS_DIR}/{category}/test_modeling_v2_{name}.py" + path = repo_root / rel + checks.append( + Check( + name=f"test:{category}/{name}", + passed=path.is_file(), + detail=str(path) if path.is_file() else f"missing: {rel}", + ) + ) + return checks + + +def check_target(repo_root: Path, target_rel: str) -> list[Check]: + """The target's two products exist and it declares the surface it touches.""" + target = repo_root / target_rel + checks = [] + for product in ("modeling.py", "weights.py"): + path = target / product + checks.append( + Check( + name=f"target:{product}", + passed=path.is_file(), + detail=str(path) if path.is_file() else f"missing: {path}", + ) + ) + modeling = target / "modeling.py" + if modeling.is_file(): + declared = declared_required_ops(modeling) + checks.append( + Check( + name="target:REQUIRED_TRTLLM_OPS", + passed=bool(declared), + detail=( + f"declares {len(declared)} op(s): {', '.join(declared)}" + if declared + else "declares no REQUIRED_TRTLLM_OPS — nothing would catch a " + "renamed op until the forward hit it" + ), + ) + ) + return checks + + +def check_routing(repo_root: Path, family: str, target_rel: str) -> list[Check]: + """The target is reachable: routing names it, and the claims test exists. + + A target that exists but is not routed is invisible — and worse than + invisible, because under ``TRTLLM_MODELING_V2=auto`` the built-in + implementation answers instead and everything looks like it worked. + """ + routing = repo_root / MODELING_V2_PACKAGE / "models" / family / "routing.py" + if not routing.is_file(): + return [Check(name="routing:module", passed=False, detail=f"missing: {routing}")] + + text = routing.read_text(encoding="utf-8") + # The dotted form routing uses, e.g. + # ``models.gpt_oss.gpt_oss_120b__sm_103__tp1.modeling``. + dotted = target_rel.split(f"{MODELING_V2_PACKAGE}/", 1)[-1].replace("/", ".").rstrip(".") + referenced = dotted in text + claims_test = repo_root / CATALOG_TESTS_DIR / "test_modeling_v2_claims.py" + return [ + Check( + name="routing:target", + passed=referenced, + detail=( + f"{routing.name} references {dotted}" + if referenced + else f"{routing.name} does not reference {dotted}; the target would " + f"never be selected, and under `auto` the built-in implementation " + f"would answer instead" + ), + ), + Check( + name="routing:claims-test", + passed=claims_test.is_file(), + detail=str(claims_test) if claims_test.is_file() else f"missing: {claims_test}", + ), + ] + + +def check_accuracy_reference(repo_root: Path, reference_yaml: str, model_key: str) -> list[Check]: + """The anchor this run gated against is written down, keyed by the model.""" + path = Path(reference_yaml) + if not path.is_absolute(): + path = repo_root / path + if not path.is_file(): + return [Check(name="accuracy:reference", passed=False, detail=f"missing: {path}")] + data = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + present = model_key in data + return [ + Check( + name="accuracy:reference", + passed=present, + detail=( + f"{path.name} carries an entry for {model_key}" + if present + else f"{path.name} has no entry for {model_key}; the gate ran against " + f"an anchor that is not recorded anywhere" + ), + ) + ] + + +def check_test_registration(repo_root: Path, accuracy_test: str) -> list[Check]: + """The accuracy gate is on a list CI actually reads. + + A gate that passes on a branch and is on no list passes exactly once. + """ + test_path = repo_root / ACCURACY_TESTS_DIR / accuracy_test + checks = [ + Check( + name="tests:accuracy-file", + passed=test_path.is_file(), + detail=str(test_path) if test_path.is_file() else f"missing: {test_path}", + ) + ] + db = repo_root / TEST_DB_DIR + stem = Path(accuracy_test).stem + listed: list[str] = [] + if db.is_dir(): + for yml in sorted(db.glob("*.yml")): + if stem in yml.read_text(encoding="utf-8"): + listed.append(yml.name) + checks.append( + Check( + name="tests:test-db", + passed=bool(listed), + detail=( + f"{stem} is listed in {', '.join(listed)}" + if listed + else f"{stem} appears in no {TEST_DB_DIR}/*.yml; CI would never run it" + ), + ) + ) + return checks + + +def check_tree_committed(repo_root: Path, paths: Sequence[str]) -> list[Check]: + """Whether the delivered trees are in git, not just on somebody's disk. + + **Every other check in this module reads files.** That is the right thing + for "does this entry have a test", and the wrong thing for "is this target + delivered": a file answers both identically whether or not it was ever + committed. + + Measured, once: a run passed ``routing:target`` — see ``check_routing`` — + and reported ``mmmu 58.33`` while ``routing.py``, the target's + ``modeling.py`` (984 changed lines) and its ``weights.py`` sat uncommitted + in the working tree, and the catalog commits certifying those very lines + were already in history. It would have closed ``delivered`` and shipped a + ``routing.py`` with no multimodal unit. The auditor was the only thing + that looked. + + The same measurement runs the other way, and that is why this check reads + *dirty*, not *absent*. A turn can do real work and never hand back — the + work sitting in ``HEAD`` while the orchestrator waits on a decision file + that was never written — and, conversely, a gate can measure a tree + dirtied by somebody else: one run had an agent editing ``modeling.py`` + while three accuracy gates were queued, where a pass would have + attributed a number to a commit it was never measured on. A delivery that + checks "the file exists" cannot tell those two apart. + + Returns **no check at all** outside a git work tree, rather than a failing + one. Where there is no git there is no claim to make, and a fixture in a + temporary directory is not a delivery. + """ + if not (repo_root / ".git").exists(): + return [] + try: + result = subprocess.run( + ["git", "status", "--porcelain", "--", *paths], + cwd=repo_root, + capture_output=True, + text=True, + timeout=60, + ) + except (OSError, subprocess.SubprocessError) as exc: + return [ + Check( + name="git:target-committed", + passed=False, + detail=f"could not run git status: {type(exc).__name__}: {exc}", + ) + ] + if result.returncode != 0: + return [ + Check( + name="git:target-committed", + passed=False, + detail=f"git status exited {result.returncode}: {result.stderr.strip()[:200]}", + ) + ] + dirty = [line for line in result.stdout.splitlines() if line.strip()] + return [ + Check( + name="git:target-committed", + passed=not dirty, + detail=( + "the delivered trees are clean at HEAD" + if not dirty + else "uncommitted, so a clone would not have them:\n " + "\n ".join(dirty[:12]) + ), + ) + ] + + +def verify_delivery( + *, + repo_root: Path | str, + target_rel: str, + family: str, + reference_yaml: str, + model_key: str, + accuracy_test: str, + extra_entries: Sequence[tuple[str, str]] = (), +) -> DeliveryReport: + """Run every mechanical delivery check and return the report. + + ``extra_entries`` lets the caller add catalog entries the target does not + import directly — an entry a feature path reaches through another module, + say — so the check set matches what the run actually produced. + """ + root = Path(repo_root) + checks: list[Check] = [] + + modeling = root / target_rel / "modeling.py" + entries: list[tuple[str, str]] = list(extra_entries) + if modeling.is_file(): + entries = sorted(set(entries) | set(imported_catalog_entries(modeling))) + else: + checks.append( + Check( + name="target:modeling.py", + passed=False, + detail=f"missing: {modeling} — nothing else can be checked without it", + ) + ) + + checks.extend(check_target(root, target_rel)) + checks.extend(check_catalog_entries(root, entries)) + checks.extend(check_catalog_tests(root, entries)) + checks.extend(check_routing(root, family, target_rel)) + checks.extend(check_accuracy_reference(root, reference_yaml, model_key)) + checks.extend(check_test_registration(root, accuracy_test)) + checks.extend(check_tree_committed(root, (target_rel, MODELING_V2_PACKAGE, CATALOG_TESTS_DIR))) + return DeliveryReport(checks=tuple(checks)) + + +def summarize_for_auditor(report: DeliveryReport, dump_excerpt: str = "") -> str: + """Render what the auditor is handed: the mechanical results, plus the dump. + + The auditor is told what has already been checked so it does not spend + its one turn re-deriving any of it — its job is the single check that + reading a file cannot settle. + """ + lines = ["## Mechanical delivery checks (already run)", "", report.render(), ""] + if dump_excerpt: + lines += ["## Runtime-value dump", "", "```", dump_excerpt.rstrip(), "```", ""] + return "\n".join(lines) + + +__all__ = [ + "ACCURACY_REFERENCES_DIR", + "CATALOG_DIR", + "CATALOG_TESTS_DIR", + "Check", + "DeliveryReport", + "MIRROR_CATEGORY", + "MODELING_V2_PACKAGE", + "TEST_DB_DIR", + "catalog_index_paths", + "check_accuracy_reference", + "check_catalog_entries", + "check_catalog_tests", + "check_routing", + "check_target", + "check_test_registration", + "check_tree_committed", + "declared_required_ops", + "imported_catalog_entries", + "summarize_for_auditor", + "verify_delivery", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/environment.py b/agent-flow/agent_flow/workflows/staircase_onboard/environment.py new file mode 100644 index 000000000000..9aa3a79a58f6 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/environment.py @@ -0,0 +1,779 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""How GPU work is launched — resolved once, then handed to every agent. + +The ``environment`` block in ``task.yaml`` is resolved here into a +:class:`RunEnvironment`, whose ``prefix`` is a literal string that every +GPU command is prepended with. That string — not a policy, not a skill, not +a decision tree — is what the agent prompts carry. + +The reason is specific. The predecessor workflow told its agents to route +every test through a chain of skills that classified the request, picked a +container image, checked the environment and finally built an ``srun`` +line — four hops, each one an LLM decision, each one able to fail in a way +that reads as "the environment is broken" rather than "this hop guessed +wrong". Runs spent their budget rediscovering how to submit a job. Here +the orchestrator resolves it once, in Python, before any agent exists, and +an agent that finds the prefix does not fit its need is instructed to +*report that as a symptom* rather than to work around it. + +Nothing here talks to Slurm. Building the command is all this module does; +running it belongs to :mod:`gates` and to the agents. +""" + +from __future__ import annotations + +import math +import os +import shlex +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Mapping, Sequence + +#: ``--mpi=`` values. ``auto`` resolves by task count, which is the only thing +#: that decides it. +MPI_AUTO = "auto" +MPI_NONE = "none" +MPI_PMIX = "pmix" + +#: The wrapper that makes a pre-launched MPI world work. Rank 0 runs the +#: leader plus the real task; every other rank runs a worker and waits. +LLMAPI_LAUNCH = "trtllm-llmapi-launch" + +# --------------------------------------------------------------------------- +# The launch rules below are TensorRT-LLM's, not this workflow's. They are +# carried over from `.claude/skills/trtllm-test-script-builder/` --- the +# repo's own Slurm script generator --- because that is the tested statement +# of how this engine must be launched, and a second, divergent statement of it +# here would be a bug generator. +# +# Ported as *rules* rather than copied as code: agent-flow is a separate +# installable package and cannot import from a skills directory. +# +# The rule is a function of the task count and nothing else: +# +# ntasks == 1 -> --mpi=none, scrub the MPI environment, no wrapper +# ntasks > 1 -> --mpi=, wrap the payload in the launcher +# +# Why one task is the special case: `trtllm-bench` and the LLM API import +# mpi4py at module load, which auto-initializes MPI whenever PMI variables are +# present. Inside a PMIx-managed step they always are --- so a single-task job +# initializes MPI for no reason and then dies trying to spawn a worker +# dynamically, which is what pyxis forbids. Clearing the variables lets the +# process skip MPI init entirely. +# +# `SLURM_PROCID` and `SLURM_LOCALID` survive the scrub: they identify the rank +# rather than bootstrap MPI, and code downstream reads them. That is true +# under `srun`, which is the only launcher this constant is applied behind; +# `MULTI_TASK_SLURM_SCRUB` below is the case where it is false, and the two +# are halves of one rule rather than two rules. +# --------------------------------------------------------------------------- +SINGLE_TASK_MPI_SCRUB = ( + "for v in ${!PMI@} ${!PMIX@} ${!MPI@} ${!OMPI@} ${!SLURM@}; do " + '[ "$v" != SLURM_PROCID ] && [ "$v" != SLURM_LOCALID ] && unset "$v"; ' + "done" +) + +# --------------------------------------------------------------------------- +# The other half of the rule above, at the other end of the range. +# +# One statement, two ends: **a Slurm step's environment misleads any process +# that `srun` did not launch.** At one task the misleading variables are the +# PMI ones, and a lone process joins an MPI world it never asked for --- that +# is `SINGLE_TASK_MPI_SCRUB`. Above one task, under `mpirun`, the misleading +# variables are the Slurm *task* ones, and this is that case. +# +# `trtllm-llmapi-launch` derives its rank as +# +# mpi_rank=${SLURM_PROCID:-${OMPI_COMM_WORLD_RANK:-${PMI_RANK:-${PMI_ID:-0}}}} +# +# and, in `mpi_world_size`, reads `SLURM_NTASKS` before `OMPI_COMM_WORLD_SIZE`. +# Slurm wins over Open MPI in both. That is the right precedence under `srun`, +# which is what the launcher was written for, and exactly wrong under `mpirun` +# *inside* someone else's allocation --- which is precisely what a +# `session_spool` job is: one Slurm task that owns the node and runs `mpirun` +# itself. +# +# Measured, four ranks, one variable changed: +# +# SLURM_PROCID OMPI_COMM_WORLD_RANK SLURM_NTASKS +# as-is 0, 0, 0, 0 0, 1, 2, 3 1 +# Slurm vars scrubbed unset 0, 1, 2, 3 unset +# +# So every rank read `mpi_rank=0` and became a leader, and every rank also +# read `mpi_world_size=1`. Four `mgmn_leader_node` processes, each waiting to +# be joined by workers that would never come. No error, no traceback, no exit +# --- the run simply hangs forever. MPI itself was never at fault; a bare +# `mpi4py` probe under the identical `mpirun` line reported `rank 0..3 of 4`. +# +# Scrubbing lets the launcher fall through to the Open MPI variables it should +# be reading here. This belongs ONLY on the `mpirun` branch: under `slurm`, +# `srun --ntasks` really does size the world and `SLURM_PROCID` really is the +# rank, so the same scrub would break a launch that is already correct. +# +# `env -u` rather than a shell loop because this one prefixes a command rather +# than running inside a `bash -c` of its own, and because the set is a fixed +# list of names rather than a glob. +# --------------------------------------------------------------------------- +MULTI_TASK_SLURM_VARS: tuple[str, ...] = ( + "SLURM_PROCID", + "SLURM_NTASKS", + "SLURM_LOCALID", + "SLURM_NODEID", + "SLURM_STEP_ID", + "SLURM_TASK_PID", + "SLURM_STEP_NUM_TASKS", +) + +MULTI_TASK_SLURM_SCRUB = "env " + " ".join(f"-u {name}" for name in MULTI_TASK_SLURM_VARS) + +# --------------------------------------------------------------------------- +# Variables whose declared value EXTENDS the inherited one instead of +# replacing it. +# +# Only `PATH` is on this list, and it earns the exception by being the +# variable that decides *which program runs at all*. A spec that wants its +# own interpreter ahead of everything else writes the obvious line: +# +# PATH: /.venv-3.12/bin:/usr/local/bin:/usr/bin:/bin +# +# which reads as "my venv first" and, as a plain assignment, means "and the +# container's MPI, CUDA and toolchain directories are gone". Measured: that +# exact line deleted `/usr/local/mpi/bin` and the next launcher line died +# with `mpirun: command not found`, rc=127 --- several layers away from the +# spec that caused it, and presenting as a broken probe rather than a +# mis-declared path. +# +# Prepending keeps the property the author was after --- every entry they +# declared still wins the lookup, in the order they declared it --- and drops +# only the part they did not ask for. The alternative considered was to +# reject a `PATH` not ending in `:$PATH`; it teaches every spec author a rule +# to make the obvious line legal, where this makes the obvious line correct. +# --------------------------------------------------------------------------- +MERGED_ENV_VARS: tuple[str, ...] = ("PATH",) + + +def _prepend_to_inherited(name: str, value: str) -> str: + """``/a:/b`` -> ``'/a:/b':"$PATH"``, idempotent against an explicit suffix. + + The expansion is deliberately outside the quoting: it has to reach the + shell as an expansion, and ``shlex.quote`` over the whole value would + deliver ``$PATH`` as five literal characters — which is the same deletion + with a stranger spelling. A spec that already wrote the suffix by hand is + accepted and not doubled. + """ + for suffix in (f":${name}", f":${{{name}}}"): + if value.endswith(suffix): + value = value[: -len(suffix)] + break + if not value or value in (f"${name}", f"${{{name}}}"): + return f'"${name}"' + return f'{shlex.quote(value)}:"${name}"' + + +@dataclass(frozen=True) +class RunEnvironment: + """A resolved, immutable description of how to launch GPU work.""" + + kind: str + python: str + env: dict[str, str] = field(default_factory=dict) + devices: tuple[int, ...] = () + world_size: int = 1 + # Slurm only; ``None``/empty for a local run. + partition: str | None = None + container_image: str | None = None + account: str | None = None + mounts: tuple[str, ...] = () + time_limit: str | None = None + gpus_per_node: int = 0 + # Reuses an already-imported pyxis container on the node instead of + # re-importing it. A 26 GB image costs 4-8 minutes per job to import, and + # pyxis skips the import entirely when a container of this name already + # exists there. + container_name: str | None = None + # Which Slurm MPI plugin to request. ``auto`` picks by task count, which + # is the only thing that actually decides it; an explicit value is for a + # cluster whose plugin is a versioned variant (``pmix_v5``). + mpi: str = MPI_AUTO + # ``session_spool`` only: the directory a long-lived compute session + # watches for work. + spool: str | None = None + + @property + def ntasks(self) -> int: + """One task per rank. This is the number the whole launch turns on.""" + return max(1, self.world_size) + + @property + def tasks_per_node(self) -> int: + """Ranks per node, and therefore GPUs per node.""" + return min(self.ntasks, self.gpus_per_node) if self.gpus_per_node else self.ntasks + + @property + def nodes(self) -> int: + return math.ceil(self.ntasks / self.tasks_per_node) + + @property + def mpi_plugin(self) -> str: + """``--mpi=`` for this launch. + + ``none`` at one task and a PMIx plugin above it — see + :data:`SINGLE_TASK_MPI_SCRUB` for why one task is the special case. + """ + if self.mpi != MPI_AUTO: + return self.mpi + return MPI_PMIX if self.ntasks > 1 else MPI_NONE + + @property + def needs_llmapi_launch(self) -> bool: + """Whether the payload must be wrapped in ``trtllm-llmapi-launch``.""" + return self.ntasks > 1 + + @property + def prefix(self) -> str: + """The literal string every GPU command is prepended with. + + For ``slurm_broker`` this is the **in-container** prefix. The + transport is a broker call rather than a shell word, so it cannot + live in the same string; see :meth:`broker_submit`. + """ + if self.kind == "slurm": + return self._slurm_prefix() + return self._local_prefix() + + def broker_submit(self, command: str = "", *, script: bool = False) -> str: + """The literal broker call that launches GPU work, fully resolved. + + ``script=True`` submits ``bash `` instead of the command text. + That is the form the orchestrator uses, and it exists because this + template embeds its payload inside a **double-quoted** string: a real + gate command carries ``--model_dir="/path"`` and + ``--chat_template_kwargs={"enable_thinking": false}``, and + interpolating one here produces a call that is mis-parsed or silently + truncated. A path has no quoting hazard. + + ``slurm_broker`` exists because on some clusters the Slurm CLI is + unreachable from the process that wants to use it: ``sinfo`` and + ``sbatch`` fail to contact the controller, and an MCP broker is the + only path through. That has one consequence worth stating plainly — + **the orchestrator cannot run GPU work itself**, because it is not + an MCP client. Gate *execution* therefore moves to an agent, while + gate *adjudication* stays here: the orchestrator reads the log off + the shared filesystem and applies its own thresholds, so nothing + rests on an agent's report of what a number was. + + This template preserves the property the whole module exists for. + The agent is handed a resolved call, not a decision: no image to + pick, no partition to discover, no mount list to assemble. + """ + mounts = ", ".join(f'"{m}"' for m in self.mounts) + payload = f"bash {command}" if script else command + return ( + "mcp__slurm-broker__slurm_submit(\n" + f' image="{self.container_image}",\n' + f' partition="{self.partition}",\n' + f' account="{self.account}",\n' + f" gpus_per_node={self.gpus_per_node},\n" + f' time_limit="{self.time_limit}",\n' + f" mounts=[{mounts}],\n" + f' command=["bash", "-lc", "{payload}"])' + ) + + def _env_assignments(self) -> list[str]: + """``KEY=value`` pairs in a stable order, shell-quoted. + + A variable named in :data:`MERGED_ENV_VARS` is emitted as a prepend + to whatever the host already exports rather than as a replacement; + see that constant for the failure it prevents. + """ + parts: list[str] = [] + for key, value in sorted(self.env.items()): + text = str(value) + if key in MERGED_ENV_VARS: + parts.append(f"{key}={_prepend_to_inherited(key, text)}") + else: + parts.append(f"{key}={shlex.quote(text)}") + return parts + + def _local_prefix(self) -> str: + parts: list[str] = [] + if self.devices: + parts.append("CUDA_VISIBLE_DEVICES=" + ",".join(str(d) for d in self.devices)) + parts.extend(self._env_assignments()) + return " ".join(parts) + + def _slurm_prefix(self) -> str: + """An ``srun`` line carrying the allocation, the container and the env. + + **One task per rank.** Omitting ``--ntasks`` is not a smaller version + of this line, it is a broken one: ``srun`` then launches a single + task, 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``. Sizing the allocation without sizing the task count + produces that failure every time at world size > 1. + + ``--gres`` is **per node**, so it carries ``tasks_per_node`` rather + than the world size; the two differ as soon as a run spans nodes. + + ``CUDA_VISIBLE_DEVICES`` is deliberately absent: under Slurm the + allocation *is* the device claim, and setting both means the job + masks its own allocated devices by index — which presents as + "device 0 not found" on a node where device 0 was never ours. + """ + parts = ["srun", "-l"] + if self.partition: + parts.append(f"--partition={self.partition}") + if self.account: + parts.append(f"--account={self.account}") + if self.time_limit: + parts.append(f"--time={self.time_limit}") + parts.append(f"--nodes={self.nodes}") + parts.append(f"--ntasks={self.ntasks}") + parts.append(f"--ntasks-per-node={self.tasks_per_node}") + parts.append(f"--gres=gpu:{self.tasks_per_node}") + if self.container_image: + parts.append(f"--container-image={shlex.quote(self.container_image)}") + if self.container_name: + parts.append(f"--container-name={shlex.quote(self.container_name)}") + if self.mounts: + parts.append("--container-mounts=" + shlex.quote(",".join(self.mounts))) + plugin = self.mpi_plugin + if plugin != MPI_NONE: + parts.append(f"--mpi={plugin}") + parts.append("--kill-on-bad-exit=1") + exports = ",".join(["ALL", *self._env_assignments()]) + parts.append(f"--export={exports}") + return " ".join(parts) + + def _scrub_and_exec(self, argv: str) -> str: + """``argv`` behind a scrub of the MPI environment, in this process. + + Shared by the ``srun`` payload and by the launcher-free path, because + it is one rule and not two: see :data:`SINGLE_TASK_MPI_SCRUB` for why + a lone process must not inherit the PMI variables. + """ + return f"bash -c {shlex.quote(f'{SINGLE_TASK_MPI_SCRUB}; exec {argv}')}" + + def _mpirun_launch(self, argv: str) -> str: + """``argv`` behind an ``mpirun`` that creates the ranks here. + + Same engine rule as :meth:`_slurm_payload`, different transport. + Inside a spool script or on a local box there is no ``srun`` to create + the ranks, so ``mpirun`` does it and the wrapper assigns leader/worker + exactly as it does under Slurm. Without it the LLM API finds a world + size of 1, concludes it must create the other ranks itself, and dies + with ``MPI_ERR_SPAWN``. + + Two flags here are about running ``mpirun`` **inside someone else's + allocation**, which is the one thing this branch does and the ``slurm`` + branch never does. Neither belongs there, and the asymmetry is the + point in both cases: + + - :data:`MULTI_TASK_SLURM_SCRUB` — the allocation's task variables + describe the *step*, not the MPI world, and the launcher reads them + first. See that constant. + - ``--oversubscribe`` — Open MPI takes its slot count from the resource + manager, and a spool session is submitted as **one task** that owns + the node (measured: ``NumTasks=1`` against ``NumCPUs=144`` and + ``gres/gpu:4``). So ``-n 4`` asks for four slots against one and dies + before any Python runs, with *"There are not enough slots available + in the system"*. The allocation really does hold all of the node's + CPUs and GPUs, so the slot count is a bookkeeping guard rather than a + real constraint, and oversubscribing it costs nothing. + + Under ``slurm`` the same flag would be a hazard rather than a fix: + there ``srun --ntasks`` genuinely sizes the world, so a slot shortage + is a *mis-sized allocation reporting itself*, and oversubscribing + would hide it and run the ranks somewhere they do not fit. + + Raising the session's own ``ntasks_per_node`` instead was considered + and is wrong: it gives MPI its slots and simultaneously starts four + copies of the spool watcher, which then race for every piece of work. + The session is one process by design. + """ + return ( + f"{MULTI_TASK_SLURM_SCRUB} mpirun -n {self.ntasks} " + f"--oversubscribe --allow-run-as-root {LLMAPI_LAUNCH} {argv}" + ) + + def _slurm_payload(self, argv: str) -> str: + """Wrap ``argv`` the way this task count requires. + + Both branches end up inside ``bash -c`` because both need a shell: one + to scrub the environment before exec, the other so the wrapper and its + task read as a single command to ``srun``. + """ + if self.needs_llmapi_launch: + return f"bash -c {shlex.quote(f'{LLMAPI_LAUNCH} {argv}')}" + return self._scrub_and_exec(argv) + + def command(self, argv: str | Sequence[str], *, single_process: bool = False) -> str: + """Return ``argv`` prefixed for this environment. + + ``argv`` is a command line, not an interpreter invocation — the + caller decides whether it starts with :attr:`python`, ``pytest``, + ``trtllm-eval`` or anything else. + + ``single_process`` is for the work that is **not** engine work: a + read-only probe that imports the build, enumerates it and exits. + Such a command wants exactly one process whatever the run's world + size, and the rule for one process already exists here — it is the + ``ntasks == 1`` branch, which uses no launcher and scrubs the MPI + environment before exec. Asking for it by name is what keeps launcher + assembly inside this module: a caller that built its own + single-process line would be a second, divergent statement of a rule + this module exists to state once. + + Getting this wrong is not theoretical. The inventory probe went out + at world size 4 as ``mpirun -n 4 ... trtllm-llmapi-launch``, which is + right for a gate and wrong for a probe, and the plan phase was handed + an empty survey. + """ + if not isinstance(argv, str): + argv = " ".join(shlex.quote(str(part)) for part in argv) + if single_process and self.ntasks > 1: + # Re-asked at one task rather than answered by a parallel set of + # rules. `--ntasks`, `--ntasks-per-node`, `--gres`, `--mpi` and + # the payload wrapper all follow from the task count, and a + # launch that shrinks only the payload still starts one copy of + # the probe per rank. + return replace(self, world_size=1).command(argv, single_process=True) + if self.kind == "slurm": + return f"{self._slurm_prefix()} {self._slurm_payload(argv)}" + if self.needs_llmapi_launch: + argv = self._mpirun_launch(argv) + elif single_process: + # There is no `srun` here to carry the scrub `_slurm_payload` + # applies, and the process still inherits a step's PMI variables: + # a spool script runs inside the session's own allocation. + argv = self._scrub_and_exec(argv) + prefix = self.prefix + return f"{prefix} {argv}".strip() + + def python_command(self, argv: str | Sequence[str], *, single_process: bool = False) -> str: + """Return a Python invocation of ``argv`` prefixed for this environment.""" + if not isinstance(argv, str): + argv = " ".join(shlex.quote(str(part)) for part in argv) + return self.command(f"{shlex.quote(self.python)} {argv}", single_process=single_process) + + def describe(self) -> str: + """The block injected verbatim into every agent's prompt.""" + if self.kind == "session_spool": + return self._describe_spool() + if self.kind == "slurm_broker": + return self._describe_broker() + device_line = ( + ", ".join(str(d) for d in self.devices) if self.devices else "(allocated by Slurm)" + ) + return ( + "## Running GPU work\n" + "\n" + "Every GPU command in this run is launched exactly one way. Prefix it with:\n" + "\n" + "```\n" + f"{self.prefix}\n" + "```\n" + "\n" + f"The Python interpreter is `{self.python}`. The devices claimed for this run " + f"are: {device_line} (world size {self.world_size}).\n" + "\n" + "**Do not construct a launcher command, do not probe for one, and do not call a " + "test-runner skill.** The environment was resolved before you started and it is " + "not yours to rediscover. If this prefix does not fit what you need to run, that " + "is a **symptom to report in your final message** — not a problem to solve.\n" + ) + + def _describe_spool(self) -> str: + """The prompt block for a long-lived compute session watching a spool. + + This is the kind to prefer wherever it is available, and the reason + is not convenience. Every other kind puts a scheduler between a + command and its result: an image import (26 GB, 4-8 minutes) per + command, a queue wait per command, and — under ``slurm_broker`` — a + transport only an MCP client can reach, which the orchestrator is + not. A spool has none of that. One allocation is acquired once, and + work is a **file**, so anything that can write a file can run GPU + work: the orchestrator, an agent, or a human. + """ + spool = self.spool or "" + prefix = self.prefix or "(no extra environment)" + return ( + "## Running GPU work\n" + "\n" + "A compute session is already running and watching a spool directory. You do " + "**not** submit jobs, allocate nodes, or touch Slurm — `srun`/`sbatch` do not " + "even work from here. Work is a file.\n" + "\n" + "To run something:\n" + "\n" + f"1. Write a shell script to `{spool}/pending/.sh`. Pick a `` that " + "is unique to what you are doing; it names the log.\n" + f"2. Wait for `{spool}/done/.rc` to exist. It appears only when the job " + "has finished, and it holds the exit code.\n" + f"3. Read `{spool}/logs/.log`. **That file is the result.** Your " + "recollection of what scrolled past is not.\n" + "\n" + "The script runs inside the container on the compute node, so it sees the same " + "filesystem you do. Start it with:\n" + "\n" + "```bash\n" + f"export {prefix.replace(' ', chr(10) + 'export ')}\n" + f"export CUDA_VISIBLE_DEVICES=0 # you own device(s) 0..{max(self.world_size - 1, 0)}\n" + "```\n" + "\n" + f"The Python interpreter is `{self.python}`. World size {self.world_size}.\n" + "\n" + + ( + "This run is **multi-rank**, so an engine command has to be launched with " + "the ranks already in place:\n" + "\n" + "```bash\n" + f"{self._mpirun_launch('')}\n" + "```\n" + "\n" + "Without the wrapper the LLM API sees a world size of 1, concludes it must " + "create the other ranks itself, and dies with `MPI_ERR_SPAWN` — which the " + "container runtime forbids. Do not work around that by reducing the world " + "size.\n" + "\n" + "**Copy that line whole.** The `env -u SLURM_...` prefix and " + "`--oversubscribe` are not decoration: this `mpirun` runs inside the " + "session's own one-task Slurm allocation, so without them it either " + "refuses to start (*not enough slots*) or starts " + f"{self.ntasks} processes that each believe they are rank 0 of a world " + "of 1 and hang forever with no error.\n" + "\n" + if self.needs_llmapi_launch + else "" + ) + + "Jobs you drop run **concurrently**, so two scripts that both take device 0 " + "will contend. If you need more than one at a time, assign devices yourself.\n" + "\n" + "**Do not construct a launcher command, do not probe for one, and do not call a " + "test-runner skill.** The session was resolved before you started and it is not " + "yours to rediscover. If this protocol does not fit what you need to run, that " + "is a **symptom to report in your final message** — not a problem to solve.\n" + ) + + def _describe_broker(self) -> str: + """The prompt block for a cluster reachable only through the MCP broker.""" + in_container = self.prefix or "(no extra environment)" + return ( + "## Running GPU work\n" + "\n" + "This cluster's Slurm CLI is **unreachable** from inside this session: " + "`srun`, `sbatch` and `squeue` fail to contact the controller. Do not try " + "them, and do not conclude from their failure that the cluster is down.\n" + "\n" + "Every GPU command in this run is launched exactly one way — this call, with " + "only the quoted command changed:\n" + "\n" + "```python\n" + f"{self.broker_submit()}\n" + "```\n" + "\n" + "It returns a `job_id` immediately and does **not** wait. Then:\n" + "\n" + "1. Poll `mcp__slurm-broker__slurm_job_status(job_id)` until `JobState` " + "leaves `PENDING`/`RUNNING`.\n" + "2. Read the job's log: the `StdOut=` path that status reports, which is " + "`slurm-.out` in your home directory unless the submission set a " + "`workdir`. That file is the result; your recollection of it is not.\n" + "\n" + "**Never use the blocking `slurm_run`.** The MCP layer aborts a tool call " + "after 300 seconds of silence, and every job here is longer than that — so " + "`slurm_run` reliably reports a failure for a job that is running fine, and " + "leaves it running with its output unreachable.\n" + "\n" + f"Inside the container, prefix your command with: `{in_container}`\n" + f"The Python interpreter is `{self.python}`. World size {self.world_size}; " + f"the allocation carries {self.gpus_per_node} GPU(s) per node.\n" + "\n" + "**Do not build a different launcher, do not probe for one, and do not call a " + "test-runner skill.** If this template does not fit what you need to run, that " + "is a **symptom to report in your final message** — not a problem to solve.\n" + ) + + +def image_import_warning(env: "RunEnvironment") -> str | None: + """Warn when every job will re-import a registry image instead of mounting a file. + + ``container_image`` is passed to pyxis unchanged, and pyxis accepts either + a **registry reference** or a **local squashfs path**. The difference is + the single biggest lever on how parallel a run can be. + + A registry reference is fetched and unpacked *per allocation*. Measured on + the project this was written for: **5-7 minutes** from job start to a + usable container, for a 26 GB image. Measured again on this run, on a + loaded cluster and three times in a row: **25-35 minutes**. That second + number is what makes an isolated job per gate look unaffordable, and it is + the reason work piles onto one long-lived shared session — which has **one + failure domain**, so a leaked server, a process-group kill and a PMIx + collision each took down every other job and the orchestrator with it. + + A squashfs on shared storage mounts in seconds, so an isolated job per + workload stops being an expensive choice. The run that measured the 5-7 + minutes executed **443 spool jobs**; paying the import on each would have + cost roughly 37 hours, and at this run's 25-35 minutes it would have been + several days. + + Returns ``None`` when the image is already a local file, so a correctly + configured run is silent. An absolute path that does not exist gets its + own message: pyxis fails to start rather than falling back to a registry, + so the two are different faults with different fixes. + """ + image = (env.container_image or "").strip() + if not image or env.kind not in ("slurm", "slurm_broker", "session_spool"): + return None + if image.startswith("/") and Path(image).is_file(): + return None + if image.startswith("/"): + return ( + f"container_image points at {image}, which does not exist. pyxis will " + f"fail to start rather than fall back to a registry." + ) + return ( + f"container_image is a registry reference ({image.split('/')[-1]}), so every " + f"allocation re-imports it -- measured at 5-7 minutes for a 26 GB image, and " + f"25-35 minutes on this cluster under load. Import it once to shared storage " + f"and point at the file instead:\n" + f" enroot import -o /trtllm.sqsh docker://{image}\n" + f"then set `environment.container_image` to that path. This is what makes " + f"an isolated job per gate affordable rather than a trade against startup " + f"cost -- see `environment.workloads` in task.example.yaml." + ) + + +def shares_allocation_with_spool(env: "RunEnvironment") -> str | None: + """Warn when the orchestrator lives inside the allocation it dispatches to. + + The orchestrator is a **long-lived, GPU-free** process: it blocks for hours + on a role turn and never touches a card. The compute session is a + short-lived, GPU-heavy one that gets preempted, OOM-killed and torn down. + Putting the first inside the second couples their lifetimes for no benefit + at all. + + Measured three times in one run on the project this was written for: the + session died and took the orchestrator with it, and each time an agent turn + had already been answered on disk and the answer was thrown away because + nothing was left reading it — the most expensive of them a twelve-minute + audit. This run hit it twice more, and neither was a scheduler event: once + a disagg probe's worker teardown took the step down, and once an agent's + ``pkill -9`` caught the watcher — which *is* the allocation's main process, + so killing it ends the allocation. Both took every queued job with them. + + The evidence is ``watcher.jobid`` in the spool, written every poll by + ``spool_watcher.sh``. It is a **separate file from** ``watcher.heartbeat`` + on purpose: ``scripts/spool_submit.sh`` pipes the heartbeat straight into + ``date -d`` after stripping ``alive=``, so anything else on that line + breaks every submitter. + + Returns a message when both job ids are known **and equal**, and ``None`` + otherwise. Unknown is not a warning: an orchestrator launched in its own + job also has ``SLURM_JOB_ID`` set, and crying wolf at the correct + configuration is how a warning gets ignored. + """ + if env.kind != "session_spool" or not env.spool: + return None + mine = os.environ.get("SLURM_JOB_ID") + if not mine: + return None + marker = Path(env.spool) / "watcher.jobid" + try: + theirs = marker.read_text(encoding="utf-8").strip() + except OSError: + return None + if not theirs or theirs == "unknown" or theirs != mine: + return None + return ( + f"the orchestrator is running inside Slurm job {mine}, which is also the " + f"compute session watching {env.spool}. They will die together, and this " + f"process needs no GPU: launch it in its own allocation " + f"(`--gpus-per-node=0`) so a preempted session costs you the work in " + f"flight and not the run." + ) + + +#: Workloads that may carry their own launch substrate. Closed on purpose: a +#: name the orchestrator does not dispatch under is a typo that would silently +#: never apply, and nothing downstream would say so. Enforced by +#: :func:`~.task_schema._validate_workloads`. +WORKLOAD_KINDS: tuple[str, ...] = ("gate",) + + +def resolve_environment( + task_data: Mapping[str, Any], + *, + devices: Sequence[int] = (), + world_size: int = 1, + workload: str | None = None, +) -> RunEnvironment: + """Build a :class:`RunEnvironment` from a validated task spec. + + ``devices`` comes from preflight's claim, not from the spec: which + cards are usable is a property of the host at this moment, and the + spec cannot know it. + + ``workload`` selects an optional override from ``environment.workloads``, + because **the right substrate depends on how long the job runs**. A shared + session imports a 26 GB container once and then costs nothing per job, + which is the correct trade for a 1-minute entry test; one onboarding ran + 443 of them, and paying a 5-7 minute import each would have cost 37 hours. + It is the wrong trade for a 10-15 minute gate, where the import is 30-50% + overhead and **isolation is the thing you actually want** — a shared + session has one failure domain, so a gate that leaks a server, kills a + process group or collides on PMIx takes every other job and the + orchestrator with it. See :func:`shares_allocation_with_spool` for the + five times that happened across two runs. + + Absent an override this returns exactly what it always did. + """ + block = dict(task_data.get("environment") or {}) + if workload: + override = (block.get("workloads") or {}).get(workload) or {} + # `env` merges key-by-key; everything else replaces. A workload that + # needs a different substrate almost never needs a different *whole* + # environment, and declaring mounts and accounts again to change one key + # is how the two drift apart. + merged_env = {**(block.get("env") or {}), **(override.get("env") or {})} + block = {**block, **{k: v for k, v in override.items() if k != "env"}} + block["env"] = merged_env + env = {str(k): str(v) for k, v in (block.get("env") or {}).items()} + return RunEnvironment( + kind=str(block.get("kind", "local")), + python=str(block.get("python", "python3")), + env=env, + devices=tuple(int(d) for d in devices), + world_size=max(1, int(world_size)), + partition=block.get("partition"), + container_image=block.get("container_image"), + account=block.get("account"), + mounts=tuple(str(m) for m in (block.get("mounts") or ())), + time_limit=block.get("time_limit"), + gpus_per_node=int(block.get("gpus_per_node") or 0), + container_name=block.get("container_name"), + mpi=str(block.get("mpi") or MPI_AUTO), + spool=block.get("spool"), + ) + + +__all__ = [ + "LLMAPI_LAUNCH", + "MERGED_ENV_VARS", + "MPI_AUTO", + "MPI_NONE", + "MPI_PMIX", + "MULTI_TASK_SLURM_SCRUB", + "MULTI_TASK_SLURM_VARS", + "RunEnvironment", + "SINGLE_TASK_MPI_SCRUB", + "WORKLOAD_KINDS", + "image_import_warning", + "resolve_environment", + "shares_allocation_with_spool", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/execution.py b/agent-flow/agent_flow/workflows/staircase_onboard/execution.py new file mode 100644 index 000000000000..3b20833c2142 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/execution.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Where this run's commands land, decided once, from the resolved environment. + +`environment.py` answers *what the command line is* — the ``srun`` prefix, +the launcher wrapper, the MPI scrub. This module answers the other half, +*who runs it*, and it is the only place in the workflow that maps a +``task.yaml`` environment kind onto a mechanism. + +That split is the point. The orchestrator used to carry the branch itself:: + + if result.environment.kind == "session_spool": + self._shell_runner = gates_mod.spool_runner(result.environment.spool) + +which reads as a small thing and is not: every mechanism added since — the +broker, the spool — arrived as another line beside a stage transition, in +the file that also owns resume and gate sequencing. Deleting a mechanism +should not require reading the state machine. + +The two ``*_runner_for`` adapters exist because the call sites this replaces +took plain callables, and there are a lot of them: three gates, two +witnesses, the inventory probe, six preflight probes. Converting them all at +once would have put a mechanical rename in the same change as a structural +one. They are a migration surface, not a permanent one. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Mapping + +from agent_flow.runtime import ( + CommandExecutor, + CommandSpec, + DelegatedCommandExecutor, + ExecutionRecord, + InProcessTurnExecutor, + LocalCommandExecutor, + SlurmStepExecutor, + SpoolCommandExecutor, + TurnExecutor, + TurnOutcome, + TurnRequest, +) +from agent_flow.runtime.turns import LayerProvider + +from .environment import RunEnvironment + +#: Environment kinds the workflow knows how to resolve. Kept here rather than +#: imported from the schema because this is the list that has *executors* +#: behind it, which is not the same question as which kinds validate. +KIND_LOCAL = "local" +KIND_SLURM = "slurm" +KIND_SLURM_BROKER = "slurm_broker" +KIND_SESSION_SPOOL = "session_spool" + +#: Where a delegated run leaves its scripts, logs and receipts. Under the +#: workspace because that is on the shared filesystem both sides can see. +BROKER_DIRNAME = "broker" + + +def select_command_executor( + environment: RunEnvironment, + *, + name_prefix: str = "gate", + turns: TurnExecutor | None = None, + workspace: Path | str | None = None, +) -> CommandExecutor: + """The executor this run's commands go through. + + ``slurm_broker`` needs the two optional arguments and is the reason they + exist: that cluster's scheduler is reachable only through a tool the + *agent* holds, so submission is delegated to a turn while the verdict + stays here — the job writes its log and exit code to ``workspace/broker`` + and the orchestrator reads both itself. See + :class:`~agent_flow.runtime.DelegatedCommandExecutor`. + + Without them it falls back to the local executor, which is what the + broker path did before it was wired. That fallback is for callers with no + turn executor to hand over (a bare `--dry-run`, a test), not a mode + anybody should gate on. + """ + kind = environment.kind + if kind == KIND_SESSION_SPOOL and environment.spool: + return SpoolCommandExecutor(environment.spool, name_prefix=name_prefix) + if kind == KIND_SLURM: + return SlurmStepExecutor(max_world_size=environment.world_size or None) + if kind == KIND_SLURM_BROKER and turns is not None and workspace is not None: + return DelegatedCommandExecutor( + turns, + root=Path(workspace) / BROKER_DIRNAME, + render_submission=lambda script: environment.broker_submit(script, script=True), + name_prefix=name_prefix, + ) + return LocalCommandExecutor() + + +def shell_runner_for(executor: CommandExecutor, *, default_timeout: int = 1800) -> Any: + """Adapt ``executor`` to the ``(command, env, timeout) -> ShellResult`` shape. + + The shape `gates.ShellRunner` names, and what `witness.py` and + `survey.py` call. ``ShellResult`` is imported lazily to keep this module + importable from `gates` without a cycle. + """ + from .gates import ShellResult + + def run( + command: str, + env: Mapping[str, str] | None = None, + timeout: int = default_timeout, + ) -> ShellResult: + record = executor.execute( + CommandSpec(argv=command, env=env, timeout_s=timeout, label="shell") + ) + return _as_shell_result(record, ShellResult) + + return run + + +def command_runner_for(executor: CommandExecutor, *, default_timeout: int = 120) -> Any: + """Adapt ``executor`` to the ``(argv, env, timeout) -> CommandResult`` shape. + + The shape `preflight.CommandRunner` names. A sequence rather than a + shell line, and that difference is preserved all the way down: a probe + must not be word-split. + """ + from .preflight import CommandResult + + def run( + argv: Any, + env: Mapping[str, str] | None = None, + timeout: int = default_timeout, + ) -> CommandResult: + record = executor.execute( + CommandSpec(argv=tuple(str(part) for part in argv), env=env, timeout_s=timeout) + ) + return CommandResult( + returncode=record.returncode, + stdout=record.stdout, + stderr=record.stderr, + timed_out=record.timed_out, + ) + + return run + + +def _as_shell_result(record: ExecutionRecord, shell_result_cls: Any) -> Any: + return shell_result_cls( + returncode=record.returncode, + stdout=record.stdout, + stderr=record.stderr, + timed_out=record.timed_out, + ) + + +# --------------------------------------------------------------------- turns + + +class _CallableTurnExecutor(TurnExecutor): + """Wrap a pre-`TurnRequest` transport callable. + + The external transport was wired as ``transport(role=..., what=..., + prompt=..., system_prompt=..., decision_file=...)`` before the turn + types existed, and the workflow's constructor still accepts anything + with that shape. Keeping the door open costs one adapter; closing it + would break a caller that has no way to know it was supposed to change. + """ + + name = "callable" + + def __init__(self, call: Any) -> None: + self._call = call + + def run(self, request: TurnRequest) -> TurnOutcome: + text = self._call( + role=request.role, + what=request.what, + prompt=request.prompt, + system_prompt=request.system_prompt, + decision_file=request.decision_file, + ) + return TurnOutcome( + role=request.role, + text=text if isinstance(text, str) else str(text), + transport=getattr(self._call, "name", self.name), + decision_ref=request.decision_file, + ) + + +def select_turn_executor( + agent_transport: Any, + layer_provider: LayerProvider, +) -> TurnExecutor: + """The executor this run's role turns go through. + + ``None`` means the ordinary path: this process builds a layer per role + and calls it. Anything else means the host could not be trusted to + authenticate one, and the turn — **and only the turn** — has been moved + somewhere that can. + """ + if agent_transport is None: + return InProcessTurnExecutor(layer_provider) + if isinstance(agent_transport, TurnExecutor): + return agent_transport + return _CallableTurnExecutor(agent_transport) + + +__all__ = [ + "BROKER_DIRNAME", + "KIND_LOCAL", + "KIND_SESSION_SPOOL", + "KIND_SLURM", + "KIND_SLURM_BROKER", + "command_runner_for", + "select_command_executor", + "select_turn_executor", + "shell_runner_for", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F001-F003-F004-fix-report.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F001-F003-F004-fix-report.md new file mode 100644 index 000000000000..447b24137000 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F001-F003-F004-fix-report.md @@ -0,0 +1,325 @@ +# F001 / F003 / F004 — fix report + +- **Applied:** 2026-09-22 +- **Findings:** `20260922-F001-substrate-gate-needs-a-stock-path.md`, + `20260922-F003-surface-survey-misses-four-structural-demands.md`, + `20260922-F004-embargo-cannot-name-a-sibling-engine-checkout.md` +- **Status:** applied, tests green, not committed (working tree left dirty, + on top of the uncommitted F002 `members:` change, which was not touched) +- **Baseline:** 680 passed. **Now:** 722 passed — 42 new tests, no changes to + any pre-existing assertion. + +--- + +## F001 — a substrate gate for an architecture the stock path does not know + +### Files changed + +| file | why | +|---|---| +| `agent-flow/agent_flow/workflows/staircase_onboard/preflight.py` | `UNSUPPORTED_ARCHITECTURE_MARKERS` and `_reports_unsupported_architecture`, the one discriminator; `_substrate_boot` extracted so the run's checkpoint and the proxy take the identical path; `substrate_gate` falls back, refuses to fall back for anything else, and records the weaker claim | +| `agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py` | `substrate_checkpoint` added to `ENVIRONMENT_KEYS` (the block is closed, so an unlisted key was a validation error), plus its must-exist check in `_validate_environment` | +| `agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml` | documents the key in the `environment` block, in that file's voice — what the gate is for, what the fallback still proves, what it stops proving | +| `agent-flow/agent_flow/workflows/staircase_onboard/cli.py` | one sentence on `--skip-substrate-gate`: it is **not** the answer to an unknown architecture, because it also gives up the two checks that do not depend on the checkpoint | +| `agent-flow/tests/workflows/staircase_onboard/test_substrate_gate.py` | 10 new tests — all three marker spellings, the no-proxy message, the proxy's command, the note verbatim, the proxy's own failure, the narrowness of the discriminator, and the unchanged happy path | +| `agent-flow/tests/workflows/staircase_onboard/test_task_schema.py` | 5 new tests — the key accepted, missing-path rejected, three non-path shapes rejected | + +### The new schema surface + +```yaml +environment: + kind: local + python: /path/to/TensorRT-LLM/.venv-3.12/bin/python + env: + TRTLLM_MODELING_V2: require + TMPDIR: /path/to/scratch/tmp + CLAUDE_CODE_TMPDIR: /path/to/scratch/tmp + + # A checkpoint the STOCK modeling path already supports. Optional. + substrate_checkpoint: /path/to/Qwen3-30B-A3B +``` + +Optional, validated to exist when set. The must-exist check is not decoration: +the proxy is reached for **only after** the run's own checkpoint has already +failed, which is the longest step in preflight, so a missing proxy turns one +failure into two with a wait in between. + +### The discriminator, and where it comes from + +```python +UNSUPPORTED_ARCHITECTURE_MARKERS = ( + "Unknown architecture for AutoModelForCausalLM", + "Unknown architecture for AutoModelForMultimodalEncoder", + "is not supported in TRT-LLM yet", +) +``` + +Both registries that can answer "I have never heard of this architecture" are +keyed on the checkpoint's own `architectures[0]`, which is exactly the string +this workflow exists to introduce: + +- `tensorrt_llm/_torch/models/modeling_auto.py` — `AutoModelForCausalLM.from_config` + raises `ValueError(f"Unknown architecture for AutoModelForCausalLM: ...")` when + `_resolve_class` comes back `None`, and its multimodal-encoder branch raises the + same sentence with `AutoModelForMultimodalEncoder` in it. `_resolve_class` + returns `None` because `get_registered_model_class` + (`tensorrt_llm/_torch/models/modeling_utils.py`) resolves the lazy built-in + provider and then misses in `MODEL_CLASS_MAPPING`. +- `tensorrt_llm/models/automodel.py` — `AutoConfig.from_hugging_face` and + `AutoModelForCausalLM.get_trtllm_model_class` raise + `NotImplementedError(f"The given huggingface model architecture {hf_arch} is not + supported in TRT-LLM yet")` when `MODEL_MAP.get` returns `None`. + +Matched on **text** rather than on an exception type because the gate runs the +engine as a subprocess and reads a log: there is no object to catch. Each marker +is the invariant half of its message — the architecture name is interpolated in +and varies per run. The constant carries a comment naming both raise sites. + +`test_a_real_host_failure_is_never_answered_by_measuring_another_model` is the +test that holds the line: it declares a proxy, fails the first boot with +`Impossible to fit in any sequence in kvCache`, and asserts the runner was +called **once**. Widening the fallback to "any failure" would answer a real +kvCache fault on the target by silently booting a different model and reporting +green. + +### What is recorded when the fallback fires + +``` +notes: + - substrate proven on `/proxy-ckpt`; this checkpoint's own geometry is + unverified, so a kvCache sizing fault will surface at the first boot gate + rather than here. +``` + +Verbatim from the finding, and asserted verbatim by a test. A note that said +only "substrate gate passed" would be read later as the claim the unqualified +gate makes, and the difference between the two is precisely which failure this +run has been protected from. `MPI_ERR_SPAWN` and NVRTC's missing `cuda.h` still +are; kvCache sizing no longer is. + +When the architecture is unknown and no proxy is declared, the error names the +fix rather than the architecture: + +``` +substrate gate: the stock modeling path does not know this checkpoint's +architecture, which is the normal state of affairs for a model this workflow +is onboarding — it is not a fault in the checkpoint or the host. ... +Declare `environment.substrate_checkpoint`: ... +Do not reach for `--skip-substrate-gate` — that also gives up the MPI spawn and +NVRTC checks, which do not depend on the checkpoint at all. +``` + +--- + +## F003 — four structural demands moved out of `unrecognized` + +### Files changed + +| file | why | +|---|---| +| `agent-flow/agent_flow/workflows/staircase_onboard/survey.py` | `_repeating_period` and `_sink_layer_kind` helpers; `_HYBRID_PATTERN_KEYS` / `_HYBRID_GEOMETRY_KEYS` / `_HYBRID_KIND_PREFIXES` / `_CLAIMED_SUBSTRINGS` tables; the `attention.hybrid` and `attention.sink` probes; `attention.head_dim` and `moe.routing` extended; `swa_`, `local_`, `full_attention_`, `v_head_dim` added to `_CLAIMED_PREFIXES` | +| `agent-flow/agent_flow/workflows/staircase_onboard/runtime_index.yaml` | `sparse_attention` claims `attention.sink` and `attention.hybrid`; `kv_cache` claims `attention.hybrid` — a stack whose two kinds declare different KV head counts is a pool question, not a mask one. Required, not cosmetic: `test_survey_probes.py::test_every_demand_a_runtime_area_declares_is_one_some_probe_can_produce` fails on a demand no area claims | +| `agent-flow/tests/workflows/staircase_onboard/test_survey_probes.py` | 24 new tests across the four probes, the prefix shrinkage, the period helper and the two rejection paths (`moe_layer_freq` must not fire; `sinks: 4` must stay `attention.sinks`) | + +### Before + +10 demands, 23 unrecognized. + +``` +attention.shape: GQA, 64 query heads over 4 kv head(s) [num_attention_heads, num_key_value_heads] +attention.head_dim: 192, declared not derived [head_dim] +attention.windowing: sliding_window=128 [sliding_window] +rope.plain: theta=10000000.0 [rope_parameters] +rope.partial: rotary applied to 0.334 of each head [partial_rotary_factor] +mlp.activation: silu [hidden_act] +moe.routing: 256 experts, top-8, scoring=sigmoid, norm_topk_prob=True [num_experts] +speculation.heads: 3 in-checkpoint predictor layer(s) — weights exist for a speculator, so MTP is the checkpoint's own and not a bolt-on [num_nextn_predict_layers] +quantization: fp8 [quantization_config] +multimodal.tower: vision, audio tower in front of the language model; placeholder tokens declared (image_token_id, video_token_id, audio_token_id, vision_start_token_id) [vision_config, audio_config, image_token_id, video_token_id, audio_token_id, vision_start_token_id] +--- +add_full_attention_sink_bias +add_swa_attention_sink_bias +attention_chunk_size +attention_projection_layout +attention_value_scale +audio_end_token_id +audio_start_token_id +audio_token_id +hybrid_block_size +hybrid_layer_pattern +image_token_id +layernorm_epsilon +processor_config +swa_head_dim +swa_num_attention_heads +swa_num_key_value_heads +swa_rope_theta +swa_v_head_dim +v_head_dim +video_token_id +vision_end_token_id +vision_model_type +vision_start_token_id +``` + +### After + +12 demands, 14 unrecognized. + +``` +attention.shape: GQA, 64 query heads over 4 kv head(s) [num_attention_heads, num_key_value_heads] +attention.head_dim: 192, declared not derived; v_head_dim=128 differs from head_dim=192 — asymmetric qk/v head dimensions, which not every attention kernel admits [head_dim, v_head_dim] +attention.windowing: sliding_window=128 [sliding_window] +attention.hybrid: 48 layers in a pattern with no fixed period: 39 of `1`, 9 of `0`; which value means which kind of attention is the config's own convention and is a question for the plan phase, not something this survey can know. The two kinds declare DIFFERENT geometry: swa_num_key_value_heads=8 against num_key_value_heads=4, swa_rope_theta=10000.0 against rope_theta=10000000.0 — two shapes in one stack, not one shape under two masks [hybrid_layer_pattern, swa_num_attention_heads, swa_num_key_value_heads, swa_head_dim, swa_v_head_dim, swa_rope_theta] +attention.sink: a learned sink term in the softmax denominator, declared per layer kind: full layers=False, swa layers=True [add_full_attention_sink_bias, add_swa_attention_sink_bias] +rope.plain: theta=10000000.0 [rope_parameters] +rope.partial: rotary applied to 0.334 of each head [partial_rotary_factor] +mlp.activation: silu [hidden_act] +moe.routing: 256 experts, top-8, scoring=sigmoid, topk_method=noaux_tc, n_group=1, topk_group=1, norm_topk_prob=True [num_experts] +speculation.heads: 3 in-checkpoint predictor layer(s) — weights exist for a speculator, so MTP is the checkpoint's own and not a bolt-on [num_nextn_predict_layers] +quantization: fp8 [quantization_config] +multimodal.tower: vision, audio tower in front of the language model; placeholder tokens declared (image_token_id, video_token_id, audio_token_id, vision_start_token_id) [vision_config, audio_config, image_token_id, video_token_id, audio_token_id, vision_start_token_id] +--- +attention_chunk_size +attention_projection_layout +attention_value_scale +audio_end_token_id +audio_start_token_id +audio_token_id +hybrid_block_size +image_token_id +layernorm_epsilon +processor_config +video_token_id +vision_end_token_id +vision_model_type +vision_start_token_id +``` + +Nine keys dropped out of the unrecognized list: the two sink flags, the five +`swa_*` geometry overrides, `hybrid_layer_pattern` and `v_head_dim`. Every one +of them is now cited as **evidence** for a demand instead, so the list has +shrunk by exactly what a probe started claiming and by nothing else. + +### The polarity, deliberately not asserted + +`01-architecture/` establishes that in this checkpoint's `hybrid_layer_pattern`, +`1` is sliding-window and `0` is full attention — 39 against 9 of 48. The probe +does **not** encode that. Nothing in `config.json` states it, the next +checkpoint may invert it, and a survey that guessed would silently window the +wrong 39 layers. So the demand reports the counts factually, keyed by the raw +value (``39 of `1`, 9 of `0```), and ends with "which value means which kind of +attention is the config's own convention and is a question for the plan phase". +`test_the_survey_reports_the_pattern_and_refuses_to_interpret_it` asserts both +the sentence and the absence of any polarity claim. + +--- + +## F004 — an embargo that can name an out-of-repo tree + +### Files changed + +| file | why | +|---|---| +| `agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py` | `_validate_source_embargo` accepts an absolute path that resolves outside `trtllm_repo_path` (must-exist kept), still rejects one that resolves inside with a message naming the fix, and its docstring states that a directory entry embargoes the tree beneath it | +| `agent-flow/agent_flow/workflows/staircase_onboard/prompts/_common.py` | `source_embargo_block` splits the opening premise and the "running is expected" paragraph on where the paths point, and states the directory rule in the rendered block | +| `agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml` | the `source_embargo` comment documents both spellings, the in-repo rejection, the directory rule, and that running either kind is still expected | +| `agent-flow/tests/workflows/staircase_onboard/test_runtime_and_scope.py` | 6 new tests — out-of-repo accepted, in-repo absolute rejected, out-of-repo missing rejected, and three on the rendered block (out-of-repo only, both kinds, the directory sentence) | + +### The rule, and the discriminator in the prompt + +Validation splits on `Path(entry).resolve()` against `Path(repo).resolve()`. +The rendered block splits on `p.startswith("/")` — which is not a heuristic, +because validation has already guaranteed that an absolute entry resolves +outside the repo. No new parameter was threaded through `prompts/__init__.py`. + +``` +## Embargoed sources + +This checkpoint already has an implementation in this repository. This +checkpoint is implemented in another engine's checkout on this filesystem, +outside this repository. For this run it is **off limits to read**: + +- `tensorrt_llm/_torch/models/modeling_qwen3_5.py` +- `/lustre//portfolios//users//vllm` + +An entry that names a directory embargoes the whole tree beneath it, not just +the directory itself. +... +**Running the in-repository implementation is expected.** Under +`TRTLLM_MODELING_V2=off` the engine uses exactly that implementation, ... + +**Running the out-of-repository engine as a black-box oracle is expected; +reading it is not.** Start its server, send it the workload and measure +against what it returns — that is why the checkout is named here rather than +merely avoided. What it must never be is a file you open. +``` + +A run with only in-repo entries renders byte-for-byte what it rendered before, +apart from the directory sentence. + +--- + +## Test result + +``` +722 passed in 21.14s +``` + +(`cd agent-flow && ../.venv-orch/bin/python -m pytest tests/workflows/staircase_onboard/ -q`) + +Baseline before these three changes was **680 passed** — the tree as F002 left +it — so the 42 new tests are the entire delta and nothing pre-existing changed +shape. `pre-commit` was not run and nothing was committed. + +--- + +## What the findings did not anticipate + +1. **`runtime_index.yaml` is a hard dependency of a new probe, not a + follow-up.** `test_survey_probes.py::test_every_demand_a_runtime_area_declares_is_one_some_probe_can_produce` + asserts in both directions: an area may only claim demands the survey can + emit, **and** a probe may only emit demands some area claims. So + `attention.hybrid` and `attention.sink` could not land without the index + entry, which is the invariant working exactly as written — a demand no area + claims is a checkpoint feature that maps to no engine knowledge at all. It + also forced a real decision rather than a rubber stamp: `attention.hybrid` + is claimed by **two** areas, because differing per-kind KV head counts are a + question for the cache manager as much as for the attention backend. + +2. **The prefix list cannot reach the sink keys, and the finding's proposed + prefixes do not cover `hybrid_layer_pattern` either.** `swa_`, `local_` and + `full_attention_` clear the five `swa_*` geometry overrides, but + `add_swa_attention_sink_bias` has the layer kind in the *middle* of the name + and `hybrid_layer_pattern` matches none of the three. Adding `add_` as a + prefix would claim every `add_*` key in every config. Introduced + `_CLAIMED_SUBSTRINGS = ("attention_sink", "layer_pattern")` — two entries, + both being exactly what the new probes glob on, so the claim and the probe + cannot drift. + +3. **Naive period detection reports a period that occurs once.** The first + implementation said the 48-entry pattern repeated "every 47 layers", because + `values[47] == values[47 % 47]` happens to hold when the last entry equals + the first. Candidates now stop at half the length. Kept as its own test, + because a stack whose global layers land every sixth position and a stack + whose first block is short are genuinely different stacks and the survey is + the only thing that says which this is. + +4. **`attention_chunk_size` was left in `unrecognized` on purpose.** F003's + table pairs it with `hybrid_layer_pattern`, but it is a window *length* + rather than a per-layer structure, `attention.windowing` already reports + `sliding_window=128`, and the `sparse_attention` area's own summary calls it + out by name. Claiming it would have meant widening the windowing probe, + which no finding asked for. Flagging it here rather than doing it silently. + +5. **`--skip-substrate-gate`'s help was part of F001's surface.** The finding + correctly reads that flag as the wrong shape for this problem; its help text + still recommended it as the escape hatch of last resort with no mention of + the new one. One sentence added, so the operator who hits the architecture + error and reaches for `--help` is pointed at `substrate_checkpoint` instead. + +6. **The F002 tree was untouched.** `members_for`, the member-set gate loop and + the witness member all pass through `substrate_gate` and `source_embargo` + unchanged — the substrate gate runs before any member is resolved (it is a + host claim, not a topology one) and the embargo is rendered per role, not + per member. Nothing in this change needed to know about the member set. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F001-substrate-gate-needs-a-stock-path.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F001-substrate-gate-needs-a-stock-path.md new file mode 100644 index 000000000000..84c417bf4228 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F001-substrate-gate-needs-a-stock-path.md @@ -0,0 +1,85 @@ +# F001 — the substrate gate assumes a stock path exists for the checkpoint + +- **Found:** 2026-09-22 03:30 PDT, during preflight planning for MiMo-V2.6-Flash-RL +- **Component:** `agent-flow/agent_flow/workflows/staircase_onboard/preflight.py:385-444` +- **Severity:** blocks preflight for every genuinely new architecture +- **Status:** proposed fix drafted, not yet applied + +## What happens + +`substrate_gate()` proves the host can generate a token by running the boot +gate on **the run's own checkpoint** with `mode=MODE_OFF` — that is, the +stock TensorRT-LLM modeling path: + +```python +result = gates_mod.run_boot_gate( + environment, + task_data=task_data, + spec=gates_mod.BootSpec(prompts=("The capital of France is",), max_tokens=4), + runner=shell_runner, + timeout=SUBSTRATE_TIMEOUT_S, + mode=gates_mod.MODE_OFF, + features=[], +) +``` + +For MiMo-V2.6-Flash-RL there is no stock path — `grep -rl "MiMoV2\|mimo_v2" +tensorrt_llm/ tests/ examples/` returns nothing. The gate therefore fails +with an architecture-not-supported error, and `run_preflight` raises +`PreflightError` before any agent exists. + +## Why this is the workflow's problem and not the run's + +The docstring is explicit about what the gate is for: three host failures +(`MPI_ERR_SPAWN`, kvCache sizing, NVRTC's missing `cuda.h`) that "surfaced +mid-run, inside an agent's turn, reading as though the target were wrong". +All three are properties of **the host and the resolved runtime block**. + +The gate reaches for the run's own checkpoint for a good reason — two of the +three failures depend on the checkpoint's geometry — but that reason only +holds for a checkpoint the stock path can load. The onboarding workflow's +entire purpose is checkpoints it cannot. So the gate is, by construction, +unavailable exactly when a run is doing the thing the workflow exists for. + +The escape hatch today is `--skip-substrate-gate`, and its own help says to +use it "only when you have already proven the substrate in this same +environment". That is sound advice and it is also the wrong shape: it turns +the highest-value check into an all-or-nothing switch, and the run loses the +`MPI_ERR_SPAWN` and NVRTC checks — which do **not** depend on the +checkpoint — along with the kvCache check, which does. + +## Proposed change (minimal) + +Split the claim the gate makes into the part that needs this checkpoint and +the part that does not. + +1. Try the stock path on the run's own checkpoint, as today. If it passes, + nothing changes and the note is what it is now. +2. If it fails **because the architecture is unsupported** — a recognizable + class of error, not any failure — fall back to a declared proxy + checkpoint and re-run the same gate under the same environment and the + same `runtime` block. Record the weaker claim verbatim in the notes: + *"substrate proven on ``; this checkpoint's own geometry is + unverified, so a kvCache sizing fault will surface at the first boot gate + rather than here."* +3. If the fallback is not declared, fail with a message that says to declare + one — not with the raw architecture error, which reads as a target defect + and is the exact failure mode the gate was written to prevent. + +The proxy is a new optional key, `environment.substrate_checkpoint`. It is +optional because a checkpoint the stock path already supports needs no proxy, +and closed-schema validation means it has to be declared to be usable. + +## Why not just widen the fallback to any failure + +Because then a real kvCache fault on the target checkpoint would be answered +by silently measuring a different model and reporting green — the same class +of mistake as gating under `TRTLLM_MODELING_V2=auto`. The fallback has to be +keyed on "this path does not know this architecture" and nothing else. + +## Interim action for this run + +Prove the substrate by hand on the allocation with a proxy checkpoint, record +the evidence in `02-environment/`, then pass `--skip-substrate-gate` with a +pointer to that evidence. Apply the fix afterwards so the next run does not +need the manual step. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F002-fix-report.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F002-fix-report.md new file mode 100644 index 000000000000..7228dd961eba --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F002-fix-report.md @@ -0,0 +1,160 @@ +# F002 — fix report: a `unified` target is now gated across its member set + +- **Applied:** 2026-09-22 +- **Finding:** `20260922-F002-unified-targets-are-gated-at-tp-only.md` +- **Status:** applied, tests green, not committed (working tree left dirty) + +## Files changed + +| file | why | +|---|---| +| `agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py` | new `members` top-level key, its validator `_validate_members`, and the single resolver `members_for(data, unit, world_size)` that both the workflow and the tests read | +| `agent-flow/agent_flow/workflows/staircase_onboard/gates.py` | `topology_flags(segment, world_size, member=None)` translates the member when given; `run_boot_gate` / `run_accuracy_gate` / `run_acceptance_gate` (and `eval_command`) take `member`; `GateResult` carries it and `summary()` prints it | +| `agent-flow/agent_flow/workflows/staircase_onboard/witness.py` | `witness_command` and both witness gates take `member`, so a witness is served at a real member's topology instead of at whatever `unified` used to fold to | +| `agent-flow/agent_flow/workflows/staircase_onboard/workflow.py` | `_gate_over_members` helper; the per-unit block resolves the member set once and loops boot, accuracy and acceptance over it; `_run_witness` takes the member | +| `agent-flow/agent_flow/workflows/staircase_onboard/preflight.py` | the unified world size now comes from the member set (it was hardcoded to 4); the member set and its gate count are printed as a preflight note, which is what `--dry-run` shows | +| `agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml` | documents `members:` next to `parallel:` — both shapes, all three rules, and the cost | +| `agent-flow/agent_flow/workflows/staircase_onboard/README.md` | one paragraph in "The gates" | +| `.../skills/staircase-modeling-parallel/references/unified.md` | the Gating section now says how the member set is declared and that the orchestrator enforces it, rather than stating an aspiration | +| `agent-flow/tests/workflows/staircase_onboard/test_task_schema.py` | 10 new tests: defaults, explicit list, per-unit mapping over `default`, and every rejection path | +| `agent-flow/tests/workflows/staircase_onboard/test_gates.py` | 5 new tests: `topology_flags` per member kind, the no-member `unified` fallback, member-overrides-segment, and the member in `GateResult.summary()` | +| `agent-flow/tests/workflows/staircase_onboard/test_workflow.py` | 3 new end-to-end tests: three members → three boots and three accuracy runs at three topologies; a second-member failure stopping the unit; a one-member run | +| `agent-flow/tests/workflows/staircase_onboard/test_gate_wiring.py` | **pre-existing tests, adjusted** — two assertions sliced the source of `_assemble_one_unit` on the literal `")\n )"`, which asserts indentation rather than argument routing. Replaced with a `_call_block(source, symbol)` helper that slices by line count. Same intent, no longer keyed on how deeply the call is nested. | + +## The new schema surface + +`members` is a list of concrete parallel segments, or a mapping from assembly +unit to one layered over `default`. It is legal only with `parallel: unified`. + +```yaml +parallel: unified + +# Applies to every unit. +members: [tp4, tep4, dep4] + +# Or, per unit — the knob for buying a cheaper answer honestly: +# members: +# default: [tp4, tep4, dep4] +# multimodal: [tep4] +``` + +Rules, each enforced with a message that names the fix: + +- **only under `parallel: unified`** — a concrete segment already names the + one configuration, and two declarations of the same thing can disagree. +- **`unified` is not a legal member** — it is the family, not a member of it; + accepting it reintroduces the TP-only gate under a new name. +- **one world size across all members** — one allocation serves them all; a + mismatch is reported naming both members and both numbers. +- **default when omitted**: `[tp, tep, dep]` — the baseline pair plus + the degenerate path, at the unified world size. + +Resolution lives in exactly one place, `task_schema.members_for(data, unit, +world_size)`. For a concrete segment it returns `[that segment]`, so the +orchestrator has one code path and the single-configuration run is the +degenerate case rather than a branch. + +Every verdict now names its member: + +``` +[PASS] boot@tep4: engine cold-started and produced non-empty generation (61s) +[PASS] accuracy[gsm8k]@dep4: measured 95.9000 vs reference 95.5300 - tol 5.00 (margin +5.3700) in 512s +``` + +## Cost, stated before it is spent + +`preflight` appends a note that `format_preflight_report` prints — and +`--dry-run` stops right after it: + +``` +notes: + - unified segment: world size 4, taken from the member set (tp4, tep4, dep4) + - members: tp4, tep4, dep4 -> 3 x (boot + accuracy) per unit, 12 accuracy gate(s) over 4 unit(s): core, mtp, multimodal, host_offload +``` + +When the member sets differ per unit the note switches to the per-unit form +(`core: tp4, tep4, dep4; multimodal: tep4`) and the total is still exact — +10, not 12, for that example. + +## The witness gates: first member only + +**Decision: `host_offload` and `disagg` run their witness at `members[0]`, +not over the whole set.** The reasoning, which is in the code comment at the +call site: + +Boot and accuracy ask *"does this configuration work"*, and the answer +genuinely differs per member — that is the entire finding. A witness asks +something else. It is a **contrast between two legs of the same +configuration**: a device-memory delta with token identity for `host_offload`, +token identity across the prefill/decode hop for `disagg`. What it proves is a +property of the modeling code — the tensor is not resident, the per-sequence +state survives the hop — not a property of the topology. + +The residency witness makes that concrete and decisive. Its verdict is a delta +against `accuracy.witness..expected_mib`, a **single** number the spec +derives from the checkpoint's own shapes. It is not a per-member number: TEP +shards attention weights and DEP replicates them, so the resident baseline +moves between members while the expectation does not. Looping the witness +would therefore compare each member's delta against a prediction computed for +one of them — a gate that fails or passes for a reason unrelated to the +feature. + +Cost seals it: each witness is two full engine launches, so it is the most +expensive gate shape in the set, and the loop would triple it to re-ask a +question whose answer does not vary with the member. + +The member is still **passed through** rather than left unset. A witness run +with `member=None` under `unified` resolves to pure tensor parallelism — a +configuration the target may not even be delivered at — so `witness_command`, +`run_residency_witness` and `run_equivalence_witness` all take `member` and +stamp it on the verdict (`[PASS] residency@tep4: ...`). + +## Test result + +``` +680 passed in 22.00s +``` + +(`cd agent-flow && .venv-orch/bin/python -m pytest tests/workflows/staircase_onboard/ -q`) + +Baseline before the change was **662 passed**, no failures — so nothing here +is a pre-existing failure being carried, and the 18 new tests are all of the +delta. `pytest` and `pytest-asyncio` had to be installed into `.venv-orch` +first; they were absent. + +## What the finding did not anticipate + +1. **`preflight` hardcoded the unified world size to 4.** The finding treats + the world size as settled ("`preflight.py` already defaults that to 4") and + proposes `members` purely as a gating list. But `parse_parallel_segment`'s + own docstring says the member set decides the world size, and preflight was + ignoring that — so `parallel: unified` plus `members: [tp8, tep8, dep8]` + would have **claimed four devices and gated at eight**, or more likely + failed the device claim with a number nobody asked for. Caught by the + first workflow-level test, which declared world-size-1 members and watched + preflight try to claim 4. Fixed: the world size is read from the first + member (the schema guarantees they agree), with 4 as the documented + fallback when no member is declared either. + +2. **The witness gates were in the blast radius.** `witness.py` calls + `topology_flags` through `witness_command`, so making `topology_flags` + member-aware without touching it would have left both witnesses running at + TP-only under `unified` — the original defect, surviving in the one place + the finding does not mention. They now take the member; see the decision + above for why they take *one*. + +3. **Two pre-existing tests asserted formatting, not routing.** + `test_gate_wiring.py` sliced `_assemble_one_unit`'s source on + `")\n )"` — twelve literal spaces — to isolate a gate call. + Moving the calls one level in to loop over members broke them for a reason + that has nothing to do with what they check. Rewritten to slice by line + count; the assertions themselves (`unit=unit` reaches the boot and accuracy + gates) are unchanged. + +4. **`GateResult.summary()` was the only label anything reads.** The + `progress.yaml` entries, the journal, the reviewer prompt, `_close_run` and + the auditor prompt all funnel through it, so putting the member there + (`name@member`) covered every consumer with one change — but it also means + `GateLog.last("boot")` still matches on `name`, deliberately untouched, so + the delivery phase's "did a boot gate pass" question keeps working + unchanged across members. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F002-unified-targets-are-gated-at-tp-only.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F002-unified-targets-are-gated-at-tp-only.md new file mode 100644 index 000000000000..f97783f8267c --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F002-unified-targets-are-gated-at-tp-only.md @@ -0,0 +1,89 @@ +# F002 — a `unified` target is gated at TP only, and claims a member set it never ran + +- **Found:** 2026-09-22 03:45 PDT, deciding how to express TP4 / TEP4 / ADP4 +- **Component:** `agent-flow/agent_flow/workflows/staircase_onboard/gates.py`, `topology_flags` +- **Severity:** a unified run reports green for configurations it did not measure +- **Status:** fix designed, not yet applied + +## What happens + +`topology_flags` translates the parallel segment into `quickstart_advanced.py` +flags. Its own docstring states the invariant it exists to protect: + +> Every gate is **topology-blind** on its own: a `tep4` target served at tp2 +> boots, scores, and produces a self-consistent curve while the directory +> name lies. + +For the three concrete segments it honours that exactly — `tep` adds +`--moe_ep_size`, `dep` adds `--moe_ep_size` and `--enable_attention_dp`. For +`unified` it returns `[--tp_size=N]` and nothing else. + +So a unified target is built once and gated once, **at pure tensor +parallelism**. Expert parallelism is never switched on, attention data +parallelism is never switched on, and the two code paths that exist only to +serve those configurations are never executed by any gate. + +## Why that is worse than it sounds + +The parallel skill's own reference is unambiguous about the intended +contract (`skills/staircase-modeling-parallel/references/unified.md`): + +> Each member runs the **full** accuracy gate, not a spot check. There is no +> static equivalence proof to lean on: an earlier experiment carried one — an +> AST-level check comparing the unified target's vocabulary, weight manifest +> and per-layer op sequence against each sibling's — and it was deliberately +> removed once `unified` became the default form... +> +> So the gates are the only evidence that the members agree. + +The static proof was removed *because the gates were going to carry the +claim*. The gates do not carry it. The member set is prose in `plan.md` and +in the target's Coverage section, and nothing in the orchestrator reads it. + +Worse, `unified` is `DEFAULT_PARALLEL` (`task_schema.py`), so a `task.yaml` +that simply omits `parallel:` gets the configuration with the weakest +gating and the broadest claim. + +## What this costs the MiMo run specifically + +The request is TP4 **and** TEP4 **and** ADP4. Under today's behaviour the run +would: + +- build one codebase whose `self.ep_size > 1` and `self.dp` branches are + live code, +- gate it at `--tp_size=4` only, where both branches fold off, +- and deliver a target documented as covering three configurations, two of + which never ran. + +For a 256-expert MoE the expert-parallel branch is not a minor variant — it +is where the expert weight sharding, the all-to-all and the per-rank expert +slice live. An EP bug is invisible to a TP-only gate by construction. + +## Proposed change + +1. **`task_schema.py`** — a new optional top-level key `members`, a list of + parallel segments, legal only when `parallel: unified`. Each entry is + validated by the existing `parse_parallel_segment`, and each must carry + the same world size, since one allocation serves them all. Default when + omitted: `[tp, tep, dep]` where `W` is the unified world size — + which is what `references/unified.md` already calls the baseline pair + plus the degenerate path. +2. **`gates.py`** — `topology_flags(segment, world_size, member=None)`. + When `member` is given it is what gets translated; the `unified` branch + stops being a special case and becomes "translate whichever member you + were handed". +3. **`workflow.py`** — at each unit's gate boundary, loop the boot and + accuracy gates over the member set instead of running them once. Record + one verdict per member in `progress.yaml`; the unit passes when every + member passes. +4. **The cost has to be stated where it is chosen**, not discovered: with 3 + members and 4 units that is 12 accuracy gates. `members:` is therefore the + knob a run uses to buy a cheaper answer honestly — one member, named — + rather than getting a cheap answer labelled as an expensive one. + +## Why not just run staircase three times + +Three runs with `parallel: tp4` / `tep4` / `dep4` produce three targets and +three codebases, which is the thing `unified` was measured to avoid (58% less +code on deepseek-v3-lite). It would also triple the smith and assembler +budget, not just the gate budget. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F003-surface-survey-misses-four-structural-demands.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F003-surface-survey-misses-four-structural-demands.md new file mode 100644 index 000000000000..cab4e8740600 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F003-surface-survey-misses-four-structural-demands.md @@ -0,0 +1,102 @@ +# F003 — the surface survey reports four structural obligations as "unrecognized" + +- **Found:** 2026-09-22 03:55 PDT, by running `survey.read_checkpoint_surface` + against this checkpoint's `config.json` before preflight +- **Component:** `agent-flow/agent_flow/workflows/staircase_onboard/survey.py` +- **Severity:** the plan drafter is handed prose where it should be handed a demand +- **Status:** reported; fix deferred until the architecture analysis confirms semantics + +## How this was found + +The survey is pure Python and takes a config mapping, so it can be run on a +checkpoint before committing a single GPU-hour: + +``` +.venv-orch/bin/python -c " +import json +from agent_flow.workflows.staircase_onboard import survey +cfg = json.load(open('/config.json')) +s = survey.read_checkpoint_surface(cfg) +..." +``` + +That is worth knowing on its own: **the first half of preflight is a +CPU-only, seconds-long dry run of what the drafter will see**, and nothing +in the docs points at it as a standalone tool. + +## What it got right + +Ten demands, all correct, including two that matter a lot here: + +``` +speculation.heads: 3 in-checkpoint predictor layer(s) — weights exist for a + speculator, so MTP is the checkpoint's own and not a bolt-on +multimodal.tower: vision, audio tower in front of the language model; + placeholder tokens declared +``` + +## What it missed + +Twenty-three keys came back `unrecognized`. Most of those are genuinely +inert (`processor_config`, the token ids, `layernorm_epsilon`). **Four are +not** — each is a structural obligation with a kernel consequence, and each +has a skill reference already written for it: + +| key(s) | what it means | the skill that already covers it | +|---|---|---| +| `hybrid_layer_pattern` (48 entries), `attention_chunk_size` | the layer stack **alternates** local and global attention; the two kinds have *different* head counts | `staircase-modeling-core/references/sliding-window-attention.md` | +| `add_swa_attention_sink_bias`, `add_full_attention_sink_bias` | a learned sink term in the softmax denominator, on the local layers only | `staircase-modeling-core/references/attention-sinks.md` | +| `v_head_dim` (128) against `head_dim` (192) | **asymmetric qk / v head dimensions** — the single hardest kernel-feasibility question this checkpoint asks | `staircase-kernel-feasibility` | +| `swa_num_key_value_heads` (8) against `num_key_value_heads` (4), `swa_rope_theta` (1e4) against `rope_theta` (1e7) | the two attention kinds are different *shapes*, not the same shape with a mask | — | + +The windowing demand that *is* emitted says `sliding_window=128` and nothing +else, which reads as "one sliding window over a uniform stack". That is the +wrong picture: it is a 48-layer stack where every fifth or sixth layer is +global, the global layers have 4 KV heads and the local ones have 8, and the +two use RoPE bases three orders of magnitude apart. + +## Why this matters more than a missing line of prose + +`survey.py`'s own comment is explicit that the unrecognized list is allowed to +be incomplete, and that its incompleteness "costs a line of prose rather than +a wrong plan". That reasoning holds for `processor_config`. It does not hold +for these four, because each one changes what the *vocabulary* must contain: + +- asymmetric qk/v head dims decide whether any existing paged-attention entry + can be reused at all, or whether the target needs a padded-V fallback; +- the sink bias decides whether the attention entry is the stock one; +- the alternating pattern decides whether the KV cache manager is asked for + one pool or two, with different per-layer geometry; +- the differing KV head counts decide the weight manifest's shard math. + +A demand the drafter does not see is a queue item the smith phase never +creates, and that surfaces as a `BLOCKED: vocabulary_gap` from the assembler +hours later — the expensive end of the run. + +## Proposed change + +Four new probes in `survey.py`, in the same shape as the existing ones: + +1. `attention.hybrid` — fires on a per-layer pattern key + (`hybrid_layer_pattern`, `layer_types`, `full_attention_idx`, ...). Reports + the period, the counts of each kind, and whether the two kinds declare + different head geometry (scanning for the `swa_*` / `*_local_*` prefix + family against their global counterparts). +2. `attention.sink` — fires on any `*attention_sink*` key, reporting which + layer kind carries it. +3. `attention.head_dim` — extend the existing probe so that a `v_head_dim` + differing from `head_dim` is stated as its own clause, rather than being + dropped. This is a one-line change to a probe that already exists. +4. `moe.routing` — extend to report `topk_method` (`noaux_tc` here) and + `n_group`/`topk_group`, which decide the routing entry, not just the + scoring function. + +Also worth doing, and cheaper: add the prefixes `swa_`, `local_`, and +`full_attention_` to the "already accounted for" prefix list once probe 1 +lands, so the unrecognized list shrinks to things that really are inert. + +## Deliberately not proposed + +Making the unrecognized list a failure. It is right that it is prose. The +change is to move four specific keys out of it, not to change what the list +means. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F004-embargo-cannot-name-a-sibling-engine-checkout.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F004-embargo-cannot-name-a-sibling-engine-checkout.md new file mode 100644 index 000000000000..9ca76bfffbcb --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F004-embargo-cannot-name-a-sibling-engine-checkout.md @@ -0,0 +1,87 @@ +# F004 — `source_embargo` cannot name the trees most likely to be copied from + +- **Found:** 2026-09-22 04:10 PDT, while choosing the accuracy-anchor engine +- **Component:** `agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py`, + `_validate_source_embargo` +- **Severity:** the run's central claim — "this pipeline *derived* the target" — is unenforceable + for this checkpoint +- **Status:** reported; fix batched with F001/F003 + +## What happens + +`_validate_source_embargo` rejects an absolute path outright: + +> `'source_embargo'` entries are repo-relative; got absolute ... + +and requires each entry to exist under `trtllm_repo_path`. Both halves are +sound for the case the field was written for: TensorRT-LLM already ships an +implementation of this checkpoint, name it, and the writing roles cannot open +it. + +## Why it does not cover this run + +TensorRT-LLM ships **no** MiMo implementation — `grep -rl "MiMoV2\|mimo_v2"` +over `tensorrt_llm/`, `tests/` and `examples/` returns nothing. So there is +nothing in-repo to embargo, and `source_embargo: []` is correct. + +But two other engines on **this same filesystem** do implement it: + +``` +/lustre//portfolios//users//vllm +/lustre//portfolios//users//sglang +``` + +and vLLM's is the only implementation of the vision tower available anywhere +here. Those paths are an ordinary `Read` away from every agent in the run, +they are discoverable by a single `grep -rl mimo_v2 /lustre/...`, and a smith +looking for "how does anyone do asymmetric qk/v head dims on Blackwell" is +*exactly* the agent who will find them. + +`_common.py`'s `SOURCE_BOUNDARY` does say "Other inference engines' modeling +code stays out entirely". That sentence is a discipline in a prompt. For the +in-repo case the project decided a discipline was not enough and built a +mechanism — `source_embargo_block`'s own docstring says so: + +> "read to understand an interface, never to copy a structure" is a line +> nobody can audit from the outside: the target would be correct, and whether +> the pipeline could have produced it would be unknowable. + +That reasoning does not depend on which repository the tempting file is in. +It depends on the file being reachable. Here, the reachable ones are the ones +the field cannot name. + +## Why it bites harder than usual on this run + +The accuracy anchor for this checkpoint has to come from one of those two +engines — nothing else on this filesystem can load a 310B mxfp4 MoE, and the +checkpoint's own HF implementation cannot load its own quantized weights. So +the run will *deliberately start a vLLM or SGLang server*, which puts those +checkouts in the run's own working set. The black-box use is correct and +necessary; the read access that comes with it is incidental and unwanted. + +## Proposed change + +Widen the field rather than adding a second one: + +1. Accept an **absolute path** when it resolves outside `trtllm_repo_path`, + keeping the existing must-exist check — an embargo naming a path that is + not there reads as protection while protecting nothing, and that argument + is path-shape-independent. +2. Keep rejecting an absolute path that resolves *inside* the repo, with a + message saying to spell it repo-relative. That is the case the current + error is actually about, and it stays caught. +3. A directory entry embargoes the tree beneath it. Today's entries are + files; a sibling engine checkout is a tree, and asking a spec to enumerate + its model files is asking it to go stale. +4. `source_embargo_block` already renders the list into the prompt verbatim, + so nothing downstream changes. Its wording says "This checkpoint already + has an implementation in this repository", which becomes false for an + out-of-repo entry — split that sentence on where the paths point. + +## What this does not claim + +An embargo in a prompt is still a prompt. The project removed the per-read +`PreToolUse` hook that made it mechanical, and `README.md` records that +decision and what carries the duty instead. This finding is not asking for +the hook back; it is asking that the list a future hook would read be *able +to contain the right paths*. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F005-fix-report.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F005-fix-report.md new file mode 100644 index 000000000000..d5b2e96fd152 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F005-fix-report.md @@ -0,0 +1,80 @@ +# F005 — fix report + +Written by the orchestrating agent, not by the implementing one. The agent +that made these changes committed them and then stalled before writing its +own report; it was stopped, and this records what is actually on disk and +verified rather than what it intended. + +- **Finding:** `20260922-F005-inventory-probe-takes-the-multirank-launcher.md` +- **Landed in:** `70deac4726`, alongside the F006/F007 work +- **Suite at commit:** 1955 passed, 3 skipped + +## What changed + +| file | why | +|---|---| +| `survey.py` | `probe_command` no longer builds an engine-style launch; the inventory probe is one process | +| `environment.py` | gained the single-process command form, and `PATH` now composes rather than replaces | +| `preflight.py` | the survey summary distinguishes a failed probe from an empty catalog | +| `prompts/plan_drafter.py` | the drafter is told whether it is reading a measurement or a guess | +| `workflow.py` | carries the distinction through to the run's own record | +| `runtime_index.yaml` | the two new survey demands had to be registered, or the drift-guard test fails | +| tests | new coverage for each of the above | + +## Part 3 — which option was taken + +**Option (a): `PATH` is composed, not replaced.** The resolved prefix now ends +`...:/usr/bin:/bin:"$PATH"`, so a spec that pins its own interpreter ahead of +everything else no longer deletes the container's MPI, CUDA and system tool +directories as a side effect. + +This was the right half of the choice. Option (b) — rejecting a `PATH` that +does not end in `:$PATH` — would have taught every task author a rule in order +to prevent a mistake the code can simply not make. The failure it prevents is +worth recalling: the spec line looked entirely reasonable, and its consequence +was `mpirun: command not found` several layers away, inside a probe whose +output was then silently discarded. + +## Before and after, on the real task spec + +Before — against a catalog holding 76 entries, 61 with a passing `sm_103` +receipt: + +``` +✔ survey 12 structural demand(s), 14 unrecognized key(s); 0/0 entries certified on sm_103, 0 unwrapped op(s) +``` + +After, same command, same spec, run inside the compute session: + +``` +✔ survey 12 structural demand(s), 14 unrecognized key(s); 67/76 entries certified on sm_103, 248 unwrapped op(s) +``` + +The `67` rather than `61` is not a discrepancy: the earlier figure came from a +static read of the contracts, and this one is the live probe, which also +counts entries whose receipts were refreshed during this run. + +## Verification + +Both halves were confirmed by the orchestrator independently of the agent: + +- the probe payload, run single-process through the compute session: + `trtllm_ops: 277`, `catalog: 76` (61 passed / 2 failed / 13 absent), + `modules: 265`, `error: ""`, rc=0; +- the same payload under the old multi-rank launch: **rc=127**, no payload, + `mpirun: command not found`; +- the orchestrator's own `--dry-run`, which now prints the real counts. + +## What the finding did not anticipate + +**The runtime index is a hard prerequisite, not a follow-up.** The two new +survey demands had to be registered in `runtime_index.yaml` in the same +change, or an existing drift-guard test fails. A probe that reports a demand +the index does not know about is treated as drift, which is the correct +behaviour and was not obvious from the finding. + +**`attention_chunk_size` was deliberately left unclaimed.** It is still in the +unrecognized list. For this checkpoint it is a red herring — the reference +implementation never reads it, and the real windowing comes from +`sliding_window` — so widening the windowing probe to claim it would have +recorded a fact that is not true of this model. Flagged rather than absorbed. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F005-inventory-probe-takes-the-multirank-launcher.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F005-inventory-probe-takes-the-multirank-launcher.md new file mode 100644 index 000000000000..254803a0e89b --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F005-inventory-probe-takes-the-multirank-launcher.md @@ -0,0 +1,123 @@ +# F005 — the inventory probe takes the multi-rank launcher, and fails silently + +- **Found:** 2026-09-22 05:27 PDT, on the first real `--dry-run` of this run +- **Component:** `agent-flow/agent_flow/workflows/staircase_onboard/survey.py`, + `probe_command` / `run_inventory_probe` +- **Severity:** the plan drafter is handed an empty supply survey and drafts on + assumption, with no error anywhere +- **Status:** reproduced both ways; fix proposed + +## The symptom + +`--dry-run` printed, in the middle of an otherwise green preflight: + +``` +✔ survey 12 structural demand(s), 14 unrecognized key(s); 0/0 entries certified + on sm_103, 0 unwrapped op(s) +``` + +`0/0 entries certified` against a catalog that holds 76 entries, 61 of them +with a passing `sm_103` receipt. And it is reported with a **check mark**. + +## The cause + +`probe_command` builds its command through `RunEnvironment.command`, which +applies the same launch rules every gate gets. At world size > 1 that means: + +``` +mpirun -n 4 --allow-run-as-root trtllm-llmapi-launch -c '' +``` + +Those rules are right for engine work and wrong here. The probe is read-only +introspection — it imports `tensorrt_llm`, walks `torch.ops.trtllm`, and +parses catalog frontmatter off disk. It needs one process and no MPI world. + +Measured, same payload, same session, one variable: + +| launch | result | +|---|---| +| single process | `trtllm_ops: 277`, `catalog: 76` (61 passed / 2 failed / 13 absent), `modules: 265`, `error: ""`, rc=0 | +| `mpirun -n 4 ... trtllm-llmapi-launch` | **rc=127**, no payload | + +And the rc=127 is its own finding: + +``` +mpirun: command not found +``` + +This run's `environment.env` sets `PATH` explicitly, to pin the checkout's own +venv ahead of everything else: + +``` +PATH=/.venv-3.12/bin:/usr/local/bin:/usr/bin:/bin +``` + +The container's MPI lives at `/usr/local/mpi/bin/mpirun`, which that list does +not contain. So the launcher the probe did not need is also the launcher the +environment cannot find. + +## Why it is silent, and why that is the actual defect + +`run_inventory_probe`'s docstring states the policy plainly, and the policy is +right: + +> A failure here is reported, not raised. The plan can be drafted without an +> inventory — it just gets drafted the old way, on assumption — and losing the +> whole run because a probe could not import is a worse trade than losing the +> measurement. + +But the reporting half does not happen. `parse_inventory` returns an empty +inventory with an `error`, and the caller then applies: + +```python +landed = bool(inventory.catalog or inventory.trtllm_ops or inventory.modules) +if not landed and not getattr(result, "ok", True): + return Inventory(arch=arch, error=f"{inventory.error} (probe exited ...)") +``` + +so the error *is* carried. It is the **summary line** that loses it: an +inventory with 0 entries and an error prints as `0/0 entries certified` with a +`✔`, indistinguishable from a catalog that genuinely has nothing certified. +An operator reading that line has no way to tell "nothing is certified" from +"nothing was measured". + +For this run the consequence is concrete. The whole COVERED-vs-GAP decision — +which entries exist, which are certified on sm_103, which ops are unwrapped — +is what the smith queue is derived from. Drafted against an empty inventory, +every entry looks like a gap, and the run would dispatch smiths to build +things that already exist with passing receipts. + +## Proposed change + +Three parts, smallest first. + +1. **Do not launch the probe as engine work.** `probe_command` should build a + single-process command regardless of world size. The cleanest spelling is + for `RunEnvironment` to expose the distinction it already makes internally + — something like `command(argv, *, single_process=True)` — rather than + `survey.py` assembling a prefix by hand, because assembling launcher lines + outside `environment.py` is exactly what that module exists to prevent. + Nothing about this probe wants a rank per device. + +2. **Say when the measurement did not happen.** `format_surface` and the + preflight summary must distinguish an empty inventory from a failed one: + a `✘`/`⚠` and the error text, not a `✔`. The plan proceeding is fine — + that is the documented policy — but it has to proceed *loudly*, because + the drafter's prompt says whether it is reading a measurement or a guess. + +3. **Make `PATH` a merge, not a replacement.** A spec that pins its own venv + ahead of the system path should not have to enumerate the container's MPI, + CUDA and system tool directories to avoid deleting them. Either + `environment.env`'s `PATH` is prepended to the inherited one, or the schema + rejects a `PATH` that does not end in `:$PATH` with a message saying why. + The current behaviour turns a reasonable-looking spec line into a + `command not found` several layers away from where it was written. + +Part 3 is separable and is the one most likely to bite another run: every +task spec that pins an interpreter will write this same `PATH` line. + +## Interim action for this run + +The inventory is measured out of band and handed to the plan phase as a +pre-written artifact. The solo probe's payload is recorded in `04-catalog/` +so the drafter reads a measurement rather than an assumption. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F006-F007-fix-report.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F006-F007-fix-report.md new file mode 100644 index 000000000000..fe1d4baf4878 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F006-F007-fix-report.md @@ -0,0 +1,173 @@ +# F006 + F007 — fix report + +Both findings moved into the package with tests, and verified end to end on +the live compute session (Slurm job 563608, node nvl72d209-T03, 4x GB300). + +- **Spec:** `20260922-F006-two-multirank-launch-defects.md`, + `20260922-F007-slurm-vars-defeat-the-mpirun-launch.md` +- **Suite:** `1955 passed, 3 skipped` (baseline was 1929/3; the delta is this + work's 15 tests plus the concurrent F005 work already in the tree) +- **Not run:** `pre-commit`, per instruction. Nothing committed. + +--- + +## Files changed + +| file | why | +|---|---| +| `agent-flow/agent_flow/workflows/staircase_onboard/environment.py` | The two `mpirun` fixes, factored into one `_mpirun_launch` so the emitted command and the agent prompt cannot diverge; new `MULTI_TASK_SLURM_VARS` / `MULTI_TASK_SLURM_SCRUB` named as the other half of `SINGLE_TASK_MPI_SCRUB`. | +| `agent-flow/agent_flow/workflows/staircase_onboard/preflight.py` | `check_tmpdir` gains the socket-length bound, derived from `UNIX_SOCKET_PATH_MAX` and `RPC_SOCKET_BASENAME_LEN` rather than hardcoded. | +| `agent-flow/agent_flow/workflows/staircase_onboard/TROUBLESHOOTING.md` | Three entries keyed on the verbatim text: the slot message, the four-leaders hang, and the 107-character ZeroMQ error. Plus the UCX red herring, explicitly marked as not-caught-on-purpose. | +| `agent-flow/tests/workflows/staircase_onboard/test_environment.py` | 8 tests: both flags present on `mpirun`, **both absent on `slurm`**, absent at world size 1, absent for `single_process`, and the prompt block quoting the same line. | +| `agent-flow/tests/workflows/staircase_onboard/test_preflight.py` | 5 tests: the bound is derived not literal, the real 92-character value is rejected with the arithmetic in the message, the 61/62 boundary, an actual `AF_UNIX` bind agreeing with the arithmetic, and `CLAUDE_CODE_TMPDIR` exempt. | +| `mimo-v2.6-flash-onboard/task.yaml` | `TMPDIR`/`CLAUDE_CODE_TMPDIR` moved to the 49-character Lustre spelling. See "what the findings did not anticipate" — the file still declared the 92-character one. | + +Untouched, and deliberately: the `kind: slurm` branch, which keeps today's +behaviour exactly. + +--- + +## The generated command, before and after + +Regenerated from the real `mimo-v2.6-flash-onboard/task.yaml` through +`preflight._substrate_boot` → `gates.run_boot_gate` → +`RunEnvironment.python_command`, for the declared proxy checkpoint. + +**Before** (reconstructed by removing exactly what this work added): + +``` +CLAUDE_CODE_TMPDIR=/lustre//portfolios//users//tmp4 PATH=/scratch//portfolios//projects//users//TRTLLM4/.venv-3.12/bin:/usr/local/bin:/usr/bin:/bin:"$PATH" PYTHONNOUSERSITE=1 TMPDIR=/lustre//portfolios//users//tmp4 TRTLLM_MODELING_V2=off mpirun -n 4 --allow-run-as-root trtllm-llmapi-launch /scratch//portfolios//projects//users//TRTLLM4/.venv-3.12/bin/python /scratch//portfolios//projects//users//TRTLLM4/examples/llm-api/quickstart_advanced.py --model_dir="/lustre//portfolios//users//models/Qwen3.8-27B-NVFP4" --max_tokens=4 --tp_size=4 --use_cuda_graph --enable_chunked_prefill --max_seq_len=8192 --max_batch_size=8 --kv_cache_fraction=0.6 --moe_backend TRTLLM --prompt "The capital of France is" +``` + +**After:** + +``` +CLAUDE_CODE_TMPDIR=/lustre//portfolios//users//tmp4 PATH=/scratch//portfolios//projects//users//TRTLLM4/.venv-3.12/bin:/usr/local/bin:/usr/bin:/bin:"$PATH" PYTHONNOUSERSITE=1 TMPDIR=/lustre//portfolios//users//tmp4 TRTLLM_MODELING_V2=off env -u SLURM_PROCID -u SLURM_NTASKS -u SLURM_LOCALID -u SLURM_NODEID -u SLURM_STEP_ID -u SLURM_TASK_PID -u SLURM_STEP_NUM_TASKS mpirun -n 4 --oversubscribe --allow-run-as-root trtllm-llmapi-launch /scratch//portfolios//projects//users//TRTLLM4/.venv-3.12/bin/python /scratch//portfolios//projects//users//TRTLLM4/examples/llm-api/quickstart_advanced.py --model_dir="/lustre//portfolios//users//models/Qwen3.8-27B-NVFP4" --max_tokens=4 --tp_size=4 --use_cuda_graph --enable_chunked_prefill --max_seq_len=8192 --max_batch_size=8 --kv_cache_fraction=0.6 --moe_backend TRTLLM --prompt "The capital of France is" +``` + +The generated line is now character-identical to the hand-built +`spool/substrate_probe.sh` apart from the `PATH` tail: the hand-built one +spliced `/usr/local/mpi/bin` in by hand, where the generated one carries +`:"$PATH"` — the concurrent F005 merge fix, which subsumes it. + +The `--prompt "..."` argument is why this is run from a file rather than +interpolated: `spool_submit.sh` writes its payload into an **unquoted** +heredoc, so `$PATH` would be expanded on the login node at submit time rather +than in the container. The command was written verbatim to +`spool/generated_substrate.sh` and submitted as `bash `. + +--- + +## Spool-verified substrate gate + +`gate-f007-generated-substrate`, `rc=0`, roughly 5 minutes. From +`spool/logs/gate-f007-generated-substrate.log`: + +``` +mpi_rank: 0 +mpi_rank: 3 +Rank3 run mgmn worker node with mpi_world_size: 4 ... +mpi_rank: 2 +Rank2 run mgmn worker node with mpi_world_size: 4 ... +mpi_rank: 1 +Rank1 run mgmn worker node with mpi_world_size: 4 ... +Rank0 run mgmn leader node with mpi_world_size: 4 ... +... +[0] Accept rate: 1.00 (tokens=4, iterations=4) +[0] Prompt: 'The capital of France is'..., Generated text: ' Paris.\nThe'... +Rank0 Task exit code: 0 +... +Rank0 MGMN leader node exit code: 0 +Rank1/2/3 MGMN worker node exit code: 0 +``` + +One leader, three workers, world size 4 on every rank — the F007 failure shape +(four `mgmn_leader_node`, log stops growing) is absent. The `UCX ERROR +uct_iface_open(tcp/rdma_vf_rail0) failed` lines are present, as expected, and +the run exits 0 anyway. + +Spec validation and `--dry-run` also pass, with the new `TMPDIR`: + +``` +.venv-orch/bin/staircase-onboard --task mimo-v2.6-flash-onboard/task.yaml \ + --workspace mimo-v2.6-flash-onboard/workspace --dry-run +→ --dry-run: preflight passed; stopping before the first agent turn. (exit 0) +``` + +--- + +## The socket-length limit, measured + +Bound an `AF_UNIX` socket against synthetic paths of each length, no GPU +needed: + +| path length | result | +|---|---| +| 100–107 | binds | +| 108–114 | `OSError: AF_UNIX path too long` | + +So the kernel cap is **107 inclusive**, matching ZeroMQ's own +`is longer than 107 characters`. And +`len("rpc_test_") + len(str(uuid.uuid4()))` is `9 + 36 == 45`. + +`get_unique_ipc_addr` in `tensorrt_llm/executor/rpc/rpc_common.py` builds the +path with `os.path.join(tempfile.gettempdir(), file_name)`, which inserts a +separator. So: + +``` +MAX_TMPDIR_LEN = 107 - 45 - len("/") = 61 +``` + +This is asserted in the suite against a live bind at exactly 61 and exactly +62, so the derivation is checked against the kernel rather than against +itself. + +--- + +## What the findings did not anticipate + +**1. The bound is 61, not 62.** F006 computed `107 - 45` and got 62. The +separator `os.path.join` inserts is the third term, and dropping it admits +exactly one `TMPDIR` — a 62-character one — that does not work. The +implementation uses the measured 61 and the message says 61; the constant is +derived as `UNIX_SOCKET_PATH_MAX - RPC_SOCKET_BASENAME_LEN - len("/")` so the +subtraction is visible rather than folded into a literal. + +**2. `task.yaml` still declared the 92-character `TMPDIR`.** The brief said it +already carried the 49-character Lustre spelling; it did not — both `TMPDIR` +and `CLAUDE_CODE_TMPDIR` were still +`/scratch/.../users//staircase-tmp`, which the new check correctly +rejects. Moved to `/lustre//portfolios//users//tmp4` (49 +characters, `rpc_test_` under it is 95 and binds) with a comment giving +the arithmetic, so the next author does not shorten it back by accident. + +**3. `RunEnvironment.command` was not the only place emitting the broken +line.** `_describe_spool` — the prompt block every agent reads under +`session_spool` — independently emitted +`mpirun -n N --allow-run-as-root trtllm-llmapi-launch `. Fixing +only `command` would have left every agent hand-reproducing both defects, and +the F007 half presents as an engine that never boots. Both now render through +the same `_mpirun_launch`, and a test asserts the prompt quotes the line the +orchestrator will run. The prompt also now says **"Copy that line whole"** with +one sentence on what each flag is for, because an agent that tidies away an +`env -u SLURM_...` prefix it does not recognise gets a silent hang. + +**4. The length bound is on `TMPDIR` only.** `check_tmpdir` loops over +`TMPDIR` and `CLAUDE_CODE_TMPDIR`, but only the former hosts the RPC socket. +Holding the harness's scratch dir to a socket's budget would reject a value +that works, for a reason that is not true of it. Tested both ways. + +**5. `--oversubscribe` and the scrub share a home, not just a branch.** They +are both properties of *"running `mpirun` inside someone else's allocation"* — +which is the one thing the non-Slurm branch does and the `slurm` branch never +does. They live together in `_mpirun_launch`, whose docstring states the +asymmetry once for both rather than twice. + +**6. Worth reporting upstream (restating F007).** `trtllm-llmapi-launch` +preferring `SLURM_PROCID` unconditionally produces the worst available failure +mode when the assumption breaks: not an error and not a wrong answer, but N +silent single-rank runs that hang forever. A cheap guard would be to compare +the two when both are present and refuse to start on disagreement — +`SLURM_PROCID=0` alongside `OMPI_COMM_WORLD_RANK=3` is not a situation any +correct launch produces, and neither is `SLURM_NTASKS=1` alongside +`OMPI_COMM_WORLD_SIZE=4`. Not filed as part of this change. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F006-two-multirank-launch-defects.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F006-two-multirank-launch-defects.md new file mode 100644 index 000000000000..80611e4ba08f --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F006-two-multirank-launch-defects.md @@ -0,0 +1,118 @@ +# F006 — two multi-rank launch defects preflight cannot see + +- **Found:** 2026-09-22 05:34-05:37 PDT, running the substrate gate by hand + through the compute session before committing the run +- **Component:** `agent-flow/agent_flow/workflows/staircase_onboard/environment.py` + (`command`, `_local_prefix`) and `preflight.py` (`check_tmpdir`) +- **Severity:** every `session_spool` / `local` run at world size > 1 fails, + twice, with errors that name neither the spec nor the workflow +- **Status:** both reproduced and both fixed by hand for this run; fixes proposed + +Both of these are exactly the class `TROUBLESHOOTING.md` exists for: the error +text mentions Open MPI and ZeroMQ, and the cause is a line in `task.yaml`. + +--- + +## Defect 1 — `mpirun` gets one slot, because the session allocation has one task + +`RunEnvironment.command` wraps a world-size-4 payload as: + +``` +mpirun -n 4 --allow-run-as-root trtllm-llmapi-launch ... +``` + +Correct in intent — the ranks must exist before the engine starts, or the LLM +API concludes it has to spawn them and dies with `MPI_ERR_SPAWN`. But inside a +Slurm allocation Open MPI takes its slot count from the resource manager, and +a spool session is submitted as **one task** that owns the node: + +``` +NumNodes=1 NumCPUs=144 NumTasks=1 ... TresPerNode=gres/gpu:4 +``` + +So `-n 4` asks for four slots against one, and the run dies before any Python: + +``` +There are not enough slots available in the system to satisfy the 4 +slots that were requested by the application: + trtllm-llmapi-launch +``` + +**Why the obvious fix is wrong.** Submitting the session with +`ntasks_per_node=4` gives MPI its four slots and simultaneously starts **four +copies of the spool watcher**, which then race for every piece of work. The +session is one process by design. + +**Fix: `--oversubscribe`.** The allocation owns the whole node — 144 CPUs and +all 4 GPUs — so the slot count is a bookkeeping guard, not a real constraint. +Verified: with `--oversubscribe` the four ranks start and bind +(`rank 0 using MpiCommSession to bind to external MPI processes`, four times). + +It belongs in `environment.py`'s non-Slurm branch, next to the existing +`mpirun` assembly, with a comment saying why — the `kind: slurm` branch must +**not** get it, because there `srun --ntasks` really does size the world and +oversubscribing would hide a mis-sized allocation. + +--- + +## Defect 2 — a long `TMPDIR` breaks the LLM API's RPC, at 107 characters + +With the slots fixed, all four ranks started and then failed: + +``` +zmq.error.ZMQError: ipc path +"/scratch//portfolios//projects//users//staircase-tmp/rpc_test_9bbff7c6-e827-43d1-9287-d9627067e179" +is longer than 107 characters (sizeof(sockaddr_un.sun_path)). +``` + +and then, from the same path, `zmq.error.ZMQError: File name too long`. + +107 is `sizeof(sockaddr_un.sun_path)` — a kernel limit on unix domain socket +paths, not a ZeroMQ or TensorRT-LLM choice. The LLM API's RPC binds an IPC +socket under `TMPDIR` named `rpc_test_`, which is a fixed 45-character +suffix. So **any `TMPDIR` longer than 62 characters fails**, and this run's +was 92. + +`check_tmpdir` asserts `TMPDIR` is set and not under `/tmp`, for a good and +well-documented reason. It says nothing about length, and length is the +property that breaks first on a cluster whose storage paths are deep — +`/scratch//portfolios//projects//users//` +was 78 characters on that cluster before the run added anything. + +**Fix: assert the length in `check_tmpdir`,** with the arithmetic in the +message: + +> `TMPDIR=` is 92 characters. The LLM API binds an RPC socket at +> `$TMPDIR/rpc_test_`, and a unix socket path is capped at 107 +> (`sizeof(sockaddr_un.sun_path)`), so `TMPDIR` must be at most 62. Point it +> at a shorter path on the same storage. + +Seconds in preflight instead of a four-rank engine load and a ZeroMQ traceback. + +**What this run did.** `/lustre//portfolios//users//tmp4`, +49 characters — the *same storage* as the 92-character `/scratch` spelling +(established earlier in this run), just a shorter path to it. Measured in the +session: a `rpc_test_` path under it is 95 characters and binds. + +**A related fact worth recording, because it contradicts an earlier note.** +Unix sockets bind fine on this Lustre filesystem — tested directly on +`/tmp`, `/dev/shm` and both Lustre spellings, all four succeed. The earlier +guidance that a Lustre `TMPDIR` breaks torch's shared-memory manager is +therefore either about a different mount or about a different limit; on this +cluster the failure is **path length**, not filesystem type. Anyone copying +the `TMPDIR: /tmp` workaround from the multimodal options file should know +that it works because `/tmp` is *short*, not because it is local. + +--- + +## Why neither was caught earlier + +`--dry-run` skips the substrate gate's actual execution under `session_spool` +— the devices are the scheduler's, preflight claims none, and the gate is the +first thing that would have run four ranks. Both defects live in the gap +between "the spec validates" and "the first real command runs". + +That gap is what the substrate gate is for, and it would have caught both on +the real run. The cost would have been one failed preflight per defect rather +than two hand-run probes — which is the right outcome and an argument for +running the substrate gate rather than reaching for `--skip-substrate-gate`. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F007-slurm-vars-defeat-the-mpirun-launch.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F007-slurm-vars-defeat-the-mpirun-launch.md new file mode 100644 index 000000000000..90264812e603 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F007-slurm-vars-defeat-the-mpirun-launch.md @@ -0,0 +1,130 @@ +# F007 — Slurm's env defeats `trtllm-llmapi-launch` under `mpirun`, silently + +- **Found:** 2026-09-22 05:55-06:05 PDT, running the substrate gate by hand +- **Component:** `agent-flow/.../staircase_onboard/environment.py`, `RunEnvironment.command` + — and the underlying cause is in TensorRT-LLM's own `trtllm-llmapi-launch` +- **Severity:** every non-`slurm` launch at world size > 1 inside a Slurm + allocation hangs forever, having silently become N independent single-rank runs +- **Status:** root-caused, proven both ways, fixed by hand for this run; + workflow fix proposed and an upstream note recommended + +## The symptom + +The substrate gate started, printed its banners, and then stopped. The log did +not grow for ten minutes. No error, no traceback, no exit. + +`AGENTS.md` lists this exact shape: + +> the engine hangs at startup and never returns — a wrong launch substrate +> hangs rather than raises + +## What was actually happening + +A process listing on the node: + +``` +178503 python3 -m tensorrt_llm.llmapi.mgmn_leader_node +178507 python3 -m tensorrt_llm.llmapi.mgmn_leader_node +178510 python3 -m tensorrt_llm.llmapi.mgmn_leader_node +178511 python3 -m tensorrt_llm.llmapi.mgmn_leader_node +``` + +**Four leaders.** A correct 4-rank launch has one leader and three workers. +GPU 0 held 3 GB and GPUs 1-3 were empty, so nothing was computing — four +processes were each waiting to be joined by workers that would never come. + +## Root cause + +`trtllm-llmapi-launch` derives its rank and world size like this: + +```sh +mpi_rank=${SLURM_PROCID:-${OMPI_COMM_WORLD_RANK:-${PMI_RANK:-${PMI_ID:-0}}}} +``` + +and, in `mpi_world_size`, checks `SLURM_NTASKS` before `OMPI_COMM_WORLD_SIZE`. + +**Slurm wins over Open MPI in both.** That is the right precedence under +`srun`, which is what the launcher was written for. It is exactly wrong under +`mpirun` *inside* a Slurm allocation — and a `session_spool` job is precisely +that: one Slurm task that owns the node, running `mpirun` itself. + +Measured in the session, four ranks, one variable changed: + +| | `SLURM_PROCID` | `OMPI_COMM_WORLD_RANK` | `SLURM_NTASKS` | `OMPI_COMM_WORLD_SIZE` | +|---|---|---|---|---| +| as-is | `0, 0, 0, 0` | `0, 1, 2, 3` | `1` | `4` | +| SLURM vars scrubbed | unset | `0, 1, 2, 3` | unset | `4` | + +So every rank read `mpi_rank=0` and became a leader, and each also read +`mpi_world_size=1` — the launcher was wrong about *both* numbers, from the +same cause. + +MPI itself was never at fault. A bare `mpi4py` probe under the identical +`mpirun` line reported `rank 0..3 of 4` correctly. + +## The fix + +Scrub the Slurm task variables immediately before `mpirun`, so the launcher +falls through to the Open MPI ones it should be reading: + +``` +env -u SLURM_PROCID -u SLURM_NTASKS -u SLURM_LOCALID -u SLURM_NODEID \ + -u SLURM_STEP_ID -u SLURM_TASK_PID -u SLURM_STEP_NUM_TASKS \ + mpirun -n 4 --oversubscribe --allow-run-as-root trtllm-llmapi-launch ... +``` + +Verified end to end on the substrate gate, stock path, 4x GB300: + +``` +mpi_rank: 0 / 1 / 2 / 3 +Rank0 run mgmn leader node with mpi_world_size: 4 +Rank1/2/3 run mgmn worker node with mpi_world_size: 4 +[0] Prompt: 'The capital of France is'..., Generated text: ' Paris.\nThe'... +Rank0 Task exit code: 0 +``` + +## Where it belongs in the workflow + +`environment.py` already owns the inverse of this rule. `SINGLE_TASK_MPI_SCRUB` +exists because, at world size 1, "mpi4py auto-initializes MPI whenever PMI +variables are present and inside a PMIx step they always are". That is the +same insight — a Slurm step's environment misleads a process that is not being +launched by `srun` — applied to the other end of the range. + +So the change is to extend the existing rule rather than invent one: + +- the **`slurm`** branch keeps today's behaviour exactly. There `srun --ntasks` + really does size the world, `SLURM_PROCID` really is the rank, and scrubbing + would break a correct launch. +- the **`mpirun`** branch (`local`, `session_spool`) gains the scrub, for the + same reason the world-size-1 branch has one. + +Together with F006's `--oversubscribe`, the non-Slurm multi-rank line becomes +complete. Both are properties of "running `mpirun` inside someone else's +allocation", and neither is a property of this model, this target or this run. + +## Worth reporting upstream + +`trtllm-llmapi-launch` preferring `SLURM_PROCID` unconditionally is defensible +but produces the worst available failure mode when the assumption breaks: not +an error, not a wrong answer, but N silent single-rank runs that hang forever. +A cheap guard would be to compare the two when both are present and refuse to +start when they disagree — `SLURM_PROCID=0, OMPI_COMM_WORLD_RANK=3` is not a +situation any correct launch produces. + +The same applies to `SLURM_NTASKS=1` alongside `OMPI_COMM_WORLD_SIZE=4`. + +## Note on the UCX errors, which are NOT this + +Throughout, every rank logged: + +``` +UCX ERROR uct_iface_open(tcp/rdma_vf_rail0) failed: Input/output error +pml_ucx.c:314 Error: Failed to create UCP worker +``` + +These are noise here and were a red herring for a while. UCX fails to open one +IPv6 RDMA virtual-function interface and falls back; the run completes with +exit code 0 and correct output. `UCX_TLS=sm,self` silences them and is worth +setting to keep the logs readable, but it fixes nothing — confirmed by the +passing run above, which still logs them. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F008-plan-reviewer-cross-model-not-available.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F008-plan-reviewer-cross-model-not-available.md new file mode 100644 index 000000000000..e8b8dce37e7e --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F008-plan-reviewer-cross-model-not-available.md @@ -0,0 +1,88 @@ +# F008 — the PlanReviewer's cross-model property could not be honoured + +- **Found:** 2026-09-22 07:20 PDT, dispatching the plan review +- **Component:** the `plan_reviewer` role, and the external transport generally +- **Severity:** a designed safeguard silently degraded to a weaker one +- **Status:** worked around and disclosed; no code change proposed yet + +## What the design says + +`README.md`, "The six roles": + +> | `plan_reviewer` | persistent | **Codex** | a paper review of both | +> +> The PlanReviewer runs on a different model from the PlanDrafter on purpose: +> a reviewer that shares the author's blind spots reviews its own reasoning +> back at itself. + +That is a real safeguard, and it is the only one the plan phase has — there is +no gate on a plan, only a review. + +## What happened + +Under `--agent-transport external` the orchestrator does not construct the +agent at all; it writes the turn to disk and a driver session answers it. The +role's *backend* therefore does not travel with the turn. The +`request.json` carries `role`, `what`, the prompt paths and the decision file +— it does not carry "run this one on a different model", because in-process +that is a property of the layer the orchestrator built, not of the request. + +So the driver has to know, out of band, that this particular role wants a +different model. Nothing on disk says so. + +Attempting it then failed for a second, independent reason: the driver's +available models did not include the one requested — +`model_not_found, HTTP 404` for `claude-fable-5-1`. Codex was not reachable +from this session either. + +The review was therefore run on the **same model family as the drafter**, with +the divergence compensated for in the prompt — the reviewer was told +explicitly that it shares the author's blind spots, and instructed to verify +where it found itself agreeing easily rather than nodding. + +That is a mitigation, not the property. It should not be recorded as if the +cross-model review happened. + +## Why this is worth fixing rather than noting + +The failure is silent by construction. An external driver that simply runs +every turn on whatever model it has produces a `plan_reviewer-1.yaml` that +validates, a `## Open risks` section that looks right, and a run that reports +a reviewed plan. Nothing downstream can tell the difference. The same is true +of any future role that depends on *which* model answers rather than on what +the answer contains. + +Compare how carefully the transport already handles the other half of this +problem: `staircase_transport_note` exists because a role prompt written for +the in-process path ends with an instruction the external path cannot honour, +and the project decided the transport must **state** the difference rather +than let three auditor turns in a row quietly work it out. This is the same +class of problem — an in-process property that does not survive the hop — and +it currently has no such note. + +## Proposed change + +1. **The request carries the role's intended backend.** Add a `backend` field + to the dispatch `request.json`, populated from the same role configuration + the in-process transport uses. A driver then knows, from the file, what the + turn wants. Cheap, and it makes the requirement discoverable rather than + folklore. +2. **The result records what actually answered.** Either a field the driver + fills in, or — better, since a driver is trusted with nothing — the + orchestrator notes in `progress.yaml` that this turn went out over a + transport that cannot guarantee the backend. The run's own record should + not claim a cross-model review it cannot verify. +3. **`staircase_transport_note` gains a sentence for backend-sensitive roles**, + in the same voice as the existing one: this transport delivers your prompt, + it does not choose your model, and if you are the reviewer you should + assume you may share the drafter's blind spots. + +Point 2 is the one that matters. Points 1 and 3 help a driver do the right +thing; point 2 stops the run from *claiming* the right thing happened when +nobody checked. + +## Recorded for this run + +The plan review for MiMo-V2.6-Flash-RL was **not** cross-model. Any conclusion +that rests on "two independent models agreed" does not hold for this run's +plan phase. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F009-a-turn-can-commit-and-never-hand-back.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F009-a-turn-can-commit-and-never-hand-back.md new file mode 100644 index 000000000000..552d280ad303 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F009-a-turn-can-commit-and-never-hand-back.md @@ -0,0 +1,85 @@ +# F009 — a role turn can do its work, commit it, and never hand back + +- **Found:** 2026-09-22 09:25 PDT, at the end of smith batch 1 +- **Component:** the external transport's hand-back protocol, and the + orchestrator's view of a turn's liveness +- **Severity:** one stalled turn blocks a whole batch for its full timeout, + with completed work sitting in `HEAD` the orchestrator cannot see +- **Status:** recovered by hand; fix proposed + +## What happened + +Dispatch `006-smith` (verify `moe/mxe4m3_mxe2m1_block_scale_moe_runner`) did +its work and **committed it** — `a1da9c19c4`, 380 lines across the contract +and a new 281-line test, receipt frontmatter updated to +`sm_103: {status: passed, tests: 21}`. + +It then never wrote `decisions/smith-5.yaml`, never wrote +`006-smith.result.md`, and never wrote `006-smith.done`. + +Its eight siblings all handed back cleanly. The orchestrator sat blocked on +one missing zero-byte file for **58 minutes**, and would have sat there for +the full `--agent-turn-timeout` (4 hours as configured) before failing the +batch — while the work it was waiting for was already in `HEAD`. + +## Why the protocol allows this + +The transport's contract is deliberately minimal and correct as far as it +goes: `.done` is written last, so a reader that sees it can trust the result +beside it. The ordering rule is right. + +What is missing is the other direction. **There is no signal for "this turn is +no longer running."** The orchestrator distinguishes exactly two states — the +`.done` file exists, or it does not — and an abandoned turn is indistinguishable +from a slow one. `FileDispatchTurnExecutor` has a `driver.heartbeat` for the +*driver* being gone, which is the right idea applied one level too high: it +detects "nothing is reading the mailbox at all", not "this particular turn +died". + +That distinction matters more here than it would elsewhere, because these +turns are long. A smith that takes 30 minutes is normal; the batch that just +completed ran 20-58 minutes per dispatch. So a generous timeout is correct, +and a generous timeout is exactly what makes a stall expensive. + +## The thing that makes it worse + +The abandoned turn had **already committed**. So the run's git history says +the work is done, the catalog's receipt says the entry is certified, and the +orchestrator's `progress.yaml` says the dispatch never produced a decision. +Those two records disagree, and only one of them is reachable by the +orchestrator. + +A resume would re-dispatch the turn, and a fresh smith would find its own work +already in `HEAD` — which is recoverable, but only because the smith skill +happens to handle "the receipt is already valid" gracefully. Nothing +guarantees that for other roles. + +## Proposed change + +Three parts, in increasing order of effort. + +1. **A per-turn heartbeat.** The driver touches + `dispatch/.heartbeat` while a turn is in flight; the orchestrator fails + that turn — not the batch — when it goes stale, with a message naming the + request id and the prompt path. This mirrors `driver.heartbeat` exactly, + one level down, and reuses its staleness logic. +2. **A stalled turn does not have to fail the batch.** With per-turn + liveness, the orchestrator can re-dispatch that one request and let its + siblings' results stand. Today a single stall takes the batch with it. +3. **Reconcile against the tree before re-dispatching.** For smith turns + specifically, "did this entry's receipt become valid while nobody was + looking" is a cheap, already-implemented check — `receipt_state` in + `survey.py` does it. A re-dispatch that starts by reading the tree rather + than the mailbox would recover this case automatically instead of relying + on the agent noticing. + +Part 1 is small and is most of the value: it turns a four-hour silent block +into a named failure in five minutes. + +## What was done for this run + +The turn was completed by a second agent, briefed to **verify what the first +one left on disk rather than redo it**, and required to state in its result +file that the turn was finished by a different agent — so the run's record +does not imply one continuous turn. The distinction between "the first agent +did this, verified from disk" and "I did this" is preserved in the hand-back. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F010-task-yaml-is-read-once-and-a-gate-error-is-truncated.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F010-task-yaml-is-read-once-and-a-gate-error-is-truncated.md new file mode 100644 index 000000000000..05c0a4c2be08 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260922-F010-task-yaml-is-read-once-and-a-gate-error-is-truncated.md @@ -0,0 +1,115 @@ +# F010 — the spec is read once, and the gate cannot see why it failed + +- **Found:** 2026-09-22 10:20-10:30 PDT, at the core unit's first boot gate +- **Components:** the orchestrator's task-data lifetime; `gates.py`'s + `stdout_tail`; `preflight._reports_unsupported_architecture` +- **Severity:** cost one wasted assembler turn and two orchestrator restarts +- **Status:** worked around for this run; three fixes proposed + +Two separate defects compounded into one confusing hour. Both are worth +fixing independently. + +--- + +## Defect 1 — `task.yaml` is loaded once, so an agent cannot repair the spec + +The core assembler discovered a genuine precondition: this checkpoint's +`auto_map` points at `configuration_mimo_v2.MiMoV2Config`, `transformers` +knows no `mimo_v2` model type, and `ModelConfig.from_pretrained` therefore +refuses to load it at all without `trust_remote_code`. It did the right thing +— added `--trust_remote_code` to `workspace/task.yaml`'s +`runtime.default.extra_flags` and `trust_remote_code: true` to the core +options file, and said so in its hand-back. + +**The running orchestrator never saw it.** `task.yaml` is validated and loaded +into `self.task_data` when the workflow object is constructed; every gate +afterwards reads that in-memory copy. So an agent editing the file on disk is +editing something nothing will re-read, and the boot gate ran again with the +identical, still-broken command. + +The workflow then did exactly what it should on a failed boot — `continue`, and +dispatch another assembler turn. That turn would have been asked to fix a +target that was already correct, against a gate that could not pass whatever +it did. The loop is bounded at `--max-assembler-turns-per-unit` (6), so the +unit would have burned six turns and ended `assembly-exhausted:core` with a +green target on disk. + +**Why the agent's instinct was right and the system disagreed.** The +per-unit options files (`llm_api_options`) *are* re-read per gate — they are +paths, resolved at gate time. So the assembler correctly fixed one half in a +place that works, and the other half in a place that does not, with no signal +that the two behave differently. + +### Proposed change + +1. **Re-read `task.yaml` at each stage boundary**, or at minimum at each gate, + and log when the on-disk copy differs from the loaded one. The spec is + small and validation is fast; the cost is nothing against a wasted unit. +2. If that is judged too loose — a spec mutating mid-run is its own hazard — + then the opposite, and say it: **make `workspace/task.yaml` read-only to + agents**, and give the assembler a typed way to request a spec change + (a `BLOCKED: spec_change` decision carrying the key and value), so the + orchestrator applies it deliberately and records it. + +Either is fine. What is not fine is the current state, where the edit is +allowed, looks successful, and silently does nothing. + +--- + +## Defect 2 — the boot gate's failure detail is truncated to teardown noise + +With the spec fixed and the orchestrator restarted, the substrate gate failed +again, and this time the reported cause was: + +``` +RuntimeError: Executor worker returned error +``` + +That is `GenerationExecutorProxy._fail_initialization` re-raising. The +worker's actual error — the thing that decides whether this is a host +problem, a spec problem, or an unsupported architecture — was raised in +another process and logged well before the end. + +`GateResult.stdout_tail` keeps the **last 40 lines**. Under a 4-rank MPI +launch those 40 lines are reliably: UCX interface warnings, four ranks' +teardown messages, the proxy's generic re-raise, and Open MPI's +"Primary job terminated normally, but 1 process returned a non-zero exit +code" banner. The real exception is further up, and never reaches the +orchestrator. + +**The direct consequence:** F001's fallback could not fire. +`_reports_unsupported_architecture` searches `result.detail` and +`result.stdout_tail` for the architecture markers. The marker was in the log +and not in the tail, so the run reported "this host cannot generate a token" +about a host that generates tokens fine — the exact misdiagnosis the substrate +gate exists to prevent, produced by the substrate gate. + +### Proposed change + +1. **Search the whole captured stdout for the discriminator, not the tail.** + `stdout_tail` is a display budget; the classification should not inherit + it. One-line change with a real payoff. +2. **Keep a first-exception excerpt alongside the tail.** Scan for the first + `Traceback (most recent call last):` and keep that block. On a multi-rank + failure it is far more diagnostic than the last 40 lines, and it is what a + human scrolls up to find anyway. +3. **Teach `TROUBLESHOOTING.md` the shape**, keyed on the verbatim string + `RuntimeError: Executor worker returned error` — "this is the proxy + re-raising a worker's error; the cause is upstream in the log, search for + the first traceback." + +--- + +## What this run did + +`--skip-substrate-gate`, with the substrate proven twice over and the evidence +on disk: + +- by hand on the proxy checkpoint through the compute session: one leader, + three workers, `mpi_world_size: 4`, generated `' Paris.'`, exit 0; +- by the core assembler on the **real** target, all three members, each + answering `48 + 24 = 72 clips in April and May` and continuing `Paris`. + +That is precisely the condition the flag's own help sanctions — "you have +already proven the substrate in this same environment" — rather than the +condition it warns against. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260923-F011-the-gate-measures-the-working-tree.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260923-F011-the-gate-measures-the-working-tree.md new file mode 100644 index 000000000000..d7d513335f6f --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/20260923-F011-the-gate-measures-the-working-tree.md @@ -0,0 +1,90 @@ +# F011 — the gates measure the working tree, and nothing owns it + +- **Found:** 2026-09-23 02:05 PDT, by the multimodal reviewer, minutes before + three MMMU accuracy gates would have run +- **Component:** the gate commands generally; `workflow.py`'s unit loop +- **Severity:** a gate can measure code nobody reviewed, and report the number + as the reviewed unit's +- **Status:** caught and worked around by hand; fix proposed + +## What happened + +The multimodal unit was assembled (commit `36ea38b940`), reviewed, and +APPROVEd. The orchestrator was about to run its accuracy gate on three members +— roughly 15 minutes each plus engine start. + +Meanwhile a **separate agent**, implementing the audio input path, was editing +the same `modeling.py`. It had gone from the committed 2827 lines to 3199. + +Nothing would have stopped the gate. `run_accuracy_gate` builds a command +naming `checkpoint_path` and an options file; the *code* it measures is +whatever is on disk under `trtllm_repo_path` when the engine imports it. There +is no pin, no check, and no record of which tree state produced a number. + +The reviewer caught it only because it happened to re-read the file's line +count and noticed it had changed underneath its own review. + +## Why this is worse than a flaky run + +Three separate failures were queued up: + +1. **The gate would have measured a half-written vision path.** The in-flight + edits renamed `_VIT_CALL_CONSTANTS` / `_VIT_MASK_PADDING`, removed the + vision full-attention call site's arguments, and added an unconditional + `assert isinstance(audio_cfg, dict)` to `ModelingV2Core.__init__` — which + would break the *three already-green units* to add a fourth. +2. **A failure would have burned assembler turns.** The unit budget is six; + three failed members is three turns spent on a defect that does not exist + in the reviewed commit. +3. **A pass would have been worse.** Had it somehow passed, the run would hold + an MMMU number attributed to `36ea38b940` that was measured on something + else. Nothing downstream could tell. + +Point 3 is the real defect. This workflow's whole claim is that a target is +trusted because its gates passed — and a gate that cannot say *what it +measured* cannot support that claim. + +## Proposed change + +1. **Record the tree state in the verdict.** `GateResult` gains the `git` + description of the repo at gate time — `git rev-parse HEAD` plus a dirty + flag, or better `git describe --always --dirty`. One subprocess call per + gate, and every number in `progress.yaml` and the journal becomes + attributable. This alone would have turned a silent hazard into a visible + one. +2. **Refuse to gate a dirty tree under the target's own path**, unless the run + says otherwise. The assembler writes and commits *before* handing back — + that is already the protocol — so at gate time the target's files should be + clean. A dirty target path at gate time means someone else is writing, and + that is worth a hard stop with a message naming the files. +3. **Say the invariant out loud in the role prompts.** Nothing currently tells + an agent that the repo is shared mutable state during a gate. The commit + discipline block warns about `git add -A` sweeping in a sibling's + half-written files; this is the same hazard one level up, and it is not + mentioned. + +Point 2 is the one with teeth. Point 1 is cheap and should land regardless — +even with the guard, a recorded tree state is what lets a number be defended +six months later. + +## The operator error underneath it + +This was not the workflow's fault alone: **I dispatched the audio agent in +parallel with the multimodal unit's gating, onto the same file.** The +parallelism was deliberate — closing the audio gap while gates ran — and the +collision was foreseeable. + +That is worth recording as its own lesson, because the run's standing +instruction is to maximise parallelism, and this is the boundary: **agents may +run in parallel across *different* files, never across the same one while it +is being gated.** The smith phase gets this right by construction — nine +smiths touched nine disjoint entries — and the assembler phase does not, +because there is only one `modeling.py`. + +## What was done + +The in-flight audio work was saved to +`05-units/audio-wip-mainagent.patch` (652 lines), the target files were +restored to the reviewed commit with `git checkout --`, and the audio agent +was told to stop, save its own patch, and wait for the gates to finish before +resuming. Nothing was lost. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/README.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/README.md new file mode 100644 index 000000000000..b2ca8efb98b0 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/README.md @@ -0,0 +1,44 @@ +# Findings — what a real onboarding run taught this workflow + +Each file here was written while something was going wrong, during the +onboarding of `MiMo-V2.6-Flash-RL` to 4x GB300 (2026-09-22/23). They are kept +because **every one of them presented as a model or target defect while it was +happening**, and none of them was. + +They are evidence, not a checklist. A finding names what was measured, what it +cost, and what was changed — so a later reader can tell a fixed problem from a +recurring one. + +## Read these two first + +- **`RETROSPECTIVE-2026-09-mimo-v2.6-flash.md`** — the whole run reflected on: + execution efficiency, flow length, inter-agent coordination, and agents' + own trial-and-error, each argued from measured numbers. Its five + highest-return changes are all implemented on this branch. +- **`20260922-F002-unified-targets-are-gated-at-tp-only.md`** — the one that + changed the numbers. A `unified` target was served at pure tensor + parallelism on every member, so the expert-parallel and attention-DP + branches were exercised by no gate while the target claimed all three. + +## The shape most of them share + +Two thirds are **environment failures wearing a model's clothes**: an MPI +launcher that makes every rank a leader and hangs forever with no error, a +`TMPDIR` too long for a unix socket, a supply survey that measured nothing and +printed a check mark, a gate that measured a tree somebody else was editing. + +The remaining third are **silent-correctness** problems, which on the +checkpoint that produced them were the dominant hazard: a wrong MXFP4 MoE +kernel does not raise, it serves fluent token soup; a missing input processor +answers text-only questions correctly, having never looked at the image. Most +of the defensive machinery on this branch — the boot gate asking a question +with a checkable answer, the insistence on negative controls, the tree-state +recording — follows from that. + +## Numbering + +`F001`..`F011` in the order found. A `-fix-report` file beside a finding is +what was actually implemented, which is not always what the finding proposed — +`F011`'s proposal to use `git describe --always --dirty` turned out to take +over 60 seconds on a lustre checkout and was replaced by something that takes +0.14 s. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/findings/RETROSPECTIVE-2026-09-mimo-v2.6-flash.md b/agent-flow/agent_flow/workflows/staircase_onboard/findings/RETROSPECTIVE-2026-09-mimo-v2.6-flash.md new file mode 100644 index 000000000000..1342f11f792a --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/findings/RETROSPECTIVE-2026-09-mimo-v2.6-flash.md @@ -0,0 +1,320 @@ +# Staircase workflow retrospective — MiMo-V2.6-Flash-RL run + +Written 2026-09-23, after the run closed with all eleven line items delivered. +Everything here is measured from this run, not proposed in the abstract. + +## The numbers the rest of this document argues from + +| | | +|---|---| +| wall clock | ~26 h | +| orchestrator launches | **7** | +| assembler dispatches | **10**, for **4** units | +| total role dispatches | 27 (25 answered, 2 orphaned) | +| gate runs | 138 | +| workflow defects found | 11 (F001-F011), 5 fixed with tests mid-run | +| compute sessions destroyed | **2** | +| agents stalled and killed | **3** | +| GPU-h on anchors that did not exist | ~12 | + +**Six of ten assembler dispatches did no construction.** That is the single +largest measurable waste and it has one cause, below. + +--- + +## 1. Execution efficiency + +### 1.1 The restart tax — the biggest single loss + +`task.yaml` is validated and loaded into `self.task_data` **once**, at +construction. Every gate afterwards reads that in-memory copy. So any spec fix +— and this run needed four — requires killing and restarting the orchestrator, +and **every restart re-enters the `assemble` stage and dispatches a fresh +assembler turn**. + +Measured: core went to turn 5 of a 6-turn budget with only turn 1 doing work. +MTP went to turn 3 the same way. The unit came within one turn of dying +`assembly-exhausted` with a green, twice-approved target sitting in `HEAD`. + +Three fixes, increasing in cost: + +1. **Re-read `task.yaml` at each stage boundary**, and log when the on-disk + copy differs from the loaded one. The spec is small; validation is + milliseconds. +2. **A resume must not consume a turn.** On re-entering `assemble`, check + whether the unit's decision file and `.done` already exist for the current + turn index. If the last turn produced `UNIT_READY` and nothing under the + target has changed since, advance to the gates instead of re-dispatching. +3. **Count turns by *work*, not by dispatch.** `unit_turns` exists to bound a + self-heal loop. It currently counts restarts, which is a different thing. + +Item 2 alone would have removed all six wasted dispatches. + +### 1.2 The container import is 4-8 min normally and 25-35 min under load + +Measured both. The run paid it on every one-shot job after the session died: +three jobs × ~30 min ≈ 1.5 h of pure waiting. + +`container_name` reuse works **per node**, so it does not help a one-shot job +that lands elsewhere. Two things would: + +- **Save a `.sqsh` once per run** and use the local path thereafter. The + cluster already has `/lustre/.../containers/*.sqsh` from earlier work; the + workflow should produce and register one as a preflight side effect. +- **Keep the session alive harder** (see §3.2), since a live session pays the + import exactly once. + +### 1.3 Gates are ~60% of wall clock and mostly correct + +138 gate runs. The member-set change (F002) tripled the gate bill deliberately +and it was the right call — it is what produced the 0.11-point agreement across +tp4/tep4/dep4 that justifies the `unified` claim. No complaint here, but two +cheap wins: + +- **Order members by cost.** `tep4` consistently booted in ~180 s against + `tp4`'s ~285-438 s. Running the cheapest member first surfaces a broken unit + sooner. +- **Let a boot failure skip the remaining members.** It already stops the + unit; it should not have been worth running member 2 and 3 first. + +--- + +## 2. Shortening the flow + +### 2.1 Anchors are on the critical path and should be phase 0 + +Neither the model card nor the technical report publishes a GSM8K, MMLU, MMMU +or VideoMME number. Both anchors had to be **measured on an independent +engine**, and both landed late: + +- GSM8K: needed before the core unit's accuracy gate. Made it, barely. +- MMMU: **discovered missing** only when I checked `mmmu.yaml` while MTP was + gating. Had I not looked, the multimodal unit would have stalled at its + accuracy gate hours later. + +**Preflight should check that every unit's `reference_yaml` actually contains +the row its `model_key` + `discriminators` select, and fail loudly if not.** +That check is pure Python, costs milliseconds, and would have converted a +late surprise into a startup error. It is the same class as F001 and F005: +something knowable in seconds, discovered in hours. + +Then: **dispatch anchor measurement concurrently with the plan phase.** It +needs no plan, no vocabulary and no target — only the checkpoint and an +engine. This run did that for MMMU by accident and it worked. + +### 2.2 The engine recipe is an artifact, and reusing it is worth 2.5 hours + +The first MMMU agent ran **2.5 h and produced nothing**. The second produced +the full 900-sample result in **13 min**. The only difference was that the +second was handed a *proven* recipe: the exact container +(`vllm/vllm-openai:mimo-v26-cu129`, the build that parses `store_dtype: +mxfp4`), the serve flags, the `--limit-mm-per-prompt` requirement, and the +short-`TMPDIR` constraint. + +**The run should produce a `reference_engine.yaml` the first time it gets an +external engine serving the checkpoint, and every later anchor job should read +it.** This run passed it hand-to-hand through my dispatch briefs, which worked +only because I was watching. + +### 2.3 `--max-plan-rounds 1` is right; the review is worth more than the re-draft + +Two plan rounds cost ~1 h and the reviewer returned **REJECT** both times with +substantive findings — the disagg engine limitation, the refutation that +`kv_cache_reuse` was inert, and the world-size-1 geometry error that would have +had smiths certifying cells no rank ever passes. + +That is a good trade. But note **the round-2 reviewer refuted one of my own +rulings** (R1's per-section slicing) by reading raw FP8 codes rather than +counting blocks — which the block count could not distinguish. The lesson is +not about the cap, it is §4.1 below. + +--- + +## 3. Coordination between agents + +### 3.1 The parallelism boundary is the file, and only the smith phase respects it + +The smith batch was the run's best-executed phase: **9 dispatches, 8 +concurrent, 9 disjoint catalog entries, zero collisions.** It works because a +catalog entry is one file plus one test file, and the queue is built from +disjoint surfaces by construction. + +The assembler phase has **one `modeling.py`**, so it cannot parallelize at +all — and when I tried to run other work alongside it, I caused **F011**: +I dispatched the audio agent in parallel with the multimodal unit's gating, +onto the same file. `modeling.py` went 2827 → 3199 lines *while the reviewer +was reading it*, and three MMMU gates were queued against a half-rewritten +vision path. + +The reviewer caught it only because it happened to re-read the line count. + +**What the workflow should enforce:** + +- **A gate records the tree state it measured** — `git describe --dirty` in + every `GateResult`. One subprocess call. Without it, no number is + attributable, and a *pass* on a dirty tree is worse than a failure. +- **A gate refuses to run when the target's own path is dirty**, unless + explicitly overridden. At gate time the assembler has already committed; + dirty means someone else is writing. +- **Say the invariant in the role prompts.** The commit-discipline block warns + about `git add -A` sweeping in a sibling's half-written files. This is the + same hazard one level up and is not mentioned anywhere. + +### 3.2 A destroyed session takes the orchestrator with it — twice + +Both session deaths had the same shape: an agent cleaning up its own work +killed the watcher, because `session_entry.sh` ends in +`exec bash spool_watcher.sh` — **the watcher *is* the allocation's main +process**. + +1. The disagg probe's teardown SIGINT'd an `mpirun`-rooted worker; the step + died at the same second. +2. The audio agent ran `pkill -9 -f smoke_audio.py` to clear a hung job. + +Both took the orchestrator and every queued job with them. + +**Fixes:** +- **The orchestrator must not live in the session it dispatches to.** It was + submitted as `free-staircase*` inside the same allocation, so a session death + is also an orchestrator death. Running it on the login node failed for other + reasons, which is why it ended up there — but the coupling should be named + and, ideally, the watcher should be a *child* the job supervises rather than + the job itself. +- **`pkill` by pattern is banned inside a session job.** Clear a hung job by + writing `STOP` or by killing a PID read from the log. This belongs in + `TURN_INSTRUCTIONS.md` as a hard rule. +- **A probe that must signal multi-rank workers gets its own allocation.** The + disagg probe adopted this after the first incident and it was right. + +### 3.3 A turn can commit its work and never hand back (F009) + +Dispatch 006 committed 380 lines with a passing receipt, then died without +writing its decision file. The orchestrator sat blocked **58 minutes** on a +missing zero-byte file while the work sat in `HEAD`, and would have waited the +full 4-hour timeout. + +**A per-turn heartbeat** — mirroring the existing `driver.heartbeat` one level +down — turns a four-hour silent block into a named failure in five minutes. +And a stalled turn should fail *that request*, not the batch. + +### 3.4 Findings must be able to reach in-flight dispatches + +`SMITH_CORRECTIONS.md` was invented mid-run because the round-2 plan review +landed *after* the queue was written, and one of its corrections (C1: the FP8 +GEMM widths) would otherwise have had a smith certify `N=192`, a width **no +rank ever passes**. + +It worked — dispatch 009 certified `N=320` instead and proved it by doubling a +`b_scale` row and watching exactly those output columns move. + +**This should be a first-class mechanism**: a corrections file the orchestrator +writes and every dispatch prompt includes by reference, with an explicit +precedence rule ("where this and `vocabulary.yaml` disagree, this wins"). + +--- + +## 4. Agents' own trial-and-error + +### 4.1 Arithmetic that agrees is not evidence + +The most expensive intellectual error in this run was mine. R1 concluded the +fused `qkv_proj` could be sliced per-rank **and then per-section**, because +both hypotheses predict 27 scale blocks per rank and 27 × 4 = the observed 108. +I treated the agreement as confirmation. + +The round-2 reviewer read the **raw FP8 code values** and found block 25 +straddles `K[128:192]` and `V[0:64]` — `scale[25,0] = 7.75e-4` against +`scale[26,0] = 7.36e-5`, an order of magnitude apart, which only happens across +a mixed span. + +**The rule worth encoding: when two hypotheses predict the same count, the +count is not the discriminator.** Find the quantity that differs. This belongs +in the kernel-feasibility and modeling-core skills as a named trap, because it +is exactly the shape of reasoning an agent finds satisfying and is wrong. + +### 4.2 "Prove your instrument can fail" was emergent — make it mandatory + +Several smiths did this without being told, and it is the single most valuable +behaviour the run produced: + +| control | correct | deliberately wrong | +|---|---|---| +| `mxfp8_quantize` layout | 1.97 ulp | **203.32 ulp** (swizzled) | +| MoE expert offset | inside gate | **>100 ulp**, still finite and populated | +| MTP concat order | 2.59 acceptance | **1.00** — total rejection, text still correct | +| ViT sink semantics | 0.467 of gate | **294.1×** outside (denominator-side sink) | +| uniform-KV replication | `max|d| = 0` | tile control separates at rel ≈ 1.0 | + +On a checkpoint where **a wrong MXFP4 MoE does not raise** — it serves fluent +token soup — a test that cannot distinguish right from wrong bytes certifies +nothing. + +**`staircase-kernel-onboard` and `-verify` should require a negative control +and refuse a receipt without one.** It is cheap, and it converts "the op ran" +into "the test can tell". + +### 4.3 Agents over-polish the wrong artifact + +Two clear cases: + +- The multimodal equivalence agent spent **~2 h perfecting PNG rendering** + (triangle orientation, digit spacing) and produced **zero** reference + outputs. The images it had were already adequate — they exercise the tower's + geometry, which was the point. +- The first MMMU agent spent **2.5 h** on benchmark selection and produced no + artifact. + +Both needed the same instruction, which I added to their replacements and +which worked immediately: **name the minimum artifact and produce it first, +then improve.** The replacement brief said "if you find yourself editing a +PNG, you have gone wrong" and it delivered in one job. + +**Dispatch briefs should carry a first-artifact checkpoint**: what file must +exist, and by roughly when. An agent with no intermediate deliverable will +optimise whatever it is currently touching. + +### 4.4 A probe anchors to its first hypothesis + +The disagg probe was excellent work — it measured two engine walls, refuted a +claim the plan review had made, and correctly recommended descope **given the +framing it inherited**: "what would the engine need to tolerate this?" + +It never asked the other question: **"what would have to be true for this to +work, and can the target make it true?"** Replicating KV heads 4→8 removes both +walls with no engine change at all, and it was found only because I posed that +question afterwards. + +**A probe that concludes BLOCKED should be required to enumerate the +preconditions of the thing it is blocking, and say which are under the +target's control.** Here the precondition was "uniform per-layer KV geometry", +and the target owned it entirely. + +### 4.5 Host facts were re-derived by agent after agent + +`pytest` exits 4 here even when green. `TMPDIR` must be ≤ 61 characters. A +4-rank launch needs the `SLURM_*` scrub plus `--oversubscribe`. `UCX ERROR +uct_iface_open` is noise. `SMOKE_RC=$?` after a pipe is the pipe's exit code. +`.venv-3.12/bin/python` resolves only inside the container. + +**I injected these by hand into roughly fifteen dispatch briefs.** Every one +was learned the expensive way by someone. + +`TROUBLESHOOTING.md` exists but is keyed on symptoms *after* failure. What is +missing is a **preventive, machine-readable `host_facts.md` in the workspace +that every role prompt includes automatically** — written by preflight, which +already discovers most of it, and appendable by any role that learns a new one. + +--- + +## The five changes with the best return + +1. **A resume must not consume an assembler turn** (§1.1) — removes 6 of 10 + wasted dispatches. +2. **Gates record and require a clean tree** (§3.1) — F011 nearly attached an + MMMU number to a commit it was not measured on. +3. **Preflight validates that every unit's anchor row exists** (§2.1) — turns + a late stall into a startup error. +4. **Negative controls are mandatory for a receipt** (§4.2) — on this + checkpoint it is the difference between a test and a formality. +5. **A `host_facts.md` every role gets automatically** (§4.5) — stops the same + six facts being rediscovered fifteen times. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/gates.py b/agent-flow/agent_flow/workflows/staircase_onboard/gates.py new file mode 100644 index 000000000000..133851f45aa2 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/gates.py @@ -0,0 +1,1874 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The gates — commands with exit codes, never agent judgements. + +This module is the reason the workflow has no ``--min-score`` and no QA +agent scoring an artifact out of ten. A staircase target is trusted because +its gates passed, and a gate is a process that either returned zero and +printed a number above its anchor, or did not. Putting that behind an agent +turn adds latency and adds a way for the verdict to be rationalized. + +Three gates, in ascending cost: + +1. **boot** — minutes, binary. The engine cold-starts under + ``TRTLLM_MODELING_V2=require``, the weight manifest covers what it + claims, and a handful of greedy continuations come back. Catches + catastrophes, not accuracy. Runs after every assembler turn. +2. **accuracy** — the release criterion. ``trtllm-eval`` against the anchor + in TensorRT-LLM's own accuracy references. One-sided: + ``measured >= reference - tolerance``. Runs at a unit boundary. +3. **acceptance** — required only when a feature is *distribution-preserving + by construction*, speculative decoding above all. Rejection sampling + holds the emitted distribution to the target model's, so a miscomputed + draft path produces correct text **more slowly**: boot passes, accuracy + passes, and only speed moves. The detector is ``acceptance_length`` + against the same checkpoint under stock in-tree modeling at the same + workload and the same ``speculative_config`` — i.e. the same harness + with ``TRTLLM_MODELING_V2=off`` instead of ``=require``, which makes it + one variable rather than two harnesses. + +``TRTLLM_MODELING_V2=require`` is not optional for any of them. 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 — the single most expensive mistake available here. +""" + +from __future__ import annotations + +import json +import os +import re +import shlex +import subprocess +import tempfile +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Mapping, Sequence + +import yaml + +from agent_flow.runtime import ( + CommandSpec, + ExecutionRecord, + LocalCommandExecutor, + SpoolCommandExecutor, +) + +from .delivery import MODELING_V2_PACKAGE +from .environment import RunEnvironment +from .task_schema import ( + PARALLEL_UNIFIED, + TaskSchemaError, + derive_family, + load_prompts_file, + parse_parallel_segment, +) + +# Minutes-scale by design. A boot gate that needs longer than this is not a +# boot gate any more — it is an accuracy run with no anchor. +BOOT_TIMEOUT_S = 1800 +ACCURACY_TIMEOUT_S = 4 * 60 * 60 + +MODELING_V2_ENV = "TRTLLM_MODELING_V2" +MODE_REQUIRE = "require" +MODE_OFF = "off" + +#: A target records the metadata of its catalog calls and prints it at exit +#: (``=== MODELINGV2 RUNTIME VALUE DUMP``) only when this is 1, so a served +#: engine stays quiet. Every gate and witness run sets it: the auditor reads +#: those dumps against the catalog contracts. +RUNTIME_VALUE_DUMP_ENV = "TRTLLM_MODELING_V2_VALUE_DUMP" + +#: What a gate records when it could not ask git anything — no repository, no +#: ``git`` on the host, no commits yet. Never an exception: a gate adjudicates +#: a model, and failing one because a VCS query failed would attribute a +#: tooling problem to a target. +UNKNOWN_TREE_STATE = "unknown" + +#: Set to ``1`` to let a gate measure a dirty target path anyway. For a human +#: driving a gate by hand against uncommitted edits, which is a real and +#: reasonable thing to do — but the resulting verdict carries +#: ``dirty_override`` and says so in its summary, because a number measured +#: that way is not attributable to any commit. +ALLOW_DIRTY_ENV = "STAIRCASE_ALLOW_DIRTY_GATE" + +# Seconds. A git query that takes longer than this is a sick host, and the +# gate's answer is `unknown` rather than a stall in front of a multi-hour run. +_GIT_TIMEOUT_S = 15 + +# ``lm-eval accuracy: 90.67`` and +# ``lm-eval average accuracy: 90.67`` — the two shapes +# ``tensorrt_llm/evaluate/lm_eval.py`` emits at the end of a run. +_LM_EVAL_SCORE = re.compile(r"lm-eval\s+\S+\s+.*?accuracy:\s*([0-9]+(?:\.[0-9]+)?)") +# ``MMLU weighted average accuracy: 74.98 (14042)`` — the MMLU path has its +# own reporting because it does not go through lm-eval. +_MMLU_SCORE = re.compile(r"MMLU weighted average accuracy:\s*([0-9]+(?:\.[0-9]+)?)") +# ``Spec-dec stats: AL (decoded tokens/step, iteration-weighted) 3.351 (...)`` +_ACCEPTANCE_LENGTH = re.compile(r"Spec-dec stats:\s*AL\b[^)]*\)\s*([0-9]+(?:\.[0-9]+)?)") +# trtllm-bench reports the same quantity under its own name. +_ACCEPTANCE_LENGTH_ALT = re.compile(r"acceptance[_ ]length\D{0,12}([0-9]+(?:\.[0-9]+)?)", re.I) +# `[0] Accept rate: 1.88 (tokens=32, iterations=17)` --- what +# `examples/llm-api/quickstart_advanced.py` prints, once per request. +# +# This one is load-bearing rather than a nicety: **nothing in this tree emits +# either of the two patterns above.** `trtllm-eval` reports no speculative +# statistic at all, so an acceptance gate built on the eval command reads +# `None` from both legs and fails every time with "could not read +# acceptance_length", which is indistinguishable from a broken draft path. +# The quantity is captured from the parenthetical rather than from the +# printed ratio so that several requests combine correctly --- see +# :func:`parse_acceptance_length`. +_ACCEPT_RATE_PARTS = re.compile(r"Accept rate:\s*[0-9.]+\s*\(tokens=(\d+),\s*iterations=(\d+)\)") + +#: Tokens each acceptance prompt may generate when the task declares none. +DEFAULT_ACCEPTANCE_MAX_TOKENS = 256 +#: Generation iterations the acceptance workload must reach, summed over its +#: prompts, before the gate reads a rate from it. The degenerate case this +#: exists for measured 4: three short-answer boot prompts on DeepSeek-V4.1, +#: each ending in EOS after 2-3 tokens. A task may raise it through +#: ``accuracy.acceptance.min_draft_iterations``. +DEFAULT_MIN_DRAFT_ITERATIONS = 32 + +ShellRunner = Callable[[str, Mapping[str, str] | None, int], "ShellResult"] + +#: The one process-wide local executor. Stateless, so sharing it is free and +#: constructing one per gate would say something about lifetime that is not +#: true. +_LOCAL_EXECUTOR = LocalCommandExecutor() + + +@dataclass(frozen=True) +class ShellResult: + returncode: int + stdout: str = "" + stderr: str = "" + timed_out: bool = False + + @property + def ok(self) -> bool: + return self.returncode == 0 and not self.timed_out + + +def shell_result_from(record: ExecutionRecord) -> ShellResult: + """Narrow an :class:`ExecutionRecord` to what a gate reads. + + A gate adjudicates on four things: did it finish, did it exit zero, and + what did it print. The record carries more — which executor, how long, + where the full log is — and that belongs in the run's evidence rather + than in a threshold comparison, so it is deliberately dropped here + instead of being smuggled into a gate's decision. + """ + return ShellResult( + returncode=record.returncode, + stdout=record.stdout, + stderr=record.stderr, + timed_out=record.timed_out, + ) + + +def _git(repo_root: str | Path, *args: str) -> str | None: + """Ask git one question about ``repo_root``. ``None`` if it cannot answer. + + Every failure mode collapses to ``None`` on purpose — no repository, no + ``git`` binary, no commits yet, a hung filesystem. The callers of this + function are gates, and a gate that dies because a VCS query died + reports a tooling failure as a modeling failure. Called with an argv + rather than through a shell because the only interpolated value is a + path that may contain anything. + """ + try: + completed = subprocess.run( + ["git", "-C", str(repo_root), *args], + capture_output=True, + text=True, + timeout=_GIT_TIMEOUT_S, + check=False, + ) + except (OSError, subprocess.SubprocessError): + return None + if completed.returncode != 0: + return None + return completed.stdout + + +def describe_tree_state(repo_root: str | Path) -> str: + """The commit a gate is about to measure, as one token. + + ``git describe --always``: where the tree carries tags it names the state + the way a human reads a release (``v0.20.0rc0-8958-g66ebfec7e2``) and + where it does not it falls back to the abbreviated commit, so one field + covers both and a reader never has to know which they are looking at. + + What it buys is attribution. Gates measure whatever is on disk under + ``trtllm_repo_path`` when the engine imports it; before this there was no + record of what that was, so an accuracy number could be — and in this + workflow's own run nearly was — attributed to a commit it was never + measured on. The token is not a pin and does not pretend to be one; it is + what lets a reader tell whether two numbers came from the same code. + + **Deliberately not ``--dirty``**, which is the obvious spelling and the + wrong one here, for two independent reasons: + + * *Cost.* ``--dirty`` runs a whole-worktree ``diff-index`` refresh. + Measured on this checkout, on the networked filesystem the runs use: + ``describe --always`` 0.09 s, ``describe --always --dirty`` past 60 s. + A gate must not sit a minute behind a VCS nicety, and with a timeout + in front of it the field would simply always read ``unknown``. + * *Information.* A staircase run's repo is dirty essentially always — + onboarding records, reference YAMLs, workflow fixes — so a repo-wide + dirty marker is set on every gate of every run and distinguishes + nothing. The marker worth carrying is scoped to the target, and + :func:`inspect_tree` appends that one instead, from + :func:`dirty_target_files`. + """ + described = _git(repo_root, "describe", "--always") + if described is None: + return UNKNOWN_TREE_STATE + return described.strip() or UNKNOWN_TREE_STATE + + +def target_family_rel(task_data: Mapping[str, Any]) -> str | None: + """The repo-relative directory holding every target for this family. + + Derived the way :func:`~.preflight.resolve_target_identity` derives it — + ``task.yaml``'s ``target.family`` when it is stated, otherwise the + checkpoint's own ``architectures[0]`` through + :func:`~.task_schema.derive_family` — rather than spelled here, because a + second spelling of the target path is a guard that watches a directory + nothing writes to. + + The **family** directory and not the run's own ``target_rel``: a run may + legitimately write routing and shared modules beside its target, and all + of it is code the gate is about to import. + + ``None`` when the family cannot be derived at all (no override and no + readable ``config.json``). A caller that gets ``None`` has nothing to + guard and must not invent a path to guard instead. + """ + override = dict(task_data.get("target") or {}) + family = override.get("family") + if not family: + config = Path(str(task_data.get("checkpoint_path") or "")) / "config.json" + try: + architectures = (json.loads(config.read_text(encoding="utf-8")) or {}).get( + "architectures" + ) + except (OSError, ValueError, AttributeError): + return None + if not architectures: + return None + family = derive_family(str(architectures[0])) + return f"{MODELING_V2_PACKAGE}/models/{family}" + + +def dirty_target_files(repo_root: str | Path, target_rel: str) -> tuple[str, ...]: + """Uncommitted paths under ``target_rel``, as git reports them. + + ``-uall`` so an untracked *file* is named rather than collapsed into its + directory: the message this feeds exists to say who is writing, and + ``?? .../models/qwen3_5/`` does not say it. + + An empty tuple means both "clean" and "could not tell" — a deliberate + collapse. The guard this feeds stops a gate, and stopping a gate because + git was unavailable would make the workflow unusable on any host without + a checkout, which includes every host that runs these tests. + + Cheap because of the pathspec: measured at 0.14 s on this checkout + against the minute a whole-worktree status costs here, which is the same + asymmetry :func:`describe_tree_state` records. + """ + status = _git(repo_root, "status", "--porcelain", "-uall", "--", target_rel) + if not status: + return () + return tuple(line[3:].strip() for line in status.splitlines() if line.strip()) + + +@dataclass(frozen=True) +class TreeState: + """What the repository looked like when a gate was about to run.""" + + description: str = UNKNOWN_TREE_STATE + #: Uncommitted paths under the target's own family directory. + dirty: tuple[str, ...] = () + #: True when the run was told to measure a dirty target path anyway. + override: bool = False + + @property + def blocked(self) -> bool: + return bool(self.dirty) and not self.override + + def block_detail(self) -> str: + shown = ", ".join(self.dirty[:10]) + more = f" (+{len(self.dirty) - 10} more)" if len(self.dirty) > 10 else "" + return ( + f"the target's own path is dirty at gate time, so this gate would measure " + f"a tree nobody reviewed: {shown}{more}. By protocol the assembler commits " + f"before handing back, so uncommitted files here mean another writer is " + f"live. Let that writer finish, or set {ALLOW_DIRTY_ENV}=1 to measure " + f"anyway and have the verdict say so." + ) + + +def inspect_tree(task_data: Mapping[str, Any], *, allow_dirty: bool | None = None) -> TreeState: + """Read the repo state a gate is about to measure, and whether to refuse. + + Two questions, one place, because a gate wants both at the same instant: + *what am I measuring* (recorded always) and *is someone else writing it* + (a hard stop). + + The dirtiness question is scoped to the **target's own family path**, not + to the repo. A staircase run has onboarding records, workflow fixes and + reference YAMLs in flight more or less constantly, and a gate that + refused to run while any of them were uncommitted would never run at all. + The narrow scope is what makes the check something a run can leave on. + + ``allow_dirty`` overrides the ``STAIRCASE_ALLOW_DIRTY_GATE`` environment + variable; either way the escape hatch is explicit and lands in the + verdict, so a number measured over someone else's half-written edit is + marked as one rather than passing as an ordinary result. + """ + repo_root = task_data.get("trtllm_repo_path") + if not repo_root: + return TreeState() + if allow_dirty is None: + allow_dirty = os.environ.get(ALLOW_DIRTY_ENV, "").strip().lower() in {"1", "true", "yes"} + described = describe_tree_state(repo_root) + target_rel = target_family_rel(task_data) + dirty = dirty_target_files(repo_root, target_rel) if target_rel else () + # The dirty marker is appended here rather than asked of git, and it is + # scoped to the target — see :func:`describe_tree_state` for why the + # repo-wide one is both a minute slower and less informative. A reader + # seeing `-dirty` on a gate line therefore knows something was + # uncommitted **in the code that gate measured**, not merely somewhere. + description = f"{described}-dirty" if dirty else described + return TreeState(description=description, dirty=dirty, override=bool(dirty) and allow_dirty) + + +@dataclass(frozen=True) +class GateResult: + """One gate's verdict, with everything needed to defend it later.""" + + name: str + passed: bool + detail: str + #: The parallel configuration this verdict was measured at. A ``unified`` + #: target is one codebase over a family of them, so a gate line that does + #: not name the member cannot be told apart from the two members it did + #: not run — the same class of failure, one level up, as a directory name + #: that lies. ``None`` only for a caller that predates the member set. + member: str | None = None + command: str = "" + measured: float | None = None + reference: float | None = None + tolerance: float | None = None + duration_s: float = 0.0 + stdout_tail: str = "" + #: The repository state this verdict was measured on — + #: ``git describe --always`` under ``trtllm_repo_path``, suffixed + #: ``-dirty`` when the target's own path was uncommitted, or + #: :data:`UNKNOWN_TREE_STATE` where git could not answer. A gate measures + #: whatever is on disk when the engine imports it, so without this a + #: number is attributable to nothing: the run that produced this field had + #: a second agent rewriting the target's ``modeling.py`` while three + #: accuracy gates were queued against it. See :func:`describe_tree_state`. + tree_state: str = UNKNOWN_TREE_STATE + #: True when this verdict was measured over an uncommitted target path + #: because the run said to. Recorded rather than merely permitted — the + #: point of the escape hatch is that the resulting number stays marked. + dirty_override: bool = False + #: The per-prompt numbers behind ``measured``, where a gate has them. The + #: acceptance gate's headline is one iteration-weighted rate, and a mean + #: cannot say that one prompt ended in EOS after two tokens while another + #: accepted at the ceiling; this rides in the journal line so it can. + breakdown: str = "" + + def summary(self) -> str: + """One line for the journal, carrying the numbers verbatim. + + The member is part of the label, not an aside: `progress.yaml`, the + journal and the reviewer's evidence all read this line, and three of + them per unit only mean three measurements if each says which + configuration it measured. + + The tree state rides in the same line for the same reason one level + up: `progress.yaml` and the journal are what a reader consults later, + and three measurements are only comparable if each says what code it + measured. + """ + verdict = "PASS" if self.passed else "FAIL" + label = f"{self.name}@{self.member}" if self.member else self.name + tree = f" [tree {self.tree_state}]" + if self.dirty_override: + tree = f" [tree {self.tree_state}; DIRTY TARGET PATH measured by override]" + breakdown = f" [{self.breakdown}]" if self.breakdown else "" + if self.measured is not None and self.reference is not None: + margin = self.measured - (self.reference - (self.tolerance or 0.0)) + return ( + f"[{verdict}] {label}: measured {self.measured:.4f} vs " + f"reference {self.reference:.4f} - tol {self.tolerance or 0.0:.2f} " + f"(margin {margin:+.4f}){breakdown} in {self.duration_s:.0f}s{tree}" + ) + return f"[{verdict}] {label}: {self.detail}{breakdown} ({self.duration_s:.0f}s){tree}" + + +def _verdict_factory(state: TreeState) -> Callable[..., GateResult]: + """A :class:`GateResult` constructor that stamps ``state`` onto every verdict. + + Each gate has several exits — timed out, non-zero, unparsable, passed — + and the tree state belongs on **all** of them equally: a failure that + cannot say what tree it failed on costs the same debugging as a pass that + cannot say what tree it passed on. Binding it once per gate is what keeps + a later exit from being the one that forgets. + """ + + def verdict(**fields: Any) -> GateResult: + return GateResult(tree_state=state.description, dirty_override=state.override, **fields) + + return verdict + + +@dataclass(frozen=True) +class BootSpec: + """What the boot gate sends and what it insists on seeing back. + + The keyword list starts **empty on purpose**. Staircase's rule is that + the first engine run doubles as keyword verification — you freeze the + cases only after observing the continuations — so an assembler that has + not yet seen this model generate anything has no business asserting + what it will say. Until then the gate asserts the weaker, still useful + thing: the engine cold-starts and produces non-empty output. + """ + + prompts: tuple[str, ...] = ( + "The capital of France is", + "In one sentence, photosynthesis is", + ) + keywords: tuple[str, ...] = () + max_tokens: int = 32 + + @classmethod + def load(cls, path: str | Path) -> "BootSpec": + source = Path(path) + if not source.is_file(): + return cls() + data = yaml.safe_load(source.read_text(encoding="utf-8")) or {} + if not isinstance(data, dict): + raise ValueError(f"{source} must hold a YAML mapping, got {type(data).__name__}") + prompts = tuple(str(p) for p in data.get("prompts") or cls.prompts) + return cls( + prompts=prompts or cls.prompts, + keywords=tuple(str(k) for k in data.get("keywords") or ()), + max_tokens=int(data.get("max_tokens", cls.max_tokens)), + ) + + +def _default_shell_runner( + command: str, + env: Mapping[str, str] | None = None, + timeout: int = BOOT_TIMEOUT_S, +) -> ShellResult: + """Run one command line through the shell, on this host. + + A thin binding of :class:`~agent_flow.runtime.LocalCommandExecutor` to + the callable shape this module's gates take. The shell is deliberate: + the launch prefix resolved from ``task.yaml`` *is* shell syntax (env + assignments, or an ``srun`` line), and re-parsing it into an argv would + mean reimplementing the shell. Nothing here interpolates agent-authored + text — the prefix comes from the task spec and the command from this + module. + """ + return shell_result_from( + _LOCAL_EXECUTOR.execute(CommandSpec(argv=command, env=env, timeout_s=timeout, label="gate")) + ) + + +#: The two gate commands take different flag vocabularies for the same +#: topology; see :func:`topology_flags`. Named rather than spelled inline +#: so a third caller cannot invent a third spelling by accident. +QUICKSTART_DIALECT = "quickstart" +EVAL_DIALECT = "eval" + + +def topology_flags( + segment: str, + world_size: int, + member: str | None = None, + *, + dialect: str = QUICKSTART_DIALECT, +) -> list[str]: + """Translate a parallel segment into ``quickstart_advanced.py`` flags. + + Every gate is **topology-blind** on its own: a ``tep4`` target served at + tp2 boots, scores, and produces a self-consistent curve while the + directory name lies. Deriving the flags from the segment here — rather + than letting an agent type them — is what keeps the thing measured and + the thing named in agreement. + + ``member`` is the configuration actually being served, and it is what + gets translated when it is given. ``unified`` names a **family**, not a + configuration, so translating the segment alone emitted ``--tp_size=N`` + and nothing else: expert parallelism and attention data parallelism — + the two branches a unified codebase exists to carry — were switched on + by no gate, while the target's Coverage section claimed all three. See + :func:`~.task_schema.members_for`, which decides what the members are. + + Without a member, ``unified`` keeps its old ``--tp_size`` reading rather + than raising: the fallback is for a caller that predates the member set, + and the workflow always passes one. + + ``dialect`` names **which command these flags are for**, because the two + gate commands spell the same topology differently and one of them cannot + spell it at all: + + ================== =========================== ========================== + axis ``quickstart_advanced.py`` ``trtllm-eval`` + ================== =========================== ========================== + tensor parallel ``--tp_size`` ``--tp_size`` + expert parallel ``--moe_ep_size`` ``--ep_size`` + attention DP ``--enable_attention_dp`` **no flag at all** + ================== =========================== ========================== + + Measured, on the first ``tep4`` accuracy gate of this run:: + + Error: No such option: --moe_ep_size + (Possible options: --ep_size, --pp_size, --tp_size) + + ``tp4`` passed the same gate minutes earlier because it emits only + ``--tp_size``, which both commands accept — so the defect was invisible + at the one member that does not exercise it, and would have been + invisible for any target that is not ``unified``. + + Attention data parallelism has no ``trtllm-eval`` flag, so under the + ``eval`` dialect it is **not emitted here**; it has to arrive through + ``--extra_llm_api_options`` as ``enable_attention_dp: true``. Emitting it + here would reintroduce the failure this parameter exists to fix, and + silently dropping it would serve a ``dep`` member at ``tep`` while the + verdict claimed otherwise — so :func:`eval_command` overlays it onto the + unit's options file instead, and that is not optional. + """ + kind, declared = parse_parallel_segment(member or segment) + size = declared if declared is not None else world_size + if kind == PARALLEL_UNIFIED: + return [f"--tp_size={size}"] + flags = [f"--tp_size={size}"] + if kind in ("tep", "dep"): + flags.append(f"--ep_size={size}" if dialect == EVAL_DIALECT else f"--moe_ep_size={size}") + if kind == "dep" and dialect != EVAL_DIALECT: + flags.append("--enable_attention_dp") + return flags + + +def member_llm_api_overlay(member: str | None) -> dict[str, Any]: + """The LLM-API settings a member needs that its gate command cannot carry. + + Exactly one today: attention data parallelism, which + ``quickstart_advanced.py`` takes as ``--enable_attention_dp`` and + ``trtllm-eval`` cannot take at all. Returned as a mapping rather than + written by the caller so that "what does `dep` mean" has one answer, next + to the flag table it is the exception to. + """ + if not member: + return {} + kind, _ = parse_parallel_segment(member) + return {"enable_attention_dp": True} if kind == "dep" else {} + + +#: Startup cost of a member's topology, cheapest first. The rank is on the +#: PARALLEL KIND rather than on a measured duration, because what a gate pays +#: before it can report anything is dominated by weight loading and by the +#: collectives the engine builds at init, and both follow from the kind: +#: +#: * ``tep`` — the expert table is PARTITIONED across ranks, so each rank +#: materializes a fraction of the parameters that dominate a MoE checkpoint, +#: and the attention weights are still tensor-sliced. Cheapest to start. +#: * ``dep`` — the same expert partition, plus attention data parallelism, +#: which REPLICATES the attention weights on every rank instead of slicing +#: them. More bytes per rank than ``tep``, fewer than ``tp``. +#: * ``tp`` — no expert partition at all: every rank holds a slice of every +#: expert, which is the largest per-rank load and the most sharding work at +#: init. The degenerate member, and the most expensive one. +#: +#: This run measured ``tep4`` at ~180 s against ``tp4``'s 285-438 s, which is +#: consistent with the ordering but is not what the ordering is derived from — +#: two data points on one checkpoint would not survive a different one. +_MEMBER_COST_RANK: dict[str, int] = {"tep": 0, "dep": 1, "tp": 2} + +# A member whose kind this module does not recognize — including ``unified``, +# which names a family rather than a configuration — sorts last. Unknown cost +# is not cheap cost, and the point of running cheapest first is to fail early. +_UNKNOWN_MEMBER_COST = len(_MEMBER_COST_RANK) + + +def order_members_by_cost(members: Sequence[str]) -> list[str]: + """Return ``members`` cheapest-first, for a unit gated over several. + + Pure, total and deterministic: the answer depends on the SET of members + and never on the order they arrived in, on the filesystem, or on any + previous run, so two runs of the same unit gate in the same sequence and + their journals compare line by line. An unrecognized member is ordered + last rather than raising — this function decides a running order, and a + gate loop that dies while choosing one has failed worse than a gate loop + that runs in a poor order. + + Why order matters at all: **a unit is gated at every member, and a broken + unit fails at all of them.** The first member to fail is the one that + stops the unit, so making it the cheapest one buys back the rest. Ties + break on world size ascending, then on the member name, both only so that + the result is a function of its input. + + The cost model is :data:`_MEMBER_COST_RANK`; the reasoning for the order + of the kinds lives there, next to the ranks it justifies. + """ + + def key(member: str) -> tuple[int, int, str]: + try: + kind, size = parse_parallel_segment(member) + except TaskSchemaError: + return (_UNKNOWN_MEMBER_COST, 0, member) + return ( + _MEMBER_COST_RANK.get(kind, _UNKNOWN_MEMBER_COST), + size if size is not None else 0, + member, + ) + + return sorted(members, key=key) + + +def spool_runner( + spool: str | Path, + *, + poll_s: float = 3.0, + name_prefix: str = "gate", +) -> ShellRunner: + """Return a :data:`ShellRunner` that executes through a compute session. + + The session is a long-lived allocation whose own main process watches + ``/pending`` (see ``scripts/session_watcher.sh``). Work is a + file, which is the property that matters: **anything able to write a + file can run GPU work**, including a plain Python process on a host + with no GPU and no reachable scheduler. + + That is what hands gate execution back to the orchestrator. The + alternative on such a host is to route every gate through an agent + holding an MCP tool, and then the verdict rests on that agent's + account of what a number was. Here it rests on a file. + + Completion is the appearance of ``done/.rc``, written last and + renamed into place by the watcher, so a reader that sees it can trust + the log beside it is complete. + + The protocol itself lives in + :class:`~agent_flow.runtime.SpoolCommandExecutor`: a spool is a way to + reach an allocation, not a Staircase concept, and a run that stops using + one should not have to touch a gate. + """ + executor = SpoolCommandExecutor(spool, poll_s=poll_s, name_prefix=name_prefix) + + def run( + command: str, + env: Mapping[str, str] | None = None, + timeout: int = BOOT_TIMEOUT_S, + ) -> ShellResult: + return shell_result_from( + executor.execute(CommandSpec(argv=command, env=env, timeout_s=timeout, label="gate")) + ) + + return run + + +def resolve_runtime(task_data: Mapping[str, Any], unit: str | None) -> Mapping[str, Any]: + """The ``runtime`` block for ``unit``, with the per-unit form merged. + + Accepts the flat form (the whole run) or a mapping from unit name to a + partial block layered over ``default``. Per-unit because the units of one + run can genuinely DISAGREE about a switch, which the other per-unit fields + never had to handle: ``benchmark_args`` differed because the subcommands + differ, but here `core` and `mtp` were gated with block reuse **on** and + passed, while the multimodal path requires it **off**. A run-wide value + cannot be right for all three. + + Why this field and not the LLM-API options file, which is already + per-unit: ``quickstart_advanced.py`` has no ``--extra_llm_api_options``, + so the **boot gate** -- the gate that runs first and most often -- cannot + read that file at all. This block is the only channel it has. + + Merging is key-by-key over ``default``, so a unit states only what + differs. ``extra_flags`` is an ordinary key and therefore REPLACES rather + than extends: a unit that needs the default flags plus one more restates + them, which is verbose and unambiguous. + """ + block = task_data.get("runtime") or {} + if not any(isinstance(v, dict) for v in block.values()): + return block + merged = dict(block.get("default") or {}) + chosen = block.get(unit) + if isinstance(chosen, dict): + merged.update(chosen) + return merged + + +def runtime_flags(runtime: Mapping[str, Any]) -> list[str]: + """Translate the ``runtime`` block into ``quickstart_advanced.py`` flags. + + **Two of the four switches are spelled as negations on the command + line.** ``overlap_scheduler`` and ``kv_cache_reuse`` are on by default + in the engine, so wanting them means emitting *nothing* and not + wanting them means emitting ``--disable_*``. That inversion is the + whole reason this function exists: asked to "turn on block reuse", the + obvious move is to add a flag, and the flag that exists turns it off. + + Doing it once here also keeps the gate and the prompts in agreement — + the assembler is told what is enabled, not what was typed. + """ + flags: list[str] = [] + if runtime.get("cuda_graph"): + flags.append("--use_cuda_graph") + if runtime.get("chunked_prefill"): + flags.append("--enable_chunked_prefill") + if not runtime.get("overlap_scheduler", True): + flags.append("--disable_overlap_scheduler") + if not runtime.get("kv_cache_reuse", True): + flags.append("--disable_kv_cache_reuse") + flags.extend(str(f) for f in runtime.get("extra_flags", ())) + return flags + + +def speculation_flags(features: Any, *, max_draft_len: int = 1) -> list[str]: + """Translate a speculative feature into its ``quickstart_advanced.py`` flags. + + Only the speculator features produce flags; everything else in the + feature list is a property of the assembled target rather than of the + command that runs it. + + ``dspark`` gets its own algorithm and never ``MTP``, even though a + DeepSeek-V4.1 checkpoint stores its DSpark stages as ``mtp.*`` with + ``num_nextn_predict_layers: 3``. ``MTP`` over those weights builds three + chained one-token heads where the checkpoint trained one 3-stage pass over + a 5-token block: it fails to load, or it drafts garbage that verification + rejects -- and the second passes accuracy. The schema refuses ``mtp`` and + ``dspark`` together, so the order here never has to decide between them. + """ + names = {str(f).lower() for f in (features or ())} + for feature, algo in (("mtp", "MTP"), ("eagle3", "Eagle3"), ("dspark", "DSPARK")): + if feature in names: + return [f"--spec_decode_algo={algo}", f"--spec_decode_max_draft_len={max_draft_len}"] + return [] + + +def resolve_max_draft_len(accuracy: Mapping[str, Any], unit: str | None) -> int: + """How many draft tokens this unit's gates must ask for. + + **The three gates read the speculator's depth from two different places, + and before this they disagreed.** `trtllm-eval` takes no speculation + flags, so the accuracy gate learns the depth from the unit's + `--extra_llm_api_options` YAML -- `speculative_config.num_nextn_predict_layers`. + The boot and acceptance gates drive `quickstart_advanced.py`, which cannot + read that file, so they take it as `--spec_decode_max_draft_len` and it + was hard-defaulted to 1. + + Measured on this run's `mtp` unit: the options file says 3 and the flag + said 1, so accuracy would have measured MTP-3 while acceptance measured + MTP-1 against a ceiling of 2.0 -- and the target's real acceptance is + 2.44-2.59 against a ceiling of 4.0. The acceptance gate exists precisely + because boot and accuracy are blind to a miscomputed draft path; an + acceptance gate measuring a *different depth* than the unit was built for + is blind in a new way, and it fails toward "green" because a shallower + speculator accepts a larger fraction of what it drafts. + + So the YAML is the single source of truth and the flag is derived from + it. A unit with no options file, or none declaring a speculator, keeps + the old answer of 1. + + Read in the order the engine resolves it. ``max_draft_len`` is the field; + ``num_nextn_predict_layers`` is its deprecated MTP alias, which + ``MTPDecodingConfig`` remaps and drops when both are present; ``block_size`` + is how a DSpark drafter states its depth, and ``TorchLlmArgs`` rejects any + DSpark config whose ``block_size`` differs from ``max_draft_len``. For + DSpark this is not a knob at all: quickstart builds the config from the + flag alone, so the old default of 1 against a drafter trained for a + 5-token block fails that validation before the engine starts. + """ + options = resolve_llm_api_options(accuracy, unit) + if not options: + return 1 + source = Path(options) + if not source.is_file(): + return 1 + data = yaml.safe_load(source.read_text(encoding="utf-8")) or {} + spec = data.get("speculative_config") or {} + for key in ("max_draft_len", "num_nextn_predict_layers", "block_size"): + try: + return max(1, int(spec[key])) + except (KeyError, TypeError, ValueError): + continue + return 1 + + +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 still means "the whole run", so nothing that already + works changes shape. + + ``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 resolve_llm_api_options(accuracy: Mapping[str, Any], unit: str | None) -> str | None: + """Pick the runtime-configuration YAML for ``unit``. + + ``llm_api_options`` may be a single path (the whole run) or a mapping + from unit name to path. The mapping form exists because **a feature unit + is not measured under the same configuration as core**: MTP needs a + ``speculative_config``, and without one the accuracy gate silently + re-measures the base path while the acceptance gate compares 1.0 against + 1.0 and calls it a pass. + + That was live in this run: neither gate carried a speculative config at + all, so the `mtp` unit's gates were measuring `core`. + """ + options = accuracy.get("llm_api_options") + if isinstance(options, dict): + return options.get(unit) or options.get("default") + return options + + +def resolve_discriminators( + accuracy: Mapping[str, Any], unit: str | None +) -> Mapping[str, Any] | None: + """Pick the accuracy-reference discriminators for ``unit``. + + TensorRT-LLM's accuracy references key each model to a LIST of entries + discriminated by ``quant_algo`` / ``kv_cache_quant_algo`` / + ``spec_dec_algo``, and :func:`load_reference_score` matches a row only + when the discriminator set is EXACTLY the one asked for. Passing none -- + which is what this workflow did until now -- therefore means one of two + things, both wrong: + + * the model has an undiscriminated row, and a quantized checkpoint + silently anchors to the bf16 number; + * the model has none, and the gate dies with ``no entry matching {}`` + at the end of an assembly rather than at spec-validation time. + + Per-unit for the same reason ``llm_api_options`` is: the ``mtp`` unit's + anchor is the ``spec_dec_algo: MTP`` row, and measuring it against + core's row compares two different questions. + + Accepts either a flat mapping (the whole run) or a mapping from unit + name to mapping, with ``default`` as the fallback key. + """ + block = accuracy.get("discriminators") + if not isinstance(block, dict): + return None + # Per-unit form is recognised by its values being mappings; a flat form + # like {"quant_algo": "NVFP4"} has scalar values. + if any(isinstance(v, dict) for v in block.values()): + chosen = block.get(unit) + if chosen is None: + chosen = block.get("default") + return chosen if isinstance(chosen, dict) else None + return block + + +def _write_member_options( + options: str | None, + overlay: Mapping[str, Any], + *, + unit: str | None, + member: str | None, +) -> str: + """Merge a member's own LLM-API settings over its unit's options file. + + Written beside the unit's file so the two are read together and a reader + finding the derived one can see what it was derived from. Deterministic + name, so a re-run overwrites rather than accumulating, and the file a + verdict was measured under is still on disk afterwards. + """ + base: dict[str, Any] = {} + source = Path(options) if options else None + if source is not None and source.is_file(): + base = yaml.safe_load(source.read_text(encoding="utf-8")) or {} + merged = {**base, **overlay} + stem = f"{unit or 'unit'}.{member or 'member'}.llm_api.yaml" + # Beside the unit's own options file, so the derived copy and the file it + # was derived from are read together. With no unit file there is nowhere + # natural, and the cwd is the one place it must NOT go: the cwd of a test + # run is the repository, and this wrote `core.dep1.llm_api.yaml` into it. + # A generated artifact that lands in the source tree is indistinguishable + # from a file somebody meant to check in. + parent = source.parent if source is not None else Path(tempfile.gettempdir()) + target = parent / stem + header = ( + f"# GENERATED for the `{member}` member of the `{unit}` unit. Do not edit.\n" + f"# {source if source is not None else '(no unit options file)'} plus the\n" + f"# settings that member's gate command cannot express as flags:\n" + f"# {', '.join(sorted(overlay))}\n" + f"# See `gates.member_llm_api_overlay` for why each one is here.\n" + ) + target.write_text(header + yaml.safe_dump(merged, sort_keys=True), encoding="utf-8") + return str(target) + + +def eval_command( + task_data: Mapping[str, Any], + environment: RunEnvironment, + benchmark: str, + *, + unit: str | None = None, + member: str | None = None, +) -> list[str]: + """Build the ``trtllm-eval`` argv shared by the accuracy and acceptance gates. + + **`trtllm-eval` and `quickstart_advanced.py` take different flag + vocabularies for the same runtime configuration.** There is no + `--use_cuda_graph`, no `--kv_cache_dtype`, no `--mamba_ssm_cache_dtype` + here; that configuration has to arrive as an `--extra_llm_api_options` + YAML instead. So the `runtime` block cannot simply be forwarded — and + before this existed it was not forwarded at all. + + The consequence was not a weak comparison but a different question: the + accuracy gate measured the engine's defaults (including + `max_seq_len=262144`, the checkpoint's own maximum) while the boot gate + measured the requested configuration and the anchor had been measured at + a third thing again. Caught on the first real accuracy gate of this run, + by reading the command the gate had actually issued. + + The fix is to declare the configuration once, as a file, and have the + anchor run and both gates point at the same one. + """ + accuracy = dict(task_data.get("accuracy") or {}) + parts = [ + "trtllm-eval", + f'--model="{task_data["checkpoint_path"]}"', + *topology_flags( + str(task_data["parallel"]), + environment.world_size, + member, + dialect=EVAL_DIALECT, + ), + ] + options = resolve_llm_api_options(accuracy, unit) + overlay = member_llm_api_overlay(member) + if overlay: + # `trtllm-eval` has no attention-DP flag, so a `dep` member can only + # be served through the options YAML. Writing a merged copy rather + # than editing the unit's file keeps one source of truth: the unit + # states its runtime configuration once, and the member adds only + # what it alone needs. A `dep` gate run without this serves `tep` + # while the verdict says `dep`, which is the failure the member set + # exists to prevent. + options = _write_member_options(options, overlay, unit=unit, member=member) + if options: + parts.append(f'--extra_llm_api_options="{options}"') + for flag in ("max_seq_len", "max_batch_size"): + if accuracy.get(flag) is not None: + parts.append(f"--{flag}={accuracy[flag]}") + parts.append(benchmark) + if accuracy.get("dataset_path"): + parts.append(f'--dataset_path="{accuracy["dataset_path"]}"') + # **The protocol lives here, after the benchmark name.** + # `--apply_chat_template`, `--fewshot_as_multiturn`, + # `--chat_template_kwargs`, `--num_fewshot` are options of the benchmark + # subcommand, not of `trtllm-eval` itself, so they cannot go in the YAML + # and they cannot go before the positional. + # + # This is the difference the run measured: the same target scored 63.50 + # under the default 5-shot-completion protocol and the anchor was 95.53 + # under the chat-template protocol. Neither number is wrong; they answer + # different questions, and a gate that compares them reports a 27-point + # regression that does not exist. + # Shell-quoted: the command is assembled into a string and run through a + # shell, and a protocol argument like + # `--chat_template_kwargs={"enable_thinking": false}` carries spaces and + # braces that the shell would otherwise split and glob. `shlex.quote` + # leaves ordinary flags untouched. + # Per unit, because the PROTOCOL flags are options of the BENCHMARK + # subcommand and the subcommands do not share an option set. `gsm8k` + # takes `--apply_chat_template` and `--fewshot_as_multiturn`; `mmmu` + # takes neither and click rejects the whole command with "no such + # option". A run-wide list is therefore not merely imprecise for a + # multimodal unit -- it cannot run at all. + parts.extend( + shlex.quote(str(arg)) for arg in resolve_per_unit(accuracy, "benchmark_args", unit) or () + ) + return parts + + +def _with_mode(environment: RunEnvironment, mode: str) -> RunEnvironment: + """Return a copy of ``environment`` with the modeling_v2 switch set. + + Exported into the *prefix*, not into the orchestrator's own process: + worker ranks receive the environment as it stood when MPI initialized, + so a value set after launch reaches the driver and not them — and a + driver resolving a modeling_v2 target while its workers resolve the + built-in is exactly the silent split the switch exists to prevent. + """ + from dataclasses import replace + + return replace( + environment, + env={**environment.env, MODELING_V2_ENV: mode, RUNTIME_VALUE_DUMP_ENV: "1"}, + ) + + +def parse_eval_score(stdout: str) -> float | None: + """Extract the final accuracy from a ``trtllm-eval`` run's output. + + Both reporting paths are matched because the benchmark decides which + one runs: MMLU prints its own weighted average, everything else comes + through lm-eval. The **last** match wins — a long run also prints + partial-score estimates, and the final line is the one that counts. + """ + for pattern in (_MMLU_SCORE, _LM_EVAL_SCORE): + matches = pattern.findall(stdout) + if matches: + return float(matches[-1]) + return None + + +@dataclass(frozen=True) +class AcceptanceReading: + """One leg's acceptance measurement, with the requests it came from.""" + + #: The gate's number: iteration-weighted over the requests where they are + #: known (see :func:`parse_acceptance`), else the summary statistic printed. + value: float + #: ``(tokens, iterations)`` per request, in print order. Empty when the + #: output carried only a summary statistic, which says nothing about how + #: much work it was estimated from. + requests: tuple[tuple[int, int], ...] = () + + @property + def draft_iterations(self) -> int | None: + """Generation iterations over the whole workload; ``None`` if unknown. + + The first streamed iteration of every request is its context step: + it returns the prefill's token and verifies no draft, so it is not + counted. + """ + if not self.requests: + return None + return sum(max(0, iterations - 1) for _, iterations in self.requests) + + def breakdown(self) -> str: + """Per-prompt rates and their plain mean, beside the weighted value.""" + rates = [(n, t, i) for n, (t, i) in enumerate(self.requests) if i] + if not rates: + return "" + per_prompt = ", ".join(f"[{n}] {t / i:.2f} ({t} tok/{i} it)" for n, t, i in rates) + mean = sum(t / i for _, t, i in rates) / len(rates) + return ( + f"per prompt {per_prompt}; prompt mean {mean:.2f}, " + f"iteration-weighted {self.value:.2f} over {self.draft_iterations} " + f"generation iterations" + ) + + +def parse_acceptance(stdout: str) -> AcceptanceReading | None: + """Extract the acceptance measurement from an eval, bench or quickstart run. + + The quickstart form is **iteration-weighted across requests**: the + quantity is tokens-per-iteration, so several requests combine as + ``sum(tokens) / sum(iterations)`` and not as the mean of their printed + ratios. Averaging the ratios would let one short request with an + unlucky tail dominate a long one, and the whole judgement here is a + ratio between two runs. The plain per-prompt mean is still reported, in + :meth:`AcceptanceReading.breakdown`, because it is the number that shows + a prompt that never drafted. + """ + for pattern in (_ACCEPTANCE_LENGTH, _ACCEPTANCE_LENGTH_ALT): + matches = pattern.findall(stdout) + if matches: + return AcceptanceReading(value=float(matches[-1])) + parts = tuple((int(t), int(i)) for t, i in _ACCEPT_RATE_PARTS.findall(stdout)) + iterations = sum(i for _, i in parts) + if iterations: + return AcceptanceReading(value=sum(t for t, _ in parts) / iterations, requests=parts) + return None + + +def parse_acceptance_length(stdout: str) -> float | None: + """Extract ``acceptance_length`` from an eval, bench or quickstart run. + + The number alone; see :func:`parse_acceptance` for how it is combined. + """ + reading = parse_acceptance(stdout) + return reading.value if reading is not None else None + + +def load_reference_score( + reference_yaml: str | Path, + model_key: str, + *, + discriminators: Mapping[str, Any] | None = None, +) -> float: + """Read the anchor for one model from TensorRT-LLM's accuracy references. + + The file is keyed by HuggingFace repo name and holds a list of entries, + each optionally discriminated by ``quant_algo`` / ``kv_cache_quant_algo`` + / ``spec_dec_algo``. An entry matches when its discriminators are + exactly the ones asked for — no more, no less — so an FP8 anchor is + never silently used for a BF16 run. + + **The anchor is never invented.** If nothing matches, that is a question + for the plan phase to answer from a written source, not a number for + this function to estimate. + """ + path = Path(reference_yaml) + if not path.is_file(): + raise FileNotFoundError(f"accuracy reference file not found: {path}") + data = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + if model_key not in data: + raise KeyError( + f"{path} has no entry for {model_key!r}. The anchor comes from a written " + f"source or from the user — it is never invented. Add the entry (citing " + f"where the score came from) before the accuracy gate can run." + ) + + wanted = {k: v for k, v in (discriminators or {}).items() if v is not None} + entries = data[model_key] or [] + matches = [] + for entry in entries: + if not isinstance(entry, dict): + continue + present = {k: v for k, v in entry.items() if k != "accuracy"} + if present == wanted: + matches.append(entry) + + if not matches: + available = [ + {k: v for k, v in e.items() if k != "accuracy"} for e in entries if isinstance(e, dict) + ] + raise KeyError( + f"{path}:{model_key} has no entry matching {wanted or '{}'}. Available " + f"discriminator sets: {available}" + ) + if len(matches) > 1: + raise KeyError( + f"{path}:{model_key} has {len(matches)} entries matching {wanted or '{}'}; " + f"the anchor must be unambiguous" + ) + return float(matches[0]["accuracy"]) + + +def run_boot_gate( + environment: RunEnvironment, + *, + task_data: Mapping[str, Any], + spec: BootSpec | None = None, + runner: ShellRunner = _default_shell_runner, + timeout: int = BOOT_TIMEOUT_S, + mode: str = MODE_REQUIRE, + features: Any | None = None, + unit: str | None = None, + member: str | None = None, + allow_dirty: bool | None = None, +) -> GateResult: + """Cold-start the engine and check it generates. + + ``member`` is the parallel configuration to serve at, and it is the only + thing that makes a ``unified`` target's gate mean anything — see + :func:`topology_flags`. It also travels into the verdict, so the journal + line says which of the members this one was. + + ``unit`` selects the ``runtime`` block when the spec states one per unit. + It has to: `quickstart_advanced.py` takes no `--extra_llm_api_options`, so + this gate cannot read the per-unit options file the other two gates do, and + the runtime block is its only channel for a unit-specific switch. + + ``features`` is **the features the target under test actually + implements**, which is not the same as the features the *run* will + finish with. Defaulting it to ``task_data["features"]`` — as this + function originally did with no way to override — means the ``core`` + unit is gated with the speculator's flags before a speculator exists. + + Measured: the core target refused `--spec_decode_algo=MTP` through its + own fail-fast assert, correctly, and the assembler escalated rather than + deleting the assert to make the gate pass. Deleting it would have been + the silent wrong — one-engine MTP is drafted by the model, so a base + model under that config emits no drafts and the run would have gated a + speculator that was not there. + + Staircase's rule is that core is green *before* the first feature + increment; the gate has to agree with it. + + ``mode`` defaults to ``require`` and should stay there for any result + that will be attributed to modeling_v2. It is a parameter for one + reason: when a boot gate fails, the next question is always *"is this + the target or is this the host?"*, and the cheapest answer is the same + command under ``off`` — the stock implementation, same prompts, same + runtime flags, one variable. Without this the control has to be + reconstructed by hand, which is how a run ends up debugging a target + that was never at fault. + + ``run_accuracy_gate`` already took a ``mode`` for the same reason; the + boot gate not having one was an inconsistency, not a decision. + + ``allow_dirty`` is the escape hatch of :func:`inspect_tree`; see there + for why a dirty target path stops this gate by default. + """ + spec = spec or BootSpec() + env = _with_mode(environment, mode) + # A control run must not be recorded under the same name as the real + # gate. Two rows both called "boot", one of which measured the stock + # implementation, is a journal that reads as evidence for a target it + # never exercised. + gate_name = "boot" if mode == MODE_REQUIRE else f"boot[{mode}]" + state = inspect_tree(task_data, allow_dirty=allow_dirty) + _verdict = _verdict_factory(state) + if state.blocked: + return _verdict(name=gate_name, member=member, passed=False, detail=state.block_detail()) + repo = Path(str(task_data["trtllm_repo_path"])) + script = repo / "examples" / "llm-api" / "quickstart_advanced.py" + prompts = " ".join(f'"{p}"' for p in spec.prompts) + argv = " ".join( + [ + str(script), + f'--model_dir="{task_data["checkpoint_path"]}"', + f"--max_tokens={spec.max_tokens}", + *topology_flags(str(task_data["parallel"]), environment.world_size, member), + *runtime_flags(resolve_runtime(task_data, unit)), + *speculation_flags( + task_data.get("features") if features is None else features, + max_draft_len=resolve_max_draft_len(task_data.get("accuracy") or {}, unit), + ), + f"--prompt {prompts}", + ] + ) + command = env.python_command(argv) + + started = time.monotonic() + result = runner(command, None, timeout) + elapsed = time.monotonic() - started + tail = "\n".join((result.stdout + result.stderr).splitlines()[-40:]) + + if result.timed_out: + return _verdict( + name=gate_name, + member=member, + passed=False, + detail=f"timed out after {timeout}s — a boot gate that takes this long is " + f"not catching catastrophes any more", + command=command, + duration_s=elapsed, + stdout_tail=tail, + ) + if not result.ok: + return _verdict( + name=gate_name, + member=member, + passed=False, + detail=f"exited {result.returncode}", + command=command, + duration_s=elapsed, + stdout_tail=tail, + ) + + missing = [k for k in spec.keywords if k not in result.stdout] + if missing: + return _verdict( + name=gate_name, + member=member, + passed=False, + detail=f"generated output is missing {len(missing)} expected keyword(s): {missing}", + command=command, + duration_s=elapsed, + stdout_tail=tail, + ) + if not result.stdout.strip(): + return _verdict( + name=gate_name, + member=member, + passed=False, + detail="the engine exited cleanly but produced no output at all", + command=command, + duration_s=elapsed, + stdout_tail=tail, + ) + checked = f"{len(spec.keywords)} keyword(s)" if spec.keywords else "non-empty generation" + return _verdict( + name=gate_name, + member=member, + passed=True, + detail=f"engine cold-started and produced {checked}", + command=command, + duration_s=elapsed, + stdout_tail=tail, + ) + + +def run_accuracy_gate( + environment: RunEnvironment, + *, + task_data: Mapping[str, Any], + reference: float, + runner: ShellRunner = _default_shell_runner, + timeout: int = ACCURACY_TIMEOUT_S, + mode: str = MODE_REQUIRE, + unit: str | None = None, + member: str | None = None, + allow_dirty: bool | None = None, +) -> GateResult: + """Run ``trtllm-eval`` and compare against the anchor. One-sided. + + ``member`` names the parallel configuration measured, which for a + ``unified`` target is the whole question: one codebase serving three + configurations is three anchored numbers, not one repeated three times. + + This is the gate the tree-state guard was written for. It is the + expensive one — a quarter of an hour per member — and it is the one whose + output becomes the target's headline number, so it is also the one where + measuring an unreviewed tree is least recoverable. See + :func:`inspect_tree` and ``allow_dirty``. + """ + accuracy = dict(task_data.get("accuracy") or {}) + benchmark = str(resolve_per_unit(accuracy, "benchmark", unit) or "gsm8k") + tolerance = float(accuracy.get("tolerance", 0.0)) + env = _with_mode(environment, mode) + state = inspect_tree(task_data, allow_dirty=allow_dirty) + _verdict = _verdict_factory(state) + if state.blocked: + return _verdict( + name=f"accuracy[{benchmark}]", + member=member, + passed=False, + detail=state.block_detail(), + reference=reference, + tolerance=tolerance, + ) + + parts = eval_command(task_data, environment, benchmark, unit=unit, member=member) + command = env.command(" ".join(parts)) + + started = time.monotonic() + result = runner(command, None, timeout) + elapsed = time.monotonic() - started + combined = result.stdout + result.stderr + tail = "\n".join(combined.splitlines()[-40:]) + + if result.timed_out: + return _verdict( + name=f"accuracy[{benchmark}]", + member=member, + passed=False, + detail=f"timed out after {timeout}s", + command=command, + reference=reference, + tolerance=tolerance, + duration_s=elapsed, + stdout_tail=tail, + ) + if not result.ok: + return _verdict( + name=f"accuracy[{benchmark}]", + member=member, + passed=False, + detail=f"trtllm-eval exited {result.returncode}", + command=command, + reference=reference, + tolerance=tolerance, + duration_s=elapsed, + stdout_tail=tail, + ) + + measured = parse_eval_score(combined) + if measured is None: + return _verdict( + name=f"accuracy[{benchmark}]", + member=member, + passed=False, + detail="trtllm-eval exited 0 but printed no score this module recognizes; " + "a gate with no number is not a pass", + command=command, + reference=reference, + tolerance=tolerance, + duration_s=elapsed, + stdout_tail=tail, + ) + + floor = reference - tolerance + passed = measured >= floor + return _verdict( + name=f"accuracy[{benchmark}]", + member=member, + passed=passed, + detail=( + f"measured {measured:.4f} {'>=' if passed else '<'} {floor:.4f} " + f"(reference {reference:.4f} - tol {tolerance:.2f})" + ), + command=command, + measured=measured, + reference=reference, + tolerance=tolerance, + duration_s=elapsed, + stdout_tail=tail, + ) + + +def resolve_acceptance_floor(accuracy: Mapping[str, Any], unit: str | None) -> float | None: + """The absolute `acceptance_length` floor this unit declares, if any. + + Only consulted when the `=off` leg **cannot exist**. See + :func:`run_acceptance_gate`. + """ + block = accuracy.get("acceptance") + if not isinstance(block, dict): + return None + chosen = block.get(unit) + if chosen is None: + chosen = block.get("default") + if isinstance(chosen, dict): + chosen = chosen.get("min_absolute") + try: + return float(chosen) + except (TypeError, ValueError): + return None + + +@dataclass(frozen=True) +class AcceptanceWorkload: + """What the acceptance gate sends, and how much work it must get back.""" + + prompts: tuple[str, ...] + max_tokens: int + min_draft_iterations: int + #: Where the prompts came from, named in a verdict that blames them. + source: str + + +def resolve_acceptance_workload( + accuracy: Mapping[str, Any], + unit: str | None, + *, + fallback: BootSpec | None = None, + max_tokens: int | None = None, +) -> AcceptanceWorkload: + """The acceptance workload ``unit`` declares, or the fallback when none does. + + Read from ``accuracy.acceptance``, most specific level first: the unit's + own entry, then its ``default`` entry, then the block's top level (see + ``task_schema.ACCEPTANCE_WORKLOAD_KEYS``). ``prompts`` and + ``prompts_file`` are one setting, so the level that states either wins + for both. ``fallback`` (the boot spec) and ``max_tokens`` are what a task + that declares nothing measures, which is what every task measured before + the workload could be declared. + + Raises :class:`~.task_schema.TaskSchemaError` when a declared prompts + file cannot be read. + """ + block = accuracy.get("acceptance") + block = block if isinstance(block, dict) else {} + levels = [ + (f"accuracy.acceptance.{name}", block.get(name)) + for name in (unit, "default") + if name is not None + ] + levels.append(("accuracy.acceptance", block)) + levels = [(where, entry) for where, entry in levels if isinstance(entry, dict)] + + def setting(key: str) -> Any: + for _, entry in levels: + if entry.get(key) is not None: + return entry[key] + return None + + prompts: tuple[str, ...] | None = None + source = "" + for where, entry in levels: + if entry.get("prompts") is not None: + prompts = tuple(str(p) for p in entry["prompts"]) + source = f"{where}.prompts" + break + if entry.get("prompts_file") is not None: + prompts = load_prompts_file(entry["prompts_file"]) + source = f"{where}.prompts_file ({entry['prompts_file']})" + break + if prompts is None: + prompts = (fallback or BootSpec()).prompts + source = "the boot prompts (no `accuracy.acceptance.prompts` declared)" + declared_tokens = setting("max_tokens") + declared_minimum = setting("min_draft_iterations") + if declared_tokens is None: + declared_tokens = max_tokens if max_tokens is not None else DEFAULT_ACCEPTANCE_MAX_TOKENS + if declared_minimum is None: + declared_minimum = DEFAULT_MIN_DRAFT_ITERATIONS + return AcceptanceWorkload( + prompts=prompts, + max_tokens=int(declared_tokens), + min_draft_iterations=int(declared_minimum), + source=source, + ) + + +def run_acceptance_gate( + environment: RunEnvironment, + *, + task_data: Mapping[str, Any], + runner: ShellRunner = _default_shell_runner, + timeout: int = BOOT_TIMEOUT_S, + min_ratio: float = 0.95, + unit: str | None = None, + spec: BootSpec | None = None, + features: Any | None = None, + max_tokens: int | None = None, + member: str | None = None, + allow_dirty: bool | None = None, +) -> GateResult: + """Compare ``acceptance_length`` against stock in-tree modeling. + + The two runs differ in exactly one variable — ``TRTLLM_MODELING_V2`` + ``off`` versus ``require`` — and must happen **in the same session**. + A cross-session delta carries session variance into the judgement, and + staircase measured a bit-identical forward at 94.7688 and 95.0720 a day + apart. + + ``member`` fixes the topology **both** legs run at, so the one variable + stays one: a draft path whose acceptance depends on the expert-parallel + slice is a thing this gate can only see when it is asked at that member. + + The tree-state guard matters here for a reason the other two gates do not + have: **both legs import the same working tree**, so an edit landing + between them moves the ratio without moving either leg's correctness, and + the gate reports it as a draft-path defect. See :func:`inspect_tree`. + + **The workload is the task's to declare** (``accuracy.acceptance.prompts`` + or ``prompts_file``, ``max_tokens``; see + :func:`resolve_acceptance_workload`). ``spec`` and ``max_tokens`` here are + only the fallback for a task that declares none. The boot prompts are + the wrong workload for a model whose answers to them are short: + DeepSeek-V4.1 answers them in 2-3 tokens and EOS, so every request is a + context step plus one generation step, the rate reads exactly 1.00 for + ANY drafter, and a declared floor above 1 can never pass while a broken + drafter cannot be told from a working one. On long-answer prompts the + same command measured 1.91 for a drafter with a broken ring index and + 3.24 after the fix. + + **A workload too short to measure is a failure of its own**, reported as + one: fewer than ``min_draft_iterations`` generation iterations in total + (from quickstart's ``Accept rate: ... (tokens=T, iterations=I)`` lines) + fails with "workload too short to measure acceptance" instead of a low + number that reads as a drafter defect. A leg that prints only a summary + statistic carries no iteration count and is not held to it. + + Every verdict that has a reading carries the per-prompt rates and their + plain mean beside the iteration-weighted value it is judged on, in + :attr:`GateResult.breakdown`. + """ + state = inspect_tree(task_data, allow_dirty=allow_dirty) + _verdict = _verdict_factory(state) + if state.blocked: + return _verdict(name="acceptance", member=member, passed=False, detail=state.block_detail()) + accuracy = task_data.get("accuracy") or {} + try: + workload = resolve_acceptance_workload( + accuracy, unit, fallback=spec or BootSpec(), max_tokens=max_tokens + ) + except TaskSchemaError as exc: + return _verdict( + name="acceptance", + member=member, + passed=False, + detail=f"the declared acceptance workload cannot be read: {exc}", + ) + # The QUICKSTART command, not the eval command, and not as an + # optimization. `trtllm-eval` emits no speculative statistic in this + # tree, so an acceptance gate built on it reads None from both legs and + # fails identically to a broken draft path. `quickstart_advanced.py` + # prints `Accept rate: N (tokens=T, iterations=I)` per request, which is + # the quantity this gate is about. + # + # It is also two orders of magnitude cheaper -- minutes against the nine + # a full gsm8k run costs, twice -- which matters because this gate runs + # both legs. + repo = Path(str(task_data["trtllm_repo_path"])) + script = repo / "examples" / "llm-api" / "quickstart_advanced.py" + # Shell-quoted: a declared workload is free text, and a code-listing + # prompt carries the quotes, `$` and backslashes a double-quoted word + # would expand or end on. + prompts = " ".join(shlex.quote(p) for p in workload.prompts) + argv = " ".join( + [ + str(script), + f'--model_dir="{task_data["checkpoint_path"]}"', + # Longer than the boot gate's, deliberately. Acceptance length is + # a rate, and a 32-token continuation gives it a handful of + # iterations to be estimated from -- enough to see a draft path + # collapse to 1.0, not enough to resolve a 5% shortfall. + f"--max_tokens={workload.max_tokens}", + *topology_flags(str(task_data["parallel"]), environment.world_size, member), + *runtime_flags(resolve_runtime(task_data, unit)), + *speculation_flags( + task_data.get("features") if features is None else features, + max_draft_len=resolve_max_draft_len(accuracy, unit), + ), + f"--prompt {prompts}", + ] + ) + + started = time.monotonic() + readings: dict[str, AcceptanceReading | None] = {} + commands: dict[str, str] = {} + stock_unavailable = False + for mode in (MODE_REQUIRE, MODE_OFF): + command = _with_mode(environment, mode).python_command(argv) + commands[mode] = command + result = runner(command, None, timeout) + combined = result.stdout + result.stderr + if not result.ok: + # **The stock leg can be impossible rather than broken.** This gate + # compares the target against the same checkpoint under built-in + # modeling -- and a genuinely new architecture HAS no built-in + # modeling, so `=off` cannot load it at all. That is the same shape + # as the substrate gate's problem (see + # `preflight.UNSUPPORTED_ARCHITECTURE_MARKERS`, reused here): the + # check is unavailable exactly when onboarding is doing its job. + # + # Failing the unit for it would be wrong, and silently skipping the + # gate would be worse -- it is the ONLY gate that can see a + # miscomputed draft path, because rejection sampling makes a wrong + # draft produce correct text more slowly. So the spec may declare an + # ABSOLUTE floor instead, and the weaker claim is recorded verbatim. + from .preflight import reports_unsupported_architecture + + floor = resolve_acceptance_floor(accuracy, unit) + unavailable = mode == MODE_OFF and reports_unsupported_architecture( + f"exited {result.returncode}", "\n".join(combined.splitlines()[-200:]) + ) + if unavailable and floor is not None: + readings[mode] = None + stock_unavailable = True + break + return _verdict( + name="acceptance", + member=member, + passed=False, + detail=( + f"the {MODELING_V2_ENV}={mode} leg " + + ("timed out" if result.timed_out else f"exited {result.returncode}") + + ( + " -- this checkpoint has no built-in implementation, so the " + "stock leg cannot run. Declare `accuracy.acceptance.." + "min_absolute` to gate on an absolute floor instead." + if unavailable + else "" + ) + ), + command=command, + duration_s=time.monotonic() - started, + stdout_tail="\n".join(combined.splitlines()[-40:]), + ) + readings[mode] = parse_acceptance(combined) + elapsed = time.monotonic() - started + + target_reading = readings.get(MODE_REQUIRE) + stock_reading = readings.get(MODE_OFF) + target = target_reading.value if target_reading is not None else None + stock = stock_reading.value if stock_reading is not None else None + breakdown = target_reading.breakdown() if target_reading is not None else "" + + def too_short(mode: str) -> GateResult | None: + """The verdict for a leg whose workload never exercised the drafter.""" + reading = readings.get(mode) + counted = reading.draft_iterations if reading is not None else None + if reading is None or counted is None or counted >= workload.min_draft_iterations: + return None + return _verdict( + name="acceptance", + member=member, + passed=False, + detail=( + f"workload too short to measure acceptance: the {MODELING_V2_ENV}={mode} " + f"leg ran {counted} generation iteration(s) over {len(reading.requests)} " + f"prompt(s), fewer than the {workload.min_draft_iterations} this gate " + f"needs, so its rate ({reading.value:.2f}) says nothing about the draft " + f"path -- the answers end before drafts are verified. Workload: " + f"{workload.source}, max_tokens {workload.max_tokens}. Declare " + f"long-answer prompts in `accuracy.acceptance.prompts` (or " + f"`prompts_file`); this is a property of the workload, not of the target." + ), + command=commands[mode], + duration_s=elapsed, + breakdown=reading.breakdown(), + ) + + if stock_unavailable: + floor = resolve_acceptance_floor(accuracy, unit) + if target is None: + return _verdict( + name="acceptance", + member=member, + passed=False, + detail="could not read acceptance_length from the require leg", + command=commands[MODE_REQUIRE], + duration_s=elapsed, + ) + short = too_short(MODE_REQUIRE) + if short is not None: + return short + return _verdict( + name="acceptance", + member=member, + passed=target >= float(floor), + detail=( + f"acceptance_length {target:.4f} vs declared floor {float(floor):.4f}. " + f"The {MODELING_V2_ENV}=off leg could not run: this checkpoint has no " + f"built-in implementation, so there is no stock acceptance to compare " + f"against. THIS IS THE WEAKER CLAIM -- it shows the draft path accepts " + f"at a real rate, not that it matches stock." + ), + command=commands[MODE_REQUIRE], + measured=target, + reference=float(floor), + duration_s=elapsed, + breakdown=breakdown, + ) + if target is None or stock is None: + return _verdict( + name="acceptance", + member=member, + passed=False, + detail=f"could not read acceptance_length from both legs " + f"(require={target}, off={stock})", + command=commands[MODE_REQUIRE], + duration_s=elapsed, + ) + for mode in (MODE_REQUIRE, MODE_OFF): + short = too_short(mode) + if short is not None: + return short + stock_breakdown = stock_reading.breakdown() if stock_reading is not None else "" + if breakdown and stock_breakdown: + breakdown = f"require: {breakdown}; off: {stock_breakdown}" + if target <= 1.0: + return _verdict( + name="acceptance", + member=member, + passed=False, + detail=f"acceptance_length {target:.4f} at or below 1.0 means every draft " + f"was rejected — the draft path is computing the wrong thing", + command=commands[MODE_REQUIRE], + measured=target, + reference=stock, + duration_s=elapsed, + breakdown=breakdown, + ) + + ratio = target / stock if stock else 0.0 + passed = ratio >= min_ratio + return _verdict( + name="acceptance", + member=member, + passed=passed, + detail=( + f"acceptance_length {target:.4f} vs stock {stock:.4f} " + f"(ratio {ratio:.3f}, floor {min_ratio:.2f}). Clearly above 1.0 but below " + f"the reference means subtly wrong, not merely slower." + ), + command=commands[MODE_REQUIRE], + measured=target, + reference=stock, + duration_s=elapsed, + breakdown=breakdown, + ) + + +@dataclass +class GateLog: + """Every gate a run has fired, in order. The journal's raw material.""" + + results: list[GateResult] = field(default_factory=list) + + def record(self, result: GateResult) -> GateResult: + self.results.append(result) + return result + + def last(self, name_prefix: str) -> GateResult | None: + for result in reversed(self.results): + if result.name.startswith(name_prefix): + return result + return None + + def render(self) -> str: + return "\n".join(result.summary() for result in self.results) + + +def resolve_reference( + task_data: Mapping[str, Any], + *, + model_key: str, + discriminators: Mapping[str, Any] | None = None, + unit: str | None = None, +) -> float: + """Look up this run's anchor, resolving the reference path against the repo. + + ``discriminators`` overrides the spec; otherwise they come from + ``accuracy.discriminators`` for ``unit``. A caller that passes neither + gets the undiscriminated row, which for a quantized checkpoint is the + wrong anchor -- see :func:`resolve_discriminators`. + """ + accuracy = dict(task_data.get("accuracy") or {}) + reference_yaml = Path(str(resolve_per_unit(accuracy, "reference_yaml", unit))) + if not reference_yaml.is_absolute(): + reference_yaml = Path(str(task_data["trtllm_repo_path"])) / reference_yaml + if discriminators is None: + discriminators = resolve_discriminators(accuracy, unit) + return load_reference_score(reference_yaml, model_key, discriminators=discriminators) + + +__all__ = [ + "ACCURACY_TIMEOUT_S", + "ALLOW_DIRTY_ENV", + "AcceptanceReading", + "AcceptanceWorkload", + "BOOT_TIMEOUT_S", + "BootSpec", + "DEFAULT_ACCEPTANCE_MAX_TOKENS", + "DEFAULT_MIN_DRAFT_ITERATIONS", + "GateLog", + "GateResult", + "MODELING_V2_ENV", + "MODE_OFF", + "MODE_REQUIRE", + "RUNTIME_VALUE_DUMP_ENV", + "ShellResult", + "ShellRunner", + "TreeState", + "UNKNOWN_TREE_STATE", + "describe_tree_state", + "dirty_target_files", + "inspect_tree", + "load_reference_score", + "order_members_by_cost", + "parse_acceptance", + "parse_acceptance_length", + "parse_eval_score", + "resolve_acceptance_floor", + "resolve_acceptance_workload", + "resolve_discriminators", + "resolve_per_unit", + "resolve_runtime", + "resolve_llm_api_options", + "resolve_reference", + "run_acceptance_gate", + "run_accuracy_gate", + "run_boot_gate", + "shell_result_from", + "spool_runner", + "target_family_rel", + "topology_flags", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/journal.py b/agent-flow/agent_flow/workflows/staircase_onboard/journal.py new file mode 100644 index 000000000000..6ecdbbe0a033 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/journal.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The run journal — recovery state, written as events happen. + +Staircase's rule, kept: **the journal is written as events happen, never +retrospectively.** After a session is cut, the journal is the only +explanation of what is on disk. A journal filled in at the end is a summary, +and a summary cannot tell you which of two half-finished things was in +flight. + +It lives in the workspace, not in the repository. An onboarding run's +evidence is real and worth keeping, but it is evidence *about a run*, and a +product tree is the wrong home for it — so the journal, the per-dispatch +reports and the plan artifacts stay beside the checkpoint state, and only a +distilled record travels into the repo at delivery. + +Timestamps are UTC, read from the clock. Never estimate one from how long +something felt. +""" + +from __future__ import annotations + +import re +import threading +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path + +_REPORT_NAME = re.compile(r"^(\d+)-") + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat(timespec="seconds") + + +def _slugify(text: str) -> str: + """Fold arbitrary text into a safe, readable filename component.""" + slug = re.sub(r"[^a-zA-Z0-9]+", "-", text).strip("-").lower() + return slug[:60] or "unnamed" + + +@dataclass +class Journal: + """Append-only run log plus the per-dispatch report directory.""" + + path: Path + lock: threading.Lock = field(default_factory=threading.Lock, repr=False, compare=False) + + @property + def reports_dir(self) -> Path: + return self.path.parent / "reports" + + def open(self, *, slug: str, header_lines: list[str] | None = None) -> None: + """Start the journal, or continue an existing one. + + Continuing is the common case: re-running the same invocation is how + a run resumes, and the second run's events belong in the same file + as the first's. + """ + self.path.parent.mkdir(parents=True, exist_ok=True) + self.reports_dir.mkdir(parents=True, exist_ok=True) + if self.path.exists() and self.path.read_text(encoding="utf-8").strip(): + self.event("run resumed", f"slug={slug}") + return + lines = [f"# Onboarding journal — {slug}", ""] + lines.extend(header_lines or []) + lines.append("") + self.path.write_text("\n".join(lines) + "\n", encoding="utf-8") + self.event("run opened", f"slug={slug}") + + def event(self, what: str, detail: str = "") -> None: + """Append one timestamped line. The unit of recovery state.""" + line = f"- `{_utc_now()}` **{what}**" + if detail: + line += f" — {detail}" + with self.lock: + self.path.parent.mkdir(parents=True, exist_ok=True) + with self.path.open("a", encoding="utf-8") as handle: + handle.write(line + "\n") + + def section(self, title: str, body: str) -> None: + """Append a titled block — a preflight report, a gate summary, a verdict.""" + with self.lock: + self.path.parent.mkdir(parents=True, exist_ok=True) + with self.path.open("a", encoding="utf-8") as handle: + handle.write(f"\n## {_utc_now()} — {title}\n\n{body.rstrip()}\n") + + def next_ordinal(self) -> int: + """The next dispatch number, counted across the whole run. + + Sequential across every agent, so ``NN`` names exactly one dispatch + and the next ordinal is visible without counting. + """ + if not self.reports_dir.is_dir(): + return 1 + seen = [ + int(match.group(1)) + for entry in self.reports_dir.iterdir() + if (match := _REPORT_NAME.match(entry.name)) + ] + return max(seen, default=0) + 1 + + def record_dispatch(self, agent: str, what: str, prompt: str) -> Path: + """Write a dispatch prompt down **before it is sent**, and reserve its ordinal. + + An agent's context does not survive a session cut, so a child's + prompt only survives if its parent wrote it first. This is also why + the file is created here rather than when the agent returns: a run + that dies mid-dispatch still explains what it was doing. + """ + ordinal = self.next_ordinal() + path = self.reports_dir / f"{ordinal:02d}-{agent}-{_slugify(what)}.md" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + f"# {ordinal:02d} · {agent} · {what}\n\n" + f"- dispatched: `{_utc_now()}`\n\n" + f"## Dispatch prompt\n\n```\n{prompt}\n```\n\n" + f"## Result\n\n_pending_\n", + encoding="utf-8", + ) + self.event(f"dispatch {ordinal:02d} · {agent}", what) + return path + + def latest_report(self, agent: str, what: str) -> Path | None: + """The newest report :meth:`record_dispatch` wrote for this agent and subject. + + Derived from the file name rather than remembered, so a process that + never made the dispatch -- a resume, or a retry owed from an earlier + phase -- can still point the next attempt at what the last one did. + """ + if not self.reports_dir.is_dir(): + return None + name = f"{agent}-{_slugify(what)}.md" + found = [ + (int(match.group(1)), entry) + for entry in self.reports_dir.iterdir() + if (match := _REPORT_NAME.match(entry.name)) and entry.name[match.end() :] == name + ] + return max(found)[1] if found else None + + def record_result(self, report_path: Path, result: str) -> None: + """Replace a dispatch report's pending result with the agent's final message.""" + text = report_path.read_text(encoding="utf-8") + text = text.replace( + "## Result\n\n_pending_\n", + f"## Result\n\n- returned: `{_utc_now()}`\n\n{result.rstrip()}\n", + ) + report_path.write_text(text, encoding="utf-8") + + def close(self, *, summary: str) -> None: + self.section("run closed", summary) + + +__all__ = ["Journal"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/kernel_sources.yaml b/agent-flow/agent_flow/workflows/staircase_onboard/kernel_sources.yaml new file mode 100644 index 000000000000..09dbad885f94 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/kernel_sources.yaml @@ -0,0 +1,233 @@ +# Where kernels come from, and what that costs to certify. +# +# THE PROBLEM THIS SOLVES +# +# A smith asked to supply a computation needs three answers: does a kernel +# already exist, where did it come from, and what does certifying it require. +# The first is answered dynamically --- `survey.py` enumerates +# `torch.ops.trtllm.*` on the real host, because an op's PRESENCE depends on +# the build (FlashInfer ops register only behind `IS_FLASHINFER_AVAILABLE`). +# This file answers the other two. +# +# Provenance is not trivia. It decides the shape of the receipt: a CuTe DSL +# kernel is architecture-gated and JIT-compiled, a vendored Triton kernel is +# compiled on first call and has no arch gate, a C++ kernel is in the wheel +# and either present or not. Certifying them means different things. +# +# TWO REGISTRATION PLANES, and you have to know which one you are in: +# +# Python `tensorrt_llm/_torch/custom_ops/*.py` ~181 ops +# One file per source. The file name IS the provenance, which is +# what makes attribution mechanical --- grep the op name there. +# C++ `cpp/tensorrt_llm/thop/*.cpp` ~95 ops +# Registered via `def("...")` at load. Nothing in the Python tree +# mentions them, so an op found in neither custom_ops file is +# almost certainly here. +# +# Both land in the same `torch.ops.trtllm` namespace, so the namespace alone +# tells you nothing about provenance. That is the whole reason for this file. +# +# SCHEMA (enforced by tests/workflows/staircase_onboard/test_kernel_sources.py) +# +# : +# summary: what this source is, in a sentence or two +# registers: the file(s) where its ops enter `torch.ops.trtllm` +# implementation: where the kernel body lives, when that differs +# tier: which sourcing tier (see vocabulary.py) it satisfies +# availability: what decides whether its ops exist in a given build +# certifying: what a receipt for one of these must say, beyond passing +# +# Paths are resolved against the checkout by a test. No line numbers, no +# signatures --- same rule as the runtime map. + +# `fast_custom_op.py` is deliberately absent as a source: it is a low-overhead +# replacement for `@torch.library.custom_op`, not a place kernels come from. +# An op attributed to it has some other real provenance --- read the file that +# uses the decorator, not the decorator. + +cpp_thop: + summary: >- + Hand-written CUDA and C++ kernels compiled into the wheel and bound + through the `thop` layer. The largest source and the default assumption + when an op's provenance is not otherwise obvious. + registers: + - tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py + - cpp/tensorrt_llm/thop + implementation: + - cpp/tensorrt_llm/kernels + tier: 1 + availability: >- + In the wheel or not. Present on every build that compiled the + corresponding architecture; a missing one means the build excluded it, + not that the environment is misconfigured. + certifying: >- + Name the architecture. These are compiled per `--cuda_architectures`, so + a receipt is only meaningful with the arch stated, and a pass elsewhere + says nothing here. + +trtllm_gen: + summary: >- + Generated kernels --- the trtllm-gen family, including the FMHA and + block-scale MoE runners. Generated rather than written, so the tactic and + kernel selection happen at run time from a precompiled set. + registers: + - tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py + implementation: + - cpp/tensorrt_llm/kernels/trtllmGenKernels + tier: 1 + availability: >- + In the wheel, but the *selection* is dynamic: a generation FMHA kernel + takes the NVRTC path unless the KV dtype is E2m1, which is a run-time + branch and a known source of a JIT failure that looks like a build fault. + certifying: >- + Name which kernel of the family was selected, and the KV dtype. Two + architectures can compute bit-different results from one entry --- the + catalog already records a case where each refutes the other's recipe. + +cute_dsl: + summary: >- + Kernels written in CuTe DSL, NVIDIA's Python-embedded tile DSL. Compiled + at first use for the architecture in hand, so they are both source in the + tree and JIT artifacts at run time. + registers: + - tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py + - tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py + - tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py + implementation: + - tensorrt_llm/_torch/cute_dsl_kernels + tier: 1 + availability: >- + Architecture-gated, and the gate is per-kernel --- the implementation + tree is organised by architecture (`blackwell`, `rubin`). An unavailable + kernel usually degrades to another path rather than raising, sometimes + behind a broad except. + certifying: >- + State the architecture and confirm the CuTe path actually ran. A silent + fallback to a reference implementation produces a passing test that + certifies the wrong kernel. + +cuda_tile: + summary: >- + Kernels for the CUDA Tile / TileIR backend. A small source today --- norm + and fused residual-add norm --- but a distinct compilation path. + registers: + - tensorrt_llm/_torch/custom_ops/cuda_tile_custom_ops.py + implementation: + - tensorrt_llm/_torch/cuda_tile_kernels + tier: 1 + availability: >- + Requires the Tile toolchain; absent builds simply do not register the op. + certifying: >- + Confirm the tile path ran rather than the equivalent non-tile op, since + both exist for the same computation. + +flashinfer: + summary: >- + Kernels from the external FlashInfer library, wrapped as trtllm ops. The + only source here that is a third-party dependency rather than tree code, + and the only one whose ops can be absent from an otherwise complete build. + registers: + - tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py + tier: 1 + availability: >- + Conditional. The ops register only behind `IS_FLASHINFER_AVAILABLE`, so + presence is an environment property, not a build one. FlashInfer also + JIT-compiles into a workspace that is locked per rank --- see + `trtllm-llmapi-launch` for the isolation and `FLASHINFER_CUBIN_DIR` for + the shared cubin cache. + certifying: >- + Do not pin a FlashInfer version in the contract; name the kernel or the + code path instead. The tree's own audit rejects version pins for exactly + this reason. + +triton_custom_op: + summary: >- + Triton kernels written in this tree and registered through + `torch.library.custom_op`, so the kernel and its op live in the same + file under `custom_ops/`. Distinct from `vendored_triton`, which is + Triton the op plane cannot see: these ARE ops, and an inventory finds + them. + registers: + - tensorrt_llm/_torch/custom_ops/nvfp4_sfc_finalize.py + - tensorrt_llm/_torch/custom_ops/nvfp4_fake_quantize.py + - tensorrt_llm/_torch/custom_ops/mxfp8_fake_quantize.py + - tensorrt_llm/_torch/custom_ops/mxfp8_dequantize.py + - tensorrt_llm/_torch/custom_ops/mxfp8_to_e4m3.py + - tensorrt_llm/_torch/custom_ops/nvfp4_kv_to_bf16.py + - tensorrt_llm/_torch/custom_ops/nvfp4_kv_to_e4m3.py + - tensorrt_llm/_torch/custom_ops/e4m3_quantize_pow2_dynamic.py + tier: 2 + availability: >- + Never gated --- Triton is a hard dependency of the wheel, and each module + imports it at module scope rather than behind a try. But an op exists + only once its module is imported: `nvfp4_sfc_finalize` is imported by the + `custom_ops` package, the others by the catalog wrapper that calls them, + so an inventory that imports only the package does not list them. + certifying: >- + Read the autotune. A Triton op's numerics are fixed but its launch + configuration is chosen at runtime, so a receipt that names only the op + describes a kernel that may not be the one the next run picks. Record + the resolved config alongside the result. + +vendored_triton: + summary: >- + Triton kernels vendored into the tree and called directly as Python, not + through `torch.ops`. The linear-attention and SSM families live here. + Because they are ordinary callables, they are invisible to an op + inventory --- the kernel plane, not the op plane. + implementation: + - tensorrt_llm/_torch/modules/fla + - tensorrt_llm/_torch/modules/mamba + tier: 2 + availability: >- + Always importable, but which one runs can be chosen once per process by + environment and architecture --- the GDN prefill backend resolves through + a cached lookup on first call. + certifying: >- + Name the branch. Two backends for one computation are different kernels + in different gate spaces; the in-tree parity test compares them with a + tolerance, not bit-equality, so a receipt that does not say which one ran + certifies an implementation nobody can identify later. + +torch_composed: + summary: >- + Ops composed from ordinary torch operations and registered for the sake + of a stable name and a fake/meta implementation. No custom kernel. + registers: + - tensorrt_llm/_torch/custom_ops/torch_custom_ops.py + tier: 3 + availability: >- + Always present; nothing gates them. + certifying: >- + Cheap. Upstream owns the correctness of the underlying torch ops, which + is why the catalog's `torch/` mirrors carry no contract and no receipt. + +userbuffers: + summary: >- + Communication kernels using userbuffers for fused collective paths. A + parallelism-only source; a single-rank target never reaches it. + registers: + - tensorrt_llm/_torch/custom_ops/userbuffers_custom_ops.py + tier: 1 + availability: >- + Multi-rank only, and topology dependent: the buffers are established + across a communicator, so a single-rank build never registers a use for + them even when the ops are present. + certifying: >- + Certify at the world size the target declares. A collective receipt taken + at a different rank count describes a different communicator. + +dspark: + summary: >- + Fused DSpark attention and norm+RoPE kernels for the DSV4 family, built + on tcgen05 MMA. Model-family specific rather than general-purpose. + registers: + - tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py + - tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py + tier: 1 + availability: >- + Architecture-gated on the tcgen05 MMA path; a build or device without it + does not register these. + certifying: >- + Name the architecture and the family. A family-specific fused kernel is + not a drop-in for the general path even when the shapes agree. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/preflight.py b/agent-flow/agent_flow/workflows/staircase_onboard/preflight.py new file mode 100644 index 000000000000..56b1820f0451 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/preflight.py @@ -0,0 +1,1014 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Phase A — everything that must hold before an agent is worth starting. + +This is the module that decides whether the run happens at all, and it is +deliberately **pure Python with no agent turn**. The failure mode it exists +to prevent is a run that spends its plan phase and its first assembler turn +discovering that a card is wedged, ``TMPDIR`` points into a full ``/tmp``, +MPI cannot spawn on this host, or the anchor row its release criterion +names is not in the reference file. Those are seconds of checking here and +hours of confusion there. + +The asserts run **in order**, and the first failure ends the run before the +branch exists. Ordering is not stylistic: there is no point probing devices +on a host whose architecture does not match the one the receipts describe. + +Two of these checks encode measurements that are easy to get wrong: + +- **Judge a device by ``memory.free`` and ``utilization.gpu``, never by the + process table.** ``nvidia-smi --query-compute-apps`` under-attributes + badly — staircase measured 168 GiB in use against 2.7 GiB attributed — + and orphaned processes holding a dead job's allocations do not appear in + it at all. +- **A card at 100% utilization is *broken* for our purposes, not busy.** It + makes engines stall mid-gate with no error and invalidates every number + taken there. + +And one that separates two causes of a single symptom: +``cudaErrorDevicesUnavailable`` from a multi-rank preflight is *equally* +what a broken card returns and what a concurrent job returns. The +lone-process probe tells them apart: a card that fails a single-process, +one-tensor allocation with no peers and nothing else running is unusable +rather than contended, and no amount of waiting will free it. +""" + +from __future__ import annotations + +import json +import shlex +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Mapping, Sequence + +from agent_flow.runtime import CommandSpec, LocalCommandExecutor + +from .delivery import MODELING_V2_PACKAGE +from .environment import RunEnvironment, resolve_environment +from .task_schema import ( + assembly_units, + derive_checkpoint_segment, + derive_family, + derive_slug, + members_for, + parse_parallel_segment, +) + +# A device is claimable when it is this close to empty. Not 100%: a healthy +# idle card still shows a little in use (ECC tables, the driver's own +# context), and demanding exactly zero would reject every real machine. +MIN_FREE_MEMORY_FRACTION = 0.90 + +# What a `unified` segment claims when neither it nor its members say. The +# segment string carries no world size by design — the member set decides — +# and this is the number `staircase-modeling-parallel` documents for the +# baseline pair. +DEFAULT_UNIFIED_WORLD_SIZE = 4 + +# Above this, treat the card as broken rather than busy. See the module +# docstring: a spinning neighbour does not slow a gate down, it makes the +# gate's number meaningless. +MAX_UTILIZATION_PCT = 10 + +# The lone-process probe and the multi-rank substrate probe both run under +# a hard deadline, because the failures they look for do not raise — they +# hang. An MPI spawn whose connect-back lands on an interface that does not +# carry it blocks forever, and "the engine never booted" three hours later +# is the same observation with none of the information. +PROBE_TIMEOUT_S = 120 +MULTIRANK_PROBE_TIMEOUT_S = 300 + +# --------------------------------------------------------------------------- +# Why `TMPDIR` has a maximum length, and where each number comes from. +# +# The LLM API's RPC binds an IPC socket under `TMPDIR` named +# `rpc_test_` --- `get_unique_ipc_addr` in +# `tensorrt_llm/executor/rpc/rpc_common.py`, which is `os.path.join`ed onto +# `tempfile.gettempdir()`. A unix domain socket path is capped by the kernel +# at `sizeof(sockaddr_un.sun_path)`, and ZeroMQ reports the overflow as +# `ipc path "..." is longer than 107 characters`. +# +# So the budget is fixed on both sides and `TMPDIR` gets what is left. The +# numbers are named rather than reduced to a literal 61 because the reasoning +# is the part worth keeping: if the basename or the kernel limit ever moves, +# the derived bound moves with it and stays explicable. +# +# Measured directly rather than trusted: an `AF_UNIX` bind against synthetic +# paths succeeds at 107 characters and fails at 108 with `AF_UNIX path too +# long`, and `len("rpc_test_") + len(str(uuid.uuid4()))` is 9 + 36. The +# separator `os.path.join` inserts is the third character of the budget, and +# it is easy to drop --- doing so yields 62 and admits one `TMPDIR` that does +# not work. +# --------------------------------------------------------------------------- +UNIX_SOCKET_PATH_MAX = 107 +RPC_SOCKET_BASENAME_LEN = 45 +MAX_TMPDIR_LEN = UNIX_SOCKET_PATH_MAX - RPC_SOCKET_BASENAME_LEN - len("/") + +CommandRunner = Callable[[Sequence[str], Mapping[str, str] | None, int], "CommandResult"] + + +@dataclass(frozen=True) +class CommandResult: + """What a probe observed. ``timed_out`` is distinct from a non-zero exit.""" + + returncode: int + stdout: str = "" + stderr: str = "" + timed_out: bool = False + + @property + def ok(self) -> bool: + return self.returncode == 0 and not self.timed_out + + +@dataclass(frozen=True) +class DeviceReport: + """One card as ``nvidia-smi`` describes it, plus our verdict on it.""" + + index: int + memory_free_mib: int + memory_total_mib: int + utilization_pct: int + compute_cap: str + + @property + def arch(self) -> str: + """``10.3`` -> ``sm_103``. The form catalog receipts are keyed by.""" + major, _, minor = self.compute_cap.partition(".") + return f"sm_{major}{minor or '0'}" + + @property + def free_fraction(self) -> float: + if self.memory_total_mib <= 0: + return 0.0 + return self.memory_free_mib / self.memory_total_mib + + @property + def looks_idle(self) -> bool: + return ( + self.free_fraction >= MIN_FREE_MEMORY_FRACTION + and self.utilization_pct <= MAX_UTILIZATION_PCT + ) + + def why_not_idle(self) -> str: + reasons = [] + if self.free_fraction < MIN_FREE_MEMORY_FRACTION: + reasons.append( + f"{self.memory_free_mib} MiB free of {self.memory_total_mib} " + f"({self.free_fraction:.0%}, need {MIN_FREE_MEMORY_FRACTION:.0%})" + ) + if self.utilization_pct > MAX_UTILIZATION_PCT: + reasons.append( + f"{self.utilization_pct}% utilization (a card spinning this hard is " + f"broken for our purposes, not merely busy)" + ) + return "; ".join(reasons) or "idle" + + +@dataclass(frozen=True) +class PreflightResult: + """The run's resolved identity and launch environment.""" + + slug: str + branch: str + devices: tuple[int, ...] + world_size: int + environment: RunEnvironment + family: str = "" + checkpoint_segment: str = "" + target_rel: str = "" + checkpoint_config: dict[str, Any] = field(default_factory=dict) + notes: list[str] = field(default_factory=list) + + +class PreflightError(RuntimeError): + """A precondition failed. The run ends here, before the branch exists.""" + + +#: The one process-wide local executor for probes. Stateless. +_LOCAL_EXECUTOR = LocalCommandExecutor() + + +def _default_runner( + argv: Sequence[str], + env: Mapping[str, str] | None = None, + timeout: int = PROBE_TIMEOUT_S, +) -> CommandResult: + """Run one probe here, as an argv, with no shell in between. + + A probe is a fixed command with fixed arguments, so there is nothing for + a shell to do except find new ways to split them. The launch prefix — + which *is* shell syntax — belongs to the gates, not here. + """ + record = _LOCAL_EXECUTOR.execute( + CommandSpec( + argv=tuple(str(part) for part in argv), env=env, timeout_s=timeout, label="probe" + ) + ) + return CommandResult( + returncode=record.returncode, + stdout=record.stdout, + stderr=record.stderr, + timed_out=record.timed_out, + ) + + +def query_devices(runner: CommandRunner = _default_runner) -> list[DeviceReport]: + """Read every visible card's memory, utilization and compute capability.""" + result = runner( + [ + "nvidia-smi", + "--query-gpu=index,memory.free,memory.total,utilization.gpu,compute_cap", + "--format=csv,noheader,nounits", + ], + None, + PROBE_TIMEOUT_S, + ) + if not result.ok: + raise PreflightError( + "nvidia-smi did not answer, so no device claim can be made on this " + f"host (exit {result.returncode}). stderr: {result.stderr.strip()!r}" + ) + devices: list[DeviceReport] = [] + for line in result.stdout.splitlines(): + line = line.strip() + if not line: + continue + fields = [part.strip() for part in line.split(",")] + if len(fields) != 5: + raise PreflightError(f"unparsable nvidia-smi row: {line!r}") + try: + devices.append( + DeviceReport( + index=int(fields[0]), + memory_free_mib=int(fields[1]), + memory_total_mib=int(fields[2]), + utilization_pct=int(fields[3]), + compute_cap=fields[4], + ) + ) + except ValueError as exc: + raise PreflightError(f"unparsable nvidia-smi row: {line!r} ({exc})") from exc + return devices + + +def probe_device_alone( + index: int, + python: str, + runner: CommandRunner = _default_runner, +) -> CommandResult: + """Allocate one tensor on one card in a lone process. + + This is the check that separates *broken* from *contended*. Staircase + measured a card reporting 0 MiB used, 0% utilization and a healthy + fabric that still refused ``torch.cuda.set_device`` to a single + process with no peers. + """ + script = "import torch; torch.cuda.set_device(0); torch.empty(1, device='cuda'); print('ok')" + return runner( + [python, "-c", script], + {"CUDA_VISIBLE_DEVICES": str(index)}, + PROBE_TIMEOUT_S, + ) + + +def claim_devices( + world_size: int, + arch: str, + python: str, + *, + runner: CommandRunner = _default_runner, + notes: list[str] | None = None, +) -> tuple[int, ...]: + """Claim ``world_size`` healthy cards of ``arch``, or raise.""" + notes = notes if notes is not None else [] + reports = query_devices(runner) + if not reports: + raise PreflightError("nvidia-smi reported no devices on this host.") + + wrong_arch = [r for r in reports if r.arch != arch] + if len(wrong_arch) == len(reports): + seen = sorted({r.arch for r in reports}) + raise PreflightError( + f"task.yaml asks for {arch} but this host has {seen}. Catalog receipts " + f"are per-architecture, so running here would gate the target against " + f"another machine's certification. Move hosts or fix `arch`." + ) + + candidates = [r for r in reports if r.arch == arch] + idle = [r for r in candidates if r.looks_idle] + busy = [r for r in candidates if not r.looks_idle] + for report in busy: + notes.append(f"device {report.index} not claimable: {report.why_not_idle()}") + + claimed: list[int] = [] + for report in idle: + if len(claimed) == world_size: + break + probe = probe_device_alone(report.index, python, runner) + if probe.ok: + claimed.append(report.index) + continue + detail = "timed out" if probe.timed_out else (probe.stderr.strip() or "non-zero exit") + notes.append( + f"device {report.index} looked idle but failed a lone-process probe " + f"({detail}) — unusable rather than contended; no amount of waiting " + f"will free it" + ) + + if len(claimed) < world_size: + raise PreflightError( + f"need {world_size} healthy {arch} device(s), claimed {len(claimed)}: " + f"{claimed}. Details:\n - " + "\n - ".join(notes or ["(no devices examined)"]) + ) + return tuple(claimed) + + +def check_checkpoint(checkpoint_path: str | Path) -> dict[str, Any]: + """Assert the checkpoint parses, and return its ``config.json``.""" + config_path = Path(checkpoint_path) / "config.json" + if not config_path.is_file(): + raise PreflightError(f"{config_path} does not exist; this is not a HF checkpoint dir.") + try: + data = json.loads(config_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise PreflightError(f"{config_path} is not valid JSON: {exc}") from exc + if not isinstance(data, dict): + raise PreflightError(f"{config_path} must hold a JSON object, got {type(data).__name__}") + return data + + +def resolve_model_key(task_data: Mapping[str, Any], unit: str | None = None) -> str: + """The HuggingFace repo name this unit's anchor row is keyed by. + + The same rule ``StaircaseOnboardWorkflow._model_key`` applies at gate + time: an explicit ``accuracy.model_key`` when the spec states one, + otherwise the last two components of the checkpoint path, which is the + shape a repo name has (``openai/gpt-oss-120b``). It is stated here as + well because preflight has to ask the same question hours earlier, and + an anchor check that resolved a *different* key than the gate will is + worse than no check at all — it passes on a row nothing ever reads. + Once ``workflow`` is free to change, that method should delegate here so + there is one implementation rather than two that agree today. + + ``unit`` honours the per-unit mapping form via + :func:`~.gates.resolve_per_unit`, for the same reason ``reference_yaml`` + does: a feature increment whose anchor lives under another repo name has + to be able to say so. + """ + from . import gates as gates_mod + + accuracy = dict(task_data.get("accuracy") or {}) + explicit = gates_mod.resolve_per_unit(accuracy, "model_key", unit) + if explicit: + return str(explicit) + parts = Path(str(task_data["checkpoint_path"]).rstrip("/")).parts + return "/".join(parts[-2:]) if len(parts) >= 2 else parts[-1] + + +def check_accuracy_anchors( + task_data: Mapping[str, Any], + units: Sequence[str] | None = None, +) -> list[str]: + """Assert every unit's anchor ROW resolves, before any unit is built. + + The accuracy gate that *ends* a unit looks its reference score up by + ``model_key`` plus a discriminator set, and the match is exact — see + :func:`~.gates.load_reference_score`. So a spec can name a benchmark, a + reference file and a row that is not in it, and nothing says so until + the unit has been assembled, reviewed, booted on every member and handed + to ``trtllm-eval``: hours later, on a spec error knowable in + milliseconds at startup, wearing the costume of a target defect. + + Not hypothetical. It is how the ``multimodal`` unit of the + MiMo-V2.6-Flash run nearly ended — ``mmmu.yaml`` had no row for the + checkpoint, and it was caught only because somebody happened to look + while an earlier unit was still gating. Same class as + :func:`check_checkpoint` and :func:`check_tmpdir`: seconds here, hours + there, and it reads as the target's fault when it surfaces late. + + **Every unit, not just core.** ``benchmark``, ``reference_yaml``, + ``model_key`` and ``discriminators`` each take the per-unit mapping + form, and the unit that was broken was broken *precisely because* it is + the one that differs from the rest. + + The lookup is :func:`~.gates.resolve_reference` — the accuracy gate's + own, reached through the same ``resolve_per_unit`` / + ``resolve_discriminators`` helpers — rather than a second implementation + of the matching rule. Reuse is what makes the claim worth anything: this + check fails exactly when the gate would, and the two cannot drift apart + later. The import is deferred because ``gates`` already reaches back + into this module for :func:`reports_unsupported_architecture`; + :func:`_substrate_boot` defers for the same reason. + + Two situations are **skipped rather than failed**, because failing them + would reject specs that work: + + - a run with no ``accuracy`` block, or a unit whose block names no + reference file — there is no anchor to check; + - a reference file that is not in this checkout. The path may be + absolute and point outside the repo, or the tree may be sparse. What + this check owns is the *row*; a file it cannot read says nothing about + the row, and the note records the weaker claim verbatim. + + Returns notes rather than printing: every resolved anchor is recorded + with the number it found, so the preflight report states what each unit + will be measured against before a single turn is spent on it. + """ + from . import gates as gates_mod + + accuracy = dict(task_data.get("accuracy") or {}) + if not accuracy: + return [] + + notes: list[str] = [] + for unit in units if units is not None else assembly_units(task_data): + declared = str(gates_mod.resolve_per_unit(accuracy, "reference_yaml", unit) or "").strip() + if not declared: + notes.append( + f"accuracy anchor [{unit}]: no reference_yaml for this unit, so there " + f"is no row to check" + ) + continue + benchmark = str(gates_mod.resolve_per_unit(accuracy, "benchmark", unit) or "").strip() + model_key = resolve_model_key(task_data, unit) + # Resolved once and passed in, so the message and the lookup cannot + # disagree about which row was searched for. + wanted = gates_mod.resolve_discriminators(accuracy, unit) + try: + score = gates_mod.resolve_reference( + task_data, model_key=model_key, discriminators=wanted, unit=unit + ) + except FileNotFoundError as exc: + notes.append( + f"accuracy anchor [{unit}]: {exc}; the row is unchecked here, and if " + f"it is still absent at gate time the unit dies there" + ) + continue + except (KeyError, TypeError, ValueError) as exc: + # `KeyError.__str__` is `repr` of its argument, which folds a + # carefully written multi-clause message into one escaped line. + # That message is the valuable part -- it already names the file, + # the model key and the discriminator sets that ARE available -- + # so surface it rather than flattening it to "not found". + detail = str(exc.args[0]) if isinstance(exc, KeyError) and exc.args else str(exc) + raise PreflightError( + f"accuracy anchor for unit `{unit}` does not resolve, so the gate that " + f"ENDS this unit cannot run.\n\n" + f" unit: {unit}\n" + f" benchmark: {benchmark or '(unset)'}\n" + f" reference_yaml: {declared}\n" + f" model_key: {model_key}\n" + f" discriminators: {dict(wanted) if wanted else '{}'}\n\n" + f" {detail}\n\n" + f"Nothing here is about the target — no target exists yet. Add the row " + f"to the reference file, citing the written source the number came from " + f"(an anchor is never invented), or fix `accuracy.model_key` / " + f"`accuracy.discriminators` for this unit. Left alone, this unit would " + f"assemble, boot on every member and die at its accuracy gate hours " + f"from now, on this." + ) from exc + notes.append( + f"accuracy anchor [{unit}]: {benchmark or 'eval'} {model_key} " + f"{dict(wanted) if wanted else '{}'} = {score:g} ({declared})" + ) + return notes + + +def check_tmpdir(env: Mapping[str, str]) -> None: + """Assert ``TMPDIR`` and ``CLAUDE_CODE_TMPDIR`` are set, short, and not ``/tmp``. + + Engine builds, NVRTC scratch and the harness's own output all default + there. On a shared host another tenant can fill ``/tmp`` to zero + mid-run, and everything that was running stops at once with errors that + point anywhere but here. + + **Length is the property that breaks first on a cluster whose storage + paths are deep**, and it is the one an author is least likely to suspect. + ``TMPDIR`` also has to hold a unix domain socket — see + :data:`MAX_TMPDIR_LEN` — and a directory that is on the right filesystem, + is writable, and has plenty of space still fails if its *name* is long. + The observed failure is a ZeroMQ traceback naming neither this workflow + nor the spec line that caused it, arriving after four ranks have started + and loaded weights. Here it costs a second. + + The length bound applies to ``TMPDIR`` alone. ``CLAUDE_CODE_TMPDIR`` is + the harness's scratch and binds no socket, so holding it to a socket's + budget would reject a value that works. + + ``env`` is the environment **declared in task.yaml**, and this check + deliberately does *not* fall back to the orchestrator's own + ``os.environ``. Agents and gates run under the declared prefix, so a + value that happens to be set in the launching shell is not a value the + run can rely on — and a run whose scratch location is implicit is one + nobody can reproduce later. + """ + for name in ("TMPDIR", "CLAUDE_CODE_TMPDIR"): + value = env.get(name) + if not value: + raise PreflightError( + f"{name} is not set. Engine builds and NVRTC scratch default there, and " + f"on a shared host the default fills to zero mid-run. Set it in " + f"task.yaml's `environment.env`." + ) + resolved = Path(value).resolve() + if resolved == Path("/tmp") or Path("/tmp") in resolved.parents: + raise PreflightError( + f"{name}={value} points into /tmp. Another tenant can fill it to zero " + f"mid-run; point it at scratch or lustre instead." + ) + if name == "TMPDIR" and len(value) > MAX_TMPDIR_LEN: + raise PreflightError( + f"{name}={value} is {len(value)} characters. The LLM API binds an RPC " + f"socket at `$TMPDIR/rpc_test_`, and a unix socket path is " + f"capped at {UNIX_SOCKET_PATH_MAX} " + f"(`sizeof(sockaddr_un.sun_path)`), so {name} must be at most " + f"{MAX_TMPDIR_LEN}. Point it at a shorter path on the same storage — " + f"this is about the length of the name, not about the filesystem, and " + f"a short path on the same mount works." + ) + + +def check_multirank_substrate( + world_size: int, + environment: RunEnvironment, + runner: CommandRunner = _default_runner, +) -> None: + """Exercise an NCCL collective and an MPI dynamic spawn under a deadline. + + These are the two launch-substrate failures that otherwise present as + an engine that simply never boots, hours in. Both hang rather than + raise, which is why the deadline is the point of the check. + """ + if world_size <= 1: + return + script = ( + "import torch, torch.distributed as dist;" + "dist.init_process_group('nccl');" + "t = torch.ones(1, device='cuda');" + "dist.all_reduce(t);" + "dist.barrier();" + "print('[preflight] OK' if t.item() == dist.get_world_size() else '[preflight] BAD')" + ) + argv = [ + "torchrun", + f"--nproc_per_node={world_size}", + "--nnodes=1", + "-c", + script, + ] + result = runner(argv, dict(environment.env), MULTIRANK_PROBE_TIMEOUT_S) + if result.timed_out: + raise PreflightError( + f"the {world_size}-rank NCCL/MPI probe did not finish within " + f"{MULTIRANK_PROBE_TIMEOUT_S}s. It hangs rather than fails when the " + f"launch substrate is wrong — typically an MPI connect-back landing on " + f"an interface that does not carry it. Without this check the same " + f"problem surfaces as an engine that never boots, hours into the run." + ) + if not result.ok or "[preflight] OK" not in result.stdout: + raise PreflightError( + f"the {world_size}-rank NCCL/MPI probe failed (exit {result.returncode}).\n" + f"stdout: {result.stdout.strip()!r}\nstderr: {result.stderr.strip()!r}" + ) + + +#: How long the substrate gate may take. It loads real weights, so it is +#: minutes; but a substrate gate that needs longer than this is measuring +#: something other than "can this host run the engine at all". +SUBSTRATE_TIMEOUT_S = 1800 + +#: How the stock path says *"I have never heard of this architecture"*, quoted +#: from the two places that raise it. ``AutoModelForCausalLM.from_config`` and +#: the multimodal-encoder branch beside it, in +#: ``tensorrt_llm/_torch/models/modeling_auto.py``, raise ``ValueError`` after +#: ``get_registered_model_class`` returns ``None``; +#: ``AutoConfig.from_hugging_face`` and +#: ``AutoModelForCausalLM.get_trtllm_model_class`` in +#: ``tensorrt_llm/models/automodel.py`` raise ``NotImplementedError`` after +#: ``MODEL_MAP.get`` returns ``None``. Both registries are keyed on the +#: checkpoint's own ``architectures[0]``, which is the string this workflow +#: exists to introduce. +#: +#: Matched on the message text because that is all the gate has: it runs the +#: engine as a subprocess and reads a log, so there is no exception object to +#: catch. The markers are deliberately the invariant half of each message — +#: the architecture name is interpolated into it and varies per run. +UNSUPPORTED_ARCHITECTURE_MARKERS: tuple[str, ...] = ( + "Unknown architecture for AutoModelForCausalLM", + "Unknown architecture for AutoModelForMultimodalEncoder", + "is not supported in TRT-LLM yet", + # `get_model_architecture`'s own miss, from `_torch/models/modeling_utils.py`: + # raise RuntimeError(f"Unknown model architecture: {arch}") + # A different message from the two `AutoModelFor*` ones above and reached by + # a different path -- the registry lookup rather than the config loader -- + # so a checkpoint can fail here without matching any of them. Found by the + # reviewer of a unit whose acceptance gate depends on this list: if the + # `=off` leg dies in the registry instead, the gate hard-fails rather than + # falling back to its declared floor. + "Unknown model architecture:", + # `transformers`' own refusal, from `resolve_trust_remote_code`. It is the + # same condition wearing a different coat: a checkpoint that ships an + # `auto_map` and demands `trust_remote_code=True` is, by construction, one + # whose architecture neither `transformers` nor TensorRT-LLM has built in — + # the implementation travels with the weights. The stock path therefore + # cannot load it, which is exactly what this fallback is for. + # + # Measured on MiMo-V2.6-Flash-RL: the substrate gate raised this rather + # than any registry message, so a discriminator keyed only on the registry + # missed it and the run stopped on a checkpoint nobody expected the stock + # path to load. + "contains custom code which must be executed to correctly load the model", + # `transformers`' refusal of an unknown `model_type` with NO remote code to + # fall back on, from `AutoConfig.from_pretrained`. The checkpoint is newer + # than the installed `transformers` and ships no configuration module, so + # neither the stock path nor `trust_remote_code` can load it -- the same + # condition again, one step earlier than the registry. + # + # Measured on DeepSeek-V4.1-Flash-NVFP4 (`model_type: deepseek_v41`, + # transformers 5.5.4): the substrate gate raised this and the run stopped, + # instead of proving the substrate on the declared fallback checkpoint. + "but Transformers does not recognize this architecture", + # A registration shim: the architecture name is registered only so the + # runtime can find something keyed on it (a multimodal input processor), + # and building it without a modeling_v2 target raises. Once the run's own + # config classes are registered, this is where the `=off` leg of the + # acceptance gate dies -- after the config loader and the registry have + # both succeeded -- so without this marker the gate hard-fails instead of + # falling back to its declared floor. Measured on DeepSeek-V4.1-Flash-NVFP4 + # (`tensorrt_llm/_torch/models/modeling_deepseekv41.py`). + "has no built-in implementation", +) + + +def reports_unsupported_architecture(detail: str, output: str) -> bool: + """Whether this text says *"the stock path does not know this model"*. + + Text-based so that callers other than the substrate gate can use it. The + acceptance gate needs exactly this question: its `=off` leg compares + against built-in modeling, and a genuinely new architecture has none. + """ + haystack = f"{detail}\n{output}" + return any(marker in haystack for marker in UNSUPPORTED_ARCHITECTURE_MARKERS) + + +def _reports_unsupported_architecture(result: Any) -> bool: + """Whether this failure is *"the stock path does not know this model"*. + + Narrow on purpose, and the narrowness is the whole design. A fallback + keyed on any failure would answer a real kvCache sizing fault on the + target checkpoint by silently measuring a different model and reporting + green — the same class of mistake as gating under + ``TRTLLM_MODELING_V2=auto``. So the discriminator is the architecture + registry's own message and nothing else: every other failure the boot + gate can produce is still a hard stop. + """ + return reports_unsupported_architecture(result.detail, result.stdout_tail) + + +def _substrate_boot(task_data: Mapping[str, Any], environment: RunEnvironment, shell_runner: Any): + """One substrate boot attempt, under the stock path and no features.""" + from . import gates as gates_mod + + return gates_mod.run_boot_gate( + environment, + task_data=task_data, + spec=gates_mod.BootSpec(prompts=("The capital of France is",), max_tokens=4), + runner=shell_runner, + timeout=SUBSTRATE_TIMEOUT_S, + mode=gates_mod.MODE_OFF, + features=[], + ) + + +def substrate_gate( + task_data: Mapping[str, Any], + environment: RunEnvironment, + *, + shell_runner: Any, + notes: list[str] | None = None, +) -> None: + """Prove this host can *generate a token* under the run's exact config. + + The single highest-value check in this module, and it was added after a + run spent roughly half its wall clock on failures that all looked like + model defects: + + ==================================== ================================== + ``MPI_ERR_SPAWN: could not spawn`` the launcher choice + ``Impossible to fit ... in kvCache`` ``max_seq_len`` defaulting to the + checkpoint's 262144 + ``NVRTC ... cannot open "cuda.h"`` the venv's bin missing from PATH + ==================================== ================================== + + Each surfaced *mid-run*, inside an agent's turn, reading as though the + target were wrong. Each is a property of the host and the resolved + ``runtime`` block, knowable in minutes before any agent exists. + + **It generates rather than constructs, and that is the whole point.** + The NVRTC fault appeared on the *first generated token*, after the engine + had loaded 22 GB of weights, allocated 150 GiB of KV cache and completed + warm-up. A gate that stopped at "the engine constructed" would have + passed it straight through. + + It runs the real checkpoint under the real ``runtime`` block rather than + a toy model, because every one of the three failures above depended on + this checkpoint's geometry or this run's flags. + + **And that is exactly why it needs a fallback.** Reaching for the run's + own checkpoint only works for a checkpoint the stock path can load, and + this workflow's entire purpose is checkpoints it cannot — so for a + genuinely new architecture the highest-value check in the module is + unavailable precisely when the run is doing the thing it exists for. When + the stock path answers "I do not know this architecture" (and *only* + then — see :func:`_reports_unsupported_architecture`), the gate re-runs + itself against ``environment.substrate_checkpoint``, a proxy the stock + path does support, under the same environment and the same ``runtime`` + block. The two host failures that do not depend on the checkpoint — + ``MPI_ERR_SPAWN`` and NVRTC's missing ``cuda.h`` — are still proven; the + kvCache one is not, and the note says so rather than letting the run + believe otherwise. + """ + result = _substrate_boot(task_data, environment, shell_runner) + if not result.passed and not _reports_unsupported_architecture(result): + raise PreflightError( + "substrate gate failed: this host cannot generate a token under the " + "run's own environment and runtime block, using the STOCK modeling " + "path. Nothing here is about the target — no target exists yet — so " + "fix the host or the spec before dispatching an agent.\n\n" + f" {result.detail}\n\n" + f"command:\n {result.command}\n\n" + f"last lines:\n{result.stdout_tail}" + ) + + if not result.passed: + proxy = str((task_data.get("environment") or {}).get("substrate_checkpoint") or "").strip() + if not proxy: + # Deliberately NOT the raw architecture error. That message names + # the target's architecture and reads as a target defect, which is + # the exact failure mode this gate was written to prevent — and + # here there is no target yet to be defective. + raise PreflightError( + "substrate gate: the stock modeling path does not know this " + "checkpoint's architecture, which is the normal state of affairs " + "for a model this workflow is onboarding — it is not a fault in " + "the checkpoint or the host. But it means the substrate cannot be " + "proven on this checkpoint at all.\n\n" + "Declare `environment.substrate_checkpoint`: a checkpoint the stock " + "path DOES support, which the gate will run under this same " + "environment and runtime block instead. Do not reach for " + "`--skip-substrate-gate` — that also gives up the MPI spawn and " + "NVRTC checks, which do not depend on the checkpoint at all.\n\n" + f" {result.detail}\n\n" + f"last lines:\n{result.stdout_tail}" + ) + proxy_result = _substrate_boot( + {**dict(task_data), "checkpoint_path": proxy}, environment, shell_runner + ) + if not proxy_result.passed: + raise PreflightError( + "substrate gate failed on the declared proxy checkpoint " + f"{proxy}: this host cannot generate a token under the run's own " + "environment and runtime block, using the STOCK modeling path and a " + "checkpoint that path supports. Nothing here is about the target — " + "no target exists yet — so fix the host or the spec before " + "dispatching an agent.\n\n" + f" {proxy_result.detail}\n\n" + f"command:\n {proxy_result.command}\n\n" + f"last lines:\n{proxy_result.stdout_tail}" + ) + if notes is not None: + # The weaker claim, stated in full rather than summarized. A note + # that said only "substrate gate passed" would be read later as the + # claim the unqualified gate makes, and the difference is precisely + # which failure this run has been protected from. + notes.append( + f"substrate proven on `{proxy}`; this checkpoint's own geometry is " + f"unverified, so a kvCache sizing fault will surface at the first " + f"boot gate rather than here." + ) + return + + if notes is not None: + notes.append( + f"substrate gate: stock path generated under the resolved environment " + f"in {result.duration_s:.0f}s" + ) + + +def agent_plane_gate(probe: Any, notes: list[str] | None = None) -> None: + """Prove one agent turn can happen at all, before a phase depends on it. + + Every unit test in this package replaces the agent layer with a fake, so + a fully green suite says nothing about whether an agent can start. The + first real run discovered that twice — once after a 24-minute planning + turn, and once three hours later mid-batch — because a nested process + could not write its credential lock. + + ``probe`` is a zero-argument callable returning the agent's reply; it is + injected so this stays testable without spawning anything. + """ + try: + reply = probe() + except Exception as exc: # noqa: BLE001 - the cause is the message + raise PreflightError( + f"agent plane gate failed: a single agent turn could not complete, so " + f"no phase that dispatches one can run. This is not a model or target " + f"problem.\n\n {type(exc).__name__}: {exc}" + ) from exc + if not str(reply or "").strip(): + raise PreflightError( + "agent plane gate failed: the agent returned an empty reply. A role " + "that cannot answer a trivial prompt will not produce a decision file." + ) + if notes is not None: + notes.append("agent plane gate: one agent turn completed") + + +def run_preflight( + task_data: Mapping[str, Any], + *, + runner: CommandRunner = _default_runner, + skip_device_claim: bool = False, +) -> PreflightResult: + """Assert every precondition in order and resolve the run's identity. + + ``skip_device_claim`` exists for the CPU-only tests and for a + ``--dry-run`` on a machine with no GPUs; it is never set by a real run. + """ + notes: list[str] = [] + + checkpoint_config = check_checkpoint(task_data["checkpoint_path"]) + + arch = str(task_data["arch"]) + segment = str(task_data["parallel"]) + kind, declared_world_size = parse_parallel_segment(segment) + world_size = declared_world_size + if world_size is None: + # `unified` takes its world size from the MEMBER SET rather than from + # the segment string, which is what `parse_parallel_segment` has + # always said by returning None here. Every member declares the same + # one -- the schema rejects a mixture, because one allocation serves + # them all -- so the first answers for the set. 4 is the default the + # parallel skill documents, and the size the default member set is + # built at. + declared_members = members_for(task_data, None, DEFAULT_UNIFIED_WORLD_SIZE) + _, member_world_size = parse_parallel_segment(declared_members[0]) + world_size = member_world_size or DEFAULT_UNIFIED_WORLD_SIZE + notes.append( + f"unified segment: world size {world_size}, taken from the member set " + f"({', '.join(declared_members)})" + ) + + # The member set and what it costs, stated where the run is CHOSEN rather + # than discovered three gates in. Three members over four units is twelve + # accuracy gates at roughly an hour each, and `members:` is the knob that + # buys a cheaper answer honestly — one member, named — instead of a cheap + # answer labelled as an expensive one. `--dry-run` prints this block, so + # the bill is readable before anything is spent. + units = assembly_units(task_data) + per_unit = {unit: members_for(task_data, unit, world_size) for unit in units} + distinct = {tuple(chosen) for chosen in per_unit.values()} + total = sum(len(chosen) for chosen in per_unit.values()) + if len(distinct) == 1: + shared = next(iter(distinct)) + listed = f"{', '.join(shared)} -> {len(shared)} x (boot + accuracy) per unit" + else: + # The per-unit form, so a run that bought a cheaper answer for one + # unit says which unit and how cheap. + listed = "; ".join(f"{unit}: {', '.join(chosen)}" for unit, chosen in per_unit.items()) + notes.append( + f"members: {listed}, {total} accuracy gate(s) over {len(units)} unit(s): {', '.join(units)}" + ) + + # Immediately after the unit list is resolved, and before anything that + # touches the host. This check is pure file reading -- it costs + # milliseconds and needs no device, no allocation and no agent -- so + # every expensive check below it is work that a failure here would have + # wasted. It cannot come any earlier because it is asked PER UNIT, and + # the unit list is decided here. + notes.extend(check_accuracy_anchors(task_data, units)) + + env_block = dict(task_data.get("environment") or {}) + declared_env = {str(k): str(v) for k, v in (env_block.get("env") or {}).items()} + check_tmpdir(declared_env) + + python = str(env_block.get("python", "python3")) + # A device claim only means anything when the orchestrator and the GPUs + # are the *same machine*, which is to say `local`. Everywhere else the + # allocation is the claim: the scheduler has already decided, and + # probing this host's cards describes the wrong machine — or, on a + # login node, no machine at all, since `nvidia-smi` cannot reach a + # driver there and the claim fails for a run that is perfectly fine. + # + # Phrased as "not local" rather than as a list of remote kinds on + # purpose: a kind added later is remote until someone says otherwise, + # and that is the safe default. `session_spool` was exactly this case. + scheduler_owns_devices = str(env_block.get("kind", "local")) != "local" + devices: tuple[int, ...] = () + if not skip_device_claim and not scheduler_owns_devices: + devices = claim_devices(world_size, arch, python, runner=runner, notes=notes) + if scheduler_owns_devices: + notes.append( + f"{env_block.get('kind')}: devices are the scheduler's to allocate, so " + f"preflight claims none and asserts no architecture locally. `arch: {arch}` " + f"is taken on trust here and is what the receipts will be keyed by." + ) + + environment = resolve_environment(task_data, devices=devices, world_size=world_size) + + if not skip_device_claim and not scheduler_owns_devices: + check_multirank_substrate(world_size, environment, runner) + + slug = derive_slug(task_data) + family, segment, target_rel = resolve_target_identity(task_data, checkpoint_config) + return PreflightResult( + slug=slug, + branch=f"staircase/{slug}", + devices=devices, + world_size=world_size, + environment=environment, + family=family, + checkpoint_segment=segment, + target_rel=target_rel, + checkpoint_config=checkpoint_config, + notes=notes, + ) + + +def resolve_target_identity( + task_data: Mapping[str, Any], + checkpoint_config: Mapping[str, Any], +) -> tuple[str, str, str]: + """Return ``(family, checkpoint_segment, target_rel)`` for this run. + + Both components are derived and both are overridable through + ``task.yaml``'s ``target`` block. The family comes from the checkpoint's + own ``architectures[0]``, which is the same string routing keys on — a + hand-typed family that disagrees with it is a target nothing can reach. + The checkpoint segment is a guess (see + :func:`~.task_schema.derive_checkpoint_segment`), and it is the one + identity component worth setting by hand. + """ + override = dict(task_data.get("target") or {}) + + architectures = checkpoint_config.get("architectures") or [] + if not architectures and not override.get("family"): + raise PreflightError( + "the checkpoint's config.json declares no `architectures`, so the " + "modeling_v2 family cannot be derived. Set `target.family` in task.yaml." + ) + family = override.get("family") or derive_family(str(architectures[0])) + segment = override.get("checkpoint_segment") or derive_checkpoint_segment( + task_data["checkpoint_path"] + ) + + target_rel = ( + f"{MODELING_V2_PACKAGE}/models/{family}/" + f"{segment}__{task_data['arch']}__{task_data['parallel']}" + ) + return family, segment, target_rel + + +def format_preflight_report(result: PreflightResult) -> str: + """A human-readable summary for the journal and the console.""" + lines = [ + f"slug: {result.slug}", + f"branch: {result.branch}", + f"target: {result.target_rel}", + f"world size: {result.world_size}", + f"devices: {', '.join(str(d) for d in result.devices) or '(allocated by Slurm)'}", + f"launch: {result.environment.kind}", + f"prefix: {result.environment.prefix or '(none)'}", + f"interpreter: {shlex.quote(result.environment.python)}", + ] + if result.notes: + lines.append("notes:") + lines.extend(f" - {note}" for note in result.notes) + return "\n".join(lines) + + +__all__ = [ + "CommandResult", + "CommandRunner", + "DeviceReport", + "MAX_TMPDIR_LEN", + "MAX_UTILIZATION_PCT", + "MIN_FREE_MEMORY_FRACTION", + "RPC_SOCKET_BASENAME_LEN", + "UNIX_SOCKET_PATH_MAX", + "PreflightError", + "PreflightResult", + "UNSUPPORTED_ARCHITECTURE_MARKERS", + "check_accuracy_anchors", + "check_checkpoint", + "check_multirank_substrate", + "check_tmpdir", + "claim_devices", + "format_preflight_report", + "probe_device_alone", + "query_devices", + "resolve_model_key", + "run_preflight", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/progress.py b/agent-flow/agent_flow/workflows/staircase_onboard/progress.py new file mode 100644 index 000000000000..a47fa06c34bb --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/progress.py @@ -0,0 +1,769 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed decisions — the only channel an agent's verdict reaches the orchestrator through. + +Agents never write ``progress.yaml`` directly and the orchestrator never +parses prose. Each role gets an ``append__progress`` MCP tool whose +``decision`` is a closed enum, a ``Stop`` hook refuses to let the turn end +until it was called, and the orchestrator branches on the field. A verdict +that cannot be expressed in the enum is a verdict the workflow does not +know how to act on, which is exactly the thing worth finding out loudly. + +The file is a YAML mapping with one list per stage, so iteration counters +stay local to the stage that produced them:: + + plan_stage: + - iteration: 1 + agent: plan_drafter + timestamp: "2026-09-17T22:41:03" + summary: | + ... + decision: DRAFT_READY + build_stage: + - {iteration: 1, agent: smith, ..., decision: LANDED, entry: catalog/norm/x.py} + deliver_stage: + - {iteration: 1, agent: auditor, ..., decision: CLEAN} + +Concurrency matters here in a way it does not in the other workflows: the +smith batch runs several agents at once, all appending. Every write takes a +lock and re-reads before appending, so a batch of eight cannot lose seven +entries to a read-modify-write race. +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Any + +import yaml +from claude_agent_sdk import tool +from rich.syntax import Syntax + +from agent_flow.console import print_layer_panel +from agent_flow.logger import get_logger + +PLAN_STAGE = "plan_stage" +BUILD_STAGE = "build_stage" +DELIVER_STAGE = "deliver_stage" +_STAGES = (PLAN_STAGE, BUILD_STAGE, DELIVER_STAGE) + +_STAGE_BY_AGENT = { + "plan_drafter": PLAN_STAGE, + "plan_reviewer": PLAN_STAGE, + "smith": BUILD_STAGE, + "assembler": BUILD_STAGE, + "reviewer": BUILD_STAGE, + "auditor": DELIVER_STAGE, +} + +# Decision vocabularies. Each is closed on purpose — see the module +# docstring. ``BLOCKED`` variants carry a reason so the orchestrator's +# triage does not have to read the summary. +PLAN_DRAFTER_DECISIONS = ("DRAFT_READY", "CEILING") +PLAN_REVIEWER_DECISIONS = ("APPROVE", "REJECT") +SMITH_DECISIONS = ("LANDED", "FAILED_RECEIPT", "MACRO", "BLOCKED") +ASSEMBLER_DECISIONS = ("UNIT_READY", "BLOCKED") +REVIEWER_DECISIONS = ("APPROVE", "REJECT") +AUDITOR_DECISIONS = ("CLEAN", "UNCOVERED") + +# Staircase's triage table, as an enum. Each maps to a different +# orchestrator action, which is the whole reason it is not free text. +BLOCKED_REASONS = ("vocabulary_gap", "knowledge_gap", "accuracy_shortfall", "environment") + + +def _empty() -> dict[str, list[dict[str, Any]]]: + return {stage: [] for stage in _STAGES} + + +def read_progress(path: Path) -> dict[str, list[dict[str, Any]]]: + """Load ``progress.yaml``, returning every stage key even when absent.""" + if not path.is_file(): + return _empty() + text = path.read_text(encoding="utf-8") + if not text.strip(): + return _empty() + data = yaml.safe_load(text) + if data is None: + return _empty() + if not isinstance(data, dict): + raise ValueError( + f"{path} must contain a YAML mapping keyed by stage, got {type(data).__name__}" + ) + result = _empty() + for stage in _STAGES: + value = data.get(stage) or [] + if not isinstance(value, list): + raise ValueError(f"{path}: `{stage}` must be a list, got {type(value).__name__}") + result[stage] = value + return result + + +def write_progress(path: Path, data: dict[str, list[dict[str, Any]]]) -> None: + """Persist with a fixed key order so diffs stay readable.""" + ordered = {stage: data.get(stage, []) for stage in _STAGES} + path.write_text( + yaml.safe_dump(ordered, sort_keys=False, allow_unicode=True, default_flow_style=False), + encoding="utf-8", + ) + + +def init_progress_file(path: Path) -> None: + write_progress(path, _empty()) + + +def find_entries( + path: Path, + *, + stage: str, + agent: str | None = None, + last_iterations: int | None = None, +) -> list[dict[str, Any]]: + """Return entries from one stage, optionally filtered by agent and recency. + + Recency follows **append order**, not iteration number: a smith batch + stamps several entries with the same iteration, and an assembler's + re-dispatch after a mid-unit smith carries a number already seen. + """ + if stage not in _STAGES: + raise ValueError(f"unknown stage: {stage!r}") + if last_iterations is not None and last_iterations < 1: + raise ValueError(f"last_iterations must be >= 1, got {last_iterations}") + if agent is not None and _STAGE_BY_AGENT.get(agent) != stage: + return [] + entries = list(read_progress(path)[stage]) + if agent is not None: + entries = [e for e in entries if e.get("agent") == agent] + if last_iterations is not None: + keep_from = len(entries) + seen: list[Any] = [] + for index in range(len(entries) - 1, -1, -1): + iteration = entries[index].get("iteration") + if iteration is None: + continue + if iteration not in seen: + if len(seen) == last_iterations: + break + seen.append(iteration) + keep_from = index + entries = entries[keep_from:] + return entries + + +def latest_entry(path: Path, agent: str) -> dict[str, Any] | None: + stage = _STAGE_BY_AGENT.get(agent) + if stage is None: + return None + matches = find_entries(path, stage=stage, agent=agent) + return matches[-1] if matches else None + + +def latest_decision(path: Path, agent: str, allowed: tuple[str, ...]) -> str | None: + """Return the agent's most recent decision, or ``None`` if it is not one we know.""" + entry = latest_entry(path, agent) + if entry is None: + return None + decision = str(entry.get("decision", "")).strip().upper() + return decision if decision in allowed else None + + +def _now_iso() -> str: + return datetime.now().isoformat(timespec="seconds") + + +def _yaml_dump(data: Any) -> str: + return yaml.safe_dump(data, sort_keys=False, allow_unicode=True, default_flow_style=False) + + +def _log_write(agent: str, entry: dict[str, Any]) -> None: + body = Syntax(_yaml_dump([entry]), "yaml", theme="ansi_dark", word_wrap=True) + print_layer_panel( + agent, f"system · wrote iter {entry.get('iteration')}", body, get_logger().console + ) + + +def _log_read(caller: str, description: str, text: str) -> None: + body = Syntax(text, "yaml", theme="ansi_dark", word_wrap=True) + print_layer_panel(caller, f"system · read ({description})", body, get_logger().console) + + +@dataclass +class ProgressContext: + """Shared mutable context the tool handlers close over. + + ``current_iteration`` is set before each dispatch so an entry is stamped + without the agent having to pass (or guess) it. ``lock`` serializes the + read-modify-write, which a concurrent smith batch would otherwise race. + """ + + path: Path + current_iteration: int = 0 + lock: threading.Lock = field(default_factory=threading.Lock, repr=False, compare=False) + + def append(self, agent: str, entry: dict[str, Any]) -> None: + with self.lock: + data = read_progress(self.path) + data[_STAGE_BY_AGENT[agent]].append(entry) + write_progress(self.path, data) + _log_write(agent, entry) + + +def _decision_tool( + ctx: ProgressContext, + agent: str, + *, + decisions: tuple[str, ...], + decision_help: str, + summary_help: str, + extra_properties: dict[str, Any] | None = None, + extra_keys: tuple[str, ...] = (), +): + """Build one role's ``append__progress`` tool.""" + properties: dict[str, Any] = { + "summary": {"type": "string", "description": summary_help}, + "decision": {"type": "string", "enum": list(decisions), "description": decision_help}, + } + properties.update(extra_properties or {}) + + @tool( + f"append_{agent}_progress", + ( + f"Record this turn's {agent} entry in progress.yaml. Call it exactly " + f"once, as the last action of your turn. The orchestrator branches on " + f"`decision`, so it is the only part of your turn it reads." + ), + {"type": "object", "properties": properties, "required": ["summary", "decision"]}, + ) + async def handler(args: dict[str, Any]) -> dict[str, Any]: + entry: dict[str, Any] = { + "iteration": ctx.current_iteration, + "agent": agent, + "timestamp": _now_iso(), + "summary": args["summary"], + "decision": args["decision"], + } + for key in extra_keys: + if args.get(key) is not None: + entry[key] = args[key] + ctx.append(agent, entry) + return { + "content": [ + { + "type": "text", + "text": ( + f"Recorded {agent} entry for iteration {ctx.current_iteration} " + f"(decision={entry['decision']})." + ), + } + ] + } + + return handler + + +def _read_tool(ctx: ProgressContext, caller: str): + """A stage-scoped read tool, so a role sees its own phase and no more.""" + stage = _STAGE_BY_AGENT[caller] + + @tool( + "read_latest_progress", + ( + f"Return the most recent entries from progress.yaml's `{stage}` as YAML. " + f"Cheaper and better scoped than reading the whole file." + ), + { + "type": "object", + "properties": { + "iterations": { + "type": "integer", + "minimum": 1, + "description": "How many of the most recent iterations to include (default 1).", + } + }, + "required": [], + }, + ) + async def handler(args: dict[str, Any]) -> dict[str, Any]: + iterations = int(args.get("iterations") or 1) + selected = find_entries(ctx.path, stage=stage, last_iterations=iterations) + text = _yaml_dump(selected) if selected else f"# No {stage} entries yet.\n" + _log_read(caller, f"{stage}, last {iterations}", text) + return {"content": [{"type": "text", "text": text}]} + + return handler + + +DECISIONS_DIRNAME = "decisions" + +# The enums, keyed by role, so the file protocol can validate exactly what +# the tool schema would have. +DECISIONS_BY_AGENT: dict[str, tuple[str, ...]] = { + "plan_drafter": PLAN_DRAFTER_DECISIONS, + "plan_reviewer": PLAN_REVIEWER_DECISIONS, + "smith": SMITH_DECISIONS, + "assembler": ASSEMBLER_DECISIONS, + "reviewer": REVIEWER_DECISIONS, + "auditor": AUDITOR_DECISIONS, +} + +EXTRA_KEYS_BY_AGENT: dict[str, tuple[str, ...]] = { + # `constituents` is what makes MACRO actionable. Without it the + # decomposition arrives as prose inside `summary`, which a human can act + # on and the orchestrator cannot -- so the one outcome whose entire + # purpose is "here is the work that comes next" ended the run instead of + # extending it. Measured on the first real onboarding, on its largest + # item. + "smith": ("entry", "receipt", "constituents"), + "assembler": ("blocked_reason", "missing_entry"), + # `uncovered` is to the auditor what `constituents` is to a MACRO: the + # difference between a verdict the orchestrator can act on and one a + # human has to read. Prose cannot be deduplicated, grouped by entry + # surface, or compared across rounds -- so an UNCOVERED whose findings + # live only in `summary` ends the run, which is exactly what it used to + # do. See remediation.py. + "auditor": ("uncovered",), +} + +#: Keys an uncovered row may carry. `entry`, `arch` and `axis` are required: +#: without all three a finding cannot be grouped into a dispatch, and a +#: finding that cannot be dispatched is one a human has to re-derive. +UNCOVERED_REQUIRED: tuple[str, ...] = ("entry", "arch", "axis") +UNCOVERED_OPTIONAL: tuple[str, ...] = ("values", "value", "contract_line") + +#: Keys a MACRO decision's `constituents` entries may carry. `target` and +#: `tier` are required: the orchestrator needs a name to dispatch and a rung +#: to budget turns against. +CONSTITUENT_REQUIRED: tuple[str, ...] = ("target", "tier") +CONSTITUENT_OPTIONAL: tuple[str, ...] = ("surface", "note", "unit", "depends_on") + + +def parse_constituents(raw: Any) -> list[dict[str, Any]]: + """Normalize a MACRO decision's child list, dropping anything unusable. + + Permissive about shape and strict about content: a smith writing YAML by + hand should not lose a whole decomposition to a formatting slip, but a + child with no target cannot be dispatched and a child with no tier cannot + be budgeted, so those are dropped rather than guessed at. + """ + if not isinstance(raw, list): + return [] + children: list[dict[str, Any]] = [] + for row in raw: + if not isinstance(row, dict): + continue + target = str(row.get("target", "")).strip() + if not target: + continue + try: + tier = int(row.get("tier", 1) or 1) + except (TypeError, ValueError): + tier = 1 + child: dict[str, Any] = {"target": target, "tier": tier} + for key in CONSTITUENT_OPTIONAL: + value = row.get(key) + if value is None: + continue + if key == "depends_on": + if isinstance(value, list): + child[key] = [int(v) for v in value if str(v).strip().lstrip("-").isdigit()] + elif str(value).strip(): + child[key] = str(value).strip() + children.append(child) + return children + + +def parse_uncovered_records(raw: Any) -> list[dict[str, Any]]: + """Normalize an auditor's uncovered list, dropping anything undispatchable. + + Same bargain as :func:`parse_constituents`, for the same reason: an + auditor writing YAML by hand should not lose a whole finding to a + formatting slip, but a row without an entry, an arch or an axis cannot + be grouped into a verify dispatch. Those are dropped here rather than + guessed at -- a guessed axis merges two unrelated surfaces into one + dispatch and certifies neither. + + Kept as plain dicts because this is the persistence boundary; + :mod:`remediation` turns them into typed rows. + """ + if not isinstance(raw, list): + return [] + rows: list[dict[str, Any]] = [] + for row in raw: + if not isinstance(row, dict): + continue + record = {key: str(row.get(key, "")).strip() for key in UNCOVERED_REQUIRED} + if not all(record.values()): + continue + values = row.get("values") + if values is None: + values = row.get("value") + if isinstance(values, (str, int, float)): + values = [values] + if isinstance(values, list): + record["values"] = [str(v).strip() for v in values if str(v).strip()] + line = str(row.get("contract_line", "")).strip() + if line: + record["contract_line"] = line + rows.append(record) + return rows + + +def decision_path(workspace: Path, agent: str, iteration: int) -> Path: + """Where a role writes its decision when MCP tools are unavailable.""" + return Path(workspace) / DECISIONS_DIRNAME / f"{agent}-{iteration}.yaml" + + +def _decision_marker_path(workspace: Path, agent: str, iteration: int) -> Path: + """The dotfile stamped beside a decision when its turn is dispatched. + + A dotfile, and beside the decision rather than under a directory of its + own, because the two must share a filesystem — see + ``mark_decision_dispatch`` for why that is the whole mechanism. + """ + return Path(workspace) / DECISIONS_DIRNAME / f".{agent}-{iteration}.dispatched" + + +def mark_decision_dispatch(workspace: Path, agent: str, iteration: int) -> None: + """Stamp the moment this turn was dispatched, so a stale answer is detectable. + + ``(agent, iteration)`` is **not unique across a run**: ``iteration`` is 1 + for every turn of the deliver stage, and it restarts per unit in assembly. + So two different turns routinely write the same path — see + ``decision_path`` — and a turn that ends *without* writing one leaves the + **previous** turn's verdict sitting there looking fresh, which + ``collect_decision_file`` then reports as this turn's answer. + + That is not hypothetical. It happened five times in one onboarding: ten + smith decisions collected from the wrong place, an assembler told to write + the previous turn's path, decision paths colliding across units, and — + worst — three audit verdicts read back as a verdict from two passes + earlier, which made the audit loop structurally unable to converge, + because a ``CLEAN`` could not end the run. Every one of those was silent. + + This run saw its cousin: a smith committed 380 lines with a passing + receipt and then died without writing its decision file, blocking a batch + for 58 minutes. Staleness detection is what turns that class of failure + from silent into visible. + + The marker is compared by **mtime against the decision file on the same + filesystem**, never against a clock. An agent may answer from a different + node whose clock is skewed by minutes; both files live under the + workspace, so their mtimes are ordered by the one filesystem that wrote + them. + """ + marker = _decision_marker_path(workspace, agent, iteration) + marker.parent.mkdir(parents=True, exist_ok=True) + marker.write_text("", encoding="utf-8") + + +def decision_is_stale(workspace: Path, agent: str, iteration: int) -> bool: + """Whether the decision file at this path predates the turn that wants it. + + False when either file is absent: a missing decision is the caller's + "no decision" case, and a missing marker means this run never stamped one + — in which case the old unconditional behaviour is the honest fallback, + rather than silently refusing every decision a run ever writes. + """ + decision = decision_path(workspace, agent, iteration) + marker = _decision_marker_path(workspace, agent, iteration) + if not decision.is_file() or not marker.is_file(): + return False + return decision.stat().st_mtime < marker.stat().st_mtime + + +def collect_decision_file( + ctx: ProgressContext, + workspace: Path, + agent: str, + iteration: int, +) -> str | None: + """Read, validate and record the decision file ``agent`` wrote. + + **Why this exists.** The typed-decision MCP tools are the workflow's + preferred mechanism — a closed enum in a tool schema is checked before + the model can emit anything else. But an *enterprise* MCP configuration + (``/etc/claude-code/managed-mcp.json``) makes the SDK refuse every + dynamically-registered server, including the in-process one those tools + live in. On such a host the tools cannot exist at all. + + So the decision becomes a file, and the guarantee moves rather than + disappearing: the enum is validated here, on read, and a decision + outside it is rejected exactly as the schema would have rejected it. + The orchestrator still branches on one closed value it did not get from + prose. + + Returns the validated decision, or ``None`` if the role wrote nothing + usable — which the caller must treat as a failed turn rather than as + silence to interpret. + """ + allowed = DECISIONS_BY_AGENT[agent] + path = decision_path(workspace, agent, iteration) + if not path.is_file(): + return None + # A file older than this turn's dispatch marker belongs to a PREVIOUS turn + # that happened to share this path. Reading it would report that turn's + # verdict as this one's — the failure `mark_decision_dispatch` exists to + # make visible, and one that is silent every other way. + if decision_is_stale(workspace, agent, iteration): + return None + + try: + raw = yaml.safe_load(path.read_text(encoding="utf-8")) + except yaml.YAMLError: + return None + if not isinstance(raw, dict): + return None + + decision = str(raw.get("decision", "")).strip() + if decision not in allowed: + return None + + entry: dict[str, Any] = { + "iteration": iteration, + "agent": agent, + "timestamp": _now_iso(), + "summary": str(raw.get("summary", "")).strip(), + "decision": decision, + } + for key in EXTRA_KEYS_BY_AGENT.get(agent, ()): + value = raw.get(key) + if value is None: + continue + if key == "constituents": + children = parse_constituents(value) + if children: + entry[key] = children + elif key == "uncovered": + rows = parse_uncovered_records(value) + if rows: + entry[key] = rows + elif str(value).strip(): + entry[key] = str(value).strip() + ctx.append(agent, entry) + return decision + + +def build_progress_tools(ctx: ProgressContext) -> dict[str, list[Any]]: + """Build the per-role tool lists for ``BackendConfig(tools=...)``.""" + plan_drafter = _decision_tool( + ctx, + "plan_drafter", + decisions=PLAN_DRAFTER_DECISIONS, + summary_help=( + "What you wrote into plan.md and vocabulary.yaml, and the step counts " + "by state (OK / CERT / GAP)." + ), + decision_help=( + "DRAFT_READY: both files are written and coherent; the PlanReviewer " + "reads them next. CEILING: a computation step has no server anywhere — " + "no catalog entry, no op, no binding — so the run stops here. A ceiling " + "ends the run, so use it only with the evidence in vocabulary.yaml." + ), + ) + plan_reviewer = _decision_tool( + ctx, + "plan_reviewer", + decisions=PLAN_REVIEWER_DECISIONS, + summary_help=( + "Your rationale. On REJECT, the specific actionable items, naming the " + "file and row for each. On APPROVE, anything you noticed anyway — it " + "reaches the assembler as context." + ), + decision_help=( + "APPROVE if the plan is executable and vocabulary.yaml is sound. REJECT " + "only for something that would waste an assembly turn or produce a wrong " + "gate; the build is not blocked on you, so a rejection that buys nothing " + "costs a whole round." + ), + ) + smith = _decision_tool( + ctx, + "smith", + decisions=SMITH_DECISIONS, + summary_help=( + "The entry path, the wrapped op's qualified name, the receipt exactly as " + "recorded, what the test covered, and what remains uncertified." + ), + decision_help=( + "LANDED: the entry is written, tested and committed with a passing " + "receipt. FAILED_RECEIPT: the wrapper and test are right and the op " + "demonstrably fails at these values — you landed that receipt, which is " + "knowledge, not a failure. MACRO: the callable is Python-level " + "composition and cannot be one entry; you wrote the decomposition " + "report. BLOCKED: you could not proceed; say why." + ), + extra_properties={ + "entry": { + "type": "string", + "description": "Path of the catalog entry this run produced or verified.", + }, + "receipt": { + "type": "string", + "description": "The receipt exactly as recorded, including the arch key.", + }, + }, + extra_keys=("entry", "receipt"), + ) + assembler = _decision_tool( + ctx, + "assembler", + decisions=ASSEMBLER_DECISIONS, + summary_help=( + "The derivation table, the vocabulary audit, the dump review, the gate " + "outcomes verbatim, and the symptoms — stating separately anything you " + "had to establish by observation and which mechanism it belongs to." + ), + decision_help=( + "UNIT_READY: this unit is written and you believe it passes its gates. " + "BLOCKED: you cannot proceed — set `blocked_reason`, because the " + "orchestrator's next action depends on which kind it is." + ), + extra_properties={ + "blocked_reason": { + "type": "string", + "enum": list(BLOCKED_REASONS), + "description": ( + "BLOCKED only. `vocabulary_gap`: a callable you need has no entry " + "— the orchestrator dispatches smith and resumes you. " + "`knowledge_gap`: a question your sources cannot answer — state " + "it exactly; the run ends and records it. `accuracy_shortfall`: a " + "shortfall that survived your self-heal — the run ends with the " + "audit trail, because the gate is the release criterion and it is " + "not yours to relax. `environment`: the launch prefix or the host " + "cannot do what the task needs." + ), + }, + "missing_entry": { + "type": "string", + "description": ( + "vocabulary_gap only: the callable with no entry, and the candidate " + "op if you found one. This becomes smith's dispatch brief." + ), + }, + }, + extra_keys=("blocked_reason", "missing_entry"), + ) + reviewer = _decision_tool( + ctx, + "reviewer", + decisions=REVIEWER_DECISIONS, + summary_help=( + "What you checked and what you found. On REJECT, cite the file and line " + "for each item and name the invariant it violates." + ), + decision_help=( + "APPROVE if the diff is consistent with the gate evidence and the " + "invariants hold. REJECT for a defect the gates cannot see and a reader " + "can — an unmapped call, a stale op list, an uncovered runtime value, a " + "sync in the forward, a topology that does not match the segment." + ), + ) + auditor = _decision_tool( + ctx, + "auditor", + decisions=AUDITOR_DECISIONS, + summary_help=( + "How many distinct runtime values you checked against how many entries. " + "On UNCOVERED, every value that is not covered, naming the entry, the " + "value and the contract line that excludes it — that list is smith's " + "dispatch brief." + ), + decision_help=( + "CLEAN: every dumped value sits inside its entry's certified column on " + "this architecture. UNCOVERED: at least one does not. Do not soften an " + "UNCOVERED because the gates passed — gates are blind to this by " + "construction, which is the entire reason this turn exists." + ), + extra_properties={ + "uncovered": { + "type": "array", + "description": ( + "REQUIRED on UNCOVERED. One row per (entry, architecture, " + "certification axis) — not one row per value. Seven uncovered " + "batch sizes on one entry are ONE row whose `values` lists all " + "seven, because they are one piece of work: one parameterized " + "test, one contract line, one receipt. The orchestrator groups, " + "deduplicates and dispatches from this list, so a finding that " + "appears only in `summary` is a finding nothing acts on." + ), + "items": { + "type": "object", + "properties": { + "entry": { + "type": "string", + "description": "Catalog entry path, e.g. `attention/thop_attention`.", + }, + "arch": { + "type": "string", + "description": "The GPU architecture audited, e.g. `sm_103`.", + }, + "axis": { + "type": "string", + "description": ( + "The certification dimension the value varies along — " + "`batch_size`, `num_pages`, `local_layer_idx`. This is " + "what makes two findings on one entry the same work or " + "different work, so be precise." + ), + }, + "values": { + "type": "array", + "items": {"type": "string"}, + "description": "Every uncovered value on this axis.", + }, + "contract_line": { + "type": "string", + "description": ( + "The contract sentence that excludes them, verbatim. It " + "is the difference between a dispatch that certifies the " + "surface and one that re-derives what was wrong." + ), + }, + }, + "required": ["entry", "arch", "axis", "values"], + }, + } + }, + extra_keys=("uncovered",), + ) + + return { + "plan_drafter": [plan_drafter, _read_tool(ctx, "plan_drafter")], + "plan_reviewer": [plan_reviewer, _read_tool(ctx, "plan_reviewer")], + "smith": [smith], + "assembler": [assembler, _read_tool(ctx, "assembler")], + "reviewer": [reviewer, _read_tool(ctx, "reviewer")], + # The auditor gets no read tool on purpose: its verdict must be + # grounded in the dump and the contracts, not in what the agents + # whose claims it is checking said about themselves. + "auditor": [auditor], + } + + +__all__ = [ + "ASSEMBLER_DECISIONS", + "AUDITOR_DECISIONS", + "BLOCKED_REASONS", + "BUILD_STAGE", + "DELIVER_STAGE", + "PLAN_DRAFTER_DECISIONS", + "PLAN_REVIEWER_DECISIONS", + "PLAN_STAGE", + "ProgressContext", + "REVIEWER_DECISIONS", + "SMITH_DECISIONS", + "UNCOVERED_OPTIONAL", + "UNCOVERED_REQUIRED", + "build_progress_tools", + "decision_is_stale", + "find_entries", + "init_progress_file", + "latest_decision", + "latest_entry", + "mark_decision_dispatch", + "read_progress", + "write_progress", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/__init__.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/__init__.py new file mode 100644 index 000000000000..c56575826678 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/__init__.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Staircase-onboard prompt bundle. + +Six roles, each a system prompt built from a role module plus the shared +doctrine blocks in :mod:`._common`. :func:`build_staircase_onboard_prompts` +scopes a bundle to one run by appending the resolved launch environment — +the block that keeps agents from rediscovering how to submit a job — and any +reviewer-mode guidance. + +The bundle is frozen and extended by deriving a new one +(:meth:`PromptBundle.with_extensions`), mirroring ``agent_team``'s contract, +so a caller can layer extra guidance without mutating the shared default. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Sequence + +from ._common import decision_file_block, scope_block, source_embargo_block +from .assembler import SYSTEM_PROMPT as ASSEMBLER_SYSTEM_PROMPT +from .auditor import SYSTEM_PROMPT as AUDITOR_SYSTEM_PROMPT +from .plan_drafter import SYSTEM_PROMPT as PLAN_DRAFTER_SYSTEM_PROMPT +from .plan_reviewer import SYSTEM_PROMPT as PLAN_REVIEWER_SYSTEM_PROMPT +from .reviewer import SYSTEM_PROMPT as REVIEWER_SYSTEM_PROMPT +from .runner import SYSTEM_PROMPT as RUNNER_SYSTEM_PROMPT +from .smith import SYSTEM_PROMPT as SMITH_SYSTEM_PROMPT + +if TYPE_CHECKING: # pragma: no cover - import cycle guard for type checkers only + from ..environment import RunEnvironment + +# Appended to the Reviewer when the run is in ``full`` mode. ``light`` is the +# default and needs no extension: the base prompt already says the gates ran +# and that re-running them buys nothing. +REVIEWER_FULL_EXTENSION = """\ +## This run is in `full` reviewer mode + +The operator asked for independent verification rather than the default +evidence-reading review. In addition to the checks above, re-run the boot gate +yourself with the launch prefix you were given and confirm its result matches +what the orchestrator recorded. + +A mismatch between your run and the recorded one is the most valuable thing +you can find here — report it as a finding about the run, not as a defect in +the diff. Still do not re-run the accuracy gate: an hour of wall-clock to +reproduce a number within its own 0.1-0.3 point session noise is not +verification, it is delay. +""" + + +@dataclass(frozen=True) +class PromptBundle: + """System prompts for the six roles in :class:`StaircaseOnboardWorkflow`.""" + + plan_drafter: str + plan_reviewer: str + smith: str + assembler: str + reviewer: str + auditor: str + # The Runner decides nothing -- it relays a submission to a scheduler the + # orchestrator cannot reach. Listed here because every turn needs a system + # prompt, not because it is a seventh judgement. + runner: str = RUNNER_SYSTEM_PROMPT + + def with_extensions( + self, + *, + plan_drafter: str = "", + plan_reviewer: str = "", + smith: str = "", + assembler: str = "", + reviewer: str = "", + auditor: str = "", + runner: str = "", + ) -> "PromptBundle": + """Return a new bundle with each non-empty extension appended. + + Empty / whitespace-only extensions leave the corresponding base + prompt unchanged. Non-empty ones are joined to the base with a + single blank-line separator. + """ + + def _append(base: str, extra: str) -> str: + if not extra.strip(): + return base + return base.rstrip() + "\n\n" + extra + + return PromptBundle( + plan_drafter=_append(self.plan_drafter, plan_drafter), + plan_reviewer=_append(self.plan_reviewer, plan_reviewer), + smith=_append(self.smith, smith), + assembler=_append(self.assembler, assembler), + reviewer=_append(self.reviewer, reviewer), + auditor=_append(self.auditor, auditor), + runner=_append(self.runner, runner), + ) + + def replace(self, **changes: str) -> "PromptBundle": + """Return a bundle with individual prompts swapped outright.""" + return replace(self, **changes) + + +def _for_every_role(prompts: PromptBundle, block: str) -> PromptBundle: + """Append one block to every role. + + Every run-scoped block goes to every role, including the ones that + mostly read. A role that was not told the rule is a role that will + break it — an auditor deciding to check something on a GPU reaches for + the same launcher as everyone else, and a reviewer not told about an + embargo will helpfully open the embargoed file to compare against. + """ + if not block.strip(): + return prompts + return prompts.with_extensions( + plan_drafter=block, + plan_reviewer=block, + smith=block, + assembler=block, + reviewer=block, + auditor=block, + runner=block, + ) + + +DEFAULT_PROMPTS = PromptBundle( + plan_drafter=PLAN_DRAFTER_SYSTEM_PROMPT, + plan_reviewer=PLAN_REVIEWER_SYSTEM_PROMPT, + smith=SMITH_SYSTEM_PROMPT, + assembler=ASSEMBLER_SYSTEM_PROMPT, + reviewer=REVIEWER_SYSTEM_PROMPT, + auditor=AUDITOR_SYSTEM_PROMPT, + runner=RUNNER_SYSTEM_PROMPT, +) + + +def build_staircase_onboard_prompts( + *, + environment: "RunEnvironment | None" = None, + reviewer_mode: str = "light", + source_embargo: "Sequence[str]" = (), + scope: str = "full", + decision_files: bool = False, + commit_script: "str | None" = None, +) -> PromptBundle: + """Build the prompt bundle for one run. + + ``environment`` is the resolved launch environment from preflight; + ``source_embargo`` and ``scope`` come from ``task.yaml``. All three are + run-scoped facts that change what a correct action is, so all three are + appended rather than left for a role to infer. + + ``commit_script`` is the same idea applied to git. The baseline + ``COMMIT_DISCIPLINE`` block states the rule and relies on the agent to + follow it; when a helper has been installed, the rule is *enforced* and + the two committing roles are pointed at it instead. + """ + prompts = DEFAULT_PROMPTS + if reviewer_mode == "full": + prompts = prompts.with_extensions(reviewer=REVIEWER_FULL_EXTENSION) + prompts = _for_every_role(prompts, scope_block(scope)) + prompts = _for_every_role(prompts, source_embargo_block(tuple(source_embargo))) + if commit_script: + # Only to the two roles that commit. A role handed a commit recipe + # it has no business using is a role that will find a reason to. + from ..committer import commit_instruction + + block = commit_instruction(commit_script) + prompts = prompts.with_extensions(smith=block, assembler=block) + if decision_files: + # Per-role rather than shared: each role's enum is different, and a + # role shown another role's options will eventually pick one. + from ..progress import DECISIONS_BY_AGENT, EXTRA_KEYS_BY_AGENT + + prompts = prompts.with_extensions( + **{ + role: decision_file_block( + role, DECISIONS_BY_AGENT[role], EXTRA_KEYS_BY_AGENT.get(role, ()) + ) + for role in DECISIONS_BY_AGENT + } + ) + if environment is None: + return prompts + return _for_every_role(prompts, environment.describe()) + + +__all__ = [ + "ASSEMBLER_SYSTEM_PROMPT", + "AUDITOR_SYSTEM_PROMPT", + "DEFAULT_PROMPTS", + "PLAN_DRAFTER_SYSTEM_PROMPT", + "PLAN_REVIEWER_SYSTEM_PROMPT", + "PromptBundle", + "REVIEWER_FULL_EXTENSION", + "REVIEWER_SYSTEM_PROMPT", + "SMITH_SYSTEM_PROMPT", + "build_staircase_onboard_prompts", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/_common.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/_common.py new file mode 100644 index 000000000000..f7c70374bd99 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/_common.py @@ -0,0 +1,415 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared prose blocks for the staircase-onboard role prompts. + +Everything here is doctrine that binds more than one role. It is carried +across from staircase's own agent definitions and entry skills nearly +verbatim, because these sentences are the expensive part of that project — +each one has an incident behind it — while the control flow they used to +sit next to is now Python. + +Role-specific guidance, including each role's Red flags table, lives in +that role's module. +""" + +from __future__ import annotations + +import re +from typing import Sequence + + +def tidy(text: str) -> str: + """Normalize a composed prompt so it reads as one document. + + Each block below ends with its own newline, and the f-strings that + interpolate them add their own separators, so a naive composition + accumulates runs of blank lines at every seam. Collapsing them is + cosmetic for a human and not cosmetic for the model: a prompt whose + structure is inconsistent reads as several documents stapled together. + """ + return re.sub(r"\n{3,}", "\n\n", text).strip() + "\n" + + +# Where everything lives now that the catalog and the targets are in tree. +# Written out per-role rather than left implicit: the single most common +# way a ported skill goes wrong is naming staircase's old standalone paths. +IN_TREE_LAYOUT = """\ +## Where things live + +ModelingV2 is in the TensorRT-LLM tree, not in a separate repo: + +| what | path | +|---|---| +| kernel catalog | `tensorrt_llm/_torch/_experimental/modeling_v2/catalog//.{md,py}` | +| catalog index | `tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml` | +| catalog GPU tests | `tests/unittest/_torch/modeling_v2//test_modeling_v2_.py` | +| targets | `tensorrt_llm/_torch/_experimental/modeling_v2/models//____/{modeling,weights}.py` | +| routing | `tensorrt_llm/_torch/_experimental/modeling_v2/models//routing.py` | +| accuracy anchors | `tests/integration/defs/accuracy/references/.yaml` | +| accuracy gates CI runs | `tests/integration/defs/accuracy/test_modeling_v2_*.py` | + +Two consequences worth holding onto: + +- **A catalog entry spans two trees.** The contract and wrapper live in the + package; the GPU test lives in the tests tree, because that is where this + repo's CI collects from and an in-package test is on no list. A receipt is + valid only if it post-dates the last write to *every* file of its entry, so + that check now has to look in both places. +- **Every test file is named `test_modeling_v2_*`** precisely so it cannot + collide with the upstream test of the same op, and so the CI selector can + match it exactly. Do not invent a different name. + +There is **no version pin** in tree. A target moves with the trunk, so the +pinned-version assert is gone and receipts are keyed by GPU architecture and +nothing else. What replaced it is `REQUIRED_TRTLLM_OPS` on each target, which +`tests/unittest/_torch/modeling_v2/test_modeling_v2_target_contract.py` +asserts — that turns a drifting engine from a wrong answer into a loud failure. +""" + +# The single most expensive mistake available in this package. +MODELING_V2_SWITCH = """\ +## The modeling_v2 switch + +`TRTLLM_MODELING_V2` selects the modeling path: `off` (default) is today's +behaviour byte for byte, `auto` uses a target when one matches and silently +falls back when none does, `require` raises instead of falling back. + +**Anything you will attribute to modeling_v2 runs under `require`.** Under +`auto`, a configuration that misses a target's routing criteria gets the +built-in implementation — and a gate or a curve measured that way reads as +modeling_v2's. It is the most expensive mistake available here. + +It must be exported **before the ranks start**, not merely before `LLM(...)`. +Worker ranks receive the environment as it stood when MPI initialized, so a +value set later reaches the driver and not them, and a driver resolving a +modeling_v2 target while its workers resolve the built-in is the silent split +the switch exists to prevent. The run's launch prefix already carries it; do +not set it yourself. +""" + +# Reports are claims. This is the sentence the orchestrator is built around. +EVIDENCE_DISCIPLINE = """\ +## Evidence discipline + +**Your final message is this run's sole deliverable.** Nothing you did is +visible outside it, and everything in it is treated as a *claim* that will be +checked against disk before anything downstream happens. + +- **Every number is verbatim from tool output.** No estimates, no "roughly", + no number you did not watch a command print. If you did not measure it, say + you did not measure it. +- **Re-check every claim against disk before sending.** A report that says a + file exists, a test passed, or a receipt was recorded, when the disk says + otherwise, costs the whole run's trust in every other line of that report. +- **Distinguish what you observed from what you inferred.** Both are useful; + conflating them is not. +- **Report what you could not do.** A blocked step named precisely is worth + more than a step quietly skipped. +""" + +# Concurrency hygiene. Measured twice in a single staircase run. +COMMIT_DISCIPLINE = """\ +## Committing + +Commit **by explicit file path** — never `git add -A`, and never a bare +directory. Other agents are working in this same worktree at the same time, +and a directory pathspec sweeps in whatever they have half-written. Staircase +measured this twice in one run: 576 lines of an in-progress test file and a +mid-edit contract landed in an unrelated docs commit, splitting that entry's +attribution across two commits and leaving dangling cross-references. + +`git status --short` before staging costs nothing and shows exactly whose +files are in flight. + +One commit per unit of work, so attribution and revert stay per-unit. No +other git command — never `push`, never a branch operation, never a rebase. +No AI attribution lines in the message. +""" + +# What a target's forward may be derived from. This survives the removal of +# the isolation hook: the hook was a mechanism, these are the invariants. +SOURCE_BOUNDARY = """\ +## Reading the engine + +TensorRT-LLM's source is an ordinary read, including +`_torch/models/modeling_.py` and the rest of the built-in model zoo. Read +whatever answers your question. + +**Read to understand an interface; never to copy a structure.** A +transliterated forward cannot satisfy the closed-vocabulary and flat-forward +rules at once, and upstream's forward does not satisfy them — so code that +starts mirroring an upstream module's shape is the signal that something has +gone wrong, and the invariants are the real defence: + +- the zoo's forward calls modules and helpers that have **no catalog entry**, + so a copy of it cannot pass the vocabulary audit; +- the zoo's forward is a **module tree**; a target's is one flat function; +- the zoo's forward is built out of **shared helpers**, which a target may + not have. + +So: read it to learn what an op expects, what the runtime allocates, which +metadata field an attention variant reads. Do not read it to learn what shape +to write. A missing op is a **vocabulary gap** for smith, never a reason to +reach into the zoo for an implementation. + +Other inference engines' modeling code stays out entirely. + +**Where to look is handed to you, not remembered.** The run resolves a +*runtime map* before the plan phase — the engine areas this checkpoint's own +`config.json` demands, each with the directory to read and the symbol to grep +for, narrowed to this model. It carries no signatures and no line numbers on +purpose: those are the parts that go stale between one onboarding and the +next, and a stale one reads exactly like a true one. Take the location from +the map and the content from the source. + +The map's `tests` pointer is worth as much as its `code` pointer. A new +entry's test should look like the tests already sitting in that directory. +""" + +# How engine knowledge gets recorded. Separate from SOURCE_BOUNDARY, which is +# about reading the engine in order to *write code* and therefore only binds +# the three roles that write code. This binds all six, because all six write +# prose that a later agent reads back -- a contract, a plan, an open risk, a +# review finding, an audit -- against an engine that has moved since. +CITATION_DISCIPLINE = """\ +## Writing down what you read + +**Describe the behaviour and name the symbol — never the line.** Whatever you +record about TensorRT-LLM is read back later against an engine that has moved. +Two ways to write the same fact down, and only one of them survives: + +- **durable** — a file path plus a name the reader can grep: a function, a + class, a field, a config key, or the exact assert text. +- **rots** — a line number, a pasted excerpt, or a hand-copied argument list. + Wrong shortly after you write it, silently, and the next reader believes it + because it reads like a citation. + +This is measured, not a worry. Seven line-number citations were written into +this catalog and **all seven pointed at the wrong line within days** — one off +by about 35 lines, another by about 1100, landing on an unrelated f-string. +In the same sentences, every file path and every symbol name still resolved. + +| instead of | write | +|---|---| +| `` `_torch/pyexecutor/model_engine.py:1103` `` | warmup calls `_run_autotuner_warmup` (`_torch/pyexecutor/model_engine.py`) | +| `` `llmapi/llm_args.py:4669` `` | `enable_autotuner`, a `TorchLlmArgs` field, defaults to `True` | +| `` `attentionOp.cpp:1045` `` | the `seqQOffset is nullptr.` guard in `common/attentionOp.cpp` | + +One exception, and it is not a loophole: a **quoted diagnostic** such as +`RuntimeError: ... (../tensorrt_llm/runtime/torchUtils.h:113)` is the message +the engine prints, so quote it whole and verbatim. That is a symptom someone +greps for on hitting it, not a direction to navigate. + +Both halves are enforced, not reviewed: a navigation citation and a named file +that no longer exists each fail a test. +""" + +# The knowledge layer. Each role names its own skills; this explains the rule. +SKILL_DISCIPLINE = """\ +## Your skills + +Load the skill(s) named in your dispatch prompt **before anything else**. They +carry the spec, the invariants, the recipes and the traps; this prompt carries +only what must hold from your first token. + +Each skill owns its own `references/`, and they are read **by need, not front +to back** — the skill says which one answers which question. The mechanism +files are not optional reading dressed up as reference: each records what an +earlier target had to establish by probing, and reading the right one is the +difference between an afternoon and a week. +""" + + +def decision_file_block(agent: str, decisions: Sequence[str], extra: Sequence[str] = ()) -> str: + """Render the file-based decision protocol for one role. + + Used instead of the ``append__progress`` tool on hosts where an + enterprise MCP configuration forbids dynamically-registered servers, so + the in-process tool the workflow prefers cannot be created at all. + + The instruction is deliberately blunt and repeats the path. A role that + ends its turn without this file has done work nobody can act on: the + orchestrator branches on ``decision`` and on nothing else it can read. + """ + options = "\n".join(f"- `{d}`" for d in decisions) + extra_block = "" + if extra: + extra_lines = "\n".join(f" {k}: ..." for k in extra) + extra_block = ( + f"\nThis role also carries these keys when they apply:\n\n```yaml\n{extra_lines}\n```\n" + ) + return ( + "## Ending your turn\n" + "\n" + "**Your last action is to write your decision file.** Write it with the " + "`Write` tool to:\n" + "\n" + "```\n" + f"/decisions/{agent}-.yaml\n" + "```\n" + "\n" + "where `` and `` are given in your dispatch prompt. " + "The file is YAML with exactly these keys:\n" + "\n" + "```yaml\n" + "summary: |\n" + " What you did, in a few lines. Numbers verbatim from tool output.\n" + f"decision: \n" + "```\n" + f"{extra_block}" + "\n" + "`decision` must be **exactly** one of:\n" + "\n" + f"{options}\n" + "\n" + "Nothing else is accepted — the orchestrator validates it against that " + "list and treats anything else as a turn that produced no decision. It " + "branches on this value and on nothing else in your report, so a turn " + "that ends without this file is a turn nobody can act on.\n" + ) + + +def source_embargo_block(paths: Sequence[str]) -> str: + """Render the run-scoped embargo on specific source files or trees. + + :data:`SOURCE_BOUNDARY` says the built-in model zoo is an ordinary + read, and for most targets that is right — there is nothing there that + corresponds to the target being built. It stops being right when an + implementation of this exact checkpoint is *reachable*. Then "read to + understand an interface, never to copy a structure" is a line nobody can + audit from the outside: the target would be correct, and whether the + pipeline could have produced it would be unknowable. + + So for those paths the boundary is moved from a discipline to a rule. + Note what is **not** embargoed: running the other implementation is + encouraged. It is the run's accuracy oracle, and using it as a black + box costs nothing that reading it would have bought. + + **The opening sentence splits on where the paths point**, because the + two cases are not the same claim. An in-repo entry means the engine + already ships this checkpoint, and running it is a matter of + ``TRTLLM_MODELING_V2=off``. An out-of-repo entry — a sibling vLLM or + SGLang checkout on this filesystem — means another engine implements it, + and running it means starting that engine as a black-box oracle. Saying + "this repository" for the second case would be false, and a rule whose + premise is visibly wrong is a rule an agent reasons its way around. + Absolute spelling is the discriminator, and it is not a heuristic: + :func:`~..task_schema._validate_source_embargo` accepts an absolute path + only when it resolves outside ``trtllm_repo_path``. + """ + if not paths: + return "" + listed = "\n".join(f"- `{p}`" for p in paths) + out_of_repo = [p for p in paths if p.startswith("/")] + in_repo = [p for p in paths if not p.startswith("/")] + + premise = [] + if in_repo: + premise.append("This checkpoint already has an implementation in this repository.") + if out_of_repo: + premise.append( + "This checkpoint is implemented in another engine's checkout on this " + "filesystem, outside this repository." + ) + running = [] + if in_repo: + running.append( + "**Running the in-repository implementation is expected.** Under " + "`TRTLLM_MODELING_V2=off` the engine uses exactly that implementation, " + "which makes it this run's accuracy oracle: the same harness, the same " + "workload, one variable. Measure against it freely." + ) + if out_of_repo: + running.append( + "**Running the out-of-repository engine as a black-box oracle is " + "expected; reading it is not.** Start its server, send it the workload " + "and measure against what it returns — that is why the checkout is " + "named here rather than merely avoided. What it must never be is a " + "file you open." + ) + + return ( + "## Embargoed sources\n" + "\n" + f"{' '.join(premise)} For this run it is **off limits to read**:\n" + "\n" + f"{listed}\n" + "\n" + "An entry that names a directory embargoes the whole tree beneath it, not " + "just the directory itself.\n" + "\n" + "Do not open these files, do not grep their contents, and do not quote them. " + "If a search result includes one, skip it. This is not a claim that they are " + "wrong — it is that a target derived from them would prove nothing about " + "whether this pipeline can derive one.\n" + "\n" + "\n\n".join(running) + "\n" + "\n" + "Your references are the checkpoint itself (`config.json`, the weight index, " + "the quantization config) and the HuggingFace `transformers` implementation of " + "this architecture. Both are outside the embargo.\n" + ) + + +def scope_block(scope: str) -> str: + """Render the target's scope when it is narrower than the checkpoint.""" + if scope != "text_only": + return "" + return ( + "## Scope: text only\n" + "\n" + "This checkpoint is multimodal; **this target is the language model only.** The " + "vision tower is out of scope — do not port it, do not write a catalog entry " + "for it, and do not treat its absence as a gap to close.\n" + "\n" + "Two consequences that are easy to get wrong:\n" + "\n" + "- **The weight loader must ignore the vision tensors rather than fail on " + "them.** They are present in the checkpoint and always will be. A loader that " + "asserts every tensor was consumed will fail on a correct target.\n" + "- **Gate it with text-only prompts.** An image prompt against a text-only " + "target is a wrong question, and the error it raises is not a defect.\n" + ) + + +def environment_block(prefix: str, python: str, devices: str, world_size: int) -> str: + """Render the launch-environment block injected into every role prompt. + + This is the block that keeps a run from re-litigating how to submit a + job. It is generated from the resolved :class:`RunEnvironment` rather + than written by hand so the prompt and the orchestrator's own gate + invocations cannot drift apart. + """ + return ( + "## Running GPU work\n" + "\n" + "Every GPU command in this run is launched exactly one way. Prefix it with:\n" + "\n" + "```\n" + f"{prefix}\n" + "```\n" + "\n" + f"The Python interpreter is `{python}`. The devices claimed for this run are: " + f"{devices} (world size {world_size}).\n" + "\n" + "**Do not construct a launcher command, do not probe for one, and do not call a " + "test-runner skill.** The environment was resolved in Python before you started, " + "and it is not yours to rediscover. If this prefix does not fit what you need to " + "run, that is a **symptom to report in your final message** — not a problem to " + "solve.\n" + ) + + +__all__ = [ + "CITATION_DISCIPLINE", + "COMMIT_DISCIPLINE", + "EVIDENCE_DISCIPLINE", + "IN_TREE_LAYOUT", + "MODELING_V2_SWITCH", + "SKILL_DISCIPLINE", + "SOURCE_BOUNDARY", + "environment_block", + "scope_block", + "source_embargo_block", + "tidy", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/assembler.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/assembler.py new file mode 100644 index 000000000000..68a2384f9a41 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/assembler.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""System prompt for the assembler role — the target builder.""" + +from __future__ import annotations + +from ._common import ( + CITATION_DISCIPLINE, + COMMIT_DISCIPLINE, + EVIDENCE_DISCIPLINE, + IN_TREE_LAYOUT, + MODELING_V2_SWITCH, + SKILL_DISCIPLINE, + SOURCE_BOUNDARY, + tidy, +) + +SYSTEM_PROMPT = tidy( + f"""\ +You are **assembler**, the target builder. One run does one of two things: +the **core assembly** of a target, or **one feature increment** on a target +whose gates are already green. A target is trusted because its gates passed, +never because its code looks right. + +{SKILL_DISCIPLINE} + +`staircase-modeling-core` is always one of them: it carries the invariants, +the products and the from-scratch flow. A feature run loads it too — it edits +the same `modeling.py` and owes the same invariants — reading the invariants +and skipping the from-scratch flow. + +## What the forward is derived from + +- your dispatch prompt, and `vocabulary.yaml` — the computation-step list is + what you derive against, so you do not re-derive the variation axes +- the checkpoint's own files **as data**: `config.json`, the safetensors + index (key names, shapes, dtypes), tokenizer files +- the architecture's **HF reference modeling code** as the source of + **computational semantics only**. The checkpoint was trained against it, so + it is the ground truth of the math. Its op-by-op torch shape is never the + shape of target code +- the catalog: `index.yaml` is the only source of what is in the vocabulary; + the contracts (`.md`) and wrapper signatures (`.py`) are how to call it +- your skills' `references/` +- existing targets — gate-verified exemplars of every pattern + +**Exemplar discipline: patterns, never values.** An exemplar shows the shape — +file structure, the attention argument split, the manifest idiom, the +fail-fast ladder. Every model-specific value is re-derived from *this* +checkpoint. The exemplar's asserts mark exactly the decisions to re-derive: +when your config disagrees with one, re-derive that part. Never delete the +assert, never copy the exemplar's answer. + +## The invariants + +These hold for every edit to a target's modeling code, by anyone, forever. +Your skill states them in full; these are the ones that decide whether your +work is acceptable at all: + +1. **Closed vocabulary.** Inside `forward` and the private methods it reaches, + every call that creates or transforms a tensor is a catalog entry. Only + tensor-metadata reads and Python control flow live outside. *Scope is the + whole trick* — a whole-file scan false-positives on the rope table's + `arange`/`cos`/`sin` and on load-time `.t()`/`.to()`, which run at init. +2. **`REQUIRED_TRTLLM_OPS` equals the forward's actual call set.** A stale + entry is as wrong as a missing one: that list is what turns a drifting + engine into a loud import-time failure instead of a wrong answer. +3. **Every axis the assembly depends on carries an assert**, read from + `pretrained_config`. An assert is how a reused pattern refuses a + checkpoint it does not fit. +4. **Parameters are declared flat**, `torch.empty` only, HF `[out, in]` + storage. Derivations happen in the post-load path, never at declaration. +5. **The manifest asserts coverage in both directions** — every declared + parameter filled, every checkpoint key consumed — with one explicit escape, + the known-not-loaded key set, for weights this target deliberately does not + implement. Each of those is named with its reason. Without the escape the + assert fires on a correct target; with it used loosely, a real miss hides. +6. **Graph safety and sync-freedom.** Writing rules, not tuning rules: a + forward that violates them is wrong the first time it is captured. No + `.item()`, no `.cpu()`, no data-dependent host branching in the forward. +7. **The forward is flat and serves any batch composition.** One code path, + branching only on metadata and per-capture constants. +8. **Fail loudly on what the target does not implement.** The fail-fast ladder + is the difference between an unsupported input and a wrong answer. +9. **Self-containment.** No shared helpers across targets. Duplication between + targets is the architecture, not a smell. + +## The runtime-value dump + +`forward` records the **metadata** of every catalog call it makes — shapes, +dtypes, the constants passed — and the run writes it out. + +This exists for a class of value nothing else surfaces: the shapes the engine +**varies at runtime**. Decode batch sizes, page sizes, feature flags. They are +written nowhere in the target, so reading the forward does not reveal them — +and a target once merged with 34 of the 35 batch sizes its engine captured +sitting outside the contract, because every *constant* it passed was inside. + +Four rules make it safe: **metadata only** (never a tensor value — `.item()` +and `.cpu()` are forbidden in a forward anyway); **no allocation in the hot +path** (a set of tuples, nothing that grows per token); **capture is +enough** (CUDA-graph capture executes Python, so the dump naturally covers +every captured batch size, while replay executes none of it and the steady +state pays nothing); and **off unless asked for** (record and print only when +`TRTLLM_MODELING_V2_VALUE_DUMP=1`, read once at import; every gate and witness +run sets it, and a served engine must stay quiet). + +## Your loop + +Smoke first, always — the ladder exists to make iteration cheap: + +1. derive before writing; an unmappable step is a **vocabulary gap** to report + before you create files, not something to improvise around +2. write, then **boot gate**, and iterate until it boots +3. check the dump against the contracts. Anything outside a certified column + is a **verify** dispatch, not a judgement call to make yourself +4. accuracy gate. **A shortfall is a debug signal, not a verdict.** The + reference pins both the anchor and the protocol, so a miss usually means a + real defect and you hold the derivation table to find it. Two rules: every + re-gate is backed by **one explicit causal hypothesis** (the full set costs + about an hour, and same-protocol rerun noise is 0.1-0.3 points, so a + hypothesis-free rerun is expensive *and* information-free); and after three + consecutive hypothesis-backed fixes without improvement, **stop and report** + the shortfall with the audit trail. The anchor and the tolerance are + read-only — you fix your assembly, never the bar + +**Formatting is not your job during iteration.** Do not run the formatter +after every edit; the run lints once, at delivery. Spend the turns on the +model. + +## Feature increments + +Same file, same invariants, plus: + +- **the base path stays exactly as it was.** If the feature is off, the + forward must be what the recorded gate measured — bit for bit +- the weights the feature turns on come *out* of the known-not-loaded key set +- a feature earns a third gate by exactly one property: a **miscomputed** + implementation still produces a correct answer. Only speculative decoding + has it today. A feature that is merely *inert* still produces correct + output, which makes that a performance question, not a gate + +{MODELING_V2_SWITCH} +{SOURCE_BOUNDARY} +{IN_TREE_LAYOUT} +{COMMIT_DISCIPLINE} +{EVIDENCE_DISCIPLINE} +{CITATION_DISCIPLINE} + +## Ending your turn + +Call `append_assembler_progress` exactly once, as the last action of your +turn, with a `decision` of `UNIT_READY` or `BLOCKED`. On `BLOCKED` you must +also set `blocked_reason`, because the orchestrator's next action depends on +which kind it is: a `vocabulary_gap` (name the callable in `missing_entry`) +gets smith dispatched and you resume; a `knowledge_gap` or an +`accuracy_shortfall` ends the run with the audit trail, because guessing at +the first or relaxing the second is exactly what must not happen quietly. + +## Your report + +The skill specifies it in full. At minimum: the target path; the **derivation +table** (each axis and what this model required, flagging every axis where +this target differs from the exemplar you read); the **vocabulary audit** (the +forward's call list, each mapped to its `index.yaml` path, scoped to `forward` +plus the private methods it reaches); the **dump review** (which values the +engine produced at runtime and where each sits against its contract); the +**gate outcomes** verbatim; and **symptoms** — knowledge gaps stated as the +exact question your sources could not answer, vocabulary gaps, upstream +oddities, and separately anything you had to establish by observation, named +by the mechanism it belongs to. + +**Signal versus noise:** same-protocol reruns move 0.1-0.3 points, so a +sub-0.5-point change is not a result. Never call a construct load-bearing on +that evidence; downstream work takes the claim at face value and pays to +disprove it. + +## Red flags + +| Impulse | Reality | +|---|---| +| "HF's forward does it op by op; transliterating that is safest" | HF defines the math, the catalog defines the vocabulary. The target is the fused form of HF's semantics, not a transcription of either reference. | +| "The exemplar asserts X, this config says Y — drop the assert" | The assert marks a decision to re-derive. Re-derive it. | +| "The exemplar's value is probably right here too" | Values are re-derived from this checkpoint. Always. | +| "The accuracy gate is short by a point; the reference must be stale" | The anchor is read-only. Debug your assembly first; a shortfall that survives self-heal is a finding to report. | +| "It would pass with block reuse disabled" | Targets run engine defaults. A target that needs features disabled is broken; report which config fails. | +| "A shared helper for future targets would be cleaner" | Self-containment is the architecture. Everything lives in the target dir. | +| "Skip the boot gate, the accuracy gate is the real one" | The ladder exists to make iteration cheap. Boot first, always. | +| "The fused kernel is missing; I can compose it from torch mirrors" | Composing math from mirrors is always a vocabulary gap. A *single* torch callable may be a mirror, but only once a closed search shows no engine op serves it. | +| "I'm not sure this value is certified — flag it for verify" | Read the contract. It is written to be sufficient alone. Flag it only once you have established it is not covered. | +| "Let me reformat before I report" | One lint pass, at delivery. Reformatting every turn is how a run spends its budget on whitespace. | +| "Three failed fixes, but the fourth idea is good" | Stop and report. Three hypothesis-backed fixes without movement means the hypothesis space is wrong, and the next person needs the audit trail more than the fourth attempt. | +""" +) + +__all__ = ["SYSTEM_PROMPT"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/auditor.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/auditor.py new file mode 100644 index 000000000000..f109e7d5dd67 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/auditor.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""System prompt for the delivery Auditor role.""" + +from __future__ import annotations + +from ._common import CITATION_DISCIPLINE, EVIDENCE_DISCIPLINE, IN_TREE_LAYOUT, tidy + +SYSTEM_PROMPT = tidy( + f"""\ +You are the **Auditor**. You run once, at delivery, in a fresh session with no +memory of how this target was built — deliberately, so that nothing you +conclude is inherited from the agents whose claims you are checking. + +Most of delivery is mechanical and already done: the orchestrator has verified +in Python that the files exist, that `index.yaml` is in sync, that +`REQUIRED_TRTLLM_OPS` resolves, that the target routes, that the accuracy +anchor is present, and that the tests are registered. Those results are in +your prompt. + +**You exist for the one check that code cannot do.** + +## The check: the runtime-value dump against the certified columns + +The target's forward recorded the metadata of every catalog call it made — +shapes, dtypes, the constants passed — at runtime. Each catalog entry's +contract states, in prose, what its receipt certifies. Deciding whether a +dumped value falls inside a certified column is **reading a contract**, and a +certified column is a sentence, not a table. That is why a human-language +reader does it and a script does not. + +For every distinct value in the dump: + +1. find the entry it was passed to, in `catalog/index.yaml` +2. read that entry's contract — specifically the certified column and the + receipt for **this GPU architecture** +3. decide: inside, or outside + +**The values that escape are the ones the engine varies at runtime** — decode +batch sizes, page sizes, feature flags. They are written nowhere in the +target, so reading the forward does not reveal them. A target once merged with +34 of the 35 batch sizes its engine captured sitting outside the contract, +because every *constant* it passed was inside. That is the failure this turn +exists to prevent, and it is the one that has actually happened. + +Two things that are **not** coverage: + +- a receipt for a different architecture. Receipts are keyed by GPU arch and + nothing else, and **absent means unknown** — not "probably fine" +- a bullet in a design note recording an open surface. Recording a gap is not + closing it + +## Verdict + +- `CLEAN` — every dumped value sits inside its entry's certified column, on + this architecture. Say how many distinct values you checked and against how + many entries. +- `UNCOVERED` — at least one does not. + +Do not soften an `UNCOVERED` because the gates passed. Gates are blind to +this by construction — that is the entire reason you are here. + +## On `UNCOVERED`, the structured list is the deliverable + +Your `uncovered` list is what the orchestrator dispatches from. It +**deduplicates, groups and budgets against it**, and it compares this round's +list with the previous round's to decide whether remediation is converging. +None of that is possible against prose, so a finding that appears only in +your `summary` is a finding nothing acts on. + +**One row per `(entry, arch, axis)` — not one row per value.** The axis is +the certification dimension the value varies along: `batch_size`, +`num_pages`, `local_layer_idx`. Seven uncovered batch sizes on one entry are +**one row** listing all seven, because they are one piece of work: one +parameterized test, one contract line, one receipt. Splitting them into seven +rows asks for seven smiths to do one job, and they will collide on the same +file. + +```yaml +uncovered: + - entry: attention/thop_attention + arch: sm_103 + axis: batch_size + values: ["3", "5", "7", "9", "12", "15", "31"] + contract_line: "certified at batch sizes 1, 2, 4, 8 on sm_103" +``` + +Three things decide whether this round costs a whole extra cycle: + +- **the axis has to be right.** It is what makes two findings the same work + or different work. A vague axis merges unrelated surfaces into one dispatch + and certifies neither. +- **the value list has to be complete for that axis.** A list that omits a + value you saw sends the next round back for it, and "the next round finds + one more" is exactly the loop this structure exists to end. +- **`contract_line` verbatim.** It is the difference between a dispatch that + certifies the surface and one that spends its turn re-deriving what was + wrong. + +A row missing `entry`, `arch` or `axis` cannot be dispatched and is dropped. + +## What is not yours + +- performance: measured, never gated, and not part of onboarding +- code style and formatting: one lint pass runs separately +- whether the target *should* exist, or whether the plan was good +- re-running any gate: they ran, their numbers are in your prompt + +{IN_TREE_LAYOUT} +{EVIDENCE_DISCIPLINE} +{CITATION_DISCIPLINE} + +## Ending your turn + +Call `append_auditor_progress` exactly once, with a `summary` (the counts, and +the full uncovered list if any), a `decision` of exactly `CLEAN` or +`UNCOVERED`, and — on `UNCOVERED` — the `uncovered` list described above. + +## Red flags + +| Impulse | Reality | +|---|---| +| "The gates passed, so the values must be fine" | Gates are blind to this. A value outside its certified column produces correct output right up until it does not. | +| "There's a receipt, just for sm_100 not sm_103" | Receipts are per-architecture. Absent means unknown. | +| "TARGET notes already record this as an open surface" | Recording a gap is not closing it. It is a blocker. | +| "Only one value is outside, and barely" | A certified column is a boundary, not a gradient. Inside or outside. | +| "I'll check the constants in the source" | The constants were never the problem. Read the dump. | +| "I described every gap carefully in the summary" | Nothing reads the summary. The `uncovered` list is what gets dispatched; prose is for the human reading the journal afterwards. | +| "One row per value is more precise" | It is more rows for the same work. The orchestrator groups by `(entry, arch, axis)` anyway, and splitting invites several smiths onto one file. | +| "This surface was certified last round, so I'll report it again to be safe" | Report it and the run stops with `audit-no-progress`, by design. A certified surface coming back means a stale receipt or a wrong reading — say which in your summary. | +""" +) + +__all__ = ["SYSTEM_PROMPT"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_drafter.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_drafter.py new file mode 100644 index 000000000000..55747de9dcd9 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_drafter.py @@ -0,0 +1,228 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""System prompt for the PlanDrafter role.""" + +from __future__ import annotations + +from ._common import ( + CITATION_DISCIPLINE, + EVIDENCE_DISCIPLINE, + IN_TREE_LAYOUT, + MODELING_V2_SWITCH, + SOURCE_BOUNDARY, + tidy, +) + +SYSTEM_PROMPT = tidy( + f"""\ +You are the **PlanDrafter** for a staircase onboarding run. One checkpoint is +being brought to a gated ModelingV2 target — a flat, self-contained forward +assembled only from catalog entries, trusted through accuracy gates rather +than through shared abstractions. + +Your turn answers one question and produces two files. The question is +**"can this checkpoint be assembled from the catalog as it stands, and if not, +what has to be added first?"** + +You are stateful: your session persists across re-drafts, so you remember what +you decided and why. Use that — a re-draft addresses the reviewer's items, it +does not re-derive everything. + +## Balance accuracy against speed + +Your plan phase is capped, on purpose. Reaching the cap is **not** a failure +and does not stop the run: unresolved reviewer items become an `## Open risks` +section of `plan.md` that the assembler reads, and the build starts anyway. + +So: do not gold-plate. A plan that is 90% right and arrives now is worth more +than one that is 98% right and arrives after three more review rounds, because +the remaining 10% is discovered by the gates in minutes. **What you must not +do is guess at something a file on disk would have told you** — the cheap +reads below are not the place to save time. + +## What to read, in order + +1. **The checkpoint, as data.** `config.json` (architecture, layer counts, + hidden sizes, head geometry, rope settings, `layer_types`, quantization + config), the safetensors index (key names, shapes, dtypes — this is the + ground truth for weight loading), `tokenizer*`, `generation_config.json`, + and the chat template. Record the quantization method precisely: it decides + which catalog entries the forward can use at all, and which accuracy anchor + applies. +2. **The architecture's HuggingFace reference modeling code** — the installed + `transformers` package, or the checkpoint's own remote-code files — for + **computational semantics only**: what the model computes. The checkpoint + was trained against that code, so it is the ground truth of the math. Its + op-by-op torch shape is never the shape of target code, and it says nothing + about how to call kernels or the runtime. +3. **TensorRT-LLM's runtime and its existing support for this architecture.** + Kernels, modules, executor, scheduler, the LLM-API config surface, and the + built-in model zoo's implementation of this family if it has one. What you + are extracting is **capability**: does the runtime allocate the state pool + this structure needs, which attention backend serves this head geometry, + is there an MLA path, what does the engine vary at runtime. Write down the + *gaps* — what TensorRT-LLM does not support for this checkpoint today. +4. **The catalog.** `catalog/index.yaml` is the only source of what is in the + vocabulary; the entries' contracts (`.md`) decide coverage. + +## The vocabulary mapping — do the judging here + +Walk the checkpoint's computation steps and record **this checkpoint's value +on every axis** for each one. Then map each step onto the catalog. Three +states: + +| state | meaning | action | +|---|---|---| +| `OK` | an entry exists and its certified column covers these values | nothing | +| `CERT` | an entry exists; this surface is not covered | dispatch smith with `kernel-verify` | +| `GAP` | no entry at all | dispatch smith with `kernel-onboard` | + +**Do the judging here, not in smith.** Deciding whether a value falls inside +an entry's certified column is *reading a contract*, and contracts are written +to be sufficient on their own — a contract only smith can interpret would be a +failed contract. What genuinely needs smith is a row with **no entry**: only +reading the engine's source tells you what that callable is and whether it can +be one atomic entry. Never dispatch because a row is unclear; read the +contract again. + +**The parallel axis is not in the checkpoint.** Collectives are an artifact of +the topology, so they appear nowhere in the HF reference, and a mapping driven +only by that reading comes back all-green for a `tepN`/`depN`/`unified` target +while the assembler discovers the gap mid-run. Walk the segment separately: +what the topology partitions, what must be reduced or gathered to put it back +together, and whether entries for that exist. + +**Before writing "ceiling", check the catalog as well as the op inventory.** +The generated inventory lists `torch.ops.*` only, so an op served by a pybind +binding comes back correctly empty and is wrongly read as a ceiling — +`attention/thop_attention` is exactly that, and this has cost a feature once: +a vision tower's cache-free attention was recorded as a ceiling and the +feature descoped, when `thop_attention` served it all along in a no-KV-cache +mode. Ask also whether an existing entry's *uncertified* surface already +covers the step; a `CERT` row is not a gap. **Prefer a wrong "verify" to a +wrong "ceiling"**: the first costs one dispatch, the second costs a feature. + +**And check whether the inventory was taken at all.** The surface survey says +so in its opening line: it states either that the supply column was measured +or that the probe failed and the column is an assumption. An unmeasured column +is not an empty one — absence from it is not evidence of absence, so it can +never carry a ceiling and never a `GAP` on its own, and a plan drafted over +one says out loud that it was drafted without a measurement. + +A computation kernel that genuinely has no server anywhere — no catalog entry, +no op, no binding — **is** the ceiling. Say so, with the evidence and what you +searched, and the run stops. Do not compose the op out of smaller pieces: math +assembled from several `catalog/torch/` mirrors is always a vocabulary gap, +never a mirror to add. + +## The accuracy anchor + +The plan names 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 later and costs days. + +If this checkpoint has no entry in the accuracy references, add one **before +assembly**, from a written source you cite: read TensorRT-LLM's own accuracy +suite for the score, the task and the knobs the matching test overrides, and +`tensorrt_llm/evaluate/` for how that task is actually scored and filtered. +Take the entry matching this checkpoint's quantization. **If no source yields +a score, stop and say so. An anchor is never invented.** + +## Your two products + +**`plan.md`** — the onboarding plan, for a reader who will execute it: + +- identity: the (checkpoint, arch, parallel) triple, and the routing criteria + that will distinguish this target (config-shape fingerprint, SM version, + parallel topology — never a tuning knob like `max_num_tokens`, and never + anything that varies per step) +- the checkpoint's derived facts: geometry, quantization, tokenizer and chat + template details, `layer_types` if hybrid +- assembly units, in order: `core` first, then one per requested feature +- TensorRT-LLM's current support and the **gaps** you found +- the accuracy anchor, its protocol, and the source you took it from +- `## Open risks` — anything you are unsure of, so the assembler inherits the + uncertainty instead of rediscovering it + +**`vocabulary.yaml`** — the typed mapping the orchestrator walks. The schema +is fixed; the orchestrator builds its dispatch queue from it and will reject +a malformed one: + +```yaml +verdict: assemblable | assemblable_after_additions | ceiling +steps: + - step: + values: {{: , ...}} + state: OK | CERT | GAP + entry: catalog//.py # OK and CERT rows + contract_line: "" + candidate: torch.ops.trtllm. # GAP rows +queue: + - {{n: 1, mode: onboard, target: , depends_on: []}} + - {{n: 2, mode: verify, target: , depends_on: [], + surface: ""}} + - {{n: 3, mode: onboard, target: , depends_on: [], + unit: }} # only a later unit needs it +ceiling: null # or {{step: ..., searched: [...], why: ...}} +``` + +Rules the schema enforces, so get them right the first time: + +- an `OK`/`CERT` row quotes the `contract_line` that decides it — that line is + the evidence, and without it nobody can audit the judgement without redoing + the reading +- a `GAP` row names a `candidate`, or the verdict is `ceiling` +- a `verify` dispatch names the `surface` — the value and the contract line + that excludes it are what smith is dispatched with, so vagueness here costs + a whole run +- **the two chapters do not overlap.** An entry smith onboards is tested at + this checkpoint's values by construction, so it never also needs verifying +- `depends_on` is for real ordering only (a verify row whose test state must + be built by an onboard row first). Everything without a dependency runs + **concurrently**, so a spurious dependency directly costs wall-clock +- `unit` names the assembly unit that **first** needs the entry: `core` (the + default when absent) or one of the spec's `features`, and nothing else. Tag + every row only a feature increment uses — a drafter's MoE, a vision tower's + attention, a host-offload gather. Every row is still dispatched up front, + alongside core's; the tag changes only what a failure blocks. An untagged + row that fails stops the run before core is assembled, while a tagged one is + deferred and retried once before its unit's first assembler turn. Leave the + tag off when core uses the entry at all or you are unsure: untagged is the + blocking direction, so it is the safe one. A row a core row `depends_on` is + core's whatever its tag says + +{MODELING_V2_SWITCH} +{SOURCE_BOUNDARY} +{IN_TREE_LAYOUT} +{EVIDENCE_DISCIPLINE} +{CITATION_DISCIPLINE} + +## Ending your turn + +Call `append_plan_drafter_progress` exactly once, with a `summary` that names +what you wrote in **both** files and the counts by state, and a `decision`: + +- `DRAFT_READY` — both files are written and coherent. The PlanReviewer reads + them next. +- `CEILING` — a computation step has no server anywhere. The run stops here; + put the step, what you searched, and why nothing upstream provides it into + `vocabulary.yaml`'s `ceiling` block and into your summary. + +## Red flags + +| Impulse | Reality | +|---|---| +| "This row is unclear, dispatch smith and let them figure it out" | Read the contract again. Judging coverage is reading a contract, and that is your job, not smith's. | +| "No op in the inventory, so it's a ceiling" | The inventory cannot see pybind bindings. Check the catalog too. A wrong ceiling costs a feature; a wrong verify costs one dispatch. | +| "I can compose this op from a few torch mirrors" | Math assembled from mirrors is always a vocabulary gap. The mirror surface is glue: layout, movement, allocation, lookup. | +| "The anchor is probably around X" | Anchors come from a written source or from the user. Never invented. | +| "The HF reference does it op by op; the plan should mirror that" | HF defines the math, the catalog defines the vocabulary. The target is the fused form of HF's semantics. | +| "I'll mark everything depends_on the previous item to be safe" | Every spurious dependency serializes a batch that would have run concurrently. Order only what genuinely must be ordered. | +| "The row's `note` says it is dspark-only" | Nothing reads notes. Put it in `unit:`, or one failed dispatch the speculator alone needed stops the run before core has started. Measured: six such rows in a 33-dispatch queue, each able to close the run. | +| "The reviewer will catch it, ship the draft" | The reviewer gets one round. What it does not catch, the gates catch — expensively. | +| "The parallel segment doesn't change the math, so the mapping is the same" | Collectives appear nowhere in the HF reference. Walk the segment separately or the assembler finds the gap mid-run. | +""" +) + +__all__ = ["SYSTEM_PROMPT"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_reviewer.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_reviewer.py new file mode 100644 index 000000000000..227544b97a4a --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_reviewer.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""System prompt for the PlanReviewer role.""" + +from __future__ import annotations + +from ._common import CITATION_DISCIPLINE, EVIDENCE_DISCIPLINE, IN_TREE_LAYOUT, tidy + +SYSTEM_PROMPT = tidy( + f"""\ +You are the **PlanReviewer** for a staircase onboarding run. The PlanDrafter +has written `plan.md` and `vocabulary.yaml`; you decide whether the build +starts on them. + +**This is a paper review.** Do not build, do not run, do not test. The gates +do that, they do it later, and they do it better than you can here. + +You are stateful across rounds: you remember what you asked for, so a +re-review checks whether those items were addressed rather than starting over. + +## You get one round, and the build is not blocked on you + +The plan phase is capped. If you REJECT and the cap is reached, your +unresolved items are written into `plan.md` as open risks and the build starts +anyway. That changes what a useful review is: + +- **REJECT only for things that would waste an assembly turn or produce a + wrong gate.** A missing contract line, a GAP row with no candidate, an + invented anchor, a queue that serializes work that could run concurrently, + a routing criterion that varies per step. +- **Do not REJECT for style, completeness, or things the gates will catch in + minutes.** "The plan could say more about X" is a note for your summary, not + a rejection. A rejection that buys nothing costs a whole round. +- Put everything else in the summary. It reaches the assembler either way. + +## What to check + +**`vocabulary.yaml` — the part with teeth.** + +- every `OK`/`CERT` row quotes a `contract_line`. Without it the coverage + claim cannot be audited without redoing the reading — check the quoted line + actually appears in that entry's contract and actually decides the question. +- every `GAP` row names a candidate, or the verdict is `ceiling`. +- a `ceiling` verdict carries the step, what was searched, and why. **A + ceiling is the one verdict that ends the run**, so it earns more scrutiny + than any other row: was the catalog checked as well as the op inventory? + Was a pybind-only op (`thop_attention` is the standing example) mistaken + for absence? Was an existing entry's *uncertified* surface considered — a + `CERT` row is not a gap? +- every `verify` dispatch names the `surface` it is dispatched for. +- the two chapters do not overlap. +- `depends_on` orders only what genuinely must be ordered. **Every spurious + dependency serializes a batch that would have run concurrently**, so this is + a real cost, not a nicety. +- the parallel segment was walked separately. Collectives appear nowhere in + the HF reference, so a mapping that is all-green for a `tepN`/`depN` target + and never mentions a reduce or a gather was not derived, it was assumed. + +**`plan.md`.** + +- the accuracy anchor names a benchmark **and a protocol**, and cites the + written source it came from. An invented anchor is an automatic REJECT; + so is a protocol mismatch (a reasoning model scored as 5-shot completion + lands near the floor whatever its ability). +- the routing criteria are things known when the class is resolved and + constant for the engine's life, and they change the *structure* of the + forward. Config shape, SM version and parallel topology qualify. + `max_num_tokens` and `cuda_graph_config.max_batch_size` do not — they are + tuning knobs. Batch composition and `num_contexts` do not — they move every + step, and a target chosen from them would be chosen once and then be wrong. +- assembly units are ordered `core` first. Core must be green before the + first feature: the accuracy tolerance is wide enough that a core defect and + a feature defect are indistinguishable when they land together. +- the plan does **not** carry the built-in model zoo's forward structure. The + drafter may read the zoo freely and should record what TensorRT-LLM supports + and where the gaps are — but a plan that reads as a transcription of + `modeling_.py`, layer by layer, is a plan for a forward that cannot pass + the closed-vocabulary audit. REJECT that. + +{IN_TREE_LAYOUT} +{EVIDENCE_DISCIPLINE} +{CITATION_DISCIPLINE} + +## Ending your turn + +Call `append_plan_reviewer_progress` exactly once with a `summary` and a +`decision` of exactly `APPROVE` or `REJECT`. On REJECT, list specific, +actionable items, naming the file and the row for each one. On APPROVE, still +list anything you noticed — it reaches the assembler as context. + +## Red flags + +| Impulse | Reality | +|---|---| +| "I'll run the tests to be sure" | Paper review. The gates run, and they run after assembly where the evidence is real. | +| "The plan is thin on X — REJECT" | Would that waste an assembly turn or produce a wrong gate? If not, it is a summary note. You have one round. | +| "The ceiling verdict looks reasonable" | It is the one verdict that ends the run. Check the catalog *and* the inventory *and* the uncertified surfaces before accepting it. | +| "Dependencies are harmless, they just order things" | Each one serializes a batch. Spurious ordering is measured in wall-clock. | +| "The anchor number looks plausible" | Plausible is not sourced. Ask which file it came from. | +| "APPROVE, and I'll note my concerns nowhere" | The summary reaches the assembler. Silence throws away the only cheap part of this review. | +""" +) + +__all__ = ["SYSTEM_PROMPT"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/reviewer.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/reviewer.py new file mode 100644 index 000000000000..733e26446a46 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/reviewer.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""System prompt for the per-unit Reviewer role.""" + +from __future__ import annotations + +from ._common import CITATION_DISCIPLINE, EVIDENCE_DISCIPLINE, IN_TREE_LAYOUT, tidy + +SYSTEM_PROMPT = tidy( + f"""\ +You are the **Reviewer** for one assembled unit — a target's core assembly, or +one feature increment on it. + +**The expensive evidence already exists.** Before you were invoked, the +orchestrator ran the boot gate and (at a unit boundary) the accuracy gate, and +their measured output is in your prompt. Your job is not to reproduce that. It +is to read the diff against the evidence and say whether this unit is sound. + +Keep the loop tight. A review that takes longer than the assembly it reviews +has inverted the economics of the whole run. + +## What to check, in order of what actually catches defects + +1. **The vocabulary audit.** Every call in `forward` and the private methods + it reaches that creates or transforms a tensor must be a catalog entry. + Scope is the whole trick: the rope table's `arange`/`cos`/`sin` and + load-time `.t()`/`.to()` run at init and are outside the rule, so a + whole-file scan false-positives on them. Check the assembler scoped it + correctly, and spot-check two or three calls against `catalog/index.yaml` + yourself rather than trusting the list. +2. **`REQUIRED_TRTLLM_OPS` equals the forward's call set** — not a superset, + not a subset. A stale name is as wrong as a missing one. +3. **The runtime-value dump against the contracts.** Any value the engine + varied at runtime that sits outside its entry's certified column is a + blocker, not a note. This is the check that has been missed before: a + target merged with 34 of the 35 batch sizes its engine captured sitting + outside the contract, because every *constant* it passed was inside. +4. **Graph safety and sync-freedom in the forward.** No `.item()`, no + `.cpu()`, no data-dependent host branching. These are writing rules — a + forward that violates them is wrong the first time it is captured, and the + boot gate may well not show it. +5. **The manifest's bidirectional coverage**, and that every key in the + known-not-loaded set is named with a reason. +6. **The topology.** Every gate is topology-blind: a `tep4` target served at + tp2 boots and scores and looks self-consistent while the directory name + lies. Read the assembly's parallel handling against the segment the target + path declares. +7. **For a feature increment: the base path is unchanged.** With the feature + off, the forward must be what the recorded gate measured, bit for bit. + +## What not to do + +- **Do not re-run the gates.** They ran. Their numbers are in your prompt, and + re-running them buys nothing but wall-clock. +- **Do not run a benchmark or a sweep.** Performance is measured, never gated, + and it is not part of onboarding at all. +- **Do not review formatting.** The run lints once, at delivery, on purpose. +- **Do not ask for refactors.** Self-containment is the architecture; + duplication between targets is not a smell, and "this could be shared" is + the opposite of the design. + +## Rejecting + +REJECT for a defect the gates cannot see and a reader can: an unmapped call, a +stale op list, an uncovered runtime value, a sync in the forward, a topology +that does not match the segment, a feature that changed the base path. + +APPROVE when the diff is consistent with the gate evidence and the invariants +hold. A gate that passed plus a clean audit is sufficient — you are not +looking for reasons to keep the unit in review. + +If the gate evidence itself looks wrong — a score that does not match the +anchor's protocol, a boot that passed without generating — say so plainly. +That is a finding about the *run*, and it matters more than anything in the +diff. + +{IN_TREE_LAYOUT} +{EVIDENCE_DISCIPLINE} +{CITATION_DISCIPLINE} + +## Ending your turn + +Call `append_reviewer_progress` exactly once, with a `summary` and a +`decision` of exactly `APPROVE` or `REJECT`. On REJECT, cite the file and the +line for every item, and say which invariant each one violates. + +## Red flags + +| Impulse | Reality | +|---|---| +| "Let me re-run the accuracy gate to be sure" | It ran. The number is in your prompt. Re-running it costs an hour and produces the same number plus 0.1-0.3 points of session noise. | +| "The forward duplicates what another target does" | Self-containment is the architecture. Duplication between targets is the design, not a smell. | +| "The whole file should be scanned for non-catalog calls" | That false-positives on init-time code and buries the real finding. Scope it to `forward` plus what it reaches. | +| "The constants are all certified, so we're clean" | The values that escape are the ones the engine varies at runtime. Read the dump. | +| "Both gates passed, so the topology is what the path says" | Every gate is topology-blind. Only reading the assembly against the segment catches it. | +| "It would read better formatted" | One lint pass, at delivery. Not yours. | +| "REJECT — I'd have written this differently" | Name the invariant it violates, or APPROVE. | +""" +) + +__all__ = ["SYSTEM_PROMPT"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/runner.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/runner.py new file mode 100644 index 000000000000..f82450e19391 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/runner.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""System prompt for the Runner — the one role that decides nothing.""" + +from __future__ import annotations + +from ._common import tidy + +SYSTEM_PROMPT = tidy( + """\ +You are the **Runner**. You exist because this cluster's scheduler is +reachable only through a tool you hold and the orchestrator does not. + +Your entire job is to get one already-written command submitted. You are a +trigger, not a witness. + +## What that means, precisely + +The orchestrator has already resolved everything: the image, the partition, +the account, the mounts, the time limit, and the command itself — which is a +script on the shared filesystem, so there is nothing to quote and nothing to +assemble. You are handed a complete call. **Make exactly that call.** + +The job writes its own log and its own exit code to the shared filesystem, +and the orchestrator reads both itself and applies its own thresholds. + +So: **nothing reads your answer.** You cannot make a gate pass by reporting +that it passed, and you cannot make one fail by reporting that it failed. If +you were to end your turn with an invented result, the run would be entirely +unaffected, because the verdict comes from files the job wrote. + +This is deliberate. It is what lets gate *execution* move to an agent while +gate *adjudication* stays in Python. + +## Do + +- Make the call you were given, unchanged. +- Poll the job's status until it leaves `PENDING`/`RUNNING`, then stop. +- Say plainly if the **submission** failed. That is the one fact this turn + can contribute that nothing else can observe. + +## Do not + +- Do not use a blocking submit. The tool layer aborts a call after a few + minutes of silence and every job here is longer than that, so blocking + reliably reports a failure for a job that is running fine — and leaves it + running with its output unreachable. +- Do not change the image, partition, account, mounts, time limit or the + command. If the call does not fit what needs to run, that is a **symptom to + report**, not a problem to solve. +- Do not read the log and tell the orchestrator what it says. It is reading + the same file. +- Do not re-submit because a job looks slow. Two jobs racing for the same + GPUs is worse than one slow job. + +## Red flags + +| Impulse | Reality | +|---|---| +| "I'll check the log and report the score" | The orchestrator reads that file itself. Your summary of it is not evidence and is not used. | +| "The job failed, so I should retry it" | A failure is a result. Retrying hides it. | +| "This partition looks busy, I'll pick another" | The environment was resolved before you started. A different partition makes the evidence belong to a different machine. | +| "Blocking is simpler than polling" | Blocking times out on a healthy job and orphans it. | +""" +) + +__all__ = ["SYSTEM_PROMPT"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/smith.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/smith.py new file mode 100644 index 000000000000..51d646410a94 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/smith.py @@ -0,0 +1,161 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""System prompt for the smith role — the catalog's op steward.""" + +from __future__ import annotations + +from ._common import ( + CITATION_DISCIPLINE, + COMMIT_DISCIPLINE, + EVIDENCE_DISCIPLINE, + IN_TREE_LAYOUT, + SKILL_DISCIPLINE, + SOURCE_BOUNDARY, + tidy, +) + +SYSTEM_PROMPT = tidy( + f"""\ +You are **smith**, the op steward of the ModelingV2 kernel catalog. Entries +are trusted because their tests passed **on this machine** — never because of +what source code or documentation claims. + +**One run, one entry.** Your dispatch prompt names exactly one, and the skill +to load for it: + +| skill | this run | +|---|---| +| `staircase-kernel-onboard` | a callable with no entry yet — read the source, fix the single-call boundary, write the three pieces | +| `staircase-kernel-verify` | an entry that exists — does its certified column cover the values this checkpoint passes? | + +{SKILL_DISCIPLINE} + +## Your knowledge boundary + +smith is the role that **must** read engine source: fixing a single-call +boundary, or deciding that a callable is a macro, cannot be done from a +contract, because the contract does not exist yet. The other roles read source +to answer specific questions; you read it to build the thing they will later +read *instead of* the source. + +Also yours: the catalog's contracts and `index.yaml`, the generated op +inventory, and the values your dispatch prompt hands you. + +Anything broken outside your scope — an engine bug, an environment problem, a +missing category — is a **symptom**: describe it in your report, do not fix it. + +## What you may write + +- `tensorrt_llm/_torch/_experimental/modeling_v2/catalog//.md` + and `.py`, plus `catalog/index.yaml` +- that entry's GPU test at + `tests/unittest/_torch/modeling_v2//test_modeling_v2_.py` +- on a macro verdict only, the decomposition report your skill specifies + +Everything else is read-only: other entries, the targets, the accuracy +references, the build config. **No dependency installs.** + +`catalog/torch/` is reserved for hand-added mirrors of torch callables — +**never write there**, including when your verdict is that a torch callable is +the answer. Report that verdict and let the orchestrator decide: a computation +mirror widens what that surface is for, and the surface is glue by design +(layout, movement, allocation, lookup). + +## Receipts + +A receipt records **a run you watched**, on this architecture, and nothing +else. Three states, and the third is the one people get wrong: + +- `passed` — the test ran here and passed +- `failed` — the wrapper and test are right and the op demonstrably fails at + these values. **Land it anyway.** That is knowledge: a real constraint on + every target that would have passed that value, found here for the cost of + one dispatch instead of in an accuracy gate +- **absent — unknown.** Not "probably fine". What a receipt does not cover is + outside every target's capability envelope + +Receipts are keyed by **GPU architecture and nothing else** — there is no +version pin in tree to drift against. What a receipt is still fresh against is +its own entry's files: **a receipt is valid only if it post-dates the last +write to every file of its entry**, and an entry now spans two trees (the +package and the tests tree), so check both, mechanically, against mtimes +rather than by recollection. A receipt that predated its own test file has +shipped here before, and it read as green. + +## Tolerance + +Never inflate one to manufacture a pass. If a comparison needs a looser +tolerance than the dtype justifies, either the reference is wrong, the +operands are wrong, or the op is wrong — and all three are findings. A +loosened tolerance carries an inline justification naming which dtype and +which accumulation order it accounts for, or it does not get loosened. + +Build the test's reference **on the spot from native torch ops** — never from +the wrapper you are testing, and never from another catalog entry. Seed any +randomness. + +{SOURCE_BOUNDARY} +{IN_TREE_LAYOUT} +{COMMIT_DISCIPLINE} +{EVIDENCE_DISCIPLINE} +{CITATION_DISCIPLINE} + +## Ending your turn + +Call `append_smith_progress` exactly once, as the last action of your turn, +with a `decision` of `LANDED`, `FAILED_RECEIPT`, `MACRO` or `BLOCKED`. The +orchestrator reads that field and nothing else to decide whether assembly can +start, so a turn that ends without it is a turn whose result cannot be acted +on. + +Your final message is the run's sole deliverable; the skill specifies its +contents. At minimum: the entry path, the wrapped op's qualified name, the +receipt exactly as recorded, what the test covered (shapes, dtypes), the files +you touched, and **what remains uncertified on this entry** — the next caller +reads that line. + +## If you return `MACRO`, the decomposition has a shape + +`MACRO` is the one outcome whose whole purpose is to say what the next work is, +and the orchestrator turns it into queue items **mechanically**. It reads a +`constituents` list and nothing else; a row without `target` and `tier` is +dropped, and a decomposition that arrives only as prose in your summary ends +the run instead of extending it. + +```yaml +decision: MACRO +summary: | + why this step is more than one kernel-launching call, measured. +constituents: + - target: quantization/dequantize_e4m3_per_tensor # REQUIRED: catalog path + tier: 1 # REQUIRED: sourcing rung + note: registered as tensorrt_llm::dequantize_... # optional + surface: fp8 e4m3 -> bf16, per-tensor scale # optional + unit: core # optional, default core + depends_on: [] # optional +``` + +`target` is the **catalog path you propose**, in `category/name` form — not the +callable, not a description of the role. Omit a constituent that an existing +entry already serves; say so in the summary instead. If the natural home is a +directory you may not write to, still give the `target` you would propose and +say why in `note` — the placement is then a decision the orchestrator can make, +rather than one it cannot see. + +## Red flags + +| Impulse | Reality | +|---|---| +| "The engine call fails; I'll fall back to torch inside the wrapper" | The wrapper exists to certify the engine op. A torch fallback certifies nothing — record `failed` instead. | +| "It passes with rtol=0.5" | Tolerance inflation manufactures green. Justify the number or record `failed`. | +| "It obviously works on this arch too" | Receipts come from runs, not plausibility. Absent means unknown. | +| "A helper composing two ops would be convenient" | Python-level composition is the exact anti-pattern the catalog bans. One entry is one op invocation; fusion is legal only *inside* that call. | +| "The contract can say 'see source'" | A caller must be able to use the op from the contract alone. Sending them to the source makes the contract a pointer instead of an interface. | +| "I'll onboard the two ops this macro needs myself" | One run, one entry. Decompose, land the report, then dispatch one child per constituent. | +| "The test passed, so the receipt is good" | Only if it post-dates every file of the entry — in both trees. Lint and index edits happen after the test ran. | +| "I'll name the test after the op, like upstream does" | `test_modeling_v2_.py`, always. The name is what keeps it from colliding with upstream's test of the same op and what the CI selector matches. | +""" +) + +__all__ = ["SYSTEM_PROMPT"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/remediation.py b/agent-flow/agent_flow/workflows/staircase_onboard/remediation.py new file mode 100644 index 000000000000..4435bdafa036 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/remediation.py @@ -0,0 +1,687 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The bounded loop that closes an `UNCOVERED` audit. + +WHAT WAS THERE BEFORE + +An auditor that found runtime values outside their entries' certified +columns returned `UNCOVERED`, and the run **ended**. The fix then happened +outside the workflow, serially, by hand: + + auditor finds a gap + -> run closes + -> a human reads the finding, certifies the entry, re-runs + -> a fresh, stateless auditor re-scans EVERYTHING + -> finds the next gap + +That is the mechanism behind an onboarding run spending audit passes 53, +54, 55, 56 and 57 rediscovering an overlapping set of gaps. Each pass paid +a full re-scan to learn one increment. + +WHY THE OBVIOUS FIX IS WRONG + + while auditor == UNCOVERED: + dispatch a smith per uncovered value + re-run the full audit + +This is worse than the manual loop, not better. It re-dispatches identical +work when the auditor repeats itself, it fans out one smith per *value* +when seven uncovered batch sizes on one entry are a single piece of work, +and it has no stopping condition that is not "the model changed its mind". + +WHAT THIS MODULE IS + +The same loop, made **bounded, deduplicated and monotone**. Four rules, +each of which exists because its absence is a specific failure: + + * **Group by entry surface, not by value.** Seven uncovered batch sizes + on one entry are one verify dispatch that extends one parameterized + GPU test and writes one receipt, not seven. The merge key is + ``(entry, arch, axis)``. + * **Give every item a stable identity.** Natural language cannot say + whether two rounds found the same gap. The identity is a content hash + over the canonicalized surface plus the commit and capture it was + observed under, so "did we already do this?" is answerable. + * **Never re-dispatch a certified identity.** If a row that was + certified comes back, something is wrong upstream -- stale receipt, + stale contract read, a wrong coverage evaluator, or a smith whose fix + did not actually cover the surface. Dispatching the same smith again + cannot fix any of those. It is reported as non-progress. + * **Require monotone convergence.** The unresolved set must shrink. If + round N's signature equals round N-1's, the loop stops with a named + verdict instead of spending another GPU hour proving it is stuck. + +Everything here is a pure function over plain data. No subprocess, no +agent, no filesystem -- the orchestrator owns those and hands in the +results. That is what makes the convergence rules testable without a GPU. +""" + +from __future__ import annotations + +import hashlib +import re +from dataclasses import dataclass, field, replace +from typing import Any, Iterable, Mapping, Sequence + +# --------------------------------------------------------------- queue states + +#: Queued, not yet dispatched. +QUEUE_PENDING = "pending" +#: Dispatched in this round; a crash here is distinguishable from a crash +#: before it, so a resume knows the work may have been lost. +QUEUE_RUNNING = "running" +#: A smith landed the contract, the parameterized test and a fresh receipt. +QUEUE_CERTIFIED = "certified" +#: A smith ran the test and it did not pass. Terminal: re-dispatching the +#: same work cannot change the answer. +QUEUE_FAILED_RECEIPT = "failed_receipt" +#: A smith could not proceed at all. Terminal for the same reason. +QUEUE_BLOCKED = "blocked" +#: Replaced by a wider item covering the same surface. +QUEUE_SUPERSEDED = "superseded" + +#: States that still owe work. +OPEN_STATES = (QUEUE_PENDING, QUEUE_RUNNING) +#: States a round must not re-dispatch. +TERMINAL_STATES = (QUEUE_CERTIFIED, QUEUE_FAILED_RECEIPT, QUEUE_BLOCKED, QUEUE_SUPERSEDED) + +VALID_STATES = OPEN_STATES + TERMINAL_STATES + +# ------------------------------------------------------------------- verdicts + +#: The unresolved set stopped shrinking, or a certified row came back. Either +#: way the loop cannot converge and a human has to look. +STOP_NO_PROGRESS = "audit-no-progress" +#: A budget ran out with work still open. +STOP_BUDGET = "audit-remediation-budget" +#: A verify smith reported FAILED_RECEIPT or BLOCKED. +STOP_VERIFICATION_FAILED = "audit-verification-failed" + +# -------------------------------------------------------------------- budgets + + +@dataclass(frozen=True) +class RemediationBudget: + """Hard limits, so the loop is a bounded cycle rather than a ``while``. + + The defaults are deliberately small. Remediation is the *cheap* half of + onboarding -- it runs catalog GPU tests, not accuracy gates -- but it is + still GPU time, and a loop that needs a third round is usually a loop + whose first round was measuring the wrong thing. + """ + + max_rounds: int = 2 + #: One attempt per identity. A second attempt at byte-identical work + #: with unchanged evidence is the definition of no progress. + max_attempts_per_item: int = 1 + max_new_items: int = 16 + max_parallel_verifiers: int = 4 + + +# ----------------------------------------------------------------------- rows + + +def _normalize_values(values: Iterable[Any]) -> tuple[str, ...]: + """Canonicalize a surface's values so two spellings of it hash alike. + + Numeric-looking values sort numerically; everything else sorts as text. + Without this, ``[3, 5, 12]`` and ``[12, 3, 5]`` are different identities + and the same gap is dispatched twice. + """ + seen: list[str] = [] + for value in values: + text = str(value).strip() + if text and text not in seen: + seen.append(text) + numeric = all(re.fullmatch(r"-?\d+(\.\d+)?", v) for v in seen) if seen else False + return tuple(sorted(seen, key=float) if numeric else sorted(seen)) + + +@dataclass(frozen=True) +class UncoveredRow: + """One finding, as the auditor reported it. + + ``axis`` is the certification dimension the value varies along -- + ``batch_size``, ``num_pages``, ``local_layer_idx``. It is what makes two + findings on one entry the same piece of work or different ones. + """ + + entry: str + arch: str + axis: str + values: tuple[str, ...] = () + #: The contract sentence that excludes the value. Carried into the smith + #: brief verbatim: it is the difference between "certify batch size 3" + #: and a dispatch that has to re-derive what was wrong. + contract_line: str = "" + + @staticmethod + def parse(raw: Any) -> "UncoveredRow | None": + """Build a row from one decision-file mapping, or ``None`` if unusable. + + Permissive about shape and strict about content: a row missing the + entry, the arch or the axis cannot be grouped or dispatched, so it + is dropped rather than guessed at. The caller reports the drop -- + silently discarding an auditor finding is the one outcome worse than + having none. + """ + if not isinstance(raw, Mapping): + return None + entry = str(raw.get("entry", "")).strip() + arch = str(raw.get("arch", "")).strip() + axis = str(raw.get("axis", "")).strip() + if not entry or not arch or not axis: + return None + values = raw.get("values") + if values is None: + values = raw.get("value") + if isinstance(values, (str, int, float)): + values = [values] + if not isinstance(values, Sequence): + values = [] + return UncoveredRow( + entry=entry, + arch=arch, + axis=axis, + values=_normalize_values(values), + contract_line=str(raw.get("contract_line", "")).strip(), + ) + + +def parse_uncovered_rows(raw: Any) -> tuple[tuple[UncoveredRow, ...], int]: + """Parse the ``uncovered`` block of an auditor decision. + + Returns the usable rows and how many were dropped, so the caller can say + so rather than silently narrowing the audit. + """ + if not isinstance(raw, Sequence) or isinstance(raw, (str, bytes)): + return (), 0 + rows: list[UncoveredRow] = [] + dropped = 0 + for candidate in raw: + row = UncoveredRow.parse(candidate) + if row is None: + dropped += 1 + else: + rows.append(row) + return tuple(rows), dropped + + +# ---------------------------------------------------------------------- items + + +@dataclass(frozen=True) +class RemediationItem: + """One verify dispatch: one entry surface, on one architecture. + + The unit the queue is made of. Several :class:`UncoveredRow` on the same + ``(entry, arch, axis)`` collapse into one of these, because they are one + piece of work: a smith extends the parameterized GPU test over the whole + surface and writes one receipt. + """ + + entry: str + arch: str + axis: str + values: tuple[str, ...] = () + contract_line: str = "" + #: The target commit and runtime capture the surface was observed under. + #: Part of the identity: the same values under a NEW capture are + #: legitimately new work (a wrapper change can open a new runtime + #: branch), while the same values under the same capture are a repeat. + target_commit: str = "" + capture_id: str = "" + status: str = QUEUE_PENDING + attempts: int = 0 + + @property + def group_key(self) -> tuple[str, str, str]: + """What makes two findings the same piece of work.""" + return (self.entry, self.arch, self.axis) + + @property + def identity(self) -> str: + """Content hash over everything that makes this work distinct. + + Stable across processes and across a resume, which is the whole + point: it is how round two answers "have I already done this?" + """ + payload = "|".join( + [ + self.entry, + self.arch, + self.axis, + ",".join(self.values), + self.target_commit, + self.capture_id, + ] + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:12] + + def as_dict(self) -> dict[str, Any]: + """The checkpoint form. ``identity`` is derived, never stored.""" + return { + "entry": self.entry, + "arch": self.arch, + "axis": self.axis, + "values": list(self.values), + "contract_line": self.contract_line, + "target_commit": self.target_commit, + "capture_id": self.capture_id, + "status": self.status, + "attempts": self.attempts, + } + + @staticmethod + def from_dict(raw: Mapping[str, Any]) -> "RemediationItem": + status = str(raw.get("status", QUEUE_PENDING)) + return RemediationItem( + entry=str(raw.get("entry", "")), + arch=str(raw.get("arch", "")), + axis=str(raw.get("axis", "")), + values=_normalize_values(raw.get("values") or ()), + contract_line=str(raw.get("contract_line", "")), + target_commit=str(raw.get("target_commit", "")), + capture_id=str(raw.get("capture_id", "")), + status=status if status in VALID_STATES else QUEUE_PENDING, + attempts=int(raw.get("attempts", 0) or 0), + ) + + def brief(self) -> str: + """The dispatch brief a verify smith reads. + + Deliberately the whole surface rather than the uncovered subset: + a parameterized test written over ``[3, 5, 7]`` and a test written + over the entry's full observed range are different tests, and only + the second one stops the next round finding value 9. + """ + values = ", ".join(self.values) or "(none recorded)" + line = ( + f"\nThe contract line that excludes them: {self.contract_line}" + if self.contract_line + else "" + ) + return ( + f"axis: {self.axis}\n" + f"uncovered values on {self.arch}: {values}{line}\n" + f"Extend the parameterized GPU test over this surface, update the " + f"contract's certified column, and write a fresh {self.arch} receipt." + ) + + +def canonicalize( + rows: Iterable[UncoveredRow], + *, + target_commit: str = "", + capture_id: str = "", +) -> tuple[RemediationItem, ...]: + """Collapse findings into one item per ``(entry, arch, axis)``. + + This is the step that turns "seven uncovered batch sizes" into one + dispatch instead of seven, and it is why the queue is a queue of + *surfaces* rather than of values. + """ + grouped: dict[tuple[str, str, str], list[UncoveredRow]] = {} + for row in rows: + grouped.setdefault((row.entry, row.arch, row.axis), []).append(row) + + items: list[RemediationItem] = [] + for (entry, arch, axis), members in grouped.items(): + values: list[str] = [] + for member in members: + values.extend(member.values) + # First non-empty contract line wins: they describe the same column. + line = next((m.contract_line for m in members if m.contract_line), "") + items.append( + RemediationItem( + entry=entry, + arch=arch, + axis=axis, + values=_normalize_values(values), + contract_line=line, + target_commit=target_commit, + capture_id=capture_id, + ) + ) + return tuple(sorted(items, key=lambda i: i.group_key)) + + +# ---------------------------------------------------------------------- merge + + +@dataclass(frozen=True) +class MergeOutcome: + """What appending a round's findings to the queue actually did.""" + + queue: tuple[RemediationItem, ...] = () + #: Genuinely new work. + appended: tuple[RemediationItem, ...] = () + #: Already queued or already running; not dispatched again. + duplicates: tuple[RemediationItem, ...] = () + #: An open item grew to cover more of the same surface. + widened: tuple[RemediationItem, ...] = () + #: A surface that was certified came back. Not re-dispatchable -- see + #: the module docstring. + regressions: tuple[RemediationItem, ...] = () + #: A surface whose previous attempt failed or was blocked. + exhausted: tuple[RemediationItem, ...] = () + + @property + def has_new_work(self) -> bool: + return bool(self.appended or self.widened) + + +def merge_into_queue( + queue: Sequence[RemediationItem], + incoming: Sequence[RemediationItem], + *, + budget: RemediationBudget = RemediationBudget(), +) -> MergeOutcome: + """Append only genuinely new work, and classify everything else. + + The classification is the useful part. "Nothing was appended" is not one + condition but four, and they need different verdicts: an auditor + repeating itself, an auditor widening a surface, a certified row coming + back, and a row whose attempt already failed. + """ + by_group: dict[tuple[str, str, str], int] = {} + working = list(queue) + for index, item in enumerate(working): + # Later entries win, so a superseded row never shadows its successor. + by_group[item.group_key] = index + + appended: list[RemediationItem] = [] + duplicates: list[RemediationItem] = [] + widened: list[RemediationItem] = [] + regressions: list[RemediationItem] = [] + exhausted: list[RemediationItem] = [] + + for item in incoming: + index = by_group.get(item.group_key) + if index is None: + if len(appended) >= budget.max_new_items: + # Over the cap: not silently dropped, reported as exhausted + # so the caller stops rather than proceeding on a partial set. + exhausted.append(item) + continue + working.append(item) + by_group[item.group_key] = len(working) - 1 + appended.append(item) + continue + + existing = working[index] + covered = set(item.values) <= set(existing.values) + + if existing.status in (QUEUE_FAILED_RECEIPT, QUEUE_BLOCKED): + exhausted.append(item) + continue + if existing.status == QUEUE_CERTIFIED: + # Same surface, already certified -> upstream is wrong somewhere. + # A wider surface under the SAME capture is the same signal: the + # first audit was incomplete. + regressions.append(item) + continue + if covered: + duplicates.append(item) + continue + if existing.attempts >= budget.max_attempts_per_item: + exhausted.append(item) + continue + + merged = replace( + existing, + values=_normalize_values(list(existing.values) + list(item.values)), + contract_line=existing.contract_line or item.contract_line, + status=QUEUE_PENDING, + ) + working[index] = merged + widened.append(merged) + + return MergeOutcome( + queue=tuple(working), + appended=tuple(appended), + duplicates=tuple(duplicates), + widened=tuple(widened), + regressions=tuple(regressions), + exhausted=tuple(exhausted), + ) + + +# ------------------------------------------------------------------ signature + + +def unresolved_signature(queue: Sequence[RemediationItem]) -> str: + """A hash of the still-open identity set. + + Round N's signature equal to round N-1's means the loop learned nothing + and would learn nothing from another round. Comparing signatures is + cheaper and stricter than comparing prose summaries, which was the only + thing available before. + """ + identities = sorted(i.identity for i in queue if i.status in OPEN_STATES) + if not identities: + return "" + return hashlib.sha256("|".join(identities).encode("utf-8")).hexdigest()[:12] + + +# ----------------------------------------------------------------- round plan + + +@dataclass(frozen=True) +class RoundDecision: + """What the orchestrator should do next.""" + + #: ``dispatch`` | ``stop`` | ``clean`` + action: str + #: The closing verdict, when ``action`` is ``stop``. + verdict: str = "" + dispatch: tuple[RemediationItem, ...] = () + #: Human-readable reason, for the journal and the console. + detail: str = "" + workers: int = 1 + + @property + def stopping(self) -> bool: + return self.action == "stop" + + +def decide_round( + merge: MergeOutcome, + *, + budget: RemediationBudget, + round_no: int, + previous_signature: str = "", +) -> RoundDecision: + """Decide whether to dispatch another verification round, or stop. + + The order of these checks is the policy, and it is deliberate: + + 1. nothing open at all -> the audit is closed; + 2. a certified surface came back -> stop, whatever else is true. It + means the evidence is untrustworthy, and more rounds cannot repair + evidence; + 3. the round budget is spent -> stop; + 4. no new work -> stop. The auditor repeated itself; + 5. the unresolved set did not change -> stop. Monotone convergence is + the contract; + 6. otherwise dispatch, capped at ``max_parallel_verifiers``. + """ + open_items = tuple(i for i in merge.queue if i.status in OPEN_STATES) + + if not open_items and not merge.regressions: + return RoundDecision(action="clean", detail="no uncovered surfaces remain") + + if merge.regressions: + names = ", ".join(f"{i.entry}[{i.axis}]" for i in merge.regressions[:4]) + return RoundDecision( + action="stop", + verdict=STOP_NO_PROGRESS, + detail=( + f"{len(merge.regressions)} surface(s) reported uncovered after being " + f"certified in an earlier round: {names}. Re-dispatching the same " + f"verification cannot fix a stale receipt, a stale contract read or a " + f"coverage evaluator that is wrong -- and one of those is what this is." + ), + ) + + if round_no > budget.max_rounds: + return RoundDecision( + action="stop", + verdict=STOP_BUDGET, + detail=( + f"remediation budget spent ({budget.max_rounds} round(s)) with " + f"{len(open_items)} surface(s) still uncovered" + ), + ) + + if merge.exhausted: + names = ", ".join(f"{i.entry}[{i.axis}]" for i in merge.exhausted[:4]) + return RoundDecision( + action="stop", + verdict=STOP_BUDGET, + detail=( + f"{len(merge.exhausted)} surface(s) are past their attempt or item budget: {names}" + ), + ) + + if not merge.has_new_work: + return RoundDecision( + action="stop", + verdict=STOP_NO_PROGRESS, + detail=( + f"the audit reported {len(merge.duplicates)} surface(s), all already " + f"queued or running. Nothing new was learned this round." + ), + ) + + signature = unresolved_signature(merge.queue) + if previous_signature and signature == previous_signature: + return RoundDecision( + action="stop", + verdict=STOP_NO_PROGRESS, + detail=( + "the unresolved set is byte-identical to the previous round's. " + "Remediation has to shrink it to continue." + ), + ) + + dispatch = tuple(i for i in merge.queue if i.status == QUEUE_PENDING) + return RoundDecision( + action="dispatch", + dispatch=dispatch, + workers=max(1, min(budget.max_parallel_verifiers, len(dispatch))), + detail=( + f"{len(merge.appended)} new and {len(merge.widened)} widened surface(s); " + f"{len(dispatch)} verify dispatch(es)" + ), + ) + + +# ------------------------------------------------------------------- outcomes + +#: Smith decisions that certify a surface, and the queue state each maps to. +_DECISION_STATES = { + "LANDED": QUEUE_CERTIFIED, + "FAILED_RECEIPT": QUEUE_FAILED_RECEIPT, + "BLOCKED": QUEUE_BLOCKED, +} + + +def apply_outcomes( + queue: Sequence[RemediationItem], + outcomes: Mapping[str, str | None], +) -> tuple[tuple[RemediationItem, ...], tuple[RemediationItem, ...]]: + """Fold a round's smith decisions back into the queue. + + ``outcomes`` maps identity to the smith's decision. Returns the new + queue and the items that failed, which the caller turns into + ``audit-verification-failed``. + + A missing or unrecognized decision counts as blocked rather than as + silence: a verify dispatch that ended without saying anything has not + certified anything, and treating it as success is how an uncertified + surface reaches a merge. + """ + updated: list[RemediationItem] = [] + failed: list[RemediationItem] = [] + for item in queue: + if item.status not in OPEN_STATES: + updated.append(item) + continue + decision = (outcomes.get(item.identity) or "").strip().upper() + state = _DECISION_STATES.get(decision, QUEUE_BLOCKED) + landed = replace(item, status=state, attempts=item.attempts + 1) + updated.append(landed) + if state in (QUEUE_FAILED_RECEIPT, QUEUE_BLOCKED): + failed.append(landed) + return tuple(updated), tuple(failed) + + +@dataclass +class RemediationState: + """The queue plus the bookkeeping the convergence rules need. + + Persisted in the run's checkpoint so a resume continues the cycle rather + than restarting it -- which is what made the manual loop re-scan + everything each time. + """ + + queue: list[RemediationItem] = field(default_factory=list) + round_no: int = 0 + signature: str = "" + + def as_dict(self) -> dict[str, Any]: + return { + "queue": [item.as_dict() for item in self.queue], + "round_no": self.round_no, + "signature": self.signature, + } + + @staticmethod + def from_dict(raw: Mapping[str, Any] | None) -> "RemediationState": + raw = raw or {} + rows = raw.get("queue") or [] + return RemediationState( + queue=[RemediationItem.from_dict(r) for r in rows if isinstance(r, Mapping)], + round_no=int(raw.get("round_no", 0) or 0), + signature=str(raw.get("signature", "")), + ) + + def render(self) -> str: + """One line per item, for the journal.""" + if not self.queue: + return "(empty)" + return "\n".join( + f" {item.identity} {item.status:<14} {item.entry}[{item.axis}] " + f"-> {', '.join(item.values) or '(none)'}" + for item in self.queue + ) + + +__all__ = [ + "OPEN_STATES", + "QUEUE_BLOCKED", + "QUEUE_CERTIFIED", + "QUEUE_FAILED_RECEIPT", + "QUEUE_PENDING", + "QUEUE_RUNNING", + "QUEUE_SUPERSEDED", + "STOP_BUDGET", + "STOP_NO_PROGRESS", + "STOP_VERIFICATION_FAILED", + "TERMINAL_STATES", + "VALID_STATES", + "MergeOutcome", + "RemediationBudget", + "RemediationItem", + "RemediationState", + "RoundDecision", + "UncoveredRow", + "apply_outcomes", + "canonicalize", + "decide_round", + "merge_into_queue", + "parse_uncovered_rows", + "unresolved_signature", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/runtime.py b/agent-flow/agent_flow/workflows/staircase_onboard/runtime.py new file mode 100644 index 000000000000..a18fb50e718b --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/runtime.py @@ -0,0 +1,496 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The runtime map: which engine area a checkpoint needs, and where it lives. + +This replaces a thing that does not work: engine knowledge written into skill +prose. The failure is not that the prose is bad, it is that prose about a +moving tree is wrong shortly after it is written and is *believed anyway*, +because it reads like knowledge. This package measured that directly — seven +line citations, all seven pointing at the wrong line within days, while every +path and symbol beside them still resolved. + +So the map carries only the two things that survive a refactor: + +- **where the code is** — directories by preference, since a directory + survives a file split; +- **what to grep for** — a class, a function, a config field. + +and one thing prose is actually good at: a sentence or two saying what the +area *is*, at the level a plan needs. No signatures, no argument lists, no +line numbers, no copied code. An agent does not learn the engine from here. +It learns where the engine is, and reads the current source. + +Three properties this has that a skill does not: + +1. **It is checked.** Every path and every doc is resolved against the + checkout by a test, so a pointer that moved is a failure rather than a + stale sentence nobody notices. +2. **It is queried, not loaded.** A dense Llama needs attention, MLP, norm + and quantization. Loading mamba, KDA, mHC and disaggregation alongside + them costs context and — worse — offers an agent detail it has no business + using. +3. **It joins to what was measured.** :func:`areas_for_demands` takes the + structural demands ``survey.py`` read off ``config.json`` and returns the + areas they imply, so the plan phase sees the areas *this* checkpoint + needs rather than a catalogue of everything. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Iterable, Mapping, Sequence + +import yaml + +#: Ships with the package; see ``pyproject.toml``'s package-data. +INDEX_PATH = Path(__file__).parent / "runtime_index.yaml" + +#: The kernel-source registry. Separate from the runtime map because it +#: answers a different question: not "where does this area live" but "where +#: did this kernel come from, and what does certifying it require". +KERNEL_SOURCES_PATH = Path(__file__).parent / "kernel_sources.yaml" + +KERNEL_SOURCE_REQUIRED: tuple[str, ...] = ("summary", "tier", "availability", "certifying") +KERNEL_SOURCE_OPTIONAL: tuple[str, ...] = ("registers", "implementation") + +#: Areas every run needs regardless of what the checkpoint declares. A target +#: always sits somewhere in the parallel topology and is always captured or +#: deliberately not captured, so these are never "not relevant" — they are +#: just not *discoverable* from a config file. +ALWAYS_RELEVANT: tuple[str, ...] = ( + "parallelism", + "cuda_graph", + "scheduling", + "kv_cache", + # Every run has an autotuner, and it is on by default. It belongs here + # rather than behind a demand because nothing in a `config.json` says + # "this model is affected by tactic selection" --- every model is. + "jit_autotune", +) + +REQUIRED_FIELDS: tuple[str, ...] = ("summary", "code", "gating") +OPTIONAL_FIELDS: tuple[str, ...] = ("symbols", "tests", "doc", "demands") + +#: ``gating`` is required, and its own fields are too. The reason is a +#: measured one: this workflow once delivered a target with three of seven +#: features "configured" and none of them "verified", because the gates that +#: ran were blind to exactly those three. Making an area declare how it is +#: proven -- and, more importantly, what stays green while it is broken -- +#: is what turns "we enabled it" into something a reviewer can check. +GATING_FIELDS: tuple[str, ...] = ("gate", "signal", "blind") + +#: Which existing gate covers an area. ``none`` is a legitimate answer and is +#: the honest one for anything the standard gates cannot see; it is also the +#: flag that says a witness has to be built before the feature can be claimed. +GATES: tuple[str, ...] = ("boot", "accuracy", "acceptance", "none") + +#: The one non-area key. It holds the shape of the engine rather than one +#: part of it, so it has its own schema: a summary and an ordered `flow` of +#: stages, each naming the areas that live there. +# +# It exists because the most common planning error is not getting an area +# wrong, it is not knowing which areas are somebody else's job. A target is +# one function in the middle of a long pipeline; most of the map is context +# for reading a log, not work the target does. +OVERVIEW_KEY = "_overview" +FLOW_FIELDS: tuple[str, ...] = ("stage", "does", "areas") + + +class RuntimeIndexError(ValueError): + """Raised when the index is malformed or points somewhere that is not there.""" + + +@dataclass(frozen=True) +class Gating: + """How a run proves this area works, and what it cannot see.""" + + gate: str + signal: str + blind: str + + @property + def needs_a_witness(self) -> bool: + """No standard gate covers it, so a claim needs something purpose-built.""" + return self.gate == "none" + + +@dataclass(frozen=True) +class RuntimeArea: + """One engine area, and where to read it.""" + + name: str + summary: str + gating: "Gating | None" = None + code: tuple[str, ...] = () + symbols: tuple[str, ...] = () + tests: tuple[str, ...] = () + doc: str | None = None + demands: tuple[str, ...] = () + + @property + def pointers(self) -> tuple[str, ...]: + """Every repo path this area names, for resolution against a checkout.""" + return (*self.code, *self.tests, *((self.doc,) if self.doc else ())) + + def render(self) -> str: + """The block an agent reads. Pointers, not content.""" + lines = [f"### {self.name}", "", " ".join(self.summary.split()), ""] + lines.append(f"- **code**: {', '.join(f'`{p}`' for p in self.code)}") + if self.symbols: + lines.append(f"- **grep for**: {', '.join(f'`{s}`' for s in self.symbols)}") + if self.tests: + lines.append(f"- **tests**: {', '.join(f'`{p}`' for p in self.tests)}") + if self.doc: + lines.append(f"- **guide**: `{self.doc}`") + if self.gating: + lines.append(f"- **proven by**: {self.gating.gate} gate — {self.gating.signal}") + lines.append(f"- **blind spot**: {self.gating.blind}") + return "\n".join(lines) + + +@dataclass(frozen=True) +class FlowStage: + """One stage of the request path, and which areas live there.""" + + stage: str + does: str + areas: tuple[str, ...] = () + + +@dataclass(frozen=True) +class Overview: + """The engine's shape: what a target is one link of.""" + + summary: str = "" + flow: tuple[FlowStage, ...] = () + + def render(self) -> str: + if not self.flow: + return "" + lines = [ + "## How the engine fits together", + "", + " ".join(self.summary.split()), + "", + "| stage | what it owns | areas |", + "|---|---|---|", + ] + for stage in self.flow: + areas = ", ".join(f"`{a}`" for a in stage.areas) or "—" + lines.append(f"| **{stage.stage}** | {' '.join(stage.does.split())} | {areas} |") + return "\n".join(lines) + "\n" + + +@dataclass(frozen=True) +class RuntimeIndex: + """The whole map, loaded once.""" + + areas: dict[str, RuntimeArea] = field(default_factory=dict) + overview: Overview = field(default_factory=Overview) + + def __iter__(self): + return iter(self.areas.values()) + + def __len__(self) -> int: + return len(self.areas) + + def get(self, name: str) -> RuntimeArea | None: + return self.areas.get(name) + + def by_demand(self, demand: str) -> tuple[RuntimeArea, ...]: + """Areas that declare ``demand`` as something that makes them relevant.""" + return tuple(a for a in self.areas.values() if demand in a.demands) + + +@dataclass(frozen=True) +class KernelSource: + """One place kernels come from, and what it costs to certify one.""" + + name: str + summary: str + tier: int + availability: str + certifying: str + registers: tuple[str, ...] = () + implementation: tuple[str, ...] = () + + @property + def pointers(self) -> tuple[str, ...]: + return (*self.registers, *self.implementation) + + def render(self) -> str: + lines = [f"### {self.name} (tier {self.tier})", "", " ".join(self.summary.split()), ""] + if self.registers: + lines.append(f"- **registers through**: {', '.join(f'`{p}`' for p in self.registers)}") + if self.implementation: + lines.append( + f"- **implemented in**: {', '.join(f'`{p}`' for p in self.implementation)}" + ) + lines.append(f"- **availability**: {' '.join(self.availability.split())}") + lines.append(f"- **certifying one**: {' '.join(self.certifying.split())}") + return "\n".join(lines) + + +def load_kernel_sources(path: str | Path = KERNEL_SOURCES_PATH) -> dict[str, KernelSource]: + """Parse and validate the kernel-source registry.""" + raw = yaml.safe_load(Path(path).read_text(encoding="utf-8")) + if not isinstance(raw, dict) or not raw: + raise RuntimeIndexError(f"{path} must be a non-empty mapping of source -> entry") + + sources: dict[str, KernelSource] = {} + errors: list[str] = [] + for name, entry in raw.items(): + if not isinstance(entry, dict): + errors.append(f"{name}: must be a mapping") + continue + unknown = sorted(set(entry) - set(KERNEL_SOURCE_REQUIRED) - set(KERNEL_SOURCE_OPTIONAL)) + if unknown: + errors.append(f"{name}: unknown field(s) {unknown}") + for required in KERNEL_SOURCE_REQUIRED: + if entry.get(required) in (None, "", []): + errors.append(f"{name}: missing required field '{required}'") + if not entry.get("registers") and not entry.get("implementation"): + errors.append(f"{name}: needs at least one of 'registers' or 'implementation'") + sources[str(name)] = KernelSource( + name=str(name), + summary=str(entry.get("summary", "")), + tier=int(entry.get("tier", 0) or 0), + availability=str(entry.get("availability", "")), + certifying=str(entry.get("certifying", "")), + registers=tuple(str(p) for p in entry.get("registers", ())), + implementation=tuple(str(p) for p in entry.get("implementation", ())), + ) + if errors: + raise RuntimeIndexError(f"{path} is malformed:\n - " + "\n - ".join(errors)) + return sources + + +def source_for_stem( + stem: str, sources: "Mapping[str, KernelSource] | None" = None +) -> KernelSource | None: + """Map a `custom_ops/.py` file back to the source it represents. + + The probe attributes an op to the *file* that registered it, because that + is what is mechanically derivable on the host. The registry is keyed by + what the source IS. This joins them through each entry's `registers` + paths, so the mapping has exactly one definition. + + ``cpp_thop`` is returned for the C++ plane, which the probe reports for + any op no Python file claims --- a conclusion, not a fallback. + """ + sources = load_kernel_sources() if sources is None else sources + for source in sources.values(): + for path in source.registers: + if Path(path).stem == stem or path.rstrip("/").endswith(stem): + return source + return sources.get(stem) + + +def format_kernel_sources(sources: "Mapping[str, KernelSource]") -> str: + """The block a smith reads before deciding where a computation comes from.""" + if not sources: + return "" + lines = [ + "## Where kernels come from", + "", + "An op's presence is measured on the host --- the surface survey above " + "enumerated what this build actually registers. This says what each " + "source IS and what certifying one of its kernels requires, because " + "that differs per source and decides the shape of your receipt.", + "", + "Both registration planes land in the same `torch.ops.trtllm` namespace, " + "so the namespace alone tells you nothing about provenance. An op found " + "in no `custom_ops/*.py` file is almost certainly registered from C++.", + "", + ] + lines.extend(source.render() + "\n" for source in sources.values()) + return "\n".join(lines).rstrip() + "\n" + + +def load_index(path: str | Path = INDEX_PATH) -> RuntimeIndex: + """Parse and validate the index. + + Validation is strict on purpose. An index is only worth consulting if a + wrong entry is a loud failure; a permissive loader turns a typo into an + area nobody can find and nobody notices is missing. + """ + raw = yaml.safe_load(Path(path).read_text(encoding="utf-8")) + if not isinstance(raw, dict) or not raw: + raise RuntimeIndexError(f"{path} must be a non-empty mapping of area -> entry") + + areas: dict[str, RuntimeArea] = {} + errors: list[str] = [] + + overview = Overview() + raw_overview = raw.pop(OVERVIEW_KEY, None) + if raw_overview is not None: + if not isinstance(raw_overview, dict): + errors.append(f"{OVERVIEW_KEY}: must be a mapping") + else: + stages = [] + for item in raw_overview.get("flow", ()): + if not isinstance(item, dict): + errors.append(f"{OVERVIEW_KEY}.flow: each stage must be a mapping") + continue + unknown_flow = sorted(set(item) - set(FLOW_FIELDS)) + if unknown_flow: + errors.append(f"{OVERVIEW_KEY}.flow: unknown field(s) {unknown_flow}") + stages.append( + FlowStage( + stage=str(item.get("stage", "")), + does=str(item.get("does", "")), + areas=tuple(str(a) for a in item.get("areas", ())), + ) + ) + overview = Overview(summary=str(raw_overview.get("summary", "")), flow=tuple(stages)) + + for name, entry in raw.items(): + if str(name).startswith("_"): + errors.append(f"{name}: unknown reserved key; only {OVERVIEW_KEY!r} is reserved") + continue + if not isinstance(entry, dict): + errors.append(f"{name}: must be a mapping, got {type(entry).__name__}") + continue + unknown = sorted(set(entry) - set(REQUIRED_FIELDS) - set(OPTIONAL_FIELDS)) + if unknown: + errors.append(f"{name}: unknown field(s) {unknown}") + for required in REQUIRED_FIELDS: + if not entry.get(required): + errors.append(f"{name}: missing required field '{required}'") + + gating = None + raw_gating = entry.get("gating") + if isinstance(raw_gating, dict): + unknown_gating = sorted(set(raw_gating) - set(GATING_FIELDS)) + if unknown_gating: + errors.append(f"{name}.gating: unknown field(s) {unknown_gating}") + for required in GATING_FIELDS: + if not raw_gating.get(required): + errors.append(f"{name}.gating: missing '{required}'") + gate = str(raw_gating.get("gate", "")) + if gate and gate not in GATES: + errors.append(f"{name}.gating.gate is {gate!r}; must be one of {list(GATES)}") + gating = Gating( + gate=gate, + signal=" ".join(str(raw_gating.get("signal", "")).split()), + blind=" ".join(str(raw_gating.get("blind", "")).split()), + ) + elif raw_gating is not None: + errors.append(f"{name}.gating: must be a mapping, got {type(raw_gating).__name__}") + if errors and name not in {e.split(":")[0] for e in errors}: + continue + areas[str(name)] = RuntimeArea( + name=str(name), + summary=str(entry.get("summary", "")), + code=tuple(str(p) for p in entry.get("code", ())), + symbols=tuple(str(s) for s in entry.get("symbols", ())), + tests=tuple(str(p) for p in entry.get("tests", ())), + doc=str(entry["doc"]) if entry.get("doc") else None, + demands=tuple(str(d) for d in entry.get("demands", ())), + gating=gating, + ) + if errors: + raise RuntimeIndexError(f"{path} is malformed:\n - " + "\n - ".join(errors)) + return RuntimeIndex(areas=areas, overview=overview) + + +def areas_for_demands( + demands: Iterable[str], + *, + index: RuntimeIndex | None = None, + always: Sequence[str] = ALWAYS_RELEVANT, +) -> tuple[RuntimeArea, ...]: + """Return the areas the checkpoint's measured demands imply. + + The join is the point. ``survey.py`` reads structural demands off + ``config.json`` — ``attention.mla``, ``moe.routing``, ``recurrent.state`` + — and those names are what each area declares itself relevant to. So a + dense Llama gets attention, MLP, norm and quantization; a hybrid Qwen + additionally gets mamba and gated_deltanet; and neither is handed the + other's areas. + + Returned in index order rather than match order, so two runs on the same + checkpoint read the same document. + """ + index = index or load_index() + wanted = set(always) + for demand in demands: + wanted.update(a.name for a in index.by_demand(demand)) + return tuple(a for a in index if a.name in wanted) + + +def format_areas( + areas: Sequence[RuntimeArea], + *, + unmatched: Sequence[str] = (), + overview: "Overview | None" = None, +) -> str: + """Render the map block for a prompt. + + The overview goes first when there is one: knowing which stage owns a + concern is what stops a target from trying to do the executor's job. + """ + if not areas: + return "" + lines = [] + if overview is not None and overview.flow: + lines += [overview.render(), ""] + lines += [ + "## Runtime map", + "", + "The engine areas this checkpoint touches, and where each one lives **in " + "this checkout right now**. Read the source; do not take a description of " + "it from anywhere, including from here. These entries say where to look " + "and what to grep for, and deliberately carry no signatures, no argument " + "lists and no line numbers — those go stale between one onboarding and " + "the next, and a stale one reads exactly like a true one.", + "", + "The `tests` path matters as much as the `code` path: a new entry's test " + "should look like the tests already sitting there.", + "", + ] + lines.extend(area.render() + "\n" for area in areas) + if unmatched: + lines.append( + "**Demands no area claims**: " + + ", ".join(f"`{d}`" for d in unmatched) + + ". That is the signal worth acting on — either the map has a gap, " + "or this checkpoint needs something the engine has no home for. Say " + "which in the plan." + ) + return "\n".join(lines).rstrip() + "\n" + + +def unmatched_demands( + demands: Iterable[str], *, index: RuntimeIndex | None = None +) -> tuple[str, ...]: + """Demands the map has no area for — a gap in the map or in the engine.""" + index = index or load_index() + return tuple(sorted({d for d in demands if not index.by_demand(d)})) + + +__all__ = [ + "ALWAYS_RELEVANT", + "FLOW_FIELDS", + "FlowStage", + "GATES", + "GATING_FIELDS", + "Gating", + "INDEX_PATH", + "KERNEL_SOURCES_PATH", + "KernelSource", + "RuntimeArea", + "RuntimeIndex", + "OVERVIEW_KEY", + "Overview", + "RuntimeIndexError", + "areas_for_demands", + "format_areas", + "format_kernel_sources", + "load_index", + "load_kernel_sources", + "source_for_stem", + "unmatched_demands", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/runtime_index.yaml b/agent-flow/agent_flow/workflows/staircase_onboard/runtime_index.yaml new file mode 100644 index 000000000000..cbdf2f1a1f22 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/runtime_index.yaml @@ -0,0 +1,677 @@ +# The runtime map: concept -> where the current code and tests live. +# +# WHAT THIS FILE IS, AND WHAT IT DELIBERATELY IS NOT +# +# It is a *map*, not a description. Each entry says what an area is at the +# level a plan needs, then names where to read the real thing. It carries no +# code, no signatures, no argument lists and no line numbers. +# +# That restraint is the whole design. Prose about engine internals rots: this +# tree already demonstrated 7 line citations that were all wrong within days +# of being written, while every path and symbol beside them still resolved. +# A map made of paths and symbols survives a refactor; a map made of copied +# detail is wrong shortly after it is written and is believed anyway, because +# it reads like knowledge. +# +# So an agent onboarding a model does not learn the engine from here. It +# learns *where the engine is* from here, and reads the current source. +# +# SCHEMA (enforced by tests/workflows/staircase_onboard/test_runtime_index.py) +# +# : +# summary: one or two sentences. What it is and why a target cares. +# code: repo-relative paths. Directories preferred over files --- a +# directory survives a file split, a file does not. +# symbols: names to grep for. A class, a function, a config field. +# tests: repo-relative paths to the tests that exercise it. This is +# half the value: a target's own test should look like the +# tests already there. +# doc: the developer guide or feature doc, when one exists. +# demands: which `survey.py` checkpoint demands make this area relevant. +# This is what lets the plan phase surface only the areas the +# checkpoint in hand actually needs. +# gating: how a run PROVES this area works. Required, because the +# alternative is what already happened once: three of seven +# features shipped "configured" and none of them "verified". +# gate -- which existing gate covers it, or `none` +# signal -- the observable that proves it actually RAN, +# named so a witness can be built against it +# blind -- what stays green while this is broken. The most +# valuable line in the entry: a gate that cannot +# see a defect is a gate that certifies it. +# +# Every path here is resolved against the checkout by a test. A pointer that +# has moved is a test failure, not a stale sentence nobody notices. + +# --------------------------------------------------------------------------- +# `_overview` is not an area. It is the shape of the engine a target sits +# inside, and it exists because the most common planning error is not getting +# an area wrong --- it is not knowing which of them are somebody else's job. +# +# A modeling_v2 target is ONE function in the middle of a long pipeline. It +# does not schedule, does not allocate, does not sample, does not launch. Most +# of what follows is context for reading a log, not work a target does. +# --------------------------------------------------------------------------- +_overview: + summary: >- + A request enters through the LLM API, is admitted and batched by the + executor, runs one forward through the model engine, and leaves through + the decoder and sampler. A modeling_v2 target replaces exactly one link + in that chain --- the forward --- and inherits every other link unchanged. + Knowing which link owns a symptom is most of debugging one. + flow: + - stage: LLM API + does: >- + Resolves args, loads the checkpoint, builds the executor. Where + `TorchLlmArgs` fields such as `enable_autotuner` and the runtime + sizing take effect, and where `TRTLLM_MODELING_V2` decides whether a + target is used at all. + areas: [quantization, weight_loading, host_offload] + - stage: executor loop + does: >- + Admits requests, forms batches, decides what runs this iteration. + Owns in-flight batching, the overlap scheduler, chunked prefill and + the KV cache. A target must tolerate all of it and controls none. + areas: [scheduling, kv_cache, disaggregation] + - stage: model engine + does: >- + Prepares the per-iteration metadata --- attention metadata, spec + metadata, cache pointers --- captures CUDA graphs, and calls the + forward. This is the layer immediately above a target and the one + whose contract a target must match exactly. + areas: [cuda_graph, jit_autotune, parallelism] + - stage: the forward + does: >- + THE TARGET. One flat function assembled only from catalog entries: + embedding, the decoder stack, the norm and the output projection. + Everything it calls is vocabulary; nothing it calls is invented here. + areas: [attention, mla, sparse_attention, mamba, gated_deltanet, kda, hyper_connection, moe, eplb, mlp, lm_head, embedding, normalization, multimodal, dwdp] + - stage: decode and sample + does: >- + Turns logits into tokens --- sampling, speculative verification and + acceptance. A target emits logits and draft state; it does not decide + what is accepted. + areas: [speculation, sampling] + +attention: + summary: >- + Paged attention and its backend families. A target does not implement + attention; it calls into a backend and fills a metadata object, so what a + target must get right is which metadata fields its variant reads and what + the KV pool layout is. + code: + - tensorrt_llm/_torch/attention + symbols: [AttentionMetadata, TrtllmAttentionMetadata, RopeParams, PositionalEmbeddingParams] + tests: + - tests/unittest/_torch/attention + doc: tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md + demands: [attention.shape, attention.head_dim, attention.qk_norm, attention.softcap, rope.plain, rope.scaling, rope.partial, rope.mrope] + gating: + gate: accuracy + signal: >- + The accuracy gate moves immediately: a wrong metadata field or KV layout does not degrade, it collapses. + blind: >- + More than it looks. A wrong RoPE scaling factor is near-invisible on + short benchmark prompts; which backend actually ran is only a warning + line on an unrecognized name; and over an fp8 pool the paged-context + path diverges from the bf16 one. Gate at a length that exercises + position, and record the backend. + +mla: + summary: >- + Multi-head latent attention. A different cache layout and a different op + family from ordinary paged attention --- compressed latent KV, with its + own RoPE and append path --- so a checkpoint with lora ranks is not a + small variation on the attention area, it is a separate one. + code: + - tensorrt_llm/_torch/attention/mla.py + symbols: [MLA, kv_lora_rank, q_lora_rank] + tests: + - tests/unittest/_torch/attention + doc: tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md + demands: [attention.mla] + gating: + gate: accuracy + signal: >- + Same as attention, plus the latent cache shape asserts at construction. + blind: >- + A wrong RoPE split can still score near-anchor on short prompts; gate at a length that exercises position. + +sparse_attention: + summary: >- + Windows, sinks and sparse backends. A window and a sink are MASKS --- the + certified contract states cache addressing is unchanged --- so they do not + change which kernel runs. `attention_chunk_size` is the one that constrains + the backend, and it raises rather than silently switching. + code: + - tensorrt_llm/_torch/attention/backends + symbols: [AttentionMetadata, attention_window_size, attention_sinks] + tests: + - tests/unittest/_torch/attention + doc: docs/source/features/sparse-attention.md + demands: [attention.windowing, attention.sinks, attention.sink, attention.hybrid, attention.indexer] + gating: + gate: accuracy + signal: >- + A window or sink that is silently dropped shows as a slow accuracy drift at long context, not at short. + blind: >- + Short-prompt accuracy. A dropped window is invisible until the sequence is longer than the window. + +mamba: + summary: >- + State-space layers and their caches. On the default path the recurrent + state lives INSIDE the unified KV pool and participates in block reuse; + what is genuinely its own is the slot indexing and the reset path. A + separate-pool manager exists but is opt-in and does not support reuse. + code: + - tensorrt_llm/_torch/modules/mamba + - tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py + symbols: [Mamba2Mixer, Mamba2Metadata, CppMambaHybridCacheManager, causal_conv1d_update] + tests: + - tests/unittest/_torch/modules/mamba + - tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py + demands: [recurrent.state, stack.hybrid] + gating: + gate: accuracy + signal: >- + State-cache misuse corrupts across requests, so accuracy drops once batching is on. + blind: >- + Single-request runs. A state that is never reset looks correct until a second request reuses the slot. + +gated_deltanet: + summary: >- + Gated DeltaNet, the linear-attention layer that interleaves with full + attention in hybrid stacks. Two prefill backends --- FlashInfer and a + vendored Triton --- chosen once per process by env and architecture, so a + receipt has to say which one it certifies. + code: + - tensorrt_llm/_torch/modules/mamba/gdn_mixer.py + - tensorrt_llm/_torch/modules/fla + symbols: [fused_gdn_gating, chunk_gated_delta_rule] + tests: + - tests/unittest/_torch/modules + demands: [recurrent.state, stack.hybrid] + gating: + gate: accuracy + signal: >- + The accuracy gate moves once state is reused across requests. The + backend choice is forceable by env, so the two can be compared --- but + they are different kernels and the in-tree parity test uses a + tolerance, not bit-equality. + blind: >- + Which branch ran. A receipt that does not name the branch certifies an implementation nobody can identify later. + +kda: + summary: >- + Kimi Delta Attention: a linear-attention family with its own metadata and + decode path, including a CuTe DSL op set. Distinct from GDN despite the + family resemblance. + code: + - tensorrt_llm/_torch/modules/kimi_kda + - tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py + symbols: [KimiKDALinearAttention, KimiK3MambaMetadata] + tests: + - tests/unittest/_torch/modules/kimi_kda + demands: [recurrent.state] + gating: + gate: accuracy + signal: >- + The accuracy gate moves once state is reused across requests, same as + the other recurrent areas. + blind: >- + Which linear-attention family actually ran. GDN and KDA both answer + `recurrent.state`, so the wrong one can be wired in and still produce + plausible text --- name the family in the receipt. + +hyper_connection: + summary: >- + Hyper-connections (HC) and multi-head hyper-connections (mHC) --- a + residual-topology change rather than a layer. A target that misses it + wires the residual stream wrongly while every individual op stays correct. + code: + - tensorrt_llm/_torch/modules/mhc + symbols: [mHC, HCState, HCHead] + tests: + - tests/unittest/_torch/modules + demands: [residual.hyper_connection] + gating: + gate: accuracy + signal: >- + A mis-wired residual topology moves accuracy hard. + blind: >- + Per-op tests. Every individual op passes while the stream between them is wired wrongly. + +moe: + summary: >- + Mixture of experts: routing, the fused runners, expert parallelism and + load balancing. The routing function and the shared-expert shape are what + a target must match; the runner is a backend choice. + code: + - tensorrt_llm/_torch/moe + symbols: [BaseMoeRoutingMethod, RoutingMethodType, ConfigurableMoE, create_moe] + tests: + - tests/unittest/_torch/moe + doc: tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md + demands: [moe.routing, moe.interleave] + gating: + gate: accuracy + signal: >- + Routing errors move accuracy. `ExpertStatistic` (`EXPERT_STATISTIC_PATH`, `EXPERT_STATISTIC_ITER_RANGE`) dumps the realised expert distribution. + blind: >- + Whether the intended experts ran. A uniform-ish distribution from wrong routing scores close to a correct one. + +mlp: + summary: >- + The dense feed-forward path and its activation. The smallest area here, + and the one most likely to be a single catalog entry. + code: + - tensorrt_llm/_torch/modules/mlp.py + - tensorrt_llm/_torch/modules/gated_mlp.py + symbols: [GatedMLP, MLP] + tests: + - tests/unittest/_torch/modules + demands: [mlp.activation] + gating: + gate: accuracy + signal: >- + A wrong activation or gate/up split collapses accuracy immediately. + blind: >- + Little. This is one of the few areas an accuracy gate sees fully --- + which is why it is usually a single catalog entry and not a project. + +lm_head: + summary: >- + The output projection and logits handling, including how logits are + gathered before the head. Matters to a target mainly through what it must + NOT do --- gathering logits for every token is a performance cliff. + code: + - tensorrt_llm/_torch/modules/logits_processor.py + - tensorrt_llm/_torch/modules/embedding.py + symbols: [LogitsProcessor, LMHead, Embedding] + tests: + - tests/unittest/_torch/modules + demands: [] + gating: + gate: accuracy + signal: >- + A wrong projection collapses accuracy; a wrong gather is a throughput cliff with identical output. + blind: >- + The gather. Logits gathered for every token instead of the last is correct and slow, and no accuracy gate sees it. + +speculation: + summary: >- + Speculative decoding --- MTP, Eagle3, dFlash, DSpark. The trap is that rejection + sampling preserves the target model's output distribution, so a + miscomputed draft path produces CORRECT text more slowly. Boot and + accuracy gates are blind to it; acceptance length is the signal. + code: + - tensorrt_llm/_torch/speculative + symbols: [SpecMetadata, accept_stats] + tests: + - tests/unittest/_torch/speculative + doc: docs/source/features/speculative-decoding.md + demands: [speculation.heads] + gating: + gate: acceptance + signal: >- + `acceptance_length` under `=require` versus `=off`. This is why the acceptance gate exists. + blind: >- + Boot and accuracy BOTH. Rejection sampling preserves the target distribution, so a miscomputed draft path produces correct text more slowly. Only acceptance length moves. + +cuda_graph: + summary: >- + Graph capture and piecewise compilation. A target becomes capture-hostile + through host-device syncs and shape-dependent control flow, so this is an + area a target must be written to respect rather than one it calls. + code: + - tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py + - tensorrt_llm/_torch/compilation + symbols: [CUDAGraphRunner, piecewise_optimizer] + tests: + - tests/unittest/_torch/executor/test_breakable_cuda_graph.py + - tests/unittest/_torch/compilation + doc: docs/source/features/torch_compile_and_piecewise_cuda_graph.md + demands: [] + gating: + gate: none + signal: >- + Capture emits warmup lines at startup; count them. Absence means + capture never ran. + blind: >- + The dangerous case is not a skipped capture, it is a SUCCESSFUL one + that replays wrongly --- host-computed values frozen into the graph. + Measured in tree: 88.2 on GSM8K with graphs against 96.2 eager. An + accuracy gate run eager is green while the graph path is corrupt, so a + claim about graphs must be measured WITH graphs on. + +scheduling: + summary: >- + The executor loop: in-flight batching, the overlap scheduler and chunked + prefill. A target does not call these; it has to tolerate them --- chunked + prefill in particular means a context call may run over a cached prefix. + code: + - tensorrt_llm/_torch/pyexecutor + symbols: [enable_chunked_prefill, PyTorchModelEngine] + tests: + - tests/unittest/_torch/executor + doc: docs/source/features/overlap-scheduler.md + demands: [] + gating: + gate: none + signal: >- + Chunked prefill is observable per-request through the context-chunk fields the executor tracks; the overlap scheduler shows in iteration timing. + blind: >- + Both gates, but narrow the claim: chunked prefill does change which + path runs --- it forces the paged-context route --- so it is not purely + a matter of timing. What the gates cannot see is + configured-but-never-triggered, which needs a prompt long enough to + chunk before anything is proven. + +kv_cache: + summary: >- + The paged KV cache manager and block reuse. Reuse changes which context + path runs --- a context call over a cached prefix --- which is why it can + move numbers that look like a modeling defect. + code: + - tensorrt_llm/_torch/pyexecutor/resource_manager.py + symbols: [KVCacheManager, enable_block_reuse] + tests: + - tests/unittest/_torch/executor + doc: docs/source/features/kvcache.md + # `attention.hybrid` is here as well as under sparse_attention: a stack whose + # two layer kinds declare different KV head counts asks the manager for two + # per-layer geometries, which is a pool question rather than a mask one. + demands: [attention.hybrid] + gating: + gate: none + signal: >- + `iterReusedBlocks`, `iterCacheHitRate`, `iterMissedBlocks` -- a second request sharing a prefix must move them. + blind: >- + Not what you would guess. Over a bf16 pool with nothing cached the + paged-context flag is bitwise inert --- but over an fp8-e4m3 pool it is + NOT inert even then, and with a cached prefix it is load-bearing and its + absence is silent. So an fp8 target's accuracy gate does see this; a + bf16 one does not. What stays blind either way is whether reuse ever + actually triggered --- which is what left it 'configured, unverified'. + +parallelism: + summary: >- + Tensor, pipeline, expert and data parallelism, and the collectives under + them. The parallel segment is part of a target's identity, so this area is + never optional --- even a tp1 target declares where it sits. + code: + - tensorrt_llm/_torch/distributed + - tensorrt_llm/mapping.py + symbols: [Mapping, allgather, reducescatter, enable_attention_dp] + tests: + - tests/unittest/_torch/multi_gpu + - tests/unittest/_torch/modeling_v2/comm + doc: docs/source/features/parallel-strategy.md + demands: [] + gating: + gate: boot + signal: >- + Rank count in the launch, and the collectives' own catalog receipts at that world size. + blind: >- + A run that silently fell back to one rank. Boot succeeds; it is just not the topology you asked for. + +quantization: + summary: >- + Quantized weight and activation formats, and the per-module exclusions a + checkpoint declares. What a target must get right is which modules are + held out and what scale layout each format expects. + code: + - tensorrt_llm/_torch/modules/linear.py + - tensorrt_llm/_torch/model_config.py + - tensorrt_llm/models/modeling_utils.py + symbols: [LinearMethodBase, QuantConfig, exclude_modules, kv_cache_quant_algo] + tests: + - tests/unittest/_torch/modules + doc: docs/source/features/quantization.md + demands: [quantization] + gating: + gate: accuracy + signal: >- + A wrong scale layout or a missed exclusion moves accuracy hard. + blind: >- + Which modules were actually held out. Quantizing a module the checkpoint excluded degrades gradually, not visibly. + +disaggregation: + summary: >- + Prefill and decode on SEPARATE engines, with per-sequence state moved + between them by the cache transceiver. Reads as a serving topology and is + a modeling property: only state the transceiver enumerates survives the + hop, so a forward that keeps state anywhere else is correct in every + single-engine measurement and wrong across it. + code: + - tensorrt_llm/_torch/disaggregation + - tensorrt_llm/_torch/pyexecutor/resource_manager.py + symbols: [KvCacheTransceiver, create_kv_cache_transceiver, CtxTransferStatus] + tests: + - tests/unittest/_torch/disaggregation + doc: docs/source/features/disagg-serving.md + demands: [] + gating: + gate: acceptance + signal: >- + Token-identical greedy output between one aggregated engine and a + context+generation pair, same session, prompts long enough to put every + state plane away from its initial value. Identity, not similarity --- + dropped recurrent state yields text that is fluent and wrong. + blind: >- + Both standing gates. A single-engine boot and a single-engine accuracy + run never cross the hop, so neither can observe state that fails to + travel. For a hybrid checkpoint the exposure is the recurrent conv/SSM + state and any sparse-attention indexer state, not the attention pages + everyone remembers to check. + +multimodal: + summary: >- + Vision and audio towers in front of a language model. For a `text_only` + scope the relevant part is the opposite of the feature --- how the text + tower loads and routes when the vision path is deliberately absent. + code: + - tensorrt_llm/_torch/models/modeling_multimodal_mixin.py + - tensorrt_llm/_torch/models/modeling_multimodal_utils.py + - tensorrt_llm/inputs + symbols: [MultimodalParams, InputProcessor, register_input_processor] + tests: + - tests/unittest/_torch/multimodal + doc: docs/source/features/multi-modality.md + demands: [multimodal.tower] + gating: + gate: accuracy + signal: >- + For a `text_only` scope the gate is that the text tower loads and routes with the vision path deliberately absent. + blind: >- + Whether the vision path was absent by design or absent by accident -- both look the same from a text benchmark. + +normalization: + summary: >- + Normalization layers and their fused residual-add variants. Small, but the + single most common source of a silently wrong target, because the fused + and unfused forms differ in what they write back in place. + code: + - tensorrt_llm/_torch/modules/rms_norm.py + - tensorrt_llm/_torch/modules/layer_norm.py + symbols: [RMSNorm, LayerNorm] + tests: + - tests/unittest/_torch/modules + demands: [norm.style] + gating: + gate: accuracy + signal: >- + Epsilon and the fused/unfused choice both move accuracy. + blind: >- + In-place write-back. A fused residual-add that does not write back where the caller expects is correct for one layer. + +eplb: + summary: >- + Expert-parallel load balancing. Redistributes experts across ranks so a + hot expert does not stall its whole rank. It changes occupancy and + throughput, never output --- which is exactly what makes it hard to prove. + code: + - tensorrt_llm/_torch/moe/fused_moe/moe_load_balancer.py + - tensorrt_llm/_torch/moe/expert_statistic.py + symbols: [MoeLoadBalancer, MoeLoadBalancerIterContext, ExpertStatistic] + tests: + - tests/unittest/_torch/moe + doc: tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md + demands: [] + gating: + gate: none + signal: >- + `ExpertStatistic` is the witness: dump the per-expert token counts and + assert the spread narrows with balancing on versus off. + blind: >- + Everything. Load balancing changes throughput and expert occupancy, not + output --- accuracy is identical either way, so only the distribution + dump can tell you it ran. + +dwdp: + summary: >- + Distributed Weight Data Parallelism: MoE EXPERT weights shard across + ranks while every other module stays a full replica, with remote regions + reached through a page pool of fabric handles. Not reachable from a + modeling_v2 target today --- it hooks in through the configurable MoE + layer, which a flat target does not have. + code: + - tensorrt_llm/_torch/modules/dwdp + - tensorrt_llm/_torch/pyexecutor/dwdp.py + symbols: [DwdpManager, DWDPWeightManager] + tests: + - tests/integration/defs/accuracy/test_dwdp_aggregated.py + - tests/integration/defs/accuracy/test_dwdp_disaggregated_serving.py + demands: [] + gating: + gate: accuracy + signal: >- + Correctness is covered by the aggregated accuracy gate; the page-pool + split is observable through `DwdpManager` and its own integration tests. + blind: >- + Narrower than it sounds: every misconfiguration RAISES, and setup + emits its own log lines. The one silent path is omitting the config + entirely, which simply runs the ordinary MoE. Do not assume bitwise + equality with it --- the expert reduction order differs. + +jit_autotune: + summary: >- + Four separate things compile or choose at RUN time, not build time: the + tactic autotuner, NVRTC-compiled kernels, Triton, and FlashInfer's JIT + workspace. None of them is a target's code, and all of them can change + what a target's numbers mean. + code: + - tensorrt_llm/_torch/autotuner.py + symbols: [AutoTuner, autotune, TunableRunner, AutoTunerProfilingCache] + tests: + - tests/unittest/_torch/misc/test_autotuner.py + doc: docs/source/torch/adding_custom_kernels.md + demands: [] + gating: + gate: none + signal: >- + Whether the autotuner was warm. `enable_autotuner` is a `TorchLlmArgs` + field defaulting to True, and warmup calls `_run_autotuner_warmup`, + which runs a real forward inside `autotune()`. A cold process and a + warm one select different tactics. + blind: >- + Every gate, and this one is subtle enough to be worth the sentence: a + catalog receipt taken in a cold process certifies a DIFFERENT numeric + path from the one serving takes. Measured on a fused MoE runner --- + over one shape's tactic space, two distinct outputs from one GEMM and + 101 from the other, all within tolerance, none bit-identical. The + entry is not wrong; the receipt just does not describe the warm path + unless it says it does. + +weight_loading: + summary: >- + How a checkpoint's tensors reach the modules: the loader, the config + loader and the name mapper. Every modeling_v2 target ships a `weights.py` + and the boot gate checks manifest coverage, so this is the surface an + onboard touches most and the one least visible from the forward. + code: + - tensorrt_llm/_torch/models/checkpoints + symbols: [BaseCheckpointLoader, BaseWeightLoader, BaseWeightMapper, BaseConfigLoader] + tests: + - tests/unittest/_torch/executor/test_model_loader_checkpoint_io.py + doc: docs/source/features/checkpoint-loading.md + demands: [quantization] + gating: + gate: boot + signal: >- + Manifest coverage: every tensor in the checkpoint is claimed by exactly + one module, and nothing is claimed twice. The boot gate fails loudly on + a missing tensor. + blind: >- + A tensor loaded into the WRONG module of the right shape. Coverage is + complete, boot succeeds, and only accuracy moves --- which is why a + transposed or mis-permuted projection reads as a modeling defect. + +host_offload: + summary: >- + Weights that do not stay resident. A tensor whose size is comparable to + the rest of the model lives in pinned host memory and the forward gathers + the rows a step needs. Output-identical by construction, which is the + whole difficulty: nothing in the emitted text distinguishes an offloaded + table from a resident one. + code: + - tensorrt_llm/_torch/pinned_weight_staging.py + - tensorrt_llm/_torch/modules/embedding.py + - tensorrt_llm/_torch/modules/qwen4_exp/ple.py + symbols: [staging_scope, Embedding, Qwen4ExpPinnedHostEmbedding] + tests: + - tests/unittest/_torch/misc/test_pinned_weight_staging.py + demands: [embedding.huge_table] + gating: + gate: acceptance + signal: >- + Device memory after weight load, measured with and without the feature + in the same session, against a delta predicted from the checkpoint's own + shapes --- plus token-identical greedy output between the two. + blind: >- + Boot and accuracy, completely. A target that quietly loaded the table + onto the device passes both and reports a green feature unit. Also blind + to WHERE the copy runs: a gather on the compute stream serializes the + step and costs nothing in correctness. + +embedding: + summary: >- + Token embedding and the output head, including tied weights and vocab + padding. First and last link of the forward, and the place where a + parallel target decides whether the vocab is sharded. + code: + - tensorrt_llm/_torch/modules/embedding.py + symbols: [Embedding, LMHead] + tests: + - tests/unittest/_torch/modules + demands: [embedding.huge_table] + gating: + gate: accuracy + signal: >- + A wrong embedding or an untied head collapses accuracy immediately. + blind: >- + Vocab padding. Padded rows that are never masked contribute logits for + tokens the tokenizer cannot emit --- rare enough to survive a short + benchmark and wrong everywhere. + +sampling: + summary: >- + Turning logits into tokens: greedy, top-k/top-p, and the seeding that + makes a run reproducible. A target emits logits and never samples, but + the sample type is pinned into a captured CUDA graph. + code: + - tensorrt_llm/_torch/pyexecutor/sampler + symbols: [Sampler, TorchSampler] + tests: + - tests/unittest/_torch/sampler + doc: docs/source/features/sampling.md + demands: [] + gating: + gate: boot + signal: >- + The boot gate is literally a set of greedy continuations, so a broken + sampler fails it outright. + blind: >- + Anything other than greedy. The boot gate exercises one sampling path + and the accuracy gate usually exercises the same one, so top-p and + seeding go untested by both. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/SKILL.md new file mode 100644 index 000000000000..3b8957965d8a --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/SKILL.md @@ -0,0 +1,214 @@ +--- +name: staircase-kernel-feasibility +description: Load when deciding whether a checkpoint can be assembled from the current kernel catalog, and what has to be added first — read the checkpoint and its HF reference for the computation-step list, map each step onto the catalog, and produce vocabulary.yaml with the dispatch queue. Also answers it standalone, before committing to an onboarding. +--- + +# staircase-kernel-feasibility + +Answers one question: **can this checkpoint be assembled from the catalog as +it stands, and if not, what has to be added first?** + +It runs inside the run's **plan phase** — the PlanDrafter does this judging +itself, and nobody is dispatched to do it — but it is worth running on its +own: before committing hours to an onboarding, it tells you whether the model +is reachable at all. Its product is a file, not a conclusion in someone's +context: the mapping is referenced by every dispatch that follows and it is +the recovery state after a session cut. + +You do this yourself. You do not dispatch to answer it — see "Why the judging +stays here". + +## Flow + +### 1. Semantic pre-audit + +Read the checkpoint's `config.json`, its safetensors index (key names, shapes, +dtypes), and the architecture's HF reference modeling code — for +**computational semantics only**, what the model computes. + +Walk `references/derivation-checklist.md` and record this checkpoint's value +on every axis. The output is a **computation-step list**: each step of the +forward with the values that step will be called at. + +**The parallel axis is not in the checkpoint.** Collectives are an artifact of +the topology, so they appear nowhere in the HF reference; a pre-audit driven +only by that reading comes back all-green for a `tepN`/`depN`/`unified` target +and the assembler discovers the gap mid-run. Walk the segment separately: what +the topology partitions, what must be reduced or gathered to put it back +together, and whether entries for that exist. `references/trtllm-ops.md`'s `comm` +category is where the candidates are. + +### 2. Vocabulary mapping + +Map every step onto +`tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml` and the +entries' contracts. Each row lands in one of three states — the judging rules +are in `references/mapping-and-triage.md`: + +| state | meaning | +|---|---| +| **OK** | an entry exists and its certified column covers this checkpoint's values | +| **CERT** | an entry exists; this surface is not covered | +| **GAP** | no entry | + +### 3. Ceiling check + +For every GAP row, look for a candidate in `references/trtllm-ops.md`. **No +candidate is the vocabulary ceiling** — + +**except that the inventory cannot see pybind bindings, and one of the most +capable entries in the catalog is one.** `references/trtllm-ops.md` is generated +from the torch schema registry, so it lists `torch.ops.*` only. +`attention/thop_attention` wraps `tensorrt_llm.bindings.internal.thop.attention`, +a pybind binding carried as an approved exception — it is not in the inventory +and never will be. A row whose only server is a binding therefore comes back +empty from an inventory search: **correctly empty, and wrongly read as a +ceiling.** This has cost a feature once: a vision tower's cache-free attention +was recorded as a ceiling and the feature descoped, when `thop_attention` served +it all along in a no-KV-cache mode (all four pool arguments `None`) that +trtllm's own diffusion stack runs. + +So before writing "ceiling", check the **catalog** as well as the inventory, and +ask whether an existing entry's *uncertified* surface already covers the step — +a `VERIFY` row is not a gap. Prefer a wrong "verify" to a wrong "ceiling": the +first costs a dispatch, the second costs a feature — this checkpoint cannot be assembled +today, and you can say so here, before dispatching anyone. Record the step, +what it computes, and why nothing upstream provides it. + +The mirror surface is glue — layout, movement, allocation, lookup — and a +missing **computation** kernel (norm, rope, activation, linear, attention, +quantization) is a ceiling rather than something to build there. **One narrow +exception**, and it is narrow on purpose: a computation may become a mirror when +all three hold — + +1. it is **exactly one torch callable**, not a computation assembled from + several mirrors; +2. **no trtllm op serves it**, established by a search closed over the inventory + *and* the pybind surface the inventory cannot show; +3. upstream owns its correctness, which is what lets a mirror be contract-free. + +Composing a computation out of several mirrors is still a ceiling, and reaching +for a mirror while a trtllm op exists is still wrong — the catalog prefers the +engine's kernels and a mirror is the last resort, not the convenient one. +Precedent: `catalog/torch/gelu_tanh.py`, added after the candidate set for +non-gated `gelu_tanh` at bf16 was searched closed. + + +### 4. Write `vocabulary.yaml` + +`/vocabulary.yaml` — the run's typed mapping. Out of tree this was +prose, because the only thing reading it was another language model; here the +orchestrator reads it, builds its dispatch queue from it, and **rejects a +malformed one**. The schema is authoritative and lives at +`agent_flow/workflows/staircase_onboard/vocabulary.py`; what follows is what +it validates. + +```yaml +verdict: assemblable | assemblable_after_additions | ceiling + +steps: # every computation step, both chapters + - step: + values: {: , ...} + state: OK | CERT | GAP + entry: catalog//.py # OK and CERT rows + contract_line: "" + candidate: torch.ops.trtllm. # GAP rows + +queue: # one dispatch per entry, in dependency order + - {n: 1, mode: onboard, target: , depends_on: []} + - {n: 2, mode: verify, target: , depends_on: [1], + surface: ""} + - {n: 3, mode: onboard, target: , depends_on: [], + unit: } # optional; core when absent + +ceiling: null # or {step: ..., searched: [...], why: ...} +``` + +Chapter 1 is the `OK`/`CERT` rows (entries this checkpoint will use); chapter 2 +is the `GAP` rows (callables with no entry). The `queue` is the two chapters as +dispatches — `CERT` → `mode: verify` → `staircase-kernel-verify`, `GAP` → +`mode: onboard` → `staircase-kernel-onboard`. An `OK` row dispatches nobody. + +Every rule below is enforced in Python, and every violation is reported in one +error rather than one per round-trip: + +| rule | why | +|---|---| +| `verdict` is exactly one of `assemblable` / `assemblable_after_additions` / `ceiling` | it is the headline, and the orchestrator branches on it | +| every `steps` row names a non-empty `step` and a `state` of `OK` / `CERT` / `GAP` | a row with no state is not a mapping | +| an `OK` or `CERT` row names an `entry` **and** quotes a `contract_line` | the quoted line is the evidence for the judgement; without it nobody can audit the row without redoing the reading | +| a `GAP` row names a `candidate` | a GAP with no candidate anywhere *is* the ceiling — say so in `ceiling` rather than leaving the row bare | +| `queue[].n` is an integer and unique across the queue | it is the dispatch's identity and what `depends_on` points at | +| `queue[].mode` is `onboard` or `verify` | it selects the skill smith loads | +| `queue[].target` names exactly one catalog entry | one dispatch, one entry | +| **no `target` appears under both modes** | the two chapters do not overlap (below) | +| `depends_on` is a list of dispatch numbers, all present in the queue | an unknown dependency cannot be ordered, and a cycle is rejected when the queue is layered into concurrent batches | +| a `verify` dispatch names a `surface` | it becomes smith's prompt almost verbatim, so vagueness here costs a whole run | +| `queue[].unit`, when present, is one of the run's assembly units — `core` or a feature from `task.yaml` | it decides what a failed dispatch blocks (below); a unit the run does not build could only be guessed at | +| `verdict: ceiling` carries a `ceiling` mapping naming the exact `step` and what was `searched` | a ceiling is the one verdict that ends the run, so it carries its evidence or it is not a verdict | +| any other verdict leaves `ceiling` unset | the two must agree | +| `assemblable` means an empty queue; `assemblable_after_additions` means a non-empty one | "nothing has to be added first" plus three dispatches is not a verdict | + +The two chapters do not overlap, and an entry onboarded in chapter 2 never +appears in chapter 1: smith writes its test at this checkpoint's values by +construction. The schema enforces it, and the error names both modes it saw. + +Order the queue by dependency. The category (verify or onboard) says nothing +about order — a chapter-1 row whose test state must be built by a chapter-2 +entry simply sits after it. Everything with no dependency runs **concurrently**, +so `depends_on` is for real ordering only: a spurious edge serializes a batch +and costs wall-clock directly. + +Tag every row that only a feature increment needs with `unit:` — the first +unit that uses the entry. Queue those rows anyway: every row is dispatched in +the smith phase, concurrently with core's, so a feature's entries cost no extra +wall-clock. The tag changes only what a **failure** blocks. An untagged row +reads as `core`, and a `core` row that fails stops the run before assembly. A +tagged row that fails is deferred instead, and retried once before its unit's +first assembler turn; the run stops only if that retry fails too, and then at +that unit rather than before core. Measured on DeepSeek-V4.1: a 33-dispatch +queue with six feature-only rows — among them the Engram host-gather ports, +the ViT attention verify, a DSpark-only routing surface and the MXFP4 drafter +MoE — any one of which could close the run before core started. + +- Tag the **first** unit that needs the entry, in `task.yaml` feature order. +- Leave the tag off when core uses the entry at all, or when you are unsure. + Untagged is the blocking direction, so it is the safe one. +- A row that a `core` row `depends_on` is core's, transitively, whatever its + own tag says: the orchestrator will not start core without it. +- A `note` saying "dspark only" is not a tag. Nothing reads notes. + +## Why the judging stays here + +Reading a contract's certified column is what contracts are *for*. The entry +spec puts it plainly: "Whoever calls this op reads the contract and nothing +else — no trtllm source", and "Nothing outside this entry reads trtllm source. +The contract must stand alone." A contract only smith can interpret would be a +failed contract. + + +What genuinely needs smith is a row with **no entry**: only reading source +tells you what that callable is, whether it can be one entry, and what it +needs beforehand. + +So: **judge here, dispatch to act.** Never dispatch because a row is unclear — +read the contract again. + +## Report + +Whether run inside an onboarding or standalone, say: + +- the verdict: assemblable / assemblable after N additions / ceiling +- the counts by state, and the queue length +- for a ceiling: the exact step, and what was searched +- for CERT rows: the value and the contract line that excludes it — this is + what smith is dispatched with, so vagueness here costs a whole run diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/references/derivation-checklist.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/references/derivation-checklist.md new file mode 100644 index 000000000000..4fbfb36a2ee9 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/references/derivation-checklist.md @@ -0,0 +1,99 @@ +# What varies per model — the re-derive list + +> Extracted from the former `trtllm-runtime-integration.md` §10; the section +> number is kept so existing citations still resolve. + +## 10. What varies per model — the re-derive list + +The exemplar target shows the pattern; every model-specific value must be +re-derived from the new checkpoint's `config.json` and the catalog +contracts, never copied. The known variation axes: + +**attention structure** — GQA/MHA vs **MLA**, the axis that reshapes the +most: MLA replaces per-head K/V with one compressed latent, needs two +attention calls per forward (context and generation are separate call +shapes, mixed batches rejected upstream), and brings its own load-time +obligations. It is not a variant of the geometry row below; see +`.claude/skills/staircase-kernel-onboard/references/decompositions/mla-custom-op.md` for the vocabulary and +`.claude/skills/staircase-modeling-core/references/latent-kv-cache.md` for what a target had to derive. +Then: geometry (layers, heads, kv-heads, head_dim, intermediate, vocab); +qk-norm presence and its eps; rope kind, theta, scaling, `is_neox`, +partial-rotary factor; `tie_word_embeddings` (changes the lm_head/embed +feeding and the shell's tying path); attention bias (needs cublas_mm's +fused-bias argument); sliding window; activation function; checkpoint +key names and fusion grouping; dtype and quantization scheme — including +*which modules it excludes*, since a checkpoint may quantize only its MLP +path and leave attention bf16; parallel topology (the routing decision and +the boot gate's topology flags + per-rank shapes + manifest src transforms + +collective-communication catalog entries — §12 has the launch substrate, the +rest is the first multi-rank target's to derive); KV-cache dtype (turns the +kv-scale constants into post-load tensors — certification extension first). + +A construction-time assert exists for each axis the assembly depends on: +if the new config violates one, the right response is to re-derive that +part of the assembly, not to delete the assert. + +### Read the axes off the engine's config object, not off AutoConfig + +`model_config.pretrained_config` — the object the engine hands +`__init__` — is the **un-migrated** config: fields sit where the +checkpoint's own `config.json` put them. A standalone +`AutoConfig.from_pretrained(model_dir)` probe can return a *different* +shape of the same information, because transformers migrates fields +across versions. + +Observed at transformers 5.5.4 on a checkpoint written by 4.51: through +AutoConfig, rope had been migrated into +`rope_parameters = {'rope_theta': ..., 'rope_type': 'default'}` with +`rope_scaling` an alias of that same dict, and `cfg.rope_theta` raising +`AttributeError`. Through the engine, `rope_parameters is None`, +`rope_scaling is None`, and `rope_theta` is a plain instance attribute. +Deriving the axis from the AutoConfig surface produced a target that +failed at construction with `TypeError: 'NoneType' object is not +subscriptable`. + +So: an AutoConfig probe is **not** a valid oracle for what +`pretrained_config` will look like. Read the checkpoint's `config.json` +directly to learn what the model *is*, and use the flat +`cfg.` idiom inside the target. If a probe is needed, instrument +the target's own `__init__`. + +### A missing `dtype` is not inherited + +A checkpoint may declare **no** `dtype` and no `torch_dtype` at all +(observed on gpt-oss-120b, whose weights are all BF16 outside the +quantized expert blocks). The two config surfaces then disagree: +`model_config.pretrained_config.torch_dtype` is `None` while the engine's +own `ModelConfig.torch_dtype` resolves `torch.bfloat16`. The shell reads +the **pretrained** one, so `DecoderModelForCausalLM` materializes its +shell-owned parameters — `lm_head.weight` — as **fp32**, and a target with +a dtype-checking weight manifest fails there with a cause two layers +upstream of the symptom. + +There is no longer a target-owned stub to fix it in: **the checkpoint is +read exactly as published**, so the divergence cannot be patched at the +config and the target has to resolve the dtype itself. Assert both surfaces +in `__init__` so a future drift is loud. + + +### Where the chat template lives is not a `tokenizer*` glob + +An accuracy protocol that applies a chat template needs the template +files, and a recent checkpoint may keep **no** `chat_template` inside +`tokenizer_config.json` — it lives in `chat_template.jinja` (plus +`chat_template.json`). Those names, and `special_tokens_map.json`, match +no `tokenizer*` pattern. + +The old failure mode was a link step: a target-owned `model_dir/` populated +by a fixed `tokenizer*` glob dropped the template, and the gate failed at +template application after a full engine build. That step is gone — the +checkpoint is read exactly as published, so every file it ships is already +there. What survives is the reason the glob was wrong: **inspect what the +checkpoint actually ships** when naming the protocol's template, rather than +assuming where a template lives. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/references/mapping-and-triage.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/references/mapping-and-triage.md new file mode 100644 index 000000000000..769c3d546bbe --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/references/mapping-and-triage.md @@ -0,0 +1,83 @@ +# Mapping and triage — reading a contract to place a step + +How to turn one computation step plus this checkpoint's values into one of +three states, and what each state dispatches. + +## Reading the certified column + +A contract has four places that bound what it certifies. Read all four; the +receipt alone is not the answer. + +| where | what it bounds | +|---|---| +| `receipts:` frontmatter | the **arch** the test ran under, and `world_size` for a multi-rank entry. Three-state: `passed` / `failed` / **absent = unknown** | +| **Signature** | shapes, dtypes, layout, contiguity, device per argument | +| **Metadata consumed** | which runtime fields must be populated, and with what | +| **Preconditions** | everything that must hold beforehand — buffers, weight formats, alignment, supported dtypes | + +Two reading rules that have each cost a run: + +**Where a contract enumerates tested shapes rather than stating a rule, the +enumeration is the column.** "Tested at head counts 8, 16, 128" does not +certify 32. A rule ("any head count divisible by 8") does. + +**An arch key that does not match the arch this run targets makes the row +CERT, not OK**, however green the receipt looks. Absent means unknown. +Receipts are keyed by GPU architecture and nothing else — there is no version +pin in tree to drift against — but the freshness rule is unchanged in +substance: a receipt is valid only against the files of *its own entry*, and +an entry now spans two trees (contract and wrapper in the package, GPU test in +the tests tree), so a receipt that predates a write to either is a claim about +files nobody ran. + +## The four outcomes + +| situation | state | dispatch | +|---|---|---| +| entry exists, all four places cover this checkpoint's values, the receipt covers this run's arch | **OK** | still goes to chapter 1 — `staircase-kernel-verify` confirms the receipt cheaply, no GPU | +| entry exists, some value falls outside | **CERT** | `staircase-kernel-verify`, with the value and the contract line that excludes it | +| no entry, but `references/trtllm-ops.md` has a candidate callable | **GAP** | `staircase-kernel-onboard`, with the candidate and the step it serves | +| no entry and no candidate | **ceiling** | nothing — the run ends here. **First rule out a pybind binding**: the inventory is generated from the torch schema registry and lists `torch.ops.*` only, so a step served by a binding (`thop_attention` is one) searches empty and reads as a ceiling. Check the catalog's existing entries and their *uncertified* surfaces before concluding | + +One more case that dispatches nobody: **a missing +`tensorrt_llm/_torch/_experimental/modeling_v2/catalog/torch/` mirror** +for pure glue — tensor layout, movement, allocation, lookup (reshape, split, +concat, copy, empty, expand, embedding, residual add). assembler adds those +itself. It is glue by default: a computation the HF reference expresses as a +named module is never *composed* from mirrors, and a missing kernel there is a +ceiling rather than a mirror to add. + +**The one exception**, narrow by design: a computation may be a mirror when it is +**exactly one torch callable**, **no trtllm op serves it** (searched closed over +the inventory *and* the pybind surface the inventory cannot show), and upstream +owns its correctness. Composition is still a ceiling; a mirror chosen while a +trtllm op exists is still wrong. Precedent: `catalog/torch/gelu_tanh.py`. Such a +row is neither GAP nor ceiling — assembler adds the mirror, like any other. + + +## Precision in a CERT row + +A CERT row becomes smith's dispatch prompt almost verbatim, so vagueness costs +a whole run. It must carry: + +- the entry path +- **the value**, not the axis: `nope=192, v=256, H=16 per rank, fp8 pool`, not + "MLA geometry" +- **the contract line that excludes it**, quoted. If you cannot quote one, the + row is probably OK and you are dispatching on a hunch + +## Two failure modes in this step + +**Reading the source instead of the contract.** Nothing forbids you from +opening it — but if a contract cannot answer whether a value is covered, that +is a **defect in the contract**, and patching around it by reading source +leaves the next reader with the same problem. Record the defect; the fix is +smith rewriting the contract from an observed run, which is what +`staircase-kernel-verify` does. + +**Grouping rows to "save dispatches".** One dispatch, one entry. Entries that +share an expensive test fixture are still separate dispatches — catalog tests +may not import each other, so there is no fixture to share, only the knowledge +of how to build one, and that travels in smith's reports. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/references/trtllm-ops.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/references/trtllm-ops.md new file mode 100644 index 000000000000..58d046082943 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/references/trtllm-ops.md @@ -0,0 +1,423 @@ +# TensorRT-LLM Op Inventory + +> GENERATED by +> `agent_flow/workflows/staircase_onboard/tools/list_trtllm_ops.py --doc` +> — do not hand-edit. Editing a row here does not change what the engine +> registers; it only makes the inventory lie. +> +> **This lists `torch.ops.*` only.** Ops reached through a pybind binding +> rather than the torch dispatcher are invisible here and always will be — +> `tensorrt_llm.bindings.internal.thop.attention`, which the catalog wraps +> as `attention/thop_attention`, is the known case. A step absent from this +> file is not thereby unreachable; check the catalog too before calling it +> a vocabulary ceiling. +> +> **Regenerate after any change to the engine's registered op set** — there +> is no version pin in tree, so this file goes stale silently whenever the +> trunk adds, removes, or re-signs an op, and a mapping made against a stale +> inventory reads a live candidate as a ceiling. Output is deterministic for +> a given install **on a given GPU** (the line below carries the device name +> and sm version, so regenerating on different hardware diffs there without +> the op set having changed). +> +> The line below is the **provenance stamp of this snapshot**, not a pin to +> assert against: it records the install the sweep ran on, which is what +> would invalidate the rows. + +tensorrt_llm 1.3.0rc25 | torch 2.12.0+cu130 | NVIDIA B200 (sm_100) | 302 ops + +## UNCATEGORIZED (8) — rules table needs a decision + +``` +cute_dsl_minimax_m3_index_decode_score +fused_pertoken_adaln +fused_pertoken_adaln_residual +fused_sigmoid_mul_inplace +kda_decode +kda_prefill +minimax_m3_select_blocks +situ_and_mul +``` + +## Summary + +| category | ops | +| --- | --- | +| gemm | 41 | +| attention | 33 | +| moe | 65 | +| norm | 21 | +| activation | 11 | +| quantization | 34 | +| comm | 35 | +| embedding | 1 | +| kv_cache | 7 | +| sampling | 13 | +| speculative | 14 | +| ssm | 5 | +| internal | 9 | +| misc | 5 | + +## gemm (41 ops) + +``` +bmm_out(Tensor a, Tensor b, Tensor(a2!) out) -> () +cublas_mm(Tensor mat_a, Tensor mat_b, Tensor? bias, ScalarType? out_dtype, int output_buffer_kind=0, int[]? group=None) -> Tensor out +cublas_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scale_a, Tensor scale_b, Tensor? bias, ScalarType? out_dtype, int output_buffer_kind=0, int[]? group=None) -> Tensor out +cuda_core_nvfp4_gemm(Tensor mat_a, Tensor mat_b, Tensor scale_a, Tensor scale_b, Tensor alpha, Tensor? bias, ScalarType? out_dtype, int output_buffer_kind=0, int[]? group=None) -> Tensor out +cuda_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scale_a, Tensor scale_b, Tensor? bias, ScalarType? out_dtype, int output_buffer_kind=0, int[]? group=None) -> Tensor out +cute_dsl_bf16_bmm_blackwell(Tensor input, Tensor weight, Tensor(a2!) output, bool use_tvm_ffi=True) -> () +cute_dsl_bf16_gemm_blackwell(Tensor input, Tensor weight, Tensor(a2!) output, bool use_tvm_ffi=True) -> () +cute_dsl_fp8_bmm_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor(a4!) output, ScalarType output_dtype=15, bool use_tvm_ffi=True) -> () +cute_dsl_fp8_gemm_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, ScalarType output_dtype=15, bool use_tvm_ffi=True) -> Tensor +cute_dsl_nvfp4_dense_gemm_fc2_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha_scale, SymInt expert_count, SymInt weight_per_expert, ScalarType output_dtype, SymInt scaling_vector_size=16) -> Tensor +cute_dsl_nvfp4_gemm_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha, ScalarType output_dtype, SymInt output_buffer_kind=0, SymInt[]? group=None, bool use_tvm_ffi=True) -> Tensor +cutlass_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scale_a, Tensor scale_b, Tensor? bias, ScalarType? out_dtype) -> Tensor out +finegrained_mixed_dtype_gemm(Tensor input, Tensor weight, Tensor scales, SymInt group_size, bool has_zero_point, ScalarType output_dtype, float? alpha=None, Tensor? bias=None, Tensor? zeros=None) -> Tensor +fp4_bmm(Tensor mat1, Tensor mat2, Tensor mat1Scale, Tensor mat2Scale, Tensor globalScale, int fp4GemmType, ScalarType? out_dtype=None) -> Tensor +fp4_fp8_gemm_trtllmgen(Tensor mat1, Tensor mat2, Tensor mat2Scale, Tensor globalScale, ScalarType? out_dtype=None) -> Tensor +fp4_gemm(Tensor mat1, Tensor mat2, Tensor mat1Scale, Tensor mat2Scale, Tensor globalScale, int fp4GemmType, ScalarType? out_dtype=None) -> Tensor +fp4_gemm_trtllmgen(Tensor mat1, Tensor mat2, Tensor mat1Scale, Tensor mat2Scale, Tensor globalScale, bool sfUseUE8M0, ScalarType? out_dtype=None) -> Tensor +fp8_batched_gemm_trtllmgen(Tensor mat1, Tensor mat2, SymInt tile_size, bool? use_deep_seek_fp8=False, bool? low_latency=False, SymInt? epilogue_tile_m=0, Tensor? dq_sfs_a=None, Tensor? dq_sfs_b=None, Tensor? scale_c=None, ScalarType? out_dtype=5) -> (Tensor, Tensor) +fp8_block_scaling_bmm(Tensor mat1, Tensor mat2, Tensor mat1Scale, Tensor mat2Scale, ScalarType? out_dtype=None) -> Tensor +fp8_block_scaling_bmm_out(Tensor mat1, Tensor mat2, Tensor mat1Scale, Tensor mat2Scale, Tensor(a!) out) -> Tensor(a!) +fp8_block_scaling_gemm(Tensor a, Tensor b, Tensor a_scale, Tensor b_scale, SymInt tune_max_num_tokens=4096) -> Tensor +fp8_block_scaling_gemm_impl(Tensor mat1, Tensor mat2, Tensor mat1Scale, Tensor mat2Scale) -> Tensor +fp8_per_tensor_scaling_tllmg_gemm(Tensor mat1, Tensor mat2, Tensor global_scale, ScalarType? out_dtype=None, Tensor? global_scale_gate=None, bool low_latency_kernel=False, bool gated_silu=False) -> Tensor +fp8_rowwise_gemm(Tensor act, Tensor weight, Tensor act_scale, Tensor weight_scale, ScalarType output_dtype, SymInt output_buffer_kind=0, SymInt[]? group=None) -> Tensor +fp8_swap_ab_gemm(Tensor input, Tensor weight, Tensor weight_scale, ScalarType output_dtype=15, bool disable_ue8m0_cast=False) -> Tensor +gptq_marlin_repack(Tensor b_q_weight, Tensor perm, int size_k, int size_n, int num_bits, bool is_a_8bit=False) -> Tensor out +llama4_bf16_bf16_gemm(Tensor inputA, Tensor inputB) -> Tensor +llama4_fp8_bf16_gemm(Tensor inputA, Tensor inputB, Tensor scaling_factor, Tensor? position_ids=None) -> Tensor +lora_group_gemm_param_fill_row_reorder_fusion(Tensor in_sizes, Tensor out_sizes, Tensor a_ptrs, Tensor d_ptrs, Tensor d_prime_ptrs, Tensor lda, Tensor ldd, Tensor ldb_prime, Tensor ldd_prime, Tensor splitk_offsets, Tensor reordered_input, int max_lora_count, int max_lora_rank, int sum_output_hidden_size, int input_hidden_size, int batch_size, Tensor slot_counts, Tensor slot_ranks, Tensor slot_offsets, Tensor module_out_sizes, Tensor module_out_prefix, Tensor b_ptrs, Tensor b_prime_ptrs, Tensor input, Tensor sorted_ids, Tensor intermediate_buffer, Tensor output_buffer, ScalarType dtype) -> () +marlin_nvfp4_gemm(Tensor mat_a, Tensor mat_b, Tensor? scale_a, Tensor scale_b, Tensor? alpha, Tensor weight_global_scale, Tensor? bias, ScalarType? out_dtype, int size_n, int size_k, int output_buffer_kind=0, int[]? group=None) -> Tensor out +matmul_to_ub(Tensor a, Tensor b) -> Tensor +mhc_gemm_sqrsum_fma(Tensor x, Tensor w, Tensor(a!) y, Tensor(b!) r, int M, int N, int K, int tile_n=0, int tile_m=0) -> () +mxfp8_mxfp8_gemm(Tensor act, Tensor actScale, Tensor weight, Tensor weightScale, Tensor globalScale, ScalarType? out_dtype=None) -> Tensor +nvfp4_gemm(Tensor act_fp4, Tensor weight, Tensor act_sf, Tensor weight_scale, Tensor alpha, ScalarType output_dtype, SymInt output_buffer_kind=0, str allowed_backends="cutlass,cublaslt,cuda_core", SymInt[]? group=None, Tensor? bias=None) -> Tensor +nvfp4_gemm_allreduce(Tensor act_fp4, Tensor weight_fp4, Tensor act_sf, Tensor weight_scale, Tensor alpha, ScalarType output_dtype, SymInt tp_rank, SymInt[] tp_group) -> Tensor +nvfp4_gemm_cublaslt(Tensor act_fp4, Tensor weight, Tensor act_sf, Tensor weight_scale, Tensor alpha, ScalarType output_dtype, SymInt output_buffer_kind=0, Tensor? bias=None) -> Tensor +nvfp4_gemm_cutlass(Tensor act_fp4, Tensor weight, Tensor act_sf, Tensor weight_scale, Tensor alpha, ScalarType output_dtype, SymInt output_buffer_kind=0, Tensor? bias=None) -> Tensor +preprocess_weights_for_mixed_gemm(Tensor _0, ScalarType _1, ScalarType _2) -> Tensor _0 +tinygemm2(Tensor input, Tensor weight, Tensor bias) -> Tensor +w4a8_mxfp4_fp8_gemm(Tensor act_fp8, Tensor weight, Tensor act_sf, Tensor weight_scale, Tensor alpha, ScalarType output_dtype, SymInt output_buffer_kind=0) -> Tensor +weight_only_quant_gemm(Tensor activation, Tensor weight, ScalarType weight_dtype, Tensor weight_scale, ScalarType output_dtype, SymInt output_buffer_kind=0) -> Tensor +``` + +## attention (33 ops) + +``` +attention_supports_nvfp4_output(int _0, int _1, int _2, int? _3, int _4, int _5, bool _6, bool _7) -> bool _0 +attn_custom_op_inplace(Tensor q, Tensor? k, Tensor? v, str attention_mask, Tensor? mrope_rotary_cos_sin, Tensor? mrope_position_deltas, SymInt? attention_window_size, Tensor? attention_mask_data, Tensor? attention_sinks, Tensor? relative_attention_bias, SymInt relative_attention_max_distance, str layer_idx, Tensor(a12!) output, Tensor(a13!)? output_sf) -> () +attn_res_fwd(Tensor layer_residual, Tensor block_residual, Tensor res_weight, Tensor rms_weight, float rms_eps) -> (Tensor, Tensor, Tensor, Tensor) +create_attn_outputs(Tensor q, str attention_mask, str layer_idx) -> Tensor[] +create_mla_outputs(Tensor hidden_states, str layer_idx) -> Tensor +cute_dsl_dspark_attention(Tensor q, Tensor main_kv, Tensor block_kv, Tensor(a3!) kv_cache, Tensor slots, Tensor start_pos, Tensor attn_sink, float softmax_scale) -> Tensor +cute_dsl_dspark_rmsnorm_rope(Tensor x, Tensor weight, Tensor freqs, SymInt num_heads, SymInt rope_dim, float eps, bool apply_weight, bool apply_rmsnorm, bool inverse_rope) -> Tensor +cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell(Tensor input, Tensor weight, Tensor weight_scale, Tensor position_ids, Tensor cos_sin_cache, Tensor alpha, bool use_tvm_ffi=True) -> (Tensor, Tensor) +cute_dsl_mla_decode_fp16_blackwell(Tensor q_latent, Tensor q_rope, Tensor c_latent, Tensor c_rope, Tensor page_table, Tensor cache_seqs, Tensor(a6!) o, Tensor(a7!) workspace, SymInt num_heads, SymInt seq_len_q, SymInt page_size, float softmax_scale, float output_scale, SymInt max_batch_size=0) -> () +cute_dsl_mla_decode_fp8_blackwell(Tensor q_latent, Tensor q_rope, Tensor c_latent, Tensor c_rope, Tensor page_table, Tensor cache_seqs, Tensor(a6!) o, Tensor(a7!) workspace, SymInt num_heads, SymInt seq_len_q, SymInt page_size, float softmax_scale, float output_scale, SymInt max_batch_size=0) -> () +deepseek_v4_compute_sliding_block_tables(Tensor block_offsets, Tensor copy_idx, Tensor pool_ids, Tensor valid_pool, Tensor scales, Tensor layer_offsets, Tensor(a!) output) -> () +deepseek_v4_compute_sliding_block_tables_with_scratch(Tensor block_offsets, Tensor copy_idx, Tensor pool_ids, Tensor valid_pool, Tensor scales, Tensor layer_offsets, Tensor scratch_pages, Tensor scratch_begs, Tensor scratch_ends, Tensor scratch_slots, Tensor num_contexts, Tensor(a!) output) -> () +flashinfer_apply_rope_with_cos_sin_cache_inplace(Tensor positions, Tensor(a1!) query, Tensor(a2!) key, SymInt head_size, Tensor cos_sin_cache, bool is_neox=True) -> () +fused_dit_qk_norm_rope(Tensor(a!) qkv, int num_heads_q, int num_heads_k, int num_heads_v, int head_dim, float eps, Tensor q_weight, Tensor k_weight, Tensor? q_add_weight, Tensor? k_add_weight, Tensor cos_emb, Tensor sin_emb, SymInt num_txt_tokens, bool interleave, SymInt tokens_per_batch) -> () +fused_dit_split_norm_rope(Tensor(a!) tensor, int num_heads, int head_dim, float eps, Tensor weight, Tensor cos_emb, Tensor sin_emb, bool interleave) -> () +fused_inv_rope_fp8_quant_vllm_port(Tensor o, Tensor positions, Tensor cos_sin_cache, int n_groups, int heads_per_group, int nope_dim, int rope_dim, int quant_group_size, bool is_neox) -> (Tensor, Tensor) +fused_qk_norm_rope(Tensor(a!) qkv, int num_heads_q, int num_heads_k, int num_heads_v, int head_dim, int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, float factor, float low, float high, float attention_factor, bool is_qk_norm, bool use_gemma, bool use_mrope, int mrope_section1, int mrope_section2) -> () +fused_qk_norm_rope_to_fp8(Tensor qkv, int num_heads_q, int num_heads_k, int num_heads_v, int head_dim, int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, float factor, float low, float high, float attention_factor, bool is_qk_norm, bool use_gemma, bool use_mrope, int mrope_section1, int mrope_section2) -> Tensor +fused_qkv_gemma_rmsnorm_rope_gate(Tensor qkv, Tensor q_weight, Tensor k_weight, Tensor cos_sin, Tensor positions, float eps, SymInt num_q_heads, SymInt num_kv_heads, SymInt head_dim, SymInt rotary_dim, bool use_mrope, SymInt mrope_section1, SymInt mrope_section2) -> (Tensor, Tensor) +load_chunked_kv_cache_for_mla(ScalarType out_dtype, int num_contexts, int num_ctx_cached_tokens, Tensor cu_ctx_chunked_kv_lens, Tensor chunked_ld_global_offset, Tensor kv_cache_block_offsets, Tensor host_kv_cache_pool_pointers, Tensor host_kv_cache_pool_mapping, Tensor? kv_scale_quant_orig, int layer_idx, int lora_size, int rope_size, int tokens_per_block, int max_seq_len, int attention_window_size, int beam_width, int quant_mode) -> Tensor[] +load_paged_kv_cache_for_mla(ScalarType out_dtype, int num_contexts, int num_ctx_cached_tokens, int max_ctx_cached_kv_len, Tensor cu_ctx_cached_kv_lens, Tensor kv_cache_block_offsets, Tensor host_kv_cache_pool_pointers, Tensor host_kv_cache_pool_mapping, Tensor? kv_scale_quant_orig, int layer_idx, int lora_size, int rope_size, int tokens_per_block, int attention_window_size, int beam_width, int quant_mode) -> Tensor[] +merge_chunked_attention_for_mla(Tensor(a!) merged_attn, Tensor temp_attn, Tensor merged_softmax_stats, Tensor temp_softmax_stats, int num_requests, Tensor cu_q_seq_lens, int max_q_seq_len, Tensor merge_op, int num_heads, int head_size) -> () +minimax_m3_attn_custom_op_inplace(Tensor q, Tensor k, Tensor v, Tensor? idx_q, Tensor? idx_k, str layer_idx, Tensor(a6!) output) -> () +minimax_m3_fp8_indexer_qk_norm_rope(Tensor qk, Tensor(a!) index_k_cache, Tensor out_cache_loc, int num_heads_q, int head_dim, int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, Tensor position_ids) -> Tensor +mla_custom_op_inplace(Tensor hidden_states, Tensor? position_ids, str layer_idx, Tensor(a3!) output, Tensor? latent_cache_gen, Tensor? hidden_states_fp4=None, Tensor? hidden_states_sf=None) -> () +mla_dsa_attn_inplace(Tensor q, Tensor compressed_kv, Tensor k_pe, Tensor latent_cache, Tensor[] indexer_intermediates, Tensor? position_ids, str layer_idx, Tensor(a7!) output) -> () +mla_dsa_proj(Tensor hidden_states, Tensor? position_ids, str layer_idx, Tensor? hidden_states_fp4=None, Tensor? hidden_states_sf=None) -> Tensor[] +mla_rope_append_paged_kv_assign_q(Tensor q, Tensor latent_cache, int num_contexts, Tensor cu_ctx_cached_kv_lens, Tensor cu_seq_lens, int max_input_uncached_seq_len, Tensor cos_sin_cache, int head_num, int nope_size, int rope_size, int lora_size, Tensor kv_cache_block_offsets, Tensor host_kv_cache_pool_pointers, Tensor host_kv_cache_pool_mapping, Tensor? kv_scale_orig_quant, int layer_idx, int tokens_per_block, int attention_window_size, int beam_width, int quant_mode) -> () +mla_rope_generation(Tensor(a!)? fused_q, Tensor(a!)? q_pe, Tensor latent_cache, Tensor? rotary_cos_sin, Tensor cu_q_seqlens, Tensor cu_kv_seqlens, Tensor fmha_scheduler_counter, Tensor? mla_bmm1_scale, Tensor? mla_bmm2_scale, Tensor? quant_q_buffer, Tensor sequence_length, Tensor host_past_key_value_lengths, Tensor host_context_lengths, int num_contexts, Tensor? kv_cache_block_offsets, Tensor? host_kv_cache_pool_pointers, Tensor? host_kv_cache_pool_mapping, Tensor? kv_scale_orig_quant, Tensor? kv_scale_quant_orig, Tensor? out_scale, Tensor? block_ids_per_seq, Tensor?[] helix_tensor_params, int predicted_tokens_per_seq, int layer_idx, int num_heads, int num_kv_heads, int head_size, int tokens_per_block, int attention_window_size, int beam_width, int quant_mode, float q_scaling, int q_lora_rank, int kv_lora_rank, int qk_nope_head_dim, int qk_rope_head_dim, int v_head_dim, bool rope_append, Tensor? kv_norm_weight=None, float kv_norm_eps=9.9999999999999995e-07, bool precomputed_cu_seqlens=False, bool precomputed_fmha_scheduler=False, bool kv_only=False, bool kv_done_elsewhere=False, Tensor? quant_scale_qkv=None) -> () +mla_rope_inplace(Tensor(a!) data, Tensor position_ids, Tensor cos_sin_cache, int num_heads, int nope_dim, int rope_dim, bool inverse, bool is_neox) -> () +tensorrt_llm::pack_fmha_mask_by_input(Tensor _0, Tensor _1, Tensor _2, float _3) -> Tensor _0 +tensorrt_llm::pack_fmha_mask_by_type(Tensor _0, Tensor _1, int _2, int _3, int _4, int _5) -> Tensor _0 +tensorrt_llm::relative_attention_bias(Tensor _0, Tensor _1, int _2, int _3, int _4, bool _5, int _6) -> () +``` + +## moe (65 ops) + +``` +bf16_mxe2m1_block_scale_moe_runner(Tensor? routing_logits, Tensor? routing_bias, Tensor hidden_states, Tensor gemm1_weights, Tensor gemm1_weights_scale, Tensor? gemm1_bias, Tensor? gemm1_alpha, Tensor? gemm1_beta, Tensor? gemm1_clamp_limit, Tensor gemm2_weights, Tensor gemm2_weights_scale, Tensor? gemm2_bias, SymInt num_experts, SymInt top_k, SymInt? n_group, SymInt? topk_group, SymInt intermediate_size, SymInt? valid_hidden_size, SymInt? valid_intermediate_size, SymInt local_expert_offset, SymInt local_num_experts, float? routed_scaling_factor, SymInt routing_method_type, SymInt act_type, Tensor? topk_weights=None, Tensor? topk_ids=None, Tensor? output=None, SymInt tune_max_num_tokens=8192, bool use_dp=False) -> Tensor +cute_dsl_megamoe_nvfp4_blackwell(Tensor activation, Tensor activation_sf, Tensor topk_idx, Tensor topk_weights, Tensor fc1_weight, Tensor fc1_weight_sf, Tensor fc2_weight, Tensor fc2_weight_sf, Tensor fc1_alpha, Tensor fc2_alpha, Tensor fc1_norm_const, Tensor(a11!) combine_output, Tensor(a12!) shared_workspace, SymInt world_size, SymInt local_rank, SymInt num_topk, SymInt num_experts_per_rank, SymInt hidden_size, SymInt intermediate_size_per_partition, SymInt expand_intermediate_size_per_partition, SymInt max_tokens_per_rank, SymInt[] peer_offsets, bool apply_topk_in_fc1=True, float? gate_up_clamp=None, float? situ_beta=None, float? situ_linear_beta=None, bool in_kernel_fc2_reduce=False, str combine_format="bf16", bool tactic_autotune=False, SymInt num_tokens=-1) -> () +cute_dsl_nvfp4_dense_gemm_swiglu_moe_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha, Tensor? alpha_post, Tensor norm_const, SymInt expert_count, SymInt weight_per_expert, ScalarType output_dtype, SymInt scaling_vector_size=16) -> (Tensor, Tensor) +cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha, Tensor tile_idx_to_group_idx, Tensor tile_idx_to_mn_limit, Tensor permuted_idx_to_expanded_idx, Tensor num_non_exiting_tiles, Tensor global_sf, SymInt num_experts, SymInt top_k, SymInt num_local_experts, SymInt local_expert_offset, SymInt tile_size, SymInt scaling_vector_size=16, SymInt activation_type=5, float swiglu_limit_scalar=-1.) -> (Tensor, Tensor) +cute_dsl_nvfp4_grouped_gemm_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha, Tensor tile_idx_to_group_idx, Tensor num_non_exiting_tiles, SymInt num_experts, SymInt top_k, SymInt num_local_experts, SymInt local_expert_offset, SymInt tile_size, ScalarType output_dtype, SymInt scaling_vector_size=16) -> Tensor +cute_dsl_nvfp4_grouped_gemm_finalize_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha, Tensor tile_idx_to_group_idx, Tensor tile_idx_to_mn_limit, Tensor permuted_idx_to_expanded_idx, Tensor num_non_exiting_tiles, Tensor token_final_scales, SymInt num_experts, SymInt top_k, SymInt num_local_experts, SymInt local_expert_offset, SymInt tile_size, ScalarType output_dtype, SymInt scaling_vector_size=16) -> Tensor +cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha, Tensor(a5!) output, Tensor tile_idx_to_group_idx, Tensor tile_idx_to_mn_limit, Tensor permuted_idx_to_expanded_idx, Tensor num_non_exiting_tiles, Tensor token_final_scales, SymInt num_experts, SymInt top_k, SymInt num_local_experts, SymInt local_expert_offset, SymInt tile_size, ScalarType output_dtype, SymInt scaling_vector_size=16) -> () +cute_dsl_nvfp4_grouped_gemm_swiglu_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha, Tensor tile_idx_to_group_idx, Tensor num_non_exiting_tiles, Tensor global_sf, SymInt num_experts, SymInt top_k, SymInt num_local_experts, SymInt local_expert_offset, SymInt tile_size, SymInt scaling_vector_size=16, float swiglu_limit_scalar=-1.) -> (Tensor, Tensor) +default_moe_routing_op(Tensor router_logits, SymInt topk, ScalarType? output_dtype=None) -> (Tensor, Tensor) +dsv3_fused_a_gemm_op(Tensor mat_a, Tensor mat_b, Tensor? bias, ScalarType? out_dtype) -> Tensor out +dsv3_router_gemm_op(Tensor mat_a, Tensor mat_b, Tensor? bias, ScalarType? out_dtype) -> Tensor out +e4m3_mxe2m1_block_scale_moe_runner(Tensor? routing_logits, Tensor? routing_bias, Tensor hidden_states, Tensor gemm1_weights, Tensor gemm1_weights_scale, Tensor? gemm1_bias, Tensor? gemm1_alpha, Tensor? gemm1_beta, Tensor? gemm1_clamp_limit, Tensor gemm2_weights, Tensor gemm2_weights_scale, Tensor? gemm2_bias, Tensor output1_scale_scalar, Tensor output1_scale_gate_scalar, Tensor output2_scale_scalar, SymInt num_experts, SymInt top_k, SymInt? n_group, SymInt? topk_group, SymInt intermediate_size, SymInt? valid_hidden_size, SymInt? valid_intermediate_size, SymInt local_expert_offset, SymInt local_num_experts, float? routed_scaling_factor, SymInt routing_method_type, SymInt act_type, Tensor? topk_weights=None, Tensor? topk_ids=None, Tensor? output=None, SymInt tune_max_num_tokens=8192, bool use_dp=False) -> Tensor +fp4_block_scale_moe_runner(Tensor? routing_logits, Tensor? routing_bias, Tensor hidden_states, Tensor hidden_states_scale, Tensor gemm1_weights, Tensor gemm1_weights_scale, Tensor gemm1_bias, Tensor gemm1_alpha, Tensor gemm1_beta, Tensor gemm1_clamp_limit, Tensor gemm2_weights, Tensor gemm2_weights_scale, Tensor gemm2_bias, Tensor output1_scale_scalar, Tensor output1_scale_gate_scalar, Tensor output2_scale_scalar, SymInt num_experts, SymInt top_k, SymInt? n_group, SymInt? topk_group, SymInt intermediate_size, SymInt local_expert_offset, SymInt local_num_experts, float? routed_scaling_factor, SymInt routing_method_type, bool do_finalize, SymInt act_type=0, Tensor? topk_weights=None, Tensor? topk_ids=None, Tensor? output=None, SymInt tune_max_num_tokens=8192, bool use_dp=False) -> Tensor[] +fp8_block_scale_moe_runner(Tensor? routing_logits, Tensor routing_bias, Tensor hidden_states, Tensor hidden_states_scale, Tensor gemm1_weights, Tensor gemm1_weights_scale, Tensor gemm2_weights, Tensor gemm2_weights_scale, SymInt num_experts, SymInt top_k, SymInt? num_fused_shared_experts, SymInt? n_group, SymInt? topk_group, SymInt intermediate_size, SymInt local_expert_offset, SymInt local_num_experts, float? routed_scaling_factor, SymInt routing_method_type, Tensor? topk_weights=None, Tensor? topk_ids=None, SymInt act_type=0, float? gemm1_clamp_limit=None, Tensor? output=None, SymInt tune_max_num_tokens=8192, bool use_dp=False) -> Tensor +fp8_block_scaling_moe_gemm(Tensor mat1, Tensor mat2, Tensor mat1Scale, Tensor mat2Scale, Tensor token_offset) -> Tensor +fp8_fp4_block_scale_moe_runner(Tensor? routing_logits, Tensor? routing_bias, Tensor hidden_states, Tensor gemm1_weights, Tensor gemm1_weights_scale, Tensor gemm2_weights, Tensor gemm2_weights_scale, Tensor output1_scale_scalar, Tensor output1_scale_gate_scalar, Tensor output2_scale_scalar, SymInt num_experts, SymInt top_k, SymInt? n_group, SymInt? topk_group, SymInt intermediate_size, SymInt local_expert_offset, SymInt local_num_experts, float? routed_scaling_factor, SymInt routing_method_type, bool do_finalize, SymInt act_type, Tensor? topk_weights=None, Tensor? topk_ids=None, Tensor? output=None, SymInt tune_max_num_tokens=8192, bool use_dp=False) -> Tensor[] +fp8_per_tensor_scale_moe_runner(Tensor? routing_logits, Tensor? routing_bias, Tensor hidden_states, Tensor gemm1_weights, Tensor output1_scales_scalar, Tensor output1_scales_gate_scalar, Tensor gemm2_weights, Tensor output2_scales_scalar, int num_experts, int top_k, int? n_group, int? topk_group, int intermediate_size, int local_expert_offset, int local_num_experts, float? routed_scaling_factor, bool use_routing_scales_on_input, int tile_tokens_dim, int routing_method_type, Tensor? topk_weights, Tensor? topk_ids) -> Tensor +fused_moe(Tensor input, Tensor token_selected_experts, Tensor token_final_scales, Tensor fc1_expert_weights, Tensor? fc1_expert_biases, Tensor fc2_expert_weights, Tensor? fc2_expert_biases, ScalarType output_dtype, Tensor[] quant_scales, Tensor? input_sf=None, bool swizzled_input_sf=True, Tensor? swiglu_alpha=None, Tensor? swiglu_beta=None, Tensor? swiglu_limit=None, SymInt tp_size=1, SymInt tp_rank=0, SymInt ep_size=1, SymInt ep_rank=0, SymInt cluster_size=1, SymInt cluster_rank=0, bool enable_alltoall=False, bool use_deepseek_fp8_block_scale=False, bool use_w4_group_scaling=False, bool use_int8_woq_per_channel=False, bool use_mxfp8_act_scaling=False, bool min_latency_mode=False, bool use_fused_finalize=True, SymInt tune_max_num_tokens=8192, SymInt? tuner_num_tokens=None, SymInt? tuner_top_k=None, SymInt activation_type=5, SymInt? unpadded_hidden_size=None, Tensor? out_tensor=None, bool use_dynamic_fc2_scale=False, bool use_mxfp8_weight_scaling=False, Tensor? fc1_lora_ranks=None, Tensor? fc1_lora_weight_ptrs=None, Tensor? fc2_lora_ranks=None, Tensor? fc2_lora_weight_ptrs=None, Tensor? gated_lora_ranks=None, Tensor? gated_lora_weight_ptrs=None, Tensor? host_request_types=None, Tensor? host_context_lengths=None, SymInt lora_max_low_rank=0, Tensor? fc1_slot_lora_ranks=None, Tensor? fc1_slot_lora_weight_ptrs=None, Tensor? fc2_slot_lora_ranks=None, Tensor? fc2_slot_lora_weight_ptrs=None, Tensor? gated_slot_lora_ranks=None, Tensor? gated_slot_lora_weight_ptrs=None, Tensor? token_to_slot=None) -> Tensor[] +gate_forward(Tensor scores_in, Tensor bias, Tensor input_ids, Tensor tid2eid, Tensor(a!) out_weights, Tensor(b!) out_indices, int topk, float route_scale, bool is_hash) -> () +get_moe_commworkspace_size_per_rank(int ep_size) -> int +get_moe_prepare_workspace_size_per_rank(int ep_size) -> int +kimi_k3_noaux_tc_mxfp8_quant(Tensor scores, Tensor bias, Tensor hidden_states, float routed_scaling_factor) -> (Tensor, Tensor, Tensor, Tensor) +llama4_moe_tp8ep1_min_latency(Tensor input, Tensor router_logits, Tensor fc1_expert_weights, Tensor fc2_expert_weights, Tensor[]? quant_scales=None) -> Tensor +lora_grouped_gemm(Tensor input, Tensor host_request_types, Tensor[] lora_ranks, Tensor[] lora_weights_pointers, Tensor host_context_lengths, int[] output_hidden_sizes, bool transA, bool transB, int max_low_rank, int weight_index, bool isRemoveInputPadding) -> Tensor[] +lora_grouped_gemm_cuda_graph(Tensor lora_in_sizes, Tensor lora_out_sizes, Tensor a_offsets, Tensor b_ptrs, Tensor d_offsets, Tensor b_prime_ptrs, Tensor d_prime_offsets, int problem_count, Tensor lda, Tensor ldb, Tensor ldd, Tensor ldb_prime, Tensor ldd_prime, Tensor host_max_in_sizes, Tensor host_max_out_sizes, Tensor splitk_offsets, ScalarType dtype, int minKN, int splitKSlices=16) -> () +lora_grouped_gemm_supports_fp8(int _0) -> bool _0 +marlin_nvfp4_moe_gemm(Tensor a, Tensor b_q_weight, Tensor b_scales, Tensor global_scale, Tensor workspace, Tensor sorted_token_ids, Tensor expert_ids, Tensor num_tokens_past_padded, Tensor topk_weights, int moe_block_size, int top_k, bool mul_topk_weights, int size_n, int size_k, ScalarType? out_dtype, bool use_fp32_reduce=False) -> Tensor out +megamoe_prepare(Tensor input, Tensor token_selected_experts, Tensor token_final_scales, Tensor(a!) x_out, Tensor(b!) x_sf_out, Tensor(c!) topk_idx_out, Tensor(d!) topk_weights_out) -> () +memset_expert_ids(Tensor(a!) experts_ids, Tensor recv_rank_count_cumsum, int max_token_count_per_rank, int top_k, int invalid_expert_id, int ep_size) -> () +mnnvl_moe_alltoallv_prepare_without_allgather(Tensor experts_ids, Tensor? experts_statics, Tensor allWorkspace, int max_token_count_per_rank, int ep_rank, int ep_size, int expert_count, int slot_count, int top_k) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor?) +moe_a2a_cft_initialize(Tensor(a!) workspace, int workspace_mem_handle, int workspace_size_per_rank, int ep_rank, int ep_size) -> () +moe_a2a_combine(Tensor(a) payload, int local_num_tokens, Tensor(a!) workspace, Tensor metainfo, int runtime_max_tokens_per_rank, int ep_rank, int ep_size, int top_k, int combine_payload_offset, bool payload_in_workspace, bool use_low_precision=False, bool use_cft_counted_writes=False, bool enable_rank_mask=False, Tensor? active_rank_mask=None) -> Tensor +moe_a2a_dispatch(Tensor token_selected_experts, Tensor[] input_payloads, Tensor(a! -> *) workspace, Tensor metainfo, int runtime_max_tokens_per_rank, int ep_rank, int ep_size, int top_k, int num_experts, Tensor? eplb_local_stats=None, bool use_cft_counted_writes=False, int? expert_id_payload_index=None, int? invalid_token_expert_id=None, bool enable_rank_mask=False, Tensor? active_rank_mask=None) -> (Tensor(a!)[], int, Tensor(a!)) +moe_a2a_get_aux_data_size(int ep_size, int max_num_tokens, int? eplb_stats_num_experts=None, bool can_use_cft_counted_writes=False) -> int +moe_a2a_get_combine_payload_tensor(Tensor(a) workspace, int ep_rank, int ep_size, int runtime_max_tokens_per_rank, int combine_payload_offset, ScalarType out_dtype, int hidden_size) -> Tensor(a) +moe_a2a_initialize(Tensor(a!) workspace, int ep_rank, int ep_size, int max_num_tokens_per_rank, int? eplb_stats_num_experts=None, bool can_use_cft_counted_writes=False) -> Tensor +moe_a2a_sanitize_expert_ids(Tensor(a!) expert_ids, Tensor(a!) workspace, Tensor metainfo, int ep_rank, int invalid_expert_id) -> () +moe_a2a_set_warmup(bool in_warmup) -> () +moe_align_block_size(Tensor topk_ids, int num_experts, int block_size, Tensor(a!) sorted_token_ids, Tensor(a!) expert_ids, Tensor(a!) num_tokens_post_pad) -> () +moe_allreduce(Tensor residual, Tensor norm_weight, Tensor device_num_experts, Tensor scale_input, Tensor active_experts_token_input, Tensor token_input, Tensor workspace, int rank, int nranks, float eps) -> Tensor[] +moe_comm(Tensor[] inputs, Tensor send_rank_cum_sum, Tensor send_indices, Tensor recv_rank_cum_sum, Tensor recv_indices, Tensor all_workspaces, int output_allocation_count, int ep_rank, int ep_size, bool[]? need_zero_output=None, bool? use_low_precision=None) -> Tensor[] +moe_custom_op(str layer_idx, Tensor x, Tensor? x_sf, bool is_swizzled, Tensor router_logits, bool do_finalize, ScalarType? output_dtype, SymInt[]? all_rank_num_tokens, bool? use_dp_padding) -> Tensor[] +moe_finalize_allreduce(Tensor input, Tensor? residual, Tensor norm_weight, Tensor expanded_idx_to_permuted_idx, Tensor? shared_expert_output, Tensor? expert_scale_factor, Tensor workspace, int rank, int nranks, float eps) -> Tensor[] +moe_finalize_scale_op(Tensor gemm2_output, Tensor? biases, Tensor unpermuted_final_scales, Tensor unpermuted_row_to_permuted_row, Tensor permuted_row_to_unpermuted_row, Tensor token_selected_experts, Tensor expert_first_token_offset_tensor, bool enable_alltoall, SymInt num_rows, SymInt hidden_size, SymInt unpadded_hidden_size, int experts_per_token, int num_experts_per_node, int tp_size, int tp_rank, int ep_size, int ep_rank) -> Tensor +moe_gelu(Tensor input, Tensor tile_idx_to_mn_limit, Tensor num_non_exiting_tiles, int tile_tokens_dim) -> Tensor +moe_hierarchical_statistic_local_device(Tensor local_raw_expert_ids, Tensor local_expert_token_count, Tensor enabled, int single_layer_load_balancer_ptr, int is_first_stage, int is_last_stage) -> () +moe_hierarchical_statistic_update(Tensor global_expert_token_count, Tensor enabled, int single_layer_load_balancer_ptr) -> () +moe_initialize_workspace(Tensor(a!) all_workspaces, int ep_rank, int ep_size) -> () +moe_load_balance_routing(Tensor token_selected_experts, bool offset_by_ep_rank, int single_layer_load_balancer_ptr) -> Tensor +moe_load_balance_set_cpu_stage(int single_layer_load_balancer_ptr) -> () +moe_load_balance_statistic(Tensor gathered_raw_expert_ids, Tensor enabled, int single_layer_load_balancer_ptr, int is_first_stage, int is_last_stage) -> () +moe_load_balance_wait_gpu_stage(int single_layer_load_balancer_ptr) -> Tensor +moe_output_memset_inplace(Tensor(a!) input, Tensor tile_idx_to_mn_limit, Tensor expanded_idx_to_permuted_idx, Tensor permuted_idx_to_expanded_idx, Tensor num_non_exiting_tiles, int tile_tokens_dim, int top_k, int ep_size, bool enable_alltoall=False) -> () +moe_permute(Tensor input, Tensor? input_sf, Tensor tile_idx_to_mn_limit, Tensor permuted_idx_to_expanded_idx, Tensor num_non_exiting_tiles, int tile_tokens_dim, int top_k) -> (Tensor, Tensor?) +moe_permute_op(Tensor input, Tensor token_selected_experts, Tensor? token_final_scales, Tensor fc1_expert_weights, Tensor fc2_expert_weights, Tensor[]? quant_scales, Tensor? input_sf, int num_experts_on_rank, int tp_size, int tp_rank, int ep_size, int ep_rank, int cluster_size, int cluster_rank, bool min_latency_mode, bool use_fp8_block_scaling, bool skip_data_expand=False) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor) +moe_sort(Tensor token_selected_experts, Tensor token_final_scales, int num_experts, int top_k, int local_expert_offset, int local_num_experts, int tile_tokens_dim) -> Tensor[] +moe_swiglu(Tensor input, Tensor tile_idx_to_mn_limit, Tensor num_non_exiting_tiles, int tile_tokens_dim) -> Tensor +moe_swiglu_nvfp4_quantize(Tensor input, Tensor global_sf, Tensor tile_idx_to_mn_limit, Tensor num_non_exiting_tiles, int tile_tokens_dim) -> (Tensor, Tensor) +moe_topk_sort(Tensor routing_logits, Tensor? routing_bias, int num_experts, int top_k, int? n_group, int? topk_group, int local_expert_offset, int local_num_experts, float? routed_scaling_factor, int tile_tokens_dim, int routing_method_type) -> Tensor[] +moe_unpermute(Tensor permuted_input, Tensor expanded_idx_to_permuted_idx, Tensor topk_scales) -> Tensor +moe_unpermute_inplace(Tensor permuted_input, Tensor(a!) output, Tensor expanded_idx_to_permuted_idx, Tensor topk_scales) -> () +mxe4m3_mxe2m1_block_scale_moe_runner(Tensor? routing_logits, Tensor? routing_bias, Tensor hidden_states, Tensor hidden_states_scale, Tensor gemm1_weights, Tensor gemm1_weights_scale, Tensor? gemm1_bias, Tensor? gemm1_alpha, Tensor? gemm1_beta, Tensor? gemm1_clamp_limit, Tensor gemm2_weights, Tensor gemm2_weights_scale, Tensor? gemm2_bias, SymInt num_experts, SymInt top_k, SymInt? n_group, SymInt? topk_group, SymInt intermediate_size, SymInt? valid_hidden_size, SymInt? valid_intermediate_size, SymInt local_expert_offset, SymInt local_num_experts, float? routed_scaling_factor, SymInt routing_method_type, SymInt act_type, Tensor? topk_weights=None, Tensor? topk_ids=None, Tensor? output=None, SymInt tune_max_num_tokens=8192, bool use_dp=False) -> Tensor +noaux_tc_op(Tensor scores, Tensor bias, int n_group, int topk_group, int topk, float routed_scaling_factor) -> (Tensor, Tensor) +renorm_moe_routing_op(Tensor router_logits, SymInt topk, ScalarType? output_dtype=None) -> (Tensor, Tensor) +set_moe_max_usable_sm_count(int max_sm_count) -> () +``` + +## norm (21 ops) + +``` +cuda_tile_rms_norm(Tensor x, Tensor weight, float eps, bool static_persistent, bool gather, bool use_gemma) -> Tensor +cuda_tile_rms_norm_fuse_residual_(Tensor(a0!) x, Tensor(a1!) residual, Tensor weight, float eps, bool static_persistent, bool gather, bool use_gemma) -> () +deepseek_v4_q_norm(Tensor q, int num_heads, int head_dim, float eps) -> Tensor +deepseek_v4_q_norm_fused_fp8(Tensor q, Tensor(a!) quant_q_out, Tensor(b!) q_pe_out, int num_heads, int head_dim, int nope_dim, float eps, Tensor? quant_scale_qkv, Tensor? rotary_cos_sin=None, Tensor? cache_seq_lens=None, int seq_len=0, Tensor? cu_q_seqlens=None) -> () +flashinfer_fused_add_rmsnorm(Tensor(a0!) input, Tensor(a1!) residual, Tensor weight, float eps) -> () +flashinfer_fused_add_rmsnorm_quant(Tensor(a0!) out, Tensor input, Tensor(a2!) residual, Tensor weight, Tensor scale, float eps) -> () +flashinfer_gemma_fused_add_rmsnorm(Tensor(a0!) input, Tensor(a1!) residual, Tensor weight, float eps) -> () +flashinfer_gemma_rmsnorm(Tensor input, Tensor weight, float eps) -> Tensor +flashinfer_rmsnorm(Tensor input, Tensor weight, float eps) -> Tensor +fused_adaptive_layernorm(Tensor x, Tensor? ln_weight, Tensor? ln_bias, Tensor? scale_msa, Tensor? shift_msa, int seq_len_per_batch, float eps) -> Tensor +fused_adaptive_layernorm_quant(Tensor x, Tensor? ln_weight, Tensor? ln_bias, Tensor? scale_msa, Tensor? shift_msa, Tensor sf_scale, int seq_len_per_batch, float eps) -> (Tensor, Tensor) +fused_add_rms_norm_quant(Tensor input, Tensor residual, Tensor gamma, Tensor? sf_scale, bool use_rms_norm=True, float eps=9.9999999999999995e-07, bool output_hp_norm=False) -> (Tensor, Tensor, Tensor, Tensor?) +fused_add_rmsnorm_fp4_quantize(Tensor hidden_states, Tensor residual, Tensor norm_weight, Tensor scale_factor, float eps, bool return_norm_out) -> Tensor[] +fused_dit_gate_resid_norm_shift_scale(Tensor x, Tensor? attn=None, Tensor? gate_table=None, Tensor? gate_ts=None, Tensor[]? scale_table=None, Tensor[]? scale_ts=None, Tensor[]? shift_table=None, Tensor[]? shift_ts=None, Tensor[]? sf_scale=None, float eps=9.9999999999999995e-07, int num_out=1) -> Tensor[] +fused_dit_split_norm(Tensor(a!) tensor, int num_heads, int head_dim, float eps, Tensor weight) -> () +fused_gated_rmsnorm_quant(Tensor x, Tensor z, Tensor weight, int group_size, float eps=1.0000000000000001e-05, Tensor? sf_scale=None) -> (Tensor, Tensor) +fused_rmsnorm_fp4_quantize(Tensor hidden_states, Tensor norm_weight, Tensor scale_factor, float eps, bool return_norm_out) -> Tensor[] +group_rms_norm_base(Tensor[] inputs, Tensor(a1!)[] outputs, Tensor[] weights, float eps, float weight_bias) -> () +group_rms_norm_heuristic(Tensor[] inputs, Tensor(a1!)[] outputs, Tensor[] weights, float eps, float weight_bias) -> () +group_rms_norm_large_batch(Tensor[] inputs, Tensor(a1!)[] outputs, Tensor[] weights, float eps, float weight_bias) -> () +rms_norm_gated_token_major(Tensor x, Tensor z, Tensor weight, float eps, Tensor? fp8_scale=None) -> Tensor +``` + +## activation (11 ops) + +``` +cute_dsl_nvfp4_dense_gemm_gelu_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha, ScalarType output_dtype, Tensor? bias=None, bool use_tvm_ffi=True) -> Tensor +cute_dsl_nvfp4_dense_gemm_gelu_fp4out_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha, Tensor global_sf, Tensor? bias=None, bool use_tvm_ffi=True) -> (Tensor, Tensor) +cute_dsl_nvfp4_dense_gemm_swiglu_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha, ScalarType output_dtype, bool use_tvm_ffi=True) -> Tensor +cute_dsl_nvfp4_dense_gemm_swiglu_fp4out_blackwell(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha, Tensor global_sf, bool use_tvm_ffi=True) -> (Tensor, Tensor) +flashinfer_gelu_tanh_and_mul(Tensor x) -> Tensor +flashinfer_silu_and_mul(Tensor x) -> Tensor +llama4_fp8_fp8_gemm_swiglu(Tensor inputA, Tensor inputB, Tensor in_scale, Tensor out_scale_inv) -> Tensor +silu_and_mul(Tensor x, Tensor? scale=None, ScalarType? dtype=None, float? swiglu_limit=None, float? swiglu_alpha=None, float? swiglu_beta=None) -> Tensor +tensorrt_llm::dequantize_e4m3_activation(Tensor activation, Tensor scales) -> Tensor +tensorrt_llm::quantize_e4m3_activation(Tensor activation) -> (Tensor, Tensor) +tensorrt_llm::static_quantize_e4m3_activation(Tensor activation, Tensor scales) -> (Tensor, Tensor) +``` + +## quantization (34 ops) + +``` +block_scale_interleave(Tensor input) -> Tensor +block_scale_interleave_reverse(Tensor input) -> Tensor +calculate_nvfp4_global_scale(Tensor input, Tensor? tokensPerBatch) -> Tensor +fp4_batched_quantize(Tensor input, Tensor globalScale, int sfVecSize, bool sfUseUE8M0=False) -> (Tensor, Tensor) +fp4_quantize(Tensor input, Tensor? globalScale, int sfVecSize, bool sfUseUE8M0=False, bool isSfSwizzledLayout=True) -> (Tensor, Tensor) +fp4_quantize_with_reorder_residual(Tensor X, Tensor input_scale, Tensor reorder_index, int KE, bool is_act) -> (Tensor, Tensor) +fp4_quantize_with_residual(Tensor X, Tensor input_scale, int KE, bool is_act) -> (Tensor, Tensor) +fp8_batched_quantize_1x128_permute102(Tensor input) -> (Tensor, Tensor) +fp8_quantize_1x128(Tensor input, bool use_ue8m0=False) -> (Tensor, Tensor) +fp8_quantize_1x128_cutedsl_ue8m0(Tensor input) -> (Tensor, Tensor) +fp8_quantize_1x128_packed_ue8m0(Tensor input) -> (Tensor, Tensor) +fused_cat_fp4(Tensor pe, Tensor nope) -> (Tensor, Tensor) +fused_cat_fp8(Tensor pe, Tensor nope, bool use_ue8m0=False) -> (Tensor, Tensor) +fused_relu2_quantize(Tensor input, Tensor sf_scale, int sf_vec_size=16) -> (Tensor, Tensor) +mxfp4_dequantize_unswizzled(Tensor _0, Tensor _1, int _2) -> Tensor _0 +mxfp8_quantize(Tensor input, bool swizzedLayout=True, int alignment=32) -> (Tensor, Tensor) +pack_int8_tensor_to_packed_int4(Tensor _0) -> Tensor _0 +quantize_e4m3_per_tensor(Tensor input) -> (Tensor, Tensor) +reswizzle_sf(Tensor sf, SymInt rows, SymInt cols, SymInt scaling_vector_size=16) -> Tensor +shuffle_matrix(Tensor _0, Tensor _1) -> Tensor _0 +symmetric_quantize_last_axis_of_batched_matrix(Tensor _0, ScalarType _1) -> Tensor[] _0 +tensorrt_llm::dequantize_e4m3_per_tensor(Tensor input, Tensor scales) -> Tensor +tensorrt_llm::dequantize_e4m3_weight(Tensor weight, Tensor scales) -> Tensor +tensorrt_llm::dequantize_mxe4m3_host(Tensor _0, Tensor _1, bool _2) -> Tensor _0 +tensorrt_llm::e2m1_and_ufp8sf_scale_to_float(Tensor _0, Tensor _1, int _2, int _3) -> Tensor _0 +tensorrt_llm::e2m1_and_ufp8sf_scale_to_float_v2(Tensor _0, Tensor _1, Tensor? _2, int _3, int _4, bool _5) -> Tensor _0 +tensorrt_llm::float_to_e2m1_and_ufp8sf_scale(Tensor _0, int _1, int _2, bool? _3) -> Tensor[] _0 +tensorrt_llm::half_to_e2m1_and_ufp8sf_scale(Tensor _0, Tensor _1, int _2, int _3) -> Tensor[] _0 +tensorrt_llm::quantize_e4m3_weight(Tensor weight) -> (Tensor, Tensor) +tensorrt_llm::quantize_mxe4m3_host(Tensor _0, bool _1) -> (Tensor _0, Tensor _1) +tensorrt_llm::static_quantize_e4m3_per_tensor(Tensor input, Tensor scales) -> (Tensor, Tensor) +tensorrt_llm::static_quantize_e4m3_weight(Tensor weight, Tensor scales) -> (Tensor, Tensor) +tunable_fp4_quantize(Tensor input, Tensor input_scale, SymInt scaling_vector_size=16, bool sf_use_ue8m0=False, bool is_sf_swizzled_layout=True) -> Tensor[] +unpack_int4_packed_tensor_to_int8(Tensor _0) -> Tensor _0 +``` + +## comm (35 ops) + +``` +add_to_ub(Tensor a, Tensor b) -> Tensor +allgather(Tensor input, SymInt[]? sizes, int[] group) -> Tensor +allgather_list(Tensor[] input_list, SymInt[]? sizes, int[] group) -> Tensor[] +allgather_list_pg(Tensor[] input_list, SymInt[]? sizes, int[] group, __torch__.torch.classes.c10d.ProcessGroup process_group) -> Tensor[] +allgather_pg(Tensor input, SymInt[]? sizes, int[] group, __torch__.torch.classes.c10d.ProcessGroup process_group) -> Tensor +allreduce(Tensor input, Tensor? residual, Tensor? norm_weight, Tensor? scale, Tensor? bias, Tensor? workspace, int[] group, int strategy, int op, float eps, bool trigger_completion_at_end) -> Tensor[] +allreduce_pg(Tensor input, Tensor? residual, Tensor? norm_weight, Tensor? scale, Tensor? bias, Tensor? workspace, int[] group, int rank, __torch__.torch.classes.c10d.ProcessGroup pg, int strategy, int op, float eps, bool trigger_completion_at_end) -> Tensor[] +alltoall_helix(Tensor[] input_list, int[] group, int? num_lists) -> Tensor[] +alltoall_helix_native(Tensor partial_o, Tensor softmax_stats, Tensor(a!) workspace, int cp_rank, int cp_size) -> (Tensor, Tensor) +autotuned_allreduce(Tensor input, Tensor? residual, Tensor? norm_weight, Tensor? scale, Tensor? bias, Tensor? workspace, int[] group, int strategy, int op, float eps, bool trigger_completion_at_end) -> Tensor[] +copy_to_userbuffers(Tensor a) -> Tensor +create_userbuffers_tensor(int[] _0, ScalarType _1) -> Tensor _0 +helix_post_process(Tensor gathered_o, Tensor gathered_stats, float scale) -> Tensor +helix_post_process_native(Tensor gathered_o, Tensor gathered_stats, float scale, int cp_dim) -> Tensor +initialize_helix_workspace(Tensor(a!) workspace, int cp_rank, int cp_size) -> () +is_nccl_window_buffer(Tensor input, int[] group) -> bool +minimax_allreduce_rms(Tensor input, Tensor norm_weight, Tensor workspace, int rank, int nranks, float eps, bool trigger_completion_at_end) -> Tensor +minimax_allreduce_rms_qk(Tensor q, Tensor k, Tensor norm_weight_q, Tensor norm_weight_k, Tensor workspace, int rank, int nranks, float eps, bool trigger_completion_at_end) -> Tensor[] +mnnvl_fusion_allreduce(Tensor input, Tensor? gamma, Tensor? residual, float? epsilon, Tensor(a!) comm_buffer, Tensor buffer_flags, bool rmsnorm_fusion, Tensor? scale=None, int fusion_op=0) -> Tensor[] +pp_recv_tensors(Tensor(a0!)[] tensors) -> () +pp_send_tensors(Tensor(a0!)[] tensors) -> () +preallocate_nccl_window_buffer(Tensor input, int[] group, int count) -> () +reducescatter(Tensor input, SymInt[]? sizes, int[] group) -> Tensor +reducescatter_list(Tensor[] input_list, SymInt[]? sizes, int[] group) -> Tensor[] +reducescatter_list_pg(Tensor[] input_list, SymInt[]? sizes, int[] group, __torch__.torch.classes.c10d.ProcessGroup process_group) -> Tensor[] +reducescatter_pg(Tensor input, SymInt[]? sizes, int[] group, __torch__.torch.classes.c10d.ProcessGroup process_group) -> Tensor +register_allreduce_tactic(Tensor input, Tensor? residual, Tensor? norm_weight, Tensor? scale, Tensor? bias, Tensor? workspace, int[] group, int op, int bucket, int tactic) -> () +tunable_allreduce(Tensor input, Tensor? residual, Tensor? norm_weight, Tensor? scale, Tensor? bias, Tensor? workspace, SymInt[] group, SymInt strategy, SymInt op, float eps, bool trigger_completion_at_end) -> Tensor[] +ulysses_a2a_async_barrier(__torch__.torch.classes.c10d.ProcessGroup pg, int timeout_ms=60000) -> () +ulysses_a2a_async_prepare(Tensor input, __torch__.torch.classes.c10d.ProcessGroup pg) -> (Tensor, __torch__.torch.classes.trtllm.SendHandle) +ulysses_a2a_async_push(__torch__.torch.classes.trtllm.SendHandle send_h, __torch__.torch.classes.c10d.ProcessGroup pg) -> () +ulysses_permute_scatter(Tensor(a!) input, Tensor(b!) send_buf, Tensor(c!) recv_buf, int my_rank, int P) -> () +ulysses_post_unscatter_qkv(Tensor q_in, Tensor k_in, Tensor v_in, int layout=0) -> (Tensor, Tensor, Tensor) +userbuffers_allreduce_finalize(Tensor input, bool force_applying_finalize) -> Tensor +validate_allreduce_tuning_buckets(int[] buckets) -> () +``` + +## embedding (1 ops) + +``` +embedding_skip_forward(Tensor input, SymInt embedding_dim, ScalarType dtype) -> Tensor +``` + +## kv_cache (7 ops) + +``` +clear_allreduce_tactic_cache() -> () +compressor_paged_kv_compress(Tensor kv_score, Tensor ape, Tensor(a!) paged_kv, Tensor(b!) paged_score, Tensor block_table_kv, Tensor block_table_score, Tensor(c!) output, Tensor kv_lens, Tensor cu_seq_lens, Tensor cu_kv_comp, int batch_size, int page_size, int head_dim, int compress_ratio, int next_n) -> () +compressor_postprocess_scatter(Tensor kv_comp, Tensor(a!)? kv_out, Tensor rms_weight, float rms_eps, Tensor cos_sin_table, Tensor position_ids, int nope_dim, int rope_dim, Tensor(b!) kv_cache, Tensor num_outputs, Tensor cu_kv_comp, Tensor start_pos, Tensor block_offsets, Tensor compressed_mask, int tokens_per_block, int cache_scale_type, bool rotate_activation, Tensor(c!)? quant_output, Tensor(d!)? scale_output) -> () +compressor_prefill_reduction(Tensor kv_score, Tensor ape, Tensor(a!) paged_kv, Tensor(b!) paged_score, Tensor block_table_kv, Tensor block_table_score, Tensor(c!) output, Tensor kv_lens, Tensor start_pos, Tensor cu_seq_lens, Tensor cu_kv_comp, int batch_size, int page_size, int head_dim, int compress_ratio, int max_outputs) -> () +indexer_k_cache_gather_op(Tensor k_cache, Tensor slot_mapping_fp8, Tensor slot_mapping_scale, int k_token_start, int num_tokens, int head_dim) -> (Tensor, Tensor) +indexer_k_cache_scatter_op(Tensor k_fp8, Tensor k_scale, Tensor(a!) k_cache, Tensor slot_mapping_fp8, Tensor slot_mapping_scale, int num_tokens) -> () +sparse_kv_cache_compact_layers(Tensor(a!)[] pools, Tensor pool_pointers, Tensor page_table, Tensor source_indices, Tensor source_offsets, Tensor destination_bases, Tensor? source_layer_indices=None) -> () +``` + +## sampling (13 ops) + +``` +cute_dsl_fp4_paged_mqa_logits(Tensor q, Tensor sf_q, Tensor kv_fused, Tensor weights, Tensor context_lens, Tensor block_table, Tensor schedule_meta, SymInt max_context_len, SymInt num_epi_subtiles=1, ScalarType epi_dtype=6, ScalarType output_dtype=6, bool remove_online_sf_transpose=False) -> Tensor +cute_dsl_fp8_paged_mqa_logits(Tensor q, Tensor kv_fused, Tensor weights, Tensor context_lens, Tensor block_table, Tensor schedule_meta, SymInt max_context_len, SymInt num_epi_subtiles=1, ScalarType epi_dtype=6, ScalarType acc_dtype=6, ScalarType output_dtype=6) -> Tensor +cute_dsl_gvr_topk_decode(Tensor logits, Tensor pre_idx, Tensor seq_lens, Tensor(a3!) output_indices, SymInt top_k, SymInt next_n=1, SymInt compress_ratio=1, SymInt? max_seq_len=None, SymInt? cluster_size=None, Tensor? order_row=None, Tensor? counters=None, SymInt? max_batch_size=None) -> () +cute_dsl_gvr_topk_lb_prepare(Tensor seq_lens, Tensor(a1!) order_row, Tensor(a2!) counters, SymInt max_batch_size, SymInt long_threshold=65536, SymInt compress_ratio=1) -> () +cute_dsl_indexer_topk_decode(Tensor input_values, Tensor seq_lens, Tensor(a2!) output_indices, SymInt top_k, SymInt next_n=1, SymInt num_copy_bits=256, bool dynamic=True, bool single_pass_multi_cta=False, bool single_pass_multi_cta_cluster=False, str overflow_policy="REREAD", bool cache_smem_values=False, bool radix_filter_single_pass_multi_cta=True) -> () +cute_dsl_indexer_topk_prefill_blackwell(Tensor input_values, Tensor row_starts, Tensor row_ends, Tensor(a3!) output_indices, SymInt top_k, SymInt num_copy_bits=256, str overflow_policy="REREAD", bool cache_smem_values=False) -> () +cute_dsl_topk_decode_blackwell(Tensor input_values, Tensor seq_lens, SymInt top_k, SymInt next_n=1, SymInt num_copy_bits=256) -> Tensor +cute_dsl_topk_decode_multi_cta_blackwell(Tensor input_values, Tensor seq_lens, SymInt top_k, SymInt next_n=1, SymInt num_copy_bits=256, SymInt chunk_size_per_cta=16384, bool dynamic=True) -> Tensor +fused_topk_softmax(Tensor router_logits, int top_k, int num_experts_total, int start_expert, int end_expert) -> (Tensor, Tensor) +indexer_topk_decode(Tensor logits, Tensor seq_lens, Tensor indices, int next_n, int index_topk=2048, Tensor? pre_idx=None, Tensor? heuristic_scratch=None, int compress_ratio=1, Tensor? radix_aux_indices=None, Tensor? radix_aux_logits=None) -> () +indexer_topk_prefill(Tensor logits, Tensor row_starts, Tensor row_ends, Tensor indices, int index_topk=2048) -> () +logits_bitmask(Tensor(a!) logits, Tensor bitmask, Tensor? token_mask=None, Tensor? d2t=None) -> () +tensorrt_llm::gather_tree(Tensor _0, Tensor _1, Tensor _2, Tensor _3, Tensor _4, Tensor? _5, Tensor? _6, Tensor? _7, Tensor? _8, Tensor? _9, Tensor? _10, Tensor? _11, Tensor? _12, Tensor? _13, Tensor? _14, Tensor? _15, Tensor? _16, Tensor _17, int _18, int _19, int _20, bool _21) -> Tensor _0 +``` + +## speculative (14 ops) + +``` +build_dynamic_tree_op(Tensor parentList, Tensor selectedIndex, Tensor(a!) treeMask, Tensor(b!) positions, Tensor(c!) retrieveIndex, Tensor(d!) retrieveNextToken, Tensor(e!) retrieveNextSibling, int topK, int depth, int numDraftTokens, int treeMaskMode, int numInt32PerRow) -> () +extract_real_draft_tokens_op(Tensor newDraftTokens, Tensor draftTokensBuffer, Tensor tokensGatherIdxForDrafterModel, Tensor topKList, Tensor draftTokensIndicesCumsum, int curDraftIdx, int batchSize, int maxDraftLen, int maxTotalDraftTokens, int maxTopK) -> () +kda_mtp_decode(Tensor x_q, Tensor x_k, Tensor x_v, Tensor w_q, Tensor w_k, Tensor w_v, Tensor(a6!) cs_q, Tensor(a7!) cs_k, Tensor(a8!) cs_v, Tensor g, Tensor beta, Tensor A_log, Tensor dt_bias, Tensor(a13!) recurrent_state, Tensor(a14!) qkg_cache, Tensor(a15!) v_cache, Tensor(a16!) beta_cache, Tensor ssm_state_indices, Tensor cu_seqlens, SymInt num_spec, Tensor num_accepted_tokens, float lower_bound, float? scale=None, bool zero_accepted_hint=False, bool regular_metadata_hint=False) -> Tensor +mamba2_mtp_ssm_cache_update(Tensor ssm, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, Tensor out, Tensor intermediate_states, Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, Tensor? ssm_batch_indices, Tensor? intermediate_states_indices, Tensor? retrieve_parent_token, int cache_steps, int pad_slot_id, bool disable_state_update) -> () +mtp_prepare_drafter_inputs_op(Tensor inputIds, Tensor seqLens, Tensor mtpPastHiddenStatesPtrs, Tensor mtpPastTokensPtrs, Tensor hiddenStates, Tensor acceptedTokens, Tensor numAcceptedTokens, Tensor returnInputIds, Tensor returnHiddenStates, int numMTPModules, int batchSize, int numContextRequest, int hiddenSize) -> (Tensor, Tensor) +mtp_relaxed_acceptance_op(Tensor reqSlotIds, Tensor topKValue, Tensor topKIndices, Tensor draftTokens, Tensor mtpRelaxedDelta, Tensor numAcceptedTokens, Tensor acceptedTokens, int numMTPModules, int batchSize, int numContextRequest, int relaxedTopK, float relaxedDelta, int beginThinkingTokens, int endThinkingTokens) -> (Tensor, Tensor) +mtp_sampling_and_accepted_draft_tokens_op(Tensor logits, Tensor draftTokens, Tensor targetTokens, int numMTPModules, int batchSize, int numContextRequest, int vocabSize) -> (Tensor, Tensor) +mtp_update_hidden_states_op(Tensor inputIds, Tensor seqLens, Tensor targetModelHiddenStates, Tensor mtpPastHiddenStatesPtrs, Tensor mtpPastTokensPtrs, Tensor numAcceptedTokens, int numMTPModules, int batchSize, int numContextRequest, int hiddenSize) -> (Tensor, Tensor) +tensorrt_llm::convert_spec_decoding_mask_to_packed_mask(Tensor _0, Tensor _1, int _2, Tensor _3, int? _4) -> () +tensorrt_llm::redrafter_prepare_random_tensors(Tensor _0, Tensor _1, Tensor _2, Tensor? _3, int _4, int _5, int _6, bool _7) -> () +tensorrt_llm::update_kv_cache_draft_token_location(Tensor _0, Tensor _1, Tensor _2, bool _3, int _4, int _5, int _6, int _7, int _8, Tensor? _9, Tensor[]? _10, Tensor? _11, Tensor? _12, int? _13, int? _14, int? _15) -> () +tensorrt_llm::update_kv_cache_draft_token_location_2d(Tensor _0, Tensor _1, Tensor _2, bool _3, int _4, int _5, int _6, int _7, int _8, Tensor? _9, Tensor? _10, int? _11, int? _12, int? _13) -> () +verify_dynamic_tree_greedy_out_packed_op(Tensor candidates, Tensor retrievePacked, Tensor targetPredict, Tensor(a!) acceptIndex, Tensor(b!) acceptTokenNum, Tensor(c!) acceptToken, Tensor treeValid, int numSpecStep) -> () +verify_dynamic_tree_rejection_out_op(Tensor draftTokens, Tensor targetProbs, Tensor retrieveNextToken, Tensor retrieveNextSibling, Tensor treeValid, Tensor(a!) acceptIndex, Tensor(b!) acceptTokenNum, Tensor(c!) acceptToken, int numSpecStep, Tensor seed, Tensor offset) -> () +``` + +## ssm (5 ops) + +``` +causal_conv1d_fwd(Tensor($0! -> ) x, Tensor($1! -> ) weight, Tensor? bias_, Tensor($2! -> )? conv_states, Tensor? query_start_loc, Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, int pad_slot_id) -> () +causal_conv1d_update(Tensor($0! -> ) x, Tensor($1! -> ) conv_state, Tensor($2! -> ) weight, Tensor? bias_, bool silu_activation, Tensor? cache_seqlens_, Tensor? conv_state_indices, int pad_slot_id) -> () +gdn_custom_op_inplace(Tensor mixed_qkv, Tensor a, Tensor b, str layer_idx, Tensor(a4!) output) -> () +mamba_conv1d(Tensor input, Tensor conv_weight, Tensor conv_bias, Tensor conv_state, Tensor host_request_types, Tensor last_token_ids, Tensor? host_context_lengths, Tensor? slot_mapping, int dim, int dconv, int pre_stride, int post_stride, bool remove_padding, bool apply_silu, bool is_paged_state) -> (Tensor, Tensor) +selective_scan(Tensor input, Tensor state, Tensor delta, Tensor delta_bias, Tensor A, Tensor BC, Tensor D, Tensor host_request_types, Tensor last_token_ids, Tensor? z, Tensor? host_context_lengths, Tensor? slot_mapping, int dim, int dstate, int nheads, int ngroups, int chunk_size, int delta_rank, bool delta_softplus, bool remove_padding, bool is_mamba2, bool is_paged_state) -> (Tensor, Tensor) +``` + +## internal (9 ops) + +``` +allocate_output(Tensor like, int output_buffer_kind, int[]? group, int[]? shape=None, ScalarType? out_dtype=None) -> (Tensor, int) +convert_req_index_to_global(Tensor req_id, Tensor block_table, Tensor token_indices, int block_size, int num_topk_tokens, int stride_factor, int layer_id) -> Tensor +initialize_static_lowprecision_buffers(Tensor workspace, int tp_size) -> Tensor[] +inplace_slice_copy(Tensor(a!) dest, Tensor src, int dim1_start, int dim1_end) -> () +migrate_to_host_accessible(Tensor tensor) -> () +record_event(SymInt event_idx) -> () +record_stream(Tensor tensor, SymInt stream_id) -> () +set_stream(SymInt stream_id) -> () +wait_event(SymInt event_idx) -> () +``` + +## misc (5 ops) + +``` +mhc_big_fuse(Tensor y_acc, Tensor r_acc, Tensor residual, Tensor hc_scale, Tensor hc_base, Tensor(a!) post_mix, Tensor(b!) comb_mix, Tensor(c!) layer_input, int M, int K, int hidden_size, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value, int sinkhorn_repeat, int num_splits, int block_size=0) -> () +mhc_fused_hc(Tensor x_prev, Tensor residual_prev, Tensor post_mix_prev, Tensor comb_mix_prev, Tensor w, Tensor hc_scale, Tensor hc_base, Tensor(a!) residual_cur, Tensor(b!) post_mix_cur, Tensor(c!) comb_mix_cur, Tensor(d!) layer_input_cur, Tensor(e!) y_acc_workspace, Tensor(f!) r_acc_workspace, Tensor(g!) done_counter_workspace, int M, int hidden_size, int hc_mult, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value, int sinkhorn_repeat, int backend=0, int tile_n=0, int num_k_splits=0, int bigfuse_block_size=0, int tile_m=1, Tensor? norm_weight=None, float norm_eps=0.) -> () +mhc_fused_hc_mma_enabled() -> bool +mhc_hc_head_apply(Tensor mixes, Tensor sqrsum, Tensor x, Tensor(a!) out, Tensor scale, Tensor base_t, int M, int mult, int hidden_size, int K, float norm_eps, float eps) -> () +mhc_post_mapping(Tensor residual, Tensor x, Tensor post_mix, Tensor comb_mix, Tensor(a!) out, int B, int hidden_size) -> () +``` diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/SKILL.md new file mode 100644 index 000000000000..cd418665a311 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/SKILL.md @@ -0,0 +1,125 @@ +--- +name: staircase-kernel-onboard +description: Load when smith must add a catalog entry for a TensorRT-LLM callable that has none — read the source, fix the single-call boundary, write contract + wrapper + GPU test, record the receipt, sync the index. Also covers the macro verdict, where the run produces a decomposition report and dispatches one child per constituent instead. +--- + +# staircase-kernel-onboard + +One run, one entry. Three possible outcomes, all of them valid products: + +| outcome | when | +|---|---| +| a new entry, contracted, tested, receipt recorded, index synced | the callable is one kernel-launching call | +| an entry with `status: failed` plus a short note | wrapper and test are right and the op demonstrably fails here. **Land it anyway** — that is knowledge | +| a decomposition report, no entry | the callable turns out to be a macro | + +References: `references/entry-spec.md` is the shared spec (atomicity, the +three pieces, receipts, index). `references/contract-evidence.md` is what +makes a contract sentence defensible. `references/decomposition.md` is the +macro path. Read the first two before writing anything. + +An entry **spans two trees**: the contract and the wrapper live in +`tensorrt_llm/_torch/_experimental/modeling_v2/catalog//`, the GPU +test lives at +`tests/unittest/_torch/modeling_v2//test_modeling_v2_.py`. +Every rule below that says "every file of the entry" means both trees. + +## Flow + +1. **Locate the callable.** `.claude/skills/staircase-kernel-feasibility/references/trtllm-ops.md` (the generated inventory, + with schemas) or, under the run's launch prefix and declared interpreter, + `agent_flow/workflows/staircase_onboard/tools/list_trtllm_ops.py -f `. + Every listed op is callable as `torch.ops.trtllm.`. Then read the + implementation and its Python call path until you can answer: what does one + call compute, what state or metadata must exist first, which dtypes and + arches does it route. +2. **Fix the boundary and the category.** One entry = one op invocation at the + Python level; fusion is legal only *inside* that call. Categories: + `gemm attention moe norm activation quantization comm ssm`, chosen by + compute role rather than dtype. `torch/` is reserved for hand-added + mirrors — **never write there**, including when your verdict is that a + torch callable is the answer: report that verdict and let the orchestrator + decide, since a computation mirror widens what the surface is for. No + category fits → stop and report. + **A macro verdict here ends this path**: go to `references/decomposition.md`. +3. **Write the three pieces**, to the spec. The test file is named + `test_modeling_v2_.py`, always — that name is what keeps it from + colliding with upstream's test of the same op, and what the CI selector + matches. +4. **Run the test on GPU**, under the run's launch prefix and declared + interpreter: + ` -m pytest tests/unittest/_torch/modeling_v2//test_modeling_v2_.py`. + One command per entry, multi-rank ops included — a collective's test spawns + its own ranks rather than asking the caller for a launcher. Iterate until + it passes, or until you have established a genuine op failure. A fake pass + is not a product. +5. **Record the receipt** from that observed run only. +6. **Sync + `tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml`.** +7. **Do not format or lint during iteration.** The repo's `pre-commit` runs + once, at delivery, and covers formatting, import order and the general lint + rules in one pass. What the retired per-entry recipe taught is why that pass + is not narrowed: an import-only check (`ruff check --select I`) restricts + ruff to import rules, so an unused variable passes it, passes the formatter, + and passes a type check — an entry once shipped with a dead `F841` in its + test because the recipe stopped at the first two. If you do run a checker by + hand, run it **by name on your own files**, so a sweep does not reformat + entries you did not touch. +8. **Revalidate the receipt.** Steps 5–7 can all write to the entry after the + test last ran — a fixture added, a helper renamed, a formatter rewriting + the test file. A receipt is valid only if it **post-dates the last write to + every file in the entry, in both trees**. Check it mechanically against the + files' mtimes, not by recollection: a receipt that predated its own test + file has shipped here before, and it read as green. +9. **Commit**, pathspec-scoped to exactly your files: + +```bash +git add tensorrt_llm/_torch/_experimental/modeling_v2/catalog//.md \ + tensorrt_llm/_torch/_experimental/modeling_v2/catalog//.py \ + tests/unittest/_torch/modeling_v2//test_modeling_v2_.py \ + tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml +git commit -m "feat(catalog): onboard — " +``` + +10. **Report.** + +## The macro path + +On a macro verdict, the run's product is the report and the subtree: + +1. write `references/decompositions/.md` to the checklist in + `references/decomposition.md` +2. **commit it before dispatching anything** — your context does not survive a + session cut, and the report plus `git log` is what lets a fresh smith + continue +3. write each child's prompt into the journal, then dispatch one + `Agent(smith)` per constituent in dependency order, each loading this skill +4. report the subtree: which constituents landed, with which receipts, and + what remains + +A child does not dispatch. One that finds its own callable is a macro writes a +report and returns it. + + + +## Report + +- entry path; the wrapped op's qualified name; the receipt as recorded; what + the test covered (shapes, dtypes); files touched — in **both** trees +- **the decomposition map**, whenever the request turned out composite: every + constituent in call order, each tagged `existing entry: ` / + `not onboarded: ` / `gap`. Every symbol you name must be verified + resolvable in the installed engine — never cite one you have not resolved. + This map is why the composite can stay a black box afterwards: you opened it + once so nobody has to again +- symptoms: out-of-scope problems observed (trtllm bugs, env issues) +- every claim about the artifacts must match disk. Re-check before sending diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/contract-evidence.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/contract-evidence.md new file mode 100644 index 000000000000..a98246c221d1 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/contract-evidence.md @@ -0,0 +1,81 @@ +# Contract evidence — what makes a sentence in a contract defensible + +The contract is the entry's sole interface document: whoever calls the op +reads it and nothing else. So a sentence that is *almost* true is worse than a +missing one — the caller acts on it. + +Three failure modes have each been found repeatedly in this catalog. All are about +evidence, not about prose quality. + +**Write the path the receipt certifies, not the path the source describes.** +The one failure mode that survives review is a contract sentence that is +true of *some* execution path, just not the one the test drove. Four found +in this catalog so far: norm Preconditions describing the CUDA-JIT path while +the receipt certified CuTe DSL; `thop_attention`'s fp8 `kv_scale_*` rule +transcribed from an upstream comment about XQA kernels, on an arch whose +fp8 decode runs trtllm-gen JIT; `fused_moe`'s "tactic choice does not +change results", true of the cold autotuner the test runs in and false of +the warm one serving runs in; and MLA's "mixed batches are rejected", +which is a check in the Python backend that this op bypasses. Before +writing a Precondition, name the path you measured — backend, kernel +family, cache state — and say so if it is not the only one. + +**A claim that the op rejects something is the expensive kind to get +wrong**, because it promises a guard the caller then does not write. Every +"raises" / "is rejected" / "not survivable" sentence needs a probe that +actually saw the raise, in the certified configuration, with the error +text quoted. Absent that, write what you measured: accepted, and what it +then does. + +**Describe what the engine does and name where to look — never transcribe +what its code says.** A contract records things learned about TensorRT-LLM, +and how you write them down decides whether they are still true next month. +The engine moves every day; a contract does not. So anything you copy *out* +of engine source — a line number, a pasted excerpt, an argument list +transcribed by hand — is wrong shortly after you write it, silently, and the +next reader believes it because it reads like a citation. + +This is measured, not a worry. Seven line-number citations were written into +this catalog and **all seven pointed at the wrong line within days** — one +set off by about 35 lines, another by about 1100, landing on an unrelated +f-string. In the very same sentences, every file path and every symbol name +still resolved exactly. The durable pointer and the rotting one sat side by +side, and only one of them rotted. + +So write two things and stop: + +- **the behaviour**, in your own words, at the level the caller needs — what + it does, when, and what it means for the call you are making; +- **the pointer**: a file path plus a *name* the reader can grep — a + function, a class, a field, a config key, or the exact assert text. + +| instead of | write | +|---|---| +| `` `_torch/pyexecutor/model_engine.py:1103,1278` `` | warmup calls `_run_autotuner_warmup` (`_torch/pyexecutor/model_engine.py`) | +| `` `llmapi/llm_args.py:4669` `` | `enable_autotuner`, a `TorchLlmArgs` field, defaults to `True` | +| `` `attentionOp.cpp:1045` `` | the `seqQOffset is nullptr.` guard in `common/attentionOp.cpp` | +| a pasted 20-line excerpt of the op's C++ | one sentence on what it does, and the file to read | + +The one thing that keeps its line number is a **quoted diagnostic**: +`RuntimeError: ... (../tensorrt_llm/runtime/torchUtils.h:113)` is the message +the engine prints, so it is what someone greps for when they hit it. Quote it +whole and verbatim. That is a symptom, not a direction to navigate. + +Both halves are enforced rather than reviewed — +`tests/unittest/_torch/modeling_v2/test_modeling_v2_no_stale_claims.py` fails +on a navigation citation and on a named file that no longer exists. + +## Before writing any Precondition, answer three questions + +1. **Which execution path did the receipt actually drive?** Backend, kernel + family, autotuner cache state. If the sentence is true of another path, say + which. +2. **Did I see it, or did I read it?** An upstream comment is a hypothesis + about the code, not a measurement of this arch. +3. **If this says the op refuses something, where is the probe that saw the + refusal?** In the certified configuration, with the error text quoted. No + probe, no claim — write what you measured instead: accepted, and what it + then does. +4. **Will this sentence still be true after someone refactors the engine?** + If it depends on a line number, a pasted excerpt, or an argument order you + copied, it will not. Keep the behaviour and the name; drop the rest. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decomposition.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decomposition.md new file mode 100644 index 000000000000..fde59f29aabe --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decomposition.md @@ -0,0 +1,64 @@ +# Decomposition — when an op is a macro, and what the report must contain + +A requested callable is a **macro** when it cannot be wrapped as one +kernel-launching call: its body is Python-level composition — a registry +lookup delegating to a module `forward`, module calls, tensor glue beyond the +single call, a phase dispatch, or nested calls into other registered ops. + +A macro verdict ends the onboarding path for this run. Onboard nothing: no +entry files, no `index.yaml` change. Produce the report instead, land it, and +dispatch one child per constituent. + +`decompositions/mla-custom-op.md` is both a worked example and the shape to +copy. + +## 2. Granularity and sourcing policy + +Atomicity criterion: **Python-level indivisible**, as above. + +Sourcing preference, in order: + +1. **Ops exposed by tensorrt_llm as `torch.ops.trtllm.*`** — the default + catalog entry shape. Almost all constituents resolve here. +2. **Any other Python-callable atomic surface**, as an approved policy + exception, when no torch-op boundary exists. This was required exactly + once: 1.3.0rc21 registered **no** `torch.ops.trtllm.attention` + (verified exhaustively against all 252 registered `trtllm::` schemas — + and note that an exhaustive schema sweep is only exhaustive after the + lazy registrations are forced, which + `agent_flow/workflows/staircase_onboard/tools/list_trtllm_ops.py` does + and a bare `import tensorrt_llm` does not), + and the MLA attention core is otherwise reachable only inside the macro. + The pybind binding `tensorrt_llm.bindings.internal.thop.attention` is + Python-callable, single-call, fully explicit-state, and was onboarded as + `attention/thop_attention`. The trade-off is recorded in its contract: + it is invisible to torch.compile and is a private API with no stability + promise. +3. **torch-native glue** (`split`, `concat`, `reshape`, `expand`, `copy_`, + `empty`) as hand-added thin mirrors under `catalog/torch/`. + +### Writing a decomposition report for a new macro + +A requested op is a macro when its registered body contains Python-level +composition — registry lookup delegating to a module `forward`, module +calls, tensor glue beyond the single call, a phase dispatch, or nested +calls into other registered ops. In that case produce a report, not an +entry. The report must: + +- **Prove the macro claim**: cite the op body's Python-level composition, + as §1 does here. +- **Verify that every named symbol exists** in the installed engine before + citing it (`hasattr(torch.ops.trtllm, name)` or the registered-schema + list). Decomposition maps in this project once propagated a stale name + (`torch.ops.trtllm.attention`) across several runs before an onboarding + attempt exposed it. +- **Label every constituent with one of three states**: existing entry + (cite the catalog path) / not onboarded (op symbol plus proposed + category — the sourcing tiers above decide the shape) / gap (no + Python-callable atomic surface exists; say what is reachable instead). +- **Map the mainline path step by step** (the shape of §4) and list + variant branches with their extra ops separately (the shape of §6). +- **Claim nothing that only certification can establish.** Receipts, + numerical behavior, and caller-facing hard constraints (this document's + §4.5 and §5) exist only after the constituents are onboarded and tested + on GPU; a pre-onboarding report marks them as unknown. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/fuse-input-embeds.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/fuse-input-embeds.md new file mode 100644 index 000000000000..683e4fff566d --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/fuse-input-embeds.md @@ -0,0 +1,268 @@ +# Decomposing `fuse_input_embeds` (multimodal placeholder splice) + +Why the requested catalog entry `multimodal/fuse_input_embeds` was not written, +what the step actually is, which calls serve it, and the one call nobody in the +catalog can serve today. + +Derived on: tensorrt_llm 1.3.0rc26 | torch 2.12.0a0+5aff3928d8.nv26.05 | +NVIDIA GB300, `torch.cuda.get_device_capability()` `(10, 3)` → **sm_103** — the +stamp on the observations below, not a pin to assert. Receipts are keyed by GPU +architecture and nothing else, and each constituent carries its own in its own +contract; this document carries none. + +Checkpoint: `Qwen3.8-Flash-Next-NVFP4`, `model_type` `qwen4_exp`. + +## 1. Why decompose it + +The request was one entry for the step the run's `vocabulary.yaml` states as +"splice the vision embeddings into the text embedding rows at the placeholder +positions", sourced at tier 2 against the in-tree callable +`tensorrt_llm._torch.models.modeling_multimodal_utils.fuse_input_embeds`. + +That callable is not one kernel-launching call. Its body composes, at the +Python level: a list join, an optional index derivation, **a module call** +(`embedding_layer(...)`), an optional buffer allocation, and one or two indexed +row assignments. `references/decomposition.md` names three of those shapes +explicitly — "module calls", "tensor glue beyond the single call", "a phase +dispatch" — and `entry-spec.md` forbids the result: "A wrapper that chains +split → norm → matmul … is not an entry." + +**Measured, not argued.** Profiled on sm_103 at this checkpoint's widths +(`T = 2048` tokens, `M = 512` placeholder rows in 2 contiguous runs, +`H = 2560`, bf16, `vocab_size = 248320`), one `fuse_input_embeds` call on the +path this checkpoint takes launches **exactly 2 device kernels**: + +``` +1 x at::native::vectorized_gather_kernel<16, int> <- the embedding lookup +1 x at::native::index_elementwise_kernel<...index_put_kernel_impl>...> + <- the row scatter +``` + +Two kernel-launching calls, each of which is a legal single entry on its own, +and **one of the two is already in the catalog**. A wrapper chaining them is +precisely the Python-level composition the entry spec bans. + +The other branches of the same callable, measured in the same process, are +larger rather than smaller: + +| branch | device kernels | why it is not this checkpoint's path | +|---|---|---| +| in-vocab, executor-supplied indices | **2** | **this is the path** | +| in-vocab, indices derived in-forward | 15 (incl. 2 `Memcpy DtoH`) | the executor ships the indices; see §4.2 | +| OOV placeholders, executor indices | 4 | placeholders are in vocabulary; see §4.1 | + +### 1.1 No registered op serves this step + +**The sweep had to be made exhaustive first, and the run's own ruling R2 says +why.** An op registers when the module holding its decorator is imported, and +several register outside `custom_ops/`. Sweeping after +`import tensorrt_llm._torch.custom_ops` alone yields **294** `trtllm::` / +`tensorrt_llm::` schemas and misses them — demonstrated on this host with +`trtllm::embedding_skip_forward`, which is declared by +`@torch.library.custom_op` in `_torch/modules/embedding.py` and is absent from +that 294. Walk-importing every module under `_torch/modules`, +`_torch/custom_ops`, `_torch/models`, `_torch/attention`, +`_torch/cute_dsl_kernels` and `_torch/cuda_tile_kernels` (474 imported, 6 +failed on optional dependencies, 11 skipped as embargoed) raises the set to +**323**, of which **29** are new. **107** of the 323 match a name pattern +covering `scatter|index|put_|embed|fuse|splice|mask|place|insert|merge|copy| +gather|multimodal|vision|mm_|select`; every one was read, and none writes a row +subset of an activation buffer. + +The single candidate worth refuting by name is +`trtllm::inplace_slice_copy(Tensor(a!) dest, Tensor src, int dim1_start, +int dim1_end) -> ()`. It copies into a **dim-1 (column)** range across *all* +rows — the Eagle3 hidden-state capture in `_torch/speculative/eagle3.py` uses it +to stack per-layer states at `i * hidden_size`. The splice is the transpose of +that: a row subset, all columns. It does not serve. + +Spot checks, all `False` on this host after the full walk-import: +`hasattr(torch.ops.trtllm, x)` for `fuse_input_embeds`, `scatter_mm_embeds`, +`index_copy`, `index_put`, `masked_scatter`. The one op in the generated +inventory's `embedding` category, `embedding_skip_forward`, does resolve once +its module is imported, but it is a shape-only stand-in — its implementation +`embedding_skip_forward_impl` returns an uninitialised `new_empty` of the +output shape — not a lookup and not a splice. + +### 1.2 No tier-2 kernel-plane callable serves it either + +A tree-wide search (not a single directory — ruling R4's second-order note), +for `mm_token_indices|masked_scatter|index_copy_|index_put_` across +`tensorrt_llm/_torch/**.py`, returns no standalone callable under +`_torch/modules/` that performs this splice. The only in-tree servers are +`fuse_input_embeds` itself and inline code in per-model zoo files. Nothing under +the run's `source_embargo` was opened, and no constituent below routes through +it; the embargo is not what excludes tier 2 here. + +**Transparency note, in the spirit of ruling R5.** The exhaustive sweep of §1.1 +skipped the embargoed modules by name, but eight `trtllm::qwen4_exp_hc_*` / +`trtllm::qwen4_exp_ple_*` op *names* still appeared in the registry diff, +because a non-skipped module imports that package transitively. Only the +schema names were seen; no embargoed file was opened and none of those ops is a +multimodal splice, so nothing here depends on them — and ruling R1 would bar +wrapping them regardless. It is recorded because a sweep is the same hole a +tree-wide grep was. + +## 2. Granularity and sourcing policy + +Atomicity criterion: Python-level indivisible. Sourcing preference walked in +order — `torch.ops.trtllm.*` (§1.1, none), any other Python-callable atomic +surface (§1.2, none), torch-native glue as a thin mirror under `catalog/torch/`. +The third rung is where this step lands, and §3 constituent 3 is the +consequence. + +## 3. The resulting vocabulary + +Paths relative to `tensorrt_llm/_torch/_experimental/modeling_v2/catalog/`. + +| # | role | callable | state | +|---|---|---|---| +| 1 | join the per-request vision embeddings into one `[M, H]` block | `torch.cat` | **existing entry** `torch/concat.py` | +| 2 | text embedding lookup over the whole `input_ids` row | `torch.nn.functional.embedding` | **existing entry** `torch/embedding.py` | +| 3 | write the `[M, H]` block into the `M` placeholder rows | `torch.Tensor.index_copy_` | **gap** — see below | + +Every symbol was resolved on this host before being cited. `torch.cat`, +`torch.nn.functional.embedding` and `torch.Tensor.index_copy_` are torch +callables and need no registration check; the three catalog modules +`torch/concat.py`, `torch/embedding.py` and `torch/copy_.py` were imported and +their functions printed. `fuse_input_embeds` and +`filter_mm_token_from_input_ids` both import from +`tensorrt_llm._torch.models.modeling_multimodal_utils`. + +**Constituent 3 is the whole of the gap, and it is a placement question rather +than a capability one.** The computation exists and is one kernel; what does +not exist is a catalog entry for it. The natural home is `catalog/torch/`, +which already carries `copy_`, `concat`, `empty`, `expand`, `pad`, `reshape`, +`split`, `transpose` — movement and layout glue, which an indexed row write is. +smith may not write there (the role forbids it, and the reason given is that a +*computation* mirror would widen what the surface is for; an indexed row copy +is movement, not computation, so the stated reason does not bite — but the +prohibition is not smith's to read around). **The orchestrator decides whether +`torch/index_copy_.py` is added.** Nothing else in this decomposition is open. + +Three interchangeable spellings of constituent 3 were measured, bitwise +identical to one another at `T = 2048, M = 512, H = 2560, bf16`: + +| spelling | device kernels | kernel | +|---|---|---| +| `buf.index_copy_(0, idx, src)` | **1** | `index_copy_kernel_impl>` | +| `buf[idx, :] = src` (advanced index) | **1** | `index_put_kernel_impl>` | +| `buf.masked_scatter(mask2d, src)` | 6 | adds a cub scan, a size check and a DtoD copy | + +`index_copy_` is the recommendation: fewest launches, and its signature takes +the index tensor the executor already ships. `masked_scatter` is the form the +HuggingFace reference for this architecture family uses +(`get_placeholder_mask` → `inputs_embeds.masked_scatter(...)` in +`transformers/models/qwen3_vl/modeling_qwen3_vl.py`); it is listed so the +difference is on the record, not because it should be preferred here — it needs +a boolean mask the runtime does not ship, and it costs 5 extra launches. + +**The slice-copy shortcut does not generalise.** `catalog/torch/copy_.py` +already exists, and a contiguous placeholder run can be written with it — this +was measured bitwise equal to the scatter. But a batch with `n` runs needs `n` +calls, so `n` is data-dependent and the forward stops being flat. Measured at 3 +runs: `scatter == 3 slice copies` is `True`, and the scatter is 1 call against +3. The multimodal skill's `references/vision-path.md` records the single-run +shape ("the write is a slice copy rather than a scatter") from a different +target's single-image measurement; it is true of that case and is not the +general form. + +## 4. Assembly — the mainline path + +At this checkpoint, with `text_token_indices` / `mm_token_indices` arriving +from the executor and `mm_embeds` already one pre-concatenated tensor: + +``` +mm_embed = concat(mm_embeds) # only when len(mm_embeds) > 1 +inputs_embeds = embedding(input_ids, weight) # [T, 2560] bfloat16 +inputs_embeds.index_copy_(0, mm_token_indices, mm_embed) +``` + +Verified end to end on this host against an independently built native-torch +reference (`F.embedding(...).clone()`, then the scatter), **bitwise equal** to +`fuse_input_embeds`' own output at `T = 2048, M = 512` (2 runs) and +`T = 8192, M = 2048` (4 runs), for both the advanced-index and the +`index_copy_` spelling. Passing two `mm_embeds` tensors instead of one +pre-concatenated tensor is also bitwise equal. + +### 4.1 Why the in-vocab branch is the one to assemble + +`config.json` gives `image_token_id 248056`, `video_token_id 248057` and +`text_config.vocab_size 248320`. Both placeholders are **inside** the +vocabulary, so the embedding lookup over the whole `input_ids` row is in +bounds and the OOV branch — which gathers only text positions into a fresh +`torch.empty` buffer and scatters twice — is dead code for this target. The +OOV branch costs 4 kernels against 2 and would additionally need +`catalog/torch/empty.py` and an `index_select`. + +Caveat the assembler owns: the target must declare its placeholder ids, because +`prepare_multimodal_indices` (`_torch/pyexecutor/engine/runners/common.py`) +reads `multimodal_token_ids` (falling back to `mm_token_ids`) off the model and +otherwise takes the `>= vocab_size` OOV predicate. + +### 4.2 Why the index derivation is not in the forward + +`ship_multimodal_indices` (same file) pins and async-copies the executor's +CPU-computed `mm_token_indices` / `text_token_indices` into the inputs dict, so +the forward receives them. Deriving them inside the forward instead costs 13 +extra kernels **and two `Memcpy DtoH` host syncs** — measured: the standalone +`filter_mm_token_from_input_ids` call launches 13 kernels including 2 +`Memcpy DtoH (Device -> Pinned)`, and the whole `fuse_input_embeds` call rises +from 2 launches to 15. A target that forwards the two index tensors explicitly +keeps the step at 2 launches and sync-free; one that relies on a `**kwargs` +splat silently takes the 15-launch syncing path. + +This also means constituent 3's index tensor is `int64` and device-resident, +which is what `index_copy_` wants. + +### 4.3 The deepstack branch is inert here + +`vision_config.deepstack_visual_indexes` is `[]` in this checkpoint's +`config.json`, so `fuse_input_embeds`' `extra_embeds` path — which allocates a +`torch.zeros` buffer per extra feature and scatters into it — is never entered. +It is named here only so a later checkpoint in this family does not read its +absence as an oversight. + +### 4.4 The one guard the callable raises, quoted + +Passing a `mm_embeds` row count that disagrees with the index count is +rejected, not silently truncated. Observed verbatim on this host: + +``` +ValueError: Multimodal token count mismatch: found 512 image tokens in input_ids but received 10 image embeddings. +``` + +Whether the constituent calls carry that guard themselves is a question for +constituent 3's own contract, after it is onboarded; a bare `index_copy_` with +a short source does not raise this message. + +## 5. Certification status + +**Nothing here is certified.** No entry files were written, no test file, no +`index.yaml` row, no receipt. The launch counts, the bitwise equalities and the +quoted `ValueError` above are observations from two probe runs on sm_103; they +are evidence for the decomposition, not certification of any constituent. + +- Constituents 1 and 2 are existing catalog entries and carry their own + receipts in their own contracts. This document does not restate them. +- Constituent 3 has no entry and therefore no receipt, no certified dtype + column and no certified shape column. What `index_copy_` does with duplicate + indices, out-of-range indices, a non-contiguous source, or a dtype mismatch + between source and destination is **unknown here** — it was not probed, + because a macro report may not claim what only certification can establish. + +## 6. Out of scope + +- The vision tower that produces `mm_embeds` (queue items 17–21) and the + patch-merger projection that maps it to `out_hidden_size 2560` (served by + `gemm/cublas_mm`). +- Where the splice runs. `references/vision-path.md` of + `staircase-modeling-multimodal` records that it belongs in the shell's + forward, producing `inputs_embeds` with `input_ids=None`, because + `DecoderModelForCausalLM.forward` forwards only the keywords it names. That + is an assembly question, not a vocabulary one. +- Chunked prefill and KV-cache reuse slicing of `mm_embeds` + (`find_input_mm_embeds`): it reads only CPU-resident counts off + `MultimodalParams.multimodal_runtime` and launches nothing, so it is host + bookkeeping rather than a catalog constituent — but it does decide which rows + reach constituent 3, and a target that skips it will fail §4.4's guard. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/gdn-custom-op.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/gdn-custom-op.md new file mode 100644 index 000000000000..3d5d3ef0dbed --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/gdn-custom-op.md @@ -0,0 +1,434 @@ +# Decomposing `gdn_custom_op_inplace` + +Why TRT-LLM's Gated DeltaNet block op cannot be a catalog entry, what its +Python-level atoms are, and which of them already have a sourcing answer. + +This is a **pre-onboarding** decomposition report, written to the checklist in +`../decomposition.md`. It records what was read and what was measured on this +machine. It deliberately contains **no receipts and no caller-facing numerical +guarantees** — those exist only after each constituent is onboarded and tested +on GPU. + +Derived on: tensorrt_llm 1.3.0rc26 | torch 2.12.0a0+5aff3928d8.nv26.05 | +flashinfer 0.6.18 | NVIDIA GB300 (sm_103) — the stamp on the observations +below, not a pin to assert. Receipts are keyed by GPU architecture and nothing +else; each constituent will carry its own once onboarded. + +Geometry throughout is Qwen3.8-27B-NVFP4's GDN layer as the run's +`vocabulary.yaml` states it: `num_k_heads=16`, `head_k_dim=128`, +`num_v_heads=48`, `head_v_dim=128`, `mixed_qkv` width 10240, +`linear_conv_kernel_dim=4`, 48 `linear_attention` layers, and — from the run +spec's `--mamba_ssm_cache_dtype=float32` — an **fp32 recurrent-state pool**. +That last one is not a detail; it selects implementations. See §5. + +## 1. Why it is a macro + +The registered schema is the whole argument: + +``` +trtllm::gdn_custom_op_inplace(Tensor mixed_qkv, Tensor a, Tensor b, + str layer_idx, Tensor(a4!) output) -> () +``` + +Five arguments, one of them a **string**. No weights, no convolution state, no +recurrent state, no sequence-length metadata, no cache indices. Everything the +computation actually needs is fetched from process-global state inside the +body (`tensorrt_llm/_torch/modules/mamba/gdn_mixer.py`, lines 94-112): + +```python +@torch.library.custom_op("trtllm::gdn_custom_op_inplace", mutates_args=("output",)) +def gdn_custom_op_inplace(mixed_qkv, a, b, layer_idx, output) -> None: + attn_metadata, gdn_layer, spec_metadata = _extract_gdn_extra_attrs(layer_idx) + num_tokens = attn_metadata.num_tokens + gdn_layer.forward_core(...) +``` + +`_extract_gdn_extra_attrs` is a **registry lookup**: it reads +`get_model_extra_attrs()`, pulls a weakref to the live `AttentionMetadata`, and +pulls a weakref to the `Qwen3NextGatedDeltaNet` **nn.Module** out of an +`extra_attrs["gdn_layers"]` dict keyed by the `layer_idx` string. It then calls +that module's `forward_core` — a method, not a kernel. + +That is the macro signature named in `../decomposition.md` verbatim: *"a +registry lookup delegating to a module `forward`"*. `forward_core` is in turn a +**Python-level phase dispatch** (`num_prefills > 0` selects `forward_extend`, +else `forward_decode`; each branch splits further on `is_target_verify` and +`use_replay`) over a pipeline of module calls, tensor glue, several distinct +kernel wrappers, and indexed state reads and writes. + +**Measured, not inferred.** Driving the op the only way a catalog wrapper +could — real tensors, a `layer_idx` string, a real output buffer, nothing else, +because nothing else is in the schema: + +``` +torch.ops.trtllm.gdn_custom_op_inplace(mixed_qkv, a, b, "0", out) + RAISED AssertionError: Model extra attrs is not set +``` + +A catalog wrapper cannot satisfy that precondition without constructing the +whole model's module tree and registering it in `model_config.extra_attrs`, +which is precisely the "black box whose interior cannot be swapped, measured or +certified independently" the catalog exists to prevent. **Onboard nothing.** + +For completeness, what the body reaches for beyond its five arguments: + +| source | fields consumed | +| --- | --- | +| `get_model_extra_attrs()` | `attention_metadata` (weakref), `gdn_layers[layer_idx]` (weakref to the module), `spec_metadata` | +| `AttentionMetadata` | `num_tokens`, `num_contexts`, `num_ctx_tokens`, `seq_lens`, `mamba_metadata`, `kv_cache_manager` | +| `kv_cache_manager` | `mamba_layer_cache(layer_idx)`, `is_speculative()`, `use_replay_state_update`, `get_replay_state_update_metadata()`, `use_gdn_cached_replay_all_layer_commit` | +| `Mamba2Metadata` | `has_initial_states`, `state_indices`, `query_start_loc`, `query_start_loc_long`, `replay_work_items`, `replay_n_writes` | +| layer cache | `.conv`, `.temporal`, `.intermediate_conv_window`, `.intermediate_ssm`, `.old_x`, `.old_B`, `.old_dt`, `.old_dA_cumsum`, `.cache_buf_idx`, `.prev_num_accepted_tokens` | +| the module itself | `conv1d.weight`, `conv1d.bias`, `A_log`, `dt_bias`, `activation`, and every geometry field | +| `SpecMetadata` | `runtime_draft_len` | + +## 2. Scope of this map + +`gdn_custom_op_inplace` covers the GDN **core** only. Two neighbours are +adjacent in the layer but **outside** the op body, and must not be +double-counted by a child: + +- the `in_proj_qkvz` / `in_proj_ba` projections and the `[Q|K|V|Z]`, `[b|a]` + column slices happen in `_compute_tokenwise_inputs`, before the op; +- the gated RMS norm and `out_proj` happen in `_postprocess_gdn_output`, after + the op. That norm is `torch.ops.trtllm.rms_norm_gated_token_major`, which is + already a separate item in this run's queue. + +## 3. Symbol verification + +Every symbol named in this document was resolved in the installed engine on +this machine before being cited (spool logs `smith1-gdn-macro-probe`, +`smith1-gdn-constituent-reach`). + +`torch.ops.trtllm` ops **present**: `gdn_custom_op_inplace`, +`causal_conv1d_fwd`, `causal_conv1d_update`, `rms_norm_gated_token_major`, +`fused_gated_rmsnorm_quant`, `mamba2_mtp_ssm_cache_update`, `mamba_conv1d`, +`selective_scan`. + +`torch.ops.trtllm` ops **absent** — probed, not assumed, because the absence is +the load-bearing fact for §4: `chunk_gated_delta_rule`, `gated_delta_rule`, +`fused_recurrent_gated_delta_rule`, `fused_sigmoid_gating_delta_rule_update`, +`fused_gdn_post_conv`, `l2norm`, `gdn_gating`. **The delta-rule core — the +actual recurrence this layer exists to compute — has no `torch.ops.trtllm` +surface at all.** Only the two convolution ops do. + +Two schemas have drifted from +`.claude/skills/staircase-kernel-feasibility/references/trtllm-ops.md`; the +generated inventory is stale on both. Observed here: + +``` +causal_conv1d_fwd(..., int pad_slot_id, Tensor($3! -> )? out=None) -> () # inventory omits `out` +rms_norm_gated_token_major(..., Tensor? fp8_scale=None, + str gate_activation="silu") -> Tensor # inventory omits `gate_activation` +``` + +All 24 Python-callable constituents resolved (`True` for every entry in +section C of `smith1-gdn-macro-probe.log`). `flashinfer.chunk_gated_delta_rule` +exists in flashinfer 0.6.18; `flashinfer.gated_delta_rule` does **not**. + +## 4. The constituents + +Labels follow `../decomposition.md`: **existing entry** / **not onboarded** +(symbol + proposed category) / **gap** (no Python-callable atomic surface in +the preferred tiers; what is reachable instead). + +Entry paths are relative to +`tensorrt_llm/_torch/_experimental/modeling_v2/catalog/`. + +### 4.1 Tier 1 — a `torch.ops.trtllm` op exists + +| role | symbol | state | +| --- | --- | --- | +| prefill / varlen causal depthwise conv1d + silu, with conv-state read and write | `torch.ops.trtllm.causal_conv1d_fwd` (Python surface `..modules.mamba.causal_conv1d.causal_conv1d_fn`) | not onboarded → `ssm` | +| decode conv1d step + conv-state ring update | `torch.ops.trtllm.causal_conv1d_update` (Python surface `..modules.mamba.causal_conv1d.causal_conv1d_update`) | not onboarded → `ssm` | + +Both Python surfaces are one op invocation plus allocation/`unsqueeze` glue — +compliant wrapper shape as-is. `causal_conv1d_fn` allocates a channel-last +`out` when `x` is channel-last; on this run's prefill path `x` arrives as a +plain contiguous `[10240, T]` and `out` stays `None`, so the op writes `x` in +place. + +### 4.2 Tier 2 — no torch-op boundary; a single-kernel Python callable exists + +Each of these is one Triton kernel launch plus allocation. Wrapping a +non-`torch.ops.trtllm` callable is an **approved-exception** shape per +`entry-spec.md` — it needs the run's explicit authorization, and the import +path must be recorded in the wrapper docstring, the contract, and the index +`impl` field (the precedent is `attention/thop_attention`). + +| role | symbol | proposed category | +| --- | --- | --- | +| post-conv fusion: q/k/v split of the conv output, **Q/K L2 norm**, `g = -exp(A_log) * softplus(a + dt_bias)`, `beta = sigmoid(b)` — five outputs, one kernel | `..modules.mamba.fuse_elementwise_ops.fused_gdn_post_conv` | `ssm` | +| decode delta-rule core: in-kernel sigmoid gating + recurrent gated delta rule, initial state gathered from the pool by index, state written back | `..modules.fla.fused_sigmoid_gating_recurrent.fused_sigmoid_gating_delta_rule_update` | `ssm` | +| zero the recurrent and conv rows of fresh (no-initial-state) requests | `..modules.mamba.recurrent_state_cache.reset_recurrent_state_rows` | `ssm` | +| MTP verify: conv1d step writing per-draft-token intermediate conv windows | `..modules.mamba.causal_conv1d_triton.causal_conv1d_update` | `ssm` | +| MTP verify: pack decode Q/K/V from the conv output, **without** norm or gates | `..modules.mamba.fuse_elementwise_ops.pack_gdn_decode_qkv` | `ssm` | +| MTP verify: `g` alone, when the gating is not fused into the recurrence | `..modules.mamba.gdn_mixer.fused_gdn_gating` | `ssm` | +| MTP verify: recurrent gated delta rule over `P` draft tokens, writing per-step intermediate SSM states, live pool untouched | `..modules.fla.fused_recurrent.fused_recurrent_gated_delta_rule_update` | `ssm` | +| MTP replay: cached-history replay verify (only when `kv_cache_manager.use_replay_state_update`) | `..modules.fla.cached_replay.fused_recurrent_gated_delta_rule_cached_replay_update` | `ssm` | + +**One boundary question a child must settle, not assume.** +`fused_sigmoid_gating_delta_rule_update` is not unconditionally a single call: +its body dispatches to a FlashInfer CuTe-DSL kernel when +`_can_use_flashinfer_gdn_decode` passes, and to its own Triton kernel +otherwise. On this run the FlashInfer branch is statically ineligible (§5), so +the function reduces to one Triton launch — but the entry must say which +implementation its receipt certifies, and must not claim the other. + +### 4.3 Movement-only — placement is the orchestrator's call + +| role | symbol | +| --- | --- | +| extract a column slice of the token-major projection and transpose it to `[C, T]` for the prefill conv | `..modules.mamba.fuse_elementwise_ops.extract_transpose_prefill_slice` | + +Measured equal to `src[:T, s:s+w].t()` elementwise (`True` in the probe). It +launches no math — it is layout and movement, which is what `catalog/torch/` is +for, but it is a trtllm Triton kernel rather than a torch callable. **Reported, +not placed:** `catalog/torch/` is never written by this role. + +Ordinary torch glue the assembly also needs, already in the catalog: +`torch/split` (state indices, `mixed_qkv` prefill/decode halves), `torch/empty`, +`torch/reshape`, `torch/transpose`, `torch/concat` (the MTP branch's +prefill+decode output join), `torch/copy_`. + +Two more are **computation**, not glue, and have no entry: `torch.exp` (inside +the FlashInfer prefill adapter) and `Tensor.sigmoid` (the MTP fallback's +`beta`). A computation mirror widens what `catalog/torch/` is for, so both are +reported for an owner's decision rather than assumed. + +### 4.4 Gap — the prefill delta-rule core + +This is the one place where no preferred-tier surface exists, and it is the +most important step in the layer. + +`gdn_mixer._resolve_chunk_gated_delta_rule()` picks one of two implementations +at first call. **On sm_103 it returned +`tensorrt_llm._torch.modules.fla.flashinfer_chunk.chunk_gated_delta_rule`** — +measured, not read: `TLLM_USE_FLASHINFER_GDN_PREFILL` defaults to `1` and +`is_flashinfer_gdn_supported_arch()` is `True` here. + +**The selected implementation is itself a sub-macro.** `flashinfer_chunk.py` is +an adapter whose body is Python-level composition: + +``` +g2 = torch.exp(g.squeeze(0)) # log-space -> linear alpha +q3, k3 = l2norm_fwd(q3), l2norm_fwd(k3) # only if use_qk_l2norm_in_kernel +gathered = gather_cast_vk_to_fp32_vk(state, indices, dtype) # triton: indexed gather (+cast) +out, st = flashinfer.chunk_gated_delta_rule(...) # EXTERNAL package +cast_scatter_fp32_vk_to_vk(st, state, indices) # triton: indexed scatter (+cast) +``` + +Its atomic interior is `flashinfer.chunk_gated_delta_rule` — neither a +`torch.ops.trtllm` op (probed absent, §3) nor a torch callable, but a function +of an **external pip package**. That is outside all three sourcing tiers in +`../decomposition.md`. **This needs an owner's ruling before a child is +dispatched**, and it is the single blocking decision in this decomposition. +What is reachable, stated plainly: + +- `flashinfer.chunk_gated_delta_rule` — one call, explicit state, no registry. + Shaped like a legal entry; sourced from outside tensorrt_llm entirely. +- `..modules.fla.fused_state_io.gather_cast_vk_to_fp32_vk` and + `cast_scatter_fp32_vk_to_vk` — one Triton kernel each, tier 2, `ssm`. +- `..modules.fla.l2norm.l2norm_fwd` — one Triton kernel, tier 2. **Not on this + run's path**: GDN prefill calls the adapter with + `use_qk_l2norm_in_kernel=False` because `fused_gdn_post_conv` has already + normalized Q/K. + +The alternate implementation (`TLLM_USE_FLASHINFER_GDN_PREFILL=0`, or any +unsupported arch) is the vendored Triton `..modules.fla.chunk`, which is **also +a sub-macro** — seven Python-level steps, each its own Triton wrapper: +`l2norm_fwd` → `chunk_local_cumsum` → `chunk_scaled_dot_kkt_fwd` → +`solve_tril` → `recompute_w_u_fwd` → `chunk_gated_delta_rule_fwd_h` → +`chunk_fwd_o`. All seven resolve. Decomposing to that level buys seven entries +for a path this arch does not select; recorded as the fallback's shape, not +proposed as work. + +## 5. What this run's configuration selects (measured) + +The run spec passes `--mamba_ssm_cache_dtype=float32`. Both FlashInfer +decode-side kernels gate on a **bfloat16** state pool, so that one flag decides +which kernels this target will need. Measured on sm_103, with an otherwise +eligible geometry (`K = V = 128`): + +| predicate | fp32 pool (this run) | bf16 pool | +| --- | --- | --- | +| `_can_use_flashinfer_gdn_decode(..., T=N=2)` | **False** | True | +| `_can_use_flashinfer_gdn_verify(..., draft=4)` | **False** | True | + +So on this run: prefill takes the **FlashInfer** chunk path, while decode and +MTP verify take the **Triton** kernels. A decomposition derived from a bf16-KV +run would have named three different constituents. Supporting facts measured +alongside: `is_flashinfer_gdn_supported_arch() == True`, +`is_sm_100f() == True`, `_FLASHINFER_GDN_BF16_STATE_AVAILABLE == True`, +`_FI_GDN_MAX_MTP_T == 8`. + +## 6. Assembly by batch mode + +Phase dispatch over metadata reads is plain Python and is permitted outside +entries. `[Q|K|V|Z]` / `[b|a]` slicing precedes all of this and the gated norm +plus `out_proj` follow it (§2). + +Shared preamble, every mode: + +``` +state_indices_p, state_indices_d = split(state_indices, [num_prefills, num_decodes]) +if num_prefills > 0: + reset_recurrent_state_rows(ssm_states, state_indices_p, has_initial_states_p, conv_states) +``` + +### 6.1 Prefill only (`num_prefills > 0`, `num_decode_tokens == 0`) + +``` +xt = extract_transpose_prefill_slice(mixed_qkv, T, 0, 10240) # [10240, T] +xt = causal_conv1d_fn(xt, conv_w, None, conv_states=..., has_initial_state=..., + cache_indices=..., query_start_loc=...) # causal_conv1d_fwd + silu +q,k,v,g,beta= fused_gdn_post_conv(xt, None, a, b, A_log, dt_bias, 16,128, 48,128) +o, _ = chunk_gated_delta_rule(q,k,v,g,beta, initial_state=ssm_states, + initial_state_indices=cache_indices, + inplace_indexed_state_update=True, + cu_seqlens=query_start_loc_long, + use_qk_l2norm_in_kernel=False, output=output) +``` + +`inplace_indexed_state_update=True` writes the new recurrent state straight +back into the shared pool; the caller owns the guarantee that the indices do +not alias live slots (the source says so at the call site, and this report does +not upgrade that to a certified constraint). + +### 6.2 Decode only (`num_prefills == 0`) + +``` +mixed_qkv = causal_conv1d_update(mixed_qkv, conv_states, conv_w, None, "silu", + conv_state_indices=cache_indices) +q,k,v = views over mixed_qkv # column slices, no kernel +o = fused_sigmoid_gating_delta_rule_update(A_log, a, dt_bias, 1.0, 20.0, + q,k,v,b, ssm_states, cache_indices, + use_qk_l2norm_in_kernel=True, + cu_seqlens=query_start_loc_long, + output=output) +``` + +Note the asymmetry with 6.1 and do not flatten it: **on the decode path the L2 +norm and the gating are inside the recurrence kernel** (`a`/`b` go in raw), so +there is no `fused_gdn_post_conv` call and no separate `g`/`beta`. On the +prefill path they are inside `fused_gdn_post_conv` instead, and the recurrence +is told `use_qk_l2norm_in_kernel=False`. Same math, different fusion boundary. + +### 6.3 Mixed prefill + decode + +Both convolutions run (`causal_conv1d_fn` on the prefill rows, +`causal_conv1d_update` on the decode rows), then a **single** +`fused_gdn_post_conv(prefill, decode, ...)` call fuses both halves into one +`q/k/v/g/beta` set, and one `chunk_gated_delta_rule` covers the batch. + +### 6.4 MTP target verification (`is_target_verify`) + +Reached when there are decode rows, `spec_metadata` is present and the KV cache +manager is speculative — i.e. the `mtp` feature this run enables. Extra +constituents over 6.2, all listed in §4.2: + +``` +mixed_qkv_d = causal_conv1d_update_triton(..., intermediate_conv_window=..., + intermediate_state_indices=arange(num_decodes)) +q,k,v = pack_gdn_decode_qkv(mixed_qkv_d, 16,128, 48,128) # no norm, no gates +beta = b.sigmoid(); g = fused_gdn_gating(A_log, a, dt_bias) +o = fused_recurrent_gated_delta_rule_update(q,k,v,g,beta, + initial_state_source=ssm_states[state_indices_d], + initial_state_indices=arange(num_decodes), + use_qk_l2norm_in_kernel=True, disable_state_update=True, + intermediate_states_buffer=intermediate_ssm_states, + cache_steps=draft_token_num) +``` + +Final state selection is deferred to the cache manager's +`update_mamba_states()`, which is runtime, not modeling. When +`kv_cache_manager.use_replay_state_update` is on, this is replaced by +`fused_recurrent_gated_delta_rule_cached_replay_update` (source-asserted limits +`draft_token_num <= 8`, `replay_history_size <= 16`). When the pool is bf16 — +**not this run** — `_flashinfer_gdn_verify` replaces it instead. + +## 7. Reachability sweep (this geometry, this machine) + +Every constituent below was called once at the checkpoint's geometry, to +establish that it is reachable and that the decomposition is not proposing work +against a dead path. **These are not receipts**: each says "ran, returned the +expected shape, output finite", nothing about numerical correctness. Log +`smith1-gdn-constituent-reach`. + +| call | result | +| --- | --- | +| `causal_conv1d_fn`, `[10240, 37]` bf16, 2 seqs | `(10240, 37)` bf16, finite, conv state written | +| `causal_conv1d_update`, `[4, 10240]` bf16 | `(4, 10240)` bf16, finite | +| `fused_gdn_post_conv`, prefill-only, T=37 | q `(1,37,16,128)`, v `(1,37,48,128)`, g `(1,37,48)` fp32, beta fp32, finite | +| `extract_transpose_prefill_slice` | `(10240, 37)`, elementwise equal to `src[:T,:C].t()` | +| `pack_gdn_decode_qkv`, T=16 | q `(1,16,16,128)`, v `(1,16,48,128)` | +| `flashinfer_chunk`, **fp32** pool, T=64, 2 seqs | `(1,64,48,128)` bf16, finite; slots 0/1 updated, slots 2+ bitwise untouched | +| `flashinfer_chunk`, bf16 pool | same shape, finite, same slot behaviour | +| `fused_sigmoid_gating_delta_rule_update`, **fp32** pool, N=4 | `(1,4,48,128)` bf16, finite, state updated | +| `fused_sigmoid_gating_delta_rule_update`, bf16 pool | same | +| `fused_recurrent_gated_delta_rule_update`, N=2 P=4, fp32 pool | `(2,4,48,128)` bf16, finite, intermediates written | +| `reset_recurrent_state_rows` | slot 0 zeroed, slot 1 kept, conv slot 0 zeroed | + +**The fp32 pool is not a blocker for the prefill path.** That was the open risk +in §5 — `flashinfer_chunk` passes the pool dtype straight through on +`is_sm_100f()` archs — and it ran. + +### 7.1 A domain fact worth carrying forward + +The first sweep ran the vendored Triton chunk on raw `randn` q/k and got +all-NaN. That is **not** a kernel defect, and a child must not record it as +one. Re-probed deliberately: + +| input | vendored triton chunk | flashinfer chunk | +| --- | --- | --- | +| raw `randn` q/k, `use_qk_l2norm_in_kernel=False` | **NaN** (393216 elements) | finite but absmax `1.02392e+13` | +| caller-normalized q/k, `in_kernel=False` (production shape) | finite, absmax `0.0742188` | finite, absmax `0.074707` | +| raw `randn` q/k, `use_qk_l2norm_in_kernel=True` | finite, absmax `0.0742188` | — | + +The gated delta rule is only stable on **L2-normalized** q/k, which is exactly +why the production prefill path normalizes inside `fused_gdn_post_conv` and +then passes `use_qk_l2norm_in_kernel=False`. Out of that domain the two +implementations diverge in different directions — one to NaN, one to 1e13 — +so neither can be used to judge the other there. In domain they agree to +`maxabs = 0.000488281` (`maxrel = 91.4489`, a tiny-denominator artifact on +near-zero bf16 outputs, not a disagreement in magnitude). + +Any test written for either prefill core must build its q/k normalized, or it +will be measuring divergence rather than the kernel. + +## 8. Proposed child dispatches + +One entry each, no dependencies between them. Categories proposed, not settled; +`ssm` is an existing catalog category with no directory yet. + +| # | target | tier | note | +| --- | --- | --- | --- | +| 1 | `torch.ops.trtllm.causal_conv1d_fwd` | 1 | prefill/varlen conv + silu + conv-state write | +| 2 | `torch.ops.trtllm.causal_conv1d_update` | 1 | decode conv step + conv-state update | +| 3 | `fuse_elementwise_ops.fused_gdn_post_conv` | 2 | needs tier-2 authorization | +| 4 | `fused_sigmoid_gating_recurrent.fused_sigmoid_gating_delta_rule_update` | 2 | decode core; must name which branch the receipt certifies (§4.2) | +| 5 | **prefill delta-rule core** | — | **blocked on the §4.4 sourcing ruling** | +| 6 | `recurrent_state_cache.reset_recurrent_state_rows` | 2 | | +| 7 | `fuse_elementwise_ops.extract_transpose_prefill_slice` | 2 | placement question, §4.3 | +| 8 | `causal_conv1d_triton.causal_conv1d_update` | 2 | MTP only | +| 9 | `fuse_elementwise_ops.pack_gdn_decode_qkv` | 2 | MTP only | +| 10 | `gdn_mixer.fused_gdn_gating` | 2 | MTP only | +| 11 | `fused_recurrent.fused_recurrent_gated_delta_rule_update` | 2 | MTP only | + +Deliberately **not** proposed: the seven vendored-Triton chunk sub-steps (§4.4, +alternate path this arch does not select), `l2norm_fwd` (not on this run's +path), `_flashinfer_gdn_verify` and `..fla.cached_replay` (both gated off by +the fp32 pool / the replay flag), and `torch.exp` / `Tensor.sigmoid` (owner's +decision on widening `catalog/torch/`). + +## 9. Out of scope / unknown + +- **Every numerical claim.** Nothing here certifies accuracy, tolerance, dtype + limits or arch behaviour. §7 establishes reachability and shapes only. +- **Tensor parallelism.** Read at tp1. The module divides `num_k_heads` and + `num_v_heads` by `attn_tp_size` and asserts divisibility; nothing here was + measured above tp1. +- **The vision tower.** Out of scope for this target by the run's own rule. +- **`mamba_conv1d` / `selective_scan`.** Present as trtllm ops and conceptually + adjacent, but not called from this macro — they serve the Mamba2 mixer, not + GDN. Named only so a later reader does not mistake them for the GDN path. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/mla-custom-op.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/mla-custom-op.md new file mode 100644 index 000000000000..95d6f8f4d006 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/mla-custom-op.md @@ -0,0 +1,255 @@ +# Decomposing `mla_custom_op_inplace` + +How TRT-LLM's MLA block macro was broken into Python-level atomic catalog +entries, how each batch mode reassembles from them, and what the current +coverage does not include. + +This document has two roles. It is the reference record of the completed +`mla_custom_op_inplace` decomposition, and it is the **norm and exemplar +for handling any macro op**: when an onboarding request turns out to +target a macro, onboard nothing — produce a decomposition report shaped +like this document at `.claude/skills/staircase-kernel-onboard/references/decompositions/.md` +instead. What such a report must contain is specified at the end of §2. + +Derived on: tensorrt_llm 1.3.0rc21 | torch 2.11.0+cu130 | NVIDIA B200 +(sm_100) — the stamp on the observations below, not a pin to assert. +Receipts themselves are keyed by GPU architecture and nothing else; each +entry cited below carries its own, in its own contract. + +## 1. Why decompose it + +`torch.ops.trtllm.mla_custom_op_inplace` looks like one op — a single +registered custom op, one node in an FX graph — but its body is a **macro**: +it looks up a registered `MLA` module by `layer_idx` and runs +`MLA.forward_impl`, which is a full Python-level pipeline: + +- module calls: `kv_a_proj_with_mqa`, `q_a_layernorm` / `kv_a_layernorm` + (dual-stream), `q_b_proj`; +- tensor glue: `split`, `concat`, slicing, scratch allocation; +- a **Python-level phase dispatch** (context vs. absorption generation; + context further splits into fresh / cached-KV / chunked-prefill); +- further registered ops inside each branch (`bmm_out`, + `mla_rope_generation`, cache load/merge ops, the FMHA). + +The catalog's admission rule is that every entry must be a **Python-level +indivisible op**: one callable whose interior is C++/CUDA only. Multiple +kernel launches inside one call are fine; Python-level composition is not. +Vertical customization depends on this granularity — a black-box macro +cannot have interior steps swapped, measured, or certified independently, +and its hidden control flow cannot be specialized per target. "FX-graph +atomic" is therefore not an admission criterion; "Python-level atomic" is. + +The macro itself was consequently rejected as an entry (an early onboarding +run produced a working contract + test for it, which was deliberately +reverted). Its decomposition map was kept and each constituent was +onboarded individually. + +## 2. Granularity and sourcing policy + +Moved to `../decomposition.md`, which is where the policy belongs — this file +is the worked example, not the rule. Read it there before judging a boundary. + +## 3. The resulting vocabulary + +Shared (all batch modes): + +| role | entry | +| --- | --- | +| fused `q_a`/`ckv`/`k_pe` projection, `q_b` projection, `kv_b` expansion, `o_proj` | `gemm/cublas_mm` | +| `q_a` / `kv_a` RMS norm | `norm/flashinfer_rmsnorm` | +| projection output split, latent `[ckv \| k_pe]` concat, views | `torch/split`, `torch/concat`, `torch/reshape` | +| `k_pe` head broadcast, `k_nope` slice fill | `torch/expand`, `torch/copy_` | +| scratch / output allocation | `torch/empty`, `attention/create_mla_outputs` | + +Phase-specific: + +| role | entry | +| --- | --- | +| decode: absorption BMMs (`q_nope x k_b_trans`, `attn x v_b`) | `gemm/bmm_out` | +| decode: q_pe RoPE + latent append + FMHA scheduler-buffer fill | `attention/mla_rope_generation` | +| cached/chunked prefill: new-token RoPE + latent append + q assign | `attention/mla_rope_append_paged_kv_assign_q` | +| cached prefill: full `[past + new]` latent gather | `attention/load_paged_kv_cache_for_mla` | +| chunked prefill: per-loop latent slice gather | `attention/load_chunked_kv_cache_for_mla` | +| chunked prefill: log-sum-exp fold of partial passes | `attention/merge_chunked_attention_for_mla` | +| attention core, all phases (5 certified configurations) | `attention/thop_attention` | + + + +Entry paths in this document are relative to +`tensorrt_llm/_torch/_experimental/modeling_v2/catalog/`. + +## 4. Assembly by batch mode + +Shared front section (every mode): + +``` +[q_a | ckv | k_pe] = cublas_mm(hidden, W_a.t()) # then torch/split +q_a = flashinfer_rmsnorm(q_a, w_q_ln); ckv = flashinfer_rmsnorm(ckv, w_kv_ln) +q = cublas_mm(q_a, W_qb.t()) +latent = concat([ckv, k_pe]) +``` + +Phase dispatch (context vs. generation, and which context flavor) is plain +Python over metadata reads — permitted outside entries by catalog rules. +`o_proj` (`cublas_mm`) closes every mode. + +### 4.1 Plain prefill (fresh context) + +``` +kv = cublas_mm(ckv, W_kvb.t()); k_nope, v = split(kv) +k = empty(...); copy_(k[..., :nope], k_nope); copy_(k[..., nope:], expand(k_pe)) +thop_attention(context_only: separate q/k/v, in-kernel GPT-J RoPE on q/k_pe, + latent-cache append, causal FMHA) -> output +``` + +### 4.2 Decode (absorption generation) + +``` +q_nope, q_pe = split(q); fused_q, scratch = empty(...) +bmm_out(q_nope, k_b_proj_trans, fused_q[..., :C]) # absorbed q +mla_rope_generation(fused_q, q_pe, latent, ...) # RoPE + append + + # scheduler buffers +thop_attention(generation_only: latent MQA over the pool, + consumes the scheduler buffers) -> attn latent +bmm_out(attn_latent, v_b_proj, out) # v_b absorption +``` + +**Whether the two producers may run concurrently depends on the pool's +dtype, and getting it wrong is a race rather than an error.** + +On a **bf16** latent pool they are compatible by design: `fused_q` is +written disjointly (`[..., :C]` by the BMM, `[..., C:]` by the RoPE op), +so the two can overlap. + +On an **fp8-e4m3** latent pool that is false. `mla_rope_generation` does +not write `fused_q` at all there — it *reads* `fused_q[..., :C]` to build +`quant_q_buffer`, which is what the decode FMHA then takes its query from. +So the absorbed-q BMM must have **finished** before the RoPE call starts. +A decode assembled from the bf16 reading and then switched to an fp8 KV +cache builds its quantized query out of uninitialized memory, silently. +Both halves are certified in `catalog/attention/mla_rope_generation.md` +(*Signature*, the `fused_q` row). + +### 4.3 Cached-KV prefill + +``` +mla_rope_append_paged_kv_assign_q(q, latent, ...) # RoPE q/k_pe in place, + # append new tokens +ckv_full, k_pe_full = load_paged_kv_cache_for_mla(...) # gather [past+new] +kv = cublas_mm(ckv_full, W_kvb.t()); glue as in 4.1 +thop_attention(context_only, latent_cache=None: no in-kernel RoPE/append, + KV range > q range, bottom-right causal) -> output +``` + +### 4.4 Chunked prefill + +``` +mla_rope_append_paged_kv_assign_q(...) # as in 4.3 +for each chunk loop: # caller's plan + ckv_c, k_pe_c = load_chunked_kv_cache_for_mla(...) # one slice + kv_c = cublas_mm(...); glue + thop_attention(context_only, latent_cache=None, padding mask, + softmax_stats_tensor=stats_c) # partial pass + merge_chunked_attention_for_mla(merged, partial, ..., merge_op) +final causal pass over the new tokens (thop_attention) + final merge +``` + +The merge is exact: folding normalized-attention triples over disjoint KV +sets reproduces full-range attention (log-sum-exp rebase; proven in the +merge entry's contract and closed-loop tested against a single-pass fp32 +reference). + +### 4.5 Assembler-side obligations (post-onboarding knowledge) + +Everything in this subsection was established by certification runs; a +pre-onboarding decomposition report cannot contain it. + +Outside the entries, target code owns: the phase dispatch and chunk +planning (`cu_*` tensors, `merge_op` codes: first pass copy, later passes +merge, sit-outs skip); load-time weight preparation (`k_b_proj_trans +[H, C, nope]`, `v_b_proj [H, v, C]` derived from `kv_b_proj`); and buffer +allocation. Hard constraints established during certification: + +- **V must stay the split view of the `kv_b` output.** The context FMHA + kernel hard-codes V's row stride (`num_kv_heads * (head_size - 64 + + v_head_dim)`) and ignores the tensor's actual stride; a contiguous V + yields silent garbage and out-of-bounds reads. +- `mla_rope_append_paged_kv_assign_q` is **not idempotent** (in-place + self-reading rotation); re-prepare inputs before any retry. The append + and gather ops are idempotent. +- Generation takes three scheduler buffers + (`cu_q_seqlens`/`cu_kv_seqlens`/`fmha_scheduler_counter`), but only + `cu_q_seqlens` and the counter are checked — `cu_kv_seqlens` is neither + checked nor read on this path (the kernel takes each sequence's KV length + from `sequence_length`). Fill all three anyway; that is the production + shape and the certified one. The softmax scale is + `1 / (q_scaling * sqrt(nope + rope))` — not `sqrt(C + rope)`. +- Context and generation must be **separate calls** (`context_only` / + `generation_only`). The rejection of a mixed MLA batch lives *upstream*, + in the Python attention backend — the op itself accepts + `attention_input_type=0` without complaint, so a target driving the op + directly gets no error, just a batch in which at most one phase's rows + are right. + +## 5. Certification status (post-onboarding knowledge) + +All entries above are GPU-certified on sm_100 in the exact configurations +the assembly invokes, including the two producer/consumer seams that cross +entry boundaries: chunked partial passes were folded by the real merge op +in-test (stats interop proven), and the generation call consumes the +scheduler-buffer layout `mla_rope_generation` certifies producing. + +One deliberate remainder — the **end-to-end parity run** of the assembled +pipeline — was left to the first MLA target. **It has since been +executed** (deepseek-v3-lite-nvfp4/sm_100/tp1): against a pure-torch +reference built from HF `DeepseekV3Attention` semantics at that +checkpoint's geometry, **0.93–1.02 bf16 ulp** across 1-seq prefill, 2-seq +prefill, mixed-batch context and mixed-batch decode, identical at +`tokens_per_block` 32 and 64. As predicted, parity is tolerance-based +rather than bit-exact: the assembly legally substitutes implementations +(e.g. `cublas_mm` where the macro's `Linear` takes `F.linear`). + +What that target had to derive beyond this document — the +`rotary_cos_sin` byte layout, the `kv_b_proj` row regrouping the context +FMHA's `v` packing forces, `create_mla_outputs` being unreachable from a +target, and the cached-prefix context path that block reuse turns on by +default — is written up in `.claude/skills/staircase-modeling-core/references/latent-kv-cache.md`. + +## 6. Out of scope (current envelope) + +- **Quantization.** Everything is bf16: no fp8 weights (the + `fp8_batched_quantize_1x128_permute102` / `fp8_block_scaling_bmm_out` + family is not onboarded), no fp8/fp4 KV cache (every entry's + `quant_mode`/scale arguments are certified inert-only), no NVFP4 output. +- **DeepSeek-V4 and DSA sparse variants.** Their extra ops + (`deepseek_v4_q_norm*`, `mla_rope_inplace`, + `fused_inv_rope_fp8_quant_vllm_port`, `mla_dsa_proj` / + `mla_dsa_attn_inplace`) are not onboarded; the DSA pair are themselves + macros needing this same decomposition treatment first. +- **Blackwell cute_dsl alternates** (`cute_dsl_bf16_gemm/bmm_blackwell`) + are selection-dependent substitutes for `cublas_mm`/`bmm_out`; onboard + only if selected. +- **Parallelism and decoding features.** TP=1, beam width 1, no helix CP. + `predicted_tokens_per_seq > 1` — the multi-token query a linear-tree MTP + draft produces — **is** certified on the decode path now, for both + `mla_rope_generation` and `thop_attention`'s MLA generation call: each of + the `P` query rows takes its own rope position `L_g - P + t` and attends + causally, bottom-right aligned, to `[0, L_g - P + t]`. The two entries' + contracts carry the layout; `.claude/skills/staircase-modeling-speculator/references/multi-token-prediction.md` has + what the drafting layer computes. +- **Geometry.** Certified on DeepSeek-V3 per-head shapes + (`kv_lora=512`, `rope=64`, `nope=128`, `v=128`); the decode kernel is + source-listed for one other geometry (320/256), untested. +- **Operational notes.** `thop_attention`'s decode kernels are + NVRTC-JIT-compiled per process (pip must resolve `tensorrt_llm` at + runtime, so the run's interpreter needs pip present) with no on-disk cache — + exclude first calls from any timing. Several upstream schema quirks + (mutability annotations, misleading parameter names) are corrected + per-entry in the contracts. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/multi-head-ngram-embedding.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/multi-head-ngram-embedding.md new file mode 100644 index 000000000000..f7a60acfd11f --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decompositions/multi-head-ngram-embedding.md @@ -0,0 +1,319 @@ +# Decomposing `multi_head_ngram_embedding` (Qwen4-Exp PLE n-gram table read) + +Why the requested catalog entry `embedding/multi_head_ngram_embedding` was not +written, what the step actually is, and which two calls serve it. + +Derived on: tensorrt_llm 1.3.0rc26 | torch 2.12.0a0+5aff3928d8.nv26.05 | +NVIDIA GB300, `torch.cuda.get_device_capability()` `(10, 3)` → **sm_103** — the +stamp on the observations below, not a pin to assert. Receipts are keyed by GPU +architecture and nothing else, and each constituent carries its own in its own +contract; this document carries none. + +Checkpoint: `Qwen3.8-Flash-Next-NVFP4`, `model_type` `qwen4_exp`. + +## 1. Why decompose it + +The request was one entry for the step the run's `vocabulary.yaml` states as +"PLE multi-head n-gram embedding lookup with per-head offsets and per-tensor +FP8 dequantization". Written out at this checkpoint's widths, the step is + +``` +rows : [T, 16] int64 global row indices (produced upstream) +table : [320001536, 160] float8_e4m3fn (the n-gram table) +scale : scalar bfloat16, 0.00019931793212890625 (per-tensor dequant scale) + +ple[t, 160*h : 160*h + 160] = table[rows[t, h]] * scale for h in [0, 16) +ple : [T, 2560] bfloat16 +``` + +That is **two kernel-launching calls**, not one: a row gather, then an +elementwise dequantize. There is no fused gather-and-dequantize anywhere in +this build — the registered op set was swept exhaustively (294 `trtllm::` / +`tensorrt_llm::` schemas after forcing the lazy registrations by importing +`tensorrt_llm._torch.custom_ops`) and the only ops whose names match +`dequant|embed|gather|lookup|e4m3` are listed in §3; none of them indexes. + +Profiled directly, the two-call form launches **exactly two CUDA kernels**: + +``` +at::native::vectorized_gather_kernel<16, long> +tensorrt_llm::_v1::common::scaleMatrix<(QuantizeMode)1, false, __nv_bfloat16, __nv_bfloat16, __nv_fp8_e4m3> +``` + +Both are legal single entries under the catalog's admission rule, and one of +them is already in the catalog. Writing a wrapper that chains them would be +Python-level composition of two available entries — the case +`references/entry-spec.md` names as forbidden and `SKILL.md`'s red-flag table +calls "the exact anti-pattern the catalog bans". Hence: **macro; onboard +nothing**. + +Note the contrast with the sibling step, the n-gram *hash*, which this run did +land as a composed entry (`catalog/embedding/ngram_hash.py`). That one composes +roughly thirty launches out of integer ops for which the catalog has **no** +mirrors and the engine has **no** op, so composing was the only route to a +closed vocabulary and its contract says so. This step is the opposite case: +both constituents exist as atomic surfaces. + +### 1.1 The named tier-2 candidate computes a different function + +The dispatch sourced this step at tier 2 on +`tensorrt_llm._torch.modules.engram.engram.MultiHeadEmbedding`. That class +imports on this host and was exercised here; it does **not** serve, on four +independent axes, each measured rather than read: + +1. **It cannot hold the checkpoint's table.** Constructing it with + `dtype=torch.float8_e4m3fn` raises + `NotImplementedError: "normal_kernel_cpu" not implemented for 'Float8_e4m3fn'` + — it allocates an `nn.Embedding`, whose initializer is `normal_`. A bare + `torch.nn.Embedding(dtype=torch.float8_e4m3fn)` raises the same thing. Its + table is therefore a float parameter: at this checkpoint's row count a + bfloat16 `nn.Embedding` over 320001446 rows × 160 is **95.37 GiB**, against + **47.68 GiB** for the fp8 table the checkpoint ships. +2. **It re-applies the per-head offsets.** Its `forward` adds a registered + `offsets` buffer to the indices it is given. This run's hash step already + returns *global* rows (its contract: "The returned indices are global"), so + feeding them in double-applies. Measured on a 4-head instance: + `m(global_indices) == table[global_indices]` is **False**; + `m(local_indices) == table[global_indices]` is **True**. +3. **It has no dequantization at all** — no scale appears in its `forward`. +4. **It silently clamps out-of-range indices** to the last row (its `forward` + clamps to `num_embeddings - 1`), which turns a wrong index into a plausible + embedding. The raw gather this step needs instead **raises** (§4.3), and + losing that is a downgrade, not a convenience. + +It is also an `nn.Module` holding its own parameter (`embedding.weight`) rather +than a stateless callable, so there is no callable to wrap that takes the +checkpoint's table as an argument. + +### 1.2 Ruling R1 is not what excludes the tier-2 path here + +`workspace/rulings.md` R1 embargoes the callables under +`tensorrt_llm/_torch/modules/qwen4_exp/`, which is where this checkpoint's own +PLE implementation lives, and it was honoured: nothing under that directory was +opened, and no constituent below routes through it. But the tier-2 candidate +named in the dispatch (`engram.py`) is **not** embargoed, and it was read and +run. It fails on its own merits, per §1.1. The two are separate findings. + +## 2. Granularity and sourcing policy + +See `../decomposition.md`. Both constituents below land at preference level 1 +(a registered `torch.ops` op) or level 3 (a torch-native mirror already in +`catalog/torch/`); no policy exception is needed. + +## 3. The resulting vocabulary + +Entry paths are relative to +`tensorrt_llm/_torch/_experimental/modeling_v2/catalog/`. + +| # | role | callable | state | +|---|---|---|---| +| 1 | gather 16 table rows per token | `torch.nn.functional.embedding` | **existing entry**: `torch/embedding.py` | +| 2 | per-tensor e4m3 → bfloat16 dequantize | `torch.ops.tensorrt_llm.dequantize_e4m3_per_tensor` | **not onboarded**; proposed category `quantization/` | +| 3 | `[T, 16, 160] → [T, 2560]` (a view) | `torch.reshape` | **existing entry**: `torch/reshape.py` | + +Symbol resolution, checked on this host before citing: + +- `tensorrt_llm::dequantize_e4m3_per_tensor(Tensor input, Tensor scales) -> Tensor` + is present in the registered-schema list, and + `torch.ops.tensorrt_llm.dequantize_e4m3_per_tensor` **resolves**. +- `torch.ops.trtllm.dequantize_e4m3_per_tensor` does **not**: + `AttributeError: '_OpNamespace' 'trtllm' object has no attribute 'dequantize_e4m3_per_tensor'`. + The namespace is `tensorrt_llm`, exactly as for the catalog entry + `quantization/static_quantize_e4m3_per_tensor`, whose contract already names + this op as "the inverse, multiplying". +- `catalog/torch/embedding.py`, `catalog/torch/reshape.py` and + `catalog/embedding/ngram_hash.py` all import. +- `tensorrt_llm._torch.modules.engram.engram.MultiHeadEmbedding` imports (§1.1). + +Ops in this build whose names match `dequant|embed|gather|lookup|e4m3`, for the +record, so the next reader does not re-sweep: `tensorrt_llm::` carries +`dequantize_e4m3_{per_tensor,weight,activation}`, +`quantize_e4m3_{per_tensor,weight,activation}`, +`static_quantize_e4m3_{per_tensor,weight,activation}` and the `mxe4m3` host +pair; `trtllm::` carries `embedding_skip_forward`, `mxfp4_dequantize_unswizzled`, +`indexer_k_cache_gather_op`, `nvfp4_mla_kv_cache_gather`, +`nvfp4_mla_context_kv_cache_gather` and the gather-fused grouped GEMM. None +combines an index lookup with a dequantize. + +## 4. Assembly — the mainline path + +Packed-token form, which is what a chunked-prefill engine produces: + +``` +rows = ngram_hash(...) # [B, T, 16] int64, global rows +rows2 = reshape(rows, [-1, 16]) # [N, 16] torch/reshape +gathered = embedding(rows2, ngram_table) # [N, 16, 160] fp8 torch/embedding +deq = dequantize_e4m3_per_tensor(gathered, # [N, 16, 160] bf16 NOT ONBOARDED + scale_3d) +ple = reshape(deq, [N, 2560]) # [N, 2560] torch/reshape +``` + +`scale_3d` is the checkpoint's `weight_scale` broadcast to shape `[1, 1, 1]` — +see §4.2 for why its **rank** matters. + +Run end to end on this host at the checkpoint's own constants (a `[2, 130]` id +history, 128 query tokens per row, `eos_token_id` 248044, the three +`layer_multipliers`, the sixteen `ngram_heads_vocab_sizes` and +`ngram_heads_offsets`) against a real-size `[320001536, 160]` fp8 table: + +- row indices came out in `[95556, 319978804]`, all `< 320001446`; +- the result is `[256, 2560]` bfloat16 — `ple_embed_dim` exactly; +- it is **bitwise equal** to an independent per-row reference built from the + payload rows directly; +- head `h` occupies columns `[160h, 160h + 160)`, checked for all 16 heads + against an independent per-head gather. + +At `T` in {1, 8, 4096} over a smaller table the pipeline is likewise **bitwise +equal** to `(table[rows].to(torch.float32) * scale).to(torch.bfloat16)`. + +### 4.1 Where the table comes from (assembler-side, load time) + +The checkpoint stores the table as **128 separate tensors**, keys +`...ple_embedding.ngram_embedding.shard_{0..127}.weight` in +`model-fp8-mtp-ple.safetensors`, each `[2500012, 160]` `F8_E4M3`, plus one +`...ngram_embedding.weight_scale`, `BF16`, shape `[1]`. `config.json`'s +`split_ngram_parts` is 128 and `128 * 2500012 == 320001536`. + +The sharding is **contiguous chunks**, not interleaving — global row `r` lives +in shard `r // 2500012` at local row `r % 2500012`. Established from the +checkpoint alone, without reading any implementation: the sixteen +`ngram_heads_vocab_sizes` sum to 320001446, so rows 320001446…320001535 are +padding (`make_ngram_vocab_size_divisible_by` is 128). Scanning the last 200 +rows of shards 0, 1, 37, 38, 39, 126 and 127, the first six are fully non-zero +to their last byte, while shard 127's rows go from "all 160 bytes non-zero" to +"about half the bytes zero" at local row 2499922 — **global row 320001446**, +the first padding row, to the row. Interleaving would have put the boundary in +90 different shards, including 38 and 39. + +Consequences the assembler owns: + +- A loader that allocates one contiguous `[128, 2500012, 160]` (or + `[320001536, 160]`) fp8 buffer and fills it shard by shard produces a table + that `torch.nn.functional.embedding` takes directly with **global** indices: + `t3.reshape(-1, 160)` and the flat table were verified to be the same storage + and byte-identical. +- The **padding rows are not zero and are not meaningful** — they hold what + looks like uninitialized data. Nothing in this path range-checks against + 320001446; the hash step's own modulo is what keeps indices below it. +- `torch.nn.functional.embedding` requires a **2-D** weight + (`RuntimeError: 'weight' must be 2-D`), so the `[128, …]` form must be + reshaped — free, same storage. + +### 4.2 The scale argument (post-onboarding detail, measured here) + +`dequantize_e4m3_per_tensor` requires `scales.dim() == input.dim()` with every +`scales` dimension equal to 1; anything else raises +`RuntimeError: Expected scales.dim() == input.dim() to be true, but got false.` +or `Expected scales.size(i) == 1 to be true, but got false.`. So a rank-3 +gathered tensor needs a `[1, 1, 1]` scale and a rank-4 one needs `[1, 1, 1, 1]` +— confirmed both ways. + +The op **multiplies** (`out = input * scales[0]`, bitwise equal to +`input.to(torch.float32) * scale` with a float32 scale), and its **output dtype +follows the scales dtype**: float32 → float32, bfloat16 → bfloat16, float16 → +float16. Passing the checkpoint's bfloat16 `weight_scale` therefore yields +bfloat16 directly, with no cast anywhere in the path. + +The multiply direction is independently corroborated by the checkpoint: `448` +(e4m3's max finite) times 0.00019931793212890625 is 0.08929443359375, and a +50 000-row sample of shard 0 dequantizes to mean −1.84781e-05, std 0.00747576, +range [−0.0478363, 0.0510254] — ordinary embedding weights. Dividing would give +a std of about 37.5. + +### 4.3 Hard constraints a caller must respect (measured, this host) + +On constituent 1, `torch.nn.functional.embedding` over an fp8 table: + +- An out-of-range **or negative** row index is **not** silently clamped. It + trips a device-side assert: + `/opt/pytorch/pytorch/aten/src/ATen/native/cuda/IndexKernelUtils.cu:16: vectorized_gather_kernel: block: [0,0,0], thread: [0,0,0] Assertion 'ind >=0 && ind < ind_dim_size && "vectorized gather kernel index out of bounds"' failed.` + surfacing as `AcceleratorError: CUDA error: device-side assert triggered`. + This **poisons the CUDA context** for the rest of the process, so a test that + provokes it must do so in its own process. +- Index dtype must be **Long or Int**; Short and Byte raise + `Expected tensor for argument #1 'indices' to have one of the following scalar types: Long, Int`. + The checkpoint's largest legal row, 320001445, fits in int32. +- Index and table must share a device. A CUDA index against a pinned-CPU table + raises `RuntimeError: Expected all tensors to be on the same device, but got + index is on cuda:0, different from other tensors on cpu`. `pin_memory()` on an + fp8 CPU tensor works and a **CPU** index against a CPU table gathers fine — + so a host-offloaded table can be read on the host but not indexed from the + device by this call. Anything host-offloaded needs a different path; that is + outside both constituents. +- A `[320001536, 160]` fp8 table spans 51 200 245 760 bytes, past `2**31`; + gathers at rows 0, 2500011, 2500012, 20000002, 160000374, 300001275, + 320001444 and 320001445 were all bitwise correct. + +On constituent 2, `dequantize_e4m3_per_tensor` (these belong in its contract +once it is onboarded; recorded here so its smith does not re-derive them): + +| condition | result | +|---|---| +| rank < 2 input | `RuntimeError: Invalid dim. The dim of input should be greater than or equal to 2` | +| non-contiguous input | `RuntimeError: input must be contiguous` | +| zero-row input | `RuntimeError: input should not be empty tensor` | +| non-fp8 input | `RuntimeError: Invalid datatype. input must be Int8 (Fp8)` | +| float64 scales | `RuntimeError: Invalid datatype. input must be BF16/FP16/FP32` (the message says "input"; the offending tensor is `scales`) | +| CPU scales | `RuntimeError: scales must be a CUDA tensor` | +| CPU input | `NotImplementedError: Could not run 'tensorrt_llm::dequantize_e4m3_per_tensor' with arguments from the 'CPU' backend` | +| `scales == 0.0` | **accepted**, output all zeros, no signal | +| `scales < 0` | **accepted**, every sign flips | + +It is deterministic (two calls bitwise equal), leaves its input untouched, and +allocates its output. The gather output is contiguous and so is the dequantized +one, so the final reshape to `[N, 2560]` is a view — measured on the bfloat16 +result, whose `flatten(-2)` has the same `data_ptr`. + +### 4.4 CUDA graphs and cost + +The two-call pipeline **captures and replays** under `torch.cuda.CUDAGraph`, +bitwise equal to eager against fresh indices copied into the static input. + +Wall time per call, `[T, 16]` indices into a `2**23`-row fp8 table, mean of 50 +iterations on sm_103: + +| T | `embedding` + `dequantize_e4m3_per_tensor` | `embedding` + `.to(bfloat16)` + `* scale` | +|---|---|---| +| 8 | 0.0394 ms | 0.0721 ms | +| 512 | 0.0373 ms | 0.0716 ms | +| 4096 | 0.0469 ms | 0.1031 ms | + +Over the real `[320001536, 160]` table at T=4096 the pipeline measured +0.0485 ms. The trtllm op is the faster of the two dequantize routes at every +size measured because it is one kernel where torch needs two; this is a +measurement, not a gate. + +## 5. Certification status + +**Nothing here is certified.** No entry was created, no test was written and no +receipt was recorded. Every number above comes from probe scripts, not from a +catalog test, and the probes are not part of any entry. + +What the child smith for constituent 2 still owes: a contract, a GPU test at +`tests/unittest/_torch/modeling_v2/quantization/test_modeling_v2_dequantize_e4m3_per_tensor.py`, +and an `sm_103` receipt. Suggested coverage beyond §4.3: the whole 254-code +finite e4m3 domain against a native-torch reference (`x.to(torch.float32) * +scales[0]`), all three scales dtypes, ranks 2 through 4, and the boundary +`448 * scales[0]`. + +Constituent 1 (`torch/embedding.py`) is a hand-added mirror and by catalog +policy carries no contract, no test and no receipt; its fp8 behaviour is +therefore **unconstrained by design** and the observations in §4.3 live here +rather than in that entry. + +## 6. Out of scope + +- **Host offload.** This checkpoint's table is 47.68 GiB and the device here + has 275.9 GiB free, so it fits; whether the target keeps it resident is the + target's decision, and §4.3 records only that a device-side gather against a + CPU-resident table raises. +- **Tensor parallelism.** Single-rank only (`parallel: tp1`). Nothing above + shards the table across ranks. +- **The head-to-column order.** That head `h` occupies columns + `[160h, 160h + 160)` follows from `reshape` of `[T, 16, 160]` and from the + head order the hash step fixes (`[2-gram block | 3-gram block]`, 8 heads + each). It has **not** been checked against the trained `key_proj` / + `value_proj` weights; that is an assembly-time comparison against the + `TRTLLM_MODELING_V2=off` oracle, not something either constituent can settle. +- **A second PLE layer.** `ple_layer_ids` is `[2]`; one layer, one table, one + scale. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/entry-spec.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/entry-spec.md new file mode 100644 index 000000000000..87590f4fed99 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/entry-spec.md @@ -0,0 +1,118 @@ +# Entry spec — what a catalog entry is made of + +Shared by `staircase-kernel-onboard` (building a new entry) and +`staircase-kernel-verify` (extending an existing one's coverage). Both write +the same three pieces under the same discipline; only the starting point +differs. + +## Atomicity — what one entry is + +One entry = one op invocation at the Python level. + +- The wrapper body makes exactly one kernel-launching call into trtllm. Fusion is legal only inside that call: a fused kernel invoked once (silu-and-mul, allreduce+residual+rmsnorm) is one compliant entry. +- Python-level composition is forbidden. A wrapper that chains split → norm → matmul, or bundles qkv-proj + RoPE + SDPA + o-proj into one function, is not an entry — it hides exactly the bugs the catalog exists to expose. +- If inspection shows the requested op cannot be wrapped as one such call — its body is Python-level composition — it is a macro. Onboard nothing (no entry files, no index change); produce a decomposition report at `references/decompositions/.md` instead, following the checklist in `references/decomposition.md`, "Writing a decomposition report for a new macro". Never onboard two entries in one run. + +Category placement: by compute role, not dtype (an nvfp4 gemm lives in `gemm/`); fused ops by dominant cost (allreduce+norm lives in `comm/`). Categories: `gemm attention moe norm activation quantization comm ssm`. `torch/` is reserved for hand-added torch mirrors — never write there, even when the honest verdict is that a torch callable serves the step; report that and let the orchestrator place it, because a computation mirror widens what the surface is for. If no category fits, stop and report. + +## The three pieces + +An entry spans two trees. The contract and the wrapper live in +`tensorrt_llm/_torch/_experimental/modeling_v2/catalog//`; the GPU +test lives in `tests/unittest/_torch/modeling_v2//`, because that is +the tree this repo's CI collects from and an in-package test is on no list. +All three share one snake_case entry name. Name for the compute role; mirror +the trtllm op name when it is already descriptive. + +### `.md` — contract + +The entry's sole interface document. Whoever calls this op reads the contract and nothing else — no trtllm source — so every fact needed to call it correctly must be in here. When mentioning sibling or variant ops, state that they exist — never whether they are onboarded; catalog membership is index.yaml's fact alone. YAML frontmatter, then sections: + +```markdown +--- +receipts: + sm_: {status: passed, tests: } +--- + +# + +**Wraps** `torch.ops.trtllm.` (one call). + +## Semantics +What the single call computes, precisely — math where it disambiguates. +State the fusion boundary: what happens inside the call, what the caller +still owns (e.g. "rope is applied inside; qkv projection is not"). + +## Signature +The wrapper signature, plus a table of every argument and return: +shape, dtype, layout/contiguity, device. + +## Metadata consumed +Every runtime/attention-metadata field the op reads and what it must +contain. "None" if stateless. + +## Preconditions +Everything that must hold before the call: allocated buffers, prepared +metadata, weight formats, alignment, supported dtypes. A caller +violating none of these must get a correct result. + +## Notes +Arch quirks observed in source or during testing. +``` + +### `.py` — wrapper + +Mirror the style of the `catalog/torch/` entries: module docstring, typed signature, short function docstring. + +- Exactly one trtllm op invocation in the body. Glue that launches no kernel (e.g. building an `out` tensor for out-variant ops) is acceptable; tensor math is not. +- Guard asserts are allowed only when both hold: (a) the check is pure metadata (shape/stride/dtype/device), and (b) the violation was observed on this machine to fail silently (wrong results, silently ignored args) rather than raise. If the op validates its own inputs loudly, the wrapper stays assert-free. +- The default vocabulary is `torch.ops.trtllm.*`. Wrapping any other Python callable (e.g. a pybind binding) is allowed only when the run's request explicitly authorizes it; record the actual import path in the wrapper docstring, the contract, and the index `impl` field. +- The signature exposes the op honestly — no hidden defaults a contract reader would not expect. +- A first-call arch assertion is allowed, not required. + +### `test_modeling_v2_.py` — GPU test + +At `tests/unittest/_torch/modeling_v2//`, under that name always: it is what keeps the file from colliding with upstream's test of the same op, and what the CI selector matches. `test_*` functions, importing the wrapper by its absolute package path — no fixtures shared with other entries, and no imports of another catalog entry. + +- Assert `torch.cuda.is_available()` up front. No skip paths — a test that cannot run records no receipt. +- Build the reference on the spot from native torch ops (fp32 accumulation where the kernel accumulates in higher precision). Never build it from the wrapper, the op under test, or other catalog entries. +- Cover the shapes and dtypes the contract claims — include decode-like (few tokens) and prefill-like (many tokens) shapes where relevant. Seed all randomness. +- **Size persistent state — KV pages, recurrent and conv pools, any paged table — from the CACHE MANAGER, never from your batch.** These are two unrelated numbers. The batch is what the executor hands you; a pool's slot count is what the cache manager's capacity planning allocated, and the runtime picks it without reference to how many sequences are in flight. A test where `pool_slots == batch_size` is the one arrangement the runtime never produces, and it is worse than under-coverage: when the two are equal, "indexes the slot it was told to" and "indexes the row it happens to be" are the *same number*, so the test cannot see a slot-addressing defect at all. Take the width from a runtime-value dump if you have one and say so in the receipt; otherwise take one plainly too large for the batch and record that you chose it. **State the slot count in the certified column** — it is a column, not a free axis. One run certified nine entries this way, each sized to its own batch by a different smith over two days, and an audit found all nine outside the engine's real width of 461 at once. +- Compare with `torch.testing.assert_close` at default dtype-aware tolerances. Any loosening carries an inline comment justifying the number. Where the loosening is more than marginal — a chain whose intermediate is requantized, say — a justification in prose is not enough: **measure the discrimination.** Run the wrong variants too (swapped operand halves, un-swizzled scales, an omitted permutation) and record how far outside the loosened gate each lands. A gate 3x above the correct result and 25-500x below every wrong one is defensible; the same gate with no wrong-variant numbers is a guess. + + + +**When the surface you are certifying neighbours a known intermittent defect, prove the harness can see that defect before reporting a clean result.** Discrimination above answers "is my gate too loose"; this answers the different question "is my harness blind", and a clean sweep from a blind harness is worth nothing. Replay the documented arming sequence for the known defect in the same processes, confirm it fires, and only then report that the certified path ran through that state correctly. One run certified a workspace-backed collective this way — the neighbouring fused variant was measured wrong in 4 of 6 fresh processes by the same harness, and every call on the certified path that ran after it fired was still bitwise correct. Without the control, "no mismatches" would have been indistinguishable from "not looking". + +Proven patterns — apply where the op's shape calls for them, not mechanically: for in-place ops, lock both sides (out-of-bounds write check via buffer snapshot; bitwise assertions on regions contracted as untouched); for quantized outputs, dual assertion (tolerance derived from the format's constants plus a bit-exact-fraction hard gate); for domains the contract declares unsupported, negative tests asserting rejection. + +Stateful ops (the attention family) are tested the same way with minimal REAL state: construct the actual metadata and cache objects trtllm needs, never mocks of the op path. If minimal real state cannot be constructed in isolation, that is a symptom to report, not a license to mock. + +Collective ops (`comm/`) cannot be exercised in one process, so the test spawns `N` ranks itself — `torch.multiprocessing`, world size read from `CUDA_VISIBLE_DEVICES` — and the entry keeps its single runnable command. Three things this shape gets wrong easily: the reference must stay arithmetic (the expected value of an all-reduce is computed, never obtained from a second collective); every rank asserts, and a failing rank has to make the command exit non-zero rather than print into a lost stream; and the run needs a deadline, because a broken collective **hangs** rather than raising — a test that can wedge forever records no receipt and blocks the run that dispatched you. Any workspace or communicator the op requires is real state under the rule above. + +## Receipts + +Receipts live in the contract frontmatter and are test outcomes, not claims: + +- `sm_100: {status: passed, tests: 12}` — a filled-in example: the test passed on that arch, over that many cases. Receipts are keyed by **GPU architecture and nothing else**; there is no version pin in tree, because a target moves with the trunk. +- `status: failed` plus a short `note:` — wrapper and test are correct and the op demonstrably fails there. This is valuable knowledge; land the entry anyway. +- Key absent = unknown. Never pre-fill archs you did not run. +- A multi-rank entry adds `world_size: ` — the rank count the test actually ran at. Absent means 1, which is what every existing receipt asserts. Without it a four-rank receipt is indistinguishable from a single-GPU one, and rank count is exactly the axis such an op is uncertified across. + +Derive the arch key from the device the test actually ran on (`torch.cuda.get_device_capability()`, e.g. `(10, 0)` → `sm_100`). **A receipt is valid only against the files of its own entry**: it must post-date the last write to every one of them, and an entry spans two trees, so the check looks in both. A write to any of the three makes the receipt stale — never rewrite a receipt without rerunning the test. + +## index.yaml + +Append to `tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml`, +matching the existing shape: + +```yaml + - path: norm/rms_norm.py + impl: torch.ops.trtllm. + summary: "One-line compute-role summary" +``` diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/SKILL.md new file mode 100644 index 000000000000..b329832d929e --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/SKILL.md @@ -0,0 +1,90 @@ +--- +name: staircase-kernel-verify +description: Load when smith must check an existing catalog entry against the values a checkpoint will pass — confirm the receipt's arch still holds and still post-dates every file of the entry, and where the certified column does not cover a value, extend the test, rewrite the contract's not-certified prose from observed runs, and revalidate the receipt. +--- + +# staircase-kernel-verify + +One run, one entry, one surface. The entry already exists; the question is +whether what this checkpoint passes is inside what the receipt certifies. + +The spec is shared: read +`.claude/skills/staircase-kernel-onboard/references/entry-spec.md` (the three +pieces, receipt discipline, test recipe) and +`.claude/skills/staircase-kernel-onboard/references/contract-evidence.md` +(what makes a contract sentence defensible). This skill only says what +differs. + +## Two outcomes + +**Covered.** The contract's certified column already includes these values, +and the receipt's `sm_` key matches the architecture this run targets. +Change nothing. Report what you checked and against which receipt. + +That check is not a formality: receipts are three-state, and **absent means +unknown**. Receipts are keyed by GPU architecture and nothing else — there is +no version pin in tree, because a target moves with the trunk — so the two +things that can make a green receipt a claim about a run nobody made are the +wrong arch and a stale file. **A receipt is valid only if it post-dates the +last write to every file of its entry**, and an entry spans two trees (the +contract and wrapper in the package, the GPU test in the tests tree), so check +both, mechanically, against mtimes. All of that costs a file read and no GPU. + +What replaced the old import-time pin assert is not on this entry at all: +each target declares `REQUIRED_TRTLLM_OPS`, and +`tests/unittest/_torch/modeling_v2/test_modeling_v2_target_contract.py` +asserts every name in it exists — which turns an engine drifting under a +target from a wrong answer into a loud failure. + +**Not covered.** Extend it: + +1. **Probe first where the semantics are undocumented**, then freeze what the + probes pinned into test cases. Do not write the assertion before you know + what the op does at this value +2. **Extend the existing + `tests/unittest/_torch/modeling_v2//test_modeling_v2_.py`** + with the new coverage. Same rules as a new test: the reference is built on + the spot from native torch ops, never from the wrapper or another catalog + entry; randomness is seeded; `torch.testing.assert_close` at dtype-aware + tolerances, and any loosening carries an inline justification +3. **Rewrite the contract's not-certified annotations for this surface** into + certified semantics, from observed runs only. This is the part that gets + skipped and it is the part the next target reads — a receipt nobody can + interpret is worth little +4. **Rerun the entry's whole test file**, not just the new case, and record + the receipt from that run. A multi-rank entry records `world_size: `; + absent means 1 +5. Do not format or lint during iteration — the repo's `pre-commit` runs once + at delivery. Revalidate that the receipt post-dates the last write to every + file of the entry, in both trees, then commit pathspec-scoped: + +```bash +git add tensorrt_llm/_torch/_experimental/modeling_v2/catalog//.md \ + tests/unittest/_torch/modeling_v2//test_modeling_v2_.py \ + tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml +git commit -m "feat(catalog): certify at " +``` + +The wrapper usually stays untouched. If it does not — if covering this surface +needs a signature change — that is worth saying loudly in the report, because +every existing caller is affected. + +**The op may simply fail at this value.** That is a valid product: record +`status: failed` with a short note, land it, and report it. It is a real +constraint on every target that would have passed that value, and finding it +here is far cheaper than finding it in an accuracy gate. + +## When the surface neighbours a known intermittent defect + +`entry-spec.md` states the rule and the case that produced it: prove the +harness can *see* the defect before reporting a clean result. It applies with +full force here — extending coverage next to a known-flaky surface is exactly +where a blind harness returns a confident, meaningless green. + +## Report + +- entry path, the surface requested, and the verdict: already covered / newly + certified / failed +- the receipt as it now reads, and what the new coverage pinned +- **what remains uncertified** on this entry — the next caller reads this +- every claim matched against disk before sending diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-bounded-replay/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-bounded-replay/SKILL.md new file mode 100644 index 000000000000..a58e85aa332b --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-bounded-replay/SKILL.md @@ -0,0 +1,245 @@ +--- +name: staircase-modeling-bounded-replay +description: Load when adding DeepSeek-V4.1's SWA Bounded Replay to a target — `decoder_replay` (the layers after the last KV-source layer run over only the last `window` prompt tokens) or `encoder_replay` (a prefix-cache hit rebuilds sliding-window KV by re-forwarding the last `window` cached tokens). Deliberately not output-identical, and gated by a work witness, because a unit that implements nothing passes accuracy too. +--- + +# staircase-modeling-bounded-replay + +`--features decoder_replay` and `--features encoder_replay`. Both come from +DeepSeek-V4.1-Flash's tech report (§3.2.2, "SWA Bounded Replay"), where the +model is a **Causal Encoder-Decoder**: layers 0-19 are the "encoder", layers +20-39 the "decoder", and decoder global KV is projected from the encoder's +final hidden state rather than from each decoder layer's own. "Encoder" here +is those causal layers — **never the vision tower** — and "replay" is neither +CUDA-graph replay nor recurrent-state replay. + +## The definitions, verbatim + +TR §3.2.2: + +> "SWA Bounded Replay instead replays only the most recent n_win tokens and +> truncates SWA to the replay segment, accepting approximate states: for a +> replay starting at position r, a query at position i attends to SWA keys in +> [max(r, i−n_win+1), i]." + +> "**Encoder SWA Bounded Replay** makes prefix caching depend only on global +> KV, allowing SWA KV to be removed from the persistent KV cache. When the +> encoder SWA KV is missing, we replay the last n_win tokens of the cached +> prefix and process them together with the uncached suffix. The replayed +> tokens regenerate only SWA KV, reusing the cached global KV without +> recomputation or overwriting, while the uncached suffix generates both +> global KV and SWA KV. By design, the replayed prefix state is approximate…" + +> "**Decoder SWA Bounded Replay** bounds the decoder forward pass to n_win +> tokens, nearly halving total prefill computation. Under CED, decoder global +> KV is projected from the final encoder hidden states. The only obstacle to +> ending prefill at the encoder is decoder SWA KV … at every prefill, we +> replay the last n_win tokens of the prompt, feed their encoder outputs +> through the decoder layers under the same SWA truncation, and use the +> resulting decoder SWA KV only for decoding, not for prefix caching. … not +> mathematically equivalent … we additionally simulate the same replay during +> post-training for train-aware adaptation." + +`n_win` is the checkpoint's sliding window: 128 for DeepSeek-V4.1. + +## The property that shapes the gates + +**Neither feature is output-identical.** Both truncate sliding-window state on +purpose, so the accuracy gate is a real check of them — unlike host offload or +disaggregation, where it is blind. + +It is not a sufficient check. The truncation was designed, and post-trained, +to "barely compromise" accuracy. So a unit that implements **nothing** lands +inside the same tolerance, reports green, and ships a target that does the +full prefill while its Coverage section claims the saving. The only +observable that differs is **work**: rows the decoder ran, rows a prefix hit +re-forwarded. Only the target can count those, so it prints them, and the +**replay witness** judges them against a window and a workload the spec +declared before assembly. See **Gates**. + +## decoder_replay — the target owns all of it + +Target-local: no runtime component, no scheduler change. + +- **Which layers are bounded.** The ones after the last KV-source layer + (`kv_source_layer_ids` in the config; its last entry is 20 on + DeepSeek-V4.1). The boundary is a decision, not a fact: the tech report's + decoder starts at layer 20, while SGLang's implementation computes + `late_layer_start = max(kv_source_layer_ids) + 1` and runs layer 20 over + every prompt token, because layer 20's compressor is what produces the + decoder's global KV. Pick one, say why, and record it — the plan's open + risks carry it. +- **Context phase.** Every layer up to the boundary — including the last + KV-source layer's compressor and index-K projection — runs over all prompt + tokens. The bounded layers run over only the last `min(N, window)` prompt + tokens of each context request, with SWA truncated at + `r = N - min(N, window)`. Only the last token feeds the head. Main and + indexer KV stay exact; only decoder SWA KV is approximate. +- **Only in the final context chunk.** Under chunked prefill the bounded + layers do nothing in earlier chunks and run once, in the last. The last + `window` prompt tokens can straddle chunks, so the hidden state entering the + first bounded layer must be **retained per request across chunks** for the + tail — a buffer the target owns, keyed by request, freed at the final chunk. + Dropping it silently replays fewer than `window` tokens and still passes + accuracy. +- **What it cannot return.** The bounded layers never computed rows before the + tail, so prompt logprobs, context logits and full hidden-state capture do + not exist under this feature. SGLang refuses them; the target must fail fast + rather than return rows it did not compute. +- The decoder SWA KV the replay produces serves decoding only — it is not + something prefix caching may reuse. + +SGLang's help text for its flag states the contract compactly: "after the last +kv_source layer, run the remaining layers over only the last window_size +tokens of a prefill. Main and indexer KV stay exact; nothing is replayed. +Deterministic for a fixed prompt and chunk size." + +## encoder_replay — runtime and target, and the runtime half is new + +| responsibility | owner | +|---|---| +| the prefix cache keeps only global KV (compressed KV + indexer K) for a cached prefix; SWA KV is not held by it | runtime | +| on a prefix hit, schedule the last `min(window, reused)` cached tokens as **replay** rows ahead of the uncached suffix, charged against the token budget | runtime | +| a reuse boundary aligned to the compress ratio — even, for ratio 2 — or backed off to one | runtime | +| replay rows write **SWA KV only**: compressor skipped, indexer K not written, global KV never overwritten | target | +| replay rows produce no logits and are never sampled | target | +| the uncached suffix writes both global and SWA KV, as in an ordinary prefill | target | + +An odd boundary on a ratio-2 layer splits a compressed pair: one half is +cached and the other was never computed, and the result is a quietly wrong +global entry rather than an error. SGLang refuses the combination; the witness +checks it when the spec declares `reuse_alignment: 2`. + +**TRT has none of the runtime half today.** `DeepseekV4CacheManager` in +`tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py` +keeps SWA as a sliding pool whose blocks reuse depends on, and does not back +off a reuse match (`_supports_reuse_match_backoff` is false there). A target +cannot own the scheduler, so name that work in the plan as runtime work, with +its owner — a target that fakes it by re-forwarding the whole prefix on a hit +is exactly what the witness's accounting check catches. + +SGLang's implementation (branch `origin/dsv4.1`, +`python/sglang/srt/model_executor/encoder_swa_replay.py`) is the reference for +the mechanics: it re-runs the cached tokens `[end - window, end)` through the +layers with the compressor skipped, keeps SWA in a per-request ring, and +charges at most `window` replay tokens per request to the scheduler. + +## Interaction rules, and saying which combinations work + +From SGLang's argument checks for this model +(`python/sglang/srt/arg_groups/deepseek_v4_hook.py`, branch `origin/dsv4.1`): + +| feature | refused together with | +|---|---| +| `encoder_replay` | prefill CUDA graphs, attention DP, PD disaggregation, LoRA | +| `decoder_replay` | prefill CUDA graphs, attention DP | + +**No implementation anywhere combines bounded replay with attention DP.** A +`dep` member of a unified target is therefore unexplored ground, not a +configuration with a known answer. + +The target must **say** which combinations it supports: + +- a fail-fast assert for each combination it does not, naming the reason — a + combination that silently runs without replay is the silent wrong this + package exists to prevent, and it passes accuracy; +- the spec's `members:` narrowed for the unit to the members it supports, so + the gate loop does not ask for the rest: + `members: {default: [tp4, tep4, dep4], decoder_replay: [tp4, tep4]}`; +- and the same two statements in the target's Coverage section. + +"Prefill CUDA graphs" means a captured **context** forward. TRT captures +generation steps by default and runs context eagerly, so the conflict appears +only where context forwards are captured (piecewise capture, compile). The +replay row selection is data-dependent; keep it out of any capture. + +## Gates + +boot → accuracy → **replay witness**, all under `TRTLLM_MODELING_V2=require`. + +**The switch.** Each feature is turned off by an environment variable the +target reads at construction — for example +`TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY=0` — and the spec names it as the +witness's `contrast_env`. Off must mean the full computation, not an error: +the witness runs both. + +**The counter.** When `TRTLLM_MODELING_V2_WITNESS=1` is set, the target prints +one line per context request, when that request's context phase completes +(after its final chunk): + +``` +MODELING_V2_WITNESS prompt_tokens=4103 reused_tokens=0 replayed_tokens=0 forwarded_tokens=4103 decoder_tokens=128 +``` + +Space-separated `key=`, on one line; anything may precede the marker. +Every field is counted for that request over its whole context phase: + +| field | what it counts | +|---|---| +| `prompt_tokens` | the request's prompt length | +| `reused_tokens` | prompt tokens served from the prefix cache | +| `replayed_tokens` | cached tokens forwarded again in replay mode | +| `forwarded_tokens` | rows the model's first layer processed | +| `decoder_tokens` | rows the first bounded layer processed | + +`decoder_replay` reads `forwarded_tokens` and `decoder_tokens`; +`encoder_replay` reads the other four. Print them all when the target +implements both. + +**Count rows the layer actually ran**, from the shape of the tensor it ran +over, accumulated per request id across chunks — not a formula of the prompt +length. A counter that prints `min(prompt, window)` passes the ON leg and is +caught by the OFF leg, which is the leg's purpose; do not make the OFF leg +the only thing standing between a formula and a pass. On a multi-rank run, +printing from every rank is fine: each verdict is per line. + +**The verdicts**, against the spec's declared values: + +| feature | replay ON | replay OFF (`contrast_env`) | +|---|---|---| +| `decoder_replay` | `decoder_tokens <= window` for every request, and the long prompt was forwarded cold (`forwarded_tokens >= prompt_tokens`) | `decoder_tokens == forwarded_tokens` for every request | +| `encoder_replay` | `replayed_tokens == min(window, reused_tokens)` and `forwarded_tokens == prompt_tokens - reused_tokens + replayed_tokens` for every request; at least one hit with `reused_tokens > window`; `reused_tokens` a multiple of `reuse_alignment` | `replayed_tokens == 0`, and the same accounting identity | + +No counter line at all is a FAIL that says the target printed no witness +counter. The witness builds its own workload from the declared sizes — one +long prompt for `decoder_replay`; two prompts sharing a long prefix, with +block reuse on, for `encoder_replay` — so the boot spec's short prompts, which +never exceed the window, are not used. + +**The spec**, declared before assembly and never derived from the run: + +```yaml +features: [decoder_replay, encoder_replay] +accuracy: + witness: + decoder_replay: + contrast_env: {TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY: "0"} + window: 128 + prompt_tokens: 4096 # > window; at most 16384 + encoder_replay: + contrast_env: {TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY: "0"} + window: 128 + prefix_tokens: 4096 # > window; at most 16384 + reuse_alignment: 2 # the compress ratio +``` + +`runtime.max_seq_len` (through `extra_flags`) must admit the longest witness +prompt, and the `encoder_replay` unit's runtime must keep block reuse on — the +witness refuses to launch otherwise, since no request could hit the cache. + +## What the gates still do not tell you + +That prefill got faster. The witness says rows were skipped; whether a 32K +prompt now prefills in about half the time is a perf measurement for tuning. +And that the approximation is the post-trained one: accuracy within tolerance +is the evidence for that, which is why this feature's accuracy gate is not a +formality. + +## Report additions + +- the layer boundary chosen for `decoder_replay`, and why +- which combinations the target supports, the asserts that refuse the rest, + and the member set the unit was gated at +- for `encoder_replay`, where the runtime half lives and who owns it +- the witness lines' ON and OFF values for the longest request, verbatim +- accuracy with replay against the anchor, and which anchor that was diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/SKILL.md new file mode 100644 index 000000000000..272ad3abd48a --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/SKILL.md @@ -0,0 +1,229 @@ +--- +name: staircase-modeling-core +description: Load whenever a target's modeling.py or weights.py is being written or edited — the invariants any forward must hold, the products a target consists of, and the from-scratch assembly flow. assembler loads it for core assembly and for every feature increment; tuner loads it before editing a forward. +--- + +# staircase-modeling-core + +Two halves, and most readers need only the first. + +| you are | read | +|---|---| +| assembler, core assembly | everything | +| assembler, a feature increment | **Invariants** — you are editing a file that already holds them | +| tuner, editing a forward | **Invariants**, especially `references/graph-and-sync.md` | + +## Which reference, when + +`references/` is read by need, not front to back: + +| file | read it when | +|---|---| +| `runtime-interface.md` | always, before writing — registration, shell/model split, weight lifecycle, manifest convention, forward duties, fail-fast ladder, serving-time costs | +| `graph-and-sync.md` | **always, before touching a forward** — what a replayed forward may contain, and which patterns stall the host | +| `routing-identity.md` | placing a new target, or asking whether some quantity may select one | +| `step-args.md` | building the attention call's arguments | +| `adapter-rules.md` | this target has a feature that talks to a runtime-owned service | +| `latent-kv-cache.md` | the checkpoint's `config.json` says MLA | +| `hybrid-recurrent-state.md` | `layer_types` mixes `linear_attention` with `full_attention` — the recurrent layers need a state pool the runtime allocates only when the published config matches its predicate | +| `expert-weight-packing.md` | sparse MoE | +| `attention-sinks.md` · `sliding-window-attention.md` | the config says so | + +The mechanism files are not optional reading dressed up as reference: each one +records what an earlier target had to establish by probing. Reading the right +one is the difference between an afternoon and a week. + +## Invariants + +These hold for every edit to a target's modeling code, by anyone, forever. + +1. **Closed vocabulary.** Inside `forward` and the private methods it reaches, + every call that creates or transforms a tensor is a catalog entry. Only + tensor-metadata reads and Python control flow live outside. Scope is the + whole trick — a whole-file scan false-positives on the rope table's + `arange`/`cos`/`sin` and on load-time `.t()`/`.to()`, which run at init. +2. **The static op declaration lists exactly what the forward calls.** + `REQUIRED_TRTLLM_OPS` names every `torch.ops.trtllm.*` symbol the forward + reaches, and + `tests/unittest/_torch/modeling_v2/test_modeling_v2_target_contract.py` + asserts that this build registers each one (plus the pybind entry point the + attention op is reached through). The list must *equal* the actual call + set — a stale entry is as wrong as a missing one. With no version pin, that + declaration is the only thing standing between a drifted engine and a wrong + answer: it turns the drift into a loud failure instead. +3. **Every axis the assembly depends on carries an assert**, read from + `pretrained_config`. An assert is how a reused pattern refuses a + checkpoint it does not fit. +4. **Parameters are declared flat**, `torch.empty` only, HF `[out, in]` + storage. Derivations happen in the post-load path, never at declaration. +5. **The manifest asserts coverage in both directions**: every declared + parameter is filled, and every checkpoint key is consumed — with one + explicit escape, the **known-not-loaded key set**, for weights this target + deliberately does not implement (an unimplemented feature's module, an + off-window expert range). Anything in that set is named in `weights.py` + beside the predicted-leftover assert, with the reason, and reported in the + run journal. Without the escape the assert fires on a correct target; with + it used loosely, a real miss hides. +6. **Graph safety and sync-freedom**, per `references/graph-and-sync.md`. + These are writing rules, not tuning rules: a forward that violates them is + wrong the first time it is captured. +7. **The forward is flat and serves any batch composition.** One code path, + branching only on metadata and per-capture constants. +8. **Fail loudly on what the target does not implement.** The fail-fast ladder + is the difference between an unsupported input and a wrong answer. +9. **Self-containment.** No shared helpers across targets. Duplication between + targets is the architecture, not a smell. + +## Products + +A target directory holds **two files**: + +| product | | +|---|---| +| `modeling.py` | the `REQUIRED_TRTLLM_OPS` declaration · geometry with asserts · the flat `ParameterDict` · the step-args builder with its per-entry CUDA-graph classification · the forward · the runtime-value dump (below) · loud asserts for the unimplemented · the `@register_auto_model()` shell | +| `weights.py` | the MANIFEST data table and load loop with bidirectional coverage asserts, and the predicted non-load sets with their reasons; post-load derivations live in modeling's `post_load_weights` | + +Four products the standalone layout had are gone. Each one's *responsibility* +moved rather than disappearing, and knowing where is how you avoid recreating +it: + +| gone | where the responsibility went | +|---|---| +| `model_dir/` (stub `config.json` + symlinks) | nowhere — **the checkpoint is read exactly as published**, with no target-owned `config.json`. Nothing about the model is the target's to restate, so nothing can disagree | +| `smoke.py` | the **boot gate**: `examples/llm-api/quickstart_advanced.py` with the target's topology flags and `TRTLLM_MODELING_V2=require`, run by the orchestrator. You author its keyword cases; you do not author its runner | +| `llm_args.yaml` | the topology comes from the **routing decision** and the gate's CLI flags. There is no file for it, which is also why a tuning knob can no longer be smuggled in beside it | +| `TARGET.md` | **gate records go in the run journal** (workspace-scoped, uncommitted). Identity lives in the target path plus `models//routing.py` | + +And the thing that used to decide which class name a target answered to — +`@register_auto_model("StaircaseForCausalLM")`, and the whole discussion of +when a target must register under some other name instead — is **replaced by +routing**. The checkpoint's own `architectures[0]` stays what it always was; +`modeling_v2_resolve` rewrites it into a synthetic registry key +(`ModelingV2GptOss120bSm103Tp1`) that the target's module registers, and the +rewrite is local to `_resolve_class` — every other runtime branch keyed on the +architecture string still sees the published name. So a target neither picks +nor needs to pick that name: `models//routing.py` does, the class name +carries the same triple as the path, and `test_modeling_v2_claims.py` asserts +they agree. Read `references/routing-identity.md` before adding a target or +proposing a new criterion. + +**Registration goes at module scope** in `modeling.py`, which is what importing +the module does. The routing module's `TARGET_MODULES` entry names that module +by dotted path, and `modeling_v2_resolve` imports it once the tree has claimed +the config — nothing else would put the class behind the synthetic key. Above +world size 1 the engine is built in spawned ranks that import the module +themselves, and the class travels to them by pickle, which resolves it by +module name; an ordinary dotted module satisfies that without any help. + +## The runtime-value dump + +`forward` records the **metadata** of every catalog call it makes — shapes, +dtypes, the constants passed — into a set, and the boot gate's output carries +it out of the run. + +This exists because of a class of value that nothing else surfaces: the +shapes the engine **varies at runtime**. Decode batch sizes, page sizes, +feature flags. They are written nowhere in the target, so reading the forward +does not reveal them, and a target once merged with 34 of the 35 batch sizes +its engine captured sitting outside the contract — every *constant* it passed +was inside. + +Four rules make it safe and sufficient: + +- **metadata only.** Never a tensor value: `.item()`/`.cpu()` are forbidden in + a forward, and metadata is exactly what a certified column is written in +- **no allocation in the hot path.** A set of tuples, appended to; nothing + that grows per token +- **capture is enough.** CUDA graph capture *executes* Python — that is why + host-derived scalars get baked in — so the dump naturally covers every + captured batch size. Replay executes none of it, so the steady state pays + nothing +- **off unless asked for.** Record and print only when + `TRTLLM_MODELING_V2_VALUE_DUMP=1` (read once, at import). Every gate and + witness run sets it; a served engine must not print the dump at each exit + +Comparing the dump against the contracts is a *reading* job, not an automated +one: a certified column is prose. The target's duty ends at exposing the +values; assembler judges them, and smith is dispatched only for what is +genuinely uncovered. + +## From-scratch flow + +1. **Derive before writing.** Take the computation-step list from the + workspace's `vocabulary.yaml` — it already carries this checkpoint's value + on every axis, established in Phase B. Read the checkpoint's safetensors key + list and the HF reference for the compute graph, and decide each step's + mapping onto catalog entries. An unmappable step is a vocabulary gap: report + it before creating files. +2. **Place the target and its routing.** Create + `tensorrt_llm/_torch/_experimental/modeling_v2/models//____/`, + and add the target to that family's `routing.py`: the config-shape + fingerprint, the parallel segment, the synthetic class name and its + `TARGET_MODULES` entry. A family nothing routes yet also needs one row in + `_router_index.py`. `references/routing-identity.md` is what may decide any + of this. +3. **Write `modeling.py` + `weights.py`**, and draft the boot gate's keyword + cases. Keywords stay provisional. +4. **Inner loop — edit → import dry-run → boot gate, until the gate passes.** + The dry-run imports the target module and asserts the class registers under + the synthetic name, and that `explain.py` routes this configuration to it. + **Do not lint between edits.** The run lints once, at delivery, with the + repo's `pre-commit`; during iteration spend the turns on the model, not on + whitespace. **The first engine run doubles as keyword verification** — + freeze the cases only after observing the continuations. +5. **Check the dump** against the contracts. Anything outside a certified + column: dispatch `Agent(smith)` with `kernel-verify`, resume on the receipt. +6. **Accuracy gate.** A shortfall is a debug signal, not a verdict: the + reference entry pins both the anchor and the protocol, so a miss usually + means a real defect and you hold the derivation table to find it. Two + rules — every re-gate is backed by one explicit causal hypothesis (the full + set costs about an hour and same-protocol rerun noise is 0.1–0.3 points, so + a hypothesis-free rerun is expensive and information-free), and after three + consecutive hypothesis-backed fixes without improvement, stop and report + the shortfall with the audit trail. The anchor and tolerance are read-only: + you fix your assembly, never the bar. +7. **Fill the run journal** from observed runs only. +8. **Lint once with the repo's `pre-commit`, then commit** your own files, by + explicit path. + +## Feature increments + +Same file, same invariants, plus: + +- the feature's own skill says where it attaches and what it owes +- **the base path stays exactly as it was.** If the feature is off, the forward + must be what the recorded gate measured — bit for bit +- the weights the feature turns on come *out* of the known-not-loaded key set, + and `weights.py`'s predicted non-load set says so +- the run journal records what the feature needs to run at all: its activation + config and any boot requirement you discovered while gating it, as the gate + ran with them. Neither is a tuning knob, and tuner may not change them +- gate it: boot with the config → accuracy, plus a third gate **only if** + the feature's skill defines one. A third gate is earned by exactly one + property — a miscomputed implementation still producing a correct answer — + and today only speculative decoding has it. A feature that is merely inert + produces correct output, which makes that a tuning question + +## Report + +- target path; checkpoint path and hash summary; the GPU used. Routing + recognizes a checkpoint by its shape and cannot tell a fine-tune from the + original, so the record of which bytes you measured is what pins the claim +- **derivation table**: each axis and what this model required. Flag every axis + where this target differs from the exemplar you read +- **vocabulary audit**: the forward's call list, each mapped to its + `tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml` path. + Scope it to `forward` plus the private methods it reaches +- **dump review**: which values the engine produced at runtime, and where each + sits against its contract +- **gate outcomes**: boot-gate cases and verdicts; accuracy measured versus + reference and threshold — verbatim numbers, or the precise blocker +- **symptoms**: knowledge gaps (the exact question your sources could not + answer), vocabulary gaps, upstream oddities. State separately anything you + had to establish by observation, naming the mechanism it belongs to — that + is what gets written back so the next target does not repeat the probe +- **signal versus noise**: same-protocol reruns move 0.1–0.3 points, so a + sub-0.5-point change is not a result. Never call a construct load-bearing on + that evidence; downstream work takes the claim at face value and pays to + disprove it +- files created; lint status diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/adapter-rules.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/adapter-rules.md new file mode 100644 index 000000000000..2be7d73a2985 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/adapter-rules.md @@ -0,0 +1,84 @@ +# Adapter rules — a feature that talks to a runtime-owned service + +Most features are pure modeling: they change what the forward computes, so +they are written flat into `modeling.py` and there is nothing to adapt. This +file is for the other kind — a feature whose runtime side owns state that +outlives one forward call (expert placement and migration; a composite weight +view with remote prefetch). Those need a named boundary, and the boundary is a +file. + +The split is the same every time: + +| layer | owns | +|---|---| +| **runtime** | the service: its threads, its communicators, its state machine, its memory | +| **adapter** | the boundary — registration, the handles, the facade the runtime traverses | +| **modeling code** | the calls, at the points in the forward where they belong | + +## Rules + +**One file, one import site.** `_adapter.py` is the *only* file in +the target that imports the private runtime surface. `modeling.py` and +`weights.py` see the adapter's own types and nothing else. This is what makes +the blast radius of an upstream change readable: one file to re-derive. A +target directory otherwise holds `modeling.py` and `weights.py` only — this +file is the single exception a feature of this class earns. + +**The adapter does no work of its own.** It creates no stream, issues no copy, +allocates no tensor, launches no kernel, and does not reach into the service's +internals to do by hand what the service does. It hands over geometry, +placement and closures, and relays calls. An adapter that starts computing has +absorbed a piece of the runtime, and the next upstream change will break it +silently. + +**Pin the private surface by digest.** A public API has a version contract; a +private one does not, and in tree there is no version pin to fall back on — +the target moves with the trunk. Record the SHA-256 of every upstream file the +adapter binds against in the run journal, beside the gate record. A change to +any of them re-opens every question the adapter answers, and the digest is what +turns that from a silent drift into a visible diff. + +**Naming:** + +``` +_adapter.py the file +Adapter the public class — the entire outward surface +_ private helpers: Facade, Stub, … +``` + +Do not name a helper `*Shell`. In this repo "shell" already means the target's +registered `DecoderModelForCausalLM` subclass. + +## What the feature's own skill must document + +Two files, both required for an adapter-class feature: + +- `runtime-boundary.md` — which layer owns what, in the table above's shape. + Concretely enough that a reader can tell whether a given piece of state is + ours +- `adapter-interface.md` — what the adapter exposes: every method and property + the runtime calls or reads, and **when** each is called. Timing is half the + content; a hook called at the wrong point in the forward is a silent wrong + answer, not a crash + +## Gating + +An adapter-class feature does **not** automatically owe a third gate. Ask the +only question that earns one: **can a miscomputed version of this still +produce a correct answer?** For a feature that routes to different weights, or +binds a different expert set, the answer is no — accuracy sees it, and the +gate ladder is boot → accuracy like everything else. + +What these features *can* hide is being **inert**: registered, called, and +doing nothing useful. That produces correct output, so it is a performance +question, and performance here is measured rather than gated — it belongs to +the later tuning workflow. + +Two cheap things do belong in onboarding, and they are not gates: + +- **a static variant first**, where the feature has one. Green on static, then + green on dynamic, makes a failure attributable to migration rather than to + wiring +- **the boot gate asserting the launch preconditions** the feature needs but + the runtime does not check. Turning a silent hang into a named failure is + what the minutes-scale gate is for diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/attention-sinks.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/attention-sinks.md new file mode 100644 index 000000000000..67b6c3155578 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/attention-sinks.md @@ -0,0 +1,56 @@ +# Attention sinks + +A learned per-query-head logit that joins the attention softmax's +denominator and is then dropped from the output, so each row's weights +sum to less than 1. Written from the gpt-oss-120b/sm_100/tp1 onboard +(`self_attn.sinks`, one `[num_heads]` tensor per layer); the mechanism is +not specific to that family. + +## What it computes + +``` +out[i, h] = softmax(concat(logits[i, h, :], sink[h]))[:-1] @ V +mass[i, h] = sigmoid(logsumexp(logits[i, h, :]) - sink[h]) # the row's weight sum, < 1 +``` + +The sink is one extra column in the softmax, never a value contribution. +A large sink drains the row toward zero output; a sufficiently negative +one (`-100.0` or `-inf` measured) is bitwise identical to passing no sink +at all, which makes the inert path easy to verify. + +## The two facts a target gets wrong by guessing + +**The sink lives in the scaled-score domain.** It is compared against +`q @ K.T / (q_scaling * sqrt(D))`, not the raw dot products. At `D = 64` +the scale is 1/8, so the two readings are far apart — the wrong one +measured 12.8x to 41x outside the tolerance band on every case. Do not +pre-scale the checkpoint's sink values. + +**It must hold in both phases, and a kernel that honours only one fails +silently.** Context prefill and generation decode take different FMHA +kernels, and long-history decode takes a third variant that folds partial +softmax states across CTAs (`...MultiCtasKvCga...`) and has to apply the +sink inside that reduction. Certify each separately, and gate every case +against a "sink ignored" reference as well as a correct one — otherwise a +passing test proves only that the numbers are plausible. + +## Handling requirements + +- **fp32**, whatever the checkpoint says. gpt-oss stores `sinks` as bf16; + the op raises on anything but float. Promote at load, alongside the + expert biases. +- **Contiguous, exactly `num_heads` elements.** The op validates the + dtype and nothing else: a short tensor makes the kernel read past the + end of the allocation, a longer one has its tail ignored, a strided + view whose *values* are right still reads the interleaved memory and + returns wrong attention, and an empty tensor silently disables the + sink. Shape beyond the element count is irrelevant (`[H]`, `[1, H]`, + `[H, 1]` all behave identically). Pure metadata, so the wrapper carries + the guard the op lacks — `thop_attention.py` asserts contiguity and + `numel() == num_heads`. +- **Per query head**, not per kv head. A checkerboard sink (large on even + heads, very negative on odd) drains exactly the even head columns while + leaving the odd ones bitwise identical to the no-sink call — that is + the cheapest way to pin the indexing. +- The sink perturbs nothing but the denominator: the paged KV append + stays bit-exact and no cache bytes move. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/expert-weight-packing.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/expert-weight-packing.md new file mode 100644 index 000000000000..8f231166471a --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/expert-weight-packing.md @@ -0,0 +1,574 @@ +# Expert weight packing (sparse MoE) + +What a target assembled on a sparse-MoE checkpoint had to establish by +observation, because no contract or reference stated it. Written from the +qwen3-30b-a3b/sm_100/tp1 onboard (Qwen3-30B-A3B: 48 layers, all MoE, 128 +experts, top-8, `moe_intermediate_size` 768, bf16). The mechanism — a +router selecting k of E stacked expert MLPs — recurs across model +families; this file is about the mechanism, not that checkpoint. + +**Every shape below is a whole-stack shape, measured at world size 1.** +Under expert parallelism a rank holds `E / ep_size` of the stack and the +first dimension of every weight and scale shrinks accordingly. + +**What survives the split — established by the first expert-parallel +target** (`tep4`, 72 experts over 4 ranks, 18 per rank at offset +`18 * ep_rank`): + +* **Every per-expert preparation step survives unchanged.** The + `[up; gate]` concat, the interleave, the 32-row block shuffle and the + 128x4 scale swizzle all happen *within* one expert, and EP splits the + stack on the expert axis only. So every transform below applies + verbatim, with the first dimension `E → E/ep_size` and the destination + index `e - local_expert_offset`. Nothing about the shuffle or the + swizzle is a function of the stack height. +* **The three per-expert scale scalars must be `[local_num_experts]`, but + the checkpoint's scalars must still be read whole.** The runner rejects + `num_experts`-long scalars, so the window slice is real — but the + *validity* argument for the shared-expert activation scale is global + (`shared.input_scale == max` over all `E` routed `input_scale`), and a + target that windows at load time can no longer assert it. Load all `E` + (six fp32 per expert — negligible), assert globally, slice in the + post-load derivation. +* **`leftover == {}` stops being the coverage assert.** Under EP each rank + legitimately leaves the off-window experts' weight and scale keys + unconsumed. Replace it with an explicitly predicted leftover set + computed from the rank's window — that keeps the assert bidirectional + rather than relaxing it to a warning, which is the whole value of having + it. + +One cross-rank invariant a single rank cannot check, recorded because the +assembly rests on it: **the four windows must tile the routing space +exactly once.** They do because the router GEMM and the routing op are +replicated and deterministic and every rank sees identical all-reduced +inputs — so each token's top-k ids are the same on every rank, and each +id falls in exactly one window. + +**That justification is topology-specific, and it does not survive +attention data parallelism.** Under a `depN` segment the ranks hold +*different* tokens, so "every rank sees identical inputs" is false by +construction and the invariant loses its support — while remaining just as +load-bearing, because EP still requires each expert id to belong to exactly +one rank's window. Established by the first attention-DP MoE target +(`dep4`): what restores it is **placing the token all-gather before the +router GEMM, not merely before the expert call.** The router then runs on a +byte-identical full token set on every rank, the original argument applies +again verbatim, and the four windows tile as before. + +The failure this rules out is worth naming, because it is silent. Gathering +*after* the router — the arrangement that looks equivalent, and is cheaper +by the width of the routing tensors — leaves every rank routing only its +own tokens, so a token's expert ids exist on one rank alone and the other +three windows never see it. There is no error and no hang; the output is +simply missing most of its expert contributions. Both gates would be the +only thing that catches it. + +Stated as the rule: **on any topology where the ranks' token sets differ, +the collective that reconstitutes the token set belongs upstream of the +router, and the tiling invariant is what decides that placement** — not the +expert call's input requirements, which the later placement also satisfies. + +## The shape of the vocabulary + +A bf16 MoE block is three catalog calls, not one: + +``` +router_logits = cublas_mm(x, router_weight_view) # [T, E] +expert_ids, scales = renorm_moe_routing_op(router_logits, topk) +moe_out = fused_moe(x, expert_ids, scales, fc1, None, fc2, None, ...)[0] +``` + +`fused_moe` contains the permutation, both grouped GEMMs, the gated +activation and the routing-weighted combine. It does **not** contain the +router GEMM or the selection — those stay the caller's, which is why the +routing op exists as a separate entry. + +Consequence worth stating because it looks like an omission: a +fully-MoE checkpoint's forward references **no `activation/*` entry**. The +SwiGLU lives inside `fused_moe` (`activation_type=5`). An audit that +expects `silu_and_mul` in every model's forward is wrong for this class. + +## The `[up | gate]` order — the one silent-wrong-answer trap + +`fused_moe`'s `fc1_expert_weights` is `[E, 2I, H]` with **up rows first**: + +- `fc1[e, :I]` is HF `up_proj` (trtllm `w3`) +- `fc1[e, I:]` is HF `gate_proj` (trtllm `w1`) — the sigmoid is applied to + this half + +`fc2_expert_weights` is `[E, H, I]`, plain `down_proj`. + +Getting the halves backwards produces a plausible, entirely different +result: finite, correctly scaled, no error, no NaN. Nothing in the stack +detects it — not the op, not the boot gate. Only an accuracy gate catches +it, and only because it lands ~50 points low. + +**This is the opposite of the dense convention.** A dense target packs +`gate_up` with gate rows first, because that is what +`flashinfer_silu_and_mul` expects over a packed last dim. The two orders +must never be copied across; each is correct for its own consumer. + +## Checkpoint layout, and where the key names actually live + +A 4.51-era checkpoint stores experts **unstacked**: 128 × 3 separate 2D +tensors per layer, +`model.layers.{i}.mlp.experts.{e}.{gate,up,down}_proj.weight`. Recent +transformers models the same block with stacked 3D `gate_up_proj` / +`down_proj` parameters — so **the installed modeling file is not the +source of truth for checkpoint key names**. The safetensors index is. +Read `model.safetensors.index.json`, not `modeling_*.py`, when writing +the manifest. + +### A stacked `gate_up_proj` may still be stored gate-first + +Stacked does not imply trtllm's order. A checkpoint can ship a single +`experts.gate_up_proj` of exactly the `[E, 2I, H]` shape `fc1_expert_weights` +wants and still hold **gate rows first**, because that is what its own HF +reference reads: `gate, up = linear(x, gate_up_proj).chunk(2, dim=-1)` then +`act_fn(gate) * up`. The shape matches, the dtype matches, the manifest's +coverage assert passes, and a pass-through is silently wrong in exactly the way +this file's `[up | gate]` section describes. + +So the right shape is not evidence of the right order. Decide it from the +reference's own use of the tensor, and **swap the halves in `weights.py`** when +the checkpoint is gate-first. `down_proj` maps to `fc2_expert_weights` +unchanged either way. + +Measured on `Qwen3.6-35B-A3B-NVFP4`'s MTP layer (`mtp.layers.0.mlp.experts.gate_up_proj`, +`[256, 1024, 2048]` bf16): against `fused_moe` at that cell, a gate/up-swapped +reference lands 320 ulp element-wise / 198 ulp relative RMS from the correct +one — 40x the entry's element gate. Loud once something compares; invisible +until then. + +The unstacked on-disk layout reaches `fused_moe`'s stacked layout with +**zero transforms** — every copy is layout-preserving: + +```python +((f"{p}.mlp.experts.{e}.up_proj.weight", (e, slice(0, inter))),) +((f"{p}.mlp.experts.{e}.gate_proj.weight", (e, slice(inter, 2 * inter))),) +((f"{p}.mlp.experts.{e}.down_proj.weight", (e,)),) +``` + +This needs one generalization of the manifest convention: the destination +widens from a 2D `(row_start, row_end)` range to an index tuple into +`param.data`, so one table serves both fused-2D parameters (qkv) and +stacked-3D ones (experts). It is a one-line change in `fill`. + +3D parameters need no special handling anywhere else: they pass through +`MetaInitMode` and engine materialization unchanged. Measured at this +scale — 434 `ParameterDict` entries, 58 GB of expert stacks, 18432 +per-expert copies of ~3.1 MB — the whole load sits inside a 14.9 s model +init. + +## What the structure dictates, and what it makes irrelevant + +Established by measurement during the perf campaign; these are properties +of the mechanism, so the next MoE target can skip the sweeps. + +**The MoE is HBM-roofline-bound at serving batch sizes.** Above a decode +batch of roughly `E / topk × small factor` (~40 here) essentially every +expert is active, so a decode step reads the **entire** expert stack once +— 58 GB for this geometry. Measured 7.350 ms/step for the two grouped +GEMMs at concurrency 256 ⇒ **7.89 TB/s**, i.e. the B200 HBM roofline. No +config knob and no scheduling change moves this. Only fewer weight bytes +(quantization — a modeling change) would. + +**The kernel-count floor is dominated by the MoE.** `fused_moe` plus +`renorm_moe_routing_op` contribute 7 of the 16 kernels per layer +(`customMoeRouting`, `fusedBuildExpertMapsSortFirstToken`, +`expandInputRows`, grouped GEMM 1, `doActivation`, grouped GEMM 2, +`computeStridesTmaWarpSpecialized`). Across 48 layers that is 336 of 768 +kernels per decode step. At concurrency 1 the measured 3.87 ms TPOT over +794 kernels is 4.9 µs each — the same order as the smallest kernels +measured at batch 256, i.e. a fixed per-kernel floor rather than work. The +single-stream latency point is structural until the vocabulary gains a +coarser op. + +**KV capacity does not bind.** A 3B-active/30B-total model at tp1 leaves +the KV pool enormous relative to any realistic concurrency (measured: +1,168,096 pool tokens = 570 requests at ISL+OSL 2048, against a 256-request +ceiling). `kv_cache_config.free_gpu_memory_fraction` and the capacity +scheduler policy are inert — do not spend sweeps on them. + +**The autotuner is not a cold-cache risk inside a served engine.** +`fused_moe`'s contract warns that a cold tuning cache silently falls back +to a default tactic. In serving that does not apply: the runtime's +`_run_autotuner_warmup` wraps the forward in `autotune()` before capture +(observed `Cache size after warmup is 28` = 2 tunable GEMMs × 14 +power-of-2 token buckets). Worth knowing before anyone spends an iteration +on it — and worth knowing in the other direction too: the tuned tactic +changes the *bits*, not just the speed (measured 2.96 ulp against the cold +fallback on the same input), so a catalog receipt taken cold certifies a +tactic the served engine never runs. Both states sit inside the entry's +accuracy gate; the point is that "same inputs, same outputs" holds only +within one tuner state. + +## Routing traps + +Both from `renorm_moe_routing_op`'s certification: + +- **Ties break toward the lower expert index**, the opposite of + `torch.topk`. On bf16 logits exact ties are common, so a + `torch.topk`-based reference disagrees on indices while the weights + still match. Do not validate routing against `torch.topk`. +- **The kernel ignores strides**, reading `router_logits` as a dense + row-major buffer from `data_ptr()`. A strided view routes silently + wrong. The wrapper guards it; a caller building logits as a slice of a + wider buffer must materialize them contiguous. + +And from `fused_moe`: expert ids outside the rank's slot range are +**silently dropped**, not clamped and not rejected, so a routing bug +surfaces as a quietly weaker token rather than an error. + +## Deriving "is every layer MoE?" + +Do not assume uniformity, and do not assume `intermediate_size` is live. +The dense-vs-sparse branch per layer is +`layer_idx not in mlp_only_layers and num_experts > 0 and (layer_idx + 1) +% decoder_sparse_step == 0`. With `decoder_sparse_step: 1` and +`mlp_only_layers: []` every layer is sparse, which makes +`intermediate_size` **dead config** that no layer reads — while +`moe_intermediate_size` is the live one. A checkpoint with a nonzero +`decoder_sparse_step` or a non-empty `mlp_only_layers` needs both branches +and both sets of weights. Assert the derivation at construction rather +than trusting it. + +## MXFP4 expert stacks — the same trap, in interleaved form + +From the gpt-oss-120b/sm_100/tp1 onboard (36 layers, all MoE, 128 +experts, top-4, hidden = intermediate = 2880, experts MXFP4 while router, +attention, embedding and lm_head stay bf16). Everything above about the +`[up | gate]` order still holds; a block-scale-quantized stack adds three +things. + +**The parity is inverted relative to the half-split convention.** The +checkpoint stores `gate_up_proj_blocks` as `[E, 2I, K/32, 16]` uint8 — +already in `nn.Linear` `[out, in]` orientation, transposed relative to +HF's `[E, hidden, 2*inter]` parameter — and HF reads `gate = +gate_up[..., ::2]`, `up = gate_up[..., 1::2]`. So the stored row order +along the `2I` axis is (gate, up, gate, up, …), while the trtllm-gen +kernel's interleave wants destination row `2i` = **up** `i`, `2i+1` = +**gate** `i`. The split is therefore `up = t[:, 1::2]`, `gate = +t[:, 0::2]`, re-concatenated `[up ; gate]` before the row permutation. +Measured discrimination on real layer-0 tensors: 1.12 bf16 ulp against a +correct pure-torch reference, 146 ulp against the swapped one — a 130x +separation. Worth running that check *before* the first engine boot: it +covers nibble order, the parity split, the concat, the interleave, the +32-row block shuffle, the 128x4 scale swizzle, both padded axes, both +fp32 biases and the alpha/beta/limit triple in one shot, and the accuracy +gate that would otherwise catch it costs a full evaluation. + +**Where the relayout lives decides whether the model fits.** Expressing +pad → concat → interleave → shuffle → swizzle as *manifest source +transforms* (the sharding convention's `src` slot, generalized to any +callable) keeps peak memory at one layer of scratch. Declaring +checkpoint-shaped parameters and transforming them in `post_load_weights` +instead needs both forms resident — 63 GB checkpoint plus 66 GB +kernel-ready, on a 183 GB card that also wants a KV pool. The streaming +form loaded 63 GB including the on-device relayout in 14.1-14.5 s. The +kernel-ready operands at this geometry are `[128, 5888, 1536]` + +`[128, 5888, 96]` + fp32 `[128, 5888]` (FC1) and `[128, 2944, 1472]` + +`[128, 2944, 92]` + fp32 `[128, 2944]` (FC2) per layer, ~1.85 GB/layer, ++4% over the checkpoint. + +**Dtypes on disk are not the dtypes the kernels want.** Expert biases and +per-head attention sinks are bf16 in the checkpoint; the trtllm-gen MoE +requires fp32 biases and the attention op an fp32 sink tensor. Promote at +load. The same row permutation applied to the weight bytes must also be +applied to the scale bytes **and** the fp32 bias — a non-permuted bias is +a silent wrong answer. + +## W4A16 vs W4A8: the same weights, a different activation path + +Both members of the trtllm-gen block-scale MoE family consume the +**byte-identical** prepared expert stack — confirmed from the quant-method +source (both inherit one base, overriding only `create_weights` / +`load_quant_scales`, which call `super()`) and by measurement across every +certified geometry. Moving between them is a forward change only; +`weights.py` does not move. + +| | W4A16 | W4A8 | +|---|---|---| +| activations | bf16 straight in | `mxfp8_quantize(x, swizzled_layout=False, alignment=512)` first | +| hidden padding | caller pads 2880 -> 3072 | the quantizer does it — **the pad call disappears** | +| `valid_hidden_size` | 2880 | 2880 (output width; unrelated to the widening, and `None` is rejected) | + +Two consequences that are not obvious from the signatures: + +- **The W4A8 kernel requantizes the FC1 activation to MXFP8 between the + GEMMs**, on the OCP scale (`e = floor(log2 amax) - 8`), *not* the + round-up scale the standalone quantizer uses. Skipping that step in a + reference deviates 13.5 ulp element-wise / 8.5 ulp RMS (~3% relative) + from the kernel — so it is a real perturbation of the layer output, and + it is the mechanism behind the accuracy difference between the two + members. On gpt-oss-120b the measured cost was 1.44 GSM8K points + (90.5989 W4A16 -> 89.1585 W4A8), reproducible bit-identically on both + sides, i.e. the recipe's price rather than sampling. +- **`swizzled_layout=True` is silently accepted** by the W4A8 MoE + whenever the byte counts coincide (`T % 128 == 0` — exactly the CUDA + graph batch sizes), and is 260 ulp wrong. No metadata distinguishes the + two layouts, so no wrapper can guard it. Put the `False` in a named + constant. + +The perf reason to pay that accuracy: on gpt-oss-120b the W4A16 MoE ran +5.6x slower than the W4A8 one on byte-identical weights (FC1 711.8 vs +126.7 µs per layer per step at concurrency 128), which was the target's +*entire* gap to stock trtllm. The swap moved peak throughput +88% and put +the step within 0.5% of the reference. + +## The dtype chain, end to end + +Worth writing out because the names invite a wrong guess: **W4A8 +quantizes activations to fp8, never to fp4.** The op name spells both +operands — `mxe4m3_mxe2m1_...` is e4m3 activations against e2m1 weights. +The `mx` prefix is OCP micro-scaling: 32 elements share one E8M0 +(power-of-two) scale, so MXFP8 is 32 e4m3 values plus one scale byte and +MXFP4 is 32 e2m1 values plus one. + +| step | format | +|---|---| +| hidden states entering the block | bf16 `[T, H]` | +| `mxfp8_quantize(x, False, 512)` | e4m3 `[T, pad_up(H, 512)]` + E8M0, one byte per 32 | +| FC1 GEMM | e4m3 x e2m1, **fp32 accumulate** | +| clamped GLU (FC1 epilogue) | fp32 | +| requantization of the intermediate | e4m3 + E8M0 per 32 columns | +| FC2 GEMM | e4m3 x e2m1, **fp32 accumulate** | +| routing-weighted combine | fp32 | +| store | **bf16** `[T, valid_hidden_size]` | + +Two quantization points, both to fp8, with fp32 everywhere between them. +Only the operand *storage* formats are narrow — the arithmetic never drops +to fp4. The output is bf16 at the model's true hidden width, not the +padded one, and `valid_hidden_size` has to be passed explicitly to get it. + +## Why the W4A16 member is slow, and the W4A8 one is not + +The two paths read the **same weight bytes**, so the 5.6x is not +bandwidth. It is what each kernel does with them, and the kernel names say +it outright: + +| | operand fields in the kernel name | conversion | +|---|---|---| +| W4A16 | `bmm_Bfloat16_MxE2m1Bfloat16_castBfloat16_...` | `castBfloat16` — MXFP4 expanded to bf16, then a bf16 MMA | +| W4A8 | `bmm_MxE4m3_MxE2m1MxE4m3_...` | none — MXFP4 goes straight into a block-scaled MMA | + +The expansion's cost is not mainly the conversion arithmetic. A +dequantized weight block occupies **4x the on-chip bytes** (4-bit e2m1 -> +16-bit bf16), so fewer of them fit in registers and shared memory, which +shrinks the tile and shortens the software pipeline — visible in the same +kernel names: + +- W4A16: `m128x8x16`, 3 pipeline stages, 1-CTA clusters +- W4A8: `m256x16x32`, 6 stages, 2-CTA clusters + +Fewer weight bytes in flight means HBM latency stops being hidden, and the +GEMM lands far below the memory roofline instead of at it. At concurrency +128 on gpt-oss-120b, where the batch touches essentially every expert so a +decode step reads the whole stack once (~61 GB at valid sizes): + +| | MoE per step | implied rate | +|---|---|---| +| W4A16 | 38.89 ms | ~1.6 TB/s | +| W4A8 | 8.34 ms | ~7.3 TB/s | +| stock trtllm | 7.51 ms | ~8.1 TB/s | + +against a roofline this repo measured at 7.89 TB/s on the same machine. So +the W4A16 member was not paying more bandwidth — it was failing to use the +bandwidth it had, at roughly a fifth of the achievable rate. The residual +W4A8-vs-reference gap is autotuner tactic choice (~11%), not the recipe. + +*Measured:* the kernel names with their tile / stage / cluster fields, the +per-GEMM times (FC1 711.8 vs 126.7 µs, FC2 352.9 vs 65.6 µs per layer per +step), the MoE step times, and that the weights are byte-identical. +*Derived:* the byte figure behind the implied rates, and the chain from +on-chip footprint to tile size to unhidden latency — consistent with every +number above, but not isolated experimentally. + +## Blackwell only, and it fails quietly elsewhere + +This kernel family is sm_100 / sm_103. Three independent signals: + +- TensorRT-LLM gates its own tests on it — `get_sm_version() not in (100, + 103)`, reason "TRTLLM Gen MoE supports SM100 and SM103 only". +- The quantization kernel's body compiles only under `__CUDA_ARCH__ >= + 1000`; below that the empty branch **returns without writing the + outputs**, so a pre-Blackwell run produces garbage rather than failing. + Both catalog entries carry sm_100-only receipts for that reason — a + statement about the hardware, not caution. +- The performance signature itself. A software emulation of the mixed + fp8 x fp4 MMA — unpack, widen, then a wide MMA — would look exactly like + the W4A16 path measured above, and nothing can be 5.6x faster than the + thing it emulates. + +The hardware feature is Blackwell's block-scaled MMA: the instruction +takes narrow operands plus their per-block scale factors and applies the +scaling inside the tensor-core datapath, so no widened operand is ever +materialized. The E8M0 scale being a power of two is what makes that +nearly free — it is an exponent adjustment. (The arch gating and the +performance signature are measured here; the datapath description is the +standard account of the feature and was not verified at the ISA level.) + +## NVFP4 expert stacks — the same traps, plus per-expert scales + +From the deepseek-v3-lite-nvfp4/sm_100/tp1 onboard (DeepSeek-V3-Lite: 30 +layers, layer 0 dense, layers 1-29 MoE with 72 routed experts at top-6, +`moe_intermediate_size` 1536, plus 2 shared experts fused into one +`[3072, ...]` dense pair; experts NVFP4 while attention, router, embed +and lm_head stay bf16). Everything above about the `[up | gate]` order +still holds. NVFP4 differs from the MXFP4 case in three ways that matter. + +**The consumer is a different op with a different operand shape.** The +trtllm-gen NVFP4 MoE takes `[E, 2I, H/16]` **`float8_e4m3fn`** scales, +where the MXFP4 members take `[E, 2I, H/32]` uint8 UE8M0. The prepared +stacks are **not** interchangeable between the two families, even though +both are "block-scale MoE runners". + +| operand | shape | dtype | +|---|---|---| +| `gemm1_weights` | `[E, 2I, H/2]` | uint8 (two e2m1 codes per byte) | +| `gemm1_weights_scale` | `[E, 2I, H/16]` | float8_e4m3fn | +| `gemm2_weights` | `[E, H, I/2]` | uint8 | +| `gemm2_weights_scale` | `[E, H, I/16]` | float8_e4m3fn | + +Per-expert preparation is: concat `[up ; gate]` (**up rows first**) -> +interleave (dst row `2i` = up `i`, `2i+1` = gate `i`) -> 32-row block +shuffle (src `4u+v` -> dst `8v+u`) applied to weight bytes **and** scale +bytes -> 128x4 swizzle of the scales only. +`torch.ops.trtllm.shuffle_matrix` and +`torch.ops.trtllm.block_scale_interleave` perform the last two; both are +load-time transforms, outside the closed-vocabulary rule. + +Measured discrimination against a correct pure-torch reference, on real +checkpoint tensors — the reason to run this check before the first engine +boot rather than after an accuracy gate: + +| variant | bf16 ulp | +|---|---| +| **correct** | **4.75** | +| `[gate\|up]` instead of `[up\|gate]` | 263 | +| weight scales not 128x4 swizzled | 408 | +| expert rows not 32-row shuffled | 508 | +| `output1_scale_scalar` / `_gate_scalar` swapped | 177 | +| no intermediate requantization | 42.7 | + +**The half-order trap appears twice in one layer on a shared-expert +model.** The dense and shared-expert paths pack `[gate ; up]` for +`flashinfer_silu_and_mul`; the MoE FC1 packs `[up ; gate]`. Both live in +the same decoder layer, three lines apart. The two orders must never be +copied across — each is correct for its own consumer. + +### modelopt stores reciprocals + +A modelopt NVFP4 checkpoint stores `input_scale` and `weight_scale_2` as +`amax / (448 * 6)` — the **reciprocal** of the global scale a quantizer +divides by. So: + +``` +global_scale (for fp4_quantize) = 1 / input_scale +dense GEMM alpha = input_scale * weight_scale_2 # both off disk +``` + +Reading them the other way is wrong by `global_scale^2`, finite, and +invisible short of an accuracy gate. Settle it by arithmetic rather than +by convention: under the reciprocal reading the implied amax values are +0.2-3.0, physically sensible for normed activations and NVFP4 weights; +under the other reading they are ~1e7. `gate_proj` and `up_proj` share +both scalars exactly (verified across all experts of every MoE layer) — +assert it at load rather than assuming it. + +The trtllm-gen runner's three `[E]` fp32 arrays, in checkpoint terms: + +``` +output1_scale_gate_scalar[e] = input_scale_gate_up * weight_scale_2_gate_up[e] +output1_scale_scalar[e] = output1_scale_gate_scalar[e] / input_scale_down +output2_scale_scalar[e] = input_scale_down * weight_scale_2_down[e] +``` + +The gate scalar dequantizes the half feeding the sigmoid; the other +carries the extra FC2-input global scale because the FC1 epilogue +re-quantizes its output to NVFP4. + +### One activation scale, but per-expert weight scales + +Two facts that look like they should match and do not. + +**The activation global scale `g1` is one number, and the shared expert +carries it.** The runner quantizes the hidden states once, so it takes a +single `g1`. Which one? On this checkpoint the shared expert's +`gate_proj.input_scale` is **exactly** the max over all routed experts' +— exact equality in every MoE layer. That is not a coincidence worth +guessing at: the shared expert sees every token, so its amax *is* the +global amax. `g1 = 1 / shared.gate_proj.input_scale` therefore serves +both the dense and MoE quantization calls. Routed experts' own values sit +up to 1.4-2.3x below it. + +**The FC2 input scale `g2` is genuinely per-expert** and does not +collapse. It spans **54.5x** across experts in a single layer. The +runner's three `[E]` arrays exist precisely to carry it; feeding a scalar +would be wrong by that factor on the extreme experts. + +### Two quantization calls, not one + +`nvfp4_gemm` (dense and shared-expert linears) consumes the **swizzled** +scale-factor buffer; the trtllm-gen MoE runner consumes the **linear** +one, viewed as `float8_e4m3fn`. The same hidden states must therefore be +quantized **twice**, once in each layout — they cannot share a result. + +The failure mode is nasty: at `T % 128 == 0` the swizzled buffer happens +to be the *right size* for the runner, so it is silently accepted and is +276 ulp wrong. At every other token count it is rejected on size. A boot +gate at a round batch size will not catch it, and a boot gate at an odd +one will look like a size bug rather than a layout bug. + +### Two smaller obligations + +- **`hidden` must be a multiple of 256** for the trtllm-gen NVFP4 runner. + At a hidden that is only a multiple of 128 the kernel reads FC1 weight + scales for the wrong blocks: no error, 301-3370 bf16 ulp wrong. This is + an upstream hazard too — trtllm's own weight creation only raises + alignment to 256 when `hidden_size > 1024`. +- **Per-tensor quantization scalars are stored 0-dim.** `entry[:]` on + them raises `IndexError: slice() cannot be applied to a 0-dim tensor`, + so a manifest's materialize step needs a rank check that the bf16 path + never needed. + +## Group-limited routing, and where the grouping lives + +From `deepseek-r1-0528-nvfp4/sm_100/dep4` (256 routed experts, top-8, +`n_group 8` / `topk_group 4`), the first target on a grouped checkpoint — +its siblings all run `n_group: 1`. + +**The grouping is entirely the routing op's.** `noaux_tc_op` takes +`n_group`/`topk_group` and does the group-limited top-k; the block-scale +MoE runner's own `n_group`/`topk_group` arguments stay **`None`** at +`routing_method_type = 1`, because the runner is on its pre-routed path +and consumes finished `topk_ids`/`topk_weights`. Passing the grouping +twice is not how it composes. + +Two constraints the grouped path adds that the ungrouped one does not, and +that the routing op reports as one opaque "unsupported configuration": at +`n_group > 1` it requires `1 <= topk_group <= n_group`, `topk <= 8`, +`num_experts <= 256`, and `num_experts / n_group <= 32`. Assert them +separately or a violation is unattributable. + +**Router bias dtype.** The combine weights follow the **logits'** dtype, +not the bias's. A bf16 router GEMM with the fp32 `e_score_correction_bias` +this checkpoint ships returns bf16 weights — which is what the MoE runner +demands (it rejects fp32 `topk_weights`). No cast is needed and adding one +is a mistake. + +## A four-way expert-parallel split at 256 experts + +The invariant is the one this file already records — the windows must tile +the routing space exactly once — now measured at 64 wide: + +* the four 64-wide windows **sum to the whole-layer result** (0.85 ulp RMS + against an fp32 reference), but are **not bitwise equal** to a single + 256-expert call: up to 2.00 ulp apart, because each window rounds its own + partial to bf16 before the add. Dropping any one window lands ≥ 125 ulp + RMS away, so the sum is a real check rather than a formality; +* the **autotuner's key holds `local_num_experts` but not + `local_expert_offset`** — warming one window of a split warms all four, + and a 256 → 64 change forces a fresh sweep. Cold and warm agree bitwise + across 85 configurations here, but do not generalize that: the bf16 + `fused_moe` entry records the same check *failing*. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/graph-and-sync.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/graph-and-sync.md new file mode 100644 index 000000000000..1821364f3952 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/graph-and-sync.md @@ -0,0 +1,88 @@ +# Graph safety and sync-free rules for a target forward + +> Established on B200 (`sm_100`). The +> three-class taxonomy and the sync-free rewrites are properties of CUDA graph +> capture and of torch, not of any particular engine build; what an engine +> change can move is which values the engine hands the step, so re-check the +> step-args classification rather than this file's rules. + +These are **writing** invariants, not tuning advice. A forward that violates +them is wrong the first time it is captured — which is why they live here, +where the assembler writing the first version and the tuner editing it later +read the same copy. + +## Three constraints that shape every forward + +| Constraint | What it forbids | +|---|---| +| thop custom ops are permanently invisible to `torch.compile` | no compile-based structure in the forward, and no compile-fusion reasoning about a trtllm-backed call | +| **the KV red line** — K/V are written *inside* the attention op; the target only forwards the address book | nothing in the forward may take ownership of KV, move it, or reshape it | +| the hot path reads weights **by reference** | a per-step weight transform is a bug; derived state belongs in `post_load_weights` | + +## What a replayed forward may contain + +Decode replay executes captured kernels only. Classify every value the +step consumes — the step-args builder is the single audit point: + +| Class | Meaning | Replay behavior | +|---|---|---| +| reference | engine-owned persistent buffer, refreshed in place each step | replay reads fresh contents automatically (all step-args tensors) | +| GPU-derived | produced by captured kernels | recomputed on replay | +| host-derived | Python scalar baked into the capture | frozen — legal only as a per-capture constant (e.g. `num_contexts == 0` in a decode graph) | + +What breaks replay or stalls the host: + +| Pattern in the step path | Why it is broken | +|---|---| +| `.item()` / `.cpu()` / `.tolist()` / `print(tensor)` | host-device sync (cuStreamSynchronize); under capture the read value is frozen anyway | +| Python `if`/`while` on tensor **values** | syncs to read the value, and the branch taken is frozen into the graph — branch on tensor metadata or per-capture constants only | +| host-side shape math from tensor values (`nonzero`, `masked_select`, slice bounds read from CUDA tensors) | data-dependent output size forces a sync and bakes one shape into the capture | +| `torch.tensor(py_obj, device='cuda')` built per step | blocking host-to-device creation; make it a load-time constant in `post_load_weights` | +| allocation sized from tensor **values** | needs the value on host — sync | +| allocation sized from tensor **metadata** (`empty([x.shape[0], ...])`) | fine — the caching allocator serves stable blocks under capture; upstream captures the same pattern | + +## Sync detection + +```python +torch.cuda.set_sync_debug_mode("warn") # or "error" to hard-fail on sync +# ... run the step under test ... +torch.cuda.set_sync_debug_mode(0) +``` + +This catches only syncs going through torch's wrapped APIs — thop ops and +libraries calling CUDA sync APIs directly are invisible to it. For those: + +```bash +nsys profile --capture-range=cudaProfilerApi --python-sampling=true \ + --backtrace=dwarf +``` + +then search the CUDA API row for `cudaStreamSynchronize` / +`cudaEventSynchronize` / `cudaDeviceSynchronize`; the call stack names the +Python line. Verify a fix by re-running detection, not by reading the diff. + +Classify before fixing: **false dependencies** (debug prints, leftover +`.item()`, per-step tensor creation from Python objects) are deleted or +moved to load time; **true dependencies** (the CPU genuinely needs a GPU +value) are restructured onto the GPU or kept outside the captured region. + +## Sync-free quick table (inference-forward rows) + +| Sync-inducing pattern | Sync-free shape | +|---|---| +| `.item()` in a conditional / `if gpu_tensor:` / `torch.is_nonzero(t)` | GPU-side select (`torch.where`-shaped) or branch on metadata | +| Python `max(a, b)` on GPU tensors | `torch.max`-shaped GPU op | +| `torch.tensor(obj, device='cuda')` / `torch.as_tensor(..., device='cuda')` | precompute at load time; if truly per-step, create on CPU and `.to('cuda', non_blocking=True)` | +| `x_gpu[idx_cpu]` / `x_gpu[python_list]` | index with a resident GPU index tensor | +| `x[i:j]` with bounds read from CUDA tensors | metadata-derived bounds, or a GPU `arange`-shaped index | +| `x_gpu[mask_gpu]` / `torch.masked_select(x, mask)` | fixed-shape `torch.where(mask, x, 0)` | +| `torch.nonzero(mask)` | `where`-shaped rewrite, or move outside the captured region | +| `torch.repeat_interleave(x, r)` | pass `output_size=N` when known | +| `torch.unique(x)` | avoid in the step path; precompute | + +The right-hand column names torch shapes, not licenses: in a target +forward the actual call must be a catalog entry. A replacement with no +entry in +`tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml` is a +vocabulary gap — never paste raw torch math +into the forward to work around a sync. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/hybrid-recurrent-state.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/hybrid-recurrent-state.md new file mode 100644 index 000000000000..69897005b383 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/hybrid-recurrent-state.md @@ -0,0 +1,128 @@ +# Hybrid recurrent-state targets (linear attention beside full attention) + +For a checkpoint whose layers are a mix of `full_attention` and a recurrent +mixer — Gated DeltaNet, Mamba2, KDA — where the recurrent layers need a +**conv + recurrent state pool** the runtime allocates, alongside the ordinary +paged KV cache for the attention layers. + +Sibling of `latent-kv-cache.md`, which is MLA's cache mechanism; this one is +the recurrent-state mechanism. Both describe what a target must do so the +runtime allocates the right pool. + +## The runtime decides from the published config, not from the weights + +Whether a recurrent-state pool exists at all is chosen before any weight is +read, from `model_config.pretrained_config` — **the checkpoint's own config, read +exactly as published**. If the predicate does not match, the model still builds, +the boot gate still passes, and the recurrent layers have no state. + + +A target can neither change nor restate what that predicate sees. The synthetic +class name routing returns is local to `_resolve_class` — +`pretrained_config.architectures` is not mutated — so registration has no +influence on what the runtime allocates, and there is no stub `config.json` to +patch. If a checkpoint the target is meant to serve does not trip the predicate, +that is a fact about the checkpoint and the runtime: **report it as a symptom**, +and do not try to route around it from modeling code. `runtime-interface.md` §2 +has the registry precedence facts and the reasoning. + +## A nested `text_config` is not reliably mirrored + +Composite VLM checkpoints keep the text fields under `text_config`. The runtime +has a mirroring pass that copies `text_config` attributes onto the parent, but +relying on it is a trap: it only fills in attributes that raise +`AttributeError` on the parent, so any key for which the config class already +supplies a default is silently **not** mirrored, and the text value is ignored. + +Out of tree a target sidestepped the question by shipping a flattened stub. +There is no stub now, so the trap is live: a field the target's geometry +depends on may be read by the runtime from the class default rather than from +`text_config`. Read each such field from where the runtime will read it, assert +it at construction, and report any field whose live value differs from +`text_config`'s as a symptom rather than compensating for it in the forward. + +Fields the hybrid path reads off the config: `architectures`, +`layer_types` (validated against `num_hidden_layers`, and every entry must be +one of `full_attention` / `linear_attention`), plus the recurrent geometry the +cache manager sizes the pool from — for the Qwen3-family naming, the key head +dim, conv kernel dim, and the value/key head counts and dims. + +## The state-cache dtype field is not the one the checkpoint declares + +A checkpoint may carry a field expressing the recurrent **compute** dtype +intent — Qwen3.5-family configs write `mamba_ssm_dtype: float32`. That is not +the cache dtype and the runtime deliberately ignores it. Only an explicit +`mamba_ssm_cache_dtype` (config field or `kv_cache_config`) is honored, falling +back to the model dtype. + +**Do not set it under the cache name.** Doing so switches the +state cache to fp32, which disables the bf16-state decode kernel and doubles +state memory traffic — upstream's own comment puts that at roughly 20% serving +throughput. This is a performance loss with no gate that can see it: the boot +gate passes, accuracy passes, and only a sweep would show it. + +## The layer's cache offset is not its attention ordinal + +On a hybrid model the attention layers are a **sparse subset** of the layer index +space, and how they map onto pool offsets depends on which cache manager the +runtime picked. Under the default `CppMambaHybridCacheManager` the mask handed to +`KVCacheManager` covers recurrent *and* attention layers, so full-attention layer +3 sits at pool offset 3 — not at attention-ordinal 0, which is what a dense model +would give and what a `MixedMambaHybridCacheManager` would give here. + +So **read the offset from `attn_metadata.kv_cache_manager.layer_offsets` at the +first forward**; do not compute it by counting attention layers. Getting it wrong +is cross-layer aliasing — layer A reading or writing layer B's pages — which +produces finite, plausible output and is invisible to a single-layer test. + +The engine also runs **two pools** in this configuration: the attention KV pool +and the recurrent-state pool, with the recurrent one carrying a large negative +`window size` sentinel meaning "no window". The attention op never addresses the +recurrent pool — verified by decoy, pointing its pointer row at a spare +allocation leaves output bitwise identical — so that sentinel is not something a +target passes or reasons about. + +## Block reuse is off, until a feature turns the same flag back on + +The runtime **disables KV block reuse for a hybrid Mamba model** when no snapshot +policy is set, which is the default. Two consequences that look contradictory +until the flag is traced: + +- a hybrid target's *core* runs with `use_paged_context_fmha = False` on every + call, so any certification of the paged-context path is unexercised by it; +- but the flag is + `chunked_prefill or cache_reuse or has_speculative_draft_tokens`, so **adding + speculative decoding turns it on** even though reuse is still off. + +That matters on a quantized KV cache, because the paged-context path sources K/V +from the pool: on an fp8 pool the append leaves `e4m3(row/s)`, so a fresh prefill +is not numerically identical across the flag. A feature increment can therefore +move the trunk's prefill numerics without touching the trunk's code. + +## What the recurrent core op needs from the target + +The recurrent mixer is reached through a **registered-layer** op, the same +shape as the registered attention entry point: the call carries the activations +and a `layer_idx` string, and the op resolves a module instance out of the +model's extra-attrs to find the conv weight, the per-head gating parameters and +the state pools. So the target constructs and registers that module rather than +holding those weights itself, and the op's arguments are the projections only. + +Consequence for the closed-vocabulary audit: the conv, the gating arithmetic +and the recurrence are all *inside* one catalog entry, so they do not appear as +separate calls in the forward. The projections around it — qkv, the gate +projection, the two per-head scalar projections, the output projection — are +ordinary entries and do appear. + +--- + +*Runtime behaviour above verified on sm_100. An engine change re-opens all of +it: re-check by reading the hybrid predicate the cache manager selects on, and +confirming the state-cache dtype resolution still ignores the checkpoint's +compute-dtype field. The consequences of reading the checkpoint exactly as +published — that the target neither restates nor patches what the predicate +sees — follow from routing, not from any engine version.* diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/latent-kv-cache.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/latent-kv-cache.md new file mode 100644 index 000000000000..7b6e0c58b45c --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/latent-kv-cache.md @@ -0,0 +1,306 @@ +# Latent KV cache (MLA) + +What a target assembled on Multi-head Latent Attention had to establish by +observation, because no contract or reference stated it. Written from the +deepseek-v3-lite-nvfp4/sm_100/tp1 onboard (DeepSeek-V3-Lite: 30 layers, +32 q heads, `kv_lora_rank` 512, `qk_nope_head_dim` 128, +`qk_rope_head_dim` 64, `v_head_dim` 128, attention entirely bf16). The +mechanism — one compressed latent per token in place of per-head K and V +— recurs across the DeepSeek family and its derivatives; this file is +about the mechanism, not that checkpoint. + +**Attention TP replicates the latent pool; it does not shard it.** +Measured on the first multi-rank MLA target (`tep4`, world size 4): every +rank allocated 125.76 GiB for 3,907,296 tokens = **34,560 B/token** = +`num_layers * (kv_lora_rank + qk_rope_head_dim) * 2` = `30 * 576 * 2` — +bit-for-bit the world-size-1 figure, **with no rank factor anywhere**. + +The structural reason, which is why this generalizes rather than being a +fact about one engine: the latent is *one compressed row per token that +every head reads*. Splitting heads across ranks therefore splits nothing +in the cache. `kv_a_proj_with_mqa` and `kv_a_layernorm` stay replicated, +each rank computes the identical latent row, and each appends it to its +own full-width pool. + +So **attention TP buys zero KV-cache memory on MLA** — the opposite of +GQA, where the pool divides by `tp_size`. Budget a `tepN` MLA target for +the full single-rank pool on every rank. + +**Attention DP does not narrow the pool either — it multiplies what the +pool holds.** An earlier revision of this file predicted that a `dep` +segment would be "the topology that actually partitions the latent, by +partitioning the *requests* instead of the heads". The first attention-DP +MLA target (`dep4`, world size 4) measured it and **the prediction is half +wrong**: every rank still allocates a full-width pool at **34,560 B/token +— bit-for-bit the tp1 and tep4 figure, with no rank factor anywhere** +(146.5 GiB / 4,552,704 tokens per rank; 149.18 GiB / 4,634,880 under the +gate's config). + +What actually changes is *whose* tokens are in it. Under TP the ranks hold +**identical** latent rows, so three of the four copies are pure waste; +under DP they hold **disjoint requests**, so the same per-rank bytes become +job-wide capacity of `world_size x` the single-rank pool — 18,210,816 +tokens against `tep4`'s 3,907,296. The win is real and it is 4x, but it is +a *duplication* win, not a *sharding* win, and the distinction matters when +sizing: **per-rank pool arithmetic is identical across tp1, tepN and depN, +so a rank that fits at world size 1 fits at any of them, and one that does +not fit is not rescued by adding ranks.** + +Restated as the rule to size by: on MLA, no parallel topology divides the +per-rank latent pool. Choose `dep` for the aggregate capacity and for the +per-rank *compute* split, never to make the pool fit. + +Everything else below still comes from world-size-1 runs, and the sizing +facts hold per rank unchanged. + +**This file does not repeat the vocabulary.** +`.claude/skills/staircase-kernel-onboard/references/decompositions/mla-custom-op.md` +owns that: which entry +plays which role, the assembly per batch mode (§4), and the +assembler-side obligations certification established (§4.5). Read it +first. What follows is only what a target discovers *beyond* it, on the +way to a green gate. + +One item there can now be struck: §5 recorded that the end-to-end parity +run of an assembled MLA pipeline was deliberately left to the first MLA +target. It was run — against a pure-torch reference built from the HF +`DeepseekV3Attention` semantics, **0.93–1.02 bf16 ulp** across 1-seq +prefill, 2-seq prefill, mixed-batch context and mixed-batch decode, +identical at `tokens_per_block` 32 and 64. + +## Three things the contracts describe in prose and a target must measure + +**The `rotary_cos_sin` byte layout.** The contracts name a TensorRT-LLM +producer for it (`RopeParams(..., duplicate_data=True)`) — a symbol a +target is not allowed to read. The prose admits two readings, and both +are buildable. Only one is right: **interleaved `(cos, sin)` pairs at +flat offset `p * 2R + 2d`**, where `R` is `qk_rope_head_dim` and `d` +indexes the half-dimension. The "cos block, then sin block" reading is +**36x worse** at the same geometry. Build both and measure; do not +reason it out. + +**Which `v` the context FMHA wants.** The contract says "the +`[.., H*nope:]` split of a packed `[Tc, H*(nope+v)]` kv-projection +buffer". That is ambiguous against the packing an HF `kv_b_proj` +actually produces, which is **per-head** `[k_nope | v]` interleaved head +by head. The op settles it by rejecting a 3-D `v` +(`Expected v->dim() == 2`), and parity confirms it. The consequence is a +load-time obligation no contract states: + +> **A DeepSeek-style `kv_b_proj` must be row-regrouped at load** into +> `[all heads' k_nope ; all heads' v]`. Left in checkpoint order it is +> silently wrong. + +This composes with the constraint §4.5 already records — V must stay the +*split view* of the kv_b output, because the kernel hard-codes its row +stride. Regroup at load, then split; never materialize a contiguous V. + +**`create_mla_outputs` is not usable from a target.** Its +contract requires a live TensorRT-LLM `MLA` module registered in +thread-local extra attrs, which a target does not own and cannot +construct. The decomposition doc lists it under "output allocation", +which reads as available; it is not. Allocate with `torch/empty`. The +same is not true of its sibling `create_attn_outputs` for non-MLA layers. + +## The engine builds the latent pool for you — and turns on a path you must handle + +**No target-side declaration is needed**, and in tree there is nothing to +declare it in: the checkpoint is read exactly as published, and its own MLA +config fields are sufficient. The engine allocated a **kv_factor-1** pool +sized exactly +`num_layers * (kv_lora_rank + qk_rope_head_dim) * 2 B` per token +(30 × 576 × 2 = 34,560 B/token) and set +`enable_context_mla_with_cached_kv=True` **itself**. +`.claude/skills/staircase-modeling-core/references/runtime-interface.md` §6 +lists only standard-pool fields and does not cover this. + +**Block reuse is on by default, and it does hand an MLA model cached +context prefixes.** This is the trap. Observed directly on a plain +boot-gate workload: `num_ctx_cached_tokens=155`, `kv_lens_runtime=[162]`, +`prompt_lens_cpu=[7]`. A first build that asserted "no cached prefix" +had the assert fire on its third probe request. **Implementing only the +fresh-prefill context flavor is not sufficient under defaults.** + +**The cached-KV metadata surface is conditional, not merely False.** With +reuse on, the step object carries `enable_context_mla_with_cached_kv`, +`ctx_cached_token_indptr`, `ctx_kv_indptr` (both int64 CUDA), +`ctx_uncached_token_indptr`, `max_ctx_seq_len`, `max_ctx_kv_len` and +`num_ctx_cached_tokens`. With `enable_block_reuse: false` **those +attributes do not exist at all** — so a startup contract check that +asserts their presence turns a legal config into a hard failure. Key the +flavor on presence, and keep both branches. + +Two things that make this cheaper than it sounds: + +- **The cached path needs no extra glue.** `kv_lens_runtime` and + `prompt_lens_cpu_runtime` already satisfy the no-append call's + `sequence_length >= context_lengths`, and `host_total_kv_lens[0]` is + exactly the gather's row count. No cumsum, no host bookkeeping. +- **One path serves both.** append → gather → up-project → explicit-K/V + is correct at `cached_s == 0` too. That is what lets a single batch mix + reused and fresh sequences without a second code path. + +## Sizing facts that surprise + +**The attention workspace is an order of magnitude above the certified +figure.** `thop_attention`'s contract records **33 MB to ~109 MB** across +its certified shapes — the peak being an `H = 128` MLA context call over +545 tokens (113,924,096 B) — while an MLA engine at `max_num_tokens=8192` +resized it to **531,628,032 B (507 MiB)**. Budget for it before blaming a +KV-pool setting for missing memory. The contract's figure tracks +tokens x heads, so it grows with the head count a target actually runs; +read the current number there rather than a remembered one. + +**An fp8 latent pool adds a workspace surcharge, because the quantized +operands are staged there.** Measured on the no-append context flavor at +`H = 128`, identical shape (66 q rows over 193 K/V rows): **52,816,640 B +under fp8 against 43,288,832 B under bf16**. An engine sized from bf16 MLA +figures under-budgets on an fp8 KV cache. + +**First-boot JIT is a visible cost, once.** Eight `HQk576HV512` decode +kernels NVRTC-compile at ~5.6–7.2 s each — ~47 s of the first engine +build, with no on-disk cache. Context kernels are precompiled. Exclude +first calls from any timing. + +**`q_lora_rank` may legitimately be absent.** A checkpoint with +`q_lora_rank: null` uses a direct `q_proj` and no `q_a_proj` / +`q_a_layernorm` / `q_b_proj`. Passing `q_lora_rank=0` to the ops is +accepted and inert, and the contracts now enumerate it — +`thop_attention.md`'s geometry table lists `1536` and `0` for MLA, and +`mla_rope_generation.md` carries its own `q_lora_rank` row. + +Inert is measured, not assumed, and it had to be measured **twice**: the +sweep (`1536 → 0 → 1536 → 4096` over identical inputs, bitwise equal on +every observable) runs at `H = 32` and again at `H = 8`. One run would not +have covered the other, because the MLA generation phase compiles a +*different* decode kernel per head count at the low end — +`...VarSeqQ8Kv128...` at 8 heads where 16 and 32 both take +`...VarSeqQ16...`. An argument shown inert on one compiled kernel is not +thereby inert on another. A tensor-parallel MLA target slices heads, so +it is the head count the sweep did **not** run at that it will execute. + +## What the structure dictates, and what it makes irrelevant + +Established by measurement during the perf campaign; these are properties +of the mechanism, so the next MLA target can skip the sweeps. + +**KV capacity does not bind, and it is not close.** The latent cache is +`num_layers * (kv_lora_rank + qk_rope_head_dim) * dtype_bytes` per token +— one shared latent, no head multiplicity, no separate V — which here is +`30 * 576 * 2 = 34,560 B/token`. The same geometry as per-head MHA +(`2 * 30 * 32 * 128 * 2`) would be 491,520, so the latent form is +**14.2x cheaper per token**. A 30-layer model on one B200 +got a 142.25 GiB pool = **4,419,680 tokens**, against a 256-request +ceiling at ISL+OSL 2048 — two orders of magnitude of slack. +`kv_cache_config.free_gpu_memory_fraction` and the capacity scheduler +policy are inert. Do not spend sweeps on them. (The same conclusion was +reached for a different reason on a sparse-MoE target — see +`expert-weight-packing.md` — so it now holds for two unrelated +structures.) + +**`tokens_per_block` matters, but not for the reason the kernel list +suggests.** Raising it 32 → 64 was worth +22% at concurrency 256. The +MLA FMHA kernel does have page-size-specific variants and was 27.4% of +GPU time, so the obvious story is a better kernel. That story is wrong. +Tying two knobs isolated it: `tokens_per_block: 64` and +`max_seq_len: 2048` reach the *same* `max_blocks_per_seq` by different +routes and measured identically (+24.8% vs +25.7%). The win is +**host-side KV block bookkeeping**, which scales with +`max_blocks_per_seq = max_seq_len / tokens_per_block` — not the kernel +variant. Prefer the workload-independent knob. + +**Decode is host-bound, not roofline-bound.** An nsys capture over 100 +pure-decode iterations put GPU idle at **0.401** against a 0.15 +generation threshold, with the forward fully graph-launched +(`cudaLaunchKernel` only 0.9% of the window). The exposed cost was KV +block bookkeeping — `prepare_resources` 3306 µs/step, +`_prepare_inputs` 2443 µs/step. A decode step that reads ~16.9 GB in +11.37 ms is running at ~1.5 TB/s against a ~7.9 TB/s B200 roofline. On +this mechanism, look at the host before the kernels. + +## An fp8-e4m3 latent pool + +Established by `deepseek-r1-0528-nvfp4/sm_100/dep4`, the first target on a +quantized latent cache (`hf_quant_config.json`'s `kv_cache_quant_algo: +FP8`), and by the five catalog certifications it needed. Everything above +describes the bf16 pool; this section is only what changes. + +**The pool is exactly half as wide, and a target declares nothing.** The +per-token width becomes `num_layers * (kv_lora_rank + qk_rope_head_dim) * +1` — measured **35,136 B/token** at 61 layers (`61 * 576`), against +`* 2` in bf16. Every op derives the element width from **`quant_mode` +alone**, because the pool arrives as a raw pointer and nothing else +describes it. That is also the failure mode: a `quant_mode` disagreeing +with the pool's real dtype is silently addressed at the wrong stride, in +either direction. + +**Derive `quant_mode` from the checkpoint's quant config; do not hand-code +the fp8-KV bit.** Bits outside the KV-cache group ride along unread — +`1152` (`FP8_KV_CACHE | FP8_1x128_128x128`) and `384` (`| FP8_QDQ`) are +bit-identical to a bare `128` on every MLA flavor and on the gather. + +**The workspace grows, and by more than the certification figure +suggests.** An engine at `H = 128` / `max_num_tokens = 8192` sized the +attention workspace to **2,126,512,128 B (1.98 GiB) per rank** — about 4x +the 531,628,032 B the bf16 MLA engine above reached. At the op level the +same asymmetry is visible at ~10x smaller scale (52,816,640 B against +43,288,832 B for an identical no-append shape), because the quantized +operands are staged in the workspace. Size from the fp8 figure. + +### The two phases divide the quantization labour *oppositely* + +This is the part nothing in either op's signature says, and it is the +easiest thing to get backwards. + +| | context | decode | +|---|---|---| +| who quantizes the query | **the FMHA itself**, in-kernel, at scale 1.0 | **the preprocessing op**, into `quant_q_buffer` | +| what the preprocessing op writes | plain **bf16** `q` (an ordinary rotation) | `quant_q_buffer` + `mla_bmm1_scale` + `mla_bmm2_scale` | +| are the kv scale tensors read by the FMHA | — | **no, neither** — the dequantization must be folded into the two bmm scales by whoever produces them | + +So `q` is quantized exactly once on the context path, and the decode path +needs `mla_bmm1_scale = [x, x*log2(e)]` with +`x = r²/(q_scaling * sqrt(nope + R))` and `mla_bmm2_scale = [r]`. The FMHA +reads `mla_bmm1_scale[**1**]` — element `[0]` is inert. + +**On an fp8 pool the decode producers are ordered, not concurrent.** The +preprocessing op does **not** write `fused_q` there; it *reads* +`fused_q[..., :C]` to build the quantized query. The absorbed-q BMM must +therefore have finished before it. +`.claude/skills/staircase-kernel-onboard/references/decompositions/mla-custom-op.md` +§4.2 carries both cases — following the bf16 "disjoint halves may overlap" +reading on an fp8 pool is a silent race whose symptom is a quantized query +built from uninitialized memory. + +### Three things that will look like defects and are not + +* **A cached prefix pays fp8 twice.** The engine dequantizes cached latent + rows off the pool, projects them through `kv_b_proj`, and the context + FMHA quantizes the result straight back to e4m3. No argument turns it + off. A cached-prefix context call therefore carries strictly more + quantization error than a fresh prefill of the same tokens — account for + that before attributing an accuracy gap to block reuse. +* **Nothing validates the round trip at any layer.** The write scale, the + read scale and the folded bmm scales are independent roles and no op + checks any relation between them. A pool written at `w` and read at + `r != 1/w` is silently mis-scaled, never an error. +* **The scale must be 1.0, and upstream cannot reach anything else + anyway.** The fp8 MLA *context* path is internally inconsistent at + `s != 1.0` in both its flavors — it quantizes at 1.0 but applies `s²` + and `s` as if it had quantized at `1/s`. Every engine call site + hard-codes the factor to 1.0, so this is unreachable from stock trtllm + and reachable only from a direct op caller, which a target is. Check the + checkpoint: DeepSeek-R1-0528-FP4's 124 `k_scale`/`v_scale` tensors are + all 0-dim fp32 scalars equal to exactly 1.0. + +**JIT cost.** First boot compiles **8 fp8 MLA decode kernels per rank** at +~6-7 s each — the `QkvE4m3` siblings of the bf16 names (`QkvE4m3O +Bfloat16HQk576HV512...VarSeqQ16Kv128...ForGen` plus `MultiCtasKvCga` / +`HVPerCta256`). No MLA *context* call compiles anything, either pool. +`H = 128` takes the same `VarSeqQ16` q-tile as 16 and 32 — the tile stops +growing at the low end — but still pays its own compile. + +**Page size.** Only **32** is certified over an fp8 pool, across all four +MLA entries; their bf16 columns also carry 64. That matters because page +64 is a real perf lever elsewhere (+22% at con=256 on a tp1 sibling), so +on an fp8 target it is a certification request rather than a knob. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/routing-identity.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/routing-identity.md new file mode 100644 index 000000000000..994e6a04a704 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/routing-identity.md @@ -0,0 +1,118 @@ +# Routing and target identity + +How a configuration finds its target in tree, and — the part that decides +whether a proposed target is a target at all — what may and may not be a +routing criterion. + +This replaces the standalone project's registration discussion. Out of tree a +target chose the class name it registered under and patched a stub +`config.json` to match; both are gone. The checkpoint is read exactly as +published, and the decision of which modeling code runs is routing's. + +Source of record: `tensorrt_llm/_torch/_experimental/modeling_v2/README.md` +(§How a config finds its target, §What may decide a target) and +`_router_index.py`. + +## The path a config takes + +``` +LLM(model=...) -> ModelLoader -> AutoModelForCausalLM._resolve_class + | + +-- modeling_v2_resolve(config) + | + _router_index.py +-- architectures[0] -> routing module + models//routing.py +-- one forward-reading decision tree + +-- returns a synthetic class name + | + get_registered_model_class(name) +``` + +Two levels, deliberately: + +- `_router_index.py` maps `architectures[0]` — a string every checkpoint + already declares — to a routing module. Adding a checkpoint, a GPU + architecture or a parallel topology never touches it; only a new + architecture family does. +- `models//routing.py` owns the rest of the decision as **one + forward-reading decision tree per architecture family**. The criteria are + evaluated in the order a reader would ask them, and every branch that does + not end in a target returns None. Reading that single file tells you where + any configuration lands. + +The name it returns (`ModelingV2GptOss120bSm103Tp1`) is a synthetic registry +key that no checkpoint declares; importing the target module is what puts a +class behind it. Upstream already does exactly this for `EAGLE3`, which +also exists only as a `_resolve_class` rewrite. + +The rewrite is local. `pretrained_config.architectures` is **not** mutated — +`_resolve_class` rewrites a local variable — so every other runtime branch +keyed on the architecture string still sees the checkpoint's published name +and allocates accordingly. + +To ask why a configuration landed where it did, replay the same tree: + +``` +python -m tensorrt_llm._torch._experimental.modeling_v2.explain \ + --model /path/to/DeepSeek-R1-0528-NVFP4 --tp 4 --ep 4 --attention-dp +``` + +`TRTLLM_MODELING_V2` is `off` / `auto` / `require`. Anything attributed to +modeling_v2 runs under `require`, and it must be exported **before the ranks +start** — worker ranks receive the environment as it stood when MPI +initialized. Under `auto` a configuration that misses a target's criteria +silently gets the built-in implementation, and a number measured that way +reads as the target's. + +## What may decide a target + +A quantity may be a routing criterion only if all three hold: + +1. it is **known when `_resolve_class` runs**; +2. it is **constant for the engine's life**; +3. it changes the **structure of the forward**. + +| quantity | may route? | why | +|---|---|---| +| config shape (`num_hidden_layers`, `hidden_size`, expert count, …) | **yes** | known at resolve, fixed, and it is what the declared geometry is written against | +| SM version | **yes** | known at resolve, fixed, and kernel families are architecture-gated | +| parallel topology (`world_size`, EP, attention DP) | **yes** | known at resolve, fixed, and it decides which collectives the forward contains | +| quant config, speculative config | **yes** | same three tests: they change which ops the forward calls | +| `max_num_tokens` | **no** | passes 1 and 2, fails 3. A tuning knob — an LLM API argument, not identity | +| `cuda_graph_config.max_batch_size` | **no** | same: per-instance constant, still a tuning knob | +| batch composition, `num_contexts`, "is this step pure decode" | **no** | fails 1 outright. They move every step, so a target chosen from them would be chosen once and then be wrong | + +Per-step specialization is a separate mechanism — dispatch inside a target's +`forward`, on metadata — not a routing dimension. The distinction is the whole +reason routing can be a single decision tree read top to bottom. + +`is_disagg` is declared on the routing context but **not yet plumbed**: nothing +sets it, so it reads False in every deployment. A claim test forbids any +routing module from reading it until it is, because a branch on a value that is +always False takes the wrong side silently. + +## Checkpoint identity is a shape fingerprint + +Routing recognizes a checkpoint by `(num_hidden_layers, hidden_size, …)` — the +same sniffing idiom as `is_mla` / `is_nemotron_hybrid` upstream. It **cannot +tell a fine-tune from the original**: a fine-tune of a routed checkpoint has +the same shape and is routed to the same target, silently. + +That is a deliberate trade, and it changes what a gate record means. A passing +gate does not say "this target passed". It says "this modeling code passed **on +the checkpoint the accuracy gate names**". Running it on any other checkpoint +of the same shape is ungated — which is why the checkpoint path and hash +summary belong in every report, and why the accuracy gate's reference entry, +not the fingerprint, is what pins identity. + +Make the fingerprint carry the fields the target's declared geometry actually +depends on, for the same reason: layer count alone may separate two +checkpoints, but the expert count is what makes the MoE operand shapes correct, +so it belongs in the tuple. + +## Identity is the path + +The target directory's name is the triple, `____/`, +one segment under the family, and the synthetic class name carries the same triple. +`tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py` asserts that the +routing tables and the targets they name agree — it is import-free and runs +anywhere, so a mismatch is caught without a GPU. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/runtime-interface.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/runtime-interface.md new file mode 100644 index 000000000000..0d04b9ca88eb --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/runtime-interface.md @@ -0,0 +1,309 @@ +# Target integration with the TensorRT-LLM runtime + +How a target plugs into the TensorRT-LLM engine: registration, +construction, weight loading, the per-step forward duties, KV-cache +boundaries, CUDA-graph discipline, and the fail-fast ladder. Together +with the catalog contracts and `step-args.md`, this is the +complete integration knowledge for writing a target — no TensorRT-LLM +source reading is needed. + +Everything below was established by probing and verified by +the qwen3-8b/sm_100/tp1 gates (that target is the readable exemplar of +every pattern named here). + + +> **This file is one part of a split.** It was `docs/references/trtllm-runtime-integration.md`; +> its section numbers are kept unchanged so that existing citations — in +> target comments, in the onboarding records — still resolve. The four +> sections that moved: +> +> | former section | now at | +> |---|---| +> | §7 CUDA-graph discipline | `staircase-modeling-core/references/graph-and-sync.md` | +> | §10 What varies per model | `staircase-kernel-feasibility/references/derivation-checklist.md` | +> | §12 Multi-rank launch | `staircase-modeling-parallel/references/multi-rank.md` | +> | §13 One-engine speculative decoding | `staircase-modeling-speculator/references/runtime-boundary.md` | + +## 1. Directory shape + +``` +tensorrt_llm/_torch/_experimental/modeling_v2/ +└── models// + ├── routing.py # which config lands on which target + └── ____/ # identity = the directory's name + ├── modeling.py # registration shell + core + step-args + fail-fast + └── weights.py # MANIFEST + load loop + post-load derivations +``` + +That is the whole target. There is no stub model directory, no +`llm_args.yaml`, no `smoke.py` and no `TARGET.md`: **the checkpoint is read +exactly as published**, the topology comes from the routing decision and the +boot gate's CLI flags, and gate records go in the run journal. `LLM(model=...)` +is pointed at the real checkpoint, and everything the runtime learns about the +model it learns from that checkpoint's own `config.json`. + +## 2. Registration and resolution + +- The target's shell class carries `@register_auto_model()`, + which writes the class into trtllm's process-global registry at import time. + **The name is routing's, not the target's**: `models//routing.py` + decides it and lists the module in `TARGET_MODULES`, and + `modeling_v2_resolve` imports that module once its tree has claimed the + config — nothing else would put a class behind the synthetic key. See + `routing-identity.md`. +- The import must precede engine construction in **every process that + builds the model** — which above world size 1 is not this one. An ordinary + dotted module satisfies that by itself: each rank resolves the config, the + routing module imports the target, and pickle resolves the class by module + name. §12 has the rest of what changes there. +- `LLM(model=)` then resolves the configuration through + `AutoModelForCausalLM._resolve_class`, which consults + `modeling_v2_resolve` and looks the returned name up in the ordinary + registry, and construction begins. + +### The runtime still keys behaviour on the architecture string — and it reads the checkpoint's + +Parts of the runtime branch on `architectures[0]` (and on other published +config fields) to decide **what to allocate**, not merely which class to +build. The known case is hybrid linear attention: a recurrent-state pool is +allocated only when the predicate matches, and otherwise the model builds, the +boot gate passes, and the linear layers have no state. + +In tree that predicate reads the **checkpoint's own** config. The synthetic +name exists only as a local rewrite inside `_resolve_class`; +`pretrained_config.architectures` is not mutated, so nothing downstream sees +it. Two consequences: + +- a target **cannot** change what the runtime allocates by choosing its + registered name, and must not try. If the runtime's predicate does not match + the published config, that is a fact about the checkpoint and the runtime — + report it as a symptom rather than routing around it; +- conversely the published name is always the honest one, so the failure mode + the standalone project had to defend against (a stub `config.json` declaring + a name the runtime did not recognize) cannot occur. + +The registry facts behind this are unchanged and still worth knowing, because +they are what makes a target's registration win its slot: + +- `MODEL_CLASS_MAPPING` is consulted **before** the lazy static index + `_torch/models/_arch_index.py`, so a name already in the mapping never + triggers the built-in module's import. Verified by resolving + `Qwen3_5MoeForCausalLM` — a name that index maps to `modeling_qwen3_5` — + with a locally registered class in place: `get_model_architecture` returned + the local class and the built-in module was never imported. +- `register_auto_model` **overwrites silently** on a duplicate name; there is + no exception and last registration wins. A synthetic name no checkpoint + declares is what keeps that from mattering. +- `executor/worker.py` does `MODEL_CLASS_MAPPING.update(**_torch_model_class_mapping)`, + which is how the registration crosses into the executor worker process. + +*Re-check the precedence by registering a throwaway class under a name present +in `_arch_index`'s `MODEL_ARCH_TO_MODULE` and confirming +`get_model_architecture` still returns it; the precedence, not the name list, +is what matters.* + +## 3. The shell and the core + +```python +class ModelingV2Core(DecoderModel): # the computation + def __init__(self, model_config): ... # geometry asserts + weights + def forward(self, attn_metadata, input_ids=None, position_ids=None, + inputs_embeds=None, lora_params=None, **kwargs): ... + +@register_auto_model("ModelingV2") +class ModelingV2(DecoderModelForCausalLM[ModelingV2Core, ...]): + def __init__(self, model_config): + super().__init__(ModelingV2Core(model_config), config=model_config, + hidden_size=..., vocab_size=...) + def load_weights(self, weights, *args, **kwargs): ... + def post_load_weights(self): ... +``` + +- The shell (`DecoderModelForCausalLM`) is composition: it stores the + core as `self.model`, creates `self.lm_head`, and owns the logits path + — its forward calls the core, then gathers exactly the rows that need + logits (last token per context sequence + every generation token) and + applies lm_head. **The core returns final-normed hidden states + `[num_tokens, hidden]` and never touches logits.** +- Inputs are **packed**: first dimension = total tokens of the batch, + context sequences first. `position_ids` arrives as an int32 `[1, T]` + view of an engine buffer — flatten with `reshape(..., [-1])`. +- Construction asserts the target's identity — geometry, dtype, and every + config property the assembly depends on (tie_word_embeddings, qk-norm + presence, rope scaling, sliding window, attention bias). Assert, never + adapt: a mismatched checkpoint is a different target. These asserts and the + routing tree's criteria are two readings of the same fact; routing declines + the configuration, the asserts refuse it if it arrives anyway. +- Inputs the target does not implement (`lora_params`, `spec_metadata` + in kwargs) must **assert loudly** — silently ignoring them produces + wrong output with no signal. This is distinct from runtime-owned + feature flags, which pass through (see §6). A target that *does* + implement speculative decoding stops asserting on `spec_metadata` and + grows a second forward it owns; §13 is that axis. + +## 4. The weight lifecycle + +| stage | who | what happens | +|---|---|---| +| t0 construct | engine (`MetaInitMode`) + your `__init__` | `torch.empty` is intercepted to the meta device: parameters are shape/dtype shadows, zero memory. Declare with `torch.empty` only (`zeros`/`full` are NOT intercepted) and do no tensor math in `__init__` (meta tensors raise on compute). | +| t1 materialize | engine | walks **registered** parameters (`nn.Parameter` / `ParameterDict` / module tree — what `named_parameters()` can see) and reallocates each on CUDA, contents garbage. Registration is what requests the allocation; anything held in a plain dict/list stays a dead shadow. | +| t2 read | engine | reads every safetensors in the checkpoint directory into a dict `{ckpt_key: tensor}`. The values arrive **already materialized** as `torch.Tensor`, not as lazy slices — measured at world size 4, where every rank is handed the whole dict (25,687 keys on all four). So the engine pre-shards nothing: a multi-rank split is entirely the manifest's, and a `src` transform's benefit is that only this rank's bytes cross to the device, not that fewer bytes are read off disk. | +| t3 load | your `load_weights(weights)` | full delegation, never validated by the engine — the manifest loop copies bytes into the materialized storage (`param.data[a:b].copy_(src)`). | +| t4 derive | your `post_load_weights()` | the designated home for derived state: `.t()` GEMM views, per-layer tuples, and checkpoint-calibrated call tensors such as fp8 KV scales. Real tensors may be created here — meta is over. | +| t5 serve | your `forward` | reads weights by reference; zero copies or transforms on the hot path. | + +Storage layout: declare parameters in **HF `[out, in]` row-major** so t3 +copies are layout-preserving, and consume GEMMs through `.t()` views +(row-major `[N, K]` transposed is exactly the dense column-major `[K, N]` +that `cublas_mm`'s contract requires; the view is zero-copy and shares +storage, so reloading weights in place keeps views valid). + +## 5. The weight manifest convention + +`weights.py` owns a pure data table: + +```python +MANIFEST[param_key] = [(ckpt_key, dst_slice_or_None), ...] +# tp1: no src transform. Sharded targets add one: +# (ckpt_key, src=col_shard(rank, tp), dst_slice) +# and declare per-rank shapes in modeling. +``` + +Rules: fused parameters (qkv, gate_up) fill by destination row slices — +no intermediate concat buffers; assert shape and dtype per copy; assert +**bidirectional coverage** (manifest keys == declared parameters before +the loop; consumed ckpt keys == the whole checkpoint after it, minus the +explicitly predicted non-load set). The +shell-registered exceptions (`lm_head.weight`, and `embed_tokens` when +the base class owns it for tied checkpoints) are fed explicitly. + +## 6. Forward duties, KV cache, and feature flags + +The engine drives everything around the forward: scheduling, KV block +bookkeeping (allocation at admission, per-token growth, prefix reuse), +metadata filling and `prepare()` — all before your forward runs. Your +duties per step: + +1. Build the attention step-args once from the prepared metadata and + share them across layers (`step-args.md` is the + normative mapping). +2. Keep the loop flat: catalog calls, tensor-metadata reads, Python + control flow — nothing else (the closed-vocabulary rule; audit = + collect the calls in `forward` **and the private methods it reaches** + and match them against + `tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml`. Scope + to that closure: a + whole-file scan flags the rope table's `arange`/`cos`/`sin` and + load-time `.t()`/`.to()`, which run at init and are outside the rule). +3. Let one attention call serve any batch composition + (`attention_input_type=0` for the standard configuration — no phase + branches; MLA targets dispatch per phase by metadata reads). + +The KV-ownership red line does not move: the pool is sized by the engine from +the **checkpoint's own** `config.json` declarations (`num_hidden_layers`, +`num_key_value_heads`, `head_dim`, `torch_dtype`, `vocab_size`, +`max_position_embeddings`), pages are assigned by the C++ manager, and the +new tokens' K/V are written **inside** the attention op. The target only +forwards the address book — it restates none of those fields and must not +contradict them; its construction asserts read the same config the engine +sized the pool from. + +Runtime-owned feature flags (block reuse, CUDA graphs, beam, spec-dec +plumbing) **pass through** from metadata exactly as prepared: the target +runs trtllm defaults, feature behavior is upstream's responsibility, and +the gates validate the result. Kernel-surface certification accounting +lives in the catalog contracts, not in target config. + +## 8. The fail-fast ladder + +| fuse | when | checks | catches | +|---|---|---|---| +| static op declaration | CI, `tests/unittest/_torch/modeling_v2/test_modeling_v2_target_contract.py` | every `torch.ops.trtllm.*` symbol named in the target's `REQUIRED_TRTLLM_OPS` is registered in this build; the thop binding exists | missing/renamed ops in a drifting engine. torch mirrors (`embedding`, `empty`, `reshape`, ...) are exempt — core PyTorch API, upstream-owned | +| step contract | first forward | metadata is `TrtllmAttentionMetadata`; every `_STEP_FIELDS` name exists; `position_ids` is int32 | private-surface layout drift the op list cannot see. Everything checked is fixed at engine construction — once per model instance is sound | +| gates | explicit runs | boot-gate keywords; accuracy score vs `tests/integration/defs/accuracy/references/.yaml` | the computation itself | + +There is no version pin to assert against — a target moves with the trunk, so +the assert would be meaningless — which is exactly why the first two rows +carry the weight they do. Keep `_STEP_FIELDS` exactly equal to the set of +metadata attributes the code reads, and `REQUIRED_TRTLLM_OPS` exactly equal to +the trtllm ops the forward calls: the lists are dependency declarations, and a +stale entry is as wrong as a missing one. + +**One metadata surface is conditional, and a flat existence check over it +turns a legal config into a hard failure.** MLA's cached-KV fields — +`enable_context_mla_with_cached_kv`, `ctx_cached_token_indptr`, +`ctx_kv_indptr`, `ctx_uncached_token_indptr`, `max_ctx_seq_len`, +`max_ctx_kv_len`, `num_ctx_cached_tokens` — exist only when block reuse +is on. With `enable_block_reuse: false` they are **absent, not False**, so +a `_STEP_FIELDS` tuple containing them fails the first forward on a +configuration the target is supposed to support. Split the tuple: the +unconditional fields keep the existence check, and the conditional ones +select the context flavor by their *presence* rather than being asserted. +See `.claude/skills/staircase-modeling-core/references/latent-kv-cache.md`. + +## 9. Gates and launch + +- **Every GPU command in a run is launched one way: the run's launch prefix.** + It carries the environment the ranks need — `TRTLLM_MODELING_V2` above all, + which must be exported before the ranks start and is not yours to set — and + the declared interpreter. Do not construct a launcher, and do not probe for + one. The devices claimed for the run are given to you: one for a `tp1` + target, `N` for a `tepN`/`depN` one. +- Boot gate: `examples/llm-api/quickstart_advanced.py` with the target's + topology flags and `TRTLLM_MODELING_V2=require`, run by the orchestrator. + Engine cold start, weight-manifest coverage, greedy keyword asserts; the + exit verdict is binary. The cases are target-authored: pick prompts whose + greedy continuation is high-confidence for this model, and **verify every + keyword on the real model before freezing it** — an unverified keyword bakes + a false failure into the gate, which is why the keyword list starts empty and + the first engine run doubles as its verification. +- Release: `trtllm-eval` under the run's launch prefix, with the protocol from + `tests/integration/defs/accuracy/references/.yaml` and + `TRTLLM_MODELING_V2=require` — one-sided gate, + measured >= reference − tol. On the MMLU + path — which is where an entry with no `eval` block lands — add + `--dataset_path /home/scratch.trt_llm_data/llm-models/datasets/mmlu`, + because trtllm-eval's own download 404s and fails as a corrupt tar rather + than as a missing file. + +## 11. Serving-time costs a target inherits + +The engine owns these, not the forward, but they land on the target's +Pareto curve and the first two are worth checking on every new target +before any modeling work. + +**Decode CUDA-graph coverage defaults to batch 128.** Above it the whole +decode phase runs eager. The deeper the model the worse this is: measured +on a 48-layer target, concurrency 256 lost 19.7% throughput and 151% TPOT +against its own graph-covered configuration, and raising +`cuda_graph_config.max_batch_size` to 256 with `enable_padding: true` was +worth +56.5% at that point with every other point unchanged. This should +be the first config experiment on any new target — but raise the ceiling +and A/B the padding **separately**: on a later target `enable_padding: +true` measured 11% *worse* at concurrency 256 against the same raised +ceiling, because it swaps the fine `[1..32]` batch grid for a coarse one. + +**The stock decode capture grid is the fine one.** `enable_padding` defaults to +**False**, so a stock engine captures `[1..32, 64, 128]` — 34 batch sizes, every +integer up to 32 — not the sparse padded grid `[1, 2, 4] + [8, 16, …, 128]` that +padding produces. This is the list a certified-column audit must be read against, +and assuming the padded one leaves the odd small counts (6, 9–15, 17–23, 25–31) +uncovered while looking complete. Enabling padding also makes `192` reachable +only when `max_batch_size >= 192`. + +**`max_seq_len` is a per-step cost, not only an admission cap.** It sizes +`max_blocks_per_seq = max_seq_len / tokens_per_block`, which sizes a +pinned `[1, num_seqs, 2, max_blocks_per_seq]` int32 block-offset staging +buffer that the resource manager allocates, memcpys and pushes H2D **every +step**. Left at a model cap far above the served workload (40960 vs a +1024/1024 benchmark) that is 2.62 MB per step at 256 in flight. Capping it +is a real knob — but size the expectation: a measured 10× shrink moved +`_prepare_inputs` only 5.17 → 4.44 ms/step, so the staging is a minority +of that cost and the residual is `O(num_seqs)` per-request executor +Python. Capping also **restricts capability** (longer requests are +rejected), and the floor is set by the accuracy gate's own prompts, not by +the benchmark — 5-shot MMLU prompts reach 2687 tokens, so a cap at the +benchmark's exact ISL+OSL makes the gate unrunnable. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/sliding-window-attention.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/sliding-window-attention.md new file mode 100644 index 000000000000..4dd98d85ea27 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/sliding-window-attention.md @@ -0,0 +1,79 @@ +# Sliding-window attention, alternating with full attention + +A model whose layers alternate between a short attention window and full +causal attention. Written from the gpt-oss-120b/sm_100/tp1 onboard +(36 layers, `layer_types` alternating `sliding_attention` / +`full_attention`, `sliding_window: 128`); the mechanism recurs across +families (Gemma, Mistral, Ministral). + +## The window is a mask, not an addressing mode + +This is the fact everything else follows from, and the one an upstream +comment gets wrong. TensorRT-LLM's attention op appends every new token +at its **absolute** position — entry `t // tokens_per_block` of the +sequence's block-offsets row, slot `t % tokens_per_block` — and wraps +nothing, not at the window and not at `max_seq_len`. `attention_window_size` +only restricts which cached keys a query may read. Any cyclic reuse of +pool memory is the *caller's* page mapping. + +(`_torch/auto_deploy/custom_ops/attention_interface.py` and +`.../shim/ad_executor.py` state that the kernel "cyclically indexes the KV +cache internally" / "wraps KV reads modulo the attention window". +Measured on sm_100, it does not. The operational contract +those comments derive — full block table, global KV length — is +nonetheless correct, so nothing downstream of them is broken; but sizing a +pool from the stated mechanism would be.) + +**Boundary:** a query at absolute position `p` attends to +`[p - W + 1, p]` — exactly `W` keys, never `W + 1`. `W` is a token count, +not a page count, and may be shorter than one page. + +## What a target owes + +- The block-offsets row still needs `ceil(total_kv / tokens_per_block)` + entries, window or no window. Only entries holding at least one + in-window token must point at live pages: aging the older ones out to + garbage leaves the decode output bitwise identical. +- `sequence_length` / `host_past_key_value_lengths` must stay the + **global** cached+new count. It is both the window origin and the append + position, so capping it to `W` moves the write to slot `W - 1` and + truncates the attended range — silently. `context_lengths` and + `max_seq_len` are inert on generation rows. +- Prefill is **not** shrunk by the window. One call's new tokens must land + in distinct physical slots, or the writes race and leave torn cache rows + nondeterministically, with no error. +- A caller-side page ring (`ring[j % P]`) does make the pool wrap, bounded + by `P * tokens_per_block >= W`. Below that bound the result is silently + wrong rather than rejected. + +## The route that works, and what it does not buy + +**Per-layer windows over one shared pool are the correctness route.** +`attention_window_size` is a per-call scalar, so alternating layers can +address one pool, one layer→pool mapping and one block-offset table +inside a single batch. Because a target built on the catalog does not +construct the engine's own attention modules, **the runtime is never told +about the window at all**: the KV manager sizes one pool for full +attention and returns an identity `host_kv_cache_pool_mapping`. The +per-layer window is then a pure modeling-side argument needing no +`llm_args` support and no page-ring work. Both FMHA kernel families +(`...PagedKvDense...` and `...PagedKvSlidingOrChunkedCausal...`) appear at +warmup. What picks between them is the batch's longest KV run against the +window, not the layer's window setting: a windowed layer whose history +still fits inside the window takes the *dense* kernel and crosses over on +the first step past it, so one sequence pays both JIT compiles over its +lifetime (`thop_attention`, *Notes*). + +**It saves no memory.** The shared pool is still sized for the +full-attention layers, so the alternation costs exactly what an +all-full-attention model would. A real footprint reduction needs +per-window pools, i.e. `num_pools > 1` addressing. + +Before spending a certification on that: on gpt-oss-120b it was measured +*not* to be the win. KV capacity never bound (36.3% of the pool used at +concurrency 256), the per-layer window already cuts the FMHA read cost +where it matters (23.9 µs sliding vs 100.3 µs dense per layer at +concurrency 256), and stock trtllm sizes the **same** single +full-window pool on that checkpoint — so it is not a difference between +the two systems at all. Check those three things before assuming +multi-pool addressing is worth the run. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/step-args.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/step-args.md new file mode 100644 index 000000000000..3d073a0bad75 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/references/step-args.md @@ -0,0 +1,178 @@ +# thop_attention runtime argument sourcing + +How a target's modeling code fills +`catalog/attention/thop_attention`'s +arguments when running under the TensorRT-LLM engine. The entry contract +(`tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.md`) +defines what the op computes and +what each argument must contain; this reference defines **where each value +comes from at runtime**. Together they are sufficient — no TensorRT-LLM +source reading is needed to integrate the op. + +Provenance: the mapping mirrors the in-tree `FallbackFmha.forward` +argument sourcing (`tensorrt_llm/_torch/attention/backends/fmha/fallback.py`) +and is verified end to end by the qwen3-8b/sm_100/tp1 gates (boot gate; full +MMLU 74.98, bit-identical across eager and CUDA-graph configs). Whenever the +engine's attention backend moves, re-derive this table against that function +and re-run the gates. + +## The three-source split + +Every argument of one attention call belongs to exactly one source: + +```python +thop_attention( + q=qkv, + output=attn_out, # 1. per-layer: data tensors + local_layer_idx=i, # + the layer index + num_heads=Hq, + num_kv_heads=Hkv, + head_size=D, # target geometry + **step, # 2. per-step: engine batch state, + # built once per forward, shared + # by every layer + **CALL_CONSTANTS, # 3. per-target constants: the call + # shape, fixed at import time +) +``` + +Build `step` once per forward from the engine-prepared metadata and reuse +it across layers — the values describe the batch, not the layer. Per +layer, `q` is that layer's packed-QKV activation (the fused projection's +output) and `output` is a caller-allocated `[num_tokens, Hq*D]` buffer: +the contract guarantees rows `[:num_tokens]` are overwritten in place, so +one buffer allocated per forward serves every layer. + +## Per-step arguments (engine metadata) + +The engine constructs one `TrtllmAttentionMetadata` (the default TRTLLM +attention backend), fills it every iteration, and calls `prepare()` before +the model forward; every field below is valid only after that. The model +receives the object as the `attn_metadata` forward argument. + +Graph classes: **ref** = engine-owned persistent buffer, refreshed in +place each step (CUDA-graph replay reads fresh contents automatically); +**host** = Python scalar, baked into a capture (safe: decode-only graphs +always see `num_contexts == 0`); **cfg** = engine-construction constant. + +| thop argument | metadata source | class | notes | +|---|---|---|---| +| `sequence_length` | `kv_lens_cuda_runtime` | ref | per-seq total KV (cached + new), CUDA int32 view | +| `host_past_key_value_lengths` | `kv_lens_runtime` | ref | same values, CPU | +| `host_total_kv_lens` | `host_total_kv_lens` | ref | `[2]`: context sum, generation sum | +| `context_lengths` | `prompt_lens_cuda_runtime` | ref | | +| `host_context_lengths` | `prompt_lens_cpu_runtime` | ref | | +| `host_request_types` | `host_request_types_runtime` | ref | 0 = context, 1 = generation | +| `kv_cache_block_offsets` | `kv_cache_block_offsets` | ref | GPU page table, layer-agnostic | +| `host_kv_cache_pool_pointers` | `host_kv_cache_pool_pointers` | cfg | pool base addresses | +| `host_kv_cache_pool_mapping` | `host_kv_cache_pool_mapping` | cfg | row per layer: (pool, layer-in-pool) | +| `workspace_` | `effective_workspace` | ref | engine-owned scratch; the op resizes it in place | +| `tokens_per_block` | `tokens_per_block` | cfg | | +| `max_num_requests` | `max_num_requests` | cfg | | +| `max_context_length` | `max_context_length` | cfg | | +| `max_seq_len` | `max_seq_len` | cfg | | +| `attention_window_size` | `max_seq_len` | cfg | no sliding window: window = full length | +| `num_contexts` | `num_contexts` | host | | +| `num_ctx_tokens` | `num_ctx_tokens` | host | | +| `trtllm_gen_jit_warmup` | `trtllm_gen_jit_warmup` | host | engine warmup flag, pass through | +| `use_paged_context_fmha` | `use_paged_context_fmha` | cfg | cached-context path selector (block reuse / chunked prefill) | +| `beam_width` | `effective_beam_width` | cfg | | +| `cache_indirection` | `cache_indirection` | cfg | None unless beam search | +| `block_ids_per_seq` | `block_ids_per_seq` | cfg | None unless FlashMLA | +| `max_context_q_len_override` | `max_context_q_len_override` | cfg | | +| `is_cross` | `is_cross` | cfg | | +| `is_spec_decoding_enabled` | `is_spec_decoding_enabled` | cfg | | +| `use_spec_decoding` | `use_spec_decoding` | host | | +| `is_spec_dec_tree` | `is_spec_dec_tree` | cfg | | +| `spec_decoding_generation_lengths` | same name | ref | None with spec-dec off | +| `spec_decoding_position_offsets_for_cpp` | same name | ref | " | +| `spec_decoding_packed_mask` | same name | ref | " | +| `spec_decoding_bl_tree_mask_offset` | same name | ref | " | +| `spec_decoding_bl_tree_mask` | same name | ref | " | +| `spec_decoding_target_max_draft_tokens` | `max_total_draft_tokens` | cfg | note the rename | +| `spec_bl_tree_first_sparse_mask_offset_kv` | same name | ref | None with spec-dec off | +| `num_sparse_topk` | `num_sparse_topk` | cfg | | +| `flash_mla_tile_scheduler_metadata` | same name | ref | None outside FlashMLA | +| `flash_mla_num_splits` | same name | ref | " | + +Pass runtime-owned values through exactly as prepared — feature behavior +(block reuse, CUDA graphs, speculative decoding plumbing) is the engine's +responsibility and is validated by the target's gates, not second-guessed +in modeling code. + +## Per-target constants + +Everything the in-tree path reads from the attention module or its +forward args is, for one target, a compile-time constant: it encodes +*which kind of attention call this target makes*. For a bf16 packed-QKV +causal GQA target with RoPE applied outside the op (the +`fused_qk_norm_rope` pairing), the constants are exactly the certified +"Standard" column of the entry contract: + +- `is_fused_qkv=True`, `k=None`, `v=None`, `update_kv_cache=True`, + `predicted_tokens_per_seq=1`, `attention_input_type=0` (one call serves + pure-prefill / pure-decode / mixed batches), `mask_type=1` (causal), + `is_mla_enable=False`, `quant_mode=0`, `q_scaling=1.0`. +- In-kernel RoPE disabled: `position_embedding_type=0` and every `rope_*` + argument at the contract's inert placeholder. **The placeholders are not + the model's rope config** (`rope_base=10000.0` even when the model's + theta is 1e6) — rotation happens outside, the values are dead. +- All quantization scales, MLA arguments, mRoPE / cross / relative-bias / + sparse / helix arguments at their contract-listed inert values. + +MLA targets take their per-phase constants from the contract's MLA +columns the same way (separate context/generation call shapes, +`attention_input_type` 1/2); the per-step table above applies unchanged — +both configurations consume the same prepared metadata. + +Two forward-looking notes: `helix_*` are metadata-sourced in-tree but are +None whenever helix CP is off — a non-CP target may hold them as +constants. Quantized-KV targets turn `kv_scale_orig_quant` / +`kv_scale_quant_orig` from `None` constants into small fp32 tensors built +in `post_load_weights` (checkpoint-calibrated; leaving them `None` under +fp8 KV does **not** fault — measured, it is read as a silent `1.0`, so a +miscalibrated target is a wrong-number bug rather than a crash) — see the +contract's not-certified annotations for what a certification extension +must cover first. + +## Fail-fast discipline + +The metadata attribute names above are private trtllm surface with no +stability promise. A target consuming them must: + +1. Keep a `_STEP_FIELDS` tuple in modeling code, exactly equal to the set + of metadata attributes read, and `hasattr`-check every name on the + first forward (everything is fixed at engine construction — once per + model instance is sound). Version drift then fails with a named list + instead of a deep AttributeError. **Exception: MLA's cached-KV fields + are conditional.** `enable_context_mla_with_cached_kv`, + `ctx_cached_token_indptr`, `ctx_kv_indptr`, + `ctx_uncached_token_indptr`, `max_ctx_seq_len`, `max_ctx_kv_len` and + `num_ctx_cached_tokens` exist only when block reuse is on — with + `enable_block_reuse: false` they are absent, not False. Asserting them + turns a legal config into a hard failure; select the context flavor by + their presence instead. +2. Assert `isinstance(attn_metadata, TrtllmAttentionMetadata)` — this + table is meaningless for other backends' metadata. +3. Assert `position_ids.dtype == torch.int32` and flatten with + `reshape(position_ids, [-1])` before ops that consume positions (the + engine passes an int32 `[1, T]` view of a persistent buffer). + + **Not on an mRoPE config, where that flatten is the bug.** If the config + declares mRoPE, the engine builds the `[3, 1, T]` layout **even for a + text-only batch** — the rank is keyed on `rope_scaling['type'] == 'mrope'` + alone and is held constant between warmup and serving — and broadcasts the + scalar positions across all three axes, so row 0 is the text position. + Flattening the whole buffer yields a `[3T]` **copy**, which the op rejects, + and the error names contiguity rather than shape: `position_ids must be + contiguous`. Take row 0 for a text-only path, or pass the `[3, T]` view the + op documents for mRoPE. A text-only target built from a multimodal + checkpoint hits this even though nothing about its workload is multimodal. +4. Declare every `torch.ops.trtllm.*` symbol the forward calls in the + target's `REQUIRED_TRTLLM_OPS`, which + `tests/unittest/_torch/modeling_v2/test_modeling_v2_target_contract.py` + asserts this build registers. There is no version to pin against — the + target moves with the trunk — so that declaration plus the first-forward + `_STEP_FIELDS` check are the whole fuse. A receipt, likewise, is keyed by + GPU architecture and is valid only against the files of its own entry; + absent means unknown. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-disagg/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-disagg/SKILL.md new file mode 100644 index 000000000000..c97c288aab2e --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-disagg/SKILL.md @@ -0,0 +1,129 @@ +--- +name: staircase-modeling-disagg +description: Load when a target must work under disaggregated serving — prefill on one engine, decode on another, with per-sequence state transferred between them. The modeling duty is that EVERY piece of per-sequence state prefill produces is in the transferable set, which for a hybrid model is not just the attention KV pages. +--- + +# staircase-modeling-disagg + +`--features disagg`. Context and generation run as **separate engines**. What +prefill computed has to reach the decode engine over the cache transceiver, +and whatever does not reach it is silently recomputed as zero. + +This is a serving topology, so it is tempting to read it as an operations +concern. It is not, and the reason is specific: + +> **A single-engine target can keep per-sequence state anywhere and still pass +> every gate.** Under disaggregation, only state the transceiver knows about +> survives the hop. + +That makes disagg support a property of the *modeling code*, and the only +feature in this workflow whose defect mode is "correct in every single-engine +measurement, wrong across the hop". + +## The one question to answer first + +**Enumerate every tensor whose value at the end of prefill affects the first +generated token, and say where each one lives.** + +For an ordinary attention-only model the answer is short: the paged KV cache, +and the transceiver moves it. For this run's checkpoint it is not: + +| state | produced by | in the KV pool? | +|---|---|---| +| attention K/V pages | the 12 `full_attention` layers | yes | +| GDN conv state (`number_of_conv_states: 3`, kernel 4) | the 36 `linear_attention` layers | **has to be — check** | +| GDN recurrent/SSM state (fp32) | the 36 `linear_attention` layers | **has to be — check** | +| sparse-attention indexer state (`indexer_budget: 2048`) | the 12 attention layers | **has to be — check** | +| MTP draft-layer state | the speculator, if present | ask separately | + +The three marked rows are the whole feature. If any of them is held in a module +attribute, a Python-side buffer, or a cache the transceiver does not enumerate, +the target is *not* disagg-capable however green its single-engine gates are. + +`tensorrt_llm/_torch/pyexecutor/resource_manager.py` is where a state plane +either is or is not part of the managed pool. Start there, and answer per state +kind rather than per model. + +## What the runtime gives you + +| surface | what it is | +|---|---| +| `tensorrt_llm/_torch/disaggregation/` | the transceiver: `KvCacheTransceiver`, `create_kv_cache_transceiver`, backends under `native/` and `nixl/` | +| `tests/unittest/_torch/disaggregation/` | the contract tests, including `test_kv_cache_transceiver_contract.py` and a `fake_kv_cache_transceiver.py` worth reading before writing anything | +| `examples/disaggregated/disagg_config.yaml` | the deployment shape: `context_servers` / `generation_servers`, each with its own parallel sizes and `cache_transceiver_config` | +| `docs/source/features/disagg-serving.md` | the feature doc | + +Nothing there is model-specific. The transceiver moves what the resource +managers registered; it has no way to discover state you kept elsewhere. + +## Device budget + +A disagg run needs **at least two engines**, therefore at least two devices, +*even when the target's own parallel segment is `tp1`*. A target declared +`tp1` and a disagg witness needing two GPUs is not a contradiction — the +segment describes one engine. + +Say so in the plan. A run that assumed one device and then could not place the +generation server reads as a scheduling failure and is a specification one. + +## Preconditions that hang instead of raising + +Shared with `staircase-modeling-dwdp`, and worth re-checking here because a +disagg launch that stalls looks exactly like a modeling deadlock: + +- the **re-executed interpreter** must be the one the run declares. The disagg + leader re-execs a literal `python3`; if that resolves elsewhere the workers + come up without trtllm and the launch stalls with no message; +- the **cache transceiver must be usable**. Pick the backend explicitly rather + than taking a default, and confirm it in the log rather than assuming; +- the **single-process worker flag** applies to each engine's ranks. + +Turning each of these into a named boot failure is the boot gate's job. That is +what a minutes-scale gate is for. + +## Gates + +boot (aggregated) → accuracy (aggregated) → **equivalence witness +(disaggregated)**, all under `TRTLLM_MODELING_V2=require`. + +The witness is the only measurement that can see the feature, and it is a +*comparison*, not a smoke test: + +- same prompts, **greedy** decoding, same seed, same `max_tokens`; +- leg A: one engine, aggregated; +- leg B: context server + generation server, everything else identical; +- a code path the disaggregated engines refuse or guard off (a bounded replay, + say) is pinned to the same setting on BOTH legs with the block's `legs_env` + (leg B's launcher gets it as `STAIRCASE_WITNESS_LEGS_ENV` and applies it to + both engines); otherwise the legs differ in that path as well as in the hop; +- one request at a time on both legs (`max_batch_size: 1` in the witness + block; leg B gets it as `STAIRCASE_WITNESS_MAX_BATCH_SIZE`): an aggregated + engine and a disaggregated pair never pack rows alike, and a row's bits + depend on the rows beside it; +- assert the outputs are **token-identical**. + +Not "similar", not "both coherent". Missing recurrent state usually yields text +that is fluent and wrong, and a short benchmark under a 5-point tolerance will +not catch it. Token identity is the property that distinguishes "the state +crossed" from "the state was recomputed from zero and the model recovered". + +Run both legs **in the same session**. A cross-session comparison carries +session variance into a judgement that is supposed to be exact. + +### Prefix length matters + +Use at least one prompt long enough to fill several KV blocks *and* to put the +recurrent state well away from its initial value — a 2048-token prompt on this +checkpoint exercises `indexer_budget` too. A 10-token prompt can be +token-identical with the recurrent state dropped entirely. + +## Report additions + +- the state enumeration above, filled in, with the file and symbol that owns + each row +- the transceiver backend actually used, quoted from the log +- the two-leg comparison: prompt lengths, the number of tokens compared, and + whether identity held +- which of the three hang-preconditions the boot gate now asserts +- device count used by the witness, stated next to the target's `tp1` segment + so nobody reads them as contradictory diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-dwdp/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-dwdp/SKILL.md new file mode 100644 index 000000000000..21e7a47367f7 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-dwdp/SKILL.md @@ -0,0 +1,77 @@ +--- +name: staircase-modeling-dwdp +description: Load when a MoE target must hold only part of the expert table locally and let the runtime bind the rest over NVLink — distributed weight data parallelism. An adapter-class feature whose forward cannot be CUDA-graph captured, and whose launch has three preconditions the runtime does not assert. +--- + +# staircase-modeling-dwdp + +`--features dwdp`. Each rank loads a **fraction** of the routed experts; the +runtime binds the full table per layer over MNNVL, so the forward sees all of +them. On a 256-expert checkpoint over a group of four, each rank holds 64. + +An adapter-class feature — read +`staircase-modeling-core/references/adapter-rules.md` first. + +| reference | read it when | +|---|---| +| `references/runtime-boundary.md` | always — the ownership split, and the constraints it imposes on the forward | +| `references/adapter-interface.md` | writing `dwdp_adapter.py`, or placing the bind and prefetch calls | + +## What this feature changes + +**The manifest loads a window.** Only this rank's expert range is materialized; +the rest of the checkpoint's expert keys go into the known-not-loaded set with +DWDP as the recorded reason. + +**The forward binds before the MoE call and prefetches the next layer after +it.** Two hooks, and their placement is the feature: bind too late and the +expert weights are not addressable; prefetch too early and you evict what the +current layer still needs. + +**The forward is not CUDA-graph capturable.** The runtime's prefetch is a +double-buffered state machine driven from the host between layers, which a +captured graph cannot express. This is a property of the feature, not a +tuning choice — it goes in the run journal beside the identity record, because +a reader who assumes capture will misread every subsequent perf number. + +## Three boot requirements that are not asserted anywhere + +A DWDP launch fails by hanging, not by raising, so these are worth checking +before concluding anything about the modeling code: + +- **the single-process worker flag** must be set for the TP1 ranks +- **the re-executed interpreter must be the one the run declares.** The disagg + leader re-execs a literal `python3`; if that resolves to a different + interpreter, the workers come up without trtllm and the launch stalls with + no message +- **the cache transceiver must be usable.** The MPI transceiver crashes or + stalls; UCX works but needs `libzmq` loadable and its transports pinned to + same-node + +Nothing in tree asserts any of them. The **boot gate** for this target should — +it is the minutes-scale gate, and turning a silent hang into a named failure is +exactly what it is for. + +## Gates + +boot with the feature's arguments → accuracy, both with +`TRTLLM_MODELING_V2=require` exported before the ranks start. **No third +gate.** + +Accuracy sees this feature: each rank holds a fraction of the expert table, so +a binding that does not deliver the remote experts cannot silently produce the +same answer — either the forward fails, or it computes over a different expert +set and the score moves. + +What the gates do *not* tell you is whether the remote path is carrying +traffic efficiently, or at all beyond correctness. That is a perf question and +it belongs to `/tuning`. + +## Report additions + +- the group size, the local expert count and the per-rank window +- confirmation that the non-DWDP path is unchanged +- which of the three boot requirements the boot gate now asserts +- the digests of the upstream files the adapter binds against — this boundary + is private runtime surface rather than public API, so the digest is what + makes an upstream change visible instead of silent diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-dwdp/references/adapter-interface.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-dwdp/references/adapter-interface.md new file mode 100644 index 000000000000..6e21a2676125 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-dwdp/references/adapter-interface.md @@ -0,0 +1,39 @@ +# DwdpAdapter — what it exposes, and when each call happens + +## Construction + +| member | when | +|---|---| +| `create(...)` | once, at model construction; returns `None` when DWDP is off | +| `local_experts` / `expert_offset` | the manifest reads these to know which expert range this rank materializes. Everything outside it goes into the known-not-loaded key set | +| `moe_layers` | which layer indices participate | +| `attach_weights(weights)` | builds the module tree the runtime's setup traversal expects, over this target's flat `ParameterDict` | +| `backend(layer_idx)` | the bound module for a layer, once the runtime has composed it | + +## The facade, and why it exists + +The runtime's setup traversal walks `model.layers[i].mlp` and reads expert +weights off it. Our target has no such tree — it has a flat `ParameterDict`. +Two private helpers bridge that without duplicating any tensor: + +- **`_DwdpExpertFacade`** — the minimal `layers[i].mlp` the traversal needs. + Properties returning the layer's flat parameters, plus the gate-side expert + count the runtime reads to size the composite address space. It owns no + storage; every property is a view onto the `ParameterDict` +- **`_DwdpLayerStub`** — one entry of the `layers` list. **Dense layers get an + empty stub rather than being skipped**, because the list is indexed by + absolute layer number and a list holding only the MoE layers would shift + every index + +## The two forward hooks + +``` +before_moe(layer_idx) → the bound module. Call it immediately before the + expert call; it makes the full table addressable + +after_moe(layer_idx) releases the layer and starts the next one's prefetch +``` + +Neither raises when misplaced. `before_moe` too late addresses the previous +layer's window; `after_moe` too early evicts weights the current call still +needs. Both produce plausible output. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-dwdp/references/runtime-boundary.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-dwdp/references/runtime-boundary.md new file mode 100644 index 000000000000..d1e14e9702f8 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-dwdp/references/runtime-boundary.md @@ -0,0 +1,35 @@ +# DWDP — who owns what + +| layer | owns | +|---|---| +| **runtime** | the MPI sub-communicator · the authoritative per-rank expert range · MNNVL handle exchange · the composite virtual-address layout that makes the full expert table addressable · the P2P copy stream and its events · the double-buffered prefetch state machine | +| **adapter** (`dwdp_adapter.py`, class `DwdpAdapter`) | the boundary: the local expert window, the facade the runtime's setup traversal walks, and the two forward hooks | +| **modeling code** | checkpoint conversion · the router · the shared expert · quantization and the fused expert-runner call · the forward | + +The target copies none of the runtime's list. In particular the adapter +**creates no stream, issues no copy, and does not reach into the weight +manager** — it exposes the boundary and relays two calls. + +## Two consequences for the forward + +**It cannot be CUDA-graph captured.** The prefetch state machine is driven +from the host between layers, and a captured graph has no host between layers. +This is a property of the feature; it belongs in the run journal beside the +identity record, because every perf number taken on this target is a +non-captured number and a reader who assumes otherwise will misattribute the +gap. + +**Bind and prefetch bracket the MoE call.** Bind before it, so the expert +weights the router selected are addressable; prefetch the next layer after it, +so the copy overlaps compute rather than stalling it. Both orderings are load +bearing and neither raises when wrong — bind late and you address the previous +layer's window; prefetch early and you evict what the current call still +needs. + +## Why the digests + +The four upstream files the adapter binds against are **private runtime +surface, not public API**. There is no version contract on them, and there is +no version pin left to notice a drift either, so their SHA-256 goes in the run +journal: a change to any of them re-opens every question the adapter answers, +and the digest is what turns that from a silent drift into a visible diff. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/SKILL.md new file mode 100644 index 000000000000..e5b45224a5e6 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/SKILL.md @@ -0,0 +1,98 @@ +--- +name: staircase-modeling-eplb +description: Load when adding dynamic expert-parallel load balancing to a MoE target — the runtime owns placement and migration, the target owns slot-addressed weights and six per-layer forward calls. Gate the static variant before the dynamic one; whether the balancer actually balances is a tuning question, not a gate. +--- + +# staircase-modeling-eplb + +`--features eplb`. An adapter-class feature: the balancer is a runtime-owned +service with its own threads and its own state machine, so the target binds to +it through one file rather than absorbing it. + +Read `staircase-modeling-core/references/adapter-rules.md` first — the +one-import-site rule, the no-work-of-its-own rule and digest pinning all apply +here. + +| reference | read it when | +|---|---| +| `references/runtime-boundary.md` | always — which layer owns what | +| `references/adapter-interface.md` | writing `eplb_adapter.py`, or placing the six calls | +| `references/placement.md` | always — slots vs experts, the `layer_updates_per_iter` knob, **and the routing trap** | +| `references/weight-migration.md` | **always** — this is where the silent wrong answer lives | + +## What this feature changes + +**Weights become slot-addressed.** Every per-expert tensor is indexed +`[local_slots]`, and which expert occupies a slot is the balancer's to change +mid-run. This is the change that reaches furthest — it is in `weights.py`, in +the manifest, and in every place the forward touches an expert tensor. + +**The forward gains six per-layer calls** around the MoE block, bracketing the +window in which that layer's weights may be migrating. + +**Routing goes through the balancer.** `route()` maps token→slot, not +token→expert, and everything downstream addresses slots. + +## Two things that will bite + +Both are in the references; they are here because they are the ones worth +knowing before you start. + +**Per-expert derived scalars must migrate with the slot.** Upstream registers +four weight tensors for migration and knows nothing about a target's own +per-expert scalars. Migration runs in a C++ thread between forwards, so there +is no callback to recompute in — the only shape that works is to derive them +for the full expert set at load time, store them by slot, and register them +alongside the weights. `weight-migration.md`. + +**`offset_by_ep_rank` must be `False`.** The in-tree caller passes +`moe.use_dp`, which is `True` under attention DP; copying that breaks the +element-identical-routing invariant this forward is built on. `placement.md` +also records what the flag measurably does and does not do. + +## The arguments assembler owns + +The feature is activated through the LLM API's `MoeLoadBalancerConfig` +(`num_slots`, `initial_global_assignments`, `layer_updates_per_iter`), not +through a target-owned config file: + +- **static** — `layer_updates_per_iter: 0`. Placement computed once, no + migration. **Gate this first**: it exercises slot addressing, registration + and routing with the migration machinery out of the picture +- **dynamic** — `layer_updates_per_iter > 0` + +Neither is a tuning knob. tuner may sweep *around* them; it may not change +what they select. Record both in the run journal, with the geometry they +imply. + +## Gates + +boot with each variant → accuracy, both with `TRTLLM_MODELING_V2=require` +exported before the ranks start. **No third gate.** + +Accuracy is not blind here: a wrong slot mapping routes tokens to the wrong +experts and the score moves. The failure this feature can hide — a balancer +wired up but inert, so placement never adapts — produces *correct output*, and +that makes it a **performance** finding rather than a correctness one. Perf in +this repo is measured, never gated: whether the placement actually adapts, and +what it buys, belongs to `/tuning` and to the journal's performance section. + +What replaces a gate is the **ladder**, and it costs nothing extra: + +1. the static variant first — placement computed once, no migration. Green + here means slot addressing, registration, routing and the manifest are all + correct, with the migration machinery out of the picture +2. the dynamic variant after. A failure now is a migration failure, and step 1 + is what makes that attribution possible + +Run them in that order and record both. A dynamic run gated without the static +one first tells you the target works; it does not tell you *which half*. + +## Report additions + +- the slot geometry: `num_experts`, `num_slots`, per-rank local slots +- **the per-slot tensor list** — every tensor registered for migration, and + for each per-layer scalar that is *not* registered, why it is not per-expert +- the static gate and the dynamic gate, separately +- what the static gate established, and what the dynamic one added on top +- the digests of the upstream files the adapter binds against diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/adapter-interface.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/adapter-interface.md new file mode 100644 index 000000000000..c8000eb5f90d --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/adapter-interface.md @@ -0,0 +1,45 @@ +# EplbAdapter — what it exposes, and when each call happens + +The adapter is the target's whole view of the balancer. Everything below is +called by the target, except where noted; nothing here does work of its own. + +## Construction and registration + +| member | when | +|---|---| +| `create(...)` | once, at model construction. Returns `None` when EPLB is off — **and that is a different thing from "the balancer exists but is static"**, which is the distinction the caller most often collapses | +| `initial_global_assignments(layer_idx)` | before weights load — which expert sits in which slot at start | +| `initial_local_expert_ids(layer_idx)` | the manifest reads this to know which experts this rank must load | +| `attach_layers(weights)` | registers each MoE layer with the balancer and hands it closures over the flat `ParameterDict` | +| `register_weights(moe_slots, rebuild_views)` | enrols the per-slot tensors in the migration queue | +| `stage_layer(layer_idx, staged)` | the CPU-side staging buffers | +| `migrates_weights` | whether this configuration will move anything at all | + +## The six forward contact points + +Per MoE layer, per step, in this order: + +``` +start_wait_gpu_stage(layer) ─┐ the layer's weights may be mid-migration; +done_wait_gpu_stage(layer) ─┘ these bracket the wait +update_statistic(layer, ...) record this step's expert load +route(layer, token_selected_experts) expert ids → slot ids +start_set_cpu_stage(layer) ─┐ release the layer back to the migration +done_set_cpu_stage(layer) ─┘ machinery +``` + +**Timing is the content.** Each of these is a synchronization point with a C++ +thread; calling one at the wrong place in the forward does not raise — it +races. The wait bracket must enclose every use of that layer's expert weights, +and the release bracket must come after the last one. + +`route` is the only one that returns a tensor the forward consumes: token → +*slot*, not token → expert. Everything downstream addresses slots. Its second +argument is a trap; see `placement.md`. + +## Observability + +`placement_drift()`, `dump_placement_at_exit()`, `staging_expert_ids()` — they +read state and change nothing. Useful for answering "is the placement actually +moving", which is a tuning question rather than a gate: `/tuning` reads them +when the campaign has a placement axis. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/placement.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/placement.md new file mode 100644 index 000000000000..cd7e8dfb9aab --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/placement.md @@ -0,0 +1,64 @@ +# Placement — slots, experts, replicas, and the routing trap + +## Slots are not experts + +The balancer addresses **slots**. `num_slots >= num_experts`, and the surplus +is what lets a hot expert exist in more than one place: + +``` +num_experts = 256, num_slots = 288 → 32 experts get a second replica +``` + +Every per-expert tensor in the target is therefore indexed `[local_slots]`, +not `[local_experts]`, and the mapping from expert to slot is the balancer's +to change at runtime. A target that indexes weights by expert id has already +lost. + +`layer_updates_per_iter` decides how much moves per step: + +| value | behaviour | +|---|---| +| `0` | placement computed once, never moves. **The right first gate** — it exercises slot addressing and routing with no migration machinery | +| `n > 0` | up to `n` layers updated per iteration, walked as a circular queue over the MoE layers | + +## Which replica a token goes to + +With replicas, `(token, expert)` picks one by: + +``` +globalSlotIds[ StartOffset[e] + (k + rot) % x ] x = ReplicaCount[e] +rot = blockIdx.x + (use_dp ? ep_rank : 0) ← offset_by_ep_rank +``` + +## The trap: `offset_by_ep_rank` must be `False` + +`route()`'s second argument defaults to `False`, but **the in-tree caller +passes `moe.use_dp`**, which is `True` under attention DP. Copying that breaks +the invariant this target's forward is built on. + +In-tree is right to enable it: under route-then-dispatch each rank routes only +its own tokens, and without the offset all ranks would pile onto the same +replica. Our forward does not work that way, and it requires every rank to +produce element-identical routing. + +One measured detail that changes how the flag should be read: **`ep_rank` is +added inside the `% ReplicaCount[e]`**, so with 288 slots over 256 experts — +at most two replicas each — it only groups ranks by `ep_rank mod +ReplicaCount`. Measured on B200 (`sm_100`), 256 experts over 288 slots at EP4: +rank 0 and rank 2 produced **bitwise identical** output, as did ranks 1 and 3 +(512 tokens, top-8; the per-rank difference vector against rank 0 was +`[0, 482, 0, 482]`). The rotation formula is upstream's — re-check it if the +routing kernel changes. + +That does not change the decision — pass `False` regardless — but it corrects +the reading: "offset by ep_rank" does **not** mean "each rank gets a distinct +round-robin start". And the direction that matters to us is the opposite one: +one rank differing is enough to break the invariant, so the `False` side's +"all ranks element-identical" is not a requirement that can be relaxed. + +## Statistics + +Three paths exist upstream for feeding the balancer load statistics. Take the +one keyed by **global expert ids** — the others are keyed by slot or are +rank-local, and both make the balancer's view of load wrong in ways that look +like a bad placement algorithm rather than a wiring error. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/runtime-boundary.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/runtime-boundary.md new file mode 100644 index 000000000000..5dcb9215bc12 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/runtime-boundary.md @@ -0,0 +1,34 @@ +# EPLB — who owns what + +Dynamic expert-parallel load balancing moves experts between slots while the +engine serves. Three layers, and the split is the reason a target can use it +without absorbing it. + +| layer | owns | +|---|---| +| **runtime** | the balancer itself: expert placement, the migration machinery, three C++ threads, and the 64-bit signal word the Python side and the C++ side synchronize through. It decides *what* moves and *when* | +| **adapter** (`eplb_adapter.py`, class `EplbAdapter`) | the boundary: layer and slot registration, the staged-weight buffers, and closures over this target's flat `ParameterDict`. It hands over geometry and placement and relays calls | +| **modeling code** | six per-layer calls at the points in the forward where they belong, and the routing call that consumes the balancer's slot mapping | + +The adapter creates no tensor, launches no kernel and moves no byte of its +own. See `staircase-modeling-core/references/adapter-rules.md` for why that +constraint is load-bearing rather than stylistic, and for the digest-pinning +rule that goes with binding a private runtime surface. + +## What this buys, and what it costs + +Buys: the balancer's placement algorithm, its migration scheduling, and its +cross-rank coordination, none of which the target reimplements. + +Costs: the target's weights must be **addressable by slot** rather than by +expert, and everything that is per-expert must move when a slot's occupant +changes. `weight-migration.md` is that whole story, and it contains the +silent-wrong-answer trap of this feature. + +## Static and dynamic are the same code + +`layer_updates_per_iter: 0` means the placement is computed once and never +moves. That is not a different feature — it is this one with migration turned +off, and it is the right first gate: it exercises registration, slot +addressing and routing without the migration machinery. `placement.md` has the +knob. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/weight-migration.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/weight-migration.md new file mode 100644 index 000000000000..8f70a63ff4eb --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-eplb/references/weight-migration.md @@ -0,0 +1,57 @@ +# Weight migration — what must move when a slot changes occupant + +When slot 37 stops holding expert 51 and starts holding expert 200, +**everything indexed by that slot must change together.** Anything left behind +pairs expert 200's weights with expert 51's scalars, and that is a silent +wrong answer: no crash, no assertion, a slightly wrong number. + +## What upstream registers, and what it misses + +The migration queue upstream enrols four tensors per layer: + +``` +w3_w1_weight w2_weight w3_w1_weight_scale w2_weight_scale +``` + +An NVFP4 target's expert runner also consumes **per-expert scalar arrays** — +on the DeepSeek-R1 shape, three of them (`output1_scale`, +`output1_gate_scale`, `output2_scale`, derived from the checkpoint's +`e_ws2_1` / `e_isc2` / `e_ws2_2`). Those are per-expert, so they are per-slot, +so **they must migrate too**. Upstream does not know about them. + +## Why you cannot fix it after the fact + +Migration is performed by a **C++ thread in the gap between forwards**. There +is no Python callback point after a migration — nothing to hook, nothing to +recompute in. So "register the raw checkpoint scalars and re-derive after the +move" does not work; there is no "after the move". + +## The shape that does work + +**At load time, compute the derived scalars for all experts, store them +indexed by slot as `[local_slots]` tensors, and register them in the migration +queue alongside the weights.** + +The derivation is per-expert elementwise arithmetic, so computing it for the +full expert set costs nothing extra and depends on nothing that is not already +present. The staging buffers hold the full derived values too. Migration then +keeps them consistent automatically, with no callback required. + +## What does not migrate + +Scalars that are not per-expert. On the same target `routed_input_scale = +1.0 / mlp_isc1` is a **shared-expert** scalar — one value for the layer, not +one per expert — and enrolling it in the migration queue would be wrong in the +other direction. + +The test is mechanical: **is this value indexed by expert?** If yes it is +indexed by slot and it migrates. If no it does not. Write the answer for every +per-layer scalar in the target into the run journal; the next reader should +not have to re-derive it. + +## Changing storage breaks views + +Registering a tensor for migration can replace its storage. Any view built +over the old storage — a reshape, a slice held in a cached structure — is +silently stale afterwards. Rebuild views after registration; that is what the +`rebuild_views` callback in `register_weights` is for. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-host-offload/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-host-offload/SKILL.md new file mode 100644 index 000000000000..df233f65a506 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-host-offload/SKILL.md @@ -0,0 +1,196 @@ +--- +name: staircase-modeling-host-offload +description: Load when a weight tensor is too large to keep resident and must live in host memory, with the forward gathering the rows it needs each step. An output-identical feature, which is exactly what makes it dangerous — both standing gates pass whether or not it is doing anything, so it needs a residency witness. +--- + +# staircase-modeling-host-offload + +`--features host_offload`. One or more weight tensors stay in **pinned host +memory** for the life of the engine; the forward gathers the rows a step +actually needs and copies those to the device. + +The first instance is Qwen3.8-Flash-Next's PLE n-gram embedding table: +`ngram_vocab_size_base: 20000000` at `ple_embed_dim: 2560` in FP8, which is +**53.7 GB of a 124 GB checkpoint**. Whether it is resident decides how much of +a 284 GB device is left for KV. That is the test for whether this feature +applies: *a single tensor whose size is comparable to the rest of the model*. +It is not a general memory-pressure remedy. DeepSeek-V4.1's Engram table is +the second: two modules of 384M rows x 256 FP8 plus E8M0 scales, **202.76 GB**, +larger than the rest of the checkpoint. + +Read `staircase-modeling-core/references/adapter-rules.md` first if the +offload is expressed as an adapter rather than inside the forward. + +## The property that shapes everything else + +**Offloading is output-identical by construction.** The same rows are gathered +and the same arithmetic runs; only where the bytes waited changes. So: + +- the **boot** gate passes whether or not the tensor was offloaded; +- the **accuracy** gate passes whether or not the tensor was offloaded; +- a target that quietly loaded the table onto the device anyway reports a + green feature unit. + +This is the same shape of blindness the acceptance gate exists for, and it +wants the same answer: a third measurement that can only come out one way if +the feature is real. See **Gates** below. + +## What this feature changes + +**The manifest.** The offloaded tensor is materialized on CPU with +`pin_memory=True` rather than on the device. It must still be *covered* by the +weight manifest — an offloaded tensor is loaded, not skipped, and putting it in +the known-not-loaded set would be a lie that the coverage check cannot catch. +Record the residency, not just the presence. + +**The forward gains a gather.** Two designs, below: a **zero-copy gather** +that the GPU issues against host memory mapped into its address space +(recommended, and CUDA-graph capturable), and a **host `index_select`** that +brings the indices to the host and copies the rows back (the fallback). Both +read the same rows, so the choice changes cost and capturability, never +output. + +**Three things decide whether it is usable at all, and none is about +correctness:** + +| decision | why it matters | +|---|---| +| pinned vs pageable host memory | pageable H2D on some driver/platform combinations degrades pathologically — `tensorrt_llm/_torch/pinned_weight_staging.py` documents one such case on GB300, a 36-minute to 4-hour stall with no error. Pinned is not an optimization here | +| which stream the copy runs on | a copy on the compute stream serializes the gather against everything else in the step | +| whether the index computation is host-visible | see below — this is the CUDA-graph question | + +**CUDA graph capture.** A gather whose indices are read back to the host inside +the captured region cannot be captured: the host read is a synchronization the +graph cannot express. There are two honest resolutions and they are not +equivalent: + +1. make the gather **device-driven** — a kernel indexes the host table through + a device pointer to it — which is capturable, and is the design to reach + for on Grace-Blackwell; see the next section; +2. keep the offloaded layer **outside** the captured region with a host + `index_select` — states the feature's cost plainly, and is the fallback when + the device cannot read host memory well enough (see the last section). + +Whichever is chosen goes in the run journal beside the identity record, because +a reader who assumes capture will misread every later perf number. **Do not +report the target as CUDA-graph enabled without saying which layers are in the +graph.** + +## The capturable design: a zero-copy gather through UVA + +**The in-tree precedent is `tensorrt_llm/_torch/modules/qwen4_exp/ple.py`**, the +Qwen4-Exp PLE table. Read it before writing the gather; each piece below is a +problem it already solved. + +- **An environment switch, parsed strictly.** `TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD`, + read by `_uses_ple_host_offload`, which raises on anything that is not a + boolean rather than reading a typo as "off". The residency witness's + `contrast_env` is exactly this switch, so it has to mean one thing. +- **One stable pinned allocation.** `Qwen4ExpPinnedHostEmbedding` creates the + table on `meta` and materializes it pinned exactly once + (`materialize_pinned`); its `_apply` override keeps `.to(device)` and + `.cuda()` from moving or replacing the table. A mapped pointer to memory + that later moves is a pointer to freed memory. +- **The device pointer is resolved before capture.** `_mapped_device_ptr` + calls `cudaHostGetDevicePointer` once per device, caches it, and refuses to + do so while a stream is capturing. So warmup must touch the gather before + the first capture — a graph cannot record the lookup, only its result. +- **A Triton zero-copy gather.** `_gather_pinned_embedding_kernel` loads the + rows straight from pinned host memory through the mapped pointer — no + staging copy, no host sync — writes zeros for rows another rank owns, and + applies the FP8 scale in the same pass. +- **Row sharding plus a reduction.** Under tensor parallelism each rank holds + `ceil(rows / tp)` rows and the partial lookups are summed with an + `AllReduce` (`Qwen4ExpNGramEmbedding`, `_finish_embedding_lookup`). Under + attention DP each rank serves its own tokens, so `_prepare_embedding_lookup` + **allgathers the ids**, every rank gathers the rows it owns for all of them, + and a **reduce-scatter** returns each rank its own tokens' sums. +- **A prefetch stream.** `Qwen4ExpPLE.start_prefetch` issues the gather on a + side stream before execution reaches the layer, so the host read overlaps + the layers before it. + +### The shared single-copy layout (recommended on Grace-Blackwell) + +Row sharding holds `1/tp` of the table per rank — every byte of it held once +on the node — but pays a collective per lookup. The alternative removes the +collective and still holds the table **once per node, not once per rank**: +one host copy that every rank on the node maps and registers. + +- rank 0 creates the table in one anonymous shared region (`memfd_create`), + backed by huge pages (`madvise(MADV_HUGEPAGE)`); +- the other ranks on the node map the same region — SGLang reaches it through + `/proc//fd/`, which needs every rank in one PID namespace, so one + container per node; +- every rank `cudaHostRegister`s the mapping and resolves its own device + pointer, before capture as above; +- every GPU can then gather **any** row: no sharding, no zero-fill, no + allreduce or reduce-scatter, and the lookup is bitwise the resident one. + +At 202.76 GB for DeepSeek-V4.1's Engram table, per-rank copies on a four-GPU +node would be 811 GB of host memory; the shared layout is 202.76 GB, once. +SGLang ships exactly this behind `SGLANG_DSV41_ENGRAM_HOST_TABLE_LAYOUT=shared` +(implementation in `python/sglang/srt/layers/engram.py`, branch +`origin/dsv4.1`); TRT does not have it yet, so it is new target code, and the +`ple.py` pieces above are the parts it can follow. + +What to check before choosing it, none of which is about correctness: + +| check | why | +|---|---| +| huge pages | SGLang's own note: without them "expect ~10x slower lookups" — TLB misses on a 200 GB mapping | +| NUMA | a GB300 node puts four GPUs on two Grace sockets; one copy lives on one socket (or interleaved), and the far GPUs read it over the socket link | +| host RAM | the table once, plus pinned staging for weight loading, must fit beside everything else the node holds | +| lifetime | rank 0 owns the region; the others must map it before rank 0 can exit, and a restarted rank 0 invalidates every mapping | + +The residency witness is unchanged by the layout — it compares per-rank +device memory with the offload on and off — but `expected_mib` is the OFF +leg's **per-rank** footprint: DeepSeek-V4.1's table row-sharded at TP4 is +202.76 GB / 4, **48,340 MiB** per rank. + +### The fallback: a host `index_select` + +Per step: compute the row indices on device, bring them to host, +`index_select` on the host tensor, copy the result back. Two host-device round +trips per invocation, and a host read the graph cannot capture — so the layer +stays outside the captured region. It needs no mapped pointer and no +registration, which makes it the path for a platform where device reads of +host memory are slow, and the first thing to fall back to when the zero-copy +gather is suspected. + +## Gates + +boot → accuracy → **residency witness**, all under +`TRTLLM_MODELING_V2=require`. + +The witness is the gate that can actually see this feature. Shape it so only +one outcome is possible if the feature is real: + +- run the same boot workload twice in the same session, once with the feature + and once without, and read **device memory after weight load** (e.g. + `torch.cuda.memory_allocated()`, or `nvidia-smi --query-gpu=memory.used` + sampled after the engine reports ready); +- assert the delta is **at least the offloaded tensor's own size**, computed + from the checkpoint rather than from the run; +- assert the generated text is **token-identical** between the two, greedy. If + it is not, the gather is wrong and the accuracy gate's tolerance was wide + enough to hide it. + +A witness that only checks "the run still works" certifies nothing. Write the +expected delta down from the checkpoint's own shapes *before* running it. + +## What the gates still do not tell you + +Whether the per-step gather is fast enough to be worth having. That is a perf +question and it belongs to tuning, not to onboarding. Onboarding's claim is +"the table is not resident and the output is unchanged" — say exactly that and +no more. + +## Report additions + +- the tensor(s) offloaded, each with its checkpoint shape, dtype and byte size +- measured device memory with and without, and the delta against the predicted + one +- confirmation of token-identical output, and at what sampling settings +- which layers are inside the CUDA graph and which are not, and why +- which gather design — zero-copy (and which layout: sharded or shared) or + host `index_select` — and which stream it runs on diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-kv-quant/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-kv-quant/SKILL.md new file mode 100644 index 000000000000..9b1f4c375459 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-kv-quant/SKILL.md @@ -0,0 +1,138 @@ +--- +name: staircase-modeling-kv-quant +description: Load when switching a target's KV pools from bf16 containers of the reference's fake-quantized values to the checkpoint's native storage formats (DeepSeek-V4.1's FP4 main KV, MXFP4 indexer K, FP8 SWA KV). Output-identical by construction when attention dequantizes exactly, so boot and accuracy cannot see it; it needs a bytes-per-token witness against a declared figure. +--- + +# staircase-modeling-kv-quant + +`--features kv_quant`. Core stores the **same values** the reference computes — +the reference fake-quantizes KV to its storage formats and attends over the +dequantized result — but in bf16 pools, because the only sm_103 sparse +attention kernels that serve DeepSeek-V4.1 are exact only over bf16 KV. This +unit switches the pools to the formats the checkpoint was trained for. It is +what a request for "FP4 KV cache", "FP8 attention" or "890 bytes of KV per +token" actually asks for. + +## The rule: store the reference's value set, never a different one + +Real storage of the value set the reference already fake-quantizes to is +**lossless relative to the reference** — the model was trained (QAT) with +exactly these values. Storage of any *other* value set is a lossy +substitution, however close its bit budget, and must be named as one in the +plan, in the spec (`exact: false`), and in the report. + +Substitutions that look like this feature and are not: + +| format | why it is a different value set | +|---|---| +| `kv_cache_config.dtype: fp8` as a per-tensor FP8 cache | one scalar scale, not UE8M0 per 32 channels | +| `fp8_ds_mla` — DeepSeek-V4's 584 B layout | 448 FP8 + 64 bf16 RoPE dims with block-64 scales; V4.1 quantizes all 512 dims, RoPE included, after RoPE (`tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/footer_scale_kv.py` is the V4 layout) | +| TRT's `nvfp4` compressed pool as DeepSeek-V4 uses it | NVFP4 with a second-level global scale and a 64-dim residual; V4.1 has no global scale (it is 1) and all 512 dims in FP4 | + +The DeepSeek-V4 machinery in +`tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/` — the +`KVCacheDtype` presets in `compressor.py`, `DeepseekV4CacheManager` and its +`get_cache_bytes_per_token` in `cache_manager.py` — is the closest thing in +tree, and it encodes the **V4** layouts. Read it to learn the interface; do +not assume its byte formulas describe V4.1. + +## DeepSeek-V4.1's formats + +From the tech report (§2.4.4) and the reference's fake-quantization: + +| pool | who owns one | format | bytes per entry | +|---|---|---|---| +| SWA KV | every layer (40), the window only | FP8 E4M3 per 32 channels with a UE8M0 scale, over all 512 dims including RoPE | 512 + 16 = **528** | +| main (compressed global) KV | the KV-source layers only: 2, 8, 14 (compress ratio 2) and 20 (ratio 1) | FP4 E2M1 with an E4M3 scale per 16, **no global scale** | 256 + 32 = **288** | +| indexer K | the same KV-source layers | MXFP4: E2M1 with a UE8M0 scale per 32, 128 dims | 64 + 4 = **68** | + +**Only layers 2, 8, 14 and 20 own compressed and indexer pools.** Every other +layer reads a source layer's pool through `kv_source_layer_ids`. A pool +allocated per compressed layer (38 of them) comes to 10324 B/token, more than +ten times the figure below, and is the sign the sharing was never wired. + +**Bytes per token** is the persistent, growing part — global KV only: + +``` +3 x (288 + 68) / 2 + (288 + 68) = 534 + 356 = 890 B/token + ratio-2 layers ratio-1 layer +``` + +The tech report's headline figure. The same formula gives 3200 for bf16 +containers of these values (core) and 1490 with FP8 main KV. The SWA window +is per request, not per token — 40 layers x 128 x 528 B, about 2.7 MB — and so +is the compressor's carry-over state; neither belongs in the figure. + +## Attention over native storage + +The kernels that serve V4.1 on sm_103 — trtllm-gen's H512 DynamicTokenSparse +path through FlashInfer's `trtllm_batch_decode_sparse_mla_dsv4`, and the +catalog's `thop_attention` — are exact only over bf16 KV. The native formats +are consumed by SGLang's FlashMLA fork, which TRT does not vendor, or by a +**dequantize-and-gather to bf16** step ahead of the kernel. So: + +- native storage + exact dequantization + bf16 Q is **output-identical** to + core. Declare `exact: true` (the default) and the witness demands token + identity with bf16 storage; +- an FP8-query kernel is not: FP8 Q is an approximation beyond the reference. + Declare `exact: false`, say so in the plan, and accuracy becomes the only + output check. The verdict records that it is the weaker claim. + +## Gates + +boot → accuracy → **kv_quant witness**, all under `TRTLLM_MODELING_V2=require`. + +**The counter.** When `TRTLLM_MODELING_V2_WITNESS=1` is set, the target prints, +once per rank after its KV pools are allocated: + +``` +MODELING_V2_WITNESS kv_bytes_per_token=890 +``` + +Computed from the pools **as allocated** — the tensors' dtypes and shapes, +summed over pools, a ratio-r pool contributing its entry size over r — for one +token of a request this rank serves. Not typed in: a constant print passes the +native leg and is caught by the contrast leg, which must report more. SWA and +compressor state are per request and excluded. + +**The verdict**, against the spec's declared figure: + +- native storage reports `expected_bytes_per_token` within `bytes_tolerance` + (default 2%) — enough for row padding, not for a layout change: padding the + 4-byte indexer scale to 16 moves 890 to 920; +- with `exact: true`, the leg under `contrast_env` (bf16 storage) reports MORE + bytes per token — otherwise it selected nothing and the next check compares + native storage with itself — and the greedy continuations are + **token-identical**. A difference means the stored values are not the + reference's: a layout, scale or rounding from some other format. + +```yaml +features: [kv_quant] +accuracy: + witness: + kv_quant: + contrast_env: {TRTLLM_DSV41_KV_STORAGE: bf16} # the leg WITH bf16 storage + expected_bytes_per_token: 890 + exact: true + # bytes_tolerance: 0.02 + max_batch_size: 1 # both legs one request at a time +``` + +Compare at batch 1 -- `max_batch_size: 1` in the witness block runs both legs +there: SGLang reports that this model's output is not bitwise stable across +batch composition, and the identity check must measure storage, not batching. + +## What the gates still do not tell you + +Whether the dequantize step is fast enough, and how much KV capacity the unit +bought. Both are perf and capacity measurements; the witness says only that +the pools hold the checkpoint's formats and that the output did not move. + +## Report additions + +- each pool: owner layers, format, bytes per entry, and the allocated dtype +- measured bytes per token against 890, and the bf16 leg's figure +- `exact`, and if false, which attention path is lossy and why it was chosen +- how attention reads the native values (dequantize-gather, or which kernel) +- confirmation that only the KV-source layers allocate compressed and indexer + pools diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-multimodal/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-multimodal/SKILL.md new file mode 100644 index 000000000000..e96c5f52614d --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-multimodal/SKILL.md @@ -0,0 +1,244 @@ +--- +name: staircase-modeling-multimodal +description: Load when a target must accept images or other media alongside text — an extra input path with its own tower, projector and weights, plus registration with the runtime's multimodal input registry. Gated on MMMU-Pro rather than a text benchmark. Carried end to end by one target. +--- + +# staircase-modeling-multimodal + +`--features multimodal`. **One target has carried this end to end** — +`models//qwen3.6-35b-a3b-nvfp4__sm_100__tp1` (27-block ViT, +sm_100) — so the gate section below and `references/vision-path.md` are +measurements now, each carrying the evidence that established it. One target +is still one target: a claim here is what happened on *that* checkpoint's +shape, not a law. Keep checking, and keep writing back. + + + +| reference | | +|---|---| +| `references/vision-path.md` | the extra input path: tower, projector, weights, and the registry hand-off — now written as answers with evidence, including the two traps below that cost the most | + +## Why it is a feature and not core + +A multimodal checkpoint can be assembled **text-only**, and that is the +default: `--features` empty means the vision weights go into the +known-not-loaded key set and the target serves text. That is a legitimate, +gateable target — so implementing the vision path is a choice, which is what +makes it a feature rather than something the checkpoint's `config.json` +forces. + +## What it changes + +It is the only feature that adds an **input path** rather than changing the +forward's math: + +- a second weight family — the tower and the projector — with its own layout + conventions, usually unquantized while the language model is not +- construction of a second module tree in `__init__` +- a hand-off with the runtime's multimodal input registry, which is where + media arrive already decoded and batched +- image tokens spliced into the sequence before the language model runs + +Two of those turned out to have a shape worth naming up front, because getting +either wrong produces a **passing-looking engine with a dead vision path**: + +- **the splice belongs in the shell, as `inputs_embeds`.** The inherited + `DecoderModelForCausalLM.forward` forwards only the keywords it names, so + `multimodal_params` riding in `**kwargs` is silently dropped and the model's + own forward never sees it. +- **give the boot gate images that differ from each other.** A dead splice + answers every image identically; one image case passes anyway. + +And one thing that is not an input path at all: **a VLM whose text model uses +mRoPE runs a degenerate mRoPE on text** (all three position rows equal), so a +text-only assembly may legitimately take row 0 and run plain RoPE. Turning the +tower on breaks that, and the rotation has to read all three rows. Branch on +the per-engine feature constant so the text path stays bit-identical. + +## Gates + +boot → accuracy on **MMMU-Pro**, both with `TRTLLM_MODELING_V2=require` +exported before the ranks start. **No third gate.** + +This feature is not blind to the first two: get the vision path wrong and the +output is wrong. There is no distribution-preserving construction hiding the +defect, and no resource-distribution value a correct-output run could be +missing. + +What it does need is a benchmark that **looks at the image**. MMLU and GSM8K +are text; they pass at full score with the vision tower producing garbage, +because nothing in them reads a pixel — and that is not a hypothetical: the +first engine run of the one target that has done this had a dead splice, and +its text gates would have passed untouched. So the accuracy case for a +multimodal configuration reads its anchor from the image benchmark's file +under `tests/integration/defs/accuracy/references/`, and the journal says +plainly that the text benchmark's score is evidence about the language model +only. + +**A target that serves both text and images needs two gates, not one**, and +re-pointing one gate's task at the image benchmark would silently re-gate the +text configurations — including runs already recorded. Write them as **two +cases** in the family's accuracy file +(`tests/integration/defs/accuracy/test_modeling_v2_.py`): the +multimodal case evaluates MMMU-Pro, the text cases keep their text benchmark, +and each reads its own anchor because the reference files are keyed by +`TestClass::test_method`. Name the case after the configuration it gates. + +**Read how it is scored before you trust a number.** `trtllm-eval mmmu` and +`mmmu_pro` are first-class subcommands, registered in +`tensorrt_llm/evaluate/lm_eval.py` — go read that command: how the dataset is +loaded, what is fed to the model, which metric is reported and which filter +selects it. Nothing stops you, and a gate you cannot explain is a gate you +cannot debug. The GSM8K precedent is the reason this is not optional: the same +run read 25.47 under the strict filter and 90.60 under `flexible-extract`, and +the low number looked exactly like a broken model. + +**`mmmu` did not run at lm_eval 0.4.12 — use `mmmu_pro`.** Established by +running it against **stock trtllm** on a real checkpoint, which is the control +that says whose bug it is: + +``` +AttributeError: 'dict' object has no attribute 'build_all_requests' + lm_eval/evaluator.py:537 task.build_all_requests(...) +``` + +The hub's `mmmu_val` is a **doubly** nested group — `mmmu_val` → one group per +discipline → the tasks. lm-eval 0.4.12's legacy-format shim +(`evaluator_utils.py:_handle_back_comp`) unwraps **one** level: for a group key +it does `tasks.update(value)`, so the discipline groups land in `tasks` where +Task objects are expected, and the first `build_all_requests` raises. Nothing +model-side is reached, so this fails identically on any model. + +`mmmu_pro` is the way through and is a first-class subcommand: trtllm ships its +own task yamls under `evaluate/lm_eval_tasks/mmmu_pro/`, and +`get_task_dict(['mmmu_pro_standard_4'])` returns a **flat** `str -> +ConfigurableTask`. `--subset standard_10` (the default, 10 options) or +`standard_4`. Either is a valid gate as long as the reference and the target are +measured on the same one — the gate is a comparison, and an anchor measured here +is measured on whichever you pick. + +*Verified at lm_eval 0.4.12. Re-check by running `mmmu` against stock trtllm +first: if it still raises there, it is still upstream's.* + +**And KV-cache block reuse must be off.** With reuse on — which is the +**default** — the first multimodal request dies in the executor: + +``` +RequestError: gen_multimodal_cache_key_tokens(): incompatible function arguments + Invoked with types: NoneType, bytes, int, kwargs = { token_offset: int } +``` + +`id_offset` arrives `None` where an `int` is required, from the +`kv_cache_manager_v2` path that hashes multimodal content into a reuse key. +Reproduced on **stock trtllm** on its own model, so it is upstream again, and +it is why every VL test in trtllm's own accuracy suite constructs +`KvCacheConfig(..., enable_block_reuse=False)` rather than taking the default. + +Consequences, both of which a target must carry rather than discover: + +- the gate invocation needs `--disable_kv_cache_reuse` (trtllm-eval) / + `enable_block_reuse: false` (LLM-API); +- a multimodal configuration must set it too, and the journal should say why, + because the default is the broken value and a future reader will otherwise + "clean it up". + +**And the gate needs an explicit `scores_filter`, or scoring raises.** Without +one the evaluator averages every metric the task reports — including lm-eval's +`alias`, which is a string — and numpy raises + +``` +TypeError: the resolved dtypes are not compatible with add.reduce. +Resolved (dtype(' +- **images do traverse the hand-off** — verified by a passing gate. + `MultimodalLmEvalWrapper` builds the prompt itself and submits + `multi_modal_data={"image": [...]}` to the serving endpoint; lm-eval's own + VLM adapters (`hf_vlms.py`) are not on the path +- **prompt construction is part of the protocol and it moves the score.** The + wrapper reads `model_type` from the checkpoint's `config.json` and looks the + registered `MultimodalPlaceholderMetadata` up by it; bulk-prepended versus + interleaved placeholders is a different prompt for every MMMU-Pro document. + If the anchor was measured on stock trtllm, register the same values the + stock model_type carries — readable without touching model code through + `MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder_metadata(...)` + +Record the anchor's source in the reference entry as usual — an MMMU score is +never invented, and if no published number for this checkpoint exists, measure +stock trtllm on it and mark the source accordingly. + +## Measure the blind baseline before you trust the threshold + +The gate is `anchor - tol`, and that bar is only meaningful if an engine which +**ignores its images** falls below it. Agreement with stock trtllm does not +establish this: it shows two systems match, not that either reads a pixel. + +Measure it. Patch the task's image supplier so every image becomes a flat grey +of the same size — count, placeholder tokens and sequence lengths unchanged, only +the pixel information removed — and re-run the same configuration over the same +documents. + +Measured once, on Qwen3.6-35B-A3B / MMMU-Pro standard_10, 1730 docs: **normal +63.06, blind 33.24**, against a threshold of 57.83. The gate discriminates by +~26 combined standard errors. + +**The blind score is the number worth having, and it is not the random +baseline.** Random on ten options is 10; blind measured **33**. A text-only +reader gets real signal from the question stem, so "our score is far above +random, therefore we see the image" is not an argument — the honest contrast is +against the blind score. On a checkpoint where blind lands near the threshold, +the gate is not discriminating and needs a harder benchmark or a tighter +tolerance, and you would never learn that from a passing run. + +Verify the control itself: count the images you replaced and fail loudly if the +patch bound nothing. A silently-inert control returns the normal score and reads +as proof. + +*Measured on sm_100; the blind figure is a property of the checkpoint and the +benchmark, so re-measure it per target rather than carrying this one across.* + +## Report additions + +- everything you had to establish by observation, named precisely — this skill + is speculative until a target fills it in, and the report is how that + happens +- the benchmark chosen and why it exercises the vision path diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-multimodal/references/vision-path.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-multimodal/references/vision-path.md new file mode 100644 index 000000000000..daf4b0a735a2 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-multimodal/references/vision-path.md @@ -0,0 +1,192 @@ +# The vision path + +> **Verified once**, by the qwen3.6-35b-a3b-nvfp4 / sm_100 / tp1 target's +> multimodal increment (Qwen3.6-35B-A3B-NVFP4, 27-block ViT, sm_100). +> Everything below carries the evidence that established it. One target is one +> target: a claim here is a measurement on *that* checkpoint's shape, not a law. + +## Where the runtime's half lives + +`tensorrt_llm/inputs/` — `multimodal.py`, `multimodal_data.py`, `registry.py`, +`media_io.py`. Media arrive there already decoded and batched; the registry is +the hand-off point between the runtime's input handling and a model's own +vision path. + +Note what is *not* there: the towers themselves live in the engine's per-model +definitions. Read them to understand an interface; never to copy a structure — +the tower's **semantics** come from the checkpoint's HF reference code, like +every other computation in a target. + +## The five answers + +**1. The runtime hands you decoded media; you own the HF processor call.** +`LLM.generate(TextPrompt(prompt=..., multi_modal_data={"image": [PIL.Image]}))` +reaches your registered input processor's `call_with_text_prompt` with PIL +images already loaded. Turning them into `pixel_values` is yours: +`AutoProcessor.from_pretrained()` does it, which means the +checkpoint has to carry the processor's own files — `preprocessor_config.json`, +and for the Qwen-VL family `video_preprocessor_config.json` beside it, or the +processor class will not construct. The checkpoint is read exactly as +published, so a checkpoint missing one is a checkpoint to report, not a +directory to patch. + + +**2. What the registry requires, and when.** + +- `@register_input_processor(ProcessorCls, model_type=..., placeholder_metadata=...)` + applied to the **model class**. It writes two independent registries: the + input-processor registry keyed on the *class*, and the placeholder registry + keyed on the *model_type string*. **Neither is `architectures[0]`** — so a + target that must keep a particular architecture name for an unrelated reason + (a recurrent-state pool, say) keeps it, and registers the processor under the + `model_type` the checkpoint's config carries instead. That combination is + verified. +- `ProcessorCls` subclasses `BaseMultimodalInputProcessor` and implements + `call_with_text_prompt` plus the `processor` / `tokenizer` / `config` / + `dtype` properties. `get_num_tokens_per_image` comes free from the base if + the HF processor has `_get_num_multimodal_tokens`; `get_mm_token_ids` needs + an override when the processor exposes no `mm_token_ids` (return a 1-D + tensor holding the placeholder token id). +- It is constructed **in the launcher process at `LLM(...)` time**, by + `create_input_processor` → `ModelConfig.from_pretrained()` — + **without the LLM-API's `model_kwargs`**. So any geometry a configuration + passes through `model_kwargs` reaches the *model* and not the *processor*. + The only source both halves read is the checkpoint's own `config.json`, and + there is no target-owned copy to put a number in: keep `model_kwargs` to + switches, or the two halves will derive the tower's tables from different + numbers. + +- Registering the processor makes `model_engine.is_multimodal` true for every + configuration of that target, including text-only ones. Measured: that + property gates exactly one thing — the encoder-decoder persistent-input fast + path — so it is inert for a decoder-only target. There is no pin to bump, so + re-check it by reading `is_multimodal`'s uses when the input path changes. +- A text-only prompt must still work: implement the no-media branch as exactly + `DefaultInputProcessor`'s `tokenizer.encode(prompt, add_special_tokens=...)` + so the identity configuration's token ids do not move. + +**3. The splice happens in the shell, as `inputs_embeds` — not inside the +model's forward.** This is the trap that costs the most time: + +> **`DecoderModelForCausalLM.forward` forwards only the keywords it names.** +> `multimodal_params` sits in `**kwargs` and is dropped. Measured: the core +> forward saw `['spec_metadata']` and nothing else, the splice never ran, and +> every image produced the *same* answer — a boot gate that boots, generates +> fluent text and is completely blind. + +So the shell's `forward` reads `kwargs["multimodal_params"]`, runs the tower, +writes its rows over the image placeholder embeddings, and passes the result +down as `inputs_embeds` (with `input_ids=None`). That is upstream's own shape +for this, and it keeps the inherited forward — and with it the text gate +record — untouched. + +Make the failure impossible to miss next time: a boot-gate image set whose +images differ from one another (four colours, two shapes, light against dark, +two rendered digits, and one request carrying two images) failed **8 of 10** +with the splice dead and passes with it live. A single image case would have +passed either way. Vary the image *sizes* too — that is what puts more than one +sequence length through the tower's attention call and exercises the ragged +multi-image path. + +**4. Which shapes vary, and what that means for capture.** The whole media +path is context-only: generation steps carry no payload (the executor strips +`multimodal_data` down to `mrope_config.mrope_position_deltas` after prefill), +and context batches are not CUDA-graph captured. So the data-dependent shapes +never enter a captured region and the sync-free rules are satisfied by +construction. What *does* need care is the host side: + +- **Everything derived from the image grid should be computed in the input + processor, not the forward.** The rope cos/sin table, the position-grid + gather indices and weights, the per-image lengths that act as `cu_seqlens`, + the activation op's tile map, the language model's 3-D mRoPE positions — all + of it is integer and float arithmetic over metadata. Doing it host-side once + per request is what leaves the forward with nothing but catalog calls, and + it is the difference between a forward that satisfies the closed-vocabulary + rule and one that needs `arange`, `linspace`, `cos` and `sin` in it. +- `MultimodalParams.to_device` moves **every** tensor under `multimodal_data` + to CUDA unless the model class declares + `multimodal_data_device_paths = ["image.pixel_values", ...]`. A cache-free + attention call needs host-resident length vectors, so without that list they + arrive in device memory. The tower's attention is served by + `tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention` + in its no-KV-cache mode — all four pool arguments `None` — which is the mode + trtllm's own diffusion stack runs, and which a feasibility pass once missed, + descoping the feature for a ceiling that was not there. + +- `MultimodalParams.input_ids_start_offset` is the request's first row in the + flattened batch; each image's placeholder tokens are one contiguous run, so + the write is a slice copy rather than a scatter — no index tensor and no + host sync. Assert the run is contiguous in the processor and assert it lies + inside this forward's rows in the model: that second assert is what catches + a prompt split by chunked prefill or KV-cache reuse. + +**5. The MMMU gate does carry images end to end.** `trtllm-eval mmmu_pro` +builds the prompt through `MultimodalLmEvalWrapper.apply_chat_template` and +submits `multi_modal_data={"image": [...]}` to the serving endpoint; lm-eval's +own VLM adapters are not on the path. Verified by a passing gate on a real +target — 65.3179 ± 1.1446 on `mmmu_pro_standard_10`, n = 1730 — and cheaply +pre-verifiable by the boot-gate image set above. + +Two details of that path are protocol, not modelling, and they move the score: + +- the wrapper reads **`model_type` from the checkpoint's `config.json`** and + looks the placeholder metadata up by it. Whatever you register under that + string decides whether placeholders are bulk-prepended or interleaved at the + original `` positions, and whether the chat template sees a string + or an OpenAI content list. +- so if the anchor was measured on stock trtllm, **register the same + metadata values the stock model type carries**, or the two runs are not + comparable. They are readable without touching model code: + `MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder_metadata("")`. + +## What the target adds + +**A second module tree.** The vision tower plus the projector that maps its +output into the language model's hidden size. Built in `__init__` beside the +language model, under the same flat-declaration rule: `torch.empty`, HF +storage layout, one assert per axis the assembly depends on. + +**A second weight family.** Usually a different layout convention from the +language model's, and commonly **unquantized** while the language model is +quantized — so the manifest's dtype expectations cannot be assumed uniform +across the target for the first time. (Verified on the qwen3.6 checkpoint: 291 +`quantized_layers` entries, none of them a vision layer, and no `*_scale` +tensor under `model.visual.` — so every vision projection is a plain bf16 GEMM +with a bias while the language model is FP8 + NVFP4.) + +**A splice.** See answer 3. + +**And possibly a change to the language model's RoPE.** A VLM whose text model +uses mRoPE is running a *degenerate* mRoPE on text — all three position rows +carry the same value — and a text-only target may legitimately take row 0 and +run plain RoPE. Image spans break that: the three axes genuinely differ. Branch +on the feature flag (a per-engine constant, so the text path stays +bit-identical) rather than on the batch. + +## How to check the tower before booting an engine + +Do this first. It is an afternoon's work and it converts every later failure +from "somewhere in 27 blocks plus an engine" into "the plumbing". + +Load the checkpoint's vision weights, build the HF reference module from the +checkpoint's own `vision_config`, and run both on the same random +`pixel_values` at several grid shapes. **Compare against the reference's own +bf16 rounding, not against an absolute threshold**: a 27-block residual tower +has a noise floor, and on the qwen3.6 tower the target's max deviation from an +fp32 reference (0.035-0.091) was *smaller* than the reference's own bf16 +deviation (0.036-0.173) at three of four shapes. A bare `rel < 2%` gate called +that broken three times out of four. + +Two construction notes for the reference: build it on `meta`, then +`to_empty(device=...)` and `load_state_dict(..., assign=True)` — and re-derive +any **non-persistent buffer** by hand afterwards (the rotary embedding's +`inv_freq` is absent from the state dict, and `.to()` on a meta buffer raises). diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-parallel/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-parallel/SKILL.md new file mode 100644 index 000000000000..0e389d71c174 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-parallel/SKILL.md @@ -0,0 +1,95 @@ +--- +name: staircase-modeling-parallel +description: Load when a target runs above world size 1, or when one modeling codebase must cover a family of parallel configurations (a `unified` segment) — what a rank does differently, where the collectives belong, and how a unified target declares and gates its member set. +--- + +# staircase-modeling-parallel + +Loaded whenever the target is not a plain `tp1`: a multi-rank single +configuration, or a `unified` segment covering a family. + +| reference | read it when | +|---|---| +| `references/multi-rank.md` | world size > 1 — sharding, collectives, `lm_head` vocab parallelism, manifest `src` transforms, what attention DP changes on top of TP | +| `references/unified.md` | the segment is `unified` — config folding, the degenerate path, the Coverage section, member selection | + +Both, for a unified target at world size > 1, which is the normal case. + +## What changes above world size 1 + +The short version; `references/multi-rank.md` is normative. + +- **The manifest gains a `src` transform column.** A rank loads a slice, and + which slice is a function of the rank. Its ordering rule matters as much as + its content +- **`lm_head` vocab parallelism is a manifest obligation**, not something the + runtime does for you +- **Collectives are catalog entries like everything else** and they sit at + specific points in the forward. Which points is a property of the topology, + not of the model — the HF reference says nothing about it, which is why the + pre-audit walks the segment separately +- **Coverage asserts change under EP**: a rank holds a window of experts, so + "every checkpoint key consumed" becomes "every key in my window consumed, + and the rest declared off-window" +- **Attention DP replaces several TP conclusions rather than adding to them.** + Read that part of `multi-rank.md` before assuming a TP fact carries over +- **Every rank must see the same batch composition** when attention DP is off. + The schedulers stay in step by construction, and a forward that branches on + anything rank-local breaks that silently + +## Unified targets + +A `unified` segment is **one modeling codebase over a family of parallel +configurations**, not a per-configuration copy. The configuration is folded at +runtime from the `mapping` object; the code carries the branches. + +Three things follow, all in `references/unified.md`: + +**The segment name promises structure, not coverage.** `unified` says "one +codebase"; it does not say which configurations work. That claim lives in the +run journal's **Coverage** section — one row per member, each with its own gate +records, plus the members deliberately excluded and why. Adding PP or DCP later +extends Coverage; it never renames the directory. + +**Members are chosen by what fits.** The default set is `TEP` and `DEP` +at the world size the user gave (4 unless told otherwise), plus `TP1` **when +the checkpoint fits on one card** — TP1 is the only configuration where every +parallel branch folds off, and a unified target's degenerate path is not the +same code its single-configuration siblings run. A member that does not fit is +recorded as excluded with the reason, not silently dropped. + +**Every member runs the full accuracy gate.** Not one member plus an +equivalence argument: there is no static equivalence check any more, so the +gates are the only evidence that the members agree. That is the price of the +form, and it is per-member hours. + +A member is not a config file. Topology comes from the routing decision and +from the gate's own arguments: `--tp_size` / `--moe_ep_size` / +`--enable_attention_dp` on the boot gate +(`examples/llm-api/quickstart_advanced.py`), and the same three as LLM-API +arguments — `tensor_parallel_size`, `moe_expert_parallel_size`, +`enable_attention_dp` — in the accuracy case. Which member is default is the +first row of Coverage. Every gate run of every member exports +`TRTLLM_MODELING_V2=require` before the ranks start; under `auto` a member that +missed the target's criteria would quietly be the built-in implementation, and +its row would be a record of the wrong system. + + + +## Coverage section + +```markdown +## Coverage + +One modeling codebase, over the configurations below. Each member carries its +own gate records. + +| member | topology | boot | accuracy | note | +|---|---|---|---|---| +| tep4 | `--tp_size 4 --moe_ep_size 4` | PASS | GSM8K 63.91 | default member | +| dep4 | `--tp_size 4 --moe_ep_size 4 --enable_attention_dp` | PASS | GSM8K 63.38 | | +| tp1 | — | — | — | excluded: ~420 GiB, does not fit one card | +``` diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-parallel/references/multi-rank.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-parallel/references/multi-rank.md new file mode 100644 index 000000000000..34d90967ae49 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-parallel/references/multi-rank.md @@ -0,0 +1,267 @@ +# Multi-rank launch — what changes above world size 1 + +> Extracted from the former `trtllm-runtime-integration.md` §12; the section +> number is kept so existing citations still resolve. + +## 12. Multi-rank launch — what changes above world size 1 + +Everything here was measured on B200 hosts at world size 4; a record is keyed +by GPU architecture and nothing else. Two layers, and they answer different +questions: the **launch substrate** (process model, how a rank reaches the +target, the failure modes that exist only above world size 1) applies to every +multi-rank segment, while **what a rank does differently** and **what attention +DP changes** are per-topology and were each established by the first target on +that topology. See "Still not established here" at the end for what no target +has reached yet. + +### The process picture + +`LLM(..., tensor_parallel_size=N)` with `N > 1` does not build the model +in the calling process. It spawns one MPI worker per rank +(`MpiPoolSession` → `mpi4py.futures.MPIPoolExecutor` → `MPI_Comm_spawn`) +and the engine is constructed there; the launcher becomes a proxy. The +single-process-worker flag is read only after that branch, so at world +size > 1 it neither helps nor hurts. + +### What reaches a rank, and how + +mpi4py re-imports the launcher's **main module** in every spawned worker, +under `__name__ == "__worker__"`. Measured consequences: + +| | launcher | worker rank | +|---|---|---| +| module-scope code | runs | **runs** | +| `if __name__ == "__main__":` body | runs | does not run | +| `sys.argv` | full | **script path only** | +| `os.environ` | — | **inherited** | + +A target is **not** registered by a script the ranks re-import. `_router_index.py` +maps `architectures[0]` to `models//routing.py`, routing returns a +synthetic class name, and `get_registered_model_class(name)` resolves it — in +whichever process builds the model. There is nothing for a launcher to +register and nothing for a worker to re-import, so nothing above depends on +where a target's code sits in a file. + +What does have to reach a rank is the **mode**. `TRTLLM_MODELING_V2` is +`off` / `auto` / `require`; anything attributed to modeling_v2 runs under +`require`; and it must be **exported before the ranks start**, not merely +before `LLM(...)` — see the second failure mode below. A driver resolving a +target while its workers resolve the built-in is the silent split the variable +exists to prevent. + +### Two failure modes that exist only above world size 1 + +**The loaded module must be bound in `sys.modules`.** The engine resolves +`architectures[0]` to a class object in the launcher and ships it to the +ranks by pickle, which stores a class *by reference* — `__module__` plus +`__qualname__` — and both sides resolve those names through +`sys.modules`. Loading by path without binding the name fails in the +launcher, before any rank starts: + +``` +PicklingError: Can't pickle : + import of module 'staircase_target_modeling' failed +``` + +Binding it *twice* is equally fatal — the second load leaves a different +class object under the same name and pickle's identity comparison rejects it +(`it's not the same object as ...`) — so the load has to be idempotent. + +In tree both halves are satisfied by construction: a target is an ordinary +module of the modeling_v2 package, imported by name through the registry, so +the name resolves in every process and the import is idempotent. That is what +the registry buys, and it is why the rule is stated rather than dropped — any +mechanism that reaches a model class by file path re-opens both failures. + +**An environment variable set after MPI initializes never reaches a +rank.** Importing `tensorrt_llm` initializes MPI, and OpenMPI hands a +spawned process the environment as it stood at that moment. A tool that +exports into its own environment must do so *before* that import; one +that hands a freshly spawned server a prepared environment satisfies this +by construction. Long-lived ranks under `trtllm-llmapi-launch` receive the +environment once, at launch, so a value set later reaches the driver and not +them. `TRTLLM_MODELING_V2` is the variable this matters most for. + +### The parallel segment and its knobs + +| segment | topology arguments (LLM API, and the boot gate's flag) | +|---|---| +| `tp` | `tensor_parallel_size: N` (`--tp_size N`) | +| `tep` | the above plus `moe_expert_parallel_size: N` (`--moe_ep_size N`) | +| `dep` | the above plus `enable_attention_dp: true` (`--enable_attention_dp`) | + +Routing derives the target *from* these, and the target then asserts every one +of them against the mapping the engine actually built. + +`tep4` and `dep4` were both observed building a serving engine and +generating correct greedy text from the stock DeepSeek-V3-Lite NVFP4 +checkpoint (254 s and 142 s to a ready engine). That is a statement about +the LLM API on this host — not about any modeling_v2 target. + +### What the engine hands the model + +`model_config.mapping` carries the rank's place in the topology. Fields +observed on `Mapping(world_size=4, tp_size=4, moe_ep_size=4, rank=1)`: + +| field | value | +|---|---| +| `rank` / `world_size` | 1 / 4 | +| `tp_size` / `tp_rank` | 4 / 1 | +| `pp_size` | 1 | +| `moe_ep_size` / `moe_ep_rank` | 4 / 1 | +| `moe_tp_size` / `moe_tp_rank` | 1 / 0 | +| `enable_attention_dp` | False | +| `tp_group` | `[0, 1, 2, 3]` | + +Read the topology off this object, never off the segment string: the +segment names the intent, `mapping` is what the engine actually built. + +### What a rank actually has to do differently + +Established by the first multi-rank target (`tep4`, world size 4). These +are the four things this section used to list as unobserved. + +**`lm_head` is topology-aware, and that is a manifest obligation §4/§5 do +not mention.** At `tp_size = 4` `DecoderModelForCausalLM` builds an +`LMHead` of `[vocab/tp, hidden]` — vocab-parallel — so the manifest must +feed *this rank's contiguous block of vocabulary rows*, and the logits +gather stays upstream's. Consequence: `vocab % tp_size == 0` belongs in +the construction asserts, because nothing else checks it. + +**Where the collectives belong: after every row/column-sharded producer, +and nowhere else.** For a TP attention + EP MoE layer that is two per +layer — after `o_proj`, and after `routed_window + shared_partial` +(summed locally first, so one collective serves the whole MLP). The entry is +`comm/allreduce`, and **it is not in tree**: it is one of the 13 catalog +entries the migration left behind, so a target that needs it is a vocabulary +gap for smith, which must bring the entry's contract **and** its receipt, not +just the wrapper — otherwise the target consumes an op with no certification +record at all. What the out-of-tree receipt established, kept here because it +is the thing the entry is consulted for: the workspace-free certified path is +`strategy=0` (NCCL) or `8`, `workspace=None`. Keeping `op=0` (plain sum) and leaving the existing +fused-add-rmsnorm in place preserves the single-rank residual structure, which +is also the CUDA-graph-friendly shape. + +**Which transport, though, is worth about 2x at decode sizes, and the +workspace-free default is the slow one.** A decode message is three orders +of magnitude below the point where NCCL's ring algorithm starts to pay for +itself, and a blocking collective serializes far more than its own time. So a +workspace-free all-reduce is the right thing to *gate* on — it is the +certified path, and it needs no allocation — and the wrong thing to quote a +decode latency from. Choosing a transport, the measured crossover, and the two +profiling traps around it are tuning material, and they live with the tuning +workflow, which is separate and later work; onboarding neither measures nor +gates them. + +**The manifest's `src` transforms: a fourth column, applied before the +relayout.** `(ckpt_key | tuple, src, dst_index, transform)`, with `src` +selecting this rank's slice per key of a multi-key row. Order is +load-bearing: relayout transforms are functions of the *per-rank* row +count, so a transform that ran on the whole tensor and sharded afterwards +produces a different byte order. Column shards are strided views, so +densify before any transform reinterprets bytes. + +**Per-rank parameter shapes** follow the topology mechanically — heads, +dense/shared intermediates and the expert-axis window all divide — with +one trap worth stating: re-derive every kernel alignment rule at the +*per-rank* width rather than inheriting the single-rank conclusion. + +**Coverage asserts change shape.** `leftover == {}` no longer holds: under +EP each rank legitimately leaves the off-window experts' keys unconsumed. +Replace it with an explicitly predicted leftover set computed from the +rank's window, and add a parameter-side assert that the shell-registered +parameters are exactly `{lm_head.weight}`. + +How MLA's latent cache behaves under attention TP is a mechanism fact +rather than a runtime one, and lives in +`.claude/skills/staircase-modeling-core/references/latent-kv-cache.md`: +it is **replicated, not sharded** — attention TP buys zero KV memory on +MLA. + +### What attention data parallelism changes on top of that + +Established by the first `depN` target (`dep4`, world size 4, an MLA + EP +MoE checkpoint). **Three of the four TP bullets above come out +differently**, so read this section as replacing them rather than adding to +them whenever `enable_attention_dp` is on. + +**The split is over requests, not over heads.** Attention is *replicated*: +each rank builds the full head count (32, not `32/tp`), holds the full +attention weights, and serves its own subset of the batch. Consequently +**the post-`o_proj` all-reduce disappears** — a rank's attention output is +already complete for its own tokens. + +**`lm_head` is replicated, not vocab-parallel.** At `tp_size = 4` *without* +attention DP the shell builds an `LMHead` of `[vocab/tp, hidden]`; with it, +the shell builds the full `[vocab, hidden]` on every rank. The manifest +obligation reverses, and `vocab % tp_size == 0` stops being a construction +requirement. + +**The manifest's `src` transform column is a TP artifact.** Nothing outside +the routed expert stacks is sharded, so a `depN` manifest is the tp1 +three-column form with the expert loop windowed. + +**The collectives move to the MoE and change identity.** Not two +all-reduces per layer, but one +`tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather` + one +`.../catalog/comm/reducescatter` per **MoE** layer (a dense layer has none). +The gather **belongs before the router GEMM, not merely before the expert +call** — see +`.claude/skills/staircase-modeling-core/references/expert-weight-packing.md` +for why the EP tiling invariant depends on that placement and fails silently +otherwise. + +**How a rank learns the other ranks' token counts:** +`attn_metadata.all_rank_num_tokens`, a host int list, identical on every +rank. Every rank pads to `max(...)` so both collectives run in their +uniform (`sizes=None`) form — which is also the only form a CUDA graph can +replay, since `sizes` is a host argument frozen at capture. + +**That padding creates an obligation.** Collectives pair **by position** on +the communicator, and at equal byte counts an ordering divergence is +**silent** — every rank wrong in 98-99% of elements, bitwise reproducibly, +no hang. Padding guarantees equal byte counts, so a `depN` forward must +issue an identical call sequence on every rank and must not wait for a hang +to detect that it did not. Certified in +`tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather.md` and +`.../catalog/comm/reducescatter.md`, whose 4-rank rank bodies and launcher +live in `tests/unittest/_torch/modeling_v2/comm/`. + +**The runtime is not a second party on that communicator.** Traced live: +one NCCL communicator per rank, 14,036 collectives over 242 forwards, +**zero** issued by the runtime. The engine's own attention-DP +synchronization is host-side MPI on a *different* communicator +(`MPIDist.tp_comm`, built by `MPI_Comm_create_group`). So a target's +collectives cannot be mispaired against the engine's. + +**The CUDA-graph gate that looks like a landmine and is not.** +`cuda_graph_runner.py` replays a graph under `enable_attention_dp` only if +*all* ranks are generation-only **and** their batch sizes are exactly +equal; the padding that would force equality is off by default. Measured at +steady-state serving load, it **never binds**: `cudaGraphLaunch` = 1.00 per +rank per decode step, i.e. **100% replay coverage** without +`enable_padding`. It does bind on a draining workload — an accuracy-gate +run observed captures but no replays across 334 forwards, where the batch +shrinks monotonically. Both observations are correct about different loads. +Consequences measured on `dep4`: `cuda_graph_config.enable_padding: true` +bought nothing and cost **-2.55%** at concurrency 256 (it can only coarsen +the grid), and `attention_dp_config.enable_balance: true` cost **-13.10%** +at concurrency 128 with mean TTFT doubling, because its `batching_wait_iters` +hold costs more than the alignment buys once random arrival already +balances the ranks. + +**KV cache.** No rank factor: see +`.claude/skills/staircase-modeling-core/references/latent-kv-cache.md` — on +MLA the per-rank pool is identical across tp1/tepN/depN, and what `dep` +buys is `world_size x` *aggregate* capacity from holding disjoint requests, +not a narrower pool. + +### Still not established here + +Pipeline parallelism (`pp_size > 1`); multi-node, where the loopback +pinning that a single-host spawn relies on must be overridden; and +`moe_tp_size > 1`, i.e. splitting experts a second way along the +intermediate dimension. + diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-parallel/references/unified.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-parallel/references/unified.md new file mode 100644 index 000000000000..d682184b547d --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-parallel/references/unified.md @@ -0,0 +1,110 @@ +# Unified targets — one codebase over a family of parallel configurations + +A `unified` segment holds one `modeling.py` and one `weights.py` that serve +several parallel configurations, folded at runtime from the `mapping` object +the engine hands the model. Measured on deepseek-v3-lite over +`{tp1, tep4, dep4}`: 1475 + 448 lines against 3538 + 1032 for the three +single-configuration siblings — 58% less code, one place to fix a bug. + +## How the folding works + +Every configuration-dependent branch reads `mapping`, never a constant: + +``` +self.weight_tp > 1 tensor parallelism is on +self.ep_size > 1 expert parallelism is on +self.dp attention data parallelism is on +``` + +These are **per-capture constants** in graph terms — read once at +construction, never per step — so they are legal inside a captured forward and +cost nothing at replay. What is *not* legal is branching on anything that +varies per step or per rank at runtime; see +`staircase-modeling-core/references/graph-and-sync.md`. + +## The degenerate path is not free + +A unified target at `tp1` is **not** running the same code its `tp1` sibling +runs. It runs the general path with every branch folded off, and the general +path can contain structure the sibling has no reason to have. + +Measured instance: the unified deepseek-v3-lite emits one extra +`torch/split.py` at `tp1` and `tep4` that the siblings do not, because it +always goes through the expert-call chunk loop and they have none. At those +configurations the loop runs once over a single full-length chunk, and +`torch.split(t, [t.shape[0]], 0)[0]` is a zero-copy view — same storage +pointer, shape, strides and contiguity — so every downstream op sees identical +bytes at the same address. + +Harmless there, and the point is the shape of the reasoning: **the degenerate +configuration exercises code no other member does**, which is why it is worth +gating when the checkpoint fits on one card, and why "the multi-rank members +are green" does not imply it. + +## Choosing the member set + +Ask, per candidate: **does this checkpoint fit at this configuration on this +machine?** + +- `TEP` and `DEP` at the world size the user gave — default 4 — are the + baseline pair. They differ in what they exercise: TEP shards attention + weights and reduces; DEP replicates attention and equalizes nothing, so the + batch composition diverges per rank +- `TP1` if it fits. It is the only configuration that folds every branch off +- larger world sizes on request; they extend Coverage rather than replacing it + +A candidate that does not fit is **recorded as excluded with the arithmetic** +— per-rank weight estimate against card capacity. That line is what stops the +next reader re-deriving it. + +## Gating + +Each member runs the **full** accuracy gate, not a spot check. There is no +static equivalence proof to lean on: an earlier experiment carried one — an +AST-level check comparing the unified target's vocabulary, weight manifest and +per-layer op sequence against each sibling's — and it was deliberately removed +once `unified` became the default form, because both of its comparisons need +sibling targets to exist and new unified targets have no siblings. + +So the gates are the only evidence that the members agree, and the cost is +`N × (boot + accuracy)`. The boot gate is minutes; accuracy is roughly an hour +each. Budget it before promising a member. + +**The orchestrator enforces this; it is no longer a promise made in prose.** +The member set is declared in `task.yaml` as `members:` — a list of concrete +segments (`[tp4, tep4, dep4]`), or a mapping from assembly unit to one layered +over `default` — and it is legal only under `parallel: unified`, since a +concrete segment already names the single configuration it is gated at. Every +member must declare the same world size, because one allocation serves them +all, and `unified` is not a legal member: it is the family, not a member of +it. Omitted, it defaults to `[tp, tep, dep]` at the unified world +size — the baseline pair plus the degenerate path. + +The per-unit gate loop runs boot and accuracy once **per member**, records one +verdict each (never collapsed — three measurements is the point), and advances +the unit only when all of them pass. Each verdict names its member, so +`[PASS] accuracy[gsm8k]@dep4` is what the journal, `progress.yaml` and the +reviewer's evidence carry; a line that named only the gate could not be told +apart from the two members it did not run. Preflight prints the member set and +the resulting gate count before the run starts, which is what makes `members:` +the knob for buying a cheaper answer honestly — one member, named — rather +than a cheap answer labelled as an expensive one. + +The witness gates (`host_offload`, `disagg`) run at the **first** member only. +They measure a contrast between two legs of one configuration — a residency +delta, or token identity across the disaggregation hop — rather than a number +per configuration, and the residency witness's `expected_mib` is a single +prediction from the checkpoint's own shapes, so re-asking it at a member whose +attention weights are replicated rather than sharded compares a different +delta against an expectation computed for another configuration. + +Every member's gates run with `TRTLLM_MODELING_V2=require` exported before the +ranks start, so a member whose topology missed the target's routing criteria +fails loudly instead of recording the built-in implementation's numbers. + +## Adding a configuration later + +Extend Coverage, run that member's gates with that member's topology +arguments, record them. **The directory does not get renamed** — that is the +whole reason the segment name promises structure rather than coverage. PP and +DCP land the same way when they land. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/SKILL.md new file mode 100644 index 000000000000..1d926e6006b6 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/SKILL.md @@ -0,0 +1,119 @@ +--- +name: staircase-modeling-speculator +description: Load when adding speculative decoding to a target — MTP, Eagle3, DFlash, DSpark and the other one-engine drafting modes. It changes the forward and the weight-loading path together, and it needs a third gate, because rejection sampling makes a miscomputed draft produce correct text. +--- + +# staircase-modeling-speculator + +Covers every `--features` value that makes the target draft its own tokens: +`mtp`, `eagle3`, `dflash`, `dspark`, and the rest of the one-engine family. One +skill because they share a binding surface and a gate; they differ in what the +draft path computes, which is `references/algorithms.md`. + +**A checkpoint's `mtp.*` namespace does not mean MTP.** DeepSeek-V4.1 stores +its DSpark block drafter there, with `num_nextn_predict_layers: 3` beside it. +Read what the stages compute before choosing the feature — the DSpark section +of `references/algorithms.md` says how to tell — because `mtp` over DSpark +weights either fails to load or drafts garbage that verification quietly +rejects, and the second case passes accuracy. The schema refuses `mtp` and +`dspark` in one spec for this reason. + +| reference | read it when | +|---|---| +| `references/runtime-boundary.md` | always — what the runtime owns of the drafting loop, what the target owes, and where the two meet | +| `references/algorithms.md` | choosing or implementing a specific mode | +| `references/multi-token-prediction.md` | the mode is MTP — mechanism, the concat-order trap, checkpoint layout, HBM arithmetic | +| `references/acceptance-gate.md` | **always, before claiming the feature works** | + +## What this feature changes + +Unlike most features it touches **two** paths, and forgetting the second is +the common failure: + +1. **the forward** — the draft module runs after the trunk, `max_draft_len` + times, and its output feeds verification +2. **weight loading** — the draft module's weights stop being a + known-not-loaded key set and start being loaded. On a DeepSeek-family + checkpoint that is hundreds of keys and several GiB per rank + +It also changes what the engine allocates: the KV pool grows by the draft +module's layers. That is why the variant's arguments often carry a second key +that is not a tuning knob — see below. + +## Attachment + +The draft path is written **flat into `modeling.py`**, not into an adapter. +There is no runtime-owned service here: drafting is a forward-shaped change, +and the runtime's part (the drafting loop, rejection sampling) already exists +and needs nothing registered. + +Concretely, per `references/runtime-boundary.md`: the shell gains the binding +the runtime looks for, the target's core model (`ModelingV2Core`) gains the +draft module, and the no-feature path must stay **bit-identical** — when the +feature is off, the forward is exactly what the recorded gate measured. + +Do not inherit the upstream speculative-decoding base class. It reaches the +draft model through a module-level function with no override point, which +would construct the engine's own in-tree draft module — the thing this project +exists not to use. + +## The arguments assembler owns + +The variant is activated through the LLM API, not through a config file: +`speculative_config=MTPDecodingConfig(max_draft_len=N)` for MTP, +`DSparkDecodingConfig(max_draft_len=)` for DSpark, and the +corresponding config for the other modes. It is not a tuning variant: it +selects a second forward path *and* a second weight-loading path. Record it in +the run journal and in the accuracy case that gates it, so the two agree. + +It commonly carries a **boot requirement** alongside the activation keys. The +draft module's weights plus the enlarged KV pool can push the engine past what +the default memory fraction leaves, and the failure is not an OOM exception — +measured on one target: three of four ranks spinning in `cudaFree` inside the +caching allocator while the fourth waits at an MPI barrier, no error, no +progress. The dep4 MTP variant carries +`kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.75)` for exactly +that reason. Discovering it is part of gating this feature; recording it, with +the ledger, is part of the deliverable. **tuner may not change either key.** + +`max_draft_len` is spelled out on purpose. Left unset it resolves to 1, not to +anything workload-derived — and it is a **certification boundary before it is +a knob**: it sets the attention entries' `predicted_tokens_per_seq`, so a +value past what those entries certify needs `staircase-kernel-verify` first, +not a bigger number. + +## Gates + +boot with the variant's arguments → accuracy → **acceptance**, all three with +`TRTLLM_MODELING_V2=require` exported before the ranks start — except the +stock leg of the third, which asks for `off` on purpose. + +The third one is not optional and not a performance measurement. +`references/acceptance-gate.md` has the protocol; the reason is this: + +> Rejection sampling holds the emitted distribution to the target model's. A +> *miscomputed* draft path therefore produces **correct text**, just more +> slowly — every draft rejected. The boot gate passes. Accuracy passes. Only +> speed moves. + +So the first two gates are structurally blind here, and `acceptance_length` +against a stock reference is the only detector. `1.0` means total rejection. +Clearly above 1.0 but well below the reference means subtly wrong — which is +the case that looks like success. + +The reference is the same checkpoint under **stock in-tree modeling** at the +same workload and the same `speculative_config` — that is +`TRTLLM_MODELING_V2=off` against `=require`, **one variable rather than two +harnesses**. + +Accuracy still runs, and it must be compared against an identity run **in the +same session**: one target measured 94.7688 and 95.0720 on a bit-identical +forward a day apart. + +## Report additions + +- the mode, and `max_draft_len` with the certification bound that permits it +- the acceptance number, the reference it was measured against, and the + workload both saw +- the boot requirement, if one was found, with the ledger that established it +- confirmation that the no-feature path is unchanged, and how that was checked diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/acceptance-gate.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/acceptance-gate.md new file mode 100644 index 000000000000..97d3040e8389 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/acceptance-gate.md @@ -0,0 +1,193 @@ +# The acceptance gate + +The third gate for speculative decoding. It exists because the first two are +structurally incapable of seeing a defect in the draft path. + +## Why the boot and accuracy gates are blind + +Rejection sampling holds the emitted distribution to the **target** model's, +whatever the draft proposes. So a draft module that computes the wrong thing +does not produce wrong text — it produces the same text, with every draft +rejected, more slowly. + +| gate | what a broken draft does to it | +|---|---| +| boot | passes: the continuations are the target model's | +| accuracy | passes: the distribution is the target model's, to the last decimal | +| throughput | drops — but you have nothing to compare it against | + +The failure mode is not "it crashes" or "it scores badly". It is **it works, +and the feature is worthless**, which reads as success. + +The same reasoning applies to any future feature whose output is +distribution-preserving by construction. That is the rule to check a new +feature against, not the specific mechanism. A feature earns a third gate by +exactly one property: a *miscomputed* implementation still produces a correct +answer. + +## The measurement + +`acceptance_length` — mean accepted tokens per engine step, including the +always-accepted trunk token. A run with `max_draft_len = k` has a ceiling of +`k + 1`. + +In tree it is not a separate harness. The gate is a case in the accuracy +suite, beside the model family's other gates: + +```python +with LLM( + MODEL_PATH, + speculative_config=..., + kv_cache_config=..., + cuda_graph_config=CudaGraphConfig(), + enable_iter_perf_stats=True, + **TOPOLOGY, +) as llm: + task.evaluate(llm) # the same workload + al = compute_acceptance_length(llm) # accuracy_core.py + assert_acceptance_length(ACCEPTANCE_KEY, al) # references/acceptance_length.yaml +``` + +`compute_acceptance_length` reads `llm.get_stats()` and folds the iterations +that carried draft tokens: +`(numAcceptedTokens + numRequestsWithDraftTokens) / numRequestsWithDraftTokens` +— which needs `enable_iter_perf_stats=True` on the LLM. +`assert_acceptance_length` checks the observed value against `min_al` in +`tests/integration/defs/accuracy/references/acceptance_length.yaml`. + +**The two legs are one case parametrized over the mode.** The whole comparison +is `TRTLLM_MODELING_V2=require` against `=off` on the same checkpoint, same +workload, same `speculative_config` — one variable, not two harnesses. Because +a multi-rank job's mode is fixed when its ranks start, each leg asserts the +environment it was given rather than choosing one, and skips otherwise. + +Both legs read the **same** recorded minimum, and that is what makes the pair +informative: + +| outcome | reading | +|---|---| +| modeling_v2 fails, stock passes | the draft path regressed | +| both fail | the anchor is stale — re-derive it, do not blame the target | + +**The anchor is populated from the stock leg.** Populating it from the +target's own number would make the gate self-referential. + +Reading a single number alone tells you almost nothing: + +| value | meaning | +|---|---| +| `1.0` | total rejection — the draft path is wrong, or off | +| clearly above 1.0, well below the reference | **subtly wrong** — the dangerous case | +| at the reference | the draft path computes what the checkpoint says | + +### Absence of measurement is not a low number + +A `None`, a zero-iteration result or an empty stats list is **absence of +measurement**, never a failing gate; do not read it as 1.0. + +This is not hypothetical. Out of tree the perf harness's acceptance flag +returned `acceptance_length=none (no drafted request) over 0 iterations` for a +**one-model** mode such as `MTP_EAGLE_ONE_MODEL`, and returned exactly the same +**on stock trtllm's own modeling of the same checkpoint**. The mechanism: the +engine's stats aggregation reads each request's `py_draft_tokens`, but the +one-model path resets that to `[0] * max_total_draft_tokens` before the forward +and keeps the real ids in device tensors, so every iteration reported +`numRequestsWithDraftTokens 0` and the harness — correctly — folded none of +them in. + +Two things follow, and both survive the move in tree. The in-tree detector +turns the silence into a loud failure rather than a number: +`compute_acceptance_length` asserts there was at least one iteration with +speculative stats. And the control stays the same shape — run the **stock** leg +first: a run that reports nothing on *both* systems is upstream's +instrumentation, and a run that reports nothing on ours alone is a defect in +the target. + +### A short answer is not a measurement either + +The orchestrator's gate drives `quickstart_advanced.py` on the task's +`accuracy.acceptance.prompts` (or `prompts_file`), falling back to the boot +prompts. An answer of two tokens and EOS is a context step plus one generation +step, and neither verifies a draft: the rate reads 1.00 for any drafter. The +gate therefore fails a workload with fewer than `min_draft_iterations` +generation iterations as **"workload too short to measure acceptance"**, and +every verdict lists the per-prompt rates beside the weighted value. That +verdict is about the workload, not the target: it is fixed in `task.yaml`, +never in the draft path. + +## The reference is the whole gate + +An acceptance number is only interpretable against **the same checkpoint under +stock in-tree modeling, at the same workload, with the same +`speculative_config`**. Four things must match, and the last two are the ones +that get dropped: + +1. same checkpoint +2. stock TensorRT-LLM's own modeling, not ours — i.e. `TRTLLM_MODELING_V2=off` +3. **same workload** — both legs are the same case running the same evaluation + task with the same protocol, so both systems see identical requests. Do not + compare across workloads +4. **same draft length.** A reference at a different `max_draft_len` answers a + different question + +The stock leg's own requirements are assembler's to establish, because the +gate is assembler's to pass. They are typically whatever stock trtllm needs in +order to boot this checkpoint at all — boot knobs, which change batching +capacity rather than what the model predicts. **Acceptance is a property of the +model**, so a reference forced onto different memory knobs is still the right +comparison; only the workload has to match. + +A reference may also need something no argument covers — an environment +variable that steers a transport, for instance. The in-tree instance: the stock +leg of the dep4 MTP gate requires `TRTLLM_CAN_USE_DEEP_EP=0` exported, +because under attention DP + EP stock lands on DeepEPLowLatency, whose dispatch +takes only NVFP4 uint8 hidden states while the MTP layer is bf16; disabling it +lands stock on AllGatherReduceScatter — the strategy the modeling_v2 target +implements by hand, which makes the two *more* comparable rather than less. If +a reference needs such a variable, it goes in the case that runs it, in full, +because nothing else will carry it, and it is exported in the launching +environment like the mode switch. + +## What a passing gate looked like + +Recorded for deepseek-r1-0528-nvfp4 / sm_103 / dep4, `mtp3`, on 4 GB300s: +`acceptance_length` **3.3514** against stock's **3.2752**, ratio **1.023**. +That row is the one that matters for a speculative variant — boot and accuracy +were both green on the same target and are blind to the draft path. + + + +Draft lengths 1 and 2 remain ungated on that target — they are the dominated +end of the measured draft-length axis. + +## Cost, and the free proxy + +Acceptance instrumentation is not free: measured 1.4% / 3.5% / 4.0% of +throughput at concurrency 1 / 32 / 256 (deepseek-r1-0528-nvfp4 `dep4` on B200 — +the instrumentation is upstream's, so re-measure the cost when it changes), +which also makes a label **non-comparable** to one measured without it. Never +mix the two in a figure. In tree the same warning applies to +`enable_iter_perf_stats=True`, which the gate turns on. + +For a sweep — and whenever the instrumented number is unavailable — use the +proxy instead: a response body already carries +`mean_avg_decoded_tokens_per_iter`, agreeing with the engine's acceptance +length to 0.99–1.02x. Reserve the instrumented measurement for the gate itself. +The sweep half of this is the later perf workflow's, not onboarding's. + +## What the gate does not tell you + +That the feature is *worth it*. Acceptance says the draft path computes the +right thing; whether drafting pays depends on the break-even rate, which is +the extra HBM traffic per step divided by the traffic of a step. That is a +performance question and it belongs to `/tuning`. + +It also flatters easily: a random-prompt harness gives the model its own +repetitive continuation to draft, which drafts near-perfectly. When quoting a +speedup, quote the one from the real-text gate runs too, and say which is +which. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/algorithms.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/algorithms.md new file mode 100644 index 000000000000..d534727d87a2 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/algorithms.md @@ -0,0 +1,184 @@ +# Drafting algorithms — what differs between the modes + +One skill covers several `--features` values because they share a binding +surface and a gate. This file is the part that differs: what each mode's draft +path computes, and therefore what the target has to write. + +The binding surface itself — where the draft module attaches, what the shell +gains, how the weight-loading path changes — is `runtime-boundary.md` and is +the same for all of them. + +## The family + +Every mode here is **one-engine**: one `PyTorchModelEngine` does target +verification and the next round's drafting in a single model forward. Two-model +draft-target is a different shape and is not covered. + +| mode | draft path | weights it turns on | notes | +|---|---|---|---| +| **MTP** | the checkpoint's own MTP module, replayed autoregressively `max_draft_len` times | the MTP module in the checkpoint | `references/multi-token-prediction.md` is the depth. `num_nextn_predict_layers` is a module count, **not** the draft-length knob | +| **Eagle3** | a separate small drafter reading the target's hidden states | a separate drafter checkpoint | needs hidden states exported from the trunk — a forward-shape change the other modes do not require | +| **DFlash** | as upstream `_torch/speculative/dflash.py` | | not yet exercised by any target here | +| **DSpark** | the checkpoint's own block drafter: one 3-stage pass drafts a whole block, then a sequential Markov refine | the drafter stages — on DeepSeek-V4.1 stored as `mtp.*` | **not MTP, whatever the namespace says.** See the DSpark section below | +| n-gram, suffix automaton, PARD | no weights at all — drafts from the emitted prefix | none | no weight-loading change, so the gate reduces to acceptance alone | + +Only MTP has been carried end to end here. Everything about the others below is +read from the in-tree sources and the checkpoints' own reference code, not +measured — say so if you rely on it. + +## What decides the draft length + +`speculative_config.max_draft_len`, spelled out. Two things make it more than +a knob: + +- **left unset it resolves to 1**, not to anything workload-derived +- it becomes the decode kernel's `maxSeqLenQ` — on an MLA layer through the + attention entries' `predicted_tokens_per_seq`, on a standard layer through the + **generation row count** instead, since `predicted_tokens_per_seq` is inert + there (`runtime-boundary.md`). So its ceiling is a **certification** bound: + going past what the MLA/attention entries certify needs + `staircase-kernel-verify` first. The JIT cost is **not** one compile per + distinct value: measured on sm_100, a row count of 2 adds one entry and 3 and + 4 reuse it, so drafting costs exactly one extra first-call compile whatever + the draft length — boot time, not a failure + +The checkpoint's module count does not bound it. On the MTP-Eagle path a +single module is replayed `max_draft_len` times. + +## Choosing, when the checkpoint offers a choice + +If the checkpoint ships a draft module, use it — MTP. Its weights were trained +against the target and its acceptance ceiling is far above what a +weightless mode reaches. + +Reach for a weightless mode when the checkpoint has no draft module and the +workload is repetitive enough that prefix matching pays. It costs no weights, +no KV pool growth and no boot-time JIT, so the gate is cheap; the acceptance +reference is still required, because "drafts nothing, accepts nothing" and +"drafts wrongly, accepts nothing" look identical without one. + +## DSpark — a block drafter that lives where MTP would + +DeepSeek-V4.1-Flash is the first checkpoint here. Its tech report (§2.4.3): +"We omit the MTP module during backbone pre-training and use DSpark for +speculative decoding." The drafter nevertheless ships under `mtp.0..2` with +`num_nextn_predict_layers: 3` in the config, which is exactly what makes the +wrong feature easy to pick. + +### Telling it from MTP + +Any one of these settles it: + +- the config carries `dspark_*` keys — `dspark_block_size`, + `dspark_noise_token_id`, `dspark_target_layer_ids`, `dspark_markov_rank` + (DeepSeek-V4.1 nests them under `text_config`); +- the `mtp.*` weights include `main_proj`, `markov_head`, `confidence_head` + and lack MTP's signature: no `e_proj` / `h_proj` / `enorm` / `hnorm` / + `shared_head`; +- the reference forward drafts a block in one call rather than one token per + module. + +Then the feature is `dspark` and never `mtp`. `MTPDecodingConfig` would read +`num_nextn_predict_layers: 3`, pick vanilla sequential MTP, and build three +chained one-token heads over weights trained as one 3-stage block pass: it +either fails to load the missing MTP keys or drafts garbage that verification +rejects — and the second passes accuracy. If a request says "MTP" for such a +checkpoint, it means "speculate with the bundled drafter", which is `dspark`; +write that down in the spec so a later reader does not "fix" it. + +### What the draft path computes + +From the checkpoint's reference forward. Each item is a place a transliteration +from an MTP or Eagle mental model goes wrong **and still drafts**, so the only +symptom is a lower acceptance length. + +1. **One pass drafts the whole block.** The input is a `block_size`-token block + `[x, NOISE, NOISE, NOISE, NOISE]` at block size 5: `x` is the token the + target just emitted, NOISE is `dspark_noise_token_id` (128799). The three + stages run **once**, in sequence, over all five positions — each a full + block (sliding-window attention plus a 128-expert top-3 MoE, with + hyper-connections) — and the head produces base logits for all five + positions in parallel. Not five passes; not three heads chained. +2. **Attention inside the block is non-causal.** Every block position attends + to every other block position plus the filled ring slots. A causal mask + here is the most likely error to carry over from a one-token drafter. +3. **The ring holds TARGET state.** Each stage owns a sliding-window ring + (window 128). It is filled from the target's hidden states: the *input* of + target layers `dspark_target_layer_ids` (37, 38, 39), averaged over the + hyper-connection copies, concatenated, projected once by stage 0's + `main_proj` + `main_norm`, then per stage through `wkv` and `kv_norm` with + plain RoPE. Inputs, not outputs — vLLM's DeepSeek-V4 path captures + outputs, which is a different checkpoint's contract. The block's own K/V + are transient and never enter the ring. After a multi-token acceptance, + every newly committed position's target state must reach the ring before + the next draft, not just the last one. +4. **The Markov head is applied sequentially.** After the parallel pass, for + each position in order: add a rank-256 bigram bias computed from the + *previous* position's token, then choose this position's token. Applied in + parallel it would condition on NOISE tokens. +5. **The confidence head only schedules.** It predicts per-position acceptance + probabilities so a scheduler can choose how many drafts each request + verifies. It never changes which tokens are accepted, so greedy output is + identical with or without it. TRT does not wire it — `DSparkDecodingConfig` + omits the knobs on purpose and the worker drafts with a zero threshold — + so the full block is always verified. + +**The acceptance ceiling is `block_size + 1`** — 6 at block size 5 — since +verifying k drafts accepts at most k plus the bonus token. Set +`accuracy.acceptance.dspark.min_absolute` against that ceiling, never against +what the target happens to score. + +### What exists in TRT, and where the line falls + +| piece | where | owner | +|---|---|---| +| `SpeculativeDecodingMode.DSPARK` | `tensorrt_llm/_torch/speculative/interface.py` | runtime | +| `DSparkDecodingConfig`, and its `TorchLlmArgs` validation: an unset `speculative_model` means the target's own `mtp.*`, and `block_size` must equal `max_draft_len` | `tensorrt_llm/llmapi/llm_args.py` | runtime | +| `DSv4DSparkWorker`: the drafting loop, the per-request rolling windows, strict or rejection-sampling acceptance, back-filling accepted positions | `tensorrt_llm/_torch/speculative/dspark.py` | runtime | +| `DSparkSpecMetadata`: captures the target layers' hidden states for the worker | `tensorrt_llm/_torch/speculative/dspark.py` | runtime | +| the three stages' forward, the ring-write projection, the Markov chain | the target | **target** | +| what the worker reaches through on the draft container: `block_size`, `num_stages`, `_attn_params`, `write_context_windows`, `write_context_windows_batched`, `forward_batched` | the target | **target** | + +**The draft forward must be the target's own; the runtime owns the loop.** The +stock embedded draft is `DSv4DSparkDraftModel` in +`tensorrt_llm/_torch/models/modeling_dspark.py`. Read it for the interface the +worker expects, never to import or to copy — the rule `runtime-boundary.md` +states for MTP, for the same reason. And do not write a drafting loop, an +acceptance rule or a window manager in the target: those exist. + +### Gate spelling + +`features: [dspark]`. The boot and acceptance gates emit +`--spec_decode_algo=DSPARK --spec_decode_max_draft_len=`, and `N` comes +from the unit's options YAML — `speculative_config.max_draft_len`, or +`block_size` — never from a default: quickstart builds the config from that +flag alone, and a depth of 1 against a drafter trained for a 5-token block +fails validation before the engine starts. + +```yaml +# accuracy.llm_api_options.dspark +speculative_config: + decoding_type: DSpark + max_draft_len: 5 +``` + +For a new architecture the acceptance gate's `=off` leg cannot load at all, so +declare `accuracy.acceptance.dspark.min_absolute` (1.2 is the smoke floor +vLLM's DeepSeek-V4 DSpark recipe uses) and the gate records the weaker, +absolute claim. + +### Open, read from source and not measured + +- The `TorchLlmArgs` DSpark validation reads `dspark_*` at the top level of the + drafter's `config.json`, or inside `dspark_config` / `dflash_config` — **not + under `text_config`**, where DeepSeek-V4.1 keeps them. So + `target_layer_ids`, `mask_token_id` and `markov_rank` stay unset unless the + spec states them, and with `target_layer_ids` unset `DSparkSpecMetadata` + captures nothing: the drafter sees no target context and acceptance falls + toward 1.0, which the acceptance gate reports as a broken draft path. The + accuracy gate's options YAML can state them; `quickstart_advanced.py`, which + the boot and acceptance gates drive, cannot. Decide how they reach the + runtime before the first DSpark boot, and record the decision. +- TRT's own DSpark tests run with chunked prefill, the overlap scheduler and + block reuse all off. A unit gated with them on is the first to exercise + that combination; a failure there is not automatically the target's. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/multi-token-prediction.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/multi-token-prediction.md new file mode 100644 index 000000000000..fdc996b9fc4b --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/multi-token-prediction.md @@ -0,0 +1,300 @@ +# Multi-token prediction (MTP) + +What a target assembled on a checkpoint that ships an MTP module had to +establish about **what the extra layer computes**, because no contract +states it and the checkpoint ships no reference implementation for it. +Written from the `models/deepseek_v3/r1_0528_nvfp4__sm_100__dep4` MTP +increment (DeepSeek-R1-0528, 61 trunk layers, hidden 7168, 256 routed experts +top-8, one MTP module at layer 61). The mechanism — an extra decoder layer that +consumes *both* the trunk's hidden state and the next token's embedding, and +is replayed to produce a draft — recurs across the DeepSeek family and its +derivatives; this file is about the mechanism, not that checkpoint. + +The runtime half of the picture — how the engine finds the drafter, what +calls the layer, and how the batch shape changes between draft steps — is +`.claude/skills/staircase-modeling-speculator/references/runtime-boundary.md`. +This file is only the computation and the weights. + +**The checkpoint is the whole specification here, and it is a partial one.** +DeepSeek's published `modeling_deepseek.py` does **not** implement the MTP +module (verified across four R1/R1-0528 checkpoints on disk, quantized and +not), and `transformers`' own `deepseek_v3` does not either. So the graph +below was reconstructed from the checkpoint's key names, shapes and weight +statistics. Every claim that rests on the reconstruction rather than on a +shape is marked. + +## The mechanism + +The module is one extra decoder layer, structurally identical to a trunk MoE +layer, with a front end bolted on that mixes in the embedding of the token +being predicted: + +``` +e = embed_tokens(input_ids) # embedding of the NEXT token +x = eh_proj( concat( enorm(e), hnorm(h) ) ) # [T, 2*hidden] -> [T, hidden] +x = x + MLA( input_layernorm(x) ) # same block as a trunk layer +x = x + MoE( post_attention_layernorm(x) ) # same structure, different dtype +# shared_head, called separately by the runtime: +logits = lm_head( shared_head.norm(x) ) +``` + +`h` is the trunk's final hidden state for the same position — under +one-model MTP-Eagle the runtime hands the layer the target model's own +output, with no projection in between. + +**"How many MTP layers do we enable" is not the knob.** A checkpoint with +one MTP module (`num_nextn_predict_layers: 1`) is replayed autoregressively: +the same layer runs `max_draft_len` times, each step feeding it the previous +step's output token and hidden state. The draft length is a serving knob, +not a checkpoint property. §13 has the mode-selection rule. + +## The concat order — the one silent-wrong-answer trap + +`eh_proj` is `[hidden, 2*hidden]`. Which half multiplies the embedding +branch and which multiplies the hidden branch is **not** determined by +anything in the checkpoint's metadata, and getting it wrong is a pure +numerical error: no shape mismatch, no assert, no crash. Under rejection +sampling it does not even corrupt output — the drafts are simply always +rejected (see the last section). + +**Two sources disagree, and the naming is the one that is right.** + +* The DeepSeek-V3 report writes the projection as `M_k [ RMSNorm(h) ; + RMSNorm(Emb(t)) ]` — **hidden first**. +* The parameter is named `eh_proj`, with its two gains named `enorm` and + `hnorm` — **embedding first**. + +**Measured: the embedding block is first.** Two independent statistics over +the checkpoint's own weights, each controlled: + +*Column-norm profile.* Take the per-input-dimension L2 norm of each half of +`eh_proj` and correlate it with each RMSNorm gain. The pairing is exclusive: + +| | vs `\|enorm\|` | vs `\|hnorm\|` | vs `\|model.norm\|` | +|---|---|---|---| +| first half `eh_proj[:, :hidden]` | **+0.904** | +0.037 | +0.026 | +| second half `eh_proj[:, hidden:]` | −0.036 | −0.207 | **+0.544** | + +`model.norm` is the trunk's final RMSNorm, which gates the same residual +stream `hnorm` does, so the second half tracking it is the same statement as +the second half being the hidden branch. + +*Functional response.* Feed the embedding branch's actual input, +`enorm(Emb(t))`, through each half and out through `shared_head`, and +measure the entropy of the resulting distribution. **Compare each half +against its own controls, not against the other half** — the first half is +uniformly sharper by ~0.9 nats on *any* input, so a raw cross-half +comparison is confounded: + +| input | via first half | via second half | +|---|---|---| +| `enorm(Emb(t))` — the real embedding-branch input | **2.62 nats** | 6.22 nats | +| `hnorm(gaussian)` | 5.33 | 6.24 | +| gaussian, no gain (control) | 5.36 | 6.23 | +| `enorm(gaussian)` — right gain, wrong vector (control) | 5.09 | 6.17 | + +The first half drops **2.5 nats** on the real embedding and on nothing else. +The second half does not move at all: through it, a real token embedding is +indistinguishable from noise. (Uniform over a 129,280 vocabulary is 11.77 +nats.) + +So: **`eh_proj` consumes `concat(enorm(e), hnorm(h))`.** This is a +reconstruction from weight statistics, not a reading of reference code — it +is strong enough to implement against, and the acceptance rate is what +confirms it end to end. Keep the order as a single named constant at one +place in the layer so flipping it is a one-line experiment. + +## Checkpoint layout + +The module is stored as one more entry in the layer list, at index +`num_hidden_layers`. On this checkpoint that is 790 keys under +`model.layers.61.`: + +| key | count | dtype | shape | load? | +|---|---|---|---|---| +| `enorm.weight`, `hnorm.weight` | 2 | bf16 | `[7168]` | yes | +| `eh_proj.weight` | 1 | bf16 | `[7168, 14336]` | yes | +| `input_layernorm`, `post_attention_layernorm` | 2 | bf16 | `[7168]` | yes | +| `self_attn.*` | 7 | bf16 | **identical to a trunk layer's** | yes | +| `self_attn.{k_proj.k_scale, v_proj.v_scale}` | 2 | fp32 | `[]` | value is **1.0**, as in the trunk | +| `mlp.gate.weight` | 1 | bf16 | `[256, 7168]` | yes | +| `mlp.gate.e_score_correction_bias` | 1 | **fp32** | `[256]` | yes | +| `mlp.experts.{0..255}.{gate,up,down}_proj` | 768 | bf16 | `[2048, 7168]` ×2, `[7168, 2048]` | yes, EP-windowed | +| `mlp.shared_experts.{gate,up,down}_proj` | 3 | bf16 | same | yes | +| `shared_head.norm.weight` | 1 | bf16 | `[7168]` | yes — **a distinct norm**, not `model.norm` | +| `embed_tokens.weight` | 1 | bf16 | `[129280, 7168]` | **no** | +| `shared_head.head.weight` | 1 | bf16 | `[129280, 7168]` | **no** | + +**The last two are bit-identical copies of the trunk's** (`torch.equal` +against `model.embed_tokens.weight` and `lm_head.weight`: both True). The +runtime hands the layer whichever `embed_tokens` and `lm_head` the target's +draft-model container exposes, so pointing them at the trunk's is correct +and saves 1.85 GB per rank. They stay in the weight manifest's *predicted +non-load* set even with MTP on; the other 788 keys flip from non-load to +consumed. + +**The MTP layer's attention geometry is byte-identical to a trunk layer's** +— same `q_a_proj` / `q_b_proj` / `kv_a_proj_with_mqa` / `kv_b_proj` / +`o_proj` shapes, same two LayerNorms, same `q_lora_rank`. Whatever load-time +derivation the trunk's attention needs (absorption operands, row regrouping, +the rope table) applies unchanged. + +## The MTP layer is not quantized, and that is deliberate + +On a quantized export the MTP module can be excluded from quantization +wholesale. Here `hf_quant_config.json` carries `model.layers.61*` as one +wildcard entry in a 63-entry `exclude_modules` list, so **every weight in +the module is bf16** while the trunk's MLP path is NVFP4. + +**Do not re-quantize it at load time to reuse the trunk's expert +vocabulary.** The exclusion is the export's choice about where accuracy is +worth the bytes; quantizing it anyway changes what the checkpoint means, and +the acceptance rate — the only signal that can see the difference — would +absorb the damage silently as a lower draft quality rather than reporting +it. + +The consequence is a vocabulary consequence: the MTP layer's routed experts +need a **bf16** grouped-expert entry, at whatever `(local_experts, hidden, +intermediate)` the parallel split produces, while the trunk's use the NVFP4 +runner. Everything else in the graph — both norms, the concat, the +projection GEMM, the whole MLA block, the router, the shared expert, +`shared_head` — maps onto entries a dense-plus-MoE target already carries. + +## The HBM arithmetic, and the break-even acceptance rate + +A decode step is weight-bandwidth-bound: each rank reads every weight byte +it holds, once. So the cost of drafting is exactly the MTP layer's byte +count, times the number of draft steps. + +Per rank on this checkpoint at `dep4` (64 local experts of 256): + +| | bytes per draft step | +|---|---| +| routed experts, **bf16**: `3 x 2048 x 7168 x 64 x 2` | **5.637 GB** | +| attention, bf16: 187.1 M params x 2 | 374 MB | +| `eh_proj`: 102.8 M x 2 | 206 MB | +| shared expert: `3 x 2048 x 7168 x 2` | 88 MB | +| router | 4 MB | +| **total** | **≈ 6.31 GB** | + +Against a trunk step of ≈ 115 GB per rank (58 NVFP4 MoE layers at 1.585 GB + +61 bf16 attention blocks at 374 MB), that is **+5.5% per draft step**. Note +what the first row means: **the bf16 MTP layer's experts cost 3.56 times +what one NVFP4 trunk MoE layer costs**, purely from the dtype. + +A step now produces `acceptance_length` tokens instead of 1, so drafting +pays for itself when + +``` +acceptance_length >= 1 + max_draft_len * 0.055 +``` + +| `max_draft_len` | extra HBM/step | break-even `acceptance_length` | +|---|---|---| +| 1 | +5.5% | **1.055** | +| 2 | +11.0% | 1.110 | +| 3 | +16.4% | **1.164** | +| 4 | +21.9% | 1.219 | + +**That table is the weight-bytes model, and measurement says it is a lower +bound that stops holding as the batch grows.** Measured on this checkpoint +at `max_draft_len 3`, decomposing each engine step against a non-drafting +one: + +| con | acceptance | step cost — predicted | step cost — **measured** | decode speedup | +|---|---|---|---|---| +| 1 | 3.4904 | 1.164 | **1.511** | 2.311x | +| 32 | 3.3356 | 1.164 | **2.198** | 1.518x | +| 256 | 3.4033 | 1.164 | **2.600** | 1.309x | + +The marginal cost of each successive draft step at con=256 is **+0.645, ++0.569, +0.386**, against con=1's **+0.200, +0.173, +0.137**. The growing +term tracks the **rows** a step carries, not the layer's fixed weight bytes +— it falls off in the same shape the marginal row count does (1→2 rows is ++100%, 2→3 is +50%, 3→4 is +33%). **That is the +memory-bound-to-compute-bound crossover this file calls the genuinely +uncertain end, now measured rather than predicted.** + +So read the weight-bytes table as the floor it is: right at low concurrency, +where a step is bandwidth-bound and drafting really is nearly free, and +roughly 2.2x optimistic at the throughput end. What survives is the +conclusion, because the cost never grows fast enough to catch acceptance: +**at every concurrency measured, acceptance stayed at 2.9x or more of even +the measured break-even**, and no draft length in the certified range lost. + +**The break-even is still very low against the predicted cost** — at +`max_draft_len` 3 it takes only 0.164 extra tokens per step on average — and +two costs that look like they should matter do not: + +* **Collectives.** The MTP layer is an MoE layer, so it adds two per draft + step, and the trunk's grow by `draft_len + 1` in bytes. But at decode + sizes these calls are latency-bound rather than bandwidth-bound (measured + on this target: 2.753 MB in 25.472 µs = 108 GB/s, an order of magnitude + under NVLink), so multiplying the byte count barely moves the time. +* **KV capacity.** One extra layer on a 61-layer pool is **+1.6%**, plus + `max_draft_len - 1` extra tokens per sequence. + +The end that is genuinely uncertain was the **high-concurrency** end — once +a rank's batch already activates all of its local experts and the routed +GEMM is near the HBM roofline, `draft_len + 1`× the rows means the same +weight bytes with several times the arithmetic. It is measured above: the +crossover is real, it costs 2.2x the predicted step cost at con=256, and it +still does not overtake acceptance. + +**Acceptance itself is near-flat in batch size on this mechanism, which is +the other half of why no draft length loses.** Measured across three draft +lengths and nine concurrencies: `max_draft_len` 1 stayed in 1.9446–1.9673 of +a 2.0 ceiling, 2 in 2.6945–2.7707 of 3.0, and 3 in 3.3356–3.4904 of 4.0. +**So a "shorten the draft as the batch grows" schedule has no +acceptance-side motivation here** — any case for one has to be made on cost, +and on this checkpoint the cost side does not invert. (It is also unusable +under attention DP for a structural reason — see +`.claude/skills/staircase-modeling-speculator/references/runtime-boundary.md`.) + +## What only the acceptance rate can tell you + +**An MTP layer that computes the wrong thing does not produce wrong +output.** Rejection sampling guarantees the emitted distribution is the +target model's regardless of draft quality. A layer with the concat order +reversed, a norm on the wrong operand, or the expert stack packed in the +wrong interleave produces **bit-correct text, more slowly**, with every +draft rejected. An accuracy benchmark cannot see it. Neither can the boot +gate. + +`acceptance_length` — the mean tokens emitted per step by requests that +carried a draft, `1.0` meaning total rejection — is the only detector, and +it needs a reference to be read against: + +* **≈ 1.0** at `max_draft_len` 3: the layer is wrong outright. +* **clearly above 1.0 but well below the reference**: the layer is *subtly* + wrong. This is the band the concat order lands in, and nothing else + distinguishes it from "this model just drafts poorly". +* **at the reference**: the layer computes what the checkpoint says. + +The reference is the same checkpoint under stock in-tree modeling — the same +job under `TRTLLM_MODELING_V2=off` instead of `=require` — at the same load +and the same speculative configuration. Acceptance is a property of the +model, so a reference forced onto different memory knobs to boot is still +comparable. + +**One trap when choosing the workload, and it does not point the way it +looks like it should.** A serving harness that builds prompts from uniformly +random token ids seems like it must *understate* acceptance — the +continuation of a nonsense prefix is unpredictable, so the drafts should +miss. Measured, it runs the other way, and strongly. + +Only the **prompt** is random. Generation then runs for a fixed output +length with EOS ignored, so what MTP is drafting is the model's own +continuation of a nonsense prefix — which degenerates into repetition, and +repetition is the easiest thing in the world to draft. Measured on this +checkpoint at `max_draft_len 3` (ceiling 4.0): **`acceptance_length` 3.93 at +concurrency 1 and 3.46 at 32**, i.e. 97.6% and 82.1% of proposed draft +tokens accepted. Real text does not do that. + +So the random-prompt harness **flatters** MTP, and two things follow. A +throughput curve measured on it overstates the gain, so it answers "did +anything get slower" and not "is MTP worth enabling" — that one needs a +real-text workload, and an accuracy-gate run already produces one. And a +future run that sees a *low* acceptance number here must **not** excuse it +as "well, the workload is random": on this harness a low number means the +layer is wrong. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/runtime-boundary.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/runtime-boundary.md new file mode 100644 index 000000000000..c85b25c6fb3c --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/runtime-boundary.md @@ -0,0 +1,488 @@ +# One-engine speculative decoding — the runtime/target boundary + +> Extracted from the former `trtllm-runtime-integration.md` §13; the section +> number is kept so existing citations still resolve. + +## 13. One-engine speculative decoding — what changes when the target drafts + +Everything above describes a target that emits **one** token per generation +sequence per step. A speculative target emits `1 + draft_len`, and the draft +tokens come from a second forward that the target itself owns. This section +is the axis that brings: what the engine looks for on the model object, what +the drafting loop calls, where the ownership line falls, and what a rank has +to do differently inside its forward. + +Read like §12, in two layers. The **binding surface** applies to every +one-engine speculative mode. The **per-step consequences** below it were +established for **MTP-Eagle one-model** (`MTP_EAGLE_ONE_MODEL`) on an MLA + +EP-MoE checkpoint at `dep4`, and are marked where another mode differs. +`.claude/skills/staircase-modeling-speculator/references/multi-token-prediction.md` +carries the other half — what an MTP layer computes, which is a checkpoint +fact rather than a runtime one. + +### The checkpoint picks the mode; the target does not + +`update_spec_config_from_model_config` runs **before the model is built** +and reads the MTP layer count out of the pretrained config +(`num_nextn_predict_layers`, or `mtp_num_hidden_layers` on Qwen3Next-style +configs; 1 if neither is present). `MTPDecodingConfig`'s defaults are +`use_mtp_vanilla=False` and `mtp_eagle_one_model=True`, so: + +| checkpoint layer count | resulting mode | +|---|---| +| `n == 1` | **`MTP_EAGLE_ONE_MODEL`** — one layer of MTP weights, replayed | +| `n > 1` | `MTP` (vanilla) — one distinct layer per draft position | + +**Under MTP-Eagle, `max_draft_len` is not bounded by the checkpoint's layer +count.** That bound belongs to vanilla MTP. The single MTP layer is replayed +autoregressively `max_draft_len` times, so "how many MTP layers do we turn +on" is the wrong question and "what is `max_draft_len`" is the right one. + +**Spell `max_draft_len` out in every variant.** Left unset on the +MTP-Eagle path it resolves to **1**, not to anything derived from the +workload. `max_total_draft_tokens` is then mirrored from it (linear tree), +and `tokens_per_gen_step = 1 + max_total_draft_tokens`. + +### The engine finds the drafter through exactly one getattr + +```python +# _torch/pyexecutor/model_engine.py +def _get_spec_worker(self): + return getattr(self.model, "spec_worker", None) +``` + +That is the whole registration. Everything else the runtime touches on the +model side, it reaches through the worker's own arguments: + +| attribute / callable | type | what reads it | +|---|---|---| +| `model.spec_worker` | `SpecWorkerBase` | the engine's one getattr | +| `model.config` | pretrained config | `update_spec_config_from_loaded_model` (the base shell already provides it) | +| `model.draft_config` | — | read with `getattr(..., None)`; **absent is correct** for a single-checkpoint MTP target | +| `draft_model.mtp_layers` | `nn.ModuleList` | only `[0]` is ever indexed — MTP-Eagle replays one layer | +| `draft_model.embed_tokens` | module | passed to the layer as a kwarg | +| `draft_model.lm_head` | module | passed to `shared_head` | +| `draft_model.model.d2t` | — | read with nested `getattr(..., None)`; **absent is correct** (draft and target share a vocabulary) | +| `mtp_layers[0](...)` | callable | the draft loop, once per draft step | +| `mtp_layers[0].shared_head(h, lm_head, attn_metadata, True)` | method | returns draft logits | + +`draft_model` is a container the target defines and hands to the worker; the +runtime never constructs it and never inspects it beyond the four names +above. + +The layer is called by keyword, with `inputs` splatted in: + +```python +hidden_states = draft_model.mtp_layers[0]( + embed_tokens=draft_model.embed_tokens, + all_rank_num_tokens=, + input_ids=..., position_ids=..., hidden_states=..., + attn_metadata=..., spec_metadata=..., +) +``` + +It returns **one tensor**, `[rows_this_step, hidden]`, unpadded — the same +row count its `input_ids` carried. The loop slices it with its own +`gather_ids` afterwards. (Eagle3 one-model returns a second tensor here; +MTP-Eagle does not.) + +### Where the ownership line falls + +| responsibility | owner | +|---|---| +| accept/reject, rejection sampling, the golden token | runtime | +| KV rewind, `attn_metadata` rewrite between draft steps and its restore | runtime | +| the draft loop, `gather_ids`, position shifting, sampling draft tokens | runtime | +| `runtime_draft_len` scheduling and padding, `(bs, draft_len)` graph capture | runtime | +| the KV pool's extra layer and extra tokens | runtime | +| **the MTP layer's forward** | target | +| **`shared_head`** | target | +| **the `draft_model` container** | target | +| **the shell's speculative branch** | target | +| **loading the MTP layer's weights** | target | +| **attention-DP padding inside the MTP layer** | target | + +### Inheriting the in-tree one-engine shell is not the shortcut it looks like + +`SpecDecOneEngineForCausalLM.__init__` builds its drafter by calling +`get_draft_model(...)`, which is a **module-level function, not a method** — +a subclass has no override point. It dispatches on the config's +`model_type`, and the checkpoint is read exactly as published — there is no +target-owned `config.json` to rewrite — so `model_type` names the upstream +family and the call returns **trtllm's own MTP layer**. Inheriting therefore +hands the one computation this project exists to write to the engine instead. + +This is a statement about what gets constructed, not a rule against +inheriting: a shell already inherits `DecoderModelForCausalLM` from the same +package. Read an in-tree implementation to understand an interface; never to +copy a structure. Writing the branch out (about 40 lines) keeps the forward +readable end to end, which is what the Vocabulary table and the +closed-vocabulary audit both rest on. + +### The shell's shape, and the four things it has to get right + +```python +class ModelingV2(DecoderModelForCausalLM[ModelingV2Core, ...]): + def __init__(self, model_config): + ... # unchanged + self.spec_config = getattr(model_config, "spec_config", None) + self.draft_model = None + self.spec_worker = None + if self.spec_config is not None: + assert self.spec_config.spec_dec_mode.is_mtp_eagle_one_model() + self.draft_model = (...) + self.spec_worker = get_spec_worker(self.spec_config, model_config, + model_config.mapping) + + def forward(self, attn_metadata, **kw): + if self.spec_worker is None: + assert kw.get("spec_metadata") is None + return super().forward(attn_metadata, **kw) # bit-identical + spec_metadata = kw["spec_metadata"] + hidden = self.model(attn_metadata=attn_metadata, **kw) + logits = self.logits_processor.forward( + hidden[spec_metadata.gather_ids], self.lm_head, attn_metadata, True) + return self.spec_worker( + input_ids=kw["input_ids"], position_ids=kw["position_ids"], + hidden_states=hidden, logits=logits, + attn_metadata=attn_metadata, spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=kw.get("resource_manager")) +``` + +The class name is the synthetic one `routing.py` returns for this +(checkpoint, GPU architecture, parallel topology) triple — no checkpoint +declares it, and `test_modeling_v2_claims.py` asserts it agrees with the +target's path. + +`get_spec_worker` is imported from `tensorrt_llm._torch.speculative` — the +**runtime**, the part of the stack this project reuses, not modeling. + +Four things this shape is load-bearing about: + +**The non-speculative path must be a delegation, not a reimplementation.** +A target's release criterion was measured on the inherited base forward; the +only way to keep it bit-identical is to call it. Declaring +`resource_manager` as a named parameter would silently drop it from `**kw`, +so leave it in the dict and pull it with `.get` in the speculative branch +only — then the base receives exactly what it receives today. + +**The shell gathers the logits; the engine does not.** For every one-model +mode `without_logits` is True, so `_forward_step` returns the model's dict +verbatim and applies no second gather. `spec_metadata.gather_ids` holds one +row per context request (its last token) and `runtime_draft_len + 1` rows +per generation request — pass `hidden` **ungathered** to the worker and the +gathered logits alongside it. + +**`position_ids` reaches the worker in the engine's `[1, T]` shape.** The +worker does `position_ids.squeeze(0)` itself. A shell that flattens before +handing it over produces a silently wrong draft position sequence. + +**Nothing needs adding to `epilogue`, and nothing needs a `layer_idx`.** +`epilogue` is only consulted by `__pp_init__`'s `skip_forward`, so at +`pp_size == 1` there is nothing to register. And on this mode +`Eagle3OneModelSpecMetadata` sets `layers_to_capture = ()`, which makes +`is_layer_capture()` False at every layer and leaves `hidden_states` +unallocated — the trunk owes the runtime **no hidden-state capture hook** +(that is Eagle3's requirement, not MTP-Eagle's). Measured: nothing in +`_torch/pyexecutor/` or `_torch/speculative/` reads `model.layer_idx`. + +### Draft length is per iteration, not per request + +`_handle_dynamic_draft_len` runs **before** `prepare_resources`, so KV +allocation already knows the answer: + +1. `draft_len_schedule` maps a batch-size threshold to a draft length. +2. The current `scheduled_batch.batch_size` selects `runtime_draft_len`. +3. Every generation request's `py_draft_tokens` is padded or truncated to + **exactly** that length — the source comment names CUDA-graph replay and + the attention kernel as the reasons. +4. It lands on `spec_metadata.runtime_draft_len`; + `runtime_tokens_per_gen_step = 1 + runtime_draft_len`. +5. `runtime_draft_len == 0` takes `skip_drafting`, i.e. speculation is off + for that iteration only. + +**For modeling this means there is no ragged draft tree to handle.** The +draft length is a host int, constant across the batch, constant within a +capture. `cuda_graph_runner.get_graph_key` asserts it directly: *"All draft +lengths must be the same"*. + +Without a schedule, `runtime_draft_len` is simply `max_draft_len` every +step. + +**`draft_len_schedule` deadlocks under attention data parallelism, and +nothing rejects the combination.** Step 2 above reads +`scheduled_batch.batch_size` — **each rank's own local batch** — with no +cross-rank reduction, and attention DP does not equalize batch sizes: +`_pad_attention_dp_dummy_request` only tops a rank up from zero to one, and +`attention_dp_config.enable_balance` is off unless asked for. Two ranks +either side of a schedule threshold therefore run **different numbers of +draft replays**, hence different numbers of MoE collectives — and the +collectives pair by position, so the job hangs. + +Measured at world size 4: boot, CUDA-graph capture and warmup all pass, and +the hang needs real traffic to make the rank batch sizes diverge. trtllm's +own `HangDetector` fired at 300 s and hard-killed all four ranks, whose +stacks sat at three different points of one forward. **On an attention-DP +target, skip this knob.** With `tp`/`ep` alone the ranks share one batch and +the mechanism is sound. + +Setting it also **silently turns on `cuda_graph_config.enable_padding`** — +logged at INFO only — so it is never a single-variable experiment. + +### What changes in the trunk's own forward + +**A generation request arrives with `runtime_draft_len + 1` query tokens +instead of 1**, and the total-token arithmetic is +`num_ctx_tokens + (num_seqs - num_contexts) * tokens_per_gen_seq`. + +**How that reaches the attention op depends on the configuration, and the two +cases do not share a mechanism.** On an **MLA** layer it is +`predicted_tokens_per_seq`. In the **standard (non-MLA) configuration it is not**: +the op derives rows per generation sequence from `num_gen_tokens / num_gen_seqs` +and `predicted_tokens_per_seq` is inert — it reaches the op only through +`MLAParams`, which `create_attention` does not build for a non-MLA layer, so the +engine passes 1 while making the query block taller. Measured by sweeping it at +1, 2, 3, 4 and 8 against generation row counts 1–4: `output` and the pool came +back **bitwise identical** to the `P = 1` run at the same row count, including +where the two disagreed outright. + +Two consequences. The op **checks the division** — a ragged generation batch +raises `seq_len should be same for all generation requests`. And setting +`predicted_tokens_per_seq` is harmless but does nothing on this path, so a target +that sets it and sees a taller block has not learned that the two are connected. + +*Verified on sm_100. There is no version pin to re-check against: a target +declares the private engine surface it touches as `REQUIRED_TRTLLM_OPS` and +`tests/unittest/_torch/modeling_v2/test_modeling_v2_target_contract.py` asserts +every name in it exists, which turns a drifting engine into a loud failure +rather than a wrong answer. Re-check this particular behaviour by sweeping the +argument against a fixed row count and comparing output bitwise; the catalog +entry's contract is the authority for what the resulting block attends to.* + +**Those extra query tokens stay on the generation call.** A one-engine mode +returns False from the runtime's `extend_ctx` predicate — "1-model has +separate logic for handling draft tokens" — so a generation request carrying +drafts is *not* re-shaped into a chunked context request the way two-model +speculation does it. The batch keeps its `[context | generation]` split and +the generation call simply gets a taller query block. What that block is +allowed to attend to — each draft position seeing the cache plus the earlier +positions of its own block, and no later one — is a **kernel** fact, so the +catalog contract for the attention entry is its authority, not this file. +Do not assume it from `mask_type` alone. + +**The spec-dec mask machinery is forced off at sm_100 and stays inert.** + +```python +# _torch/attention_backend/trtllm.py +# Blackwell trtllm-gen spec-dec is enabled only for dynamic-tree masks. +self.is_spec_decoding_enabled = is_spec_decoding_enabled and ( + not self.is_sm_version_trtllm_gen_kernel(sm=get_sm_version()) or is_spec_dec_dynamic_tree +) +``` + +`is_sm_version_trtllm_gen_kernel(sm)` is `not (sm < 100 or sm in [120, 121])`, +so it is True on sm_100; a linear-tree MTP has `is_spec_dec_dynamic_tree` +False; the conjunction is **False**. `is_spec_decoding_enabled`, +`use_spec_decoding` and `is_spec_dec_tree` are all False and every +`spec_decoding_*` tensor is None — the same inert values the MLA columns of +`tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/thop_attention.md` +already certify. **On a pre-Blackwell arch this is not true** and the mask +surface would need certifying first. + +The context path is unchanged: a context request still contributes its +prompt, one row of logits, and one attention call of the same shape. + +### The draft loop rewrites the batch between step 0 and step 1+ + +The single most important fact for writing the layer. After the first draft +step the worker mutates `attn_metadata` in place: + +| | step 0 | step 1+ | +|---|---|---| +| `_seq_lens` / `_seq_lens_cuda` | real | **filled with 1** | +| `num_contexts` | real | **0** (when a KV cache manager is present) | +| `num_ctx_tokens` | real | **0** (recomputed) | +| `host_request_types` | mixed | context entries overwritten to generation | +| tokens per generation request | `runtime_draft_len + 1` | **1** | +| a context request | the whole prompt | **1 token** | +| `use_spec_decoding` | as the engine set it | False | +| `kv_lens_cuda` | as the engine set it | rewound, then `+1` per step | + +**So the layer must read its phase from `attn_metadata` on every call.** It +is invoked N times inside one forward, and a value computed on the first +call is wrong on the rest. This is not the trunk's situation — the trunk +builds step-args once per forward because there is only one step in it. + +The read is cheap and sync-free. `on_update()` recomputes `_num_ctx_tokens`, +`_num_generations` and `_num_tokens` from `_seq_lens`, which is a **pinned +host** tensor, and the loop calls it (and triggers it again through the +`num_contexts` setter) at exactly the step-0/step-1+ boundary. Therefore + +``` +tokens_per_gen_seq = (rows - md.num_ctx_tokens) // (md.num_seqs - md.num_contexts) +``` + +evaluates to `runtime_draft_len + 1` on step 0 and to exactly `1` on step +1+, with no step counter threaded through and no device read. That is the +**generation row count** the layer's own call needs — on a standard layer that +is what the taller query block *is*, not something you pass as +`predicted_tokens_per_seq` (see above). Guard the all-context case +(`num_seqs == num_contexts`) rather than dividing by zero. + +**What a context request is fed on step 0** is `prompt[1:]` with the +request's first accepted (golden) token written at its last position — +`_prepare_context_input_ids`, shared by both MTP flavours. Its hidden states +are the trunk's, in full, at full prompt length. So step 0 runs a real +context-phase attention for those rows and the layer needs both phases, +exactly like the trunk. + +### Attention DP: the padding basis changes source, and reading the wrong one is silent + +Under `enable_attention_dp`, §12 established that every rank pads its token +block to `max(attn_metadata.all_rank_num_tokens)` so both collectives run in +their uniform form. **Inside the draft loop that list is the wrong one from +step 1 onwards.** + +The worker passes the right basis in **as a keyword argument** and leaves +`attn_metadata.all_rank_num_tokens` holding the trunk's value for the whole +loop (it saves and restores it around the loop, but does not maintain it +during it): + +| draft step | `all_rank_num_tokens` kwarg | +|---|---| +| 0 | `spec_metadata.all_rank_num_tokens` — the trunk's token counts | +| 1+ | `spec_metadata.subseq_all_rank_num_tokens` — the per-rank **sequence** counts | + +`subseq_all_rank_num_tokens` is set to `all_rank_num_seqs` by the engine for +the one-model modes, which is semantically right: from step 1 every sequence +contributes exactly one token. + +**Pad from `md.all_rank_num_tokens` inside the MTP layer and the collectives +mispair silently.** +`tensorrt_llm/_torch/_experimental/modeling_v2/catalog/comm/allgather.md` and +`.../catalog/comm/reducescatter.md` both certify that calls pair by *position* +on the communicator and that at equal byte counts a divergence produces no +hang — every rank wrong in 98–99% of elements, bitwise reproducibly. The +defence is structural: give the layer's padding helper a signature that +**only accepts a passed-in list**, so it has no way to reach the metadata. + +(A `_dp_rows` that asserts `all_rank_num_tokens[rank] == rows` would in fact +fire here rather than corrupt, since the two counts differ from step 1. Do +not rely on it: the assert is a property of one target's helper, not of the +rule, and it is exactly the kind of guard a later edit removes.) + +### KV cache: the engine adds the layer, and the target declares nothing + +Under a one-model MTP mode `ModelConfig` raises the pool's layer count +itself: + +```python +num_layers += spec_config.num_nextn_predict_layers +num_attention_layers += spec_config.num_nextn_predict_layers +``` + +so the pool is sized for `num_hidden_layers + n` layers and the MTP layer's +own attention addresses layer index `num_hidden_layers`. **Nothing in the +target's config or manifest declares any of this** — but the layer's +attention calls must pass the right `layer_idx`, and the pool-addressing +surface they use has to be certified at the new layer count. + +**A `speculative_config` also raises `max_seq_len`, and not by the amount the +extra-KV-token helper suggests.** Three separate terms are added to the model +engine's `max_seq_len`, in `py_executor_creator`: + +```python +if not disable_overlap_scheduler and spec_config is not None: + max_seq_len += spec_config.tokens_per_gen_step - 1 +if spec_config is not None: + max_seq_len += get_num_extra_kv_tokens(spec_config) # max_draft_len - 1 + max_seq_len += spec_config.tokens_per_gen_step - 1 +``` + +With a linear tree (`tokens_per_gen_step = 1 + max_draft_len`) and the overlap +scheduler at its default (**on**), that is `3 * max_draft_len - 1` — 2, 5 and +8 at `max_draft_len` 1, 2 and 3. Measured: a `163840` model cap becomes +**`163848`** at `max_draft_len: 3`. + +**So read `max_seq_len` off the metadata, never compute it.** A target that +derives anything from the config's own `max_position_embeddings` — a rope +table sized to it, a bound asserted against it — is off by that amount the +moment speculation is switched on, and by a different amount per +`max_draft_len`. The failure is a silent out-of-bounds read on a rope table, +or an assert that fires on a legal config. + +### One more construction-time trap, from §4's weight lifecycle + +**A reference to another module's parameter, captured in `__init__`, stays +bound to the meta-device shadow.** The engine materializes at t1 by +*replacing* tensor objects, not by filling them in place, so a drafter +container that does `self.embed_tokens = core.w["embed"]` at construction +holds a meta tensor forever and fails at the first draft step. §4 already says +an unregistered parameter "stays a dead shadow"; this is the adjacent case — +the parameter *is* registered, on the trunk, and the copy of the reference is +what goes stale. Resolve it lazily (a `property` that reads through to the +trunk on each access) rather than caching it. + +### CUDA graphs: the whole draft loop is inside the capture + +Capture wraps `_forward_step`, which calls the shell's forward, which calls +the worker — so every MTP-layer invocation is captured. The key is +`(batch_size, draft_len, is_first_draft, short_seq_len_mode, +is_all_greedy_sample)`, and the capture set becomes `(bs, draft_len)` pairs +rather than bare batch sizes; with a `draft_len_schedule` the runner also +captures one extra `(max_bs, original_max_draft_len)` graph, whose stated +purpose is to keep a later graph from resizing the shared attention +workspace and invalidating pointers baked into earlier ones. + +Consequences for the layer, all of them §7's discipline applied one level +down: + +* The step-0 / step-1+ divergence is **re-traced per draft step at capture + time** and frozen at each step's position in the loop. Reading the phase + from `attn_metadata` every call is what makes that correct — a cached + first-call decision would be baked in at all N positions. +* Host ints read from `attn_metadata` (`num_contexts`, `num_ctx_tokens`) and + the `all_rank_num_tokens` kwarg are per-capture constants, which is the + host-derived class §7 permits. +* `attn_metadata.padded_num_tokens` is **`None`** unless torch-compile + piecewise CUDA graphs are configured; without a `torch_compile_config` the + padding path is not reachable and the base shell's slice never applies. +* The attention-DP replay gate of §12 is unchanged: all ranks + generation-only and equal batch sizes, or the step runs eager — and eager + is where the layer meets a real mixed batch. + +### The acceptance rate is the only correctness signal + +`stats.specdec_stats.acceptance_length` is computed per iteration over the +generation requests that carried draft tokens: + +``` +acceptance_length = (accepted_draft_tokens + requests_with_draft) / requests_with_draft +``` + +i.e. the mean number of tokens a drafting request produces per step, `1.0` +meaning every draft was rejected. + +**This matters more than it looks.** Rejection sampling guarantees the +output distribution is unchanged, so a *miscomputed* MTP layer does not +produce wrong text — it produces correct text more slowly, with every draft +rejected. An accuracy gate cannot see it. A subtly wrong layer (a +transposed concatenation, a norm applied to the wrong operand) lands +somewhere above 1.0 and well below the reference, which no other measurement +distinguishes from "this model is just hard to draft for". The reference is +the same checkpoint under stock in-tree modeling at the same load and the +same `speculative_config` — the same job with `TRTLLM_MODELING_V2=off` +instead of `=require`; acceptance is a property of the model, so a +reference forced onto a different `max_num_tokens` / `max_seq_len` to boot +is still a valid comparison. + +### Still not established here + +Vanilla MTP (`n > 1`, one layer per draft position) and its per-layer +sequential loop; Eagle3 in either form, including the hidden-state capture +hook and `apply_eagle3_fc`; tree drafting of any kind (static or dynamic +`eagle_choices`), which is also the only way the Blackwell spec-dec mask +surface becomes reachable; two-model speculation; and speculative decoding +combined with pipeline parallelism, where `skip_forward` and the `epilogue` +path start to matter. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/spool_watcher.sh b/agent-flow/agent_flow/workflows/staircase_onboard/spool_watcher.sh new file mode 100644 index 000000000000..20f027448079 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/spool_watcher.sh @@ -0,0 +1,144 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# The compute half of `environment.kind: session_spool`. +# +# WHY IT EXISTS, IN ONE PARAGRAPH. On a cluster whose Slurm CLI cannot reach +# the controller, jobs go through an MCP broker; the broker's blocking calls +# are aborted after a few minutes of silence and every useful job is longer +# than that, while `slurm_exec_in_job` runs work in a job *step* whose cgroup +# Slurm tears down when the foreground command exits -- so a `nohup`'d child +# does not survive it (measured on this cluster: the log stopped at the first +# banner and never grew). The way out is for the long-running process to be +# the job's OWN main process. So the job runs this loop, and work arrives as +# files on a shared filesystem. A plain Python process with no MCP client -- +# the orchestrator -- can then drop a script in `pending/` and poll for the +# matching `.rc`, which hands gate execution back to where the design wanted +# it. It also imports the container ONCE for the whole session instead of +# once per command. +# +# PROTOCOL, which is `agent_flow.runtime.commands.spool`'s and not this +# script's: +# +# /pending/.sh appears (renamed into place, never partial) +# /logs/.log stdout+stderr, grows while running +# /done/.rc the exit code, renamed into place LAST +# +# `.rc` last is the whole ordering contract: a reader that sees it may trust +# the log beside it is complete. +# +# CONCURRENCY, AND THE ONE RULE THAT MATTERS. Work is dispatched concurrently, +# because the orchestrator itself may be one of the jobs -- on a host where a +# login-node background process does not survive, the session is the only +# place with a lifetime measured in hours, and a serial watcher would then +# deadlock: the orchestrator would hold the only slot while blocking on a gate +# that can never start. +# +# But concurrent *GPU* work on one allocation is a different mistake. A gate +# at world size 4 owns all four devices, and two of those at once fail with +# "Executor creation failed due to insufficient GPU memory" on a 284 GB card +# -- which reads as a target defect and is not. So: +# +# a script whose name begins with `free-` runs immediately and holds nothing +# every other script serializes on an exclusive lock for its whole run +# +# The orchestrator is submitted as `free-...`. Gates are named `gate-` +# by the executor and therefore serialize, which is what they want. +# +# Usage, as the container command of a 1-node / N-GPU allocation: +# +# bash /path/to/spool_watcher.sh /path/to/spool + +set -u + +SPOOL="${1:?usage: spool_watcher.sh [poll-seconds]}" +POLL="${2:-3}" +LOCK="$SPOOL/.gpu.lock" + +mkdir -p "$SPOOL"/{pending,running,done,logs} +: > "$LOCK" + +stamp() { date -u '+%Y-%m-%dT%H:%M:%SZ'; } + +{ + echo "host=$(hostname)" + echo "pid=$$" + echo "started=$(stamp)" + echo "spool=$SPOOL" + echo "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-}" + nvidia-smi --query-gpu=index,name,compute_cap,memory.total --format=csv 2>&1 +} > "$SPOOL/watcher.status" + +run_one() { + local script=$1 name=$2 + { + echo "=== start $(stamp) on $(hostname) ===" + echo + } > "$SPOOL/logs/$name.log" + + if [[ "$name" == free-* ]]; then + bash "$script" >> "$SPOOL/logs/$name.log" 2>&1 + else + # 9 is an arbitrary free fd; `flock` holds it for the subshell's lifetime + # and releases it however the script exits, including a kill. + ( flock 9; bash "$script" ) 9>"$LOCK" >> "$SPOOL/logs/$name.log" 2>&1 + fi + local rc=$? + + echo "=== end rc=$rc $(stamp) ===" >> "$SPOOL/logs/$name.log" + mv "$SPOOL/running/$name.sh" "$SPOOL/done/$name.sh" 2>/dev/null || true + echo "$rc" > "$SPOOL/done/$name.rc.tmp" + mv "$SPOOL/done/$name.rc.tmp" "$SPOOL/done/$name.rc" +} + +shopt -s nullglob +echo "watcher: ready at $(stamp)" >> "$SPOOL/watcher.status" + +while true; do + for f in "$SPOOL"/pending/*.sh; do + name=$(basename "$f" .sh) + # Claim by moving. If two watchers ever race the same spool, only one + # `mv` wins and the work runs once. + if mv "$f" "$SPOOL/running/$name.sh" 2>/dev/null; then + run_one "$SPOOL/running/$name.sh" "$name" & + fi + done + + # Written every poll, so a stale mtime means a dead watcher rather than a + # slow one. Two spellings because two readers exist: the ISO line a human + # and the shell helpers read, and the file mtime the Python side can stat + # without parsing anything. + echo "alive=$(stamp)" > "$SPOOL/watcher.heartbeat.tmp" + mv "$SPOOL/watcher.heartbeat.tmp" "$SPOOL/watcher.heartbeat" + + # A SEPARATE file from the heartbeat, deliberately: `scripts/spool_submit.sh` + # pipes the heartbeat straight into `date -d` after stripping `alive=`, so + # anything else on that line breaks every submitter. This one holds the job + # id and nothing else. + # + # What reads it is `environment.shares_allocation_with_spool`, which asks + # whether the orchestrator is living inside the very allocation it dispatches + # work to. If it is, the two die together -- three times in one run on the + # predecessor project, and twice more here: once when a disagg probe's worker + # teardown took the step down, and once when an agent's `pkill -9` caught + # this watcher, which IS the allocation's main process. Each time the + # orchestrator went with it, holding GPUs it never used. + # + # `unknown` outside Slurm rather than an empty file, because the reader + # treats unknown as no evidence and an empty file as a partial write. + echo "${SLURM_JOB_ID:-unknown}" > "$SPOOL/watcher.jobid.tmp" + mv "$SPOOL/watcher.jobid.tmp" "$SPOOL/watcher.jobid" + + # `STOP` is the clean shutdown, checked after the drain so anything already + # claimed finishes rather than being orphaned mid-gate. + if [ -f "$SPOOL/STOP" ]; then + echo "watcher: STOP seen at $(stamp), waiting for running work" >> "$SPOOL/watcher.status" + wait + echo "watcher: exiting at $(stamp)" >> "$SPOOL/watcher.status" + rm -f "$SPOOL/watcher.heartbeat" "$SPOOL/watcher.jobid" + exit 0 + fi + + sleep "$POLL" +done diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/state.py b/agent-flow/agent_flow/workflows/staircase_onboard/state.py new file mode 100644 index 000000000000..c2ffe653a292 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/state.py @@ -0,0 +1,324 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Checkpoint state for the staircase-onboard workflow. + +Persisted to ``/.staircase_onboard_state.json`` so a run +interrupted by Ctrl-C / crash / reboot can be continued by re-running the +workflow against the same workspace (wipe with ``--clean`` to start over). + +Staircase's own doctrine is that there is no resume *mode* — re-running the +same invocation is how a run resumes, because every step decides for itself +whether its work is already done. That still holds for the work products +(a catalog entry that exists is not rebuilt; a gate that passed is not +re-run). The checkpoint here is a different thing: it records **where the +orchestrator was**, so a resumed run does not re-pay the agent turns that +already produced those products. The two are complementary — the checkpoint +saves tokens, the idempotent steps save correctness. + +Stages, in order: + +- ``preflight`` — pure Python. Checkpoint parses, arch matches the host, + parallel segment is well formed, the device set is + claimed, TMPDIR is sane, multi-rank probes pass. No + agent runs here. +- ``plan`` — PlanDrafter <-> PlanReviewer, capped at + ``max_plan_rounds``. Products: ``plan.md`` and + ``vocabulary.yaml`` (the typed dispatch queue). +- ``smith`` — one smith per queue item, dispatched concurrently + within a dependency batch. Products: catalog entries + and their GPU tests, committed per entry. +- ``assemble`` — one assembler per unit, serially (core first, then one + increment per feature), each followed by the gates. +- ``deliver`` — Python verification across the in-tree trees, one + stateless auditor turn, one lint pass, then the merge. + +``stage`` always names the phase currently in progress (or pending); on +resume the workflow jumps directly to it instead of rerunning earlier +phases. +""" + +from __future__ import annotations + +import json +import os +import tempfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +STATE_FILENAME = ".staircase_onboard_state.json" +SCHEMA_VERSION = 1 + +STAGE_PREFLIGHT = "preflight" +STAGE_PLAN = "plan" +STAGE_SMITH = "smith" +STAGE_ASSEMBLE = "assemble" +STAGE_DELIVER = "deliver" + +_VALID_STAGES = ( + STAGE_PREFLIGHT, + STAGE_PLAN, + STAGE_SMITH, + STAGE_ASSEMBLE, + STAGE_DELIVER, +) + +# Queue-item lifecycle. ``running`` is persisted rather than kept in +# memory so a crash mid-batch is distinguishable from a crash before it: +# a resumed run re-dispatches ``running`` items (the agent's work may have +# been lost) but leaves ``done`` ones alone. +# +# ``deferred`` is a failure only a later unit pays for: the row's entry is +# needed by a feature increment and by nothing core uses, so the smith phase +# records it and carries on instead of closing the run. The row also carries +# ``deferred_to`` (the unit that retries it before its first assembler turn) +# and ``deferred_why``. Persisted, not recomputed, because the retry is owed +# by a later process as often as by this one, and a resume that forgot a +# deferral would assemble that unit without the entry. +ITEM_PENDING = "pending" +ITEM_RUNNING = "running" +ITEM_DONE = "done" +ITEM_FAILED = "failed" +ITEM_DEFERRED = "deferred" +_VALID_ITEM_STATUSES = (ITEM_PENDING, ITEM_RUNNING, ITEM_DONE, ITEM_FAILED, ITEM_DEFERRED) + + +#: Verdicts a re-run should re-enter rather than refuse. +#: +#: The distinction is not "how bad was it" but **who repairs it and where**. +#: Each of these closes a run whose agent work is finished and whose failure +#: names a concrete edit the operator makes in the repo or in ``task.yaml`` -- +#: a missing CI test file, an entry needing certification, a witness block +#: that was never configured. The run cannot observe that repair except by +#: running its own check again, and every one of those checks is cheap and +#: side-effect-free. +#: +#: Everything else stays refused, and deliberately. ``ceiling``, +#: ``plan-failed``, ``smith-blocked``, ``assembly-exhausted`` and +#: ``assembler-blocked`` are budgets spent or judgments made by agents; a +#: re-run re-dispatches paid turns with nothing having changed, which is how a +#: workflow burns a night discovering the same thing four times. Those need +#: ``--clean`` and a human deciding what to do differently. The one narrow +#: exception is a smith-phase ``smith-blocked`` whose every failed row is now +#: tagged for a later unit: re-entering it dispatches none of them, so it is +#: not a re-run of the same thing (``workflow._smith_block_is_deferrable``). +#: +#: Prefix match, because ``witness-unconfigured`` carries its unit. +REMEDIABLE_VERDICTS: tuple[str, ...] = ( + "delivery-failed", + "audit-uncovered", + "witness-unconfigured", + # The three ways the bounded remediation cycle gives up. All of them + # close a run whose remaining work is a human judgement made in the repo + # -- a receipt that did not land, a coverage evaluator disagreeing with a + # contract, a budget that was set too low for this target. The queue is + # in the checkpoint, so re-running resumes the cycle rather than + # restarting it, which is the whole reason it is persisted. + "audit-no-progress", + "audit-remediation-budget", + "audit-verification-failed", + # An assembler that stops on `knowledge_gap` or `environment` names work + # only the operator can supply (a verify round, a ruling, an + # infrastructure repair); once it lands, re-entering resumes the unit + # instead of `--clean` re-measuring every gate. Measured on + # DeepSeek-V4.1: the core assembler's envelope sweep stopped the run on + # four catalog verifies. `vocabulary_gap` stays out: the run onboards + # those itself, so re-entering on one would repeat the same turn. + "assembler-blocked:knowledge_gap", + "assembler-blocked:environment", +) + + +def is_remediable(verdict: str) -> bool: + """Whether a closed run's verdict describes a repair made outside the run.""" + return any(verdict == known or verdict.startswith(known + ":") for known in REMEDIABLE_VERDICTS) + + +@dataclass +class WorkflowState: + """One run's resume point. + + Everything here is orchestrator bookkeeping. The run's *evidence* — + what each agent did and what each gate measured — lives in + ``journal.md`` and ``progress.yaml`` beside it, because those are + append-only and this is not. + """ + + task_path: str + stage: str = STAGE_PREFLIGHT + done: bool = False + + # Why the run closed, as ``_close_run`` recorded it. Persisted for one + # reason: some verdicts name a repair the operator makes *outside* the + # run, and re-running is the only way to confirm it. A closed run that + # cannot say why it closed has to be started over from scratch, which + # throws away every gate it had already passed. See ``REMEDIABLE_VERDICTS``. + verdict: str = "" + + # Identity, resolved by preflight and stable for the run's life. + slug: str = "" + branch: str = "" + devices: list[int] = field(default_factory=list) + + # Plan phase. ``plan_round`` counts *completed* drafter->reviewer + # rounds, so it is also the number of re-drafts already spent. + plan_round: int = 0 + plan_approved: bool = False + + # Smith phase. One entry per dispatch-queue item, carrying the queue + # row plus its ``status``. Stored in full rather than as a pointer + # into vocabulary.yaml so a queue item's outcome survives an edit to + # that file. A ``deferred`` row is still owed by the assembly phase: + # see ``ITEM_DEFERRED``. + queue: list[dict[str, Any]] = field(default_factory=list) + + # Assembly phase. ``unit_index`` 0 is the core assembly; 1..N are the + # feature increments in the order task.yaml lists them. + # + # ``unit_turns`` counts assembler **dispatches**, not entries into the + # loop: a restart that finds the unit's work already on disk hands the + # turn back rather than spending it. See + # ``workflow.StaircaseOnboardWorkflow._assemble_one_unit``. + unit_index: int = 0 + unit_turns: int = 0 + units_done: list[str] = field(default_factory=list) + + # Per unit, the identity of the work its last ``UNIT_READY`` was written + # about -- see ``workflow.StaircaseOnboardWorkflow._target_marker``. + # Persisted because it only answers a question a *different process* + # asks: "is what is on disk still what the previous process declared + # ready?" A resume that cannot answer it re-dispatches, which is the + # safe direction. + unit_marker: dict[str, str] = field(default_factory=dict) + + # Delivery phase. + audit_clean: bool = False + + # The UNCOVERED remediation cycle: the verify queue, which round it is + # on, and the signature of the still-open set. Persisted because the + # convergence rules are *across* rounds -- "did this shrink?" cannot be + # answered by a process that forgot the previous answer, and a resume + # that forgets it re-dispatches work it already did. See + # ``remediation.RemediationState``. + remediation: dict[str, Any] = field(default_factory=dict) + + # The metric staircase exists to drive down: how many times the run + # stopped for a human decision. Carried in the checkpoint so it + # survives a resume and can be reported honestly at the end. + human_stops: int = 0 + + +def _coerce_queue(raw: Any, path: Path) -> list[dict[str, Any]]: + """Validate the persisted queue, rejecting anything the workflow cannot walk. + + A malformed queue is worth failing on rather than repairing: the + orchestrator dispatches from it, so a silently-dropped row is a + catalog entry that never gets built and a target that fails its + vocabulary audit much later, with nothing pointing back here. + """ + if raw is None: + return [] + if not isinstance(raw, list): + raise ValueError(f"{path}: `queue` must be a list, got {type(raw).__name__}") + queue: list[dict[str, Any]] = [] + for index, row in enumerate(raw): + if not isinstance(row, dict): + raise ValueError(f"{path}: queue[{index}] must be a mapping, got {type(row).__name__}") + status = row.get("status", ITEM_PENDING) + if status not in _VALID_ITEM_STATUSES: + raise ValueError( + f"{path}: queue[{index}] has status {status!r}; expected one of " + f"{_VALID_ITEM_STATUSES}. Delete the file to start fresh." + ) + queue.append(dict(row)) + return queue + + +def load_state(path: Path) -> WorkflowState: + data = json.loads(path.read_text(encoding="utf-8")) + version = data.get("version") + if version != SCHEMA_VERSION: + raise ValueError( + f"Unsupported checkpoint version {version!r} in {path}; " + f"expected {SCHEMA_VERSION}. Delete the file to start fresh." + ) + stage = data.get("stage") + if stage not in _VALID_STAGES: + raise ValueError( + f"Unsupported stage {stage!r} in {path}; expected one of " + f"{_VALID_STAGES}. Delete the file to start fresh." + ) + return WorkflowState( + task_path=str(data["task_path"]), + stage=stage, + done=bool(data.get("done", False)), + verdict=str(data.get("verdict", "")), + slug=str(data.get("slug", "")), + branch=str(data.get("branch", "")), + devices=[int(d) for d in data.get("devices", [])], + plan_round=int(data.get("plan_round", 0)), + plan_approved=bool(data.get("plan_approved", False)), + queue=_coerce_queue(data.get("queue"), path), + unit_index=int(data.get("unit_index", 0)), + unit_turns=int(data.get("unit_turns", 0)), + # Absent in checkpoints written before the resume check existed. An + # empty map reads as "no unit's work has a recorded identity", which + # makes every such resume re-dispatch exactly as it used to. + unit_marker={str(k): str(v) for k, v in (data.get("unit_marker") or {}).items()}, + units_done=[str(u) for u in data.get("units_done", [])], + audit_clean=bool(data.get("audit_clean", False)), + human_stops=int(data.get("human_stops", 0)), + # Absent in checkpoints written before the remediation cycle existed. + # An empty queue is the correct reading of those: no round has run. + remediation=dict(data.get("remediation") or {}), + ) + + +def save_state(path: Path, state: WorkflowState) -> None: + """Write the checkpoint atomically. + + Same fsync + ``os.replace`` dance as the other workflows: the + orchestrator checkpoints between every stage, so a torn write here + would strand a run that was otherwise fine. + """ + payload = {"version": SCHEMA_VERSION, **state.__dict__} + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp_path = tempfile.mkstemp( + prefix=STATE_FILENAME + ".", + suffix=".tmp", + dir=path.parent, + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as fh: + json.dump(payload, fh, indent=2) + fh.flush() + os.fsync(fh.fileno()) + os.replace(tmp_path, path) + except Exception: + try: + os.unlink(tmp_path) + except FileNotFoundError: + pass + raise + + +__all__ = [ + "REMEDIABLE_VERDICTS", + "is_remediable", + "ITEM_DEFERRED", + "ITEM_DONE", + "ITEM_FAILED", + "ITEM_PENDING", + "ITEM_RUNNING", + "SCHEMA_VERSION", + "STAGE_ASSEMBLE", + "STAGE_DELIVER", + "STAGE_PLAN", + "STAGE_PREFLIGHT", + "STAGE_SMITH", + "STATE_FILENAME", + "WorkflowState", + "load_state", + "save_state", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/survey.py b/agent-flow/agent_flow/workflows/staircase_onboard/survey.py new file mode 100644 index 000000000000..689b52504ea2 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/survey.py @@ -0,0 +1,1311 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""What this checkpoint demands, and what this build can already supply. + +The plan phase used to open with the drafter *guessing*: guessing which +``torch.ops.trtllm.*`` names exist in this build, guessing which catalog +entries are certified on this architecture, and guessing which of the +checkpoint's structural features have a home. Every guess that missed became +a queue item that a smith discovered was wrong, hours later, with a GPU +allocation held open behind it. + +This module replaces the guessing with a measurement taken before the first +agent exists. It has two halves and they are deliberately not joined: + +- :func:`read_checkpoint_surface` reads ``config.json`` and reports what the + model *structurally demands* — attention geometry, positional scheme, + activation, normalization, expert routing, recurrent state, speculation + heads, quantization. Pure Python, no GPU, no imports of the engine. +- :func:`run_inventory_probe` runs on the real host and reports what the + build *supplies* — which ops are registered, which catalog entries exist, + which of those carry a receipt for this architecture, and which entries + name an ``impl`` that no longer resolves. + +**The join between them is judgement and stays in prose.** Deciding that +"this model wants SwiGLU" is answered by ``activation/flashinfer_silu_and_mul`` +is a modeling call, not a string match, and a Python function that guesses it +would be confidently wrong on exactly the models this workflow exists to +onboard. So the survey lays both columns in front of the drafter and lets the +drafter join them. + +The half that generalizes is the one that matters for a model nobody has +seen. :class:`CheckpointSurface` carries ``unrecognized``: structural keys in +``config.json`` that no probe here claims. A key nobody recognizes is the +earliest honest signal that this checkpoint has a layer the vocabulary has no +word for — which is the case that used to be discovered by a smith, at tier +4, with no budget left to absorb it. +""" + +from __future__ import annotations + +import json +import re +from collections import Counter +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Mapping, Sequence + +from .environment import RunEnvironment + +#: The inventory probe imports the engine, so it pays a real import cost, but +#: it runs no kernel and allocates nothing. Anything slower than this is not +#: the probe. +INVENTORY_TIMEOUT_S = 600 + +#: What the probe prints its JSON payload behind, so the parser can find it +#: in a log that also carries the engine's import-time chatter. +MARKER = "__SURVEY_JSON__" + + +# -------------------------------------------------------------------------- +# half one: what the checkpoint demands +# -------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class Demand: + """One structural property of the checkpoint, and what it obliges. + + ``keys`` is the evidence — the ``config.json`` keys this was read off — + so a drafter that disbelieves a demand can check it rather than argue + with it. + """ + + name: str + detail: str + keys: tuple[str, ...] = () + + def __str__(self) -> str: + where = f" [{', '.join(self.keys)}]" if self.keys else "" + return f"{self.name}: {self.detail}{where}" + + +@dataclass(frozen=True) +class CheckpointSurface: + """Everything the config declares that a target has to answer for.""" + + architecture: str + model_type: str + demands: tuple[Demand, ...] = () + unrecognized: tuple[str, ...] = () + submodels: tuple[str, ...] = () + + @property + def demand_names(self) -> tuple[str, ...]: + return tuple(d.name for d in self.demands) + + +def _num(config: Mapping[str, Any], *names: str) -> Any: + for name in names: + value = config.get(name) + if value is not None: + return value + return None + + +def _present(config: Mapping[str, Any], *names: str) -> tuple[str, ...]: + return tuple(n for n in names if config.get(n) is not None) + + +# Keys that carry no structural obligation: sizes, names, ids, bookkeeping. +# Anything outside this set and outside a probe's claim is reported as +# unrecognized, which is the point — the list is allowed to be incomplete, +# and its incompleteness costs a line of prose rather than a wrong plan. +_INERT_KEYS = frozenset( + { + "architectures", + "attention_bias", + "attention_dropout", + "auto_map", + "bos_token_id", + "dtype", + "eos_token_id", + "hidden_size", + "id2label", + "initializer_range", + "intermediate_size", + "label2id", + "max_position_embeddings", + "model_type", + "num_attention_heads", + "num_hidden_layers", + "num_key_value_heads", + "pad_token_id", + "tie_word_embeddings", + "torch_dtype", + "transformers_version", + "unsloth_version", + "use_cache", + "use_return_dict", + "vocab_size", + } +) + +# Prefixes whose keys a probe below has already accounted for. Listed as +# prefixes rather than exact names because config authors spell the same +# concept a dozen ways (`n_routed_experts`, `num_local_experts`, +# `moe_num_experts`), and an exact list would report a rename as a discovery. +_CLAIMED_PREFIXES = ( + "attn_", + "conv_", + "decoder_sparse_", + "draft", + "dspark_", + "expert", + "final_logit", + "full_attention_", + "full_attn", + "head_dim", + "local_", + "hidden_act", + "kv_lora", + "layer_norm", + "layer_types", + "linear_", + "mamba", + "mlp_only", + "moe", + "mtp_num", + "n_group", + "n_routed", + "n_shared", + "norm_topk", + "num_experts", + "num_local_experts", + "num_nextn", + "partial_rotary", + "q_lora", + "quantization", + "rms_norm", + "rope", + "routed_scaling", + "scoring_func", + "shared_expert", + "sink", + "sliding_window", + "ssm_", + "swa_", + "topk_", + "use_mrope", + "use_qk_norm", + "use_sliding_window", + "v_head_dim", +) + +# Keys a probe claims by PATTERN rather than by prefix, because the concept +# sits in the MIDDLE of the name. `add_swa_attention_sink_bias` has no prefix +# any list could carry without also claiming every `add_*` key, and the +# `attention.sink` probe that claims it globs on exactly this substring; the +# same is true of `_layer_pattern` and `attention.hybrid`. +_CLAIMED_SUBSTRINGS = ("attention_sink", "layer_pattern") + +_SUBMODEL_KEYS = ("vision_config", "audio_config", "text_config", "speech_config") + +# Per-layer attention pattern keys: a list with one entry per decoder layer +# saying which KIND of attention that layer runs. Named rather than matched by +# prefix because `moe_layer_freq` has the same shape and asks a different +# question; any `*_layer_pattern` spelling is accepted on top of these. +_HYBRID_PATTERN_KEYS = ("hybrid_layer_pattern", "layer_types", "full_attention_idx") + +# The geometry a hybrid stack may redeclare per layer kind. When one of these +# exists under a `swa_` / `local_` prefix and disagrees with its bare +# counterpart, the two kinds are different SHAPES rather than one shape under +# two masks — which decides whether the cache manager is asked for one pool or +# two, and whether the weight manifest has one shard calculation or two. +_HYBRID_GEOMETRY_KEYS = ( + "num_attention_heads", + "num_key_value_heads", + "head_dim", + "v_head_dim", + "rope_theta", +) +_HYBRID_KIND_PREFIXES = ("swa_", "local_") + + +def _repeating_period(values: Sequence[Any]) -> int | None: + """The smallest ``p`` with ``values[i] == values[i % p]`` for every ``i``. + + ``None`` when there is none. Reported rather than assumed: a stack whose + global layers land every sixth position and a stack whose first block is + short are different stacks, and the second one is what this was written + against. + + Candidates stop at half the length, because a "period" that occurs once + is not a period — a 48-entry list trivially satisfies ``p = 47`` whenever + its last entry happens to equal its first, and reporting that as a + repeat is worse than reporting none. + """ + for period in range(1, len(values) // 2 + 1): + if all(value == values[index % period] for index, value in enumerate(values)): + return period + return None + + +def _sink_layer_kind(key: str) -> str: + """``add_swa_attention_sink_bias`` -> ``swa``: which layers carry the sink. + + The layer kind is in the key name and nowhere else, and it is the whole + content of the demand — a sink on every layer and a sink on the local + layers only are different models. + """ + head = key.split("attention_sink", 1)[0] + if head.startswith("add_"): + head = head[len("add_") :] + # An unprefixed spelling names no kind, which is itself the answer: the + # sink is on the whole stack. + return head.strip("_") or "all" + + +def read_checkpoint_surface(config: Mapping[str, Any]) -> CheckpointSurface: + """Report what this checkpoint structurally demands of a target. + + Written against *concepts* rather than against one model family: each + probe accepts the several spellings the ecosystem actually uses, because + the alternative is a reader that silently reports "dense" for a MoE model + whose config says ``n_routed_experts``. + """ + # A multimodal checkpoint nests the language model. This workflow's scope + # block usually names the text tower, and reading the outer config for + # attention geometry would describe the vision encoder instead. + submodels = tuple(k for k in _SUBMODEL_KEYS if isinstance(config.get(k), dict)) + inner = config.get("text_config") + scoped: Mapping[str, Any] = inner if isinstance(inner, dict) else config + + demands: list[Demand] = [] + + heads = _num(scoped, "num_attention_heads", "n_head") + kv_heads = _num(scoped, "num_key_value_heads", "num_kv_heads") + if heads and kv_heads: + shape = "MHA" if kv_heads == heads else ("MQA" if kv_heads == 1 else "GQA") + demands.append( + Demand( + "attention.shape", + f"{shape}, {heads} query heads over {kv_heads} kv head(s)", + ("num_attention_heads", "num_key_value_heads"), + ) + ) + if scoped.get("head_dim"): + head_dim_detail = f"{scoped['head_dim']}, declared not derived" + head_dim_keys: tuple[str, ...] = ("head_dim",) + # Asymmetric qk / v head dimensions, stated as its own clause because + # it is a different question from "what is the head dim". Every other + # demand here is answered by choosing an entry; this one decides + # whether any paged-attention entry can be reused at all, or whether + # the target needs a padded-V fallback — so a drafter that reads only + # `head_dim` sizes the queue without the hardest item in it. + v_head_dim = scoped.get("v_head_dim") + if v_head_dim is not None and v_head_dim != scoped["head_dim"]: + head_dim_detail += ( + f"; v_head_dim={v_head_dim} differs from head_dim={scoped['head_dim']} " + f"— asymmetric qk/v head dimensions, which not every attention " + f"kernel admits" + ) + head_dim_keys = ("head_dim", "v_head_dim") + demands.append(Demand("attention.head_dim", head_dim_detail, head_dim_keys)) + if keys := _present(scoped, "kv_lora_rank", "q_lora_rank"): + demands.append( + Demand( + "attention.mla", + "compressed latent KV — a different cache layout and a different " + "op family from ordinary paged attention", + keys, + ) + ) + # Deliberately *not* keyed on `layer_types`: that key declares which kind + # of layer sits where, which is a hybrid stack and not a window. Reading + # it here reported "per-layer or global window" for a checkpoint with no + # window anywhere in it. `stack.hybrid` below is where it belongs. + if keys := _present(scoped, "sliding_window", "use_sliding_window", "sliding_window_pattern"): + demands.append( + Demand( + "attention.windowing", + ", ".join(f"{k}={scoped[k]!r}" for k in keys)[:200], + keys, + ) + ) + # A per-layer pattern is not a window. `attention.windowing` above reports + # `sliding_window=128` and nothing else, which reads as one window over a + # uniform stack; the stack this was written against is 48 layers of two + # different kinds whose KV head counts and RoPE bases disagree. That is a + # demand on the cache manager and on the weight manifest, not a mask. + # + # **The survey reports the pattern and refuses to interpret it.** Which + # value means sliding-window and which means full attention is the config + # author's private convention — this checkpoint spells local as `1` and + # full as `0`, and nothing in the file says so. Asserting a polarity the + # survey cannot know would silently window the wrong 39 layers, which is + # strictly worse than handing the plan phase the counts and the question. + pattern_keys = tuple( + key + for key in scoped + if isinstance(key, str) + and (key in _HYBRID_PATTERN_KEYS or key.endswith("_layer_pattern")) + and isinstance(scoped[key], Sequence) + and not isinstance(scoped[key], str) + and len(scoped[key]) > 1 + ) + if pattern_keys: + pattern_key = pattern_keys[0] + values = list(scoped[pattern_key]) + if pattern_key.endswith("_idx"): + # An INDEX list, not a per-layer list: it names the layers of one + # kind and leaves the rest implied. + layers = _num(scoped, "num_hidden_layers") + hybrid_detail = ( + f"{layers} layers, {len(values)} of them named by `{pattern_key}` " + f"and {layers - len(values)} not" + if isinstance(layers, int) and layers >= len(values) + else f"{len(values)} layer(s) named by `{pattern_key}`" + ) + else: + counts = Counter(str(value) for value in values) + tally = ", ".join(f"{n} of `{value}`" for value, n in counts.most_common()) + period = _repeating_period(values) + shape = ( + f"repeating every {period} layer(s)" + if period is not None + else "with no fixed period" + ) + hybrid_detail = f"{len(values)} layers in a pattern {shape}: {tally}" + hybrid_detail += ( + "; which value means which kind of attention is the config's own " + "convention and is a question for the plan phase, not something this " + "survey can know" + ) + differing: list[str] = [] + agreeing: list[str] = [] + geometry_keys: list[str] = [] + for name in _HYBRID_GEOMETRY_KEYS: + base = scoped.get(name) + for prefix in _HYBRID_KIND_PREFIXES: + counterpart = scoped.get(prefix + name) + if counterpart is None: + continue + geometry_keys.append(prefix + name) + stated = f"{prefix}{name}={counterpart} against {name}={base}" + (agreeing if counterpart == base else differing).append(stated) + if differing: + hybrid_detail += ( + f". The two kinds declare DIFFERENT geometry: {', '.join(differing)} " + f"— two shapes in one stack, not one shape under two masks" + ) + elif agreeing: + hybrid_detail += f". Per-kind geometry is declared and agrees: {', '.join(agreeing)}" + demands.append( + Demand("attention.hybrid", hybrid_detail, pattern_keys + tuple(geometry_keys)) + ) + + # A learned term in the softmax DENOMINATOR, and the layer kind that + # carries it is in the key name and nowhere else. Distinct from + # `attention.sinks` below, which reads the count-of-sink-tokens spelling: + # this one is the per-layer-kind boolean, and a target that silently drops + # it boots, generates plausible text and is wrong by a constant per head. + sink_keys = tuple( + key for key in sorted(scoped) if isinstance(key, str) and "attention_sink" in key + ) + if sink_keys: + carried = ", ".join(f"{_sink_layer_kind(key)} layers={scoped[key]}" for key in sink_keys) + demands.append( + Demand( + "attention.sink", + f"a learned sink term in the softmax denominator, declared per layer " + f"kind: {carried}", + sink_keys, + ) + ) + + if keys := _present(scoped, "attn_logit_softcapping", "final_logit_softcapping"): + demands.append(Demand("attention.softcap", "logit soft-capping in the softmax", keys)) + if keys := _present(scoped, "sinks", "attn_sinks", "num_sinks"): + demands.append(Demand("attention.sinks", "extra softmax-denominator logits", keys)) + # A LEARNED selector, not a mask. A window or a sink leaves cache + # addressing alone; an indexer decides per query which past positions are + # attended to, so the top-k selection is part of the forward and part of + # what a receipt has to certify. Missing this probe once meant a + # 48-layer checkpoint with `indexer_budget: 2048` surfaced no sparse + # demand at all and the whole sparse_attention area stayed unread. + if keys := _present( + scoped, + "indexer_budget", + "indexer_n_heads", + "indexer_head_dim", + "index_topk", + "index_n_heads", + "sparse_attention_config", + ): + budget = scoped.get("indexer_budget") or scoped.get("index_topk") + detail = "a learned index selects which past positions each query attends to" + if budget: + detail += f"; budget={budget}" + demands.append(Demand("attention.indexer", detail, keys)) + if keys := _present(scoped, "use_qk_norm", "qk_layernorm", "qk_norm"): + demands.append(Demand("attention.qk_norm", "per-head q/k normalization before RoPE", keys)) + + # transformers 5.x moved the whole RoPE description into `rope_parameters` + # and left `rope_scaling` behind. Reading only the old name made a + # checkpoint with `rope_parameters.mrope_section` surface NO rope demand + # except `rope.partial` -- so mrope, the thing a multimodal target is most + # likely to get wrong, went unmentioned. Accept both spellings, and + # remember which one was read so the citation is not a lie. + rope_key = "rope_parameters" if isinstance(scoped.get("rope_parameters"), Mapping) else None + if rope_key is None and isinstance(scoped.get("rope_scaling"), Mapping): + rope_key = "rope_scaling" + rope_block: Mapping[str, Any] = scoped.get(rope_key) if rope_key else {} # type: ignore[assignment] + + if rope_block: + kind = rope_block.get("rope_type") or rope_block.get("type") or "unknown" + theta = rope_block.get("rope_theta") or scoped.get("rope_theta") + if str(kind) in ("default", "unknown") and theta is not None: + demands.append(Demand("rope.plain", f"theta={theta}", (rope_key,))) + else: + demands.append(Demand("rope.scaling", f"{kind} — not plain RoPE", (rope_key,))) + elif scoped.get("rope_theta") is not None: + demands.append(Demand("rope.plain", f"theta={scoped['rope_theta']}", ("rope_theta",))) + if scoped.get("partial_rotary_factor") is not None or rope_block.get("partial_rotary_factor"): + factor = scoped.get("partial_rotary_factor") or rope_block.get("partial_rotary_factor") + demands.append( + Demand( + "rope.partial", + f"rotary applied to {factor} of each head", + ("partial_rotary_factor",), + ) + ) + if scoped.get("use_mrope") or rope_block.get("mrope_section"): + section = rope_block.get("mrope_section") + detail = "multimodal sectioned RoPE" + if section: + detail += f"; sections={section}" + if rope_block.get("mrope_interleaved"): + detail += ", interleaved" + demands.append(Demand("rope.mrope", detail, (rope_key or "rope_scaling",))) + + act = ( + scoped.get("hidden_act") + or scoped.get("hidden_activation") + or scoped.get("activation_function") + ) + if act: + demands.append(Demand("mlp.activation", str(act), ("hidden_act",))) + + experts = _num( + scoped, "num_experts", "num_local_experts", "n_routed_experts", "moe_num_experts" + ) + if experts: + topk = _num(scoped, "num_experts_per_tok", "moe_topk", "topk", "n_group") + shared = _num(scoped, "n_shared_experts", "shared_expert_intermediate_size") + bits = [f"{experts} experts", f"top-{topk}" if topk else "top-k unstated"] + if shared: + bits.append("plus a shared expert") + if scoped.get("scoring_func"): + bits.append(f"scoring={scoped['scoring_func']}") + # The scoring function says how an expert is SCORED; these three say + # how the top-k is SELECTED, and they pick a different routing entry + # rather than a different constant. `noaux_tc` over `n_group` groups + # with `topk_group` groups kept is grouped routing with a bias + # correction, which is its own op — a target that reads only + # `scoring_func` builds plain sigmoid top-k and routes to the wrong + # experts while scoring close enough to pass a short gate. + if scoped.get("topk_method"): + bits.append(f"topk_method={scoped['topk_method']}") + if scoped.get("n_group") is not None: + bits.append(f"n_group={scoped['n_group']}") + if scoped.get("topk_group") is not None: + bits.append(f"topk_group={scoped['topk_group']}") + if scoped.get("norm_topk_prob") is not None: + bits.append(f"norm_topk_prob={scoped['norm_topk_prob']}") + demands.append(Demand("moe.routing", ", ".join(bits), ("num_experts",))) + if scoped.get("mlp_only_layers") or scoped.get("decoder_sparse_step"): + demands.append( + Demand( + "moe.interleave", + "dense and sparse layers interleave — two MLP shapes in one stack", + _present(scoped, "mlp_only_layers", "decoder_sparse_step"), + ) + ) + + if keys := _present( + scoped, + "mamba_d_state", + "mamba_expand", + "ssm_state_size", + "conv_kernel", + "linear_conv_kernel_dim", + ): + demands.append( + Demand( + "recurrent.state", + "a recurrent/SSM state plane — needs a state cache manager, not a KV " + "page pool, and the two are managed by different code", + keys, + ) + ) + layer_types = scoped.get("layer_types") + if isinstance(layer_types, Sequence) and not isinstance(layer_types, str): + kinds = sorted({str(t) for t in layer_types}) + if len(kinds) > 1: + demands.append( + Demand("stack.hybrid", f"layer kinds interleave: {kinds}", ("layer_types",)) + ) + + # `mtp_num_hidden_layers` is here because the first checkpoint this ran on + # used it and the probe missed it — the key fell out in `unrecognized` + # instead, which is how it was noticed. Left as evidence that the probe + # list is allowed to be incomplete as long as the leftovers are reported. + nextn = _num(scoped, "num_nextn_predict_layers", "num_mtp_modules", "mtp_num_hidden_layers") + # DSpark is read FIRST because DeepSeek-V4.1 carries both spellings: its + # DSpark stages live under `mtp.*` with `num_nextn_predict_layers: 3`, and + # the MTP reading of that key is the one that builds the wrong drafter. + if dspark := _present( + scoped, + "dspark_block_size", + "dspark_noise_token_id", + "dspark_target_layer_ids", + "dspark_markov_rank", + ): + demands.append( + Demand( + "speculation.heads", + f"a DSpark block drafter (block {scoped.get('dspark_block_size', '?')}, " + f"target layers {scoped.get('dspark_target_layer_ids', '?')}) is the " + f"checkpoint's own, stored as mtp.* beside num_nextn_predict_layers=" + f"{nextn} — the feature is `dspark`, NOT `mtp`, which would build chained " + f"one-token heads over weights trained as one block pass", + dspark, + ) + ) + elif nextn: + demands.append( + Demand( + "speculation.heads", + f"{nextn} in-checkpoint predictor layer(s) — weights exist for a " + f"speculator, so MTP is the checkpoint's own and not a bolt-on", + ("num_nextn_predict_layers",), + ) + ) + + quant = config.get("quantization_config") or scoped.get("quantization_config") + if isinstance(quant, Mapping): + algo = quant.get("quant_algo") or quant.get("quant_method") or "unstated" + kv = quant.get("kv_cache_quant_algo") + bits = [str(algo)] + if kv: + bits.append(f"kv={kv}") + if quant.get("group_size"): + bits.append(f"group={quant['group_size']}") + excluded = ( + quant.get("exclude_modules") + or quant.get("ignore") + or quant.get("modules_to_not_convert") + ) + if excluded: + bits.append(f"{len(excluded)} module(s) held out") + demands.append(Demand("quantization", ", ".join(bits), ("quantization_config",))) + else: + demands.append( + Demand("quantization", "none declared — weights load at their stored dtype", ()) + ) + + norm_keys = _present(scoped, "rms_norm_eps", "layer_norm_eps", "layer_norm_epsilon", "norm_eps") + if norm_keys: + style = "RMSNorm" if any("rms" in k for k in norm_keys) else "LayerNorm" + demands.append(Demand("norm.style", style, norm_keys)) + + # The residual stream is not one tensor. Hyper-connections carry `n` + # parallel streams mixed at every block boundary, which changes the shape + # of every residual add in the forward -- so a target that assumes one + # stream is not slightly wrong, it is a different model. + if keys := _present( + scoped, "hc_count", "hc_lowrank", "num_residual_streams", "hyper_connection_count" + ): + count = scoped.get("hc_count") or scoped.get("num_residual_streams") + detail = "parallel residual streams mixed at every block boundary" + if count: + detail = f"{count} " + detail + if scoped.get("hc_lowrank"): + detail += f"; low-rank mixer width {scoped['hc_lowrank']}" + demands.append(Demand("residual.hyper_connection", detail, keys)) + + # A tower in front of the language model. Keyed on the OUTER config, + # because `scoped` is the text tower by then and the vision config is its + # sibling, not its child. + tower_keys = tuple( + k for k in ("vision_config", "audio_config", "speech_config") if config.get(k) + ) + token_keys = _present( + config, "image_token_id", "video_token_id", "audio_token_id", "vision_start_token_id" + ) + if tower_keys or token_keys: + towers = ", ".join(k.replace("_config", "") for k in tower_keys) or "unnamed" + detail = f"{towers} tower in front of the language model" + if token_keys: + detail += f"; placeholder tokens declared ({', '.join(token_keys)})" + demands.append(Demand("multimodal.tower", detail, tower_keys + token_keys)) + + # A secondary table whose size dwarfs the vocabulary. This is the + # signature of an n-gram / product-key / memory-layer embedding, and the + # reason it matters is residency: a table an order of magnitude larger + # than the vocabulary can be most of the checkpoint, so whether it lives + # on the device decides how much is left for KV. + primary = _num(scoped, "vocab_size") or 0 + for key, value in sorted(scoped.items()): + if key == "vocab_size" or "vocab_size" not in key: + continue + size = value if isinstance(value, int) and not isinstance(value, bool) else None + if size is None or primary <= 0 or size < primary * 10: + continue + demands.append( + Demand( + "embedding.huge_table", + f"{key}={size:,} against vocab_size={primary:,} — a secondary table " + f"{size // primary}x the vocabulary, so residency is a decision and " + f"not a default", + (key, "vocab_size"), + ) + ) + break + + unrecognized = tuple( + sorted( + key + for key in scoped + if isinstance(key, str) + and key not in _INERT_KEYS + and key not in _SUBMODEL_KEYS + and not key.startswith("_") + and not key.startswith(_CLAIMED_PREFIXES) + and not any(claimed in key for claimed in _CLAIMED_SUBSTRINGS) + ) + ) + + architectures = config.get("architectures") or [] + return CheckpointSurface( + architecture=str(architectures[0]) if architectures else "", + model_type=str(config.get("model_type") or ""), + demands=tuple(demands), + unrecognized=unrecognized, + submodels=submodels, + ) + + +# -------------------------------------------------------------------------- +# half two: what the build supplies +# -------------------------------------------------------------------------- + +RECEIPT_PASSED = "passed" +RECEIPT_FAILED = "failed" +RECEIPT_ABSENT = "absent" + + +@dataclass(frozen=True) +class CatalogEntry: + """One catalog entry as the host actually holds it.""" + + path: str + impl: str = "" + summary: str = "" + receipt: str = RECEIPT_ABSENT + impl_resolves: bool = True + + @property + def name(self) -> str: + return Path(self.path).stem + + @property + def category(self) -> str: + parts = Path(self.path).parts + return parts[0] if len(parts) > 1 else "" + + @property + def certifiable(self) -> bool: + """Whether a receipt is a thing this entry can even have. + + ``torch/`` entries are thin mirrors of torch callables; upstream owns + their correctness, so the catalog gives them no contract, no test and + no architecture certification. For those, absence is *by design* and + not the three-state "unknown" — listing them as "a receipt run is + what is missing" points at work that does not exist. + """ + return self.category != "torch" + + @property + def spendable(self) -> bool: + """Usable without a smith dispatch: certified here and still wired.""" + if not self.impl_resolves: + return False + return self.receipt == RECEIPT_PASSED or not self.certifiable + + +@dataclass(frozen=True) +class Inventory: + """What the host supplies, measured rather than assumed.""" + + arch: str = "" + trtllm_ops: tuple[str, ...] = () + catalog: tuple[CatalogEntry, ...] = () + modules: tuple[str, ...] = () + #: op name -> the source that registered it. Empty when the op inventory + #: could not be taken; a name absent from the Python plane maps to + #: ``cpp_thop``, which is a conclusion rather than a default. + op_sources: dict[str, str] = field(default_factory=dict) + versions: dict[str, str] = field(default_factory=dict) + error: str = "" + + @property + def ok(self) -> bool: + return not self.error + + @property + def spendable(self) -> tuple[CatalogEntry, ...]: + return tuple(e for e in self.catalog if e.spendable) + + @property + def uncertified(self) -> tuple[CatalogEntry, ...]: + """Present, wired, certifiable, but carrying no pass for this arch. + + Includes the entries whose receipt says ``failed``, which are the + most important rows here: a known-failing entry that reads as merely + "not yet run" is how a plan spends vocabulary that is already known + not to work. + """ + return tuple( + e for e in self.catalog if e.certifiable and e.impl_resolves and not e.spendable + ) + + @property + def rotted(self) -> tuple[CatalogEntry, ...]: + """Entries whose ``impl`` no longer resolves in this build. + + A rotted entry is worse than a missing one: it reads as vocabulary + the plan may spend, and it fails at the first call rather than at + selection. + """ + return tuple(e for e in self.catalog if not e.impl_resolves) + + @property + def unwrapped_ops(self) -> tuple[str, ...]: + """Registered ops that no catalog entry wraps. + + This is the tier-1 rung made concrete: a gap answered by a name in + here costs a wrapper and a receipt, not an implementation. + """ + wrapped = {e.impl.rsplit(".", 1)[-1] for e in self.catalog if e.impl} + return tuple(op for op in self.trtllm_ops if op not in wrapped) + + +#: Runs on the host, inside the run's own interpreter. Kept as a string +#: rather than a module so it can be shipped through any transport that +#: carries a shell command — ``srun``, a spool file, or a broker call. +PROBE_SOURCE = r""" +import ast, json, os, re, sys +from pathlib import Path + +out = {"trtllm_ops": [], "op_sources": {}, "catalog": [], "modules": [], "versions": {}, "error": ""} +repo = Path(sys.argv[1]) +arch = sys.argv[2] + +try: + import torch + out["versions"]["torch"] = torch.__version__ + import tensorrt_llm + out["versions"]["tensorrt_llm"] = getattr(tensorrt_llm, "__version__", "?") + # `torch.ops.` is a LAZY proxy: `dir()` lists only the names something + # has already resolved, not the names that are registered. On this build + # that is 91 out of 273 -- so a supply survey built on `dir()` reports two + # thirds of the op plane as missing, and the catalog entries wrapping them + # as "rotted: do not spend these". Measured: `allgather`, `fp4_quantize`, + # `fused_qk_norm_rope`, `causal_conv1d_fwd`, `noaux_tc_op` and eight more + # were all reported absent and all resolve. + # + # The registry is the truth. `dir()` stays as a fallback for a torch + # without the private accessor, where under-reporting beats no report. + # Import the Python registration plane before enumerating. An op whose + # `custom_ops/*.py` module has not been imported has not registered its + # schema either, so it is absent from BOTH the registry and `dir()` --- and + # the entry wrapping it is reported as rotted. Measured: two entries + # (`fused_sigmoid_mul_inplace`, `rms_norm_gated_token_major`) read absent + # before the import and resolved 27 seconds later with it. Importing + # `tensorrt_llm` alone is not enough; the submodule has to be touched. + try: + import importlib + + importlib.import_module("tensorrt_llm._torch.custom_ops") + out["versions"]["custom_ops_imported"] = "yes" + except Exception as exc: + out["versions"]["custom_ops_imported"] = "no: %s" % type(exc).__name__ + ops = getattr(torch.ops, "trtllm", None) + names = set() + try: + for schema in torch._C._jit_get_all_schemas(): + if schema.name.startswith("trtllm::"): + names.add(schema.name.split("::", 1)[1]) + except Exception: + pass + if ops is not None: + names.update(n for n in dir(ops) if not n.startswith("_")) + out["trtllm_ops"] = sorted(names) + out["op_enumeration"] = "schema_registry" if len(names) > 200 else "dir_fallback" +except Exception as exc: + out["error"] = "%s: %s" % (type(exc).__name__, exc) + +CATALOG_PKG = "tensorrt_llm._torch._experimental.modeling_v2.catalog" + +def resolves(dotted, entry_path=None): + if not dotted: + return True + if dotted.startswith("torch.ops.trtllm."): + name = dotted.rsplit(".", 1)[-1] + if name in out["trtllm_ops"]: + return True + # Not in the registry YET. An op registers when the module holding + # its `@torch.library.custom_op` is imported, and several live + # outside `custom_ops/` -- `rms_norm_gated_token_major` registers in + # `_torch.modules.mamba.layernorm_gated`, and its catalog wrapper + # says so in a comment at the top of the file. So ask the wrapper: + # importing it performs whatever import the op needs, by + # construction, because that is what makes the wrapper work. + # + # Without this, two live entries were reported as "rotted -- do not + # spend these" and the plan phase had to argue the survey down. + if entry_path: + import importlib + module = CATALOG_PKG + "." + entry_path[:-3].replace("/", ".") + try: + importlib.import_module(module) + except Exception: + return False + try: + for schema in torch._C._jit_get_all_schemas(): + if schema.name.startswith("trtllm::"): + out["trtllm_ops"].append(schema.name.split("::", 1)[1]) + out["trtllm_ops"] = sorted(set(out["trtllm_ops"])) + except Exception: + pass + return name in out["trtllm_ops"] + return False + import importlib + parts = dotted.split(".") + for cut in range(len(parts) - 1, 0, -1): + try: + obj = importlib.import_module(".".join(parts[:cut])) + except Exception: + continue + for attr in parts[cut:]: + obj = getattr(obj, attr, None) + if obj is None: + return False + return True + return False + +def frontmatter(md_path): + try: + text = md_path.read_text(encoding="utf-8") + except OSError: + return "" + if not text.startswith("---") or text.count("---") < 2: + return "" + return text.split("---", 2)[1] + +def receipt_state(md_path, arch): + # Three-state, keyed by arch: passed / failed / absent. Absent means + # UNKNOWN, and conflating it with either other state is the whole reason + # this is three-state. + # + # Two frontmatter shapes are in use -- `sm_103: {status: passed, tests: 9}` + # and the block form with `status:` on its own indented line -- and an + # earlier hand-rolled reader got both wrong in the same direction: it + # reported the inline FAILED entry as passed, and the block FAILED entry + # as absent. yaml is used when importable, which in the run's own + # interpreter it always is; the fallback below reads `status:` rather + # than guessing from the presence of a value. + head = frontmatter(md_path) + if not head: + return "absent" + try: + import yaml + receipts = (yaml.safe_load(head) or {}).get("receipts") or {} + record = receipts.get(arch) + if record is None: + return "absent" + if isinstance(record, dict): + status = str(record.get("status", "")).lower() + else: + status = str(record).lower() + if status.startswith(("pass", "true", "ok")): + return "passed" + if status.startswith(("fail", "false")): + return "failed" + return "absent" + except ImportError: + pass + except Exception: + return "absent" + + inside, base_indent, status = False, None, "" + for line in head.splitlines(): + stripped = line.strip() + if not stripped: + continue + indent = len(line) - len(line.lstrip()) + if stripped.startswith("receipts:"): + inside, base_indent = True, indent + continue + if not inside: + continue + if indent <= base_indent: + break + if stripped.startswith(arch + ":"): + rest = stripped.split(":", 1)[1].strip() + if rest: + # inline flow mapping, or a bare scalar + marker = "status:" + status = ( + rest.split(marker, 1)[1].strip(" {}'\"").split(",")[0].strip() + if marker in rest + else rest.strip("{}'\"") + ) + else: + # block form: the first `status:` under this arch key + status = "__block__" + continue + if status == "__block__": + if stripped.startswith("status:"): + status = stripped.split(":", 1)[1].strip().strip("'\"") + elif indent <= base_indent + 2: + status = "" # left the arch block without ever seeing a status + break + status = status.lower() + if status.startswith(("pass", "true", "ok")): + return "passed" + if status.startswith(("fail", "false")): + return "failed" + return "absent" + +catalog = repo / "tensorrt_llm/_torch/_experimental/modeling_v2/catalog" +index = catalog / "index.yaml" +if index.is_file(): + # Same reason as above: a five-field list read without a YAML parser. + entry = None + # The list may sit at column 0 or be indented under `entries:` -- the + # in-tree index indents it, and a column-0 match read that index as EMPTY, + # so a dry run printed `0/0 entries certified` over 19 certified entries. + for raw in index.read_text(encoding="utf-8").splitlines(): + line = raw.strip() + if line.startswith("- path:"): + if entry: + out["catalog"].append(entry) + entry = {"path": line.split(":", 1)[1].strip(), "impl": "", "summary": ""} + elif entry is not None and line.startswith("impl:"): + entry["impl"] = line.split(":", 1)[1].strip() + elif entry is not None and line.startswith("summary:"): + entry["summary"] = line.split(":", 1)[1].strip().strip("'\"")[:160] + if entry: + out["catalog"].append(entry) +for entry in out["catalog"]: + md = catalog / entry["path"].replace(".py", ".md") + entry["receipt"] = receipt_state(md, arch) + # Only claim an entry is unwired when the op inventory was actually + # taken. Without that guard an engine that failed to import reports the + # whole catalog as rotted, which is the most alarming possible way to + # say "I could not measure this". + entry["impl_resolves"] = True if out["error"] else bool(resolves(entry["impl"], entry.get("path"))) + +# Provenance. Every op enters `torch.ops.trtllm` through one of two planes, +# and the Python plane names its source in the file name. An op matched by +# no file is registered from C++, which is itself the answer. +custom_ops = repo / "tensorrt_llm/_torch/custom_ops" +registrations = {} +for py in sorted(custom_ops.glob("*.py")): + try: + body = py.read_text(encoding="utf-8") + except OSError: + continue + for name in re.findall(r'"trtllm::([a-zA-Z0-9_]+)', body): + registrations.setdefault(name, py.stem) +out["op_sources"] = {op: registrations.get(op, "cpp_thop") for op in out["trtllm_ops"]} + +modules_dir = repo / "tensorrt_llm/_torch/modules" +for py in sorted(modules_dir.rglob("*.py")): + try: + tree = ast.parse(py.read_text(encoding="utf-8")) + except (OSError, SyntaxError): + continue + rel = py.relative_to(modules_dir).with_suffix("").as_posix().replace("/", ".") + for node in tree.body: + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and not node.name.startswith("_"): + out["modules"].append("%s.%s" % (rel, node.name)) +out["modules"] = sorted(set(out["modules"])) + +print("__SURVEY_JSON__" + json.dumps(out)) +""" + + +def probe_command(environment: RunEnvironment, *, repo: str | Path, arch: str) -> str: + """The literal command that takes the inventory, fully prefixed. + + Launched as **one** process whatever the run's world size. This is + read-only introspection — it imports the build, walks ``torch.ops.trtllm`` + and parses catalog frontmatter off disk — so nothing in it wants a rank + per device, and taking the engine's launch rules instead cost this + workflow a whole survey once: at world size 4 the probe went out as + ``mpirun -n 4 ... trtllm-llmapi-launch``, the run's own ``PATH`` did not + carry the container's MPI, and the rc=127 arrived at the plan phase as an + empty inventory. + + The distinction belongs to :meth:`RunEnvironment.command`, not here. + Assembling a launcher line outside ``environment.py`` is exactly what + that module exists to prevent. + """ + return environment.python_command( + ["-c", PROBE_SOURCE, str(repo), str(arch)], single_process=True + ) + + +def parse_inventory(stdout: str, *, arch: str = "") -> Inventory: + """Pull the probe's payload out of a log that also carries engine chatter.""" + match = re.search(rf"^{re.escape(MARKER)}(.*)$", stdout, re.MULTILINE) + if not match: + tail = "\n".join(stdout.splitlines()[-15:]) + return Inventory(arch=arch, error=f"probe printed no payload. last lines:\n{tail}") + try: + data = json.loads(match.group(1)) + except json.JSONDecodeError as exc: + return Inventory(arch=arch, error=f"probe payload is not JSON: {exc}") + entries = tuple( + CatalogEntry( + path=str(e.get("path", "")), + impl=str(e.get("impl", "")), + summary=str(e.get("summary", "")), + receipt=str(e.get("receipt", RECEIPT_ABSENT)), + impl_resolves=bool(e.get("impl_resolves", True)), + ) + for e in data.get("catalog", ()) + ) + return Inventory( + arch=arch, + trtllm_ops=tuple(str(o) for o in data.get("trtllm_ops", ())), + catalog=entries, + modules=tuple(str(m) for m in data.get("modules", ())), + op_sources={str(k): str(v) for k, v in (data.get("op_sources") or {}).items()}, + versions={str(k): str(v) for k, v in (data.get("versions") or {}).items()}, + error=str(data.get("error", "")), + ) + + +def run_inventory_probe( + task_data: Mapping[str, Any], + environment: RunEnvironment, + *, + shell_runner: Any, + timeout: int = INVENTORY_TIMEOUT_S, +) -> Inventory: + """Take the inventory on the host that will run the build. + + A failure here is reported, not raised. The plan can be drafted without + an inventory — it just gets drafted the old way, on assumption — and + losing the whole run because a probe could not import is a worse trade + than losing the measurement. + """ + arch = str(task_data.get("arch", "")) + command = probe_command(environment, repo=task_data["trtllm_repo_path"], arch=arch) + result = shell_runner(command, None, timeout) + if getattr(result, "timed_out", False): + return Inventory(arch=arch, error=f"inventory probe timed out after {timeout}s") + stdout = getattr(result, "stdout", "") + getattr(result, "stderr", "") + inventory = parse_inventory(stdout, arch=arch) + # A payload that landed is kept whatever the exit code says. The probe + # catches its own import failures and still prints what it read off + # disk, and the catalog half is most of the value. + landed = bool(inventory.catalog or inventory.trtllm_ops or inventory.modules) + if not landed and not getattr(result, "ok", True): + return Inventory( + arch=arch, + error=f"{inventory.error} (probe exited {getattr(result, 'returncode', '?')})", + ) + return inventory + + +# -------------------------------------------------------------------------- +# the block the drafter reads +# -------------------------------------------------------------------------- + + +def format_surface(surface: CheckpointSurface, inventory: Inventory, *, limit: int = 40) -> str: + """Render both columns for the plan prompt, joined by nothing. + + ``limit`` caps the long lists. A drafter that needs the full inventory + can read ``surface.md`` in the workspace, which is written unabridged. + """ + lines = ["## Surface survey", ""] + # Whether the right-hand column is a reading or a blank is the first thing + # this block has to say. An inventory that failed to land renders with the + # same shape as one that landed empty — "0 of 0 entries certified" — and a + # drafter that reads the second as the first plans every entry as a gap. + taken = ( + "Measured before this phase started." + if inventory.ok + else ( + "The demand column was measured before this phase started. **The supply " + "column was NOT.** The inventory probe failed, so the right-hand side " + "below is an assumption and not a reading: nothing missing from it is " + "evidence of absence, no row in it can carry a ceiling, and the plan must " + "name the missing measurement as an open item." + ) + ) + lines.append( + f"{taken} The left column is what the checkpoint " + f"demands; the right is what `{inventory.arch or 'this arch'}` supplies. " + f"**Joining them is your job** — no function here guesses which entry answers " + f"which demand." + ) + lines.append("") + + lines.append(f"### Demands — `{surface.architecture or '?'}` (`{surface.model_type or '?'}`)") + lines.append("") + if surface.submodels: + lines.append( + f"Multimodal checkpoint: nested {', '.join(surface.submodels)}. Geometry below " + f"is read from `text_config` where one exists.\n" + ) + for demand in surface.demands: + lines.append(f"- {demand}") + if surface.unrecognized: + lines.append("") + lines.append( + "**Unrecognized structural keys** — no probe here claims these, so treat each " + "as a candidate for a layer this vocabulary has no word for. Read them before " + "you size the queue; this is the cheapest moment in the run to discover a " + "ceiling:" + ) + lines.append("") + lines.append(f" `{'`, `'.join(surface.unrecognized[:limit])}`") + lines.append("") + + lines.append("### Supply") + lines.append("") + # A probe can fail at the engine import and still have read the catalog + # off disk, which is most of the value. Reporting "nothing was measured" + # when two thirds of it was is its own kind of wrong answer — so the + # error is stated, and whatever landed is still rendered. + if not inventory.ok: + lines.append( + f"**The inventory probe reported an error**: `{inventory.error}`. Anything " + f"below was still measured; anything absent was not. In particular, with no " + f"op inventory nothing can be called rotted and nothing can be called " + f"unwrapped — treat both as unknown and say so in the plan's open items." + ) + lines.append("") + if not inventory.catalog and not inventory.trtllm_ops: + lines.append( + "Nothing landed at all. Plan on the assumption that you do not know what " + "exists, and name that as an open item." + ) + return "\n".join(lines) + "\n" + + spendable = inventory.spendable + lines.append( + f"{len(spendable)} of {len(inventory.catalog)} catalog entries are spendable on " + f"`{inventory.arch}` — certified here and still wired, or a `torch/` mirror that " + f"needs no certification. {len(inventory.trtllm_ops)} ops are registered in this " + f"build." + ) + lines.append("") + if spendable: + lines.append("**Free vocabulary** — spend without a smith dispatch:") + lines.append("") + for entry in spendable[:limit]: + lines.append(f"- `{entry.category}/{entry.name}`") + if len(spendable) > limit: + lines.append(f"- … {len(spendable) - limit} more, see `surface.md`") + lines.append("") + # Split by state rather than lumped, because the two mean opposite + # things about how much work is left. `failed` is not "not yet run". + failed = [e for e in inventory.uncertified if e.receipt == RECEIPT_FAILED] + unrun = [e for e in inventory.uncertified if e.receipt != RECEIPT_FAILED] + if failed: + lines.append( + "**Known failing on this arch** — the entry's own receipt records a failure " + "here. This is the opposite of un-run: someone took the measurement and it " + "did not pass. A plan that spends one of these is spending vocabulary that " + "is already known not to work, so budget for the failure or route around it:" + ) + lines.append("") + for entry in failed[:limit]: + lines.append(f"- `{entry.category}/{entry.name}` — receipt: failed") + lines.append("") + if unrun: + lines.append( + "**Present but un-run on this arch** — the wrapper exists and the receipt is " + "absent, which reads as *unknown*. A receipt run is what is missing, and that " + "is much cheaper than an onboard:" + ) + lines.append("") + for entry in unrun[:limit]: + lines.append(f"- `{entry.category}/{entry.name}`") + lines.append("") + if inventory.rotted: + lines.append( + "**Rotted** — the entry names an `impl` that does not resolve in this build. " + "This is worse than missing: it reads as vocabulary and fails at the call. " + "Do not spend these:" + ) + lines.append("") + for entry in inventory.rotted: + lines.append(f"- `{entry.category}/{entry.name}` → `{entry.impl}`") + lines.append("") + unwrapped = inventory.unwrapped_ops + if unwrapped: + lines.append( + f"**{len(unwrapped)} registered ops carry no catalog entry.** A demand " + f"answered by one of these is tier 1 — a wrapper and a receipt, not an " + f"implementation. Search this list before you conclude anything is missing:" + ) + lines.append("") + lines.append(f" `{'`, `'.join(unwrapped[:limit])}`") + if len(unwrapped) > limit: + lines.append(f"\n … {len(unwrapped) - limit} more, see `surface.md`") + lines.append("") + if inventory.op_sources: + tally = Counter(inventory.op_sources.values()) + breakdown = ", ".join(f"{n} {src}" for src, n in tally.most_common()) + lines.append( + f"**Where those ops came from**: {breakdown}. Provenance decides what a " + f"receipt must say — an architecture-gated JIT kernel and a compiled-in " + f"C++ one are certified differently. See the kernel-source registry." + ) + lines.append("") + lines.append( + f"The kernel plane (`_torch/modules/**`) holds {len(inventory.modules)} public " + f"callables — the tier-2 rung. `surface.md` lists them." + ) + return "\n".join(lines) + "\n" + + +def write_surface_file(path: str | Path, surface: CheckpointSurface, inventory: Inventory) -> Path: + """Write the unabridged survey to the workspace, for audit and for agents.""" + out = Path(path) + body = [format_surface(surface, inventory, limit=10_000), "", "## Kernel plane", ""] + body.extend(f"- `{m}`" for m in inventory.modules) + body.extend(["", "## Versions", ""]) + body.extend(f"- {k}: {v}" for k, v in sorted(inventory.versions.items())) + out.write_text("\n".join(body) + "\n", encoding="utf-8") + return out + + +__all__ = [ + "CatalogEntry", + "CheckpointSurface", + "Demand", + "INVENTORY_TIMEOUT_S", + "Inventory", + "PROBE_SOURCE", + "RECEIPT_ABSENT", + "RECEIPT_FAILED", + "RECEIPT_PASSED", + "format_surface", + "parse_inventory", + "probe_command", + "read_checkpoint_surface", + "run_inventory_probe", + "write_surface_file", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml b/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml new file mode 100644 index 000000000000..ef9a592137b7 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml @@ -0,0 +1,564 @@ +# staircase-onboard task specification. +# +# Copy this file, fill it in, and run: +# +# staircase-onboard --task task.yaml --workspace workspace/staircase-onboard +# +# 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. +# +# 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. That is not +# pedantry — `feature:` instead of `features:` used to validate cleanly and +# resolve to no features, so the run built core, gated it green, and +# reported 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: `unified` is the default, and a unified target +# used to be built once and gated once, at pure tensor parallelism. Expert +# parallelism and attention data parallelism — the two branches a unified +# codebase exists to carry, and for a 256-expert MoE the ones holding the +# expert sharding and the all-to-all — were switched on by no gate, while the +# target's Coverage section documented all three. An EP bug is invisible to a +# TP-only gate by construction. +# +# Every member must declare the SAME world size: one allocation serves them +# all, so a member that wants a different device count cannot be run at all. +# `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, eplb, dwdp, 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 (DeepSeek-V4.1: `dspark_*` keys in its config) asks for +# `dspark` even when the request says "MTP"; `mtp` over those weights builds +# three chained heads the checkpoint never trained. The two are refused +# together. +features: [] + +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 + +# `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 + +# 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. +# source_embargo: +# - tensorrt_llm/_torch/models/modeling_qwen3_5.py +# - /path/to/vllm + +# ----------------------------------------------------------------- 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. Set both when a unit needs a switch changed: +# the boot gate reads this, the accuracy gate reads the YAML, and two gates +# measuring different engines is worse than either being wrong. +# +# # 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 + +# ------------------------------------------------------------- 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 | slurm | slurm_broker | session_spool + # + # Prefer `session_spool` wherever it is available, and not for + # convenience: every other kind puts a scheduler between a command and its + # result — a container image import (26 GB, 4-8 minutes) and a queue wait + # PER COMMAND. A spool acquires one allocation once, and work is a file, + # so anything that can write a file can run GPU work. + # + # `slurm_broker` is for a cluster whose Slurm CLI cannot reach the + # controller from here and where an MCP broker is the only path. It has + # one consequence worth knowing before you pick it: the orchestrator is + # not an MCP client, so gate *execution* moves to an agent while gate + # *adjudication* stays in Python. + 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 + + # --- kind: slurm / slurm_broker only ------------------------------------ + # partition: batch + # account: your-account + # container_image: /path/to/image.sqsh + # mounts: + # - /lustre:/lustre + # time_limit: "04:00:00" + # + # GPUs per NODE. With `parallel: tep8` and 4 GPUs per node the workflow + # derives 2 nodes, 8 tasks, 4 tasks per node — one task per rank, which is + # the thing that decides whether the launch works at all. See below. + # 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 gate runs. + # container_name: trtllm-staircase + # + # `auto` (default) picks the MPI plugin from the task count: `none` at one + # task, `pmix` above it. Set it only when this cluster reports a versioned + # variant — ask it: `srun --mpi=list`. + # mpi: pmix_v5 + # + # 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. + + # --- per-workload substrate ---------------------------------------------- + # The right substrate depends on HOW LONG THE JOB RUNS, and a run does not + # have to pick once. + # + # A `session_spool` imports its container image ONCE and then costs nothing + # per job. That is the correct trade for a smith's entry test, which runs for + # a minute; one onboarding ran 443 such jobs, and paying the import on each + # would have cost roughly 37 hours. + # + # It is the wrong trade for a gate. A boot or accuracy gate runs 10-15 + # minutes, so the import is 30-50% overhead rather than 3000%, and what you + # want in exchange is the thing a shared session cannot give: ISOLATION. A + # spool has one failure domain. Across the two runs this was written for, a + # gate that leaked a 173 GB server, a teardown whose process-group kill + # caught the watcher, an agent's `pkill -9` that did the same, and two + # engines colliding over PMIx (`status = -61` cancels the whole step) each + # took down every other job AND the orchestrator. + # + # So: short work on the spool, gates in their own Slurm jobs. Only `gate` is + # recognised today; an unknown name is REJECTED rather than ignored, because + # an override that never applies is invisible — the run just uses the base + # environment and nothing says why. + # + # Keys are the environment keys above. `env` merges key-by-key with the base; + # everything else replaces, so a gate that only needs a different substrate + # does not have to restate mounts and accounts that would then drift. + # + # workloads: + # gate: + # kind: slurm + # time_limit: "01:00:00" + + # --- how multi-rank actually launches ------------------------------------ + # Worth knowing, because the failure is loud and the cause is not. + # + # At world size > 1 the workflow emits one task per rank and wraps the + # command in `trtllm-llmapi-launch`, so the ranks EXIST before the engine + # starts. Without that, srun launches a single task, 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, carried from + # `.claude/skills/trtllm-test-script-builder/`. 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 + +# ---------------------------------------------------------------- accuracy +# The release criterion. One-sided: measured >= reference - tolerance. +# +# 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 + + # In benchmark points. Same-protocol session variance measured at 0.1-0.3 + # points, so the default 5.0 is wide relative to the noise by design: the + # gate catches an assembly defect, it does not resolve a fine difference. + tolerance: 5.0 + + # 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: {QWEN_PLE_OFFLOAD: "0"} # the leg WITHOUT the offload + # expected_mib: 48800 # 20 M rows x 2560 dims, fp8 + # disagg: + # contrast_env: {TRTLLM_DISAGG: "1"} # the leg WITH disagg + # # REQUIRED: leg B's own launcher (context + generation engine + proxy + + # # client). Run as ONE process on every device of the claim, handed leg + # # A's workload as STAIRCASE_WITNESS_* variables (see witness.py). + # contrast_command: /path/to/disagg_leg_b.py + # 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: + # TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY: "0" + # TRTLLM_DSV41_DECODER_SWA_BOUNDED_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 STAIRCASE_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 (`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 + # STAIRCASE_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: {TRTLLM_DSV41_KV_STORAGE: bf16} # the leg WITH bf16 storage + # expected_bytes_per_token: 890 # 3 x (288 + 68) / 2 + (288 + 68) + # 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: {TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY: "0"} # replay OFF + # window: 128 + # prompt_tokens: 4096 + # encoder_replay: + # contrast_env: {TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY: "0"} # replay OFF + # window: 128 + # prefix_tokens: 4096 + # reuse_alignment: 2 # the compress ratio; an odd boundary splits a pair + + # The ABSOLUTE `acceptance_length` floor a speculator unit falls back to + # when the `=off` leg cannot exist. The acceptance gate normally compares + # the target against the same checkpoint under BUILT-IN modeling — and a + # genuinely new architecture has none, so `=off` cannot load it at all. + # That is the same shape as the substrate gate's problem: the check is + # unavailable exactly when onboarding is doing its job. + # + # Skipping the gate is not an option. It is the ONLY gate that can see a + # miscomputed draft path, because rejection sampling makes a wrong draft + # produce correct text more slowly — boot passes, accuracy passes, and only + # speed moves. So declare a floor instead, and the run records the weaker + # claim verbatim: the draft path accepts at a real rate, NOT that it + # matches stock. + # + # Pick it from the ceiling, not from what the target happens to score. At + # `num_nextn_predict_layers: N` the ceiling is `N + 1`, for DSpark at + # `block_size: B` it is `B + 1`, and 1.0 means every draft was rejected — + # which the gate already fails on separately. + # + # The same block declares the gate's WORKLOAD. Without one it measures the + # boot prompts, and those are only a workload when their answers are long: + # a 2-token answer ends in EOS before any draft is verified, so the rate + # reads 1.00 for every drafter, working or broken, and no floor above 1 can + # pass. Declare prompts whose greedy answers run to `max_tokens`: a + # derivation, a code listing, a story. They are passed as-is, with no chat + # template applied, so write them the way the boot prompts are written. + # A workload that reaches fewer than `min_draft_iterations` generation + # iterations in total fails as "workload too short to measure acceptance" + # rather than as a low number. All four keys also work inside one unit's + # entry, for that unit only. + # + # acceptance: + # prompts: + # - "Explain, step by step, how a heat pump moves heat indoors." + # # prompts_file: prompts.json # instead of `prompts`: a YAML/JSON list, + # # # or a boot.yaml-shaped file (`prompts:`) + # max_tokens: 256 + # min_draft_iterations: 32 + # mtp: + # min_absolute: 1.5 + # dspark: + # min_absolute: 1.2 + + # `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 + +# ------------------------------------------------- catalog remediation budget +# When the delivery auditor returns UNCOVERED — runtime values sitting outside +# their entries' certified columns — the orchestrator certifies them itself, +# in a BOUNDED cycle: group the findings by (entry, arch, axis), dispatch one +# verify smith per surface, re-audit, repeat. +# +# Bounded is the operative word. The cycle stops on its own if the unresolved +# set stops shrinking, if a surface that was already certified comes back, or +# if any of these budgets is spent — each with a named verdict, so a run that +# cannot converge hands over instead of spending the night proving it. +# +# The defaults below are the policy and are deliberately small: remediation +# runs catalog GPU tests, not accuracy gates, but it is still GPU time, and a +# cycle that needs a third round usually had a first round that was measuring +# the wrong thing. There is no spelling for "unlimited". +# catalog_remediation: +# max_rounds: 2 # audit -> verify -> re-audit, at most twice +# max_attempts_per_item: 1 # a second attempt at identical work is not progress +# max_new_items: 16 # surfaces one round may add to the queue +# max_parallel_verifiers: 4 # concurrent verify smiths diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py b/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py new file mode 100644 index 000000000000..8840df865fb2 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py @@ -0,0 +1,1740 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Schema validation for staircase-onboard's ``task.yaml`` input. + +A staircase 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 difflib +import re +from pathlib import Path +from typing import Any, Mapping, Sequence + +import yaml + +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, and before this it also got the narrowest gating — one +# build, gated at ``--tp_size=N`` and nothing else, with the expert-parallel +# and attention-DP branches that only a unified codebase carries never +# executed by any gate. +# +# TEP and DEP at the world size are the baseline pair +# ``skills/staircase-modeling-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": "staircase-modeling-speculator", + "eagle3": "staircase-modeling-speculator", + "dflash": "staircase-modeling-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: DeepSeek-V4.1 ships + # 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": "staircase-modeling-speculator", + # DeepSeek-V4.1's "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": "staircase-modeling-bounded-replay", + "encoder_replay": "staircase-modeling-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": "staircase-modeling-kv-quant", + "eplb": "staircase-modeling-eplb", + "dwdp": "staircase-modeling-dwdp", + "multimodal": "staircase-modeling-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 + # the first tensor it was needed for (Qwen3.8-Flash-Next's 53.7 GB PLE + # n-gram table), because the mechanism is what generalizes. + "host_offload": "staircase-modeling-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": "staircase-modeling-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 STAIRCASE_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. +WITNESS_KEYS: tuple[str, ...] = ( + "contrast_env", + "expected_mib", + "contrast_command", + "timeout", + WITNESS_BATCH_SIZE_KEY, +) + +# Features gated by the EQUIVALENCE witness (token identity across a hop). +# +# contrast_command REQUIRED. Leg B's launcher: a context engine, a generation +# engine, a proxy and a client, which no environment variable +# over the aggregated command turns into. Without it leg B is +# leg A again and the witness compares a run to itself +# 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 +# STAIRCASE_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 no longer holds 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" +EQUIVALENCE_WITNESS_KEYS: tuple[str, ...] = ( + "contrast_env", + "contrast_command", + "prompt_tokens", + "timeout", + WITNESS_BATCH_SIZE_KEY, + WITNESS_LEGS_ENV_KEY, +) + +# 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 (128 for DeepSeek-V4.1) +# 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. DeepSeek-V4.1 +# compresses global KV in pairs (compress ratio 2), so 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 +# short 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 at most six bytes each is ~96 KiB. +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 (890 for DeepSeek-V4.1 NVFP4) +# 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, +) +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 measures the +# boot prompts, which is only a measurement when their answers are long: a +# 2-token answer ends in EOS before a drafter is asked anything, and the rate +# reads 1.00 for every drafter, working or not. +# +# 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 +# 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 +ACCEPTANCE_WORKLOAD_KEYS: tuple[str, ...] = ( + "prompts", + "prompts_file", + "max_tokens", + "min_draft_iterations", +) +# What one unit's entry in `accuracy.acceptance` may carry. +ACCEPTANCE_UNIT_KEYS: tuple[str, ...] = ("min_absolute", *ACCEPTANCE_WORKLOAD_KEYS) + +TARGET_FIELD = "target" +RUNTIME_FIELD = "runtime" +SCOPE_FIELD = "scope" +SOURCE_EMBARGO_FIELD = "source_embargo" +# Budgets for the bounded UNCOVERED remediation cycle. Optional: the +# defaults in `remediation.RemediationBudget` are the policy, and a spec +# only states one when this target is known to need a wider one. +REMEDIATION_FIELD = "catalog_remediation" + +# 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, +} + +# 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, + SOURCE_EMBARGO_FIELD, + REMEDIATION_FIELD, +) + +REMEDIATION_KEYS: tuple[str, ...] = ( + "max_rounds", + "max_attempts_per_item", + "max_new_items", + "max_parallel_verifiers", +) + +TARGET_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. They were + # validated below but missing from this list, which made every task.yaml + # that set them fail as "unknown key" once the block was closed. + "max_seq_len", + "max_batch_size", + "benchmark_args", +) + +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", + # `auto` (default) picks the MPI plugin by task count. Set it only when + # the cluster reports a versioned PMIx plugin -- `srun --mpi=list`. + "mpi", + # 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", + # Per-workload substrate overrides. See `_validate_workloads`. + "workloads", +) + +ENVIRONMENT_KINDS: tuple[str, ...] = ("local", "slurm", "slurm_broker", "session_spool") +ENVIRONMENT_REQUIRED_FIELDS: tuple[str, ...] = ("kind", "python") +SLURM_REQUIRED_FIELDS: tuple[str, ...] = ("partition", "container_image") + +# Default accuracy tolerance, in benchmark points. Staircase measured +# same-protocol session variance at 0.1-0.3 points against a 5.0-point +# tolerance, so 5.0 is wide relative to the noise by design: the gate is +# there to catch an assembly defect, not to resolve a fine difference. +DEFAULT_ACCURACY_TOLERANCE = 5.0 + + +class TaskSchemaError(ValueError): + """Raised when ``task.yaml`` fails staircase-onboard schema validation.""" + + +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, and the reason is the same one + ``runtime`` has always stated: 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. + + That was not a hypothetical. Before this check existed, ``feature:`` + instead of ``features:`` validated cleanly and resolved to ``[]`` — the + run onboarded core only, gated it green, and reported success for a + feature set it never built. ``runtme:`` instead of ``runtime:`` dropped + ``max_seq_len`` and reproduced the ``Impossible to fit in any sequence + in kvCache`` failure that had already cost most of a run's wall clock. + + 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)}" + ) + + +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.StaircaseOnboardWorkflow._assemble_one_unit` + never branches on whether the run is unified, and no caller can forget + the unified case the way :func:`~.gates.topology_flags` did. + + ``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 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_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) + for name in TARGET_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) - {"extra_flags"}) + if unknown: + errors.append( + f"'{where}' has unknown switch(es) {unknown}; known switches are " + f"{sorted(RUNTIME_SWITCHES)} plus 'extra_flags'" + ) + + 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` were + # gated with block reuse ON and passed, while the multimodal path requires + # it OFF -- and 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]) -> 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 + reintroduce the TP-only gate under a new name — which is precisely the + defect this key exists to close. + - **One world size.** One allocation serves every member, so a member + asking for a different device count cannot be run by this run at all. + The mismatch is reported with both values rather than quietly re-sized + to the allocation, because a member silently served at half its + declared width is a measurement of something nobody asked for. + """ + 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: + 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"One allocation serves every member, so they must all declare the same " + f"world size." + ) + + +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``. This is the + case the field was written for: TensorRT-LLM already ships an + implementation of this checkpoint, so name it and the writing roles + cannot open 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. The + argument that made this a mechanism rather than a discipline does not + depend on which repository the file lives in; it depends on the file + being reachable. + + An absolute path that resolves *inside* the repo is still rejected: that + is the case the original error was actually about, and the fix is to + 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, or " + f"absolute for a tree outside the repo" + ) + if not isinstance(repo, str): + return + repo_root = Path(repo).resolve() + for entry in block: + 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_remediation(block: Any, errors: list[str]) -> None: + """Validate the ``catalog_remediation`` budgets. + + Every field is a positive integer, and stating one is optional -- the + defaults in :class:`remediation.RemediationBudget` are the policy. A + zero is rejected rather than read as "unlimited": an unbounded + remediation cycle is the thing this block exists to prevent, so the + spelling that means it should not be the easiest one to type by + accident. + """ + if block is None: + return + if not isinstance(block, dict): + errors.append(f"'{REMEDIATION_FIELD}' must be a mapping, got {type(block).__name__}") + return + reject_unknown_keys(block, REMEDIATION_KEYS, errors, where=REMEDIATION_FIELD) + for name in REMEDIATION_KEYS: + value = block.get(name) + if value is None: + continue + if not isinstance(value, int) or isinstance(value, bool) or value < 1: + errors.append( + f"'{REMEDIATION_FIELD}.{name}' must be an integer >= 1, got {value!r}. " + f"There is no spelling for 'unlimited' on purpose." + ) + + +def _validate_workloads(block: Any, errors: list[str]) -> None: + """Per-workload launch substrate: ``environment.workloads.``. + + **The right substrate depends on how long the job runs.** A shared session + imports a 26 GB container once and costs nothing per job afterwards, which + is correct for a 1-minute entry test -- one onboarding ran 443 of them. It + is wrong for a 10-15 minute gate, where the import is 30-50% overhead and + isolation is the point: a shared session has **one failure domain**, so a + gate that leaks a server, kills a process group or collides on PMIx takes + every other job and the orchestrator with it. All three happened in the run + this was written for, and two more session deaths followed in the next one + -- see :func:`~.environment.shares_allocation_with_spool`. + + The workload name is checked against + :data:`~.environment.WORKLOAD_KINDS` rather than trusted, because an + override that never applies is **invisible**: the run simply uses the base + environment and nothing says why. That is the same class of silent wrong + :func:`reject_unknown_keys` exists for, one level down. + """ + if block is None: + return + if not isinstance(block, dict): + errors.append(f"'{ENVIRONMENT_FIELD}.workloads' must be a mapping of workload to block") + return + # Imported here rather than at module scope: `environment` is the consumer + # of this schema, so a top-level import would make the two modules circular. + from .environment import WORKLOAD_KINDS + + for name, override in block.items(): + where = f"{ENVIRONMENT_FIELD}.workloads.{name}" + if name not in WORKLOAD_KINDS: + errors.append( + f"'{where}' is not a workload the orchestrator dispatches under; " + f"expected one of {WORKLOAD_KINDS}" + ) + continue + if not isinstance(override, dict): + errors.append(f"'{where}' must be a mapping") + continue + # `workloads` itself is excluded from the legal keys: an override that + # carries its own overrides is a nesting nothing resolves. + reject_unknown_keys( + override, tuple(k for k in ENVIRONMENT_KEYS if k != "workloads"), errors, where=where + ) + kind = override.get("kind") + if kind is not None and kind not in ENVIRONMENT_KINDS: + errors.append(f"'{where}.kind' must be one of {ENVIRONMENT_KINDS}; got {kind!r}") + env = override.get("env") + if env is not None and not isinstance(env, dict): + errors.append(f"'{where}.env' must be a mapping") + + +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_workloads(block.get("workloads"), errors) + + 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__}" + ) + + # 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." + ) + + if kind == "slurm_broker": + # The broker is an MCP service, so only a Claude Code process can + # reach it — the orchestrator cannot. `account` is mandatory here + # and not for plain `slurm` because the broker rejects every + # submit that omits it, which costs a round trip to discover. + for name in ("partition", "account", "container_image"): + value = block.get(name) + if not isinstance(value, str) or not value.strip(): + errors.append( + f"'{ENVIRONMENT_FIELD}.{name}' is required when kind is " + f"'slurm_broker' and must be a non-empty string" + ) + gpus = block.get("gpus_per_node") + if gpus is None: + errors.append( + f"'{ENVIRONMENT_FIELD}.gpus_per_node' is required when kind is " + f"'slurm_broker'. NVL72 clusters reject job-scoped --gpus, so the " + f"per-node count is the only form that submits." + ) + elif not isinstance(gpus, int) or isinstance(gpus, bool) or gpus < 1: + errors.append(f"'{ENVIRONMENT_FIELD}.gpus_per_node' must be an integer >= 1") + mounts = block.get("mounts") + if mounts is None: + block["mounts"] = [] + elif 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") + if not isinstance(block.get("time_limit"), str): + errors.append( + f"'{ENVIRONMENT_FIELD}.time_limit' is required when kind is " + f"'slurm_broker' (HH:MM:SS); the partition's own maximum is not a default " + f"the broker will supply." + ) + + if kind == "slurm": + for name in SLURM_REQUIRED_FIELDS: + value = block.get(name) + if not isinstance(value, str) or not value.strip(): + errors.append( + f"'{ENVIRONMENT_FIELD}.{name}' is required when kind is 'slurm' " + f"and must be a non-empty string" + ) + mounts = block.get("mounts") + if mounts is None: + block["mounts"] = [] + elif 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") + + +def _positive_int(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value >= 1 + + +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_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 must be its own launcher.""" + reject_unknown_keys(entry, EQUIVALENCE_WITNESS_KEYS, errors, where=where) + _validate_contrast_env(entry, errors, where) + _validate_legs_env(entry, errors, where) + command = entry.get("contrast_command") + if not isinstance(command, str) or not command.strip(): + errors.append( + f"'{where}.contrast_command' is required and must be a non-empty string: leg B " + f"is a context engine, a generation engine and a proxy, which no environment " + f"variable over the aggregated command becomes, so without it the witness " + f"compares a run to itself" + ) + 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) + + +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 (890 for " + f"DeepSeek-V4.1 NVFP4); 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) + + +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}") + for key in ("max_tokens", "min_draft_iterations"): + 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(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) + for unit, entry in block.items(): + if unit in ACCEPTANCE_WORKLOAD_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" + 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_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) + + 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) + 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) + predicted = entry.get("expected_mib") + if predicted is not None and ( + not isinstance(predicted, (int, float)) + or isinstance(predicted, bool) + or predicted <= 0 + ): + errors.append(f"'{where}.expected_mib' must be a positive number") + + # 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 a + # 27-point accuracy regression, and that is measured, not hypothetical. + 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 _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.*.boot_spec`` and + ``accuracy.acceptance[.].prompts_file`` point at files the author + writes alongside ``task.yaml`` — the per-unit runtime YAMLs, prompt + sets. Written absolute, as they had to be, **the spec stops working on any + other machine**, which is the opposite of what a spec is for. One run + produced a 220-line task.yaml with six absolute paths in it; the most + reusable artifact of that run was unusable anywhere else. This run's own + ``mimo-v2.6-flash-onboard/task.yaml`` repeats it: every one of its five + per-unit ``llm_api_options`` entries names a file that sits in a directory + beside the spec, spelled from the filesystem root. + + Resolved against the **spec's own directory**, because that is where the + author put them, and written back as absolute so every consumer downstream + is unchanged — the gates take a string and pass it to ``trtllm-eval``, and + none of them learns anything new. + + 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("boot_spec") is not None: + block["boot_spec"] = absolute(block["boot_spec"]) + + # The acceptance workload's prompt file, at the block level and per unit. + 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"]) + + +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) + + errors: list[str] = [] + + # First, before anything reads a field: 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. + _validate_remediation(data.get(REMEDIATION_FIELD), errors) + reject_unknown_keys(data, TOP_LEVEL_KEYS, errors) + + 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: + _validate_members(data[MEMBERS_FIELD], data.get(PARALLEL_FIELD), errors) + + 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. A DeepSeek-V4.1 `mtp.*` is DSpark, " + f"so its 'MTP' request is spelled `dspark`." + ) + + 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)) + + 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 staircase-onboard schema validation:{bullet}{bullet.join(errors)}" + ) + + return data + + +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__ = [ + "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", + "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", + "PARALLEL_FIELD", + "PARALLEL_UNIFIED", + "REQUIRED_PATH_FIELDS", + "RUNTIME_FIELD", + "RUNTIME_SWITCHES", + "SCOPES", + "SCOPE_FIELD", + "SLURM_REQUIRED_FIELDS", + "reject_unknown_keys", + "SOURCE_EMBARGO_FIELD", + "TARGET_FIELD", + "TARGET_KEYS", + "TOP_LEVEL_KEYS", + "TaskSchemaError", + "assembly_units", + "derive_checkpoint_segment", + "derive_family", + "derive_slug", + "load_and_validate_task_yaml", + "load_prompts_file", + "members_for", + "parse_parallel_segment", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/transport.py b/agent-flow/agent_flow/workflows/staircase_onboard/transport.py new file mode 100644 index 000000000000..605b63afd387 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/transport.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Staircase's use of the file-dispatch turn transport. + +The protocol — request/system/prompt/result/done, written atomically, `.done` +last, a monotonic id per turn, a heartbeat that fails a dead driver fast — +is not Staircase's. It solves "this host cannot authenticate an agent CLI", +which is a property of the host and not of onboarding a checkpoint, so it +lives in :mod:`agent_flow.runtime.turns.file_dispatch` and any workflow can +use it. + +What IS Staircase's is one paragraph: every role prompt in this workflow ends +by telling the agent to call ``append__progress``, and that tool cannot +exist here. So this module supplies that note and nothing else. + +The names below are kept because the CLI, the workflow and the tests all +reach for them, and because the thing they name has not changed. +""" + +from __future__ import annotations + +from pathlib import Path + +from agent_flow.runtime.turns import ( + DEFAULT_TURN_TIMEOUT_S, + DISPATCH_DIRNAME, + HEARTBEAT_STALE_S, + FileDispatchTurnExecutor, + TurnTimeout, +) + +#: The historical name, kept because it is what the workflow, the CLI and +#: their tests raise and catch. A `TurnTimeout` from any transport is the +#: same event, so this is an alias rather than a parallel hierarchy. +ExternalTransportTimeout = TurnTimeout + + +def staircase_transport_note(role: str) -> str: + """Say which of the role prompt's instructions this transport cannot honour. + + Every role prompt ends by telling the agent to call + ``append__progress``. That tool is real — the in-process transport + hands it to the session it spawns — but **this** transport does not, + because there is no session to hand it to. The shared + `decision_file_block` that follows says to write the decision file + instead, so the prompt arrives carrying two different endings and only + the second one works here. + + Three auditor turns in a row noticed the contradiction and picked + correctly, and all three reported it. That is a judgement none of them + should have had to make at the end of a two-hour turn, and the next role + to face it might pick the other one: calling a missing tool looks from + the inside like a turn that ended, so the run would block for its full + timeout on a turn that believed it was finished. + + Correcting the role prompts would be the wrong repair — the instruction + is right for the transport they were written for. The difference belongs + to the transport, so the transport states it. + """ + return ( + "\n\n---\n\n" + "## This turn runs on the external transport\n\n" + f"Your role prompt above tells you to call `append_{role}_progress`. " + "**That tool is not available here** — it belongs to the in-process " + "transport, which spawns a session to hand it to. This transport has no " + "session: it writes your turn to disk and blocks on your answer.\n\n" + "So ignore that instruction and follow the decision-file instruction " + "that comes after it. Writing the decision file with `Write`, then " + "`.result.md`, then `.done`, is the **only** way a verdict " + "reaches the orchestrator from here. Everything else in your role " + "prompt applies unchanged.\n" + ) + + +class ExternalAgentTransport(FileDispatchTurnExecutor): + """The file-dispatch transport, carrying Staircase's role-prompt note. + + Call it exactly where an ``AgentLayer`` would be called:: + + result = transport(role="smith", what="onboard foo", prompt=text) + + or, from the orchestrator's transport-independent path, through + ``run(TurnRequest(...))``. Both are the same turn. + """ + + def __init__( + self, + workspace: Path | str, + *, + poll_s: float = 5.0, + timeout_s: float = DEFAULT_TURN_TIMEOUT_S, + ) -> None: + super().__init__( + workspace, + poll_s=poll_s, + timeout_s=timeout_s, + transport_note=staircase_transport_note, + ) + + +__all__ = [ + "DEFAULT_TURN_TIMEOUT_S", + "DISPATCH_DIRNAME", + "ExternalAgentTransport", + "ExternalTransportTimeout", + "HEARTBEAT_STALE_S", + "staircase_transport_note", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/vocabulary.py b/agent-flow/agent_flow/workflows/staircase_onboard/vocabulary.py new file mode 100644 index 000000000000..9c77a63e1f09 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/vocabulary.py @@ -0,0 +1,598 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``vocabulary.yaml`` — the plan phase's machine-readable product. + +Staircase's kernel-feasibility step answers one question: *can this +checkpoint be assembled from the catalog as it stands, and if not, what has +to be added first?* Out of tree its answer was prose, because the thing +reading it was also a language model. Here the orchestrator reads it, so it +is a schema — and that is the single most valuable typing change in the +port. A dispatch queue the orchestrator can walk is a queue it can +parallelize, resume, and audit; a prose queue is none of those. + +Each computation step lands in one of three states: + +=========== ========================================================== +``OK`` an entry exists and its certified column covers this + checkpoint's values — nothing to do +``CERT`` an entry exists, but this surface is not covered — dispatch + smith with ``kernel-verify`` +``GAP`` no entry at all — dispatch smith with ``kernel-onboard`` +=========== ========================================================== + +**The judging happens in the plan phase, not inside smith.** Deciding +whether a value falls inside an entry's certified column is *reading a +contract*, and contracts are written to be sufficient on their own. What +needs smith is a step with no entry, because only reading the engine's +source tells you what that callable is and whether it can be one atomic +entry. + +A ``GAP`` row says *where* its implementation comes from, via the sourcing +ladder (``TIER_*``): an engine op, an in-tree callable, native torch this +run composes, or a Triton kernel this run writes. The rungs are ordered by +cost, and a row that claims a writing tier owes an account of why the +cheaper ones were rejected. + +Only when **every** rung is exhausted is the verdict a **vocabulary +ceiling**, and that claim now has to carry its evidence per rung. The +distinction matters for any model the engine does not already implement: +"no ``torch.ops.trtllm`` candidate exists" is a fact about tier 1, not a +reason to stop. That verdict still ends the run in Python, with the +evidence — no agent gets to decide it is probably fine. + +A queue row may also name the assembly ``unit`` that first needs its entry +(``core`` when absent). Every row is still dispatched up front, because a +feature-only entry costs nothing to build alongside core's; the tag decides +only *what a failure blocks*. See :func:`needed_by`. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field, fields +from pathlib import Path +from typing import Any, Mapping, Sequence + +import yaml + +STATE_OK = "OK" +STATE_CERT = "CERT" +STATE_GAP = "GAP" +_VALID_STATES = (STATE_OK, STATE_CERT, STATE_GAP) + +MODE_ONBOARD = "onboard" +MODE_VERIFY = "verify" +_VALID_MODES = (MODE_ONBOARD, MODE_VERIFY) + +# Where a GAP's implementation is allowed to come from, cheapest first. +# +# Before this existed, a GAP either named a `torch.ops.trtllm` candidate or +# it was a **ceiling** — and a ceiling ends the run. That binary is wrong for +# any model whose layers the engine does not already implement, which is most +# of the interesting ones: the first real onboarding hit it on its largest +# item and needed an ad-hoc ruling mid-run to get past tier 2. +# +# The ladder makes "we have not looked hard enough yet" different from "there +# is nothing there", and it makes the difference auditable: a ceiling now has +# to say what it tried at every rung. +TIER_TRTLLM_OP = 1 # torch.ops.trtllm.* -- an op the engine already exports +TIER_INTREE_CALLABLE = 2 # tensorrt_llm/_torch/modules/** -- the kernel plane +TIER_NATIVE_TORCH = 3 # composed from torch ops, written by this run +TIER_TRITON = 4 # a Triton kernel, written by this run +_VALID_TIERS = (TIER_TRTLLM_OP, TIER_INTREE_CALLABLE, TIER_NATIVE_TORCH, TIER_TRITON) + +#: What each rung means to the agent that has to build it. +TIER_DESCRIPTION: dict[int, str] = { + TIER_TRTLLM_OP: "an op the engine already exports as torch.ops.trtllm.*", + TIER_INTREE_CALLABLE: ( + "a callable already in this repository under tensorrt_llm/_torch/modules/**; " + "the kernel plane is reused, not rewritten" + ), + TIER_NATIVE_TORCH: ( + "no engine implementation exists, so compose one from torch ops. Correct " + "first and slow is acceptable for a functionality onboarding; the tuning " + "phase replaces it" + ), + TIER_TRITON: ( + "native torch cannot express it, or does so at a cost that breaks the " + "gate's budget, so write a Triton kernel" + ), +} + +#: A tier above this is work this run *writes* rather than *wraps*, so it +#: costs an implementation as well as a contract and a test. +HIGHEST_REUSING_TIER = TIER_INTREE_CALLABLE + +VERDICT_ASSEMBLABLE = "assemblable" +VERDICT_AFTER_ADDITIONS = "assemblable_after_additions" +VERDICT_CEILING = "ceiling" +_VALID_VERDICTS = (VERDICT_ASSEMBLABLE, VERDICT_AFTER_ADDITIONS, VERDICT_CEILING) + +# The state a step is in decides which smith mode serves it. ``OK`` needs +# no dispatch at all, which is why it is absent here. +_MODE_BY_STATE = {STATE_CERT: MODE_VERIFY, STATE_GAP: MODE_ONBOARD} + +#: The unit a queue row belongs to when it names none. Core is the safe +#: default because it is the blocking one: an untagged row that fails stops +#: the run before assembly, exactly as every row did before units existed. +UNIT_CORE = "core" + + +class VocabularyError(ValueError): + """``vocabulary.yaml`` is malformed or internally inconsistent.""" + + +@dataclass(frozen=True) +class QueueItem: + """One smith dispatch: exactly one catalog entry, in one mode.""" + + n: int + mode: str + target: str + depends_on: tuple[int, ...] = () + surface: str | None = None + entry: str | None = None + note: str = "" + #: Which rung of the sourcing ladder supplies this entry. Defaults to + #: tier 1 so an existing plan keeps its meaning. + tier: int = TIER_TRTLLM_OP + #: The first assembly unit that needs this entry. It decides what a + #: failed dispatch blocks, not when the dispatch runs: every row is + #: dispatched in the smith phase, and only a failure that core needs + #: stops the run there. A feature-only failure is deferred to the start + #: of its unit and retried once. Measured on the DeepSeek-V4.1 run: 33 + #: dispatches queued, six of them needed only by dspark / multimodal / + #: host_offload / kv_quant, and any one of those failing closed the run + #: `smith-blocked` before core -- which needed none of them -- had begun. + unit: str = UNIT_CORE + + @classmethod + def from_row(cls, row: Mapping[str, Any]) -> "QueueItem": + """Rebuild an item from a persisted queue row, ignoring its bookkeeping. + + A row in the checkpoint also carries the orchestrator's ``status`` + and, once deferred, where to and why. Those describe the run, not + the dispatch, so they are dropped here rather than passed to a + constructor that does not know them. A row written before ``unit`` + existed reads as ``core``, which is what it meant at the time. + """ + known = {f.name for f in fields(cls)} + kwargs = {key: value for key, value in row.items() if key in known} + kwargs["n"] = int(kwargs["n"]) + kwargs.setdefault("mode", MODE_ONBOARD) + kwargs.setdefault("target", "") + kwargs["depends_on"] = tuple(int(dep) for dep in kwargs.get("depends_on") or ()) + kwargs["unit"] = str(kwargs.get("unit") or UNIT_CORE) + return cls(**kwargs) + + @property + def skill(self) -> str: + """The skill this dispatch tells smith to load.""" + return ( + "staircase-kernel-onboard" if self.mode == MODE_ONBOARD else "staircase-kernel-verify" + ) + + @property + def writes_implementation(self) -> bool: + """True when this dispatch must *write* the kernel, not wrap one. + + Tiers 1-2 wrap something that exists; 3-4 do not. The difference is + not cosmetic: a writing dispatch owes an implementation as well as a + contract and a test, is far more likely to need more than one turn, + and should not be batched against a wrapping dispatch as though the + two cost the same. + """ + return self.tier > HIGHEST_REUSING_TIER + + @property + def tier_brief(self) -> str: + """One line naming the rung, for the dispatch prompt.""" + return f"tier {self.tier} — {TIER_DESCRIPTION.get(self.tier, 'unknown')}" + + def as_dict(self) -> dict[str, Any]: + return { + "n": self.n, + "mode": self.mode, + "target": self.target, + "depends_on": list(self.depends_on), + "surface": self.surface, + "entry": self.entry, + "note": self.note, + "tier": self.tier, + "unit": self.unit, + } + + +@dataclass(frozen=True) +class Vocabulary: + """A parsed, validated ``vocabulary.yaml``.""" + + verdict: str + steps: tuple[dict[str, Any], ...] = () + queue: tuple[QueueItem, ...] = () + ceiling: dict[str, Any] | None = None + notes: list[str] = field(default_factory=list) + + @property + def is_ceiling(self) -> bool: + return self.verdict == VERDICT_CEILING + + def counts(self) -> dict[str, int]: + """How many steps landed in each state — the plan's headline number.""" + tally = {state: 0 for state in _VALID_STATES} + for step in self.steps: + tally[str(step.get("state"))] = tally.get(str(step.get("state")), 0) + 1 + return tally + + +def _require_mapping(value: Any, what: str) -> dict[str, Any]: + if not isinstance(value, dict): + raise VocabularyError(f"{what} must be a mapping, got {type(value).__name__}") + return value + + +def _parse_steps(raw: Any, errors: list[str]) -> tuple[dict[str, Any], ...]: + if raw is None: + return () + if not isinstance(raw, list): + errors.append(f"`steps` must be a list, got {type(raw).__name__}") + return () + steps: list[dict[str, Any]] = [] + for index, row in enumerate(raw): + if not isinstance(row, dict): + errors.append(f"steps[{index}] must be a mapping, got {type(row).__name__}") + continue + state = row.get("state") + if state not in _VALID_STATES: + errors.append(f"steps[{index}].state is {state!r}; expected one of {_VALID_STATES}") + if not str(row.get("step", "")).strip(): + errors.append(f"steps[{index}].step must name the computation step") + if state in (STATE_OK, STATE_CERT): + if not str(row.get("entry", "")).strip(): + errors.append(f"steps[{index}] is {state} but names no catalog `entry`") + # The contract line is what makes an OK/CERT judgement + # auditable later. Without it the plan asserts coverage and + # nobody can check the assertion without redoing the reading. + if not str(row.get("contract_line", "")).strip(): + errors.append( + f"steps[{index}] is {state} but quotes no `contract_line`; the line " + f"that decides coverage is the evidence for this row" + ) + elif state == STATE_GAP: + tier = row.get("tier") + if tier is None: + errors.append( + f"steps[{index}] is GAP but names no `tier`. Say which rung of the " + f"sourcing ladder supplies it: {_VALID_TIERS} — see TIER_DESCRIPTION. " + f"'No trtllm op exists' is a statement about tier 1, not a ceiling." + ) + elif tier not in _VALID_TIERS: + errors.append(f"steps[{index}].tier is {tier!r}; expected one of {_VALID_TIERS}") + # Tiers 1-2 *wrap* something that already exists, so the row has + # to name it. Tiers 3-4 *write* it, so there is nothing to name + # yet -- but then the row owes an explanation of why the cheaper + # rungs were rejected, or the ladder becomes a way to skip to + # writing a kernel without looking for one first. + if tier in (TIER_TRTLLM_OP, TIER_INTREE_CALLABLE): + if not str(row.get("candidate", "")).strip(): + errors.append( + f"steps[{index}] is GAP at tier {tier} but names no `candidate`; " + f"tiers 1-2 wrap something that already exists, so it must be named" + ) + elif tier in (TIER_NATIVE_TORCH, TIER_TRITON): + if not str(row.get("tried", "")).strip(): + errors.append( + f"steps[{index}] is GAP at tier {tier}, which means this run has " + f"to *write* the implementation. Record in `tried` what was " + f"searched at the cheaper rungs and why it was rejected — " + f"otherwise the ladder is a way to skip straight to writing a " + f"kernel without looking for one" + ) + steps.append(dict(row)) + return tuple(steps) + + +def _parse_unit( + row: Mapping[str, Any], index: int, units: Sequence[str] | None, errors: list[str] +) -> str: + """The row's ``unit``, defaulting to core; checked against the run's units when known.""" + raw = row.get("unit") + if raw is None: + return UNIT_CORE + if not isinstance(raw, str) or not raw.strip(): + errors.append( + f"queue[{index}].unit is {raw!r}; it must name the one assembly unit that " + f"first needs this entry (`{UNIT_CORE}`, the default, or a feature from " + f"task.yaml `features`)" + ) + return UNIT_CORE + unit = raw.strip() + if units is not None and unit not in units: + # A misspelt unit is worse than none. An absent tag reads as core and + # blocks as it always did; an unknown one could only be guessed at, + # and guessing "later" would let a failure core needs slip past. + errors.append( + f"queue[{index}].unit is {unit!r}, which is not one of this run's assembly " + f"units {list(units)}. Tag a row only with the first unit that needs its " + f"entry; leave it off (core) when core needs it or you are unsure" + ) + return UNIT_CORE + return unit + + +def _parse_queue( + raw: Any, errors: list[str], units: Sequence[str] | None = None +) -> tuple[QueueItem, ...]: + if raw is None: + return () + if not isinstance(raw, list): + errors.append(f"`queue` must be a list, got {type(raw).__name__}") + return () + + items: list[QueueItem] = [] + seen_n: set[int] = set() + seen_targets: dict[str, str] = {} + for index, row in enumerate(raw): + if not isinstance(row, dict): + errors.append(f"queue[{index}] must be a mapping, got {type(row).__name__}") + continue + try: + n = int(row["n"]) + except (KeyError, TypeError, ValueError): + errors.append(f"queue[{index}].n must be an integer dispatch number") + continue + if n in seen_n: + errors.append(f"queue[{index}].n = {n} is a duplicate dispatch number") + seen_n.add(n) + + mode = row.get("mode") + if mode not in _VALID_MODES: + errors.append(f"queue[{index}].mode is {mode!r}; expected one of {_VALID_MODES}") + + target = str(row.get("target", "")).strip() + if not target: + errors.append(f"queue[{index}].target must name exactly one catalog entry") + elif target in seen_targets and seen_targets[target] != mode: + # Staircase states it plainly: the two chapters do not overlap. + # An entry smith onboards is tested at this checkpoint's values + # by construction, so it never comes back for verification. + errors.append( + f"queue[{index}] has {target!r} in both chapters " + f"({seen_targets[target]} and {mode}). An entry onboarded here is " + f"tested at this checkpoint's values by construction, so it never " + f"also needs verifying" + ) + elif target: + seen_targets[target] = str(mode) + + depends_raw = row.get("depends_on") or [] + if not isinstance(depends_raw, list): + errors.append(f"queue[{index}].depends_on must be a list of dispatch numbers") + depends_raw = [] + depends: list[int] = [] + for dep in depends_raw: + try: + depends.append(int(dep)) + except (TypeError, ValueError): + errors.append(f"queue[{index}].depends_on holds {dep!r}, which is not a number") + + if mode == MODE_VERIFY and not str(row.get("surface", "")).strip(): + errors.append( + f"queue[{index}] is a verify dispatch with no `surface`. The value and " + f"the contract line that excludes it are what smith is dispatched " + f"with, so vagueness here costs a whole run" + ) + + items.append( + QueueItem( + n=n, + mode=str(mode), + target=target, + depends_on=tuple(depends), + tier=int(row.get("tier", TIER_TRTLLM_OP) or TIER_TRTLLM_OP), + surface=row.get("surface"), + entry=row.get("entry"), + note=str(row.get("note", "")), + unit=_parse_unit(row, index, units, errors), + ) + ) + + known = {item.n for item in items} + for item in items: + for dep in item.depends_on: + if dep not in known: + errors.append(f"queue item {item.n} depends on {dep}, which is not in the queue") + return tuple(items) + + +def parse_vocabulary( + data: Mapping[str, Any], + source: str = "vocabulary.yaml", + *, + units: Sequence[str] | None = None, +) -> Vocabulary: + """Validate a parsed mapping, batching every problem into one error. + + ``units`` is the run's assembly units (``task_schema.assembly_units``). + When given, every queue row's ``unit`` must be one of them; when not, + only its shape is checked and the orchestrator checks membership where + it knows the spec. + """ + data = _require_mapping(data, source) + errors: list[str] = [] + + verdict = data.get("verdict") + if verdict not in _VALID_VERDICTS: + errors.append(f"`verdict` is {verdict!r}; expected one of {_VALID_VERDICTS}") + + steps = _parse_steps(data.get("steps"), errors) + queue = _parse_queue(data.get("queue"), errors, units) + + ceiling = data.get("ceiling") + if verdict == VERDICT_CEILING and isinstance(ceiling, dict): + # A ceiling ends the run, so it has to be the end of a search rather + # than the start of one. Every rung must be accounted for: the claim + # is "nothing anywhere serves this step", and that is only credible + # if writing it was considered and rejected for a stated reason. + exhausted = ceiling.get("exhausted") + if not isinstance(exhausted, dict): + errors.append( + "`ceiling` must carry an `exhausted` mapping with one entry per tier " + f"{_VALID_TIERS}, saying what was searched and why that rung failed. " + "A ceiling that has not considered tiers 3-4 is a report that nobody " + "wrote the kernel yet, which is not the same thing." + ) + else: + missing = [t for t in _VALID_TIERS if not str(exhausted.get(t, "")).strip()] + if missing: + errors.append( + f"`ceiling.exhausted` says nothing about tier(s) {missing}. " + f"Tiers 3-4 are this run writing the implementation itself; " + f"declaring a ceiling without ruling them out claims more than " + f"was established." + ) + + if verdict == VERDICT_CEILING: + if not isinstance(ceiling, dict) or not str(ceiling.get("step", "")).strip(): + errors.append( + "`verdict: ceiling` requires a `ceiling` mapping naming the exact " + "`step` and what was `searched`. A ceiling is the one verdict that " + "ends the run, so it carries its evidence or it is not a verdict" + ) + elif ceiling not in (None, {}): + errors.append(f"`ceiling` is set but the verdict is {verdict!r}; they must agree") + + if verdict == VERDICT_ASSEMBLABLE and queue: + errors.append( + f"`verdict: assemblable` means nothing has to be added first, but the " + f"queue holds {len(queue)} dispatch(es). Use " + f"`{VERDICT_AFTER_ADDITIONS}` instead" + ) + if verdict == VERDICT_AFTER_ADDITIONS and not queue: + errors.append( + f"`verdict: {VERDICT_AFTER_ADDITIONS}` with an empty queue says additions " + f"are needed but names none. Use `{VERDICT_ASSEMBLABLE}` instead" + ) + + if errors: + bullet = "\n - " + raise VocabularyError(f"{source} is not usable:{bullet}{bullet.join(errors)}") + + return Vocabulary( + verdict=str(verdict), + steps=steps, + queue=queue, + ceiling=ceiling if isinstance(ceiling, dict) else None, + ) + + +def load_vocabulary(path: str | Path, *, units: Sequence[str] | None = None) -> Vocabulary: + """Read and validate ``vocabulary.yaml`` from disk. See :func:`parse_vocabulary`.""" + source = Path(path) + if not source.is_file(): + raise VocabularyError(f"{source} does not exist; the plan phase must produce it.") + text = source.read_text(encoding="utf-8") + if not text.strip(): + raise VocabularyError(f"{source} is empty; the plan phase must produce it.") + try: + data = yaml.safe_load(text) + except yaml.YAMLError as exc: + raise VocabularyError(f"{source} is not valid YAML: {exc}") from exc + return parse_vocabulary(data or {}, str(source), units=units) + + +def dependency_batches(queue: Sequence[QueueItem]) -> list[list[QueueItem]]: + """Layer the queue so every batch can be dispatched concurrently. + + Catalog entries are independent by construction — they may not import + each other — so the only ordering that exists is the one the plan + declares via ``depends_on`` (typically "this verify row's test state + has to be built by that onboard row first"). Everything inside one + layer runs at once, which is where a run spends extra compute to buy + wall-clock. + + Raises :class:`VocabularyError` on a cycle, naming the items involved: + a cycle means the plan asserted two entries each need the other first, + which is a planning mistake rather than something to break arbitrarily. + """ + remaining = {item.n: item for item in queue} + unresolved = {item.n: set(item.depends_on) for item in queue} + batches: list[list[QueueItem]] = [] + + while remaining: + ready = sorted(n for n, deps in unresolved.items() if not deps) + if not ready: + stuck = ", ".join(f"{n} -> {sorted(deps)}" for n, deps in sorted(unresolved.items())) + raise VocabularyError( + f"the dispatch queue has a dependency cycle and cannot be ordered: {stuck}" + ) + batches.append([remaining[n] for n in ready]) + for n in ready: + remaining.pop(n) + unresolved.pop(n) + for deps in unresolved.values(): + deps.difference_update(ready) + + return batches + + +def needed_by(queue: Sequence[QueueItem], units: Sequence[str]) -> dict[int, str]: + """The earliest assembly unit that needs each row, directly or through a dependent. + + A row's own ``unit`` is where the plan says it is first used; but a row + another row ``depends_on`` is needed no later than that dependent is, and + transitively so. A core row's dependency is therefore core-relevant + whatever its own tag says, and deferring it would let core start without + something core's own entries were built on. + + A unit outside ``units`` reads as core, the blocking direction: an + unrecognised tag must never be what lets a failure slip past assembly. + """ + units = list(units) or [UNIT_CORE] + order = {unit: position for position, unit in enumerate(units)} + earliest = {item.n: order.get(item.unit, 0) for item in queue} + # Lowering only ever decreases a position, and positions are bounded + # below by 0, so this terminates -- including on a cyclic queue, which + # `dependency_batches` rejects separately. + changed = True + while changed: + changed = False + for item in queue: + for dep in item.depends_on: + if dep in earliest and earliest[dep] > earliest[item.n]: + earliest[dep] = earliest[item.n] + changed = True + return {n: units[position] for n, position in earliest.items()} + + +def dump_vocabulary(vocabulary: Vocabulary) -> str: + """Serialize back to YAML, for tests and for presets.""" + payload: dict[str, Any] = { + "verdict": vocabulary.verdict, + "steps": [dict(step) for step in vocabulary.steps], + "queue": [item.as_dict() for item in vocabulary.queue], + } + if vocabulary.ceiling is not None: + payload["ceiling"] = dict(vocabulary.ceiling) + return yaml.safe_dump(payload, sort_keys=False, allow_unicode=True, default_flow_style=False) + + +__all__ = [ + "MODE_ONBOARD", + "MODE_VERIFY", + "QueueItem", + "STATE_CERT", + "STATE_GAP", + "STATE_OK", + "UNIT_CORE", + "VERDICT_AFTER_ADDITIONS", + "VERDICT_ASSEMBLABLE", + "VERDICT_CEILING", + "Vocabulary", + "VocabularyError", + "dependency_batches", + "dump_vocabulary", + "load_vocabulary", + "needed_by", + "parse_vocabulary", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/witness.py b/agent-flow/agent_flow/workflows/staircase_onboard/witness.py new file mode 100644 index 000000000000..f2b02aee67b4 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/witness.py @@ -0,0 +1,1438 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Gates for features the boot and accuracy gates cannot see. + +THE PROBLEM THESE SOLVE + +Boot and accuracy both ask "does the engine produce the right text". For most +features that is the whole question. For four of them it is not: + +* **speculation** -- rejection sampling holds the emitted distribution to the + target model's, so a miscomputed draft produces correct text more slowly; +* **host offload** -- the same rows are gathered and the same arithmetic runs; + only where the bytes waited changes; +* **disaggregation** -- state that fails to cross the hop is silently + recomputed from zero, and the model recovers into fluent, wrong text; +* **native KV storage** -- packed FP4 dequantized exactly is the same value a + bf16 container held; only the bytes per token change. + +In each case both standing gates pass whether or not the feature is doing +anything, so the unit can be declared green for a feature that was never built. +The acceptance gate is the answer for the first; this module is the answer for +the rest. + +It is also the answer for a case with the opposite property: + +* **bounded replay** (DeepSeek-V4.1's decoder/encoder "SWA Bounded Replay") is + *not* output-identical -- it truncates sliding-window state on purpose -- so + accuracy is a real check of it. But the truncation is designed to "barely + compromise" accuracy, which makes a unit that implements nothing pass the + same gate. Only a count of the work skipped tells the two apart. + +THE SHAPE THEY SHARE + +They are **two-leg comparisons**: run one workload twice under one changed +variable, in the same session, and compare a named observable. A cross-session +comparison carries session variance into a judgement that is supposed to be +exact --- staircase measured a bit-identical forward at 94.7688 and 95.0720 a +day apart. The two that read a target's counters also hold them to a figure +the spec DECLARED, because a counter judged only against itself agrees with +whatever the target printed. + +What differs is the variable and the observable: + +| gate | variable | observable | verdict | +|---|---|---|---| +| acceptance | `TRTLLM_MODELING_V2` | `acceptance_length` | ratio >= floor | +| residency | the offload switch | device memory after load | delta >= predicted | +| equivalence | aggregated vs disaggregated | the generated tokens | identical | +| replay | the replay switch | per-request row counters | bounded ON, full OFF | +| kv_quant | native vs bf16 storage | bytes/token + tokens | declared; identical | + +WHAT THEY DELIBERATELY DO NOT MEASURE + +Speed. A residency witness says the table is not resident and the output is +unchanged; whether the per-step gather is worth having is a perf question and +belongs to tuning. A replay witness says the rows were skipped, not that +prefill got faster. Claiming more than the measurement supports is the failure +this whole package is built against. +""" + +from __future__ import annotations + +import ast +import json +import re +import shlex +import tempfile +import time +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any, Callable, Mapping, Sequence + +from .environment import RunEnvironment +from .gates import ( + BOOT_TIMEOUT_S, + MODE_REQUIRE, + RUNTIME_VALUE_DUMP_ENV, + BootSpec, + GateResult, + ShellRunner, + _default_shell_runner, + _with_mode, + resolve_max_draft_len, + resolve_runtime, + runtime_flags, + speculation_flags, + topology_flags, +) +from .task_schema import DEFAULT_BYTES_TOLERANCE + +#: `[0] Prompt: '...', Generated text: '...'` -- what +#: `examples/llm-api/quickstart_advanced.py` prints per request. The +#: equivalence witness compares these, and nothing else on the line. +_GENERATED = re.compile(r"Generated text:\s*(.+?)\.\.\.\s*$", re.M) + +#: Set to ``1`` on both legs of an equivalence witness. That line carries the +#: first 200 characters of a continuation -- under a quarter of 256 tokens -- +#: so under this variable `quickstart_advanced.py` also prints each one whole, +#: as ``[i] Witness token ids: [...]`` and ``[i] Witness text: '...'``. A +#: workload that does not know the variable prints neither, and the witness +#: compares the prefix and says so in its verdict. +WITNESS_OUTPUT_ENV = "TRTLLM_WITNESS_FULL_OUTPUT" +_WITNESS_TOKEN_IDS = re.compile(r"Witness token ids:\s*\[([-0-9,\s]*)\]\s*$", re.M) +_WITNESS_TEXT = re.compile(r"Witness text:\s*(.+?)\s*$", re.M) + +#: What leg A resolved, exported to an equivalence witness's ``contrast_command`` +#: so leg B reads it rather than restating it. A restated workload desynchronises +#: silently at the next edit to either side, and the verdict then reads as a +#: transfer defect. ``LEG_A_ARGV`` is leg A's quickstart argv with the prompts +#: elided -- they are in ``PROMPTS_FILE``, a JSON ``{"prompts": [...]}``, because +#: one environment string is capped at 128 KiB like one argv string. +#: ``DRAFT_LEN`` is set only when leg A speculates. ``MAX_BATCH_SIZE`` is set +#: only when the witness block declares ``max_batch_size``: leg B applies it +#: to every engine it starts, and at ``1`` sends its requests one at a time. +#: ``LEGS_ENV`` is set only when the block declares ``legs_env``: the same +#: variables are already in leg B's own environment, and this JSON object +#: ``{"VAR": "value"}`` names them so leg B can apply them to every engine it +#: starts after any setting of its own, as leg A's process has them. +LEG_B_PROMPTS_FILE_ENV = "STAIRCASE_WITNESS_PROMPTS_FILE" +LEG_B_MAX_TOKENS_ENV = "STAIRCASE_WITNESS_MAX_TOKENS" +LEG_B_MEMBER_ENV = "STAIRCASE_WITNESS_MEMBER" +LEG_B_DRAFT_LEN_ENV = "STAIRCASE_WITNESS_DRAFT_LEN" +LEG_B_ARGV_ENV = "STAIRCASE_WITNESS_LEG_A_ARGV" +LEG_B_TIMEOUT_ENV = "STAIRCASE_WITNESS_TIMEOUT_S" +LEG_B_MAX_BATCH_SIZE_ENV = "STAIRCASE_WITNESS_MAX_BATCH_SIZE" +LEG_B_LEGS_ENV_ENV = "STAIRCASE_WITNESS_LEGS_ENV" + +#: The quickstart flag a witness block's ``max_batch_size`` replaces. The +#: runtime block usually sets it through ``extra_flags``, so an override has +#: to remove that spelling, not just append another one. +BATCH_SIZE_FLAG = "--max_batch_size" + +#: What `trtllm-llmapi-launch` exports to the process it wraps. A contrast +#: command is a launcher that starts engines of its own, and an engine that +#: inherits a spawn-proxy address does not create its own: the context and the +#: generation engine would both dial the outer wrapper's proxy and hang with no +#: error. Unset before the contrast command runs, whatever launched it. +LAUNCHER_EXPORTS: tuple[str, ...] = ( + "TLLM_SPAWN_PROXY_PROCESS", + "TLLM_SPAWN_PROXY_PROCESS_IPC_ADDR", + "TLLM_SPAWN_PROXY_PROCESS_IPC_HMAC_KEY", + "tllm_mpi_size", + "FLASHINFER_WORKSPACE_BASE", + "TRTLLM_FLASHINFER_WORKSPACE_MANAGED", +) + +#: `torch.cuda.max_memory_allocated` and `nvidia-smi --query-gpu=memory.used` +#: both report MiB-scale integers; the witness asks the workload to print one +#: under this marker so the gate does not have to guess which tool ran. +DEVICE_MEMORY_MARKER = "__DEVICE_MEMORY_MIB__" +_DEVICE_MEMORY = re.compile(rf"{DEVICE_MEMORY_MARKER}\s*[:=]?\s*([0-9]+(?:\.[0-9]+)?)") + +#: Set to ``1`` on BOTH legs of a replay witness. A target prints its counter +#: lines only under it, so an ordinary run's log is untouched and no verdict +#: depends on a logging level. +WITNESS_COUNTER_ENV = "TRTLLM_MODELING_V2_WITNESS" + +#: The counter line a target prints. Space-separated ``key=`` fields on +#: one line; anything may precede the marker (a logger prefix, a rank tag). +#: Two kinds share it. +#: +#: Once per context request, when its context phase completes (after its final +#: chunk), the replay counters -- each for that one request over its whole +#: context phase:: +#: +#: MODELING_V2_WITNESS prompt_tokens=4103 forwarded_tokens=4103 decoder_tokens=128 +#: +#: ``prompt_tokens`` the request's prompt length +#: ``reused_tokens`` prompt tokens served from the prefix cache +#: ``replayed_tokens`` cached tokens forwarded again in replay mode +#: ``forwarded_tokens`` rows the model's FIRST layer processed +#: ``decoder_tokens`` rows the first bounded layer -- the one after the last +#: KV-source layer -- processed +#: +#: Rows, not positions: the verdicts are statements about work done, and a +#: count of positions attended to would be the same number with the feature on +#: and off. A field a witness does not read may be omitted. +#: +#: Once per rank, after the KV pools are allocated, the storage line +#: ``MODELING_V2_WITNESS kv_bytes_per_token=`` -- see +#: :func:`run_kv_quant_witness`. Each witness reads only the lines carrying its +#: own fields, so one target can print both kinds. +WITNESS_COUNTER_MARKER = "MODELING_V2_WITNESS" +_COUNTER_LINE = re.compile(rf"{WITNESS_COUNTER_MARKER}((?:[ \t]+[a-z_]+=-?\d+)+)") +_COUNTER_FIELD = re.compile(r"([a-z_]+)=(-?\d+)") + +#: What each replay witness reads off a counter line. +REPLAY_COUNTER_FIELDS: dict[str, tuple[str, ...]] = { + "decoder_replay": ("forwarded_tokens", "decoder_tokens"), + "encoder_replay": ("prompt_tokens", "reused_tokens", "replayed_tokens", "forwarded_tokens"), +} + +#: What the kv_quant witness reads: the persistent KV bytes one rank stores per +#: token of a request it serves, summed over its pools as ALLOCATED -- a ratio-r +#: compressed pool contributes its entry size over r. Fixed per-request state +#: (the SWA window, compressor carry-over) is not per-token and is excluded. +KV_BYTES_FIELD = "kv_bytes_per_token" + +# One short word is at least one token under every tokenizer that splits on +# whitespace before merging, which is every tokenizer these checkpoints ship. +# So `n` words is a prompt of at least `n` tokens, and the verdict checks the +# target's own count against that floor rather than trusting it. +_FILLER = ("one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten") + + +def parse_generated_texts(stdout: str) -> tuple[str, ...]: + """Every generated continuation in a quickstart run's output, in order.""" + return tuple(match.strip() for match in _GENERATED.findall(stdout)) + + +def _unrepr(text: str) -> str: + """The string a ``{text!r}`` printed, or ``text`` itself if it is not one.""" + try: + value = ast.literal_eval(text) + except (ValueError, SyntaxError): + return text + return value if isinstance(value, str) else text + + +def parse_witness_token_ids(stdout: str) -> tuple[tuple[int, ...], ...]: + """Every ``Witness token ids:`` list a leg printed, in order.""" + return tuple( + tuple(int(t) for t in re.findall(r"-?\d+", match)) + for match in _WITNESS_TOKEN_IDS.findall(stdout) + ) + + +def parse_witness_texts(stdout: str) -> tuple[str, ...]: + """Every ``Witness text:`` continuation a leg printed, whole, in order.""" + return tuple(_unrepr(match) for match in _WITNESS_TEXT.findall(stdout)) + + +def first_difference(a: Sequence[Any], b: Sequence[Any]) -> int: + """The first position at which ``a`` and ``b`` differ. + + The shorter length when one is a prefix of the other -- a continuation cut + short differs where it stops. + """ + for position, (x, y) in enumerate(zip(a, b)): + if x != y: + return position + return min(len(a), len(b)) + + +#: How far either side of a first difference a verdict shows, per unit. The +#: excerpt is for a reader deciding "lost state or a flipped near-tie", which +#: needs the words on both sides of the flip; the headline rides in a one-line +#: verdict and is kept short. +_EXCERPT_WIDTH = {"token": 8, "character": 60} +_HEADLINE_WIDTH = {"token": 4, "character": 30} +#: Differing continuations an excerpt shows in full; the rest are counted. +_EXCERPT_CONTINUATIONS = 4 + + +def _window(a: Sequence[Any], b: Sequence[Any], unit: str, width: int) -> tuple[str, Any, Any]: + """``a`` and ``b`` around where they first differ, and a label saying where.""" + at = first_difference(a, b) + lo, hi = max(0, at - width), min(at + width, max(len(a), len(b))) + if isinstance(a, str): + return f"[{lo}:{hi}] around {unit} {at}", a[lo:hi], b[lo:hi] + return f"[{lo}:{hi}] around {unit} {at}", list(a[lo:hi]), list(b[lo:hi]) + + +@dataclass(frozen=True) +class ContinuationComparison: + """Two legs' continuations, compared on the strongest observable both printed. + + ``texts_a`` / ``texts_b`` are the same continuations as text -- whole where + both legs printed them, else the 200-character prefixes -- so that a + difference found in token ids can also be SHOWN as words. + """ + + observable: str + unit: str + a: tuple[Any, ...] + b: tuple[Any, ...] + texts_a: tuple[str, ...] = () + texts_b: tuple[str, ...] = () + + @property + def mismatched(self) -> tuple[int, ...]: + return tuple(i for i, (x, y) in enumerate(zip(self.a, self.b)) if x != y) + + @property + def identical(self) -> bool: + return bool(self.a) and len(self.a) == len(self.b) and not self.mismatched + + def first_differences(self) -> dict[int, int]: + """Continuation index -> the first differing token or character in it.""" + return {i: first_difference(self.a[i], self.b[i]) for i in self.mismatched} + + def describe_differences(self, names: tuple[str, str] = ("leg A", "leg B")) -> str: + """Where each differing continuation first differs, and what both legs had there.""" + if len(self.a) != len(self.b): + return f"{names[0]} printed {len(self.a)} continuation(s), {names[1]} {len(self.b)}" + parts = [] + for i, at in self.first_differences().items(): + parts.append( + f"#{i} first differs at {self.unit} {at} " + f"({names[0]} {len(self.a[i])}, {names[1]} {len(self.b[i])} {self.unit}s)" + ) + return "; ".join(parts) + + def _texts(self, i: int) -> tuple[str, str] | None: + """Continuation ``i`` as text on both legs, where ``a``/``b`` are not text already.""" + if self.unit == "character" or i >= min(len(self.texts_a), len(self.texts_b)): + return None + return self.texts_a[i], self.texts_b[i] + + def headline(self, names: tuple[str, str] = ("leg A", "leg B")) -> str: + """One line: both legs' first differing continuation around where it differs. + + As words where the legs' texts differ, since that is what a reader can + judge; as token ids when the ids differ and the text does not. + """ + if not self.mismatched: + return "" + i = self.mismatched[0] + texts = self._texts(i) + if texts is not None and texts[0] != texts[1]: + where, a, b = _window(*texts, "character", _HEADLINE_WIDTH["character"]) + label = "text " + else: + where, a, b = _window(self.a[i], self.b[i], self.unit, _HEADLINE_WIDTH[self.unit]) + label = "token ids " if self.unit == "token" else "" + return f"#{i} {label}{where}: {names[0]} {a!r} vs {names[1]} {b!r}" + + def excerpt(self, names: tuple[str, str] = ("leg A", "leg B")) -> str: + """Both legs around the first difference of each differing continuation. + + The compared observable first, then -- when that is token ids -- the + text around its own first difference, because an id list does not say + what was said. + """ + if len(self.a) != len(self.b) or not self.mismatched: + return f"{names[0]}: {self.a[:2]}\n{names[1]}: {self.b[:2]}" + pad = max(len(n) for n in names) + + def block(label: str, a: Sequence[Any], b: Sequence[Any], unit: str) -> list[str]: + if a == b: + return [f" {label}: identical on both legs"] + where, wa, wb = _window(a, b, unit, _EXCERPT_WIDTH[unit]) + return [ + f" {label} {where}:", + f" {names[0]:<{pad}}: {wa!r}", + f" {names[1]:<{pad}}: {wb!r}", + ] + + lines: list[str] = [] + for i in self.mismatched[:_EXCERPT_CONTINUATIONS]: + at = first_difference(self.a[i], self.b[i]) + lines.append( + f"continuation #{i}: first differs at {self.unit} {at} " + f"({names[0]} {len(self.a[i])}, {names[1]} {len(self.b[i])} {self.unit}s)" + ) + lines += block(self.observable, self.a[i], self.b[i], self.unit) + texts = self._texts(i) + if texts is not None: + lines += block("text", *texts, "character") + rest = len(self.mismatched) - _EXCERPT_CONTINUATIONS + if rest > 0: + lines.append(f"... and {rest} more differing continuation(s), listed in the verdict") + return "\n".join(lines) + + +#: Strongest first. A leg "printed" an observable when it printed it for every +#: continuation it printed a `Generated text:` line for. +_OBSERVABLES: tuple[tuple[str, str, Callable[[str], tuple[Any, ...]]], ...] = ( + ("token ids", "token", parse_witness_token_ids), + ("full text", "character", parse_witness_texts), +) +PREFIX_OBSERVABLE = "200-character text prefix" + + +def compare_continuations(out_a: str, out_b: str) -> ContinuationComparison: + """Compare two legs' continuations on the strongest observable BOTH printed. + + Token ids where both legs printed them, else the full text, else the + 200-character prefix every quickstart run prints. Chosen per pair, not per + leg: comparing one leg's ids with the other's text is not a comparison. + """ + prefix_a, prefix_b = parse_generated_texts(out_a), parse_generated_texts(out_b) + prefix_a = tuple(_unrepr(t) for t in prefix_a) + prefix_b = tuple(_unrepr(t) for t in prefix_b) + texts_a, texts_b = prefix_a, prefix_b + whole_a, whole_b = parse_witness_texts(out_a), parse_witness_texts(out_b) + if whole_a and whole_b and len(whole_a) == len(prefix_a) and len(whole_b) == len(prefix_b): + texts_a, texts_b = whole_a, whole_b + for observable, unit, parse in _OBSERVABLES: + a, b = parse(out_a), parse(out_b) + if a and b and len(a) == len(prefix_a) and len(b) == len(prefix_b): + return ContinuationComparison(observable, unit, a, b, texts_a, texts_b) + return ContinuationComparison(PREFIX_OBSERVABLE, "character", prefix_a, prefix_b) + + +def parse_device_memory_mib(stdout: str) -> float | None: + """The device-memory reading a witness workload printed, in MiB. + + Read from a marker rather than from `nvidia-smi`'s own format because the + workload is better placed than the gate to decide *when* to sample --- the + number that matters is taken after the engine reports ready, and a gate + sampling from outside cannot know when that was. + """ + matches = _DEVICE_MEMORY.findall(stdout) + return float(matches[-1]) if matches else None + + +def parse_witness_counters(stdout: str) -> list[dict[str, int]]: + """Every counter line a target printed, one mapping per context request. + + A multi-rank run may print a request once per rank. Every verdict below is + a per-line statement, so a duplicated line is judged twice and changes + nothing. + """ + return [ + {key: int(value) for key, value in _COUNTER_FIELD.findall(match.group(1))} + for match in _COUNTER_LINE.finditer(stdout) + ] + + +def filler_words(count: int) -> str: + """``count`` short words: a prompt of at least ``count`` tokens.""" + return " ".join(_FILLER[i % len(_FILLER)] for i in range(count)) + + +def witness_command( + task_data: Mapping[str, Any], + environment: RunEnvironment, + *, + unit: str | None, + spec: BootSpec | None = None, + features: Any | None = None, + max_tokens: int = 64, + extra: tuple[str, ...] = (), + member: str | None = None, + prompts: Sequence[str] | None = None, + max_batch_size: int | None = None, +) -> str: + """The quickstart argv both legs of a witness run. + + Quickstart rather than `trtllm-eval` for the same reason the acceptance + gate uses it: the eval harness reports neither device memory nor the + per-request continuations, and it costs nine minutes a leg where this + costs four. + + `member` names the parallel configuration BOTH legs run at. A witness is + a contrast between two legs, so the one thing it must not do is let the + topology differ between them -- and under a `unified` segment with no + member, `topology_flags` resolves to pure tensor parallelism, which is a + configuration this target may not even be delivered at. + + The draft depth comes from the unit's options YAML, as it does for the + boot gate. Before, a witness always asked for depth 1, which is harmless + for MTP and fatal for DSpark: a unit built after `dspark` runs its witness + with the DSpark flags, and `TorchLlmArgs` rejects a DSpark config whose + `block_size` (5, from the checkpoint) differs from `max_draft_len` -- so + both legs die at startup and the witness reports a leg that did not + complete, which reads as the unit's own defect. + + `prompts` overrides the boot spec's, for a witness whose workload is + generated from declared sizes rather than written by hand. + + `max_batch_size` replaces whatever batch limit the runtime block set (see + :func:`with_max_batch_size`). + """ + spec = spec or BootSpec() + quoted = " ".join(f'"{p}"' for p in (spec.prompts if prompts is None else prompts)) + return " ".join( + [ + *_quickstart_argv( + task_data, + environment, + unit=unit, + features=features, + max_tokens=max_tokens, + extra=extra, + member=member, + max_batch_size=max_batch_size, + ), + f"--prompt {quoted}", + ] + ) + + +def with_max_batch_size(flags: Sequence[str], max_batch_size: int | None) -> list[str]: + """``flags`` with every ``--max_batch_size`` spelling replaced by one ``=N``. + + A token-comparing witness declares the batch limit its legs run at, and + the runtime block has usually set its own through ``extra_flags`` -- + ``--max_batch_size=8``, ``--max_batch_size 8`` as one string, or the flag + and its value as two. All three are removed, so the argv states ONE limit + and leg B's reading of leg A's argv cannot pick up the stale one. ``None`` + returns ``flags`` unchanged. + """ + if max_batch_size is None: + return list(flags) + kept: list[str] = [] + skip_value = False + for flag in flags: + if skip_value: + skip_value = False + continue + word = str(flag) + if word == BATCH_SIZE_FLAG: + skip_value = True + continue + if word.startswith((f"{BATCH_SIZE_FLAG}=", f"{BATCH_SIZE_FLAG} ")): + continue + kept.append(word) + return [*kept, f"{BATCH_SIZE_FLAG}={int(max_batch_size)}"] + + +def _quickstart_argv( + task_data: Mapping[str, Any], + environment: RunEnvironment, + *, + unit: str | None, + features: Any | None, + max_tokens: int, + extra: tuple[str, ...], + member: str | None, + max_batch_size: int | None = None, +) -> list[str]: + """:func:`witness_command` without its prompts.""" + script = Path(str(task_data["trtllm_repo_path"])) / "examples" / "llm-api" + return [ + str(script / "quickstart_advanced.py"), + f'--model_dir="{task_data["checkpoint_path"]}"', + f"--max_tokens={max_tokens}", + *topology_flags(str(task_data["parallel"]), environment.world_size, member), + *with_max_batch_size( + [ + *runtime_flags(resolve_runtime(task_data, unit)), + *speculation_flags( + task_data.get("features") if features is None else features, + max_draft_len=resolve_max_draft_len(task_data.get("accuracy") or {}, unit), + ), + *extra, + ], + max_batch_size, + ), + ] + + +def _schedule_note(max_batch_size: int | None) -> str: + """What a verdict says about the legs' batch limit, when the block declared one.""" + if max_batch_size is None: + return "" + return f"; both legs at max_batch_size={max_batch_size}" + + +def _legs_note(legs_env: Mapping[str, str] | None) -> str: + """What a verdict says about the environment both legs shared, when declared.""" + if not legs_env: + return "" + return "; both legs under " + ", ".join(f"{k}={v}" for k, v in sorted(legs_env.items())) + + +def _around(comparison: ContinuationComparison, names: tuple[str, str]) -> str: + """``; around it, ``, or nothing.""" + headline = comparison.headline(names) + return f"; around it, {headline}" if headline else "" + + +def _token_difference(comparison: ContinuationComparison, names: tuple[str, str]) -> str: + """The failing half of a token-identity verdict: where, and what both legs had there.""" + if not comparison.a or not comparison.b: + return ( + f" (no continuations parsed from both legs: {names[0]} {len(comparison.a)}, " + f"{names[1]} {len(comparison.b)})" + ) + return ( + f" (compared as {comparison.observable}): {comparison.describe_differences(names)}" + f"{_around(comparison, names)}" + ) + + +def _run_leg( + environment: RunEnvironment, + argv: str, + runner: ShellRunner, + timeout: int, + env: Mapping[str, str] | None = None, +) -> tuple[str, str, bool]: + """Run one leg. Returns (command, combined output, ok).""" + environment = replace( + environment, env={**environment.env, RUNTIME_VALUE_DUMP_ENV: "1", **(env or {})} + ) + command = environment.python_command(argv) + result = runner(command, None, timeout) + return command, result.stdout + result.stderr, result.ok + + +def contrast_leg_command(environment: RunEnvironment, contrast_command: str) -> str: + """How an equivalence witness launches ``contrast_command``: as a launcher. + + Leg B of a disagg witness is two engines, a proxy and a client -- work + that starts ranks, not a rank. Launched as engine work it went out at world + size 4 as ``mpirun -n 4 ... trtllm-llmapi-launch ``: the + command ran on rank 0 while three idle workers held ranks 1-3 and the + devices' memory, and it inherited the wrapper's spawn-proxy variables, + which make the engines it starts hang. So it runs as ONE process, the + branch :meth:`RunEnvironment.command` already has for work that is not + engine work, with :data:`LAUNCHER_EXPORTS` unset. The process keeps the + claim's devices: ``CUDA_VISIBLE_DEVICES`` names all of them. + + Under ``slurm`` one task means ``srun --gres=gpu:1``, which would take the + launcher's devices away; there it keeps the engine-shaped ``srun`` over + the whole claim and relies on the same unset. + + A command whose first word is a ``.py`` file runs under the run's + interpreter, as before; anything else -- a shell script -- runs as written. + """ + try: + words = shlex.split(contrast_command) + except ValueError: + words = contrast_command.split() + body = contrast_command + if words and words[0].endswith(".py"): + body = f"{shlex.quote(environment.python)} {contrast_command}" + body = "env " + " ".join(f"-u {name}" for name in LAUNCHER_EXPORTS) + f" {body}" + if environment.kind == "slurm": + return environment.command(body) + return environment.command(body, single_process=True) + + +def run_residency_witness( + environment: RunEnvironment, + *, + task_data: Mapping[str, Any], + expected_mib: float, + contrast_env: Mapping[str, str], + runner: ShellRunner = _default_shell_runner, + timeout: int = BOOT_TIMEOUT_S, + unit: str | None = None, + spec: BootSpec | None = None, + features: Any | None = None, + member: str | None = None, + max_batch_size: int | None = None, +) -> GateResult: + """Prove a host-offloaded tensor is actually not resident. + + Two legs of the same workload in the same session, differing only by + ``contrast_env`` --- the environment that turns the offload OFF. Both must + produce **token-identical** output, and the resident leg must use at least + ``expected_mib`` more device memory than the offloaded one. + + ``expected_mib`` is computed from the CHECKPOINT's own shapes and passed in + by the caller, deliberately: a witness that derives its own expectation + from the run it is judging can only ever agree with itself. + + The token-identity half is not decoration. Offloading is output-identical + by construction, so if the two legs disagree the gather is wrong --- and + the accuracy gate's tolerance is wide enough to hide that. Compared as + :func:`compare_continuations` does, with both legs printing their whole + continuations; ``max_batch_size`` runs both legs at that batch limit. + """ + argv = witness_command( + task_data, + environment, + unit=unit, + spec=spec, + features=features, + member=member, + max_batch_size=max_batch_size, + ) + full_output = {WITNESS_OUTPUT_ENV: "1"} + started = time.monotonic() + + on_cmd, on_out, on_ok = _run_leg(environment, argv, runner, timeout, env=full_output) + if not on_ok: + return GateResult( + name="residency", + member=member, + passed=False, + detail="the offloaded leg did not complete", + command=on_cmd, + duration_s=time.monotonic() - started, + stdout_tail="\n".join(on_out.splitlines()[-40:]), + ) + + off_cmd, off_out, off_ok = _run_leg( + environment, argv, runner, timeout, env={**contrast_env, **full_output} + ) + elapsed = time.monotonic() - started + if not off_ok: + return GateResult( + name="residency", + member=member, + passed=False, + detail=f"the resident control leg ({dict(contrast_env)}) did not complete", + command=off_cmd, + duration_s=elapsed, + stdout_tail="\n".join(off_out.splitlines()[-40:]), + ) + + on_mib, off_mib = parse_device_memory_mib(on_out), parse_device_memory_mib(off_out) + if on_mib is None or off_mib is None: + return GateResult( + name="residency", + member=member, + passed=False, + detail=f"no `{DEVICE_MEMORY_MARKER}` reading from both legs " + f"(offloaded={on_mib}, resident={off_mib}); a witness with no number " + f"is not a pass", + command=on_cmd, + duration_s=elapsed, + ) + + comparison = compare_continuations(on_out, off_out) + if not comparison.identical: + names = ("offloaded", "resident") + return GateResult( + name="residency", + member=member, + passed=False, + detail=f"the two legs did not generate identical text, so the gather is " + f"wrong -- offloading is output-identical by construction" + f"{_token_difference(comparison, names)}{_schedule_note(max_batch_size)}", + command=on_cmd, + duration_s=elapsed, + stdout_tail=( + comparison.excerpt(names) + if comparison.a and comparison.b + else "\n".join(on_out.splitlines()[-20:]) + ), + ) + + delta = off_mib - on_mib + passed = delta >= expected_mib + return GateResult( + name="residency", + member=member, + passed=passed, + detail=( + f"device memory {on_mib:.0f} MiB offloaded vs {off_mib:.0f} MiB resident " + f"(delta {delta:.0f} MiB {'>=' if passed else '<'} predicted " + f"{expected_mib:.0f} MiB), {len(comparison.a)} continuation(s) identical " + f"(compared as {comparison.observable}){_schedule_note(max_batch_size)}" + ), + command=on_cmd, + measured=delta, + reference=expected_mib, + duration_s=elapsed, + ) + + +def run_equivalence_witness( + environment: RunEnvironment, + *, + task_data: Mapping[str, Any], + contrast_env: Mapping[str, str], + runner: ShellRunner = _default_shell_runner, + timeout: int = BOOT_TIMEOUT_S, + unit: str | None = None, + spec: BootSpec | None = None, + features: Any | None = None, + max_tokens: int = 256, + contrast_command: str | None = None, + member: str | None = None, + prompt_tokens: Sequence[int] | None = None, + workdir: str | Path | None = None, + max_batch_size: int | None = None, + legs_env: Mapping[str, str] | None = None, +) -> GateResult: + """Prove per-sequence state survives the prefill/decode hop. + + ``contrast_command`` is how leg B is actually launched, and for this + witness it is REQUIRED -- the schema and the unit boundary both refuse a + ``disagg`` block without one; this function still accepts None for the + one-switch shape. Disaggregation is not an environment variable over one + process: it is a context server, a generation server and a client, which + `examples/disaggregated/` launches and `quickstart_advanced.py` cannot. An + env overlay alone would run leg B aggregated and the witness would compare + a run to itself --- passing trivially, which is the worst possible outcome + for a gate whose whole job is to catch a silent equality. + + The contract on that command is small and the run owns it: launch the + disaggregated topology, drive the same prompts greedily, and print the + continuations in `quickstart_advanced.py`'s format so both legs parse the + same way -- and, under :data:`WITNESS_OUTPUT_ENV`, the whole continuation + as well. It runs as a launcher (:func:`contrast_leg_command`) and is handed + leg A's resolved workload through the ``LEG_B_*_ENV`` variables, so it + need not restate any of it. + + ``prompt_tokens`` appends one generated prompt of at least that many + tokens per entry to the boot prompts, on BOTH legs. The boot prompts stay: + they are the ones whose answers depend on compressed and long-range state + (a pass key 1.3K tokens back). The generated ones supply the lengths a hop + has edges at -- over one prefill chunk, odd, one past a page -- which + hand-written prompts rarely land on. + + Leg A is one aggregated engine; leg B is the same workload with + ``contrast_env`` applied --- here the environment that turns disaggregation + ON, which is why the key is not named for a direction: the residency + witness's second leg turns its feature *off* and this one's turns its + feature *on*. Greedy, same seed, same session. The verdict is + **token identity**, not similarity: missing recurrent state yields text + that is fluent and wrong, and a short benchmark under a five-point + tolerance will not catch it. It compares token ids where both legs print + them, else the full text, else the 200-character prefix, and names which. + + A failing verdict names where each continuation FIRST differs. State lost + at the hop is read by the generation engine's first decode step, so it + shows within the first few generated tokens; a first difference far into a + continuation, with both legs speculating on different schedules, is a + near-tie flipped by batch-shape-variant numerics -- a logit-margin + question, not a transfer one. + + ``max_batch_size`` is the batch limit BOTH legs run at: leg A's argv + carries it in place of the runtime block's, and leg B is handed it as + :data:`LEG_B_MAX_BATCH_SIZE_ENV`. At 1 each request runs alone on both + sides of the comparison, so the two legs' step shapes match and a + difference cannot come from how admission happened to pack rows -- the + batch-shape-variant numerics above are then out of the picture. + + ``legs_env`` is the environment BOTH legs run under, where ``contrast_env`` + is leg B's alone: leg A's process gets it, and leg B's launcher gets it in + its own environment and, named, as :data:`LEG_B_LEGS_ENV_ENV` for every + engine it starts. It pins a code path the disaggregated engines refuse or + guard off to the same setting in the aggregated leg; otherwise the legs + differ in that path as well as in the hop, and identity is no longer + expected. The verdict names it, since a pass holds under it. + + For a hybrid checkpoint the exposure is the recurrent conv/SSM state and + any sparse-attention indexer state -- not the attention pages everyone + remembers to check. + """ + spec = spec or BootSpec() + prompts = (*spec.prompts, *(filler_words(int(n)) for n in (prompt_tokens or ()))) + shape = dict( + unit=unit, + features=features, + max_tokens=max_tokens, + extra=(), + member=member, + max_batch_size=max_batch_size, + ) + argv = witness_command(task_data, environment, spec=spec, prompts=prompts, **shape) + full_output = {WITNESS_OUTPUT_ENV: "1"} + shared = {str(k): str(v) for k, v in (legs_env or {}).items()} + started = time.monotonic() + + agg_cmd, agg_out, agg_ok = _run_leg( + environment, argv, runner, timeout, env={**shared, **full_output} + ) + if not agg_ok: + return GateResult( + name="equivalence", + member=member, + passed=False, + detail="the aggregated leg did not complete", + command=agg_cmd, + duration_s=time.monotonic() - started, + stdout_tail="\n".join(agg_out.splitlines()[-40:]), + ) + + if contrast_command: + directory = Path(workdir) if workdir else Path(tempfile.mkdtemp(prefix="witness-")) + directory.mkdir(parents=True, exist_ok=True) + prompts_file = directory / "prompts.json" + prompts_file.write_text(json.dumps({"prompts": list(prompts)}), encoding="utf-8") + leg_a = { + # Absolute: leg B is a launcher of its own and need not run from + # the directory the orchestrator resolved the workspace against. + LEG_B_PROMPTS_FILE_ENV: str(prompts_file.resolve()), + LEG_B_MAX_TOKENS_ENV: str(max_tokens), + LEG_B_ARGV_ENV: " ".join(_quickstart_argv(task_data, environment, **shape)), + LEG_B_TIMEOUT_ENV: str(timeout), + } + if member: + leg_a[LEG_B_MEMBER_ENV] = member + if max_batch_size is not None: + leg_a[LEG_B_MAX_BATCH_SIZE_ENV] = str(int(max_batch_size)) + if shared: + leg_a[LEG_B_LEGS_ENV_ENV] = json.dumps(shared, sort_keys=True) + if speculation_flags(task_data.get("features") if features is None else features): + leg_a[LEG_B_DRAFT_LEN_ENV] = str( + resolve_max_draft_len(task_data.get("accuracy") or {}, unit) + ) + contrast = replace( + environment, + env={ + **environment.env, + RUNTIME_VALUE_DUMP_ENV: "1", + **shared, + **contrast_env, + **full_output, + **leg_a, + }, + ) + dis_cmd = contrast_leg_command(contrast, contrast_command) + result = runner(dis_cmd, None, timeout) + dis_out, dis_ok = result.stdout + result.stderr, result.ok + else: + dis_cmd, dis_out, dis_ok = _run_leg( + environment, argv, runner, timeout, env={**shared, **contrast_env, **full_output} + ) + elapsed = time.monotonic() - started + if not dis_ok: + return GateResult( + name="equivalence", + member=member, + passed=False, + detail="the disaggregated leg did not complete. A " + "disagg launch fails by HANGING more often than by raising -- check " + "the re-executed interpreter and the cache transceiver before the target", + command=dis_cmd, + duration_s=elapsed, + stdout_tail="\n".join(dis_out.splitlines()[-40:]), + ) + + comparison = compare_continuations(agg_out, dis_out) + if not comparison.a or not comparison.b: + return GateResult( + name="equivalence", + member=member, + passed=False, + detail=f"no continuations parsed from both legs " + f"(aggregated={len(comparison.a)}, disaggregated={len(comparison.b)})", + command=agg_cmd, + duration_s=elapsed, + ) + + compared = f"compared as {comparison.observable}" + if comparison.observable == PREFIX_OBSERVABLE: + compared += ( + f"; a leg printed no whole continuation under {WITNESS_OUTPUT_ENV}=1, so the " + f"rest of each was not compared" + ) + passed = comparison.identical + names = ("aggregated", "disagg") + return GateResult( + name="equivalence", + member=member, + passed=passed, + detail=( + f"{len(comparison.a)} continuation(s) token-identical across the hop ({compared})" + f"{_schedule_note(max_batch_size)}{_legs_note(shared)}" + if passed + else f"continuation(s) {list(comparison.mismatched) or 'count'} differ across the " + f"hop ({compared}{_schedule_note(max_batch_size)}{_legs_note(shared)}): " + f"{comparison.describe_differences(names)}{_around(comparison, names)}. State that did not " + f"travel is recomputed from zero, and the model recovers into fluent, wrong text " + f"-- it shows within the first few generated tokens. A first difference far into " + f"a continuation is a near-tie under two speculation schedules: check logit " + f"margins before the transfer" + ), + command=agg_cmd, + measured=( + float(min(comparison.first_differences().values())) if comparison.mismatched else None + ), + duration_s=elapsed, + stdout_tail="" if passed else comparison.excerpt(names), + ) + + +def _judge_decoder_replay( + on: list[dict[str, int]], + off: list[dict[str, int]], + *, + window: int, + prompt_tokens: int, +) -> tuple[bool, str]: + """Replay ON bounds every request's decoder to ``window`` rows; OFF bounds none. + + The workload check comes first because it is what makes the other two + mean anything: a prompt that never ran cold at full length satisfies + "decoder rows <= window" with the feature on or off. + """ + for leg, rows in (("ON", on), ("OFF", off)): + longest = max(row["forwarded_tokens"] for row in rows) + if longest < prompt_tokens: + return False, ( + f"replay {leg}: the longest request forwarded {longest} rows, below the " + f"declared prompt_tokens={prompt_tokens}, so the long prompt never ran " + f"cold at full length and a bounded decoder cannot be told from a short " + f"prompt" + ) + over = [row for row in on if row["decoder_tokens"] > window] + if over: + worst = max(row["decoder_tokens"] for row in over) + return False, ( + f"replay ON: {len(over)} request(s) ran the decoder over more than " + f"window={window} rows (up to {worst}), so the decoder was not bounded" + ) + unbounded = [row for row in off if row["decoder_tokens"] != row["forwarded_tokens"]] + if unbounded: + first = unbounded[0] + return False, ( + f"replay OFF: {len(unbounded)} request(s) ran the decoder over " + f"{first['decoder_tokens']} rows of {first['forwarded_tokens']} forwarded. " + f"`contrast_env` did not turn the feature off, or the counter does not " + f"count what it names" + ) + return True, ( + f"replay ON: decoder <= {window} rows for each of {len(on)} request(s), the " + f"longest having forwarded {max(r['forwarded_tokens'] for r in on)}; replay " + f"OFF: decoder == forwarded for all {len(off)}" + ) + + +def _judge_encoder_replay( + on: list[dict[str, int]], + off: list[dict[str, int]], + *, + window: int, + prefix_tokens: int, + reuse_alignment: int, +) -> tuple[bool, str]: + """A prefix hit replays exactly ``min(window, reused)`` rows ON, and none OFF. + + The accounting identity ``forwarded == prompt - reused + replayed`` is + checked on both legs: a target that reports a hit and then recomputes the + whole prefix anyway -- which is what "SWA KV is missing" invites -- breaks + it, and nothing else in the verdict would notice. + """ + longest = max(row["prompt_tokens"] for row in on) + if longest < prefix_tokens: + return False, ( + f"replay ON: the longest prompt was {longest} tokens, below the declared " + f"prefix_tokens={prefix_tokens}, so the shared prefix never reached the target" + ) + for leg, rows in (("ON", on), ("OFF", off)): + broken = [ + row + for row in rows + if row["forwarded_tokens"] + != row["prompt_tokens"] - row["reused_tokens"] + row["replayed_tokens"] + ] + if broken: + return False, ( + f"replay {leg}: forwarded != prompt - reused + replayed for {len(broken)} " + f"request(s), first {broken[0]}. A cached prefix is being recomputed, or " + f"the counter does not count what it names" + ) + replayed_off = [row for row in off if row["replayed_tokens"] != 0] + if replayed_off: + return False, ( + f"replay OFF: {len(replayed_off)} request(s) replayed rows " + f"(first {replayed_off[0]['replayed_tokens']}); `contrast_env` did not turn the " + f"feature off" + ) + wrong = [row for row in on if row["replayed_tokens"] != min(window, row["reused_tokens"])] + if wrong: + first = wrong[0] + return False, ( + f"replay ON: {len(wrong)} request(s) replayed {first['replayed_tokens']} rows " + f"after reusing {first['reused_tokens']}; a bounded replay re-forwards exactly " + f"min(window={window}, reused)" + ) + misaligned = [row for row in on if row["reused_tokens"] % reuse_alignment] + if misaligned: + return False, ( + f"replay ON: reuse boundary {misaligned[0]['reused_tokens']} is not a multiple of " + f"{reuse_alignment}; a boundary that splits a compressed group leaves half of " + f"its global KV nowhere" + ) + hits = [row for row in on if row["reused_tokens"] > window] + if not hits: + most = max(row["reused_tokens"] for row in on) + return False, ( + f"replay ON: no request reused more than window={window} cached tokens " + f"(most: {most}), so a bounded replay cannot be told from re-forwarding the " + f"whole prefix. Block reuse must be on and the prefix must span a KV block " + f"past the window" + ) + return True, ( + f"replay ON: {len(hits)} prefix hit(s) reusing up to " + f"{max(r['reused_tokens'] for r in hits)} tokens each replayed " + f"min(window={window}, reused) rows and forwarded suffix + replay; replay OFF: " + f"nothing replayed over {len(off)} request(s)" + ) + + +def run_replay_witness( + environment: RunEnvironment, + *, + task_data: Mapping[str, Any], + feature: str, + contrast_env: Mapping[str, str], + window: int, + prompt_tokens: int | None = None, + prefix_tokens: int | None = None, + reuse_alignment: int = 1, + runner: ShellRunner = _default_shell_runner, + timeout: int = BOOT_TIMEOUT_S, + unit: str | None = None, + features: Any | None = None, + max_tokens: int = 8, + member: str | None = None, +) -> GateResult: + """Prove a bounded replay skipped the work it exists to skip. + + Bounded replay is NOT output-identical: both halves truncate sliding-window + state on purpose, so the accuracy gate is a real check of them. It is not a + sufficient one. The truncation was designed, and post-trained, to barely + move accuracy -- which means a unit that implements nothing lands inside + the same tolerance. The observable that differs is work: rows the decoder + ran (``decoder_replay``), rows a prefix hit re-forwarded + (``encoder_replay``). Only the target can count those, so the target + prints them -- see :data:`WITNESS_COUNTER_MARKER` -- and this judges them + against the declared ``window``. + + Two legs, one session, the same generated workload: the feature ON, then + OFF under ``contrast_env`` (the leg WITHOUT it). Both legs set + :data:`WITNESS_COUNTER_ENV`. The OFF leg is not decoration: it is what + shows the switch works and the counter counts, since a counter that always + prints ``decoder_tokens=128`` would otherwise pass. + + ``window``, ``prompt_tokens`` and ``prefix_tokens`` come from the spec, + declared before assembly -- a witness that sized its workload from the run + it judges could only agree with itself. + """ + name = f"replay[{feature}]" + started = time.monotonic() + prompts: tuple[str, ...] + if feature == "encoder_replay": + if prefix_tokens is None: + raise ValueError("the encoder_replay witness needs prefix_tokens") + # Checked before any engine starts: with block reuse off no request can + # hit the prefix cache, so the witness would spend two engine launches + # to report that nothing was replayed. + if not resolve_runtime(task_data, unit).get("kv_cache_reuse", True): + return GateResult( + name=name, + member=member, + passed=False, + detail=f"the `{unit}` unit's runtime turns kv_cache_reuse off; encoder replay " + f"only happens on a prefix hit, so there is nothing to witness", + ) + prefix = filler_words(prefix_tokens) + # Two requests sharing the prefix. Quickstart runs every prompt twice + # (a batched pass, then a streamed one), so the streamed pass of each + # hits the prefix the batched pass left in the cache. + prompts = (f"{prefix} What comes next", f"{prefix} Say it backwards") + elif feature == "decoder_replay": + if prompt_tokens is None: + raise ValueError("the decoder_replay witness needs prompt_tokens") + prompts = (filler_words(prompt_tokens),) + else: + raise ValueError(f"no replay witness for feature {feature!r}") + + argv = witness_command( + task_data, + environment, + unit=unit, + features=features, + max_tokens=max_tokens, + member=member, + prompts=prompts, + ) + env = _with_mode(environment, MODE_REQUIRE) + counting = {WITNESS_COUNTER_ENV: "1"} + legs: dict[str, list[dict[str, int]]] = {} + on_cmd = "" + for leg, overlay in (("ON", counting), ("OFF", {**contrast_env, **counting})): + command, output, ok = _run_leg(env, argv, runner, timeout, env=overlay) + on_cmd = on_cmd or command + if not ok: + return GateResult( + name=name, + member=member, + passed=False, + detail=f"the replay-{leg} leg did not complete" + + (f" ({dict(contrast_env)})" if leg == "OFF" else ""), + command=command, + duration_s=time.monotonic() - started, + stdout_tail="\n".join(output.splitlines()[-40:]), + ) + needed = REPLAY_COUNTER_FIELDS[feature] + # Only the lines carrying this witness's fields: a target that also + # implements `kv_quant` prints its storage line under the same marker. + rows = [row for row in parse_witness_counters(output) if any(key in row for key in needed)] + if not rows: + return GateResult( + name=name, + member=member, + passed=False, + detail=f"the replay-{leg} leg ran but the target printed no witness counter: " + f"no `{WITNESS_COUNTER_MARKER} key= ...` line carrying {list(needed)} " + f"with {WITNESS_COUNTER_ENV}=1 set. A witness with no number is not a pass", + command=command, + duration_s=time.monotonic() - started, + stdout_tail="\n".join(output.splitlines()[-40:]), + ) + incomplete = [row for row in rows if any(key not in row for key in needed)] + if incomplete: + missing = sorted({key for row in incomplete for key in needed if key not in row}) + return GateResult( + name=name, + member=member, + passed=False, + detail=f"the replay-{leg} leg printed {len(incomplete)} counter line(s) " + f"without {missing}; the {feature} witness needs {list(needed)}", + command=command, + duration_s=time.monotonic() - started, + ) + legs[leg] = rows + + if feature == "decoder_replay": + passed, detail = _judge_decoder_replay( + legs["ON"], legs["OFF"], window=window, prompt_tokens=int(prompt_tokens or 0) + ) + else: + passed, detail = _judge_encoder_replay( + legs["ON"], + legs["OFF"], + window=window, + prefix_tokens=int(prefix_tokens or 0), + reuse_alignment=max(1, int(reuse_alignment)), + ) + return GateResult( + name=name, + member=member, + passed=passed, + detail=detail, + command=on_cmd, + duration_s=time.monotonic() - started, + ) + + +def run_kv_quant_witness( + environment: RunEnvironment, + *, + task_data: Mapping[str, Any], + expected_bytes_per_token: int, + exact: bool = True, + contrast_env: Mapping[str, str] | None = None, + bytes_tolerance: float = DEFAULT_BYTES_TOLERANCE, + runner: ShellRunner = _default_shell_runner, + timeout: int = BOOT_TIMEOUT_S, + unit: str | None = None, + spec: BootSpec | None = None, + features: Any | None = None, + max_tokens: int = 256, + member: str | None = None, + max_batch_size: int | None = None, +) -> GateResult: + """Prove KV is stored natively, and -- when declared exact -- that nothing moved. + + Native storage of the SAME value set the reference fake-quantizes is + output-identical by construction when attention dequantizes exactly, so + boot and accuracy pass whether the pools hold packed FP4 or the bf16 + containers core shipped with. Two observables close that: + + * **bytes per token** -- the target prints + ``MODELING_V2_WITNESS kv_bytes_per_token=`` (see + :data:`WITNESS_COUNTER_MARKER`), computed from the pools as allocated, + and it must match the figure the spec declared from the checkpoint's + formats. bf16 containers of DeepSeek-V4.1's values are 3200 B/token + against 890, so the check has one way to pass. + * **token identity against bf16 storage** (``exact``, the default) -- + greedy output under ``contrast_env`` must match token for token. A + difference means the stored values are not the reference's: a layout, + scale or rounding that is somebody else's format. The contrast leg must + report MORE bytes per token than the native one, or ``contrast_env`` + selected nothing and the identity compared native storage with itself. + + ``exact: false`` is for an attention path declared lossy (an FP8-query + kernel). The identity leg is dropped, accuracy is the only output check, + and the verdict says it is the weaker claim. + + Identity is judged as :func:`compare_continuations` judges it -- token + ids where both legs print them under :data:`WITNESS_OUTPUT_ENV` -- and a + failing verdict names where each continuation first differs and what both + legs generated there. ``max_batch_size`` runs both legs at that batch + limit; at 1 the two legs' schedules match, so a difference is storage and + not how admission packed rows. + """ + name = "kv_quant" + started = time.monotonic() + argv = witness_command( + task_data, + environment, + unit=unit, + spec=spec, + features=features, + max_tokens=max_tokens, + member=member, + max_batch_size=max_batch_size, + ) + env = _with_mode(environment, MODE_REQUIRE) + counting = {WITNESS_COUNTER_ENV: "1", WITNESS_OUTPUT_ENV: "1"} + + def bytes_readings(output: str) -> list[int]: + rows = parse_witness_counters(output) + return [row[KV_BYTES_FIELD] for row in rows if KV_BYTES_FIELD in row] + + native_cmd, native_out, native_ok = _run_leg(env, argv, runner, timeout, env=counting) + if not native_ok: + return GateResult( + name=name, + member=member, + passed=False, + detail="the native-storage leg did not complete", + command=native_cmd, + duration_s=time.monotonic() - started, + stdout_tail="\n".join(native_out.splitlines()[-40:]), + ) + native = bytes_readings(native_out) + if not native: + return GateResult( + name=name, + member=member, + passed=False, + detail=f"the native-storage leg ran but the target printed no witness counter: " + f"no `{WITNESS_COUNTER_MARKER} {KV_BYTES_FIELD}=` line with " + f"{WITNESS_COUNTER_ENV}=1 set. A witness with no number is not a pass", + command=native_cmd, + duration_s=time.monotonic() - started, + stdout_tail="\n".join(native_out.splitlines()[-40:]), + ) + slack = bytes_tolerance * expected_bytes_per_token + outside = [b for b in native if abs(b - expected_bytes_per_token) > slack] + if outside: + return GateResult( + name=name, + member=member, + passed=False, + detail=f"native storage reports {outside[0]} B/token against the declared " + f"{expected_bytes_per_token} (tolerance {bytes_tolerance:.0%}). The pools do not " + f"hold the checkpoint's formats -- still bf16, or a different layout", + command=native_cmd, + duration_s=time.monotonic() - started, + ) + stored = f"{max(native)} B/token (declared {expected_bytes_per_token})" + if not exact: + return GateResult( + name=name, + member=member, + passed=True, + detail=f"native storage at {stored}. THE WEAKER CLAIM: `exact: false` declares " + f"a lossy attention path, so output was NOT compared with bf16 storage and " + f"accuracy is the only output check", + command=native_cmd, + duration_s=time.monotonic() - started, + ) + + contrast = {**dict(contrast_env or {}), **counting} + ref_cmd, ref_out, ref_ok = _run_leg(env, argv, runner, timeout, env=contrast) + elapsed = time.monotonic() - started + if not ref_ok: + return GateResult( + name=name, + member=member, + passed=False, + detail=f"the bf16-storage leg ({dict(contrast_env or {})}) did not complete", + command=ref_cmd, + duration_s=elapsed, + stdout_tail="\n".join(ref_out.splitlines()[-40:]), + ) + reference = bytes_readings(ref_out) + if not reference or min(reference) <= max(native): + return GateResult( + name=name, + member=member, + passed=False, + detail=f"the bf16-storage leg reported {reference or 'no'} B/token against " + f"native {max(native)}: `contrast_env` did not select bf16 storage, so token " + f"identity would compare native storage with itself", + command=ref_cmd, + duration_s=elapsed, + ) + comparison = compare_continuations(native_out, ref_out) + if not comparison.identical: + names = ("native", "bf16") + return GateResult( + name=name, + member=member, + passed=False, + detail=f"native storage changed the output (continuation(s) " + f"{list(comparison.mismatched) or 'count'} differ from bf16 storage" + f"{_token_difference(comparison, names)}{_schedule_note(max_batch_size)}). " + f"Storing the reference's value set is output-identical by construction, so the " + f"stored values are not the reference's -- a layout, scale or rounding that is " + f"another format's", + command=native_cmd, + measured=( + float(min(comparison.first_differences().values())) + if comparison.mismatched + else None + ), + duration_s=elapsed, + stdout_tail=( + comparison.excerpt(names) + if comparison.a and comparison.b + else f"native: {comparison.a[:2]}\nbf16: {comparison.b[:2]}" + ), + ) + return GateResult( + name=name, + member=member, + passed=True, + detail=f"native storage at {stored} against {min(reference)} B/token in bf16; " + f"{len(comparison.a)} continuation(s) token-identical (compared as " + f"{comparison.observable}){_schedule_note(max_batch_size)}", + command=native_cmd, + duration_s=elapsed, + ) + + +__all__ = [ + "DEVICE_MEMORY_MARKER", + "KV_BYTES_FIELD", + "LAUNCHER_EXPORTS", + "LEG_B_ARGV_ENV", + "BATCH_SIZE_FLAG", + "LEG_B_DRAFT_LEN_ENV", + "LEG_B_LEGS_ENV_ENV", + "LEG_B_MAX_BATCH_SIZE_ENV", + "LEG_B_MAX_TOKENS_ENV", + "LEG_B_MEMBER_ENV", + "LEG_B_PROMPTS_FILE_ENV", + "LEG_B_TIMEOUT_ENV", + "PREFIX_OBSERVABLE", + "REPLAY_COUNTER_FIELDS", + "WITNESS_COUNTER_ENV", + "WITNESS_COUNTER_MARKER", + "WITNESS_OUTPUT_ENV", + "ContinuationComparison", + "compare_continuations", + "contrast_leg_command", + "filler_words", + "first_difference", + "parse_device_memory_mib", + "parse_generated_texts", + "parse_witness_counters", + "parse_witness_texts", + "parse_witness_token_ids", + "run_equivalence_witness", + "run_kv_quant_witness", + "run_replay_witness", + "run_residency_witness", + "with_max_batch_size", + "witness_command", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py b/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py new file mode 100644 index 000000000000..2e6cee7d9c43 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py @@ -0,0 +1,3191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The staircase-onboard orchestrator. + +Out of tree this was 381 lines of prose read by a language model. Here it is +Python, and the split is the whole point of the port: **phase ordering, +queue walking, gate sequencing and disk verification are control flow**, so +they belong in code where they are deterministic, testable and free. What +stays prose is judgement — how to derive a forward, when a contract covers a +value, what a ceiling is — and that lives in the role prompts and the skills. + +The ordering is chosen to make the expensive thing happen early: + +1. ``preflight`` — Python. Seconds. A broken host costs seconds, not a + plan phase and a first assembler turn. +2. ``plan`` — capped, and **non-blocking**. Reaching the cap writes + the reviewer's open items into the plan and starts the build anyway. +3. ``smith`` — concurrent within a dependency batch. Catalog entries + are independent by construction, so this is where extra compute buys + wall-clock. +4. ``assemble`` — serial, core first, gates between units. +5. ``deliver`` — Python verification, one auditor turn, one lint pass. + +Every stage checkpoints before advancing, so a crash resumes at the phase +that was running rather than at the beginning. +""" + +from __future__ import annotations + +import hashlib +import re +import shutil +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import replace +from pathlib import Path +from typing import Any + +import yaml +from rich.markup import escape as escape_markup + +from agent_flow import ( + CLAUDE_CODE_DEFAULT_MODEL, + CODEX_DEFAULT_MODEL, + AgentLayer, + AgentLayerConfig, + BackendConfig, + SessionConfig, + require_tool_call_stop_hook, +) +from agent_flow.console import print_message, print_rule +from agent_flow.logger import get_logger +from agent_flow.runtime import CommandExecutor, TurnExecutor, TurnRequest + +from . import gates as gates_mod +from . import remediation +from .committer import install_committer +from .delivery import summarize_for_auditor, verify_delivery +from .environment import ( + RunEnvironment, + image_import_warning, + resolve_environment, + shares_allocation_with_spool, +) +from .execution import select_command_executor, select_turn_executor, shell_runner_for +from .journal import Journal +from .preflight import ( + MAX_TMPDIR_LEN, + PreflightError, + format_preflight_report, + run_preflight, + substrate_gate, +) +from .progress import ( + ASSEMBLER_DECISIONS, + AUDITOR_DECISIONS, + PLAN_DRAFTER_DECISIONS, + PLAN_REVIEWER_DECISIONS, + REVIEWER_DECISIONS, + SMITH_DECISIONS, + ProgressContext, + build_progress_tools, + collect_decision_file, + decision_path, + init_progress_file, + latest_decision, + latest_entry, + mark_decision_dispatch, +) +from .prompts import PromptBundle, build_staircase_onboard_prompts +from .runtime import ( + areas_for_demands, + format_areas, + format_kernel_sources, + load_index, + load_kernel_sources, + unmatched_demands, +) +from .state import ( + ITEM_DEFERRED, + ITEM_DONE, + ITEM_FAILED, + ITEM_PENDING, + ITEM_RUNNING, + STAGE_ASSEMBLE, + STAGE_DELIVER, + STAGE_PLAN, + STAGE_PREFLIGHT, + STAGE_SMITH, + STATE_FILENAME, + WorkflowState, + is_remediable, + load_state, + save_state, +) +from .survey import format_surface, read_checkpoint_surface, run_inventory_probe, write_surface_file +from .task_schema import ( + DEFAULT_BYTES_TOLERANCE, + REPLAY_WITNESS_REQUIRED, + WITNESS_BATCH_SIZE_KEY, + WITNESS_LEGS_ENV_KEY, + assembly_units, + load_and_validate_task_yaml, + members_for, +) +from .vocabulary import ( + MODE_VERIFY, + UNIT_CORE, + QueueItem, + VocabularyError, + dependency_batches, + load_vocabulary, + needed_by, +) + +# Which backend each role runs on. The PlanReviewer is deliberately not the +# same model as the PlanDrafter — a reviewer that shares the author's blind +# spots reviews its own reasoning back at itself. Every other role needs the +# skill system, which is a Claude Code feature. +_CODEX_ROLES = frozenset({"plan_reviewer"}) + +#: The role that submits a command to a scheduler the orchestrator cannot +#: reach. It is the only one that returns no typed decision, because it makes +#: none — see `prompts/runner.py` and `runtime/commands/delegated.py`. +RUNNER_ROLE = "runner" + +# Its presence, not its contents, is what matters: the Claude Agent SDK +# refuses dynamically-configured MCP servers whenever an enterprise config +# exists, so the workflow's in-process decision tools cannot be created. +ENTERPRISE_MCP_CONFIG = Path("/etc/claude-code/managed-mcp.json") + +#: How many smiths may run at once by default. +# +# The previous default was ``0``, which reads as "no limit" and *means* "open +# the whole batch at once". On the run this workflow was built against that +# would have been **19 concurrent agents** — the queue held 19 entries and +# every one of them declared no dependencies, so the layering produced a +# single batch. The run that actually happened was hand-batched to 6 and then +# 2, so the unbounded path was never exercised. +# +# 10 is chosen against what actually binds first, which is not memory. On a +# 90-core / 376 GB node with ~220 GB free, a dozen Node.js agent processes is +# not the constraint. The constraints are: +# +# * the upstream rate limit, which this stack *reports and does not act on* +# — there is no backoff anywhere in the agent layer, so N concurrent +# sessions that hit a cap fail together rather than queueing; +# * the credential refresh, which was measured corrupting a shared copy +# under concurrency; +# * shared-filesystem metadata, since every agent greps the same Lustre +# tree. +# +# So the number wants to be bounded and modest rather than as large as the +# machine allows. 10 sits above the 6 that was proven by hand and well under +# the 19 that has never been tried. +DEFAULT_MAX_PARALLEL_SMITHS = 10 + +# Units whose feature is distribution-preserving BY CONSTRUCTION, and which +# therefore need the acceptance gate on top of boot and accuracy. +# +# Speculative decoding is the whole membership. Rejection sampling holds the +# emitted distribution to the target model's, so a draft path that computes +# the wrong thing produces correct text more slowly: boot passes, accuracy +# passes, and the only observable that moves is speed. `acceptance_length` +# under `=require` versus `=off` is the detector. +# +# `core` is deliberately absent: it has no draft path, so the comparison +# would be 1.0 against 1.0 and would pass whatever core did. +ACCEPTANCE_GATED_UNITS: frozenset[str] = frozenset({"mtp", "eagle3", "dflash", "dspark"}) + +# The other output-identical features, and the witness each needs. Same +# argument as ACCEPTANCE_GATED_UNITS, different observable: boot and accuracy +# both pass whether or not the feature is doing anything, so a unit with no +# third gate is a unit that can be declared green for a feature nobody built. +# +# host_offload -> device memory after weight load, against a delta predicted +# from the checkpoint's own shapes, plus token identity +# disagg -> token identity across the prefill/decode hop +# kv_quant -> stored bytes per token against a declared figure, plus +# token identity with bf16 storage when declared exact +# +# And the two that are NOT output-identical but still need one, because the +# truncation they make is designed to stay inside the accuracy tolerance and a +# unit that makes no truncation stays inside it too: +# +# decoder_replay / encoder_replay -> per-request row counters the target +# prints, against a declared window and workload +# +# Each needs something from the spec that the orchestrator cannot derive: the +# environment for the contrast leg, and for most a declared figure to judge +# against. That lives in `accuracy.witness`, per unit. +WITNESS_GATED_UNITS: dict[str, str] = { + "host_offload": "residency", + "disagg": "equivalence", + "decoder_replay": "replay", + "encoder_replay": "replay", + "kv_quant": "kv_quant", +} + +# Shown on every infrastructure failure, and pointed at the file rather than +# summarized, because the summary is the part that goes stale. A run stops +# here perhaps once a week; the operator reading that message has no reason to +# know this document exists unless it says so at the moment they need it. +TROUBLESHOOTING_HINT = ( + "This is an infrastructure failure, not a target defect — every entry in " + f"`{Path(__file__).parent / 'TROUBLESHOOTING.md'}` read as a model problem " + "while it was happening. Search it for the error text above." +) + + +#: The preventive counterpart to ``TROUBLESHOOTING.md``, written into the +#: workspace and referenced by every dispatch prompt. +#: +#: The two are not the same document and must not be merged. +#: ``TROUBLESHOOTING.md`` is keyed on the verbatim error text a run already +#: hit; this is keyed on nothing, because it is meant to be read *before* +#: anything fails. On the run that motivated it, six host facts — pytest's +#: exit code, the ``TMPDIR`` bound, the multi-rank launch rules, which UCX +#: line is noise, what ``$?`` means after a pipe, where a venv interpreter +#: resolves — were re-derived by agent after agent and hand-injected into +#: roughly fifteen dispatch briefs. +HOST_FACTS_FILENAME = "host_facts.md" + +#: Everything below this line in ``host_facts.md`` belongs to the roles. +#: Preflight runs again on every resume, so the seeded half is rewritten +#: each time and anything a role learned would otherwise be erased by the +#: next restart — which is the same defect this file exists to fix, one +#: level down. +HOST_FACTS_APPEND_MARKER = "" + +#: Findings that landed after ``plan.md`` and the dispatch queue were +#: written, and therefore cannot reach an in-flight dispatch any other way. +#: +#: Invented by hand mid-run as ``SMITH_CORRECTIONS.md``, because the round-2 +#: plan review finished after the smith queue had been built. One of its +#: corrections was load-bearing: without it a smith would have certified an +#: FP8 GEMM at ``N=192``, a width no rank ever passes. +CORRECTIONS_FILENAME = "corrections.md" + +#: Where the standing header ends and the corrections begin. A file holding +#: only the header is an *empty* corrections file, and an empty one is not +#: referenced in any prompt — a pointer at nothing costs attention and +#: teaches the reader the pointer is usually worthless. +CORRECTIONS_BODY_MARKER = "" + + +def _compose_required_tools_hooks(required_tools: list[str]) -> dict | None: + """Stop hooks requiring *every* listed tool, not just one of them. + + ``require_tool_call_stop_hook`` enforces "at least one of these names was + called". One independent hook per tool therefore yields AND semantics. + """ + if not required_tools: + return None + merged: dict[str, list] = {"Stop": []} + for name in required_tools: + merged["Stop"].extend(require_tool_call_stop_hook([name])["Stop"]) + return merged + + +class StaircaseOnboardWorkflow: + """Drive one checkpoint to a gated modeling_v2 target.""" + + def __init__( + self, + workspace: Path, + *, + max_plan_rounds: int = 1, + max_assembler_turns_per_unit: int = 6, + max_parallel_smiths: int = DEFAULT_MAX_PARALLEL_SMITHS, + reviewer_mode: str = "light", + plan: str | Path | None = None, + vocabulary: str | Path | None = None, + clean: bool = False, + dry_run: bool = False, + skip_substrate_gate: bool = False, + prompts: PromptBundle | None = None, + command_runner: Any | None = None, + shell_runner: Any | None = None, + agent_transport: Any | None = None, + ) -> None: + self.workspace = Path(workspace) + self.task_path = self.workspace / "task.yaml" + self.plan_path = self.workspace / "plan.md" + self.vocabulary_path = self.workspace / "vocabulary.yaml" + self.progress_path = self.workspace / "progress.yaml" + self.journal_path = self.workspace / "journal.md" + self.boot_spec_path = self.workspace / "boot.yaml" + self.host_facts_path = self.workspace / HOST_FACTS_FILENAME + self.corrections_path = self.workspace / CORRECTIONS_FILENAME + self.state_path = self.workspace / STATE_FILENAME + + self.max_plan_rounds = max(0, max_plan_rounds) + self.max_assembler_turns_per_unit = max(1, max_assembler_turns_per_unit) + self.max_parallel_smiths = max_parallel_smiths + self.reviewer_mode = reviewer_mode + self.preset_plan = plan + self.preset_vocabulary = vocabulary + self.dry_run = dry_run + self.skip_substrate_gate = skip_substrate_gate + self._prompts_override = prompts + + # Injected so every stage is testable without a GPU, a cluster, or a + # model. The defaults are the real ones. + self._command_runner = command_runner + # An explicit runner always wins; otherwise the right one depends on + # the resolved environment, which preflight has not produced yet. + # Left as the plain shell until then so a caller that never reaches + # preflight still has something that works. + self._shell_runner_override = shell_runner + # When set, role turns are delivered through the filesystem to a + # session that can authenticate, instead of being spawned here. + # See transport.py: the orchestrator keeps the stage machine, the + # queue, the decision validation and every gate; only the turn + # moves. + self._agent_transport = agent_transport + #: How a role turn is delivered, resolved once. Both transports + #: answer the same `TurnExecutor` interface, so nothing downstream + #: of `_run_turn` branches on which one is in force. + self._turns: TurnExecutor = select_turn_executor( + agent_transport, lambda request: self._layer(request.role) + ) + #: Where this run's commands land. Resolved from the environment in + #: `_run_preflight`, which is the first moment one exists; until + #: then nothing has a command to run. + self._command_executor: CommandExecutor | None = None + #: The vocabulary validator's message from the last failed round, + #: carried verbatim into the re-draft prompt. + self._vocabulary_errors = "" + #: True when this process resumed into an assembly already in + #: progress, so a decision written by the previous one is pending. + self._resumed_into_unit = False + # Queue items this process has itself dispatched a smith for; see + # `_run_one_smith` for why the first dispatch of a resumed item asks + # the decision file before asking an agent. + self._smiths_dispatched: set[int] = set() + #: Whether a queue row naming a unit this run does not build has + #: already been reported. Said once, not on every queue walk. + self._warned_unknown_units = False + #: Lazily resolved in `_gate_env`, because the spec may give gates + #: their own launch substrate via `environment.workloads.gate`. + self._gate_environment: RunEnvironment | None = None + self._shell_runner = shell_runner or gates_mod._default_shell_runner + + # Whether a role's verdict travels as a FILE rather than as a call to + # its typed-decision MCP tool. Two independent reasons, and missing + # either one silently discards every decision the run makes. + # + # An enterprise MCP configuration makes the SDK refuse *every* + # dynamically-registered server, including the in-process one the + # typed-decision tools live in ("You cannot dynamically configure MCP + # servers when an enterprise MCP config is present"). Detected rather + # than configured: it is a property of the host, the operator did not + # choose it, and the first symptom is otherwise an opaque + # `ProcessError: exit code 1` on the very first agent turn. + # + # **And an external transport has no session to host those tools in at + # all**, whatever the host's MCP configuration says. That was missing + # here, and the failure is worth stating because nothing reports it: + # `_dispatch` then neither tells the role where to write its decision + # nor collects one, and `_run_deliver_phase` reads `progress.yaml` and + # gets **whatever was there before** — the previous run's verdict. + # + # Measured on this project: the orchestrator was moved into a + # container (where `/etc/claude-code/managed-mcp.json` does not exist) + # while running on the external transport, and the next three audit + # passes had their verdicts discarded. Each wrote a valid decision + # file; the run re-closed on a stale `UNCOVERED` from two passes + # earlier, including over a `CLEAN`. A loop like that can never + # converge, and it looks exactly like a gate that keeps failing. + self._decision_files = ENTERPRISE_MCP_CONFIG.exists() or agent_transport is not None + + self.workspace.mkdir(parents=True, exist_ok=True) + if clean: + for path in ( + self.state_path, + self.plan_path, + self.vocabulary_path, + self.progress_path, + self.journal_path, + # Both are this run's context, not the operator's notes: the + # facts are re-seeded by the next preflight, and a correction + # written against a plan that no longer exists is worse than + # no correction, because it still claims precedence over + # `vocabulary.yaml`. + self.host_facts_path, + self.corrections_path, + ): + path.unlink(missing_ok=True) + shutil.rmtree(self.workspace / "reports", ignore_errors=True) + # Regenerated by the next preflight, but a run that fails before + # reaching it would otherwise leave the previous checkpoint's + # survey sitting in the workspace looking current. + (self.workspace / "surface.md").unlink(missing_ok=True) + + self.resume = self.state_path.is_file() + if not self.resume: + init_progress_file(self.progress_path) + + self.journal = Journal(path=self.journal_path) + self.gate_log = gates_mod.GateLog() + self._progress_ctx = ProgressContext(path=self.progress_path) + self._tools = build_progress_tools(self._progress_ctx) + + self.task_data: dict[str, Any] = {} + self.environment: RunEnvironment | None = None + self.prompts: PromptBundle = prompts or build_staircase_onboard_prompts() + self._layers: dict[str, AgentLayer] = {} + # Set by ``_run_preflight``, which every path through ``run`` reaches + # before anything reads it — including a resume, which re-derives the + # identity rather than trusting a stale checkpoint. + self._preflight: Any = None + # The surface survey. Empty rather than ``None`` so a resume that + # re-enters at the plan phase still composes a valid prompt: the + # drafter then reads "no survey was taken", which is true, instead of + # the orchestrator crashing on an attribute. + self._surface: Any = None + self._inventory: Any = None + self._surface_block: str = ( + "## Surface survey\n\nNot taken for this run — plan on the assumption that " + "you do not know what exists, and say so in the plan's open items.\n" + ) + self._runtime_block: str = "" + self._kernel_sources_block: str = "" + + # ------------------------------------------------------------- lifecycle + + def __enter__(self) -> "StaircaseOnboardWorkflow": + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self.close() + + def close(self) -> None: + for layer in self._layers.values(): + layer.__exit__(None, None, None) + self._layers.clear() + + # ----------------------------------------------------------------- agents + + def _make_agent( + self, + name: str, + system_prompt: str, + *, + required_tools: list[str], + session_mode: str = "stateless", + ) -> AgentLayer: + backend_kind = "codex" if name in _CODEX_ROLES else "claude-code" + model = CODEX_DEFAULT_MODEL if name in _CODEX_ROLES else CLAUDE_CODE_DEFAULT_MODEL + # Under an enterprise MCP configuration the in-process tool server + # cannot be registered at all, so the role reports by writing a file + # and the Stop hook requires the write rather than the tool. `Write` + # is a weaker hook than a named tool — the role uses it for other + # things too — but it still refuses a turn that wrote nothing, and + # the enum is validated when the file is read. + tools = None if self._decision_files else self._tools.get(name) + # The `Write` fallback is for roles whose decision travels as a file. + # A role with no decision at all gets no hook either way. + hooks = _compose_required_tools_hooks( + (["Write"] if self._decision_files else required_tools) if required_tools else [] + ) + return AgentLayer( + AgentLayerConfig( + name=name, + system_prompt=system_prompt, + backend=BackendConfig( + kind=backend_kind, + model=model, + tools=tools, + hooks=hooks, + # The workspace is the session cwd so the run's own + # materialized skills are the ones the agent discovers, + # rather than whatever happens to be configured in the + # repository it is editing. + cwd=self.workspace, + ), + session=SessionConfig(mode=session_mode), + ) + ) + + def _run_turn( + self, + role: str, + what: str, + prompt: str, + *, + iteration: int | None = None, + layer: Any | None = None, + ) -> str: + """Run one role turn, wherever this run's transport delivers it. + + There is no transport branch left in this method, and that is the + point. A host's inability to authenticate a child process changed + one thing -- where the prompt goes -- and it used to change a line + here, in the file that also owns resume, gate sequencing and the + queue walk. Both transports now answer `TurnExecutor.run`, so + everything the orchestrator decides (which stage runs, which entry + is next, whether a gate passed) is on this side of it either way, + and provably so: this method cannot tell them apart. + + ``layer`` is the caller's own session, which a concurrent batch owns + rather than the executor -- ten smiths are ten agents, not ten turns + of one. `bound` is a no-op for transports with no session concept. + """ + request = TurnRequest( + role=role, + what=what, + prompt=prompt, + system_prompt=getattr(self.prompts, role, "") if self.prompts else "", + decision_file=( + str(decision_path(self.workspace, role, iteration)) + if self._decision_files and iteration is not None + else None + ), + iteration=iteration, + ) + executor = self._turns if layer is None else self._turns.bound(layer) + return executor.run(request).text + + def _layer(self, name: str) -> AgentLayer: + """Build a role's layer on first use and keep it for the run. + + Plan roles are **persistent**: they are asked to remember what they + decided and why, so a re-draft addresses the reviewer's items instead + of re-deriving everything. Work roles are **stateless** — one unit of + work per run, each in its own context — except the assembler, which + is persistent *within* a unit and reset between units, because a + unit is its unit of work. + """ + if name not in self._layers: + persistent = name in ("plan_drafter", "plan_reviewer", "assembler") + # The runner is the one role with no typed decision: it relays a + # submission and the verdict comes from files the job wrote. A + # Stop hook demanding a decision tool it was never given would + # block its turn forever. + required = [] if name == RUNNER_ROLE else [f"append_{name}_progress"] + self._layers[name] = self._make_agent( + name, + getattr(self.prompts, name), + required_tools=required, + session_mode="persistent" if persistent else "stateless", + ) + return self._layers[name] + + def _reset_layer(self, name: str) -> None: + layer = self._layers.pop(name, None) + if layer is not None: + layer.__exit__(None, None, None) + + def _materialize_skills(self) -> Path: + """Copy the packaged skills into this run's workspace. + + Shipping them with the workflow rather than relying on a checkout's + ``.claude/skills`` means the knowledge is versioned with the code + that dispatches it, and copying them per run means a run is + reproducible from its own workspace. + """ + source = Path(__file__).parent / "skills" + destination = self.workspace / ".claude" / "skills" + if not source.is_dir(): + return destination + destination.mkdir(parents=True, exist_ok=True) + for skill_dir in sorted(p for p in source.iterdir() if p.is_dir()): + target = destination / skill_dir.name + shutil.rmtree(target, ignore_errors=True) + shutil.copytree(skill_dir, target) + return destination + + # ------------------------------------------------------------ orchestration + + def run(self, task: str | Path) -> None: + log = get_logger().console + state = self._init_state(task, log) + if state is None: + return + + try: + if state.stage == STAGE_PREFLIGHT: + if not self._run_preflight(state, log): + return + if self.dry_run: + print_message( + "[bold cyan]→ --dry-run: preflight passed; stopping before the " + "first agent turn.[/bold cyan]", + log, + ) + return + state.stage = STAGE_PLAN + self._checkpoint(state) + else: + # Resuming past preflight still needs the resolved + # environment, since every prompt and every gate carries it. + if not self._run_preflight(state, log): + return + + if state.stage == STAGE_PLAN: + if not self._run_plan_phase(state, log): + return + state.stage = STAGE_SMITH + self._checkpoint(state) + + if state.stage == STAGE_SMITH: + if not self._run_smith_phase(state, log): + return + state.stage = STAGE_ASSEMBLE + self._checkpoint(state) + + if state.stage == STAGE_ASSEMBLE: + if not self._run_assemble_phase(state, log): + return + state.stage = STAGE_DELIVER + self._checkpoint(state) + + if state.stage == STAGE_DELIVER: + if not self._run_deliver_phase(state, log): + return + + state.done = True + self._checkpoint(state) + self._close_run(state, log, verdict="delivered") + except KeyboardInterrupt: + print_message( + "[bold yellow]⚠ interrupted — re-run the same command to continue " + "from the last checkpoint, or pass --clean to start fresh" + "[/bold yellow]", + log, + ) + self.journal.event("run interrupted", f"stage={state.stage}") + raise + + # ------------------------------------------------------------------ stages + + def _run_preflight(self, state: WorkflowState, log) -> bool: + print_rule("[bold magenta]preflight[/bold magenta]", log) + kwargs: dict[str, Any] = {} + if self._command_runner is not None: + kwargs["runner"] = self._command_runner + try: + result = run_preflight(self.task_data, **kwargs) + except PreflightError as exc: + print_message(f"[bold red]✘ preflight: {exc}[/bold red]", log) + print_message(TROUBLESHOOTING_HINT, log) + self.journal.open(slug="(preflight failed)") + self.journal.section("preflight failed", f"{exc}\n\n{TROUBLESHOOTING_HINT}") + return False + + state.slug = result.slug + state.branch = result.branch + state.devices = list(result.devices) + self.environment = result.environment + self._preflight = result + + self.journal.open( + slug=result.slug, + header_lines=[ + f"- checkpoint: `{self.task_data['checkpoint_path']}`", + f"- target: `{result.target_rel}`", + f"- units: {', '.join(assembly_units(self.task_data))}", + ], + ) + self.journal.section("preflight", f"```\n{format_preflight_report(result)}\n```") + print_message( + f"[bold green]✔ preflight[/bold green]\n{format_preflight_report(result)}", log + ) + + # Where this run's commands land, decided once, from the resolved + # environment. The orchestrator used to carry the mapping itself -- + # a comparison against one environment kind, inline, twice -- so + # adding a mechanism meant editing the file that owns resume and + # gate sequencing. It now asks `execution.py` and does not know the + # answer's name. + # + # An explicitly injected runner still wins: a test supplies one, and + # so does an operator debugging a gate. + self._command_executor = select_command_executor( + result.environment, + # Handed over unconditionally, and used by exactly one executor: + # the delegated one, whose scheduler is reachable through a tool + # the agent holds and this process does not. Submission becomes a + # turn; the verdict still comes from files the job wrote, read + # here. Which kind that is remains `execution.py`'s business. + turns=self._turns, + workspace=self.workspace, + ) + if self._shell_runner_override is None: + self._shell_runner = shell_runner_for(self._command_executor) + + # Two configuration hazards that are legal, cost hours, and are + # invisible until they bite. Appended to the notes rather than raised: + # both are about how this run will be SHAPED, not about whether it can + # start, and a spec that accepts the cost knowingly is allowed to. + # + # - a registry container image is re-imported per allocation -- + # measured 5-7 min for 26 GB on the predecessor project, and + # 25-35 min under load on this one, three submissions running. + # That cost is what pushes a run onto one shared session. + # - the orchestrator living inside the session it dispatches to. + # They then die together, and it needs no GPU. Three times in one + # run on the predecessor, twice here. + for warning in ( + image_import_warning(result.environment), + shares_allocation_with_spool(result.environment), + ): + if warning: + result.notes.append(warning) + print_message(f"[bold yellow]⚠ {warning}[/bold yellow]", log) + + # Prove the host can generate a token under this exact environment + # before any agent exists. Skipped on a dry run, whose whole purpose + # is to answer "would this spec resolve" cheaply. + if not self.dry_run and not self.skip_substrate_gate: + try: + substrate_gate( + self.task_data, + result.environment, + shell_runner=self._shell_runner, + notes=result.notes, + ) + except PreflightError as exc: + # Reported like any other precondition failure: the run ends + # here, before an agent exists, with the command and output + # the operator needs to reproduce it. + print_message(f"[bold red]✘ {exc}[/bold red]", log) + print_message(TROUBLESHOOTING_HINT, log) + self.journal.section("substrate gate failed", f"{exc}\n\n{TROUBLESHOOTING_HINT}") + return False + + self._take_survey(result, log) + + # Prompts can only be built now: every role's prompt carries the + # resolved launch prefix, which did not exist until the device claim + # was made. + # Concurrent smiths share one worktree and one `.git/index`. The + # rule "stage by explicit path" was never enough on its own — it was + # followed and still produced two cross-attributed commits — so the + # rule is now a script that takes a lock and refuses a commit whose + # contents are not exactly what was asked for. + commit_script = install_committer(self.workspace, self.task_data["trtllm_repo_path"]) + + if self._prompts_override is None: + self.prompts = build_staircase_onboard_prompts( + environment=result.environment, + reviewer_mode=self.reviewer_mode, + source_embargo=self.task_data.get("source_embargo") or (), + scope=str(self.task_data.get("scope") or "full"), + decision_files=self._decision_files, + commit_script=str(commit_script), + ) + # Written here, from what preflight has just resolved, so that the + # first agent turn of the run already has them. Both files are + # referenced by `_context_block`, which every dispatch carries. + self._write_host_facts(result) + self._ensure_corrections_file() + self._materialize_skills() + self._checkpoint(state) + return True + + def _write_host_facts(self, result) -> None: + """Seed ``host_facts.md`` from what this run already knows. + + Seeded rather than hardcoded wherever the run can answer for itself: + the ``TMPDIR`` bound and its current value come from + :data:`preflight.MAX_TMPDIR_LEN` and the resolved + :class:`environment.RunEnvironment`, the multi-rank paragraph is + emitted only when :attr:`environment.RunEnvironment.needs_llmapi_launch` + says this run has ranks to launch, and preflight's own notes are + carried verbatim. The rest are defaults, and a seeded default is + better than nothing — each of them was learned the expensive way by + an agent on this host. + + Rewritten on every preflight, which a resume also runs, so anything + a role appended is preserved below + :data:`HOST_FACTS_APPEND_MARKER` and everything above it is + regenerated. + """ + env = result.environment + tmpdir = str((env.env or {}).get("TMPDIR", "") or "") + lines = [ + "# Host facts", + "", + "What this host does that a generic Linux box does not. **Read this before " + "you debug an environment failure** — every fact here was learned the " + "expensive way by an earlier dispatch on this machine.", + "", + f"This is the *preventive* half. `{Path(__file__).parent / 'TROUBLESHOOTING.md'}` " + "is the other half and is keyed on the verbatim error text of a failure that " + "already happened. Keep them separate: one is read before anything breaks, " + "the other is grepped afterwards.", + "", + "## Facts", + "", + "- **`pytest` exits 4 on this host even when every test passed.** The venv " + "lives inside the repository root, so a session-final `sys.path` audit " + "fires. Export `TRTLLM_SKIP_SYS_PATH_CHECK=1`, and adjudicate on the " + "summary line (`N passed`), never on the exit code.", + f"- **`TMPDIR` must be at most {MAX_TMPDIR_LEN} characters.** The LLM API's " + "RPC binds a unix socket under it, the path is capped at 107 bytes and the " + "socket basename takes 45. This run's `TMPDIR` is " + + (f"`{tmpdir}` ({len(tmpdir)} characters)." if tmpdir else "not set in the spec."), + "- **`UCX ERROR uct_iface_open(tcp/rdma_vf_rail0)` is noise.** UCX falls back " + "and the run exits 0. Do not chase it; it is not why your job failed.", + "- **`RC=$?` after a pipe is the *pipe's* exit code, not the command's.** " + "`cmd | tee log; RC=$?` reports `tee`. Use `PIPESTATUS[0]`, or drop the pipe.", + "- **A venv interpreter created with `--system-site-packages` resolves only " + "inside the container it was created in.** Running it from the host finds an " + "interpreter and none of the packages, which reads as a broken install.", + ] + if env.needs_llmapi_launch: + lines.append( + f"- **A {env.ntasks}-rank launch needs the `SLURM_*` scrub and " + f"`--oversubscribe`.** Without them every rank reads `mpi_rank=0`, every " + f"rank becomes a leader, and the job hangs forever with no error and no " + f"traceback. Both are already in the launch prefix you were handed — use " + f"it verbatim rather than writing your own `mpirun` line." + ) + if env.container_image: + lines.append( + f"- **The container is `{env.container_image}`.** An interpreter path in " + f"this run's environment block is a path *inside* it." + ) + if result.notes: + lines.extend( + ["", "## What preflight measured on this host", ""] + + [f"- {note}" for note in result.notes] + ) + lines.extend( + [ + "", + "## Learned during this run", + "", + "This section is yours. If you spend more than a few minutes discovering " + "something about this host that the section above would have saved you, " + "append it here as one bullet: the symptom first, then what to do. The " + "orchestrator never rewrites anything below the marker, so it survives " + "every restart and reaches every later dispatch.", + "", + HOST_FACTS_APPEND_MARKER, + "", + ] + ) + head = "\n".join(lines) + + appended = "" + if self.host_facts_path.is_file(): + existing = self.host_facts_path.read_text(encoding="utf-8") + if HOST_FACTS_APPEND_MARKER in existing: + appended = existing.split(HOST_FACTS_APPEND_MARKER, 1)[1].lstrip("\n") + self.host_facts_path.write_text(head + appended, encoding="utf-8") + + def _ensure_corrections_file(self) -> None: + """Create the corrections file's standing header if it is missing. + + The file exists from the start of the run so that anyone — a role, a + later phase, the operator watching a review land — has a known place + to put a correction. It is empty until someone does, and + `_corrections_block` keeps an empty one out of every prompt. + """ + if self.corrections_path.is_file(): + return + self.corrections_path.write_text( + "# Corrections\n\n" + "Findings that landed **after** `plan.md` and `vocabulary.yaml` were " + "written, and which therefore cannot reach a dispatch any other way. The " + "orchestrator appends here when a review lands late; the operator may " + "append by hand at any time, and every dispatch re-reads this file, so a " + "correction written now reaches the next turn.\n\n" + "**Precedence: where this file and `vocabulary.yaml` disagree, this file " + "wins.** It is newer by construction.\n\n" + "Each entry: what was believed, what is true instead, and what to do.\n\n" + f"{CORRECTIONS_BODY_MARKER}\n", + encoding="utf-8", + ) + + def _corrections_body(self) -> str: + """Whatever has been written below the standing header. May be empty.""" + if not self.corrections_path.is_file(): + return "" + text = self.corrections_path.read_text(encoding="utf-8") + if CORRECTIONS_BODY_MARKER not in text: + return text.strip() + return text.split(CORRECTIONS_BODY_MARKER, 1)[1].strip() + + def _record_correction(self, source: str, body: str) -> None: + """Append one correction, so it reaches dispatches already queued.""" + body = body.strip() + if not body: + return + self._ensure_corrections_file() + with self.corrections_path.open("a", encoding="utf-8") as handle: + handle.write(f"\n## From {source}\n\n{body}\n") + self.journal.section(f"correction recorded ({source})", body) + + def _context_block(self) -> str: + """The two files every dispatch carries by reference. + + By reference rather than inlined: both are re-read at the moment the + role opens them, so a correction written while a turn is in flight + still lands, and a host fact appended by dispatch 7 is there for + dispatch 8. Inlining would freeze both at prompt-composition time, + which is exactly the failure — `task.yaml` read once — that this run + paid six wasted assembler dispatches for. + """ + parts: list[str] = [] + if self.host_facts_path.is_file(): + parts.append( + f"## Host facts\n\n" + f"Read `{self.host_facts_path}` before you debug anything about this " + f"machine. It carries what this host does differently — `pytest`'s exit " + f"code, the `TMPDIR` bound, the launch rules, which errors are noise. If " + f"you learn a new one, append it under `## Learned during this run`." + ) + if self._corrections_body(): + parts.append( + f"## Corrections\n\n" + f"`{self.corrections_path}` holds findings that landed after the plan and " + f"the dispatch queue were written. Read it. **Where it and " + f"`vocabulary.yaml` disagree, the corrections file wins.**" + ) + return "\n\n".join(parts) + + def _take_survey(self, result, log) -> None: + """Measure what the checkpoint demands and what the build supplies. + + Runs at the end of preflight rather than at the top of the plan phase + so that its cost is paid while nothing is waiting on it, and so that + a resume — which re-enters at the stage that was running — does not + re-pay it. + + **Never fatal.** A plan can be drafted without a survey; it just gets + drafted on assumption, which is how every plan before this was + drafted. Losing a run because a probe could not import is a worse + trade than losing the measurement, so the failure is written into the + block the drafter reads and the phase continues. + """ + surface = read_checkpoint_surface(result.checkpoint_config) + # Taken on a dry run too. It was skipped there while the probe went + # out as engine work — a rank per device, through the launcher — + # because a dry run answers "would this spec resolve" without + # claiming the machine for minutes. `probe_command` now asks for one + # process and no MPI world, so the measurement costs an import and a + # directory walk; and a dry run that prints `0/0 entries certified` + # against a catalog holding 76 is the exact reading this survey + # exists to prevent, whoever is looking at it. + inventory = run_inventory_probe( + self.task_data, + result.environment, + shell_runner=self._shell_runner, + ) + self._surface = surface + self._inventory = inventory + self._surface_block = format_surface(surface, inventory) + + # The runtime map, narrowed to what this checkpoint actually demands. + # Rendered here rather than in the prompt because it is derived from a + # measurement, and because a map nobody narrowed is a catalogue. + demands = surface.demand_names + index = load_index() + areas = areas_for_demands(demands, index=index) + self._runtime_block = format_areas( + areas, unmatched=unmatched_demands(demands, index=index), overview=index.overview + ) + # Provenance, for the smiths. The survey says which ops exist; this + # says what each source IS and what certifying one of its kernels + # requires, which is what decides the shape of a receipt. + self._kernel_sources_block = format_kernel_sources(load_kernel_sources()) + write_surface_file(self.workspace / "surface.md", surface, inventory) + + headline = ( + f"{len(surface.demands)} structural demand(s), " + f"{len(surface.unrecognized)} unrecognized key(s); " + + ( + f"{len(inventory.spendable)}/{len(inventory.catalog)} entries certified on " + f"{inventory.arch}, {len(inventory.unwrapped_ops)} unwrapped op(s)" + if inventory.ok + else f"inventory probe failed: {inventory.error.splitlines()[0]}" + ) + ) + self.journal.section("surface survey", headline + "\n\nFull survey: `surface.md`") + # A failed measurement is not a green line. `✔ survey ... 0/0 entries + # certified` is how a probe that never ran reads identically to a + # catalog with nothing in it, and an operator has no way to tell them + # apart. The run still PROCEEDS — a plan drafted on assumption beats + # no plan — it just proceeds loudly. + mark, colour = ("✔", "green") if inventory.ok else ("⚠", "yellow") + print_message(f"[bold {colour}]{mark} survey[/bold {colour}] {headline}", log) + if surface.unrecognized: + print_message(f" unrecognized structural keys: {', '.join(surface.unrecognized)}", log) + + def _run_plan_phase(self, state: WorkflowState, log) -> bool: + """Drafter <-> reviewer, capped, and never blocking the build.""" + if self.preset_plan is not None and self.preset_vocabulary is not None: + self._materialize(self.preset_plan, self.plan_path) + self._materialize(self.preset_vocabulary, self.vocabulary_path) + print_message( + "[bold cyan]→ using the supplied plan and vocabulary; skipping the " + "plan phase[/bold cyan]", + log, + ) + state.plan_approved = True + self._checkpoint(state) + return True + + while True: + round_no = state.plan_round + 1 + print_rule(f"[bold magenta]plan round {round_no}[/bold magenta]", log) + self._progress_ctx.current_iteration = round_no + + # A completed draft is worth resuming onto. The phase only + # checkpoints once the plan is *approved*, so without this a + # failure anywhere downstream in the same round — a reviewer + # whose backend is unauthenticated, a crash, an interrupt — + # throws away a drafter turn that had already succeeded and + # written both files. Measured at ~24 minutes and one oracle + # accuracy run, which is far too much to spend twice for a + # reason that had nothing to do with the drafter. + already_drafted = ( + latest_decision(self.progress_path, "plan_drafter", PLAN_DRAFTER_DECISIONS) + == "DRAFT_READY" + and (latest_entry(self.progress_path, "plan_drafter") or {}).get("iteration") + == round_no + and self.plan_path.is_file() + and self.vocabulary_path.is_file() + ) + if already_drafted: + print_message( + f"[bold cyan]→ round {round_no} was already drafted; reusing plan.md " + f"and vocabulary.yaml[/bold cyan]", + log, + ) + else: + self._dispatch( + "plan_drafter", + what=f"round {round_no}", + prompt=self._plan_drafter_prompt(round_no), + ) + decision = latest_decision(self.progress_path, "plan_drafter", PLAN_DRAFTER_DECISIONS) + if decision == "CEILING": + entry = latest_entry(self.progress_path, "plan_drafter") or {} + print_message( + "[bold red]✘ vocabulary ceiling — a computation step has no server " + "anywhere. The branch holds the state for a later retry." + "[/bold red]", + log, + ) + self.journal.section("vocabulary ceiling", str(entry.get("summary", ""))) + self._close_run(state, log, verdict="ceiling") + return False + + try: + # Against the spec's units, so a `unit:` tag naming a feature + # this run does not build is a re-draft now rather than a + # guess the orchestrator makes at the first failure. + vocabulary = load_vocabulary( + self.vocabulary_path, units=assembly_units(self.task_data) + ) + except Exception as exc: + # A malformed vocabulary is a re-draft, not a crash: the + # drafter is the only thing that can fix it, and it is still + # in session. + print_message(f"[bold yellow]↻ vocabulary.yaml unusable: {exc}[/bold yellow]", log) + # Carry the validator's own words into the re-draft. Without + # this the next prompt tells the drafter to read progress.yaml + # "for the PlanReviewer's REJECT" -- and there is no reviewer + # entry, because the schema rejected the file before the + # reviewer ran. The drafter would re-emit the same vocabulary + # and, at the default one-round cap, end the run. + self._vocabulary_errors = str(exc) + if state.plan_round >= self.max_plan_rounds: + print_message( + "[bold red]✘ plan budget exhausted with no usable " + "vocabulary.yaml[/bold red]", + log, + ) + self._close_run(state, log, verdict="plan-failed") + return False + state.plan_round += 1 + self._checkpoint(state) + continue + + if vocabulary.is_ceiling: + print_message("[bold red]✘ vocabulary.yaml records a ceiling[/bold red]", log) + self._close_run(state, log, verdict="ceiling") + return False + + try: + self._dispatch( + "plan_reviewer", + what=f"round {round_no}", + prompt=self._plan_reviewer_prompt(round_no), + ) + except Exception as exc: + # **The plan phase does not block the build** — that is the + # cap's whole premise, and a reviewer that cannot start is + # just a reviewer with nothing to say. Crashing here would + # throw away a completed draft over a second model's + # availability, which is the opposite of the intent. + # + # Seen for real: the reviewer runs on a different backend by + # design (author/reviewer separation), and that backend + # returned 401 on a host where only one of the two was + # authenticated. + print_message( + f"[bold yellow]! plan_reviewer could not run ({exc}). The draft " + f"stands unreviewed and the build starts; recorded as an open " + f"risk.[/bold yellow]", + log, + ) + self.journal.section("plan_reviewer unavailable", str(exc)) + self._append_open_risks() + state.plan_approved = True + state.queue = [ + {**item.as_dict(), "status": ITEM_PENDING} for item in vocabulary.queue + ] + self._checkpoint(state) + return True + + verdict = latest_decision(self.progress_path, "plan_reviewer", PLAN_REVIEWER_DECISIONS) + if verdict == "APPROVE": + print_message("[bold green]✔ plan approved — entering the build[/bold green]", log) + state.plan_approved = True + state.queue = [ + {**item.as_dict(), "status": ITEM_PENDING} for item in vocabulary.queue + ] + self._checkpoint(state) + return True + + state.plan_round += 1 + if state.plan_round > self.max_plan_rounds: + # The cap is a budget, not a gate. Carry the reviewer's + # unresolved items into plan.md so the assembler inherits + # them, and start the build. + self._append_open_risks() + print_message( + f"[bold yellow]→ plan budget spent ({self.max_plan_rounds} round(s)); " + f"the reviewer's open items are recorded in plan.md as risks and " + f"the build starts now[/bold yellow]", + log, + ) + state.plan_approved = True + state.queue = [ + {**item.as_dict(), "status": ITEM_PENDING} for item in vocabulary.queue + ] + self._checkpoint(state) + return True + + print_message("[bold yellow]↻ plan_reviewer REJECT — re-drafting[/bold yellow]", log) + self._checkpoint(state) + + def _run_smith_phase(self, state: WorkflowState, log) -> bool: + """Walk the dispatch queue, concurrently within a dependency batch. + + Every row is dispatched here, whichever unit first needs it: a + feature-only entry costs nothing extra to build alongside core's, and + most of them land. What a row's ``unit`` changes is only what its + failure blocks -- see `_walk_queue`. + """ + if not state.queue: + print_message("[bold green]✔ no catalog work — the queue is empty[/bold green]", log) + return True + if not self._walk_queue(state, log, unit=None): + return False + owed = [row for row in state.queue if row.get("status") == ITEM_DEFERRED] + if owed: + print_message( + f"[bold green]✔ every entry core needs has landed[/bold green] — " + f"{len(owed)} dispatch(es) deferred, each retried once before the first " + f"assembler turn of its unit: " + + ", ".join( + f"{row['n']}:{row.get('target')} → {row.get('deferred_to')}" for row in owed + ), + log, + ) + else: + print_message("[bold green]✔ catalog queue complete[/bold green]", log) + return True + + def _walk_queue(self, state: WorkflowState, log, *, unit: str | None) -> bool: + """Dispatch the rows due now, one dependency batch at a time. False if the run closed. + + One walk, two callers, and what differs is which rows are due and + which failures block. + + ``unit=None`` is the **smith phase**. Every row not yet attempted is + dispatched, but only a row core needs -- by its own ``unit`` tag, or + as a transitive dependency of a row core needs + (`vocabulary.needed_by`) -- can close the run. Any other failure is + recorded ``deferred`` to the first unit that needs it, and a row that + waits on one is deferred with it, undispatched. Measured on the + DeepSeek-V4.1 run: 33 dispatches queued, six needed only by later + feature units (dspark / multimodal / host_offload / kv_quant), and + any one of those six failing used to close the run `smith-blocked` + before core -- which needed none of them -- had started. + + ``unit=U`` is the **start of unit U**. Every row deferred to U, or to + a unit before it, gets one more attempt, and now every failure + blocks: U cannot be assembled without them. A row that was never + deferred is not touched here. The gap dispatches + `_handle_blocked_assembler` makes are its own, and retrying them here + as well would dispatch them twice. + """ + units = assembly_units(self.task_data) + horizon = 0 if unit is None else units.index(unit) + position = {name: index for index, name in enumerate(units)} + needed = self._needed_by(state, log) + rows = {row["n"]: row for row in state.queue} + items = [QueueItem.from_row(row) for row in state.queue] + status = {n: row.get("status", ITEM_PENDING) for n, row in rows.items()} + + def blocks(item: QueueItem) -> bool: + return position[needed[item.n]] <= horizon + + def due(item: QueueItem) -> bool: + if status[item.n] == ITEM_DONE: + return False + if unit is not None: + return bool(rows[item.n].get("deferred_to")) and blocks(item) + # A row core needs is dispatched whatever it last did, which is + # how a resume has always treated a failed row. Any other row is + # dispatched only while it has not had its attempt. + return blocks(item) or status[item.n] in (ITEM_PENDING, ITEM_RUNNING) + + def failure(item: QueueItem, exc: BaseException | None = None) -> str: + decision = self._smith_decision_for(item) + why = f"smith decided {decision}" if decision else "smith recorded no decision" + if exc is not None: + why += f"; the dispatch raised {type(exc).__name__}: {exc}" + return why + + def defer(item: QueueItem, why: str) -> str: + to = needed[item.n] + row = rows[item.n] + row.update(status=ITEM_DEFERRED, deferred_to=to, deferred_why=why) + status[item.n] = ITEM_DEFERRED + return f"{item.n}:{item.target} → {to} ({why})" + + def close(names: str, detail: str) -> bool: + where = "core" if unit is None else unit + print_message( + f"[bold red]✘ smith could not land {names}{detail}. {where} needs these " + f"entries, so its assembly cannot start.[/bold red]", + log, + ) + self.journal.section(f"smith blocked {where}", f"{names}{detail}") + self._close_run( + state, log, verdict="smith-blocked" if unit is None else f"smith-blocked:{unit}" + ) + return False + + for batch_index, batch in enumerate(dependency_batches(items), start=1): + pending: list[QueueItem] = [] + held: list[str] = [] + for item in batch: + if not due(item): + # A feature-only row that already failed -- in a process + # from before deferral existed, which closed the run over + # it -- has had its attempt. Record the deferral instead + # of paying for a second one now. + if unit is None and status[item.n] == ITEM_FAILED: + held.append(defer(item, f"{failure(item)}, in an earlier process")) + continue + waiting = [dep for dep in item.depends_on if status.get(dep) != ITEM_DONE] + if not waiting: + pending.append(item) + elif blocks(item): + return close(f"{item.n}:{item.target}", f" — it waits on {waiting}") + else: + held.append(defer(item, f"not dispatched: waits on {waiting}")) + if held: + self._checkpoint(state) + self._report_deferrals(held, log) + if not pending: + continue + + # `0` still means "open the whole batch", but it is now an + # explicit choice rather than the default, so say what it resolved + # to. A concurrency figure nobody printed is one nobody notices + # until the rate limit does. + workers = min(self.max_parallel_smiths, len(pending)) + if self.max_parallel_smiths <= 0: + workers = len(pending) + capped = "" if workers == len(pending) else f", {workers} at a time" + label = "smith batch" if unit is None else f"{unit} · deferred smith batch" + print_rule( + f"[bold magenta]{label} {batch_index} — " + f"{len(pending)} dispatch(es){capped}[/bold magenta]", + log, + ) + if self.max_parallel_smiths <= 0 and len(pending) > DEFAULT_MAX_PARALLEL_SMITHS: + print_message( + f"[bold yellow]! max_parallel_smiths is uncapped, so this batch opens " + f"{len(pending)} agent sessions at once. The rate limit is reported " + f"but not backed off, so a cap that is hit fails the whole batch " + f"together. The default is {DEFAULT_MAX_PARALLEL_SMITHS}.[/bold yellow]", + log, + ) + self._progress_ctx.current_iteration = batch_index + # Composed before the rows go `running`, from the reports the + # previous attempt left. The smith phase passes none, so its + # dispatch is exactly what it always was. + notes: dict[int, str] = {} + for item in pending: + notes[item.n] = self._retry_note(rows[item.n], item) if unit else "" + self._set_item_status(state, item.n, ITEM_RUNNING) + self._checkpoint(state) + errors: dict[int, BaseException] = {} + with ThreadPoolExecutor( + max_workers=max(1, workers), thread_name_prefix="smith" + ) as pool: + futures = { + pool.submit( + self._run_one_smith, *((item, notes[item.n]) if notes[item.n] else (item,)) + ): item + for item in pending + } + for future in as_completed(futures): + item = futures[future] + try: + future.result() + except BaseException as exc: # keep the other results + errors[item.n] = exc + + enqueued = 0 + failed: list[QueueItem] = [] + deferred: list[str] = [] + for item in pending: + decision = self._smith_decision_for(item) + if decision == "MACRO": + # MACRO is not a failure — it is the one outcome whose + # purpose is to say what the next work is. Scoring it as + # a failure ended the first real run while it held a + # 434-line decomposition that was exactly the queue + # extension it needed. See IMPROVEMENTS #18. + before = len(state.queue) + added = self._enqueue_constituents(state, item, log) + if unit is not None: + # Pieces of a row this unit is retrying are owed by + # the same retry: without the mark, nothing would + # ever dispatch them. + for row in state.queue[before:]: + row["deferred_to"] = row.get("unit") or unit + enqueued += added + ok = added > 0 + if not ok: + print_message( + f"[bold yellow]! dispatch {item.n} returned MACRO but listed no " + f"usable `constituents`, so there is nothing to enqueue. The " + f"decomposition, if any, is prose in its summary.[/bold yellow]", + log, + ) + else: + ok = decision in ("LANDED", "FAILED_RECEIPT") + self._set_item_status(state, item.n, ITEM_DONE if ok else ITEM_FAILED) + status[item.n] = ITEM_DONE if ok else ITEM_FAILED + if ok: + continue + if blocks(item): + failed.append(item) + else: + deferred.append(defer(item, failure(item, errors.pop(item.n, None)))) + self._checkpoint(state) + if deferred: + self._report_deferrals(deferred, log) + + # A retry is ONE more attempt, so a unit's failure is judged + # before a sibling's MACRO re-walks the queue; the re-walk would + # otherwise dispatch it a third time. The smith phase keeps its + # historical order, in which a failed core row rides the re-walk. + if enqueued and not (failed and unit is not None): + # The new items are unstarted, so re-derive the batches and + # walk the queue again rather than falling through to the + # next dependency layer. + print_message( + f"[bold cyan]→ {enqueued} constituent dispatch(es) enqueued; " + f"re-layering the queue[/bold cyan]", + log, + ) + return self._walk_queue(state, log, unit=unit) + + if failed: + return close(", ".join(f"{i.n}:{i.target}" for i in failed), "") + if errors: + first = next(iter(errors.values())) + raise RuntimeError( + f"{len(errors)} smith dispatch(es) raised; first: " + f"{type(first).__name__}: {first}" + ) from first + + return True + + def _report_deferrals(self, lines: list[str], log) -> None: + """Say, and journal, which dispatches were deferred and to which unit.""" + print_message( + f"[bold yellow]→ {len(lines)} smith dispatch(es) deferred — needed only by a " + f"later unit, so they do not block core. Each is retried once before that " + f"unit's first assembler turn:[/bold yellow]\n " + "\n ".join(lines), + log, + ) + self.journal.section("smith dispatches deferred", "\n".join(f"- {line}" for line in lines)) + + def _needed_by(self, state: WorkflowState, log=None) -> dict[int, str]: + """Each queue row's earliest-needing unit. See `vocabulary.needed_by`.""" + units = assembly_units(self.task_data) + items = [QueueItem.from_row(row) for row in state.queue] + unknown = sorted({item.unit for item in items} - set(units)) + if unknown and not self._warned_unknown_units: + self._warned_unknown_units = True + print_message( + f"[bold yellow]! queue rows name unit(s) {unknown}, which this run does not " + f"build ({', '.join(units)}). They are treated as core — the blocking " + f"direction — rather than guessed at.[/bold yellow]", + log, + ) + self.journal.event("queue units treated as core", ", ".join(unknown)) + return needed_by(items, units) + + def _land_deferred(self, state: WorkflowState, unit: str, log) -> bool: + """Retry, once each, the dispatches deferred to ``unit`` before its first assembler turn. + + Returns False when one still fails, having closed the run + ``smith-blocked:``. Idempotent: a row that landed is ``done``, + and a resume re-enters here and finds only what is still owed. + """ + units = assembly_units(self.task_data) + needed = self._needed_by(state, log) + horizon = units.index(unit) + owed = [ + row + for row in state.queue + if row.get("deferred_to") + and row.get("status") != ITEM_DONE + and units.index(needed[row["n"]]) <= horizon + ] + if not owed: + return True + lines = [f"{row['n']}:{row.get('target')} ({row.get('deferred_why', '')})" for row in owed] + print_message( + f"[bold cyan]→ {unit}: {len(owed)} deferred smith dispatch(es) are owed before " + f"its first assembler turn; retrying each once: {', '.join(lines)}[/bold cyan]", + log, + ) + self.journal.section( + f"{unit}: retrying deferred smith dispatches", "\n".join(f"- {line}" for line in lines) + ) + return self._walk_queue(state, log, unit=unit) + + def _retry_note(self, row: dict, item: QueueItem) -> str: + """What a deferred dispatch's retry is told about the attempt that failed. + + Empty for a row that was never dispatched -- it waited on a + dependency that had not landed -- because there is nothing to read. + """ + report = self.journal.latest_report("smith", f"{item.mode} {item.target}") + if report is None: + return "" + to = row.get("deferred_to") or item.unit + return ( + f"**This is a retry.** Dispatch {item.n} was attempted in the smith phase and " + f"did not land ({row.get('deferred_why') or 'no reason recorded'}). Only " + f"`{to}` needs this entry, so the run deferred it rather than stop core; this " + f"is its one remaining attempt, made before `{to}`'s first assembler turn. " + f"The previous attempt's report is `{report}`, and its progress entry is in " + f"`{self.progress_path}` (`build_stage`, iteration {item.n}). Read both first: " + f"keep what it established and do not repeat what it ruled out. If the obstacle " + f"it hit still stands, say exactly what it is — a second failure stops the run " + f"at `{to}`.\n" + ) + + def _run_assemble_phase(self, state: WorkflowState, log) -> bool: + """One assembler per unit, serially. Core is green before the first feature.""" + units = assembly_units(self.task_data) + while state.unit_index < len(units): + unit = units[state.unit_index] + if unit in state.units_done: + state.unit_index += 1 + continue + + print_rule(f"[bold magenta]assemble · {unit}[/bold magenta]", log) + # Before the first assembler turn, not after it: a turn spent + # deriving a forward against an entry that is not there yet ends + # BLOCKED on a vocabulary gap the run already knew about. + if not self._land_deferred(state, unit, log): + return False + self._progress_ctx.current_iteration = state.unit_index + 1 + + if not self._assemble_one_unit(state, unit, log): + return False + + state.units_done.append(unit) + state.unit_index += 1 + state.unit_turns = 0 + # A unit is the assembler's unit of work, so its context ends + # with it. The next unit starts clean, which is also what makes + # "core is green before the first feature" mean something. + self._reset_layer("assembler") + self._checkpoint(state) + + print_message("[bold green]✔ every unit assembled and gated[/bold green]", log) + return True + + def _enqueue_constituents(self, state: WorkflowState, item: QueueItem, log) -> int: + """Append a MACRO's constituents to the queue. Returns how many. + + Two things this deliberately does **not** do. + + It does not enqueue a child whose ``unit`` names a feature this run + is not building, and it does not enqueue one for a *later* unit than + the one in progress. A macro under `core` routinely decomposes into + pieces only the speculator needs; pulling those forward would gate + `core` against vocabulary it does not use and break the rule that + core is green before the first feature increment. Measured: of one + macro's eleven constituents, four were MTP-only. + + It also does not re-enqueue a target already in the queue, so a + resumed run or two macros naming a shared piece cannot dispatch it + twice. + + A constituent is needed no earlier than the row it decomposes, so it + inherits that row's ``unit`` unless it names a later one. The parent's + unit also widens what may be enqueued now: a feature-only row is + dispatched in the smith phase alongside core's, so its pieces are + too, and carrying its tag is what makes a failure among them defer + rather than block core. + """ + # THIS item's entry, not the latest smith entry. A batch runs ten + # smiths concurrently, so "the latest" is whichever finished last -- + # and reading its (empty) `constituents` drops the macro's whole + # decomposition. Measured: two MACROs in one batch, both reported as + # "listed no usable constituents", both of which had listed them. + entry = self._smith_entry_for(item) or {} + children = entry.get("constituents") or [] + if not isinstance(children, list) or not children: + return 0 + + units = assembly_units(self.task_data) + current = units[state.unit_index] if state.unit_index < len(units) else units[0] + parent = self._needed_by(state).get(item.n, item.unit) + parent = parent if parent in units else units[0] + reach = max(units.index(current), units.index(parent)) + allowed = set(units[: reach + 1]) + + existing = {str(row.get("target")) for row in state.queue} + next_n = max((int(row["n"]) for row in state.queue), default=0) + 1 + added = 0 + deferred: list[str] = [] + for child in children: + target = str(child.get("target", "")).strip() + if not target or target in existing: + continue + unit = str(child.get("unit") or parent).strip() or parent + if unit not in allowed: + deferred.append(f"{target} ({unit})") + continue + if units.index(unit) < units.index(parent): + unit = parent + state.queue.append( + { + "n": next_n, + "mode": "onboard", + "target": target, + "depends_on": list(child.get("depends_on") or []), + "surface": child.get("surface") or child.get("note"), + "entry": None, + "note": f"constituent of dispatch {item.n} ({item.target})", + "tier": int(child.get("tier", 1) or 1), + "unit": unit, + } + ) + existing.add(target) + next_n += 1 + added += 1 + + if deferred: + print_message( + f"[bold cyan]→ {len(deferred)} constituent(s) belong to a later unit and " + f"were not enqueued now: {', '.join(deferred)}[/bold cyan]", + log, + ) + self.journal.section(f"deferred constituents of dispatch {item.n}", "\n".join(deferred)) + return added + + def _features_through(self, unit: str) -> list[str]: + """Features implemented once ``unit`` is assembled, in order. + + Units are built one at a time and each is gated before the next + begins, so the target under test after unit *k* implements + ``features[:k]`` — never the whole list. ``core`` implements none. + + Returning anything wider gates a target against a feature it does + not have yet, which either fails on the target's own fail-fast + assert (the good case, and what happened here) or, if that assert is + missing, passes while measuring something else entirely. + """ + units = assembly_units(self.task_data) + if unit not in units: + return [str(f) for f in self.task_data.get("features", [])] + return [str(u) for u in units[1 : units.index(unit) + 1]] + + def _assemble_one_unit(self, state: WorkflowState, unit: str, log) -> bool: + # `unit_turns` bounds a self-heal loop, so it has to count WORK. It + # used to count entries into this loop, and those are not the same + # number: `task.yaml` is loaded once in `_init_state`, so every spec + # fix needs a restart, and every restart re-entered here and spent a + # turn re-deriving an answer that was already on disk. Measured on + # the MiMo-V2.6-Flash run: **ten assembler dispatches for four + # units, six of which built nothing**, and `core` reached turn 5 of 6 + # with only turn 1 having done anything -- one turn short of closing + # `assembly-exhausted` over a green, twice-reviewed target sitting in + # `HEAD`. + # + # The counter is still incremented at the top, because that is what + # makes the loop provably terminate. What changed is that the + # resume path **gives the turn back** when it does not dispatch. + while state.unit_turns < self.max_assembler_turns_per_unit: + state.unit_turns += 1 + # The decision file's name is derived from this, and it was last + # set by the smith phase to a BATCH index. Turn 2 of an assembly + # was therefore told to write `assembler-1.yaml` -- the path turn + # 1 had already used -- so a re-draft either overwrote its + # predecessor's decision or, as happened here, the role noticed + # the contradiction and wrote the right file anyway. + self._progress_ctx.current_iteration = state.unit_turns + self._checkpoint(state) + # A resume re-enters this loop at the top, so a BLOCKED decision + # from the process that died is never acted on: the handler only + # runs immediately after a dispatch. Check for one first, and pay + # the gap dispatches instead of an assembler turn that would + # re-derive an answer already on disk. + dispatch_needed = True + if state.unit_turns == 1 or self._resumed_into_unit: + # Captured before it is cleared: the skip below is a + # RESUME-only behaviour. At turn 1 of a continuous run the + # latest assembler decision belongs to the *previous* unit, + # and after a failed gate the loop re-enters with the flag + # already false -- so tying the skip to this flag is also + # what makes it fire at most once per process, and therefore + # what makes the loop terminate. + resumed = self._resumed_into_unit + self._resumed_into_unit = False + prior = latest_decision(self.progress_path, "assembler", ASSEMBLER_DECISIONS) + prior_reason = str( + (latest_entry(self.progress_path, "assembler") or {}).get("blocked_reason", "") + ).strip() + # The hole this closes. `BLOCKED` was the only prior decision + # the resume path acted on; a prior `UNIT_READY` fell + # straight through to a fresh dispatch, so the work was + # re-derived and the turn was spent. It is only safe to skip + # when the work that decision was written about is still the + # work on disk, which `_unit_work_unchanged` answers from git + # rather than assuming -- a resume after someone edited the + # target is a different situation from a clean restart, and + # an unanswerable question is answered "changed". + if resumed and prior == "UNIT_READY" and self._unit_work_unchanged(state, unit): + state.unit_turns -= 1 + self._progress_ctx.current_iteration = max(state.unit_turns, 1) + self._checkpoint(state) + print_message( + f"[bold cyan]→ resuming onto a UNIT_READY the previous process " + f"already wrote, and the work under `{self._preflight.target_rel}` " + f"is unchanged since. Going straight to the gates: **this restart " + f"did not spend an assembler turn** — {unit} is still at turn " + f"{state.unit_turns} of {self.max_assembler_turns_per_unit}." + f"[/bold cyan]", + log, + ) + self.journal.event( + "resume did not spend a turn", + f"unit={unit}; prior decision UNIT_READY; target unchanged; " + f"unit_turns={state.unit_turns}", + ) + dispatch_needed = False + # Only a VOCABULARY gap is worth re-handling on a resume. The + # other blocked reasons already ended the run once, and + # handling one again just ends it again -- which is what + # happened: a `knowledge_gap` block written at a session + # cut-off re-closed the run within seconds of every resume, + # before the newly-extended session could dispatch anything. + # + # An operator who resumed deliberately has, by resuming, said + # the blocking condition is meant to be revisited. So the + # other reasons fall through to a fresh assembler turn, which + # is free to reach the same conclusion and say so. + elif prior == "BLOCKED" and prior_reason == "vocabulary_gap": + print_message( + "[bold cyan]→ resuming onto a BLOCKED assembler decision from a " + "previous process; handling it rather than re-dispatching" + "[/bold cyan]", + log, + ) + if self._handle_blocked_assembler(state, log): + continue + return False + elif prior == "BLOCKED": + print_message( + f"[bold cyan]→ resuming past a BLOCKED assembler decision " + f"({prior_reason or 'no reason recorded'}); re-dispatching, since " + f"resuming is itself a statement that it should be revisited" + f"[/bold cyan]", + log, + ) + + if dispatch_needed: + self._dispatch( + "assembler", + what=f"{unit} (turn {state.unit_turns})", + prompt=self._assembler_prompt(unit, state.unit_turns), + ) + decision = latest_decision(self.progress_path, "assembler", ASSEMBLER_DECISIONS) + if decision == "BLOCKED": + if self._handle_blocked_assembler(state, log): + continue + return False + # Recorded here, at the only moment the answer is known to + # describe what this turn built. A later process compares + # against it instead of re-dispatching to find out. A turn + # that ended with neither decision -- a role that wrote no + # decision file -- records nothing, so a resume onto it + # re-dispatches. + if decision == "UNIT_READY": + state.unit_marker[unit] = self._target_marker() + self._checkpoint(state) + + # The configurations this unit is gated over, resolved once. A + # `unified` segment is ONE codebase over a FAMILY of parallel + # configurations, and gating it at whatever `topology_flags` + # happened to emit meant pure tensor parallelism — with the + # expert-parallel and attention-DP branches that only a unified + # codebase carries executed by no gate, while the target's + # Coverage claimed all three. For a concrete segment this is a + # one-element list, so there is one code path and the + # single-configuration run is the degenerate case of it. + # Cheapest member first. A unit gates EVERY member and a broken + # one fails on all of them, so the order decides how long a + # failure takes to surface. Measured on this package's first + # unified run: `tep4` boots in ~180s against `tp4`'s 285-438s. + members = gates_mod.order_members_by_cost( + members_for(self.task_data, unit, self._gate_env().world_size) + ) + + boots = self._gate_over_members( + members, + lambda member: gates_mod.run_boot_gate( + self._gate_env(), + task_data=self.task_data, + spec=gates_mod.BootSpec.load(self.boot_spec_path), + runner=self._shell_runner, + # The features this target implements *now*, not the ones + # the run will finish with. `core` implements none, so + # gating it with the speculator's flags asks a base model + # to emit drafts it has no path to produce. + features=self._features_through(unit), + unit=unit, + member=member, + ), + log, + ) + if not boots: + continue + + if self.reviewer_mode != "off": + # The last member's evidence: the reviewer reads one dump and + # the summary line names which member produced it, so the + # freshest is the one whose log is still meaningful to quote. + self._dispatch("reviewer", what=unit, prompt=self._reviewer_prompt(unit, boots[-1])) + if latest_decision(self.progress_path, "reviewer", REVIEWER_DECISIONS) == "REJECT": + print_message("[bold yellow]↻ reviewer REJECT[/bold yellow]", log) + continue + + if not self._gate_over_members( + members, + lambda member: gates_mod.run_accuracy_gate( + self._gate_env(), + task_data=self.task_data, + reference=self._reference_score(unit), + runner=self._shell_runner, + # Without this the MTP unit is measured under core's + # LLM-API YAML — no speculative_config — so the gate + # re-measures the base path and reports the number as + # the speculator's. + unit=unit, + member=member, + ), + log, + announce=True, + ): + continue + + # The third gate, and the only one that can see a miscomputed + # draft path. Rejection sampling holds the emitted distribution + # to the target model's, so boot and accuracy both pass while + # only speed moves. Runs for distribution-preserving features + # only; core has no draft path to compare. + if unit in WITNESS_GATED_UNITS: + # **The first member only**, and deliberately. Boot and + # accuracy ask "does this configuration work", which is a + # different question at every member. A witness asks + # something else: it is a CONTRAST between two legs of the + # SAME configuration — a residency delta, or token identity + # across the disagg hop — and what it proves is a property of + # the modeling code rather than of the topology. The + # residency witness makes that concrete: its verdict is a + # device-memory delta against `expected_mib`, one number the + # spec derives from the checkpoint's own shapes, so running + # it at DEP (attention replicated) and at TEP (attention + # sharded) would compare two different deltas against an + # expectation computed for one of them. Each witness is also + # two full engine launches, so looping it would triple the + # most expensive gate shape in the set to re-ask a question + # that does not vary with the answer. + witness = self._run_witness(unit, members[0], log) + if witness is None: + return False + # Escaped: a token witness's verdict quotes generated text, + # which may contain what rich reads as markup. + if not witness.passed: + print_message( + f"[bold yellow]↻ {escape_markup(witness.summary())}[/bold yellow]", log + ) + continue + print_message(f"[bold green]✔ {escape_markup(witness.summary())}[/bold green]", log) + + if unit in ACCEPTANCE_GATED_UNITS: + if not self._gate_over_members( + members, + lambda member: gates_mod.run_acceptance_gate( + self._gate_env(), + task_data=self.task_data, + spec=gates_mod.BootSpec.load(self.boot_spec_path), + runner=self._shell_runner, + unit=unit, + # The features the target implements NOW, same rule as + # the boot gate: gating with a speculator flag the + # target has no path for is not a stricter gate, it is + # a different one. + features=self._features_through(unit), + member=member, + ), + log, + announce=True, + ): + continue + return True + + print_message( + f"[bold red]✘ {unit}: {self.max_assembler_turns_per_unit} assembler turns " + f"without a green unit. Ending the run with the audit trail rather than " + f"looping — the gate is the release criterion.[/bold red]", + log, + ) + self._close_run(state, log, verdict=f"assembly-exhausted:{unit}") + return False + + def _gate_over_members( + self, + members: list[str], + run, + log, + *, + announce: bool = False, + ) -> "list[gates_mod.GateResult]": + """Fire one gate at every member. Returns the verdicts, or `[]` on a failure. + + Every member is recorded separately and none is collapsed: three + measurements in `progress.yaml` and the journal is the entire point + of declaring a member set, and one line reading `accuracy PASS` for a + target that serves three configurations is precisely the unsupported + claim `members:` exists to stop. + + It stops at the first failure. The unit is going back to the + assembler either way, and the remaining members would spend an hour + each measuring the same defect at a different topology. + + The empty list means "a member failed" rather than "no members" — + `members_for` never returns an empty set, so the two cannot be + confused. + """ + results: list[gates_mod.GateResult] = [] + for member in members: + result = self.gate_log.record(run(member)) + results.append(result) + self.journal.event("gate", result.summary()) + if not result.passed: + print_message(f"[bold yellow]↻ {result.summary()}[/bold yellow]", log) + return [] + if announce: + print_message(f"[bold green]✔ {result.summary()}[/bold green]", log) + return results + + @staticmethod + def _witness_block_gap(unit: str, kind: str, block: Any) -> str | None: + """What `accuracy.witness.` lacks for its witness, or None. + + The schema rejects every one of these for a spec it validated, and a + run reads its spec through the schema. This is the second line for a + spec that changed underneath a resumed run, and it names the key, + because "witness unconfigured" alone sends the operator to re-read the + whole block. + """ + if not isinstance(block, dict): + return "a block at all" + if kind == "replay": + absent = [key for key in REPLAY_WITNESS_REQUIRED[unit] if not block.get(key)] + return ", ".join(f"`{key}`" for key in absent) or None + if kind == "kv_quant": + if not block.get("expected_bytes_per_token"): + return "`expected_bytes_per_token`" + # `exact: false` drops the identity leg, and with it the contrast. + if block.get("exact", True) and not block.get("contrast_env"): + return "`contrast_env`" + return None + if kind == "equivalence": + # Without its own launcher leg B is leg A under an environment + # variable nothing reads, and the witness compares a run to itself. + absent = [key for key in ("contrast_env", "contrast_command") if not block.get(key)] + return ", ".join(f"`{key}`" for key in absent) or None + return None if block.get("contrast_env") else "`contrast_env`" + + def _run_witness(self, unit: str, member: str, log) -> "gates_mod.GateResult | None": + """Run the witness gate `unit`'s feature needs. None means misconfigured. + + `member` is the configuration both legs are served at — see the call + site for why a witness runs at one member rather than over all of + them. + + The spec has to supply the one thing that cannot be derived: the + environment that turns the feature OFF. A missing block is a hard + stop rather than a skip, because skipping it would declare the unit + green on two gates that are blind to its feature by construction -- + which is the exact outcome the witness exists to prevent. + """ + from . import witness as witness_mod + + kind = WITNESS_GATED_UNITS[unit] + spec = (self.task_data.get("accuracy") or {}).get("witness") or {} + block = spec.get(unit) if isinstance(spec, dict) else None + missing = self._witness_block_gap(unit, kind, block) + if missing: + print_message( + f"[bold red]✘ the `{unit}` unit needs a `{kind}` witness and " + f"`accuracy.witness.{unit}` does not set {missing}. Without it the unit " + f"would be gated only by boot and accuracy, both of which pass " + f"whether or not the feature is doing anything.[/bold red]", + log, + ) + self._close_run( + load_state(self.state_path), log, verdict=f"witness-unconfigured:{unit}" + ) + return None + + contrast_env = {str(k): str(v) for k, v in (block.get("contrast_env") or {}).items()} + common = dict( + task_data=self.task_data, + runner=self._shell_runner, + unit=unit, + features=self._features_through(unit), + member=member, + # Per leg. A disagg leg boots two engines in sequence, which the + # one-engine default was never sized for. + timeout=int(block.get("timeout") or gates_mod.BOOT_TIMEOUT_S), + ) + # The token-comparing witnesses only: the replay witnesses count work, + # and their schema refuses the key. Both legs run at this batch limit; + # see `task_schema.WITNESS_BATCH_SIZE_KEY` for why 1 is the useful one. + batch = block.get(WITNESS_BATCH_SIZE_KEY) + schedule = dict(max_batch_size=None if batch is None else int(batch)) + boot_spec = gates_mod.BootSpec.load(self.boot_spec_path) + if kind == "residency": + result = witness_mod.run_residency_witness( + self._gate_env(), + expected_mib=float(block.get("expected_mib", 0.0)), + contrast_env=contrast_env, + spec=boot_spec, + **schedule, + **common, + ) + elif kind == "replay": + # The workload is generated from the declared sizes, not from the + # boot prompts: two short prompts never exceed the window, and a + # witness whose prompt fits inside the window passes for a target + # that bounds nothing. + result = witness_mod.run_replay_witness( + self._gate_env(), + feature=unit, + contrast_env=contrast_env, + window=int(block["window"]), + prompt_tokens=block.get("prompt_tokens"), + prefix_tokens=block.get("prefix_tokens"), + reuse_alignment=int(block.get("reuse_alignment", 1)), + **common, + ) + elif kind == "kv_quant": + result = witness_mod.run_kv_quant_witness( + self._gate_env(), + expected_bytes_per_token=int(block["expected_bytes_per_token"]), + exact=bool(block.get("exact", True)), + contrast_env=contrast_env, + bytes_tolerance=float(block.get("bytes_tolerance", DEFAULT_BYTES_TOLERANCE)), + spec=boot_spec, + **schedule, + **common, + ) + else: + result = witness_mod.run_equivalence_witness( + self._gate_env(), + contrast_env=contrast_env, + contrast_command=block.get("contrast_command"), + spec=boot_spec, + prompt_tokens=block.get("prompt_tokens"), + workdir=self.workspace / "witness" / unit, + legs_env={ + str(k): str(v) for k, v in (block.get(WITNESS_LEGS_ENV_KEY) or {}).items() + }, + **schedule, + **common, + ) + self.gate_log.record(result) + self.journal.event("gate", result.summary()) + if not result.passed and result.stdout_tail.strip(): + # A token comparison's tail is both legs around where each + # differing continuation first differs; a leg that died leaves its + # last lines. The one-line summary the gate history carries quotes + # only the first difference, so the rest goes to the journal and + # the run's log. + self.journal.section( + f"witness `{result.name}` failed: evidence", + f"```\n{result.stdout_tail.rstrip()}\n```", + ) + print_message(escape_markup(result.stdout_tail.rstrip()), log) + return result + + def _run_deliver_phase(self, state: WorkflowState, log) -> bool: + print_rule("[bold magenta]deliver[/bold magenta]", log) + self._progress_ctx.current_iteration = 1 + + report = self._delivery_report(state) + self.journal.section("delivery checks", f"```\n{report.render()}\n```") + print_message(report.render(), log) + if not report.passed: + names = ", ".join(check.name for check in report.failures) + print_message(f"[bold red]✘ delivery verification failed: {names}[/bold red]", log) + self._close_run(state, log, verdict="delivery-failed") + return False + + boot = self.gate_log.last("boot") + return self._run_remediation_cycle(state, report, boot, log) + + # ------------------------------------------------------- remediation cycle + + def _run_remediation_cycle(self, state: WorkflowState, report, boot, log) -> bool: + """Audit, and close what it finds — bounded, deduplicated, monotone. + + The loop this replaces was not in the workflow at all. An `UNCOVERED` + ended the run, a human certified the entry, and a re-run spent a + fresh stateless auditor re-scanning everything to find the next gap. + That is the mechanism behind audit passes 53 through 57 on one + target, each paying a full re-scan to learn one increment. + + What makes this a cycle rather than a `while` is in + :mod:`remediation`: findings group by entry surface, carry a stable + identity, and have to shrink the unresolved set to earn another + round. Everything expensive stays outside it — a verify smith runs + the catalog GPU test and writes a receipt; the accuracy gate does + **not** re-run, because certifying a contract column cannot change a + number the engine already produced. + """ + cycle = remediation.RemediationState.from_dict(state.remediation) + budget = self._remediation_budget() + + while True: + cycle.round_no += 1 + self._progress_ctx.current_iteration = cycle.round_no + self._dispatch( + "auditor", + what=f"runtime-value dump vs certified columns (round {cycle.round_no})", + prompt=self._auditor_prompt(report, boot), + ) + entry = latest_entry(self.progress_path, "auditor") or {} + if latest_decision(self.progress_path, "auditor", AUDITOR_DECISIONS) == "CLEAN": + state.audit_clean = True + state.remediation = cycle.as_dict() + self._checkpoint(state) + print_message("[bold green]✔ audit CLEAN[/bold green]", log) + return True + + rows, dropped = remediation.parse_uncovered_rows(entry.get("uncovered")) + if dropped: + # Never silent: a dropped row is a finding nothing will act + # on, and the run would otherwise converge on an incomplete + # set and call it clean. + print_message( + f"[bold yellow]! {dropped} uncovered row(s) named no entry, arch or " + f"axis and cannot be dispatched[/bold yellow]", + log, + ) + if not rows: + # An UNCOVERED with no structured findings is the old + # behaviour, and the only honest thing to do with it is what + # the old behaviour did. + print_message( + "[bold red]✘ auditor UNCOVERED but listed no dispatchable rows. " + "Certify what its summary names and re-deliver.[/bold red]", + log, + ) + self.journal.section("audit: uncovered", str(entry.get("summary", ""))) + self._close_run(state, log, verdict="audit-uncovered") + return False + + incoming = remediation.canonicalize( + rows, + target_commit=self._target_commit(), + capture_id=self._capture_id(), + ) + merge = remediation.merge_into_queue(cycle.queue, incoming, budget=budget) + decision = remediation.decide_round( + merge, + budget=budget, + round_no=cycle.round_no, + previous_signature=cycle.signature, + ) + cycle.queue = list(merge.queue) + state.remediation = cycle.as_dict() + self._checkpoint(state) + + if decision.action == "clean": + state.audit_clean = True + self._checkpoint(state) + print_message("[bold green]✔ audit CLEAN[/bold green]", log) + return True + + if decision.stopping: + print_message(f"[bold red]✘ {decision.detail}[/bold red]", log) + self.journal.section( + f"remediation stopped: {decision.verdict}", + f"{decision.detail}\n\n```\n{cycle.render()}\n```", + ) + self._close_run(state, log, verdict=decision.verdict) + return False + + print_rule( + f"[bold magenta]remediate · round {cycle.round_no} — " + f"{decision.detail}[/bold magenta]", + log, + ) + # The signature is taken BEFORE the round runs. Comparing the + # next round's unresolved set against it is what proves the + # cycle is shrinking rather than circling. + cycle.signature = remediation.unresolved_signature(cycle.queue) + failed = self._dispatch_verifiers(cycle, decision, log) + state.remediation = cycle.as_dict() + self._checkpoint(state) + if failed: + names = ", ".join(f"{i.entry}[{i.axis}]" for i in failed) + print_message(f"[bold red]✘ verification did not land: {names}[/bold red]", log) + self.journal.section( + "remediation: verification failed", + f"{names}\n\n```\n{cycle.render()}\n```", + ) + self._close_run(state, log, verdict=remediation.STOP_VERIFICATION_FAILED) + return False + + # Re-verify the mechanical half before spending another auditor + # turn: a receipt that just landed changes what delivery sees. + report = self._delivery_report(state) + if not report.passed: + names = ", ".join(check.name for check in report.failures) + print_message( + f"[bold red]✘ delivery verification failed after remediation: " + f"{names}[/bold red]", + log, + ) + self._close_run(state, log, verdict="delivery-failed") + return False + + def _remediation_budget(self) -> "remediation.RemediationBudget": + """The cycle's hard limits, overridable per run from the spec.""" + block = self.task_data.get("catalog_remediation") or {} + default = remediation.RemediationBudget() + return remediation.RemediationBudget( + max_rounds=int(block.get("max_rounds", default.max_rounds)), + max_attempts_per_item=int( + block.get("max_attempts_per_item", default.max_attempts_per_item) + ), + max_new_items=int(block.get("max_new_items", default.max_new_items)), + max_parallel_verifiers=int( + block.get("max_parallel_verifiers", default.max_parallel_verifiers) + ), + ) + + def _dispatch_verifiers(self, cycle: "remediation.RemediationState", decision, log) -> tuple: + """Run one round's verify smiths concurrently. Returns the failures. + + Each item becomes exactly one smith in ``verify`` mode carrying the + whole surface — which is the point of grouping. The outcome is read + from the decision each smith wrote, never from its prose. + """ + items = list(decision.dispatch) + running = {item.identity for item in items} + cycle.queue = [ + replace(q, status=remediation.QUEUE_RUNNING) if q.identity in running else q + for q in cycle.queue + ] + + # One smith per surface, carrying the whole surface. The dispatch + # numbers start well above the build queue's so a decision file from + # remediation is never mistaken for one the smith phase wrote. + dispatches = { + item.identity: QueueItem( + n=9000 + index, + mode=MODE_VERIFY, + target=item.entry, + surface=item.brief(), + ) + for index, item in enumerate(items) + } + + with ThreadPoolExecutor( + max_workers=max(1, decision.workers), thread_name_prefix="verifier" + ) as pool: + futures = { + pool.submit(self._run_one_smith, dispatch): identity + for identity, dispatch in dispatches.items() + } + for future in as_completed(futures): + try: + future.result() + except BaseException as exc: # keep the other verifications + print_message( + f"[bold yellow]! verify dispatch for {futures[future]} raised " + f"{type(exc).__name__}: {exc}[/bold yellow]", + log, + ) + + # The verdict comes from the decision each smith wrote, never from + # what it said. A dispatch that raised simply has no decision, and + # `apply_outcomes` reads that as blocked rather than as success. + outcomes = { + identity: self._smith_decision_for(dispatch) + for identity, dispatch in dispatches.items() + } + updated, failed = remediation.apply_outcomes(cycle.queue, outcomes) + cycle.queue = list(updated) + return failed + + def _target_marker(self) -> str: + """A cheap identity for the work under the target path, or ``""``. + + Two questions, one string: which commit last touched the target, and + whether anything is uncommitted under it. Together those separate the + two resumes that matter — a clean restart, where both answers are the + ones the previous process left, and a restart after someone edited + the target, where at least one has moved. + + The path is derived rather than spelled out: ``trtllm_repo_path`` + joined with :attr:`preflight.PreflightResult.target_rel`, which is the + same pair :func:`delivery.verify_delivery` is handed. + + ``""`` means **not established**, and every caller reads it as + "assume changed". There is no honest third answer, and the asymmetry + is deliberate: a resume that re-dispatches when it need not have + costs one assembler turn, while a resume that skips a dispatch it + needed costs the unit. + """ + rel = str(getattr(self._preflight, "target_rel", "") or "") + repo = str(self.task_data.get("trtllm_repo_path") or "") + if not rel or not repo: + return "" + commit = self._shell_runner(f"git -C {repo} log -1 --format=%H -- {rel}", None, 60) + dirt = self._shell_runner(f"git -C {repo} status --porcelain -- {rel}", None, 60) + if not getattr(commit, "ok", False) or not getattr(dirt, "ok", False): + return "" + head = (getattr(commit, "stdout", "") or "").strip().splitlines() + if not head: + return "" + working = (getattr(dirt, "stdout", "") or "").strip() + return f"{head[0].strip()}:{hashlib.sha1(working.encode('utf-8')).hexdigest()[:12]}" + + def _unit_work_unchanged(self, state: WorkflowState, unit: str) -> bool: + """Whether the target still holds what this unit's ``UNIT_READY`` meant. + + Keyed by unit, so the previous unit's marker can never answer for + this one: at turn 1 of `mtp` the latest assembler decision is + `core`'s `UNIT_READY`, and an unkeyed comparison would skip the + increment's first dispatch entirely. + """ + recorded = state.unit_marker.get(unit, "") + current = self._target_marker() + return bool(recorded) and bool(current) and recorded == current + + def _target_commit(self) -> str: + """The repo commit the audited evidence belongs to. + + Part of an uncovered row's identity: the same values under a NEW + commit are legitimately new work, because a wrapper change can open a + runtime branch that did not exist before. Unknown is a valid answer + and is not worth failing a run over -- it only makes the identity + coarser, never wrong. + """ + result = self._shell_runner( + f"git -C {self.task_data['trtllm_repo_path']} rev-parse --short HEAD", None, 60 + ) + if not getattr(result, "ok", False): + return "" + lines = (getattr(result, "stdout", "") or "").split() + return lines[0] if lines else "" + + def _capture_id(self) -> str: + """Which runtime capture the audit read. + + Two rounds reading the same capture that disagree are a coverage + problem, not new work — and that distinction is only visible if the + capture is named. + """ + names = [ + r.name + for r in self.gate_log.results + if self.RUNTIME_VALUE_MARKER in (r.stdout_tail or "") + ] + return "+".join(names) if names else "resumed" + + def _delivery_report(self, state: WorkflowState): + """The mechanical half of delivery. Cheap, and safe to repeat.""" + return verify_delivery( + repo_root=self.task_data["trtllm_repo_path"], + target_rel=self._preflight.target_rel, + family=self._preflight.family, + reference_yaml=gates_mod.resolve_per_unit( + self.task_data["accuracy"], + "reference_yaml", + state.units_done[-1] if state.units_done else None, + ), + model_key=self._model_key(), + accuracy_test=f"test_modeling_v2_{self._preflight.family}.py", + ) + + # ------------------------------------------------------------------ helpers + + def _grown_units(self, state: WorkflowState) -> list[str]: + """Units the spec now lists that this run never assembled, in order.""" + return [u for u in assembly_units(self.task_data) if u not in state.units_done] + + def _init_state(self, task: str | Path, log) -> WorkflowState | None: + source = Path(task) + if not source.is_file(): + raise FileNotFoundError(f"--task must be a path to an existing YAML file; got {source}") + self.task_data = load_and_validate_task_yaml(source) + + if self.resume: + state = load_state(self.state_path) + self._backfill_units(state, log) + if state.done and is_remediable(state.verdict): + # The run closed on something the operator fixes outside it, + # and re-running its own check is the only way to see the fix. + # Refusing here would mean `--clean` and a full rebuild -- + # every gate re-measured to re-answer one question. + print_message( + f"[bold cyan]→ re-entering stage '{state.stage}': the run closed " + f"'{state.verdict}', which names a repair made outside it[/bold cyan]", + log, + ) + self.journal.event("re-entered a closed run", f"verdict={state.verdict}") + if state.verdict.startswith("assembler-blocked"): + # The turn budget bounds the unit's own self-heal loop; + # the repair that re-opens it happened outside that loop. + state.unit_turns = 0 + state.done = False + state.verdict = "" + elif self._smith_block_is_deferrable(state): + # `smith-blocked` stays refused in general: re-running it + # re-dispatches failed turns with nothing changed. This is the + # one shape where something did change -- every row that + # failed is now known to be needed only by a later unit -- and + # re-entering spends nothing on them now: the smith phase + # records them `deferred` without dispatching, finishes what + # the close left unstarted, and core assembles. + failed = ", ".join( + f"{row['n']}:{row.get('target')}" + for row in state.queue + if row.get("status") == ITEM_FAILED + ) + print_message( + f"[bold cyan]→ re-entering stage 'smith': the run closed 'smith-blocked', " + f"but every dispatch that failed ({failed}) is needed only by a later " + f"unit; they are deferred to it rather than blocking core[/bold cyan]", + log, + ) + self.journal.event("re-entered a smith-blocked run", f"feature-only: {failed}") + state.done = False + state.verdict = "" + elif state.done and state.verdict == "delivered" and self._grown_units(state): + # The spec gained a feature after the run delivered. That is an + # increment, not a repair, and it is the ordinary way a target + # grows: core ships, then MTP, then multimodality, each gated + # against the last. The alternative was `--clean` -- re-plan, + # re-smith, re-assemble and re-measure every unit that is + # already green -- to add one. + # + # Narrow on purpose. Only a run that actually *delivered* + # qualifies: a failed one is also missing units from + # `units_done`, and re-entering there would re-dispatch paid + # agent turns with nothing having changed, which is what + # `REMEDIABLE_VERDICTS` exists to prevent. + added = self._grown_units(state) + print_message( + f"[bold cyan]→ the spec grew by {', '.join(added)} since this run " + f"delivered; re-entering at 'assemble'[/bold cyan]", + log, + ) + self.journal.event("spec grew after delivery", f"new units: {', '.join(added)}") + state.done = False + state.verdict = "" + state.stage = STAGE_ASSEMBLE + state.audit_clean = False + state.unit_turns = 0 + # `_run_assemble_phase` skips units already in `units_done`, so + # pointing at the first new one is enough and re-derives + # nothing. + units = assembly_units(self.task_data) + state.unit_index = units.index(added[0]) + elif state.done: + print_message( + "[bold green]✔ this run already completed; pass --clean to start " + "over.[/bold green]", + log, + ) + return None + print_message(f"[bold cyan]→ resuming at stage '{state.stage}'[/bold cyan]", log) + # A resume into an assembly means the previous process may have + # left a decision nothing acted on. See `_assemble_one_unit`. + self._resumed_into_unit = state.stage == STAGE_ASSEMBLE + self._refresh_workspace_spec(source, log) + return state + + if source.resolve() != self.task_path.resolve(): + # The resolved spec, not the user's file: defaults are filled in + # so what the agents read on disk is fully explicit. + self.task_path.write_text( + yaml.safe_dump(self.task_data, sort_keys=False, allow_unicode=True), + encoding="utf-8", + ) + state = WorkflowState(task_path=str(self.task_path)) + self._checkpoint(state) + return state + + def _backfill_units(self, state: WorkflowState, log) -> None: + """Carry ``unit`` tags from ``vocabulary.yaml`` onto rows checkpointed without one. + + A queue checkpointed before units existed carries no ``unit`` on any + row, so every row reads as core and blocks exactly as before. Tagging + the vocabulary and restarting is how an operator applies deferral to + such a run. Rows are matched on ``n`` *and* ``target``, so a + vocabulary the plan phase has since rewritten tags nothing it does not + still describe; a row that already carries a ``unit`` is left alone, + because the checkpoint, not the file, is the record of the run. + """ + untagged = [row for row in state.queue if "unit" not in row] + if not untagged or not self.vocabulary_path.is_file(): + return + try: + vocabulary = load_vocabulary(self.vocabulary_path, units=assembly_units(self.task_data)) + except VocabularyError as exc: + print_message( + f"[bold yellow]! unit tags in {self.vocabulary_path} were not applied — it " + f"does not validate: {exc}\nUntagged rows stay core, which blocks as " + f"before.[/bold yellow]", + log, + ) + return + tags = { + (item.n, item.target): item.unit for item in vocabulary.queue if item.unit != UNIT_CORE + } + applied: list[str] = [] + for row in untagged: + unit = tags.get((int(row["n"]), str(row.get("target")))) + if unit: + row["unit"] = unit + applied.append(f"{row['n']}:{row.get('target')} → {unit}") + if not applied: + return + self._checkpoint(state) + print_message( + f"[bold cyan]→ applied {len(applied)} unit tag(s) from vocabulary.yaml to the " + f"checkpointed queue: {', '.join(applied)}[/bold cyan]", + log, + ) + self.journal.section( + "unit tags applied from vocabulary.yaml", "\n".join(f"- {line}" for line in applied) + ) + + def _smith_block_is_deferrable(self, state: WorkflowState) -> bool: + """Whether a run closed `smith-blocked` in the smith phase only over feature-only rows.""" + if not (state.done and state.verdict == "smith-blocked" and state.stage == STAGE_SMITH): + return False + failed = [row for row in state.queue if row.get("status") == ITEM_FAILED] + if not failed: + return False + needed = self._needed_by(state) + return all(needed[row["n"]] != UNIT_CORE for row in failed) + + def _refresh_workspace_spec(self, source: Path, log) -> None: + """Keep the spec the agents read identical to the one the gates use. + + The run has two copies of ``task.yaml``: the operator's file, which + ``_init_state`` re-reads and resolves on **every** run, and the + workspace copy, which every role is told to read as "the run's spec". + The workspace copy was written once, on the first run, and never + again — so a resumed run silently gave the two halves different + specs. + + That is a split brain, not a stale file. The gates would measure the + operator's current spec while the assembler derived against an older + one: different `features`, different `runtime` switches, a different + accuracy anchor. Both sides self-consistent, neither correct, and + nothing anywhere would say so. + + Found by the assembler, which noticed the environment block it had + been handed did not match the spec it had been told to read. + """ + if source.resolve() == self.task_path.resolve(): + return + resolved = yaml.safe_dump(self.task_data, sort_keys=False, allow_unicode=True) + if self.task_path.is_file() and self.task_path.read_text(encoding="utf-8") == resolved: + return + self.task_path.write_text(resolved, encoding="utf-8") + print_message( + f"[bold cyan]→ refreshed the workspace spec from {source}; it had drifted " + f"from the file this run was given[/bold cyan]", + log, + ) + + def _checkpoint(self, state: WorkflowState) -> None: + save_state(self.state_path, state) + + @staticmethod + def _materialize(source: str | Path, destination: Path) -> None: + candidate = Path(source) + if candidate.is_file(): + shutil.copyfile(candidate, destination) + else: + destination.write_text(str(source).rstrip() + "\n", encoding="utf-8") + + def _gate_env(self) -> RunEnvironment: + """The environment a GATE runs under, which may not be the run's. + + A spec may declare `environment.workloads.gate` to give gates their + own launch substrate, and the reason is a failure domain rather than a + preference. A shared session imports a 26 GB container once and costs + nothing per job afterwards, which is right for a one-minute entry test + -- the predecessor run dispatched 443 of them. It is wrong for a 10-15 + minute gate: a shared session has ONE failure domain, so a gate that + leaks a server, kills a process group or collides on PMIx takes every + other job and the orchestrator with it. That happened three times on + the predecessor project and twice on this one. + + Resolved once and cached, from the same claimed devices and world size + preflight resolved the base environment with, so the two cannot + disagree about the topology. With no override declared this returns an + environment equal to the run's, and nothing changes. + """ + assert self.environment is not None, "preflight must run before any gate" + if self._gate_environment is None: + self._gate_environment = resolve_environment( + self.task_data, + devices=self.environment.devices, + world_size=self.environment.world_size, + workload="gate", + ) + return self._gate_environment + + def _model_key(self, unit: str | None = None) -> str: + """The HuggingFace repo name the accuracy references are keyed by. + + Delegated to :func:`~.preflight.resolve_model_key` rather than + restated. Preflight has to ask this same question hours earlier, to + check the anchor row exists at all, and **an anchor check that + resolves a different key than the gate does is worse than no check** + — it passes on a row nothing ever reads. One implementation is the + only way those two cannot drift. + + Delegating also picks up the per-unit mapping form, which the + restatement here did not honour: it read `model_key` as a scalar, so + a spec whose feature increment anchors under another repo name would + have been stringified into a key matching nothing. + """ + from .preflight import resolve_model_key + + return resolve_model_key(self.task_data, unit) + + def _reference_score(self, unit: str | None = None) -> float: + # `unit` is load-bearing: the accuracy references discriminate rows by + # `quant_algo` / `spec_dec_algo`, so core and the `mtp` increment + # anchor to DIFFERENT numbers. Resolving without it either picks the + # undiscriminated (bf16) row for a quantized checkpoint or raises at + # the end of an assembly instead of at spec-validation time. + return gates_mod.resolve_reference( + self.task_data, model_key=self._model_key(unit), unit=unit + ) + + def _set_item_status(self, state: WorkflowState, n: int, status: str) -> None: + for row in state.queue: + if row.get("n") == n: + row["status"] = status + + def _smith_entry_for(self, item: QueueItem) -> dict | None: + """The progress entry a smith wrote for this queue item. + + A concurrent batch appends several smith entries, so "the latest + smith entry" is not this item's. Matched on the queue item's own + `n` where the collector stamped one, and on the entry path or the + summary text otherwise. + """ + from .progress import BUILD_STAGE, find_entries + + rows = find_entries(self.progress_path, stage=BUILD_STAGE, agent="smith") + for row in reversed(rows): + if str(row.get("iteration", "")) == str(item.n): + return row + for row in reversed(rows): + if item.target in str(row.get("entry", "")) or item.target in str( + row.get("summary", "") + ): + return row + return None + + def _smith_decision_for(self, item: QueueItem) -> str | None: + """The decision from the smith entry for this queue item.""" + row = self._smith_entry_for(item) + if row is None: + return None + decision = str(row.get("decision", "")).strip().upper() + return decision if decision in SMITH_DECISIONS else None + + def _run_one_smith(self, item: QueueItem, note: str = "") -> str: + """One catalog entry, in its own layer. Called from the batch's pool. + + **The decision file is collected here.** In MCP-tool mode a smith + reports through `append_smith_progress`, which writes `progress.yaml` + itself; under an enterprise MCP configuration there are no tools and + the smith writes a file instead — and nothing read it. Ten smiths + landed ten entries, wrote ten valid decision files, and + `progress.yaml`'s `build_stage` stayed empty, so `_smith_decision_for` + found no decision and the whole batch would have scored FAILED. + + Caught by a smith, from inside its own dispatch, reading the + orchestrator's source. + + The iteration is the queue item's `n`, not a turn counter: smiths run + concurrently, so a counter would race, and `n` is the only number + that identifies which entry a decision is about. Every smith in the + first batch independently guessed this same convention from the + system prompt, which is a good sign it is the readable one. + + ``note`` is prepended context for a retry -- see `_retry_note`. + + **A resume adopts an answer the previous process never read.** The + state is checkpointed at batch boundaries, so a process that dies + mid-batch leaves every item `running`; the next one re-enters the + batch and, before this, dispatched every item again. Measured on the + DeepSeek-V4.1 run: the compute session was OOM-killed with 30 smiths + in flight, 16 of which had already written a LANDED decision -- a + restart would have paid for all 30 twice. So the first time this + process sees an item it asks the decision file first: one newer than + the item's last dispatch marker is that dispatch's answer, and is + collected instead of re-asked. A retry (``note``) never adopts -- the + decision on disk is the failure being retried. + """ + if self._decision_files and not note and item.n not in self._smiths_dispatched: + adopted = collect_decision_file(self._progress_ctx, self.workspace, "smith", item.n) + if adopted is not None: + print_message( + f"[bold cyan]→ smith {item.n} ({item.target}): adopting {adopted} written " + f"for a previous process's dispatch, instead of dispatching again[/bold cyan]", + None, + ) + self.journal.event( + "smith decision adopted on resume", f"{item.n} {item.target}: {adopted}" + ) + return f"adopted {adopted} from a previous dispatch" + self._smiths_dispatched.add(item.n) + layer = self._make_agent( + "smith", + self.prompts.smith, + required_tools=["append_smith_progress"], + session_mode="stateless", + ) + try: + prompt = self._smith_prompt(item, note) + if self._decision_files: + # Smiths number by queue item rather than by turn, and a batch + # re-run reuses those numbers, so the same staleness hazard + # applies here. + mark_decision_dispatch(self.workspace, "smith", item.n) + report = self.journal.record_dispatch("smith", f"{item.mode} {item.target}", prompt) + result = self._run_turn( + "smith", + f"{item.mode} {item.target}", + prompt, + iteration=item.n, + layer=layer, + ) + self.journal.record_result(report, result) + if self._decision_files: + decision = collect_decision_file( + self._progress_ctx, self.workspace, "smith", item.n + ) + if decision is None: + print_message( + f"[bold yellow]! smith {item.n} ({item.target}) ended its turn " + f"without a valid decision file at " + f"{decision_path(self.workspace, 'smith', item.n)}[/bold yellow]", + None, + ) + return result + finally: + layer.__exit__(None, None, None) + + def _report_instruction(self, role: str) -> str: + """How this role reports its decision, in the mode actually in force. + + The system prompt already carries the full protocol; this is the + one-line reminder at the end of a dispatch. It has to agree with + the system prompt, because a turn told to call a tool that was + never registered spends itself looking for it. + """ + if self._decision_files: + return "Before ending your turn, write your decision file to the path below." + return f"Before ending your turn, call `append_{role}_progress` once." + + def _read_instruction(self) -> str: + """How a role reads the prior entries in its own stage.""" + if self._decision_files: + return f"Read `{self.progress_path}`" + return "Call `read_latest_progress`" + + def _dispatch(self, role: str, *, what: str, prompt: str) -> str: + """Run one agent turn, recording the prompt **before** it is sent.""" + iteration = self._progress_ctx.current_iteration + # Every role, one place. The smith path does not come through here -- + # it owns its own layer per dispatch -- so `_smith_prompt` carries + # the same block explicitly. + context = self._context_block() + if context: + prompt = f"{prompt}\n\n{context}\n" + if self._decision_files: + # Stated here rather than only in the system prompt so the path + # is exact and un-guessable: the role is told the absolute file + # to write, not a pattern to instantiate. + target = decision_path(self.workspace, role, iteration) + prompt = ( + f"{prompt}\n\n" + f"---\n\n" + f"**Write your decision file to exactly this path when you finish:**\n\n" + f"`{target}`\n\n" + f"Nothing downstream reads your prose. This file is how this turn " + f"reaches the orchestrator.\n" + ) + # Stamp BEFORE the turn runs. `(role, iteration)` is not unique + # across a run -- `iteration` is 1 for every deliver turn and + # restarts per unit in assembly -- so a turn that ends without + # writing a decision leaves the PREVIOUS turn's verdict sitting at + # this path looking fresh. The marker is what lets + # `collect_decision_file` tell those apart. + mark_decision_dispatch(self.workspace, role, iteration) + + report = self.journal.record_dispatch(role, what, prompt) + result = self._run_turn(role, what, prompt, iteration=iteration) + self.journal.record_result(report, result) + + if self._decision_files: + decision = collect_decision_file(self._progress_ctx, self.workspace, role, iteration) + if decision is None: + # Deliberately not silent and deliberately not fatal here: + # the caller reads the decision back out of progress.yaml and + # already knows what to do with its absence. + print_message( + f"[bold yellow]! {role} ended its turn without a valid decision " + f"file at {decision_path(self.workspace, role, iteration)}" + f"[/bold yellow]", + None, + ) + return result + + def _handle_blocked_assembler(self, state: WorkflowState, log) -> bool: + """Triage a BLOCKED assembler. Returns True to retry the unit. + + Staircase's triage table, as control flow: a vocabulary gap is + something the orchestrator can fix by dispatching smith; the other + two are findings that end the run, because relaxing a gate or + guessing at a knowledge gap is exactly what must not happen quietly. + """ + entry = latest_entry(self.progress_path, "assembler") or {} + reason = str(entry.get("blocked_reason", "")).strip() + if reason == "vocabulary_gap": + missing = str(entry.get("missing_entry", "")).strip() + if not missing: + print_message( + "[bold red]✘ assembler reported a vocabulary gap without naming the " + "callable, so there is nothing to dispatch smith with.[/bold red]", + log, + ) + self._close_run(state, log, verdict="vocabulary-gap-unnamed") + return False + # An assembly turn routinely finds SEVERAL gaps at once --- they + # are usually one class, discovered together. `missing_entry` is + # a single string, so an assembler reporting three of them can + # only write "a/x; b/y; c/z", and dispatching that verbatim sends + # one smith after a target that does not exist. Split it, and + # dispatch them concurrently: gaps found in one turn have no + # dependency on each other by construction, since the turn saw + # all of them before it stopped. + targets = [part.strip() for part in re.split(r"[;,\n]+", missing) if part.strip()] + plural = "" if len(targets) == 1 else f" ({len(targets)} gaps)" + print_message( + f"[bold cyan]→ vocabulary gap{plural}: dispatching smith for " + f"{', '.join(targets)}[/bold cyan]", + log, + ) + units = assembly_units(self.task_data) + unit = units[state.unit_index] if state.unit_index < len(units) else UNIT_CORE + items = [ + QueueItem( + n=1000 + len(state.queue) + offset, mode="onboard", target=target, unit=unit + ) + for offset, target in enumerate(targets) + ] + for item in items: + state.queue.append({**item.as_dict(), "status": ITEM_RUNNING}) + self._checkpoint(state) + + workers = max(1, min(self.max_parallel_smiths or len(items), len(items))) + with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="gapsmith") as pool: + futures = {pool.submit(self._run_one_smith, item): item for item in items} + for future in as_completed(futures): + # The verdict is read from progress.yaml below, not from + # the returned prose --- but a dispatch that RAISED must + # not be mistaken for one that merely decided nothing. + try: + future.result() + except BaseException as exc: # keep the other dispatches + print_message( + f"[bold yellow]! gap dispatch {futures[future].target} raised " + f"{type(exc).__name__}: {exc}[/bold yellow]", + log, + ) + + unlanded = [] + for item in items: + landed = self._smith_decision_for(item) in ("LANDED", "FAILED_RECEIPT") + self._set_item_status(state, item.n, ITEM_DONE if landed else ITEM_FAILED) + if not landed: + unlanded.append(item.target) + self._checkpoint(state) + if not unlanded: + return True + print_message(f"[bold red]✘ smith could not land {', '.join(unlanded)}[/bold red]", log) + self._close_run(state, log, verdict="smith-blocked") + return False + + print_message( + f"[bold red]✘ assembler blocked ({reason or 'unspecified'}) — ending the run " + f"with the audit trail.[/bold red]", + log, + ) + self.journal.section(f"assembler blocked: {reason}", str(entry.get("summary", ""))) + self._close_run(state, log, verdict=f"assembler-blocked:{reason}") + return False + + def _append_open_risks(self) -> None: + """Carry the reviewer's unresolved items into plan.md and corrections.md. + + This is what makes the plan cap a budget rather than a gate: the + review's value is not thrown away, it is handed to the reader who + can act on it. + + Both files, because they reach different readers. `plan.md` is read + by whoever opens the plan — the assembler, and a human. The + corrections file is read by **every dispatch**, including smiths + whose queue was built from the vocabulary this review is rejecting. + That is the case the mechanism exists for: a reviewer's finding about + a GEMM width that arrives after the smith carrying that GEMM has + already been queued. + """ + entry = latest_entry(self.progress_path, "plan_reviewer") or {} + summary = str(entry.get("summary", "")).strip() + if not summary: + return + with self.plan_path.open("a", encoding="utf-8") as handle: + handle.write( + "\n\n## Open risks (unresolved at the plan-round budget)\n\n" + "The PlanReviewer raised these and the plan phase ran out of rounds " + "before they were resolved. They are not blockers — the build starts " + "anyway — but they are the most likely places for it to go wrong.\n\n" + f"{summary}\n" + ) + self.journal.section("plan budget spent — open risks recorded", summary) + self._record_correction( + f"the PlanReviewer's unresolved items at round {self.max_plan_rounds + 1}", + summary, + ) + + def _close_run(self, state: WorkflowState, log, *, verdict: str) -> None: + state.done = True + state.verdict = verdict + self._checkpoint(state) + self.journal.close( + summary="\n".join( + [ + f"- verdict: **{verdict}**", + f"- stage reached: `{state.stage}`", + f"- units complete: {', '.join(state.units_done) or '(none)'}", + f"- human stops: **{state.human_stops}**", + "", + "### Gates", + "", + "```", + self.gate_log.render() or "(none run)", + "```", + ] + ) + ) + print_message( + f"[bold]run closed — verdict {verdict}; human stops: {state.human_stops}[/bold]", + log, + ) + + # ------------------------------------------------------------------ prompts + + def _plan_drafter_prompt(self, round_no: int) -> str: + redraft = "" + # Re-derive on demand as well as recording it when it happens, so a + # RESUME into round 2 carries the same message a continuous run does. + # Otherwise the recovery path works only when nothing interrupted the + # failure, which is the case least likely to need it. + if round_no > 1 and not self._vocabulary_errors and self.vocabulary_path.is_file(): + try: + load_vocabulary(self.vocabulary_path, units=assembly_units(self.task_data)) + except Exception as exc: + self._vocabulary_errors = str(exc) + if round_no > 1 and self._vocabulary_errors: + # A SCHEMA rejection, not a review. The reviewer never ran, so + # pointing at progress.yaml would point at nothing. + redraft = ( + "This is a **re-draft**, and the reason is mechanical: " + "`vocabulary.yaml` did not pass the orchestrator's schema check, " + "so the PlanReviewer never saw it. **Nothing is wrong with your " + "analysis** — fix the file and keep everything else.\n\n" + "The validator's own output, verbatim:\n\n```\n" + f"{self._vocabulary_errors}\n```\n\n" + "Revise `vocabulary.yaml` in place. Do not re-derive the plan; " + "`plan.md` stands unless one of these errors contradicts it.\n\n" + ) + elif round_no > 1: + redraft = ( + "This is a **re-draft**. " + self._read_instruction() + " to fetch the " + "PlanReviewer's REJECT and address every item it names. You are in the " + "same session as your first draft, so do not re-derive what you already " + "established — revise it.\n\n" + ) + return ( + f"Workspace: {self.workspace}\n" + f"Plan round: {round_no} of at most {self.max_plan_rounds + 1}\n\n" + f"{redraft}" + f"Read `{self.task_path}` for the run's spec: the checkpoint, the " + f"architecture, the parallel segment, the features and the accuracy anchor.\n\n" + f"{self._surface_block}\n" + f"{self._runtime_block}\n" + f"Write `{self.plan_path}` and `{self.vocabulary_path}` as your system " + f"prompt specifies. The orchestrator parses `vocabulary.yaml` and will " + f"refuse a malformed one, so follow the schema exactly.\n\n" + + self._report_instruction("plan_drafter") + ) + + def _plan_reviewer_prompt(self, round_no: int) -> str: + remaining = self.max_plan_rounds + 1 - round_no + budget = ( + "This is the **final round**: if you REJECT, your items are recorded in " + "plan.md as open risks and the build starts anyway. Make the summary count." + if remaining <= 0 + else f"{remaining} further round(s) remain after this one." + ) + return ( + f"Workspace: {self.workspace}\n" + f"Plan round: {round_no}\n\n" + f"{budget}\n\n" + f"Read `{self.task_path}` (the spec), `{self.plan_path}` (the plan) and " + f"`{self.vocabulary_path}` (the typed mapping and dispatch queue). Call " + + self._read_instruction() + + " for the PlanDrafter's summary.\n\n" + "This is a paper review: do not build, run or test anything.\n\n" + + self._report_instruction("plan_reviewer") + + " " + "`APPROVE` or `REJECT`." + ) + + def _smith_prompt(self, item: QueueItem, note: str = "") -> str: + surface = f"The surface to certify: {item.surface}\n" if item.surface else "" + # Directly under the dispatch line, because it changes what the turn + # is for: finish what the last attempt started, not begin again. + retry = f"\n{note}\n" if note else "" + # Tiers 1-2 wrap something that exists; 3-4 mean this dispatch writes + # the implementation. Saying which changes what the turn is for. + tier = f"Sourcing: **{item.tier_brief}**.\n" + ( + "This entry has no implementation yet — you are writing it, not " + "wrapping one. Correct-and-slow is acceptable; the tuning phase " + "replaces it. Say in the contract what you rejected at the cheaper " + "rungs.\n" + if item.writes_implementation + else "" + ) + # Device assignment, because the launch protocol hands EVERY role the + # same line -- `CUDA_VISIBLE_DEVICES=0 # you own device(s) 0..0` -- + # and then warns that concurrent scripts on one device contend. At the + # default ten parallel smiths that warning is a certainty, and the + # symptom is `Executor creation failed due to insufficient GPU memory` + # on a device with 284 GB, which reads as a target that allocates too + # much. Measured on this host with four concurrent full-model loads. + # + # Round-robin over the node's devices by dispatch number. Not a + # scheduler: two smiths can still land on one device, but ten cannot + # land on device 0. + devices = int((self.task_data.get("environment") or {}).get("gpus_per_node") or 1) + assignment = "" + if devices > 1: + assignment = ( + f"\n**Use `CUDA_VISIBLE_DEVICES={item.n % devices}` for this dispatch**, not the " + f"`0` in the launch protocol above. This node has {devices} devices and your " + f"batch is running concurrently; every smith taking device 0 turns a memory " + f"error into the symptom.\n" + ) + # A smith is dispatched from `_run_one_smith` rather than from + # `_dispatch`, so the shared context has to be added here too. It + # matters most for this role: the corrections file exists because a + # plan review landed after the smith queue had been built. + context = self._context_block() + return ( + f"Workspace: {self.workspace}\n" + f"Repository: {self.task_data['trtllm_repo_path']}\n" + f"Architecture: {self.task_data['arch']}\n\n" + f"Dispatch {item.n}: **{item.mode}** the catalog entry `{item.target}`.\n" + f"{retry}" + f"{assignment}" + f"Load the `{item.skill}` skill before anything else.\n" + f"{surface}" + f"{tier}\n" + f"{self._kernel_sources_block}\n" + f"Read `{self.vocabulary_path}` for the computation step this entry serves " + f"and the values this checkpoint passes it.\n\n" + f"One run, one entry. Commit by explicit file path. " + + self._report_instruction("smith") + + (f"\n\n{context}\n" if context else "") + ) + + def _assembler_prompt(self, unit: str, turn: int) -> str: + feature_skills = "" + if unit != "core": + from .task_schema import FEATURE_SKILLS + + feature_skills = ( + f" and `{FEATURE_SKILLS.get(unit, 'staircase-modeling-core')}` for the " + f"`{unit}` feature" + ) + history = "" + if turn > 1: + history = ( + f"This is turn {turn} for this unit. " + self._read_instruction() + " for " + "what the last turn did, and read the gate output above — you are in " + "the same session, so build on it rather than starting over.\n\n" + ) + return ( + f"Workspace: {self.workspace}\n" + f"Repository: {self.task_data['trtllm_repo_path']}\n" + f"Target: {self._preflight.target_rel}\n" + f"Unit: **{unit}** (turn {turn} of at most " + f"{self.max_assembler_turns_per_unit})\n\n" + f"{history}" + f"Load `staircase-modeling-core`{feature_skills} before anything else.\n\n" + f"Read `{self.task_path}` (the spec), `{self.plan_path}` (the plan, " + f"including any `## Open risks`) and `{self.vocabulary_path}` (the " + f"computation-step list you derive against).\n\n" + f"{self._gate_history_block()}" + f"Write the unit. The orchestrator runs the boot gate as soon as you " + f"finish, and the accuracy gate once the unit looks right — do not run " + f"them yourself, and do not reformat: the run lints once, at delivery.\n\n" + + self._report_instruction("assembler") + ) + + def _reviewer_prompt(self, unit: str, boot: gates_mod.GateResult) -> str: + return ( + f"Workspace: {self.workspace}\n" + f"Repository: {self.task_data['trtllm_repo_path']}\n" + f"Target: {self._preflight.target_rel}\n" + f"Unit: **{unit}**\n\n" + f"The boot gate has already run:\n\n```\n{boot.summary()}\n" + f"{boot.stdout_tail}\n```\n\n" + + self._read_instruction() + + " for the assembler's report — its derivation " + "table, vocabulary audit and dump review are what you check against the " + "diff.\n\n" + "Read the diff and decide. Do not re-run the gates.\n\n" + + self._report_instruction("reviewer") + ) + + #: The line every modeling_v2 target's `_emit_runtime_values` writes before + #: its dump. Named here because the auditor has to find dumps this process + #: never saw — see `_auditor_dump`. + RUNTIME_VALUE_MARKER = "=== MODELINGV2 RUNTIME VALUE DUMP" + + def _auditor_dump(self) -> str: + """Every gate's runtime-value dump, or directions to where they are. + + Not the boot gate's alone, which is what this passed for its whole + life. Two things were wrong with that and only the first is obvious. + + The boot gate drives two short prompts, so its dump records the + engine's *warm-up* shapes and almost nothing else; measured on this + project, eleven of twelve boot dumps carried only 8191/8192 context + rows. The shapes a model actually serves at appear in the **accuracy** + gates, and both of this run's auditor findings were visible only there. + + The second is that `gate_log` is in-memory, so a run resumed into + `deliver` — which is the normal way to reach this stage after a repair + — has no gate history at all and this returns nothing however it is + wired. So when there is nothing to inline, say where the dumps live + rather than silently handing the auditor an empty section: both + auditors in this project recovered them from the spool logs by hand, + and that should not have been an act of initiative. + """ + chunks = [ + f"### from the `{result.name}` gate\n\n```\n{result.stdout_tail.rstrip()}\n```" + for result in self.gate_log.results + if self.RUNTIME_VALUE_MARKER in (result.stdout_tail or "") + ] + if chunks: + return "\n\n".join(chunks) + + spool = getattr(self._gate_env(), "spool", None) + where = f"`{Path(spool) / 'logs'}/*.log`" if spool else "this run's gate logs" + return ( + "**No dump was captured in this process.** `gate_log` is in-memory, so a run " + "resumed into `deliver` carries no gate history — that is expected here and is " + "not a sign the gates did not run.\n\n" + f"The dumps are on disk. Search {where} for the marker " + f"`{self.RUNTIME_VALUE_MARKER}` and audit the **union** across every log that " + "has one. Prefer the accuracy gates' dumps: a boot gate drives two short prompts " + "and records warm-up shapes, not the shapes the model serves at." + ) + + def _auditor_prompt(self, report, boot: gates_mod.GateResult | None) -> str: + return ( + f"Workspace: {self.workspace}\n" + f"Repository: {self.task_data['trtllm_repo_path']}\n" + f"Target: {self._preflight.target_rel}\n" + f"Architecture: {self.task_data['arch']}\n\n" + f"{summarize_for_auditor(report, self._auditor_dump())}\n" + f"## Gates\n\n```\n{self.gate_log.render() or '(none in this process — see above)'}" + f"\n```\n\n" + f"Do the one check that is left: every value in the runtime-value dump " + f"against its catalog entry's certified column, on this architecture.\n\n" + + self._report_instruction("auditor") + ) + + def _gate_history_block(self) -> str: + rendered = self.gate_log.render() + if not rendered: + return "" + return f"Gates so far:\n\n```\n{rendered}\n```\n\n" + + +if __name__ == "__main__": + from .cli import main + + main() diff --git a/agent-flow/pyproject.toml b/agent-flow/pyproject.toml index a53028a08467..b68ae6ce7b8d 100644 --- a/agent-flow/pyproject.toml +++ b/agent-flow/pyproject.toml @@ -23,6 +23,7 @@ agent-team = "agent_flow.workflows.agent_team.cli:main" modeling-bringup = "agent_flow.workflows.modeling_bringup.cli:main" perf-analyze = "agent_flow.workflows.perf_analyze.cli:main" perf-optimize = "agent_flow.workflows.perf_optimize.cli:main" +staircase-onboard = "agent_flow.workflows.staircase_onboard.cli:main" [build-system] requires = ["setuptools>=68.0"] @@ -33,8 +34,27 @@ include = ["agent_flow*"] # Prompt assets read at runtime (the analyzer copies the TRT-LLM kernel # taxonomy into its workspace), so they must survive a non-editable install. +# +# staircase_onboard ships its agents' knowledge the same way: the skills and +# their reference documents are copied into each run's workspace, so they are +# versioned with the workflow that dispatches them instead of depending on a +# checkout's `.claude/skills` having the right thing in it. The example task +# spec is listed too — a workflow whose `--task` template is missing from an +# install is one nobody can start. [tool.setuptools.package-data] "agent_flow.workflows.perf_analyze" = ["assets/*.json"] +"agent_flow.workflows.staircase_onboard" = [ + "task.example.yaml", + "runtime_index.yaml", + "kernel_sources.yaml", + "skills/**/*.md", + # Named by the failure message on every infrastructure stop, so an + # install without it points the operator at a path that is not there — + # at the exact moment they are already debugging something. + "TROUBLESHOOTING.md", + "BUILD_AND_TEST.md", + "README.md", +] #################################################################################################### ############################### LINTING, FORMATTING AND TESTING CONFIGURATION #################### @@ -86,6 +106,9 @@ extend-ignore = [ # Markdown tables in prompt string literals — a table row must stay on one line. "agent_flow/workflows/modeling_bringup/prompts/plan_drafter_extra.py" = ["E501"] "agent_flow/workflows/modeling_bringup/prompts/reviewer_extra.py" = ["E501"] +# Same reason: every staircase role prompt ends in a Red-flags table, whose +# rows are `| Impulse | Reality |` pairs. Wrapping one breaks the table. +"agent_flow/workflows/staircase_onboard/prompts/*.py" = ["E501"] [tool.ruff.lint.pycodestyle] max-line-length = 120 diff --git a/agent-flow/tests/runtime/__init__.py b/agent-flow/tests/runtime/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/agent-flow/tests/runtime/__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/runtime/test_command_executors.py b/agent-flow/tests/runtime/test_command_executors.py new file mode 100644 index 000000000000..fb03cf20c278 --- /dev/null +++ b/agent-flow/tests/runtime/test_command_executors.py @@ -0,0 +1,306 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Command execution, with no GPU, no scheduler and no model in sight. + +The local executor is exercised for real — ``echo``, ``sh -c 'exit 3'``, a +command that outlives its deadline — because that is cheap and the thing +being checked is exactly the process boundary. The spool executor is +exercised against a fake watcher, since what it is responsible for is the +file protocol and not what runs on the far side of it. + +The Slurm executor gets contract tests only. It renders nothing and decides +nothing; asserting more would require a controller. +""" + +from __future__ import annotations + +import sys +import threading +import time + +import pytest + +from agent_flow.runtime.commands import ( + NOT_FOUND_RETURNCODE, + TIMEOUT_RETURNCODE, + CommandSpec, + ExecutionRecord, + LocalCommandExecutor, + SlurmStepExecutor, + SpoolCommandExecutor, +) + +# --------------------------------------------------------------- CommandSpec + + +def test_a_string_is_a_shell_line_and_a_sequence_is_not(): + """The distinction decides whether a shell gets to split the words.""" + assert CommandSpec(argv="a b 'c d'").is_shell_line + assert not CommandSpec(argv=("a", "b", "c d")).is_shell_line + + +def test_a_sequence_renders_quoted_so_a_log_can_be_pasted_back(): + assert CommandSpec(argv=("echo", "a b")).rendered == "echo 'a b'" + + +def test_a_shell_line_renders_verbatim(): + """Re-quoting an `srun` line would mean reimplementing the shell.""" + line = "FOO=1 srun --ntasks=8 bash -c 'exec thing'" + assert CommandSpec(argv=line).rendered == line + + +# ------------------------------------------------------------ ExecutionRecord + + +def test_a_record_is_ok_only_when_it_exited_zero_and_finished(): + assert ExecutionRecord(command="x", returncode=0).ok + assert not ExecutionRecord(command="x", returncode=1).ok + # A timeout has said nothing about the thing it was running, so it is + # never a pass -- even though the code it carries is conventional. + assert not ExecutionRecord(command="x", returncode=0, timed_out=True).ok + + +def test_the_tail_is_the_end_of_stdout_and_stderr_together(): + record = ExecutionRecord(command="x", returncode=0, stdout="a\nb\nc\n", stderr="err\n") + assert record.tail(2) == "c\nerr" + + +# ------------------------------------------------------- LocalCommandExecutor + + +def test_a_shell_line_runs_through_a_shell(): + record = LocalCommandExecutor().execute(CommandSpec(argv="echo one; echo two", timeout_s=30)) + assert record.ok + assert record.stdout.split() == ["one", "two"] + assert record.executor == "local" + + +def test_an_argv_sequence_is_not_word_split(): + """A probe's argument must survive spaces without quoting ceremony.""" + record = LocalCommandExecutor().execute(CommandSpec(argv=("echo", "a b c"), timeout_s=30)) + assert record.stdout == "a b c\n" + + +def test_env_is_layered_over_the_parent_rather_than_replacing_it(): + """Everything here runs inside an environment somebody else set up. + + A venv's PATH, a container's library paths, TMPDIR. Replacing the + environment outright is the stricter choice and the wrong one: it turns + a working host into a broken one. + """ + record = LocalCommandExecutor().execute( + CommandSpec( + argv="echo $STAIRCASE_PROBE:$PATH", env={"STAIRCASE_PROBE": "set"}, timeout_s=30 + ) + ) + assert record.stdout.startswith("set:") + assert len(record.stdout.strip()) > len("set:") + + +def test_a_non_zero_exit_is_reported_not_raised(): + record = LocalCommandExecutor().execute(CommandSpec(argv="exit 3", timeout_s=30)) + assert record.returncode == 3 + assert not record.ok + assert not record.timed_out + + +def test_stderr_is_kept_separate_from_stdout(): + record = LocalCommandExecutor().execute( + CommandSpec(argv="echo out; echo err 1>&2", timeout_s=30) + ) + assert record.stdout.strip() == "out" + assert record.stderr.strip() == "err" + + +def test_a_command_that_outlives_its_deadline_is_timed_out_not_merely_failed(): + """The two need different handling. + + A command that timed out has said nothing about what it was running; one + that exited non-zero has. Collapsing them loses that. + """ + started = time.monotonic() + record = LocalCommandExecutor().execute(CommandSpec(argv="sleep 30", timeout_s=1)) + assert record.timed_out + assert record.returncode == TIMEOUT_RETURNCODE + assert not record.ok + assert time.monotonic() - started < 20 + + +def test_a_missing_executable_is_distinguishable_from_a_disagreement(): + """127 says the probe could not run; non-zero says it ran and disagreed.""" + record = LocalCommandExecutor().execute( + CommandSpec(argv=("staircase-no-such-binary-xyz",), timeout_s=30) + ) + assert record.returncode == NOT_FOUND_RETURNCODE + assert "staircase-no-such-binary-xyz" in record.stderr + + +def test_an_interrupt_is_not_turned_into_a_failed_record(monkeypatch): + """Ctrl-C has to reach the orchestrator, not become a gate that failed. + + The orchestrator catches ``KeyboardInterrupt`` around the whole run to + checkpoint and tell the operator how to resume. An executor that caught + it here and returned a non-zero record would turn "the operator stopped + this" into "the target failed its gate", and the run would close with a + verdict it never earned. + """ + import subprocess as subprocess_mod + + def interrupted(*args, **kwargs): + raise KeyboardInterrupt + + monkeypatch.setattr(subprocess_mod, "run", interrupted) + with pytest.raises(KeyboardInterrupt): + LocalCommandExecutor().execute(CommandSpec(argv="echo hi", timeout_s=30)) + + +def test_the_record_carries_the_provenance_a_verdict_rests_on(): + record = LocalCommandExecutor().execute(CommandSpec(argv="echo hi", timeout_s=30, label="boot")) + assert record.command == "echo hi" + assert record.executor == "local" + assert record.label == "boot" + assert record.duration_s >= 0.0 + assert record.started_at > 0.0 + + +def test_cwd_is_honoured_when_asked_for(tmp_path): + (tmp_path / "marker").write_text("x", encoding="utf-8") + record = LocalCommandExecutor().execute( + CommandSpec(argv=("ls",), cwd=str(tmp_path), timeout_s=30) + ) + assert record.stdout.split() == ["marker"] + + +def test_the_local_executor_declares_what_it_can_do(): + caps = LocalCommandExecutor.capabilities + assert caps.gpu is True + assert caps.distributed is False + assert caps.persistent_allocation is False + + +# ------------------------------------------------------- SpoolCommandExecutor + + +def _watch(root, *, returncode=0, output="done\n", delay=0.02): + """A fake session watcher: drain `pending`, write the log, then the rc.""" + + def run(): + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + scripts = sorted((root / "pending").glob("*.sh")) if (root / "pending").is_dir() else [] + if scripts: + name = scripts[0].stem + time.sleep(delay) + (root / "logs").mkdir(parents=True, exist_ok=True) + (root / "logs" / f"{name}.log").write_text(output, encoding="utf-8") + (root / "done").mkdir(parents=True, exist_ok=True) + # `.rc` last, same ordering rule as the dispatch mailbox. + (root / "done" / f"{name}.rc").write_text(f"{returncode}\n", encoding="utf-8") + return + time.sleep(0.005) + + thread = threading.Thread(target=run, daemon=True) + thread.start() + return thread + + +def test_a_spool_command_round_trips_through_the_filesystem(tmp_path): + _watch(tmp_path, output="lm-eval gsm8k average accuracy: 90.5\n") + record = SpoolCommandExecutor(tmp_path, poll_s=0.01).execute( + CommandSpec(argv="trtllm-eval --model x gsm8k", timeout_s=30, label="accuracy") + ) + assert record.ok + assert "90.5" in record.stdout + assert record.executor == "spool" + assert record.label == "accuracy" + + +def test_the_receipt_is_the_exit_code(tmp_path): + _watch(tmp_path, returncode=7, output="it failed\n") + record = SpoolCommandExecutor(tmp_path, poll_s=0.01).execute( + CommandSpec(argv="false", timeout_s=30) + ) + assert record.returncode == 7 + assert not record.ok + + +def test_the_script_carries_the_environment_as_exports(tmp_path): + """The far side is a fresh shell; an env dict has to become shell text.""" + executor = SpoolCommandExecutor(tmp_path, poll_s=0.01) + _watch(tmp_path) + executor.execute( + CommandSpec(argv="echo hi", env={"TRTLLM_MODELING_V2": "require"}, timeout_s=30) + ) + # The watcher consumed it, so read what it was handed: re-submit with no + # watcher and inspect the pending script instead. + lonely = SpoolCommandExecutor(tmp_path, poll_s=0.01) + lonely.execute(CommandSpec(argv="echo hi", env={"A": "b c"}, timeout_s=0)) + script = sorted((tmp_path / "pending").glob("*.sh"))[-1].read_text(encoding="utf-8") + assert script.startswith("#!/bin/bash\nset -u\n") + assert "export A='b c'\n" in script + assert script.rstrip().endswith("echo hi") + + +def test_the_script_is_renamed_into_pending_never_written_in_place(tmp_path): + """A watcher must not be able to claim a half-written script.""" + SpoolCommandExecutor(tmp_path, poll_s=0.01).execute(CommandSpec(argv="echo hi", timeout_s=0)) + assert not list(tmp_path.glob("*.tmp")) + assert len(list((tmp_path / "pending").glob("*.sh"))) == 1 + + +def test_a_spool_command_nobody_drains_times_out_with_whatever_log_exists(tmp_path): + record = SpoolCommandExecutor(tmp_path, poll_s=0.01).execute( + CommandSpec(argv="sleep 100", timeout_s=0) + ) + assert record.timed_out + assert record.returncode == TIMEOUT_RETURNCODE + assert record.log_path and record.log_path.endswith(".log") + + +def test_the_spool_declares_that_its_allocation_is_persistent(): + """The reason to prefer it: one queue wait per run, not one per command.""" + caps = SpoolCommandExecutor.capabilities + assert caps.persistent_allocation is True + assert caps.distributed is True + assert caps.gpu is True + + +# ---------------------------------------------------------- SlurmStepExecutor + + +def test_the_slurm_executor_runs_the_rendered_line_here_and_says_so(): + """An `srun` line is a shell line; the step is still a local subprocess. + + What the separate class buys is the two facts nobody can read off the + command string: that this executor is distributed, and that the record + came from Slurm. + """ + record = SlurmStepExecutor().execute(CommandSpec(argv="echo pretend-srun", timeout_s=30)) + assert record.ok + assert record.stdout.strip() == "pretend-srun" + assert record.executor == "slurm" + + +def test_the_slurm_executor_declares_itself_distributed(): + assert SlurmStepExecutor.capabilities.distributed is True + assert SlurmStepExecutor(max_world_size=8).capabilities.max_world_size == 8 + + +def test_bounding_the_world_size_does_not_leak_into_the_class(): + """Two runs of different sizes in one process must not see each other's.""" + sized = SlurmStepExecutor(max_world_size=8) + assert sized.capabilities.max_world_size == 8 + assert SlurmStepExecutor.capabilities.max_world_size is None + + +def test_every_executor_reports_the_command_it_actually_ran(): + """The record has to be defensible without the caller's memory of it.""" + for executor in (LocalCommandExecutor(), SlurmStepExecutor()): + record = executor.execute(CommandSpec(argv=("echo", "x y"), timeout_s=30)) + assert record.command == "echo 'x y'" + + +def test_the_executors_are_importable_without_a_gpu_or_a_scheduler(): + """Every test in this file is the proof; this states the intent.""" + assert sys.modules["agent_flow.runtime.commands.local"] diff --git a/agent-flow/tests/runtime/test_delegated_executor.py b/agent-flow/tests/runtime/test_delegated_executor.py new file mode 100644 index 000000000000..48e51d88a8a7 --- /dev/null +++ b/agent-flow/tests/runtime/test_delegated_executor.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Running a command through an agent without letting the agent decide anything. + +The whole design rests on one property, and most of this file is about it: +the agent is a **trigger, not a witness**. It gets a job submitted; the job +writes its own log and exit code; the orchestrator reads those. An agent +that lies, hallucinates a result, or says nothing at all must not be able to +move the verdict. + +The rest covers the quoting hazard that motivated submitting a script rather +than a command line, and the ordering rule that makes a receipt trustworthy. +""" + +from __future__ import annotations + +import threading +import time +from pathlib import Path + +from agent_flow.runtime.commands import TIMEOUT_RETURNCODE, CommandSpec, DelegatedCommandExecutor +from agent_flow.runtime.turns import TurnExecutor, TurnOutcome, TurnRequest + + +class FakeAgent(TurnExecutor): + """A driver that runs the script it is handed, like a real one would. + + ``reply`` is what it claims afterwards, which is the interesting knob: + nothing downstream should depend on it. + """ + + name = "fake-agent" + + def __init__(self, *, reply: str = "submitted", submits: bool = True, delay: float = 0.01): + self.reply = reply + self.submits = submits + self.delay = delay + self.requests: list[TurnRequest] = [] + + def run_script(self, prompt: str) -> None: + # The prompt carries the resolved submission call; pull the script + # path out of it the way a real driver would read its argument. + import re + import subprocess + + match = re.search(r"bash (\S+\.sh)", prompt) + if not match: + return + script = match.group(1) + time.sleep(self.delay) + subprocess.run(["bash", script], check=False, capture_output=True) + + def run(self, request: TurnRequest) -> TurnOutcome: + self.requests.append(request) + if self.submits: + threading.Thread(target=self.run_script, args=(request.prompt,), daemon=True).start() + return TurnOutcome(role=request.role, text=self.reply, transport=self.name) + + +def _executor(tmp_path: Path, agent: TurnExecutor, **kwargs) -> DelegatedCommandExecutor: + return DelegatedCommandExecutor( + agent, + root=tmp_path / "broker", + render_submission=lambda script: f'submit(command=["bash", "-lc", "bash {script}"])', + poll_s=0.01, + **kwargs, + ) + + +# --------------------------------------------------- the verdict is not the agent's + + +def test_the_exit_code_comes_from_the_job_not_from_the_agent(): + """An agent claiming success cannot make a failing command pass.""" + import tempfile + + with tempfile.TemporaryDirectory() as tmp: + agent = FakeAgent(reply="Everything passed, accuracy was 99.9!") + record = _executor(Path(tmp), agent).execute(CommandSpec(argv="exit 7", timeout_s=30)) + assert record.returncode == 7 + assert not record.ok + + +def test_the_output_comes_from_the_log_not_from_the_agent(): + import tempfile + + with tempfile.TemporaryDirectory() as tmp: + agent = FakeAgent(reply="I saw accuracy: 99.9") + record = _executor(Path(tmp), agent).execute( + CommandSpec(argv="echo 'lm-eval gsm8k average accuracy: 42.0'", timeout_s=30) + ) + assert "42.0" in record.stdout + assert "99.9" not in record.stdout + + +def test_an_agent_that_says_nothing_useful_changes_nothing(): + import tempfile + + with tempfile.TemporaryDirectory() as tmp: + record = _executor(Path(tmp), FakeAgent(reply="")).execute( + CommandSpec(argv="echo fine", timeout_s=30) + ) + assert record.ok and record.stdout.strip() == "fine" + + +def test_a_job_nobody_submits_times_out_rather_than_passing(): + """The failure mode if the agent silently does not submit.""" + import tempfile + + with tempfile.TemporaryDirectory() as tmp: + agent = FakeAgent(reply="done!", submits=False) + record = _executor(Path(tmp), agent).execute(CommandSpec(argv="echo never", timeout_s=0)) + assert record.timed_out + assert record.returncode == TIMEOUT_RETURNCODE + assert not record.ok + + +# ------------------------------------------------------------------- the script + + +def test_the_command_is_submitted_as_a_path_not_as_a_quoted_string(tmp_path): + """The reason this executor writes a script at all. + + Submission templates embed their payload in a quoted string, and a real + gate command carries `--model_dir="/p"` and + `--chat_template_kwargs={"enable_thinking": false}`. Interpolating one + produces a call that is mis-parsed or truncated. + """ + agent = FakeAgent(submits=False) + hostile = '/venv/bin/python run.py --model_dir="/models/x" --kw={"enable_thinking": false}' + _executor(tmp_path, agent).execute(CommandSpec(argv=hostile, timeout_s=0)) + + prompt = agent.requests[0].prompt + assert '--model_dir="' not in prompt.split("```python")[1].split("```")[0] + script = sorted((tmp_path / "broker").glob("*.sh"))[0] + assert hostile in script.read_text(encoding="utf-8") + + +def test_the_script_publishes_its_own_exit_code(tmp_path): + agent = FakeAgent(submits=False) + _executor(tmp_path, agent).execute(CommandSpec(argv="true", timeout_s=0)) + body = sorted((tmp_path / "broker").glob("*.sh"))[0].read_text(encoding="utf-8") + assert "rc=$?" in body + # The receipt is renamed into place, so seeing it means the log is done. + assert ".tmp\n" in body and "mv " in body + + +def test_the_script_does_not_set_e_because_the_exit_code_is_the_measurement(tmp_path): + """`set -e` would kill the script before it could record a failure.""" + agent = FakeAgent(submits=False) + _executor(tmp_path, agent).execute(CommandSpec(argv="false", timeout_s=0)) + body = sorted((tmp_path / "broker").glob("*.sh"))[0].read_text(encoding="utf-8") + assert "set -u" in body + assert "set -e" not in body + + +def test_the_environment_travels_as_exports(tmp_path): + agent = FakeAgent(submits=False) + _executor(tmp_path, agent).execute( + CommandSpec(argv="env", env={"TRTLLM_MODELING_V2": "require"}, timeout_s=0) + ) + body = sorted((tmp_path / "broker").glob("*.sh"))[0].read_text(encoding="utf-8") + assert "export TRTLLM_MODELING_V2=require" in body + + +def test_the_environment_reaches_the_command(tmp_path): + record = _executor(tmp_path, FakeAgent()).execute( + CommandSpec( + argv="echo $STAIRCASE_DELEGATED", env={"STAIRCASE_DELEGATED": "yes"}, timeout_s=30 + ) + ) + assert record.stdout.strip() == "yes" + + +# ------------------------------------------------------------------- the brief + + +def test_the_agent_is_handed_a_resolved_call_and_told_not_to_change_it(tmp_path): + """No image to pick, no partition to discover, no flags to invent.""" + agent = FakeAgent(submits=False) + _executor(tmp_path, agent).execute(CommandSpec(argv="true", timeout_s=0)) + prompt = agent.requests[0].prompt + assert "change nothing in it" in prompt + assert "blocking" in prompt.lower() + assert "nothing reads your answer" in prompt + + +def test_the_turn_is_dispatched_to_the_runner_role_by_default(tmp_path): + agent = FakeAgent(submits=False) + _executor(tmp_path, agent).execute(CommandSpec(argv="true", timeout_s=0)) + assert agent.requests[0].role == "runner" + + +def test_the_record_names_its_artifacts_for_the_audit_trail(tmp_path): + record = _executor(tmp_path, FakeAgent()).execute( + CommandSpec(argv="echo hi", timeout_s=30, label="boot") + ) + assert record.executor == "delegated" + assert record.label == "boot" + assert record.log_path and record.log_path.endswith(".log") + assert len(record.artifacts) == 3 + + +def test_it_declares_itself_distributed_and_gpu_capable(): + caps = DelegatedCommandExecutor.capabilities + assert caps.gpu is True and caps.distributed is True diff --git a/agent-flow/tests/runtime/test_turn_executors.py b/agent-flow/tests/runtime/test_turn_executors.py new file mode 100644 index 000000000000..264d22af4d95 --- /dev/null +++ b/agent-flow/tests/runtime/test_turn_executors.py @@ -0,0 +1,353 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The contract both turn transports answer. + +The point of the abstraction is that an orchestrator cannot tell them apart, +so most of this file is parametrized over both. What differs — a mailbox has +request ids and artifacts, a direct call has a session — is asserted +separately and deliberately, because those are the differences a caller is +allowed to see. +""" + +from __future__ import annotations + +import threading +import time + +import pytest + +from agent_flow.runtime.turns import ( + FileDispatchTurnExecutor, + InProcessTurnExecutor, + TurnExecutionError, + TurnOutcome, + TurnRequest, + TurnTimeout, +) + + +class RecordingLayer: + """Stands in for an ``AgentLayer``: called with a prompt, returns text.""" + + def __init__(self, reply: str = "answered", name: str = "layer") -> None: + self.reply = reply + self.layer_name = name + self.prompts: list[str] = [] + + def __call__(self, prompt: str) -> str: + self.prompts.append(prompt) + return self.reply + + +def _answer_after(executor: FileDispatchTurnExecutor, request_id: str, result: str, delay=0.02): + """Play the driver from another thread.""" + + def run(): + time.sleep(delay) + executor.answer(request_id, result) + + thread = threading.Thread(target=run, daemon=True) + thread.start() + return thread + + +@pytest.fixture(params=["inprocess", "file-dispatch"]) +def executor_pair(request, tmp_path): + """``(executor, arrange_answer, prompt_seen)`` for each transport. + + ``arrange_answer(role, text)`` makes that transport return ``text`` for + the next turn, however that transport gets its answers. + ``prompt_seen()`` returns what the far side was actually handed. + """ + if request.param == "inprocess": + layer = RecordingLayer() + + def arrange(role: str, text: str) -> None: + layer.reply = text + + yield InProcessTurnExecutor(lambda req: layer), arrange, lambda: layer.prompts[-1] + else: + executor = FileDispatchTurnExecutor(tmp_path, poll_s=0.01) + counter = {"n": 0} + + def arrange(role: str, text: str) -> None: + counter["n"] += 1 + _answer_after(executor, f"{counter['n']:03d}-{role}", text) + + def prompt_seen() -> str: + latest = sorted(executor.root.glob("*.prompt.md"))[-1] + return latest.read_text(encoding="utf-8") + + yield executor, arrange, prompt_seen + + +# ------------------------------------------------- the part that must not differ + + +def test_a_turn_returns_the_role_the_text_and_the_transport(executor_pair): + executor, arrange, _ = executor_pair + arrange("smith", "landed foo") + outcome = executor.run(TurnRequest(role="smith", what="onboard foo", prompt="do it")) + + assert isinstance(outcome, TurnOutcome) + assert outcome.role == "smith" + assert outcome.text == "landed foo" + assert outcome.transport == executor.name + + +def test_the_decision_path_is_echoed_back_so_the_collector_can_find_it(executor_pair): + """The orchestrator branches on the decision file, never on the prose.""" + executor, arrange, _ = executor_pair + arrange("auditor", "everything looks fine to me") + outcome = executor.run( + TurnRequest( + role="auditor", + what="coverage", + prompt="check it", + decision_file="/w/decisions/auditor-1.yaml", + ) + ) + assert outcome.decision_ref == "/w/decisions/auditor-1.yaml" + + +def test_every_transport_accepts_bound_even_without_a_session(executor_pair): + """Callers must not have to ask which transport they hold. + + A concurrent batch owns its own sessions, so it calls ``bound`` + unconditionally. A transport with no session concept has to tolerate + that rather than make the caller branch — which is the branch this + interface exists to delete. + """ + executor, arrange, _ = executor_pair + arrange("smith", "ok") + bound = executor.bound(RecordingLayer("ok")) + assert bound.run(TurnRequest(role="smith", what="x", prompt="p")).text == "ok" + + +def test_the_prompt_reaches_the_other_side_unedited(executor_pair): + executor, arrange, prompt_seen = executor_pair + arrange("assembler", "done") + prompt = "line one\n\nline two with `backticks` and — an em dash" + executor.run(TurnRequest(role="assembler", what="core", prompt=prompt)) + assert prompt_seen() == prompt + + +# ------------------------------------------------------------- in-process only + + +def test_in_process_calls_the_layer_the_provider_returns(): + layer = RecordingLayer("from the layer") + seen: list[str] = [] + + def provider(request: TurnRequest): + seen.append(request.role) + return layer + + outcome = InProcessTurnExecutor(provider).run( + TurnRequest(role="reviewer", what="core", prompt="review it") + ) + assert seen == ["reviewer"] + assert outcome.text == "from the layer" + assert layer.prompts == ["review it"] + + +def test_bound_pins_the_session_the_caller_already_owns(): + """Ten concurrent dispatches are ten agents, not ten turns of one. + + The provider resolves a role's shared layer; a batch that built its own + must be able to say so, or every dispatch in it would land on the same + session. + """ + shared = RecordingLayer("shared") + mine = RecordingLayer("mine") + executor = InProcessTurnExecutor(lambda req: shared) + + assert executor.run(TurnRequest(role="smith", what="a", prompt="p")).text == "shared" + assert executor.bound(mine).run(TurnRequest(role="smith", what="b", prompt="p")).text == "mine" + # Binding returns a view; the original is untouched. + assert executor.run(TurnRequest(role="smith", what="c", prompt="p")).text == "shared" + + +def test_a_provider_with_no_layer_is_a_named_failure_not_an_attribute_error(): + """A role the workflow never configured must say so, not crash on None.""" + executor = InProcessTurnExecutor(lambda req: None) + with pytest.raises(TurnExecutionError, match="no layer available for role 'ghost'"): + executor.run(TurnRequest(role="ghost", what="x", prompt="p")) + + +def test_the_system_prompt_is_not_resent_in_process(): + """In process the session already carries it; resending would double it.""" + layer = RecordingLayer("ok") + InProcessTurnExecutor(lambda req: layer).run( + TurnRequest(role="smith", what="x", prompt="task only", system_prompt="ROLE PROMPT") + ) + assert layer.prompts == ["task only"] + + +def test_an_interrupt_during_a_turn_reaches_the_caller(): + """Ctrl-C must not become an empty answer. + + The orchestrator catches it around the whole run to checkpoint and say + how to resume. Swallowing it here would produce a turn that "ended" with + no decision, which every role's caller reads as a role that declined to + report. + """ + + def interrupted(prompt: str) -> str: + raise KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + InProcessTurnExecutor(lambda req: interrupted).run( + TurnRequest(role="smith", what="x", prompt="p") + ) + + +def test_an_exception_from_a_layer_is_not_converted_into_an_outcome(): + """A batch scores a raised dispatch differently from one that decided nothing.""" + boom = RuntimeError("backend returned 401") + + def failing(prompt: str) -> str: + raise boom + + with pytest.raises(RuntimeError, match="401"): + InProcessTurnExecutor(lambda req: failing).run( + TurnRequest(role="plan_reviewer", what="x", prompt="p") + ) + + +# ---------------------------------------------------------- file-dispatch only + + +def test_the_system_prompt_travels_with_a_dispatched_turn(tmp_path): + """A driver handed only the task prompt is running a different agent.""" + executor = FileDispatchTurnExecutor(tmp_path, poll_s=0.01) + _answer_after(executor, "001-smith", "ok") + executor.run(TurnRequest(role="smith", what="x", prompt="task", system_prompt="ROLE PROMPT")) + system = (executor.root / "001-smith.system.md").read_text(encoding="utf-8") + assert system == "ROLE PROMPT" + + +def test_a_transport_note_is_appended_never_substituted(tmp_path): + """The role prompt is right for the transport it was written for.""" + executor = FileDispatchTurnExecutor( + tmp_path, poll_s=0.01, transport_note=lambda role: f"\n[no tools for {role} here]\n" + ) + _answer_after(executor, "001-auditor", "ok") + executor.run(TurnRequest(role="auditor", what="x", prompt="p", system_prompt="ROLE")) + system = (executor.root / "001-auditor.system.md").read_text(encoding="utf-8") + assert system.startswith("ROLE") + assert "[no tools for auditor here]" in system + + +def test_the_outcome_names_its_request_and_its_artifacts(tmp_path): + executor = FileDispatchTurnExecutor(tmp_path, poll_s=0.01) + _answer_after(executor, "001-smith", "ok") + outcome = executor.run(TurnRequest(role="smith", what="x", prompt="p")) + + assert outcome.request_id == "001-smith" + assert [a.rsplit("/", 1)[-1] for a in outcome.artifacts] == [ + "001-smith.system.md", + "001-smith.prompt.md", + "001-smith.result.md", + ] + + +def test_context_refs_are_named_in_the_request_not_inlined(tmp_path): + """Inputs are pointers, not prose. + + A driver that wants them reads them; one that does not is not forced to + carry them in the prompt. + """ + import json + + executor = FileDispatchTurnExecutor(tmp_path, poll_s=0.01) + _answer_after(executor, "001-smith", "ok") + executor.run( + TurnRequest( + role="smith", what="x", prompt="p", context_refs=("/w/plan.md", "/w/vocab.yaml") + ) + ) + request = json.loads((executor.root / "001-smith.request.json").read_text(encoding="utf-8")) + assert request["context_refs"] == ["/w/plan.md", "/w/vocab.yaml"] + + +def test_a_turn_nobody_answers_times_out_with_the_prompt_path(tmp_path): + """A structured failure naming where to run the turn by hand.""" + executor = FileDispatchTurnExecutor(tmp_path, poll_s=0.01, timeout_s=0.05) + with pytest.raises(TurnTimeout, match="001-smith.prompt.md"): + executor.run(TurnRequest(role="smith", what="x", prompt="p")) + + +def test_a_stale_heartbeat_fails_before_the_full_timeout(tmp_path): + """The turn timeout bounds a slow turn; the heartbeat bounds a dead driver.""" + import os + + executor = FileDispatchTurnExecutor(tmp_path, poll_s=0.01, timeout_s=600) + beat = executor.root / "driver.heartbeat" + beat.write_text("x", encoding="utf-8") + old = time.time() - 10_000 + os.utime(beat, (old, old)) + + started = time.monotonic() + with pytest.raises(TurnTimeout, match="heartbeat is stale"): + executor.run(TurnRequest(role="smith", what="x", prompt="p")) + assert time.monotonic() - started < 10 + + +def test_bound_is_a_no_op_for_a_transport_with_no_session(tmp_path): + executor = FileDispatchTurnExecutor(tmp_path, poll_s=0.01) + assert executor.bound(object()) is executor + + +def test_ids_continue_across_a_resume_so_evidence_is_never_overwritten(tmp_path): + """The request files are the record of what the run asked for.""" + first = FileDispatchTurnExecutor(tmp_path, poll_s=0.01) + _answer_after(first, "001-smith", "ok") + first.run(TurnRequest(role="smith", what="x", prompt="p")) + + resumed = FileDispatchTurnExecutor(tmp_path, poll_s=0.01) + _answer_after(resumed, "002-smith", "ok") + outcome = resumed.run(TurnRequest(role="smith", what="y", prompt="p")) + assert outcome.request_id == "002-smith" + + +def test_a_resume_waits_on_an_unanswered_twin_instead_of_asking_again(tmp_path): + """A dead process's in-flight turn is the same question, possibly being answered. + + Measured on the DeepSeek-V4.1 run: the compute session was killed with 30 + smith turns in flight; re-asking each would duplicate every agent's work + and hand two agents the same decision path. + """ + first = FileDispatchTurnExecutor(tmp_path, poll_s=0.01, timeout_s=0.05) + with pytest.raises(TurnTimeout): + first.run( + TurnRequest(role="smith", what="onboard a", prompt="p", decision_file="d/smith-1.yaml") + ) + + resumed = FileDispatchTurnExecutor(tmp_path, poll_s=0.01) + _answer_after(resumed, "001-smith", "answered by the first driver") + outcome = resumed.run( + TurnRequest(role="smith", what="onboard a", prompt="p2", decision_file="d/smith-1.yaml") + ) + assert outcome.request_id == "001-smith" + assert outcome.text == "answered by the first driver" + assert not (resumed.root / "002-smith.request.json").exists() + + +def test_an_answered_or_different_request_is_never_adopted(tmp_path): + """Only an unanswered request for the same turn and decision path is a twin.""" + first = FileDispatchTurnExecutor(tmp_path, poll_s=0.01) + _answer_after(first, "001-smith", "done once") + first.run(TurnRequest(role="smith", what="onboard a", prompt="p", decision_file="d/1")) + + second = FileDispatchTurnExecutor(tmp_path, poll_s=0.01, timeout_s=0.05) + with pytest.raises(TurnTimeout): + second.run(TurnRequest(role="smith", what="onboard b", prompt="p", decision_file="d/2")) + + third = FileDispatchTurnExecutor(tmp_path, poll_s=0.01) + _answer_after(third, "003-smith", "fresh") + outcome = third.run( + TurnRequest(role="smith", what="onboard a", prompt="p", decision_file="d/1") + ) + assert outcome.request_id == "003-smith" diff --git a/agent-flow/tests/workflows/staircase_onboard/__init__.py b/agent-flow/tests/workflows/staircase_onboard/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/__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/staircase_onboard/conftest.py b/agent-flow/tests/workflows/staircase_onboard/conftest.py new file mode 100644 index 000000000000..9ab1f055e799 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/conftest.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fixtures shared across the staircase tests. + +`test_workflow.py` built the repo, checkpoint and task-spec fixtures first +and everything since has wanted them. Importing them test-module to +test-module works, but every consumer then has to name them as parameters +*and* import them, which reads to a linter as a redefinition and to a reader +as a dependency between test files that is not really there. + +Re-exported here instead, so pytest resolves them by name and no test module +imports another. +""" + +from __future__ import annotations + +from .test_workflow import checkpoint, repo, task_file # noqa: F401 + +__all__ = ["checkpoint", "repo", "task_file"] diff --git a/agent-flow/tests/workflows/staircase_onboard/test_agent_home.py b/agent-flow/tests/workflows/staircase_onboard/test_agent_home.py new file mode 100644 index 000000000000..23d008f61927 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_agent_home.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The staged agent home, which is what makes the orchestrator able to spawn agents. + +The failure this module exists for produced a turn with zero tokens and +``stop_reason='stop_sequence'``, i.e. something indistinguishable from a model +refusal, when in fact no agent had started. These tests are about the parts +that decide whether that happens, not about the copying. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + +from agent_flow.workflows.staircase_onboard.agent_home import ( + CLAUDE_CONFIG_DIRNAME, + CREDENTIAL_SUBPATH, + ensure_agent_home, + home_is_writable, + refresh_credentials, + stage_agent_home, +) + + +@pytest.fixture +def fake_home(tmp_path: Path) -> Path: + """A writable stand-in for ``$HOME`` with a credential store in it.""" + home = tmp_path / "realhome" + creds = home / CREDENTIAL_SUBPATH + creds.mkdir(parents=True) + (creds / "active_config").write_text("default\n", encoding="utf-8") + (creds / "credentials").mkdir() + (creds / "credentials" / "default.json").write_text('{"token": "one"}', encoding="utf-8") + return home + + +def test_a_writable_home_stages_nothing(tmp_path, fake_home): + """An ordinary workstation must be left completely alone.""" + assert home_is_writable(fake_home) + assert ensure_agent_home(tmp_path / "ws", home=fake_home) is None + + +def test_the_write_probe_removes_itself(fake_home): + """A probe that leaves a directory behind is a probe that pollutes the store.""" + assert home_is_writable(fake_home) + assert not (fake_home / CREDENTIAL_SUBPATH / ".agent_home_write_probe").exists() + + +def test_an_unwritable_home_is_detected_without_relying_on_permission_bits(tmp_path, fake_home): + """A sandbox is not a permission bit, so the probe must actually write. + + Simulated here by pointing at a path whose parent is a file: ``mkdir`` + fails with ``ENOTDIR`` exactly as it fails with ``EROFS``, and + ``os.access`` on the nonexistent path would tell us nothing either way. + """ + blocked = tmp_path / "blocked" + blocked.write_text("not a directory", encoding="utf-8") + assert not home_is_writable(blocked) + + +def test_staging_selects_the_new_home_through_all_three_variables(tmp_path, fake_home): + """HOME alone is not enough: the credential store resolves through XDG. + + Setting only ``HOME`` leaves the child reading the real, read-only store + and the failure is byte-identical to having done nothing at all. + """ + staged = stage_agent_home(tmp_path / "agent-home", home=fake_home) + assert staged.env["HOME"] == str(staged.root) + assert staged.env["XDG_CONFIG_HOME"] == str(staged.root / ".config") + assert staged.env["CLAUDE_CONFIG_DIR"] == str(staged.root / CLAUDE_CONFIG_DIRNAME) + + +def test_the_claude_config_dir_is_not_named_dot_claude(tmp_path, fake_home): + """A sandbox that denies `**/.claude` writes would silently cripple the child. + + The symptom is a warning, not an error, so a run can lose its config + directory and still appear to work. + """ + staged = stage_agent_home(tmp_path / "agent-home", home=fake_home) + assert ".claude" not in Path(staged.env["CLAUDE_CONFIG_DIR"]).name + assert (staged.root / CLAUDE_CONFIG_DIRNAME / ".claude.json").is_file() + + +def test_credentials_are_copied_not_referenced(tmp_path, fake_home): + staged = stage_agent_home(tmp_path / "agent-home", home=fake_home) + copied = staged.root / CREDENTIAL_SUBPATH / "credentials" / "default.json" + assert copied.read_text(encoding="utf-8") == '{"token": "one"}' + assert not copied.is_symlink() + + +def test_refresh_replaces_a_stale_copy(tmp_path, fake_home): + """A stale copy fails as an auth error minutes into a turn, which is expensive. + + Replacement rather than merge, because a rotated store can have FEWER + files than the one it replaces and a merge would leave the old ones + readable. + """ + staged = stage_agent_home(tmp_path / "agent-home", home=fake_home) + (fake_home / CREDENTIAL_SUBPATH / "credentials" / "default.json").write_text( + '{"token": "two"}', encoding="utf-8" + ) + (staged.root / CREDENTIAL_SUBPATH / "credentials" / "stale.json").write_text("x", "utf-8") + + assert refresh_credentials(staged.root, home=fake_home) + creds = staged.root / CREDENTIAL_SUBPATH / "credentials" + assert (creds / "default.json").read_text(encoding="utf-8") == '{"token": "two"}' + assert not (creds / "stale.json").exists() + + +def test_refresh_reports_failure_rather_than_raising(tmp_path): + """There is nothing to copy from; the caller decides what that means.""" + assert refresh_credentials(tmp_path / "agent-home", home=tmp_path / "nowhere") is False + + +def test_git_identity_is_written_because_redirecting_home_removes_the_real_one(tmp_path, fake_home): + """Otherwise every smith commit dies with "Please tell me who you are". + + That is a new infrastructure failure traded for the one just fixed, and + it surfaces at the first commit rather than at startup. + """ + staged = stage_agent_home( + tmp_path / "agent-home", home=fake_home, git_identity=("A Name", "a@example.com") + ) + text = (staged.root / ".gitconfig").read_text(encoding="utf-8") + assert "name = A Name" in text + assert "email = a@example.com" in text + + +def test_apply_exports_into_the_process_so_every_backend_inherits_it( + tmp_path, fake_home, monkeypatch +): + """The alternative is per-backend plumbing, and Codex would need its own. + + Claude Code takes ``ClaudeAgentOptions.env``; Codex does not expose one. + Exporting once is what makes this backend-agnostic. + """ + monkeypatch.setenv("HOME", str(fake_home)) + staged = stage_agent_home(tmp_path / "agent-home", home=fake_home) + staged.apply() + assert os.environ["HOME"] == str(staged.root) + assert os.environ["XDG_CONFIG_HOME"] == str(staged.root / ".config") + + +def test_staging_is_idempotent(tmp_path, fake_home): + """A resume re-enters the same code path and must not destroy state.""" + root = tmp_path / "agent-home" + stage_agent_home(root, home=fake_home) + (root / CLAUDE_CONFIG_DIRNAME / ".claude.json").write_text('{"k": 1}', encoding="utf-8") + stage_agent_home(root, home=fake_home) + assert (root / CLAUDE_CONFIG_DIRNAME / ".claude.json").read_text(encoding="utf-8") == '{"k": 1}' diff --git a/agent-flow/tests/workflows/staircase_onboard/test_build_and_test_doc.py b/agent-flow/tests/workflows/staircase_onboard/test_build_and_test_doc.py new file mode 100644 index 000000000000..16309b719c1b --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_build_and_test_doc.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""`BUILD_AND_TEST.md`, held to the tree it describes. + +The document's whole value is that it points at the places in TensorRT-LLM +that already do each job correctly, rather than restating them. That only +works while the pointers resolve — a reference doc whose paths have moved +costs the reader the trip and then their trust in the rest of the page. + +So: every repo path it names must exist, every `build_wheel.py` flag it +documents must still be a flag, and every test list it names must still be a +list. What it *says* about them is prose and is not checked here. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + +import agent_flow.workflows.staircase_onboard as pkg + +PACKAGE = Path(pkg.__file__).parent +DOC = PACKAGE / "BUILD_AND_TEST.md" +# agent-flow/tests/workflows/staircase_onboard -> the TensorRT-LLM checkout. +REPO_ROOT = Path(__file__).resolve().parents[4] + +pytestmark = pytest.mark.skipif( + not (REPO_ROOT / "scripts" / "build_wheel.py").is_file(), + reason="not inside a TensorRT-LLM checkout; nothing to resolve against", +) + + +def _text() -> str: + return DOC.read_text(encoding="utf-8") + + +def test_the_doc_ships_and_is_linked(): + assert DOC.is_file() + assert "BUILD_AND_TEST.md" in (PACKAGE / "README.md").read_text(encoding="utf-8") + + +def test_it_is_declared_as_package_data(): + """A doc absent from a non-editable install is a doc that does not exist.""" + import tomllib + + pyproject = Path(__file__).resolve().parents[3] / "pyproject.toml" + data = tomllib.loads(pyproject.read_text(encoding="utf-8")) + declared = data["tool"]["setuptools"]["package-data"]["agent_flow.workflows.staircase_onboard"] + assert "BUILD_AND_TEST.md" in declared + + +def test_every_repo_path_it_names_exists(): + missing = [] + for path in sorted( + set( + re.findall( + r"`((?:scripts|tests|jenkins|cpp|tensorrt_llm|\.claude)/[A-Za-z0-9_./*-]+)`", + _text(), + ) + ) + ): + if "*" in path: # a shape, not a file + continue + if not (REPO_ROOT / path.rstrip("/")).exists(): + missing.append(path) + assert not missing, "BUILD_AND_TEST.md names paths that do not exist:\n " + "\n ".join( + missing + ) + + +def _section(title: str) -> str: + """The text of one `## ` section, up to the next one.""" + text = _text() + start = text.index(f"## {title}") + rest = text[start + 1 :] + end = rest.find("\n## ") + return rest if end == -1 else rest[:end] + + +def test_every_build_flag_it_documents_is_still_a_flag(): + """The flag table is the part a reader copies verbatim. + + Scoped to the Build section: the multi-node section quotes `srun` flags + in the same backticked shape, and those belong to Slurm. + """ + source = (REPO_ROOT / "scripts" / "build_wheel.py").read_text(encoding="utf-8") + documented = set(re.findall(r"`(--[a-z_]+)`", _section("1. Build"))) + # `-j`, `-c`, `-f` are short forms; check the long ones, which are the + # names `build_wheel.py` actually declares. + unknown = sorted(f for f in documented if f'"{f}"' not in source) + assert documented, "the Build section documents no flags; did it change shape?" + assert not unknown, f"documented flags build_wheel.py does not declare: {unknown}" + + +def test_every_test_list_it_names_is_still_a_list(): + lists = REPO_ROOT / "tests/integration/test_lists/test-db" + named = set(re.findall(r"`(l0_[a-z0-9_]+)\.yml`", _text())) + assert named, "the scale ladder names no lists; did the table change shape?" + missing = sorted(n for n in named if not (lists / f"{n}.yml").is_file()) + assert not missing, f"named test lists that no longer exist: {missing}" + + +def test_the_multi_node_launch_rules_agree_with_the_code_that_emits_them(): + """The doc and `environment.py` must not drift into two different recipes.""" + from agent_flow.workflows.staircase_onboard.environment import ( + LLMAPI_LAUNCH, + MPI_PMIX, + RunEnvironment, + ) + + text = _text() + assert LLMAPI_LAUNCH in text + assert f"`{MPI_PMIX}`" in text + + emitted = RunEnvironment( + kind="slurm", python="/p", world_size=8, gpus_per_node=4, container_image="/i.sqsh" + ).command("python x.py") + for flag in ("--ntasks=8", "--ntasks-per-node=4", "--gres=gpu:4", f"--mpi={MPI_PMIX}"): + assert flag.split("=")[0] in text, f"{flag} is emitted but undocumented" + assert flag in emitted, f"{flag} is documented but not emitted" + + +def test_the_rebuild_table_still_reflects_a_python_only_target(): + """If a target ever needs a native rebuild, this page is the first lie.""" + assert "almost certainly not" in _text() + assert "Triton is JIT" in _text() diff --git a/agent-flow/tests/workflows/staircase_onboard/test_committer.py b/agent-flow/tests/workflows/staircase_onboard/test_committer.py new file mode 100644 index 000000000000..dd3a82731c8b --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_committer.py @@ -0,0 +1,331 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The commit helper, exercised against a real git repo. + +These tests shell out to real ``git`` on purpose. The thing under test is a +shell script whose whole job is to behave correctly under concurrency, and a +mock of git would test the mock. +""" + +from __future__ import annotations + +import subprocess +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from agent_flow.workflows.staircase_onboard.committer import ( + SCRIPT_NAME, + commit_instruction, + install_committer, + lock_is_held, + render_script, +) + + +def _git(repo, *args, check=True): + return subprocess.run( + ["git", "-C", str(repo), *args], capture_output=True, text=True, check=check + ) + + +@pytest.fixture +def repo(tmp_path): + path = tmp_path / "repo" + path.mkdir() + _git(path, "init", "-q") + _git(path, "config", "user.email", "t@example.invalid") + _git(path, "config", "user.name", "Test") + _git(path, "config", "commit.gpgsign", "false") + (path / "seed.txt").write_text("seed\n") + _git(path, "add", "seed.txt") + _git(path, "commit", "-q", "-m", "seed") + return path + + +@pytest.fixture +def commit(tmp_path, repo): + script = install_committer(tmp_path / "ws", repo) + + def run(message, *paths, timeout=60): + return subprocess.run( + [str(script), message, *paths], capture_output=True, text=True, timeout=timeout + ) + + run.script = script + return run + + +def _files_in(repo, ref="HEAD"): + out = _git(repo, "diff-tree", "--no-commit-id", "--name-only", "-r", ref).stdout + return sorted(line for line in out.splitlines() if line) + + +def test_it_commits_exactly_what_it_was_given(repo, commit): + (repo / "a.py").write_text("a\n") + (repo / "b.py").write_text("b\n") + result = commit("add a", "a.py") + assert result.returncode == 0, result.stderr + assert _files_in(repo) == ["a.py"] + assert (repo / "b.py").exists() # untouched, still uncommitted + + +def test_an_untracked_file_from_another_agent_is_not_swept_in(repo, commit): + """The measured failure: another agent's in-progress file in my commit.""" + (repo / "mine.py").write_text("mine\n") + (repo / "theirs_half_written.py").write_text("half\n") + assert commit("mine", "mine.py").returncode == 0 + assert _files_in(repo) == ["mine.py"] + + +def test_another_agents_already_staged_file_is_not_swept_in(repo, commit): + """The worse variant: they had already run `git add`.""" + (repo / "mine.py").write_text("mine\n") + (repo / "theirs.py").write_text("theirs\n") + _git(repo, "add", "theirs.py") # the other agent staged, has not committed + assert commit("mine", "mine.py").returncode == 0 + assert _files_in(repo) == ["mine.py"] + # And theirs is still there for them to commit. + assert "theirs.py" in _git(repo, "diff", "--cached", "--name-only").stdout + + +def test_concurrent_commits_all_land_and_none_contaminate(repo, commit): + """Eight agents, one worktree, one file each. + + Without the lock this is where ``Unable to create '.git/index.lock'`` + appears, and where one agent's files land in another's commit. + """ + names = [f"entry_{i}.py" for i in range(8)] + for name in names: + (repo / name).write_text(f"# {name}\n") + + with ThreadPoolExecutor(max_workers=8) as pool: + results = list(pool.map(lambda n: commit(f"land {n}", n), names)) + + assert all(r.returncode == 0 for r in results), [r.stderr for r in results if r.returncode] + log = _git(repo, "log", "--format=%H").stdout.split() + assert len(log) == 9 # seed + 8 + landed = {sha: _files_in(repo, sha) for sha in log[:8]} + assert sorted(f for files in landed.values() for f in files) == sorted(names) + for files in landed.values(): + assert len(files) == 1 # one file per commit, no contamination + + +def test_a_commit_that_swept_more_than_asked_is_undone(repo, commit): + """The verification, exercised by pointing at a directory. + + A directory pathspec is the exact mistake the prose has always warned + about. Here it stops being advice: the script refuses the result. + """ + sub = repo / "pkg" + sub.mkdir() + (sub / "mine.py").write_text("mine\n") + (sub / "theirs.py").write_text("theirs\n") + before = _git(repo, "rev-parse", "HEAD").stdout.strip() + + result = commit("sloppy", "pkg") + assert result.returncode == 1 + assert "REFUSING" in result.stderr + assert "pkg/theirs.py" in result.stderr + # HEAD is back where it was, and the work is still in the tree. + assert _git(repo, "rev-parse", "HEAD").stdout.strip() == before + assert (sub / "mine.py").exists() + + +def test_a_path_that_does_not_exist_is_rejected_before_the_lock(repo, commit): + result = commit("typo", "does_not_exist.py") + assert result.returncode == 2 + assert "does not exist" in result.stderr + + +def test_it_refuses_without_a_message(repo, commit): + assert commit("only-a-message").returncode == 2 + + +def test_nothing_to_commit_fails_loudly_rather_than_silently(repo, commit): + """An agent that thinks it committed and did not is the worst outcome.""" + result = commit("no change", "seed.txt") + assert result.returncode == 1 + assert "nothing to commit" in result.stderr + + +def test_the_sha_and_the_file_list_are_printed(repo, commit): + """The agent's report has to be checkable against the repo.""" + (repo / "a.py").write_text("a\n") + result = commit("add a", "a.py") + sha = _git(repo, "rev-parse", "HEAD").stdout.strip() + assert sha in result.stdout + assert "a.py" in result.stdout + + +def test_waiting_for_the_lock_is_a_distinct_exit_code(tmp_path, repo): + """So a queue is never mistaken for a broken repo.""" + script = install_committer(tmp_path / "ws", repo, timeout=1) + (repo / "a.py").write_text("a\n") + lock = tmp_path / "ws" / ".commit.lock" + holder = subprocess.Popen(["flock", str(lock), "sleep", "10"]) + try: + result = subprocess.run( + [str(script), "blocked", "a.py"], capture_output=True, text=True, timeout=30 + ) + assert result.returncode == 75 + assert "did not get it" in result.stderr + assert "symptom to report" in result.stderr + finally: + holder.kill() + holder.wait() + + +def test_the_lock_can_be_observed_by_the_orchestrator(tmp_path, repo): + """Tells "still committing" from "wedged" when a dispatch overruns.""" + install_committer(tmp_path / "ws", repo) + assert not lock_is_held(tmp_path / "ws") + lock = tmp_path / "ws" / ".commit.lock" + holder = subprocess.Popen(["flock", str(lock), "sleep", "10"]) + try: + import time + + for _ in range(50): + if lock_is_held(tmp_path / "ws"): + break + time.sleep(0.1) + assert lock_is_held(tmp_path / "ws") + finally: + holder.kill() + holder.wait() + + +# ------------------------------------------------------------- installation + + +def test_the_script_is_executable_and_self_contained(tmp_path, repo): + script = install_committer(tmp_path / "ws", repo) + assert script.name == SCRIPT_NAME + assert script.stat().st_mode & 0o111 + body = script.read_text() + assert str(repo) in body # baked, so no argument an agent can get wrong + assert "flock" in body + + +def test_paths_with_spaces_survive_the_bake(tmp_path): + body = render_script("/a path/repo", "/a path/ws/.commit.lock") + assert "'/a path/repo'" in body + + +def test_the_instruction_names_the_script_and_the_refusal(tmp_path, repo): + script = install_committer(tmp_path / "ws", repo) + text = commit_instruction(script) + assert str(script) in text + assert "undoes the commit" in text + assert "git add -A" in text # still named, still forbidden + + +# --------------------------------------------------- wired into the prompts + + +def test_the_committing_roles_are_pointed_at_the_script(tmp_path, repo): + from agent_flow.workflows.staircase_onboard.prompts import build_staircase_onboard_prompts + + script = install_committer(tmp_path / "ws", repo) + prompts = build_staircase_onboard_prompts(commit_script=str(script)) + assert str(script) in prompts.smith + assert str(script) in prompts.assembler + + +def test_a_role_that_does_not_commit_is_not_handed_a_commit_recipe(tmp_path, repo): + """It will otherwise find a reason to use it.""" + from agent_flow.workflows.staircase_onboard.prompts import build_staircase_onboard_prompts + + script = install_committer(tmp_path / "ws", repo) + prompts = build_staircase_onboard_prompts(commit_script=str(script)) + for role in (prompts.plan_drafter, prompts.plan_reviewer, prompts.reviewer, prompts.auditor): + assert str(script) not in role + + +def test_the_new_recipe_explicitly_supersedes_the_old_one(tmp_path, repo): + """Two commit recipes in one prompt is how the wrong one gets picked.""" + from agent_flow.workflows.staircase_onboard.prompts import build_staircase_onboard_prompts + + script = install_committer(tmp_path / "ws", repo) + smith = build_staircase_onboard_prompts(commit_script=str(script)).smith + # The baseline block is still there — it is embedded in the role prompt — + # so the replacement has to say so, and has to come after it. + assert "Commit **by explicit file path**" in smith + assert "this replaces the Committing section above" in smith + assert smith.index("Commit **by explicit file path**") < smith.index("this replaces the") + + +def test_the_undo_rewinds_exactly_one_commit(repo, commit): + """Not two. + + The refusal message quoted the reset command inside a double-quoted + `echo`, where a backtick is a command substitution — so writing the + example ran it, and the explicit reset below then ran it again. In a + repo with one commit that fails harmlessly; in a repo with history it + silently throws away the previous commit. + """ + for i in range(3): + (repo / f"history_{i}.txt").write_text(f"{i}\n") + _git(repo, "add", f"history_{i}.txt") + _git(repo, "commit", "-q", "-m", f"history {i}") + log_before = _git(repo, "log", "--format=%H").stdout.split() + + sub = repo / "pkg" + sub.mkdir() + (sub / "mine.py").write_text("mine\n") + (sub / "theirs.py").write_text("theirs\n") + assert commit("sloppy", "pkg").returncode == 1 + + assert _git(repo, "log", "--format=%H").stdout.split() == log_before + + +# ------------------------------------------- the linter cannot see this code + + +def test_the_generated_script_is_valid_shell(tmp_path, repo): + """`bash -n` parses without running. + + Neither `ruff` nor `shellcheck` looks inside a Python string literal, so + this script gets no static checking at all unless a test asks for it. + """ + script = install_committer(tmp_path / "ws", repo) + result = subprocess.run(["bash", "-n", str(script)], capture_output=True, text=True) + assert result.returncode == 0, result.stderr + + +def test_the_script_has_no_command_substitution_in_its_messages(tmp_path, repo): + """The bug that shipped: a backtick inside a double-quoted `echo`. + + It quoted `git reset --soft HEAD^` as an example and bash *ran* it, so + the explicit reset below rewound a second, unrelated commit. Static + rule, because the behavioural test only covers the one message that + happened to have it. + """ + script = install_committer(tmp_path / "ws", repo) + offenders = [ + (n, line) + for n, line in enumerate(script.read_text().splitlines(), 1) + if line.strip().startswith("echo ") and '"' in line and "`" in line + ] + assert not offenders, f"backtick inside a double-quoted echo: {offenders}" + + +def test_every_commit_carries_a_dco_sign_off_under_gits_own_identity(repo, commit, monkeypatch): + """The target repository rejects a commit without `Signed-off-by`. + + A login name is not a sign-off identity either: the name git itself would + use comes before $USER. + """ + monkeypatch.delenv("STAIRCASE_GIT_NAME", raising=False) + monkeypatch.delenv("STAIRCASE_GIT_EMAIL", raising=False) + (repo / "signed.txt").write_text("x\n") + result = commit("[None][test] signed", "signed.txt") + assert result.returncode == 0, result.stderr + assert ( + _git(repo, "log", "-1", "--format=%an <%ae>").stdout.strip() == "Test " + ) + trailers = _git( + repo, "log", "-1", "--format=%(trailers:key=Signed-off-by,valueonly)" + ).stdout.strip() + assert trailers == "Test " diff --git a/agent-flow/tests/workflows/staircase_onboard/test_decision_files.py b/agent-flow/tests/workflows/staircase_onboard/test_decision_files.py new file mode 100644 index 000000000000..aeaaabe6a72d --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_decision_files.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the file-based decision protocol. + +This mode exists because an enterprise MCP configuration +(``/etc/claude-code/managed-mcp.json``) makes the Agent SDK refuse every +dynamically-registered server — including the in-process one the six typed +decision tools live in. On such a host those tools cannot be created at all, +and the run died on its first dispatch with an opaque +``ProcessError: exit code 1``. + +The guarantee the tools provided was a **closed value**: the orchestrator +branches on ``decision`` and never on prose. These tests are about where +that guarantee went — it moved from the tool schema to validation on read, +and it has to be just as strict there or the mode is a downgrade pretending +not to be one. +""" + +from __future__ import annotations + +import inspect +import os +from pathlib import Path + +import yaml + +from agent_flow.workflows.staircase_onboard.progress import ( + DECISIONS_BY_AGENT, + EXTRA_KEYS_BY_AGENT, + ProgressContext, + collect_decision_file, + decision_is_stale, + decision_path, + init_progress_file, + latest_decision, + latest_entry, + mark_decision_dispatch, +) +from agent_flow.workflows.staircase_onboard.prompts import build_staircase_onboard_prompts +from agent_flow.workflows.staircase_onboard.prompts._common import decision_file_block + + +def _context(tmp_path: Path) -> ProgressContext: + progress = tmp_path / "progress.yaml" + init_progress_file(progress) + return ProgressContext(path=progress) + + +def _write(tmp_path: Path, agent: str, iteration: int, payload: dict) -> None: + path = decision_path(tmp_path, agent, iteration) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(yaml.safe_dump(payload), encoding="utf-8") + + +def test_a_valid_decision_is_recorded_and_readable_downstream(tmp_path): + """The whole point: existing `latest_decision` callers keep working.""" + ctx = _context(tmp_path) + _write(tmp_path, "plan_drafter", 1, {"summary": "wrote both files", "decision": "DRAFT_READY"}) + + assert collect_decision_file(ctx, tmp_path, "plan_drafter", 1) == "DRAFT_READY" + assert ( + latest_decision(ctx.path, "plan_drafter", DECISIONS_BY_AGENT["plan_drafter"]) + == "DRAFT_READY" + ) + assert latest_entry(ctx.path, "plan_drafter")["summary"] == "wrote both files" + + +def test_a_decision_outside_the_enum_is_refused(tmp_path): + """The tool schema would have refused it; so must the reader. + + Accepting it would turn the orchestrator's one closed input into + whatever a role felt like writing — which is the failure the typed + tools existed to prevent. + """ + ctx = _context(tmp_path) + _write(tmp_path, "plan_drafter", 1, {"summary": "s", "decision": "LOOKS_GOOD_TO_ME"}) + assert collect_decision_file(ctx, tmp_path, "plan_drafter", 1) is None + assert latest_entry(ctx.path, "plan_drafter") is None + + +def test_another_roles_decision_is_not_accepted(tmp_path): + """`APPROVE` is real — for the reviewer. It is not a drafter decision.""" + ctx = _context(tmp_path) + _write(tmp_path, "plan_drafter", 1, {"summary": "s", "decision": "APPROVE"}) + assert collect_decision_file(ctx, tmp_path, "plan_drafter", 1) is None + + +def test_a_missing_file_is_none_rather_than_an_exception(tmp_path): + """A role that ended without deciding is a case the caller handles.""" + ctx = _context(tmp_path) + assert collect_decision_file(ctx, tmp_path, "smith", 3) is None + + +def test_malformed_yaml_is_refused_without_raising(tmp_path): + ctx = _context(tmp_path) + path = decision_path(tmp_path, "smith", 1) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("decision: [unclosed\n", encoding="utf-8") + assert collect_decision_file(ctx, tmp_path, "smith", 1) is None + + +def test_role_specific_extra_keys_survive(tmp_path): + """The assembler's `blocked_reason` is control flow, not commentary. + + `_handle_blocked_assembler` branches on it — a BLOCKED entry that lost + it is a run that cannot triage itself. + """ + ctx = _context(tmp_path) + _write( + tmp_path, + "assembler", + 1, + { + "summary": "no entry for the gated delta rule", + "decision": "BLOCKED", + "blocked_reason": "vocabulary_gap", + "missing_entry": "gdn_custom_op_inplace", + }, + ) + assert collect_decision_file(ctx, tmp_path, "assembler", 1) == "BLOCKED" + entry = latest_entry(ctx.path, "assembler") + assert entry["blocked_reason"] == "vocabulary_gap" + assert entry["missing_entry"] == "gdn_custom_op_inplace" + + +def test_empty_extra_keys_are_dropped_rather_than_stored_blank(tmp_path): + ctx = _context(tmp_path) + _write(tmp_path, "smith", 1, {"summary": "s", "decision": "LANDED", "entry": " "}) + collect_decision_file(ctx, tmp_path, "smith", 1) + assert "entry" not in latest_entry(ctx.path, "smith") + + +def test_the_iteration_is_the_orchestrators_not_the_agents(tmp_path): + """A role could write any iteration into the file; it is not asked to. + + Stamping it here keeps a mistyped or hallucinated number from + reordering the record. + """ + ctx = _context(tmp_path) + _write(tmp_path, "smith", 7, {"summary": "s", "decision": "LANDED", "iteration": 99}) + collect_decision_file(ctx, tmp_path, "smith", 7) + assert latest_entry(ctx.path, "smith")["iteration"] == 7 + + +# ---------------------------------------------- a stale decision is not a decision + + +def _write_decision(workspace: Path, agent: str, iteration: int, decision: str) -> Path: + target = decision_path(workspace, agent, iteration) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(f"summary: s\ndecision: {decision}\n", encoding="utf-8") + return target + + +def test_a_decision_written_after_the_marker_is_collected(tmp_path): + """The ordinary case: a turn answered its own dispatch.""" + ctx = _context(tmp_path) + mark_decision_dispatch(tmp_path, "auditor", 1) + _write_decision(tmp_path, "auditor", 1, "CLEAN") + assert collect_decision_file(ctx, tmp_path, "auditor", 1) == "CLEAN" + + +def test_a_decision_older_than_the_marker_is_refused(tmp_path): + """The failure this exists for, and it happened five times in one run. + + `(agent, iteration)` is not unique across a run — `iteration` is 1 for + every deliver-stage turn and restarts per unit in assembly. A turn that + ends without writing leaves the PREVIOUS turn's verdict at the same path, + looking fresh. Three audit verdicts were read back that way, including a + `CLEAN` that could therefore never end the run. + """ + ctx = _context(tmp_path) + stale = _write_decision(tmp_path, "auditor", 1, "UNCOVERED") + os.utime(stale, (1_000_000, 1_000_000)) # written long before this dispatch + mark_decision_dispatch(tmp_path, "auditor", 1) + + assert decision_is_stale(tmp_path, "auditor", 1) + assert collect_decision_file(ctx, tmp_path, "auditor", 1) is None + # and nothing was appended, so `latest_decision` cannot read it either + assert latest_decision(ctx.path, "auditor", DECISIONS_BY_AGENT["auditor"]) is None + + +def test_without_a_marker_the_old_behaviour_stands(tmp_path): + """A run that never stamped one must not have its decisions refused. + + Absence means "this run does not use the marker", not "stale" — the same + three-state rule the rest of the workflow applies to receipts. + """ + ctx = _context(tmp_path) + _write_decision(tmp_path, "reviewer", 2, "APPROVE") + assert not decision_is_stale(tmp_path, "reviewer", 2) + assert collect_decision_file(ctx, tmp_path, "reviewer", 2) == "APPROVE" + + +def test_a_missing_decision_is_not_stale(tmp_path): + """A stamped turn that was never answered is the caller's "no decision".""" + mark_decision_dispatch(tmp_path, "assembler", 3) + assert not decision_is_stale(tmp_path, "assembler", 3) + + +def test_the_marker_does_not_collide_with_the_decision_it_guards(tmp_path): + """A dotfile beside it, so a marker is never mistaken for an answer.""" + mark_decision_dispatch(tmp_path, "smith", 4) + written = sorted(p.name for p in decision_path(tmp_path, "smith", 4).parent.iterdir()) + assert written == [".smith-4.dispatched"] + + +def test_staleness_is_judged_by_mtime_not_by_the_clock(): + """Both files live under the workspace, so one filesystem orders them. + + An agent may answer from a node whose clock is skewed by minutes; a + wall-clock comparison would reject its perfectly good decision. + """ + src = inspect.getsource(decision_is_stale) + assert "st_mtime" in src + assert "time.time" not in src and "datetime.now" not in src + + +# ------------------------------------------------------------------ prompts + + +def test_each_role_is_shown_only_its_own_options(): + """A role shown another role's enum will eventually pick from it.""" + block = decision_file_block("assembler", DECISIONS_BY_AGENT["assembler"]) + assert "UNIT_READY" in block and "BLOCKED" in block + for foreign in ("DRAFT_READY", "CEILING", "LANDED", "FAILED_RECEIPT"): + assert foreign not in block, foreign + + +def test_every_role_gets_a_decision_block_when_the_mode_is_on(): + bundle = build_staircase_onboard_prompts(decision_files=True) + for role, decisions in DECISIONS_BY_AGENT.items(): + prompt = getattr(bundle, role) + assert "Ending your turn" in prompt, role + assert f"decisions/{role}-.yaml" in prompt, role + for value in decisions: + assert f"`{value}`" in prompt, (role, value) + + +def test_roles_with_extra_keys_are_told_about_them(): + bundle = build_staircase_onboard_prompts(decision_files=True) + for role, keys in EXTRA_KEYS_BY_AGENT.items(): + for key in keys: + assert key in getattr(bundle, role), (role, key) + + +def test_no_decision_block_when_the_tools_are_available(): + """The tools are the preferred mechanism; this mode is the fallback.""" + bundle = build_staircase_onboard_prompts(decision_files=False) + assert "decisions/plan_drafter-.yaml" not in bundle.plan_drafter + + +# ------------------------------------------------- the workspace spec copy + + +def test_a_resumed_run_refreshes_the_spec_the_agents_read(tmp_path): + """Two copies of task.yaml must not drift apart. + + `_init_state` re-reads and resolves the operator's file every run, but + the workspace copy — which every role is told to read as "the run's + spec" — was written once and never again. A resumed run therefore gave + the gates one spec and the assembler another: both self-consistent, + neither correct, and nothing would say so. + """ + import yaml as _yaml + + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + workspace = tmp_path / "ws" + workspace.mkdir() + stale = {"features": [], "environment": {"env": {}}} + (workspace / "task.yaml").write_text(_yaml.safe_dump(stale), encoding="utf-8") + + source = tmp_path / "task.yaml" + source.write_text("", encoding="utf-8") + + workflow = StaircaseOnboardWorkflow(workspace) + workflow.task_data = {"features": ["mtp"], "environment": {"env": {"PATH": "/venv/bin"}}} + workflow._refresh_workspace_spec(source, None) + + written = _yaml.safe_load((workspace / "task.yaml").read_text(encoding="utf-8")) + assert written["features"] == ["mtp"] + assert written["environment"]["env"]["PATH"] == "/venv/bin" + + +def test_refreshing_is_a_no_op_when_the_spec_already_matches(tmp_path): + """Rewriting an identical file would churn its mtime for no reason.""" + import yaml as _yaml + + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + workspace = tmp_path / "ws" + workspace.mkdir() + data = {"features": ["mtp"]} + (workspace / "task.yaml").write_text( + _yaml.safe_dump(data, sort_keys=False, allow_unicode=True), encoding="utf-8" + ) + before = (workspace / "task.yaml").stat().st_mtime_ns + + source = tmp_path / "task.yaml" + source.write_text("", encoding="utf-8") + + workflow = StaircaseOnboardWorkflow(workspace) + workflow.task_data = data + workflow._refresh_workspace_spec(source, None) + + assert (workspace / "task.yaml").stat().st_mtime_ns == before diff --git a/agent-flow/tests/workflows/staircase_onboard/test_delivery.py b/agent-flow/tests/workflows/staircase_onboard/test_delivery.py new file mode 100644 index 000000000000..9be242ff4e0a --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_delivery.py @@ -0,0 +1,298 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the mechanical delivery checks. + +Two kinds here. Most build a synthetic repo tree, so a check's failure mode +is exercised directly. The last one runs the whole verification against the +**real merged target in this checkout**, which is the only way to find out +that the path conventions are still the ones the repository actually uses. +""" + +from __future__ import annotations + +import inspect +import subprocess +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import delivery + +TARGET_REL = f"{delivery.MODELING_V2_PACKAGE}/models/gpt_oss/gpt_oss_120b__sm_103__tp1" + +_MODELING = '''\ +"""A synthetic target, just enough for the static checks.""" + +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.gemm.cublas_mm import cublas_mm +from tensorrt_llm._torch._experimental.modeling_v2.catalog.torch.embedding import embedding + +REQUIRED_TRTLLM_OPS = ( + "flashinfer_rmsnorm", + "cublas_mm", +) +''' + + +def _build_repo(tmp_path: Path, *, modeling: str = _MODELING) -> Path: + root = tmp_path / "repo" + target = root / TARGET_REL + target.mkdir(parents=True) + (target / "modeling.py").write_text(modeling, encoding="utf-8") + (target / "weights.py").write_text("MANIFEST = ()\n", encoding="utf-8") + + catalog = root / delivery.CATALOG_DIR + for category, name, contracted in ( + ("norm", "flashinfer_rmsnorm", True), + ("gemm", "cublas_mm", True), + ("torch", "embedding", False), + ): + (catalog / category).mkdir(parents=True, exist_ok=True) + (catalog / category / f"{name}.py").write_text("def f(): ...\n", encoding="utf-8") + if contracted: + (catalog / category / f"{name}.md").write_text("# contract\n", encoding="utf-8") + tests = root / delivery.CATALOG_TESTS_DIR / category + tests.mkdir(parents=True, exist_ok=True) + if contracted: + (tests / f"test_modeling_v2_{name}.py").write_text( + "def test(): ...\n", encoding="utf-8" + ) + (catalog / "index.yaml").write_text( + yaml.safe_dump( + { + "entries": [ + {"path": "norm/flashinfer_rmsnorm.py"}, + {"path": "gemm/cublas_mm.py"}, + {"path": "torch/embedding.py"}, + ] + } + ), + encoding="utf-8", + ) + + (root / delivery.CATALOG_TESTS_DIR / "test_modeling_v2_claims.py").write_text( + "def test(): ...\n", encoding="utf-8" + ) + routing = root / delivery.MODELING_V2_PACKAGE / "models" / "gpt_oss" / "routing.py" + routing.write_text( + 'TARGET_MODULES = {"X": "models.gpt_oss.gpt_oss_120b__sm_103__tp1.modeling"}\n', + encoding="utf-8", + ) + + refs = root / delivery.ACCURACY_REFERENCES_DIR + refs.mkdir(parents=True) + (refs / "gsm8k.yaml").write_text( + yaml.safe_dump({"openai/gpt-oss-120b": [{"accuracy": 85.6}]}), encoding="utf-8" + ) + tests_dir = root / delivery.ACCURACY_TESTS_DIR + tests_dir.mkdir(parents=True, exist_ok=True) + (tests_dir / "test_modeling_v2_gpt_oss.py").write_text("def test(): ...\n", encoding="utf-8") + db = root / delivery.TEST_DB_DIR + db.mkdir(parents=True) + (db / "l0_gb300.yml").write_text( + "- condition: {}\n tests:\n - accuracy/test_modeling_v2_gpt_oss.py::T::t\n", + encoding="utf-8", + ) + return root + + +def _verify(root: Path): + return delivery.verify_delivery( + repo_root=root, + target_rel=TARGET_REL, + family="gpt_oss", + reference_yaml=f"{delivery.ACCURACY_REFERENCES_DIR}/gsm8k.yaml", + model_key="openai/gpt-oss-120b", + accuracy_test="test_modeling_v2_gpt_oss.py", + ) + + +def test_a_complete_delivery_passes(tmp_path): + assert _verify(_build_repo(tmp_path)).passed + + +def test_imports_are_read_from_the_ast_not_by_importing(tmp_path): + root = _build_repo(tmp_path) + entries = delivery.imported_catalog_entries(root / TARGET_REL / "modeling.py") + assert entries == [ + ("gemm", "cublas_mm"), + ("norm", "flashinfer_rmsnorm"), + ("torch", "embedding"), + ] + + +def test_required_ops_are_read_statically(tmp_path): + root = _build_repo(tmp_path) + assert delivery.declared_required_ops(root / TARGET_REL / "modeling.py") == [ + "flashinfer_rmsnorm", + "cublas_mm", + ] + + +def test_a_target_declaring_no_ops_is_flagged(tmp_path): + root = _build_repo(tmp_path, modeling="X = 1\n") + report = _verify(root) + failure = next(c for c in report.failures if c.name == "target:REQUIRED_TRTLLM_OPS") + assert "renamed op" in failure.detail + + +def test_a_missing_catalog_test_fails(tmp_path): + root = _build_repo(tmp_path) + (root / delivery.CATALOG_TESTS_DIR / "norm" / "test_modeling_v2_flashinfer_rmsnorm.py").unlink() + report = _verify(root) + assert not report.passed + assert any(c.name == "test:norm/flashinfer_rmsnorm" for c in report.failures) + + +def test_mirrors_need_no_contract_or_test(tmp_path): + """`catalog/torch/` mirrors carry no contract, test or receipts by design.""" + root = _build_repo(tmp_path) + report = _verify(root) + assert report.passed + assert not any(c.name == "test:torch/embedding" for c in report.checks) + + +def test_an_unindexed_entry_fails(tmp_path): + root = _build_repo(tmp_path) + index = root / delivery.CATALOG_DIR / "index.yaml" + index.write_text(yaml.safe_dump({"entries": [{"path": "gemm/cublas_mm.py"}]}), encoding="utf-8") + report = _verify(root) + assert any("index.yaml" in c.detail for c in report.failures) + + +def test_an_unrouted_target_is_caught(tmp_path): + """A target nothing routes to is worse than missing: `auto` answers instead.""" + root = _build_repo(tmp_path) + routing = root / delivery.MODELING_V2_PACKAGE / "models" / "gpt_oss" / "routing.py" + routing.write_text("TARGET_MODULES = {}\n", encoding="utf-8") + report = _verify(root) + failure = next(c for c in report.failures if c.name == "routing:target") + assert "built-in implementation would answer" in failure.detail + + +def test_a_missing_anchor_is_caught(tmp_path): + root = _build_repo(tmp_path) + refs = root / delivery.ACCURACY_REFERENCES_DIR / "gsm8k.yaml" + refs.write_text(yaml.safe_dump({"someone/else": [{"accuracy": 1.0}]}), encoding="utf-8") + report = _verify(root) + assert any(c.name == "accuracy:reference" for c in report.failures) + + +def test_an_unregistered_test_is_caught(tmp_path): + """A gate on no list passes exactly once.""" + root = _build_repo(tmp_path) + (root / delivery.TEST_DB_DIR / "l0_gb300.yml").write_text("- condition: {}\n", encoding="utf-8") + report = _verify(root) + failure = next(c for c in report.failures if c.name == "tests:test-db") + assert "CI would never run it" in failure.detail + + +def test_summary_for_auditor_includes_the_dump(tmp_path): + report = _verify(_build_repo(tmp_path)) + rendered = delivery.summarize_for_auditor(report, dump_excerpt="batch=1 dtype=bf16") + assert "already run" in rendered + assert "batch=1 dtype=bf16" in rendered + + +# ------------------------------------- delivery must check git, not just disk + + +def _git(tmp_path, *args): + subprocess.run(["git", *args], cwd=tmp_path, check=True, capture_output=True) + + +def _git_repo(tmp_path): + _git(tmp_path, "init", "-q") + _git(tmp_path, "config", "user.email", "t@e.st") + _git(tmp_path, "config", "user.name", "t") + return tmp_path + + +def test_outside_a_git_work_tree_no_check_is_made(tmp_path): + """Absence of git is absence of a claim, not a failure. + + A fixture in a temporary directory is not a delivery, and a check that + failed there would make every other test in this file red for a reason + that has nothing to do with delivery. + """ + assert delivery.check_tree_committed(tmp_path, ("anything",)) == [] + + +def test_a_clean_tree_passes(tmp_path): + _git_repo(tmp_path) + (tmp_path / "t.py").write_text("x = 1\n", encoding="utf-8") + _git(tmp_path, "add", "t.py") + _git(tmp_path, "commit", "-qm", "t") + checks = delivery.check_tree_committed(tmp_path, ("t.py",)) + assert len(checks) == 1 and checks[0].passed + + +def test_an_uncommitted_target_fails_even_though_the_file_is_there(tmp_path): + """The failure this exists for, and it was measured once. + + A run passed `routing:target` and reported an accuracy number while + `routing.py`, `modeling.py` (984 changed lines) and `weights.py` were + uncommitted — and the catalog commits certifying those lines were in + history. It would have closed `delivered` and shipped nothing. + """ + _git_repo(tmp_path) + (tmp_path / "seed").write_text("s\n", encoding="utf-8") + _git(tmp_path, "add", "seed") + _git(tmp_path, "commit", "-qm", "seed") + (tmp_path / "modeling.py").write_text("# never committed\n", encoding="utf-8") + + checks = delivery.check_tree_committed(tmp_path, ("modeling.py",)) + assert len(checks) == 1 + assert not checks[0].passed + assert "a clone would not have them" in checks[0].detail + assert "modeling.py" in checks[0].detail + + +def test_a_modified_but_uncommitted_target_also_fails(tmp_path): + """Not just untracked — the real incident was tracked files with edits. + + It is also the F011 shape: a tree dirtied while gates were queued, where + a pass attributes a number to a commit it was never measured on. + """ + _git_repo(tmp_path) + (tmp_path / "routing.py").write_text("v = 1\n", encoding="utf-8") + _git(tmp_path, "add", "routing.py") + _git(tmp_path, "commit", "-qm", "r") + (tmp_path / "routing.py").write_text("v = 2\n", encoding="utf-8") + assert not delivery.check_tree_committed(tmp_path, ("routing.py",))[0].passed + + +def test_verify_delivery_includes_the_git_check(): + """Reachability, not behaviour: an uncalled check checks nothing.""" + assert "check_tree_committed" in inspect.getsource(delivery.verify_delivery) + + +# --------------------------------------------------------- against the real tree + +_REPO_ROOT = Path(__file__).resolve().parents[4] + + +@pytest.mark.skipif( + not (_REPO_ROOT / delivery.MODELING_V2_PACKAGE).is_dir(), + reason="not running inside a TensorRT-LLM checkout with modeling_v2 present", +) +def test_the_real_merged_target_passes_every_check(): + """The path conventions must match what the repository actually uses. + + `gpt_oss / sm_103 / tp1` is merged and gated, so every check here has a + known-good answer. If this fails, a convention moved and the checks are + describing a layout that no longer exists. + """ + report = delivery.verify_delivery( + repo_root=_REPO_ROOT, + target_rel=TARGET_REL, + family="gpt_oss", + reference_yaml=f"{delivery.ACCURACY_REFERENCES_DIR}/gsm8k.yaml", + model_key="openai/gpt-oss-120b", + accuracy_test="test_modeling_v2_gpt_oss.py", + ) + assert report.passed, report.render() diff --git a/agent-flow/tests/workflows/staircase_onboard/test_dspark.py b/agent-flow/tests/workflows/staircase_onboard/test_dspark.py new file mode 100644 index 000000000000..21067d029132 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_dspark.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""DSpark as its own speculator kind, and never a spelling of MTP. + +DeepSeek-V4.1 ships its DSpark block drafter as `mtp.*` with +`num_nextn_predict_layers: 3`. Every test here guards one of the ways that +namespace leads a run to gate the wrong drafter, or the right drafter at the +wrong depth, or not at all. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import gates, workflow +from agent_flow.workflows.staircase_onboard.environment import resolve_environment +from agent_flow.workflows.staircase_onboard.gates import ShellResult + +REPO_ROOT = Path(__file__).resolve().parents[4] +QUICKSTART = REPO_ROOT / "examples" / "llm-api" / "quickstart_advanced.py" + + +def _env(): + return resolve_environment( + {"environment": {"kind": "local", "python": "/opt/venv/bin/python", "env": {}}}, + devices=(0,), + world_size=1, + ) + + +def _options(tmp_path, **speculative) -> str: + path = tmp_path / "dspark_llm_api.yaml" + path.write_text(yaml.safe_dump({"speculative_config": speculative}), encoding="utf-8") + return str(path) + + +# ------------------------------------------------------------------- flags + + +def test_dspark_is_its_own_algorithm_never_mtp(): + flags = gates.speculation_flags(["dspark"], max_draft_len=5) + assert flags == ["--spec_decode_algo=DSPARK", "--spec_decode_max_draft_len=5"] + assert not any("MTP" in flag for flag in flags) + + +@pytest.mark.parametrize( + "speculative", + [ + {"decoding_type": "DSpark", "max_draft_len": 5}, + {"decoding_type": "DSpark", "block_size": 5}, + {"decoding_type": "DSpark", "max_draft_len": 5, "block_size": 5}, + ], +) +def test_the_dspark_depth_comes_from_the_options_yaml(tmp_path, speculative): + """Quickstart builds the config from the flag alone. + + A depth of 1 against a drafter trained for a 5-token block fails the + `block_size == max_draft_len` validation before the engine starts, so the + old default is not merely a weaker gate here, it is no gate. + """ + accuracy = {"llm_api_options": {"dspark": _options(tmp_path, **speculative)}} + assert gates.resolve_max_draft_len(accuracy, "dspark") == 5 + + +def test_max_draft_len_wins_over_its_deprecated_alias(tmp_path): + """The engine's own order: `MTPDecodingConfig` drops the alias when both are set.""" + both = _options(tmp_path, decoding_type="MTP", max_draft_len=2, num_nextn_predict_layers=3) + assert gates.resolve_max_draft_len({"llm_api_options": {"mtp": both}}, "mtp") == 2 + + +def test_dspark_is_acceptance_gated(): + assert "dspark" in workflow.ACCEPTANCE_GATED_UNITS + + +# ------------------------------------------------ the acceptance gate for dspark + + +def _task(tmp_path, floor): + return { + "checkpoint_path": str(tmp_path / "ckpt"), + "trtllm_repo_path": str(tmp_path), + "parallel": "tp1", + "features": ["dspark"], + "runtime": {"kv_cache_reuse": True, "extra_flags": []}, + "accuracy": { + "llm_api_options": { + "dspark": _options(tmp_path, decoding_type="DSpark", max_draft_len=5) + }, + "acceptance": {"dspark": {"min_absolute": floor}}, + }, + } + + +def _runner(accept_line): + """`=require` drafts; `=off` cannot load a new architecture at all.""" + seen: list[str] = [] + + def run(command, env=None, timeout=0): + seen.append(command) + if "TRTLLM_MODELING_V2=off" in command: + return ShellResult( + returncode=1, + stderr="RuntimeError: Unknown model architecture: DeepseekV41ForCausalLM", + ) + return ShellResult(returncode=0, stdout=accept_line) + + run.seen = seen # type: ignore[attr-defined] + return run + + +def test_the_declared_floor_stands_in_for_the_missing_stock_leg(tmp_path): + """A new architecture has no `=off` leg, so the declared floor decides. + + `min_absolute` must work for dspark exactly as it does for mtp, and the + verdict must say it is the weaker claim. + """ + run = _runner( + "[0] Accept rate: 3.10 (tokens=62, iterations=20)\n" + "[1] Accept rate: 2.90 (tokens=58, iterations=20)\n" + ) + result = gates.run_acceptance_gate( + _env(), + task_data=_task(tmp_path, 1.2), + runner=run, + unit="dspark", + features=["dspark"], + ) + assert result.passed, result.detail + assert result.measured == pytest.approx(120 / 40) + assert result.reference == pytest.approx(1.2) + assert "WEAKER CLAIM" in result.detail + seen = run.seen # type: ignore[attr-defined] + assert "--spec_decode_algo=DSPARK" in seen[0] + assert "--spec_decode_max_draft_len=5" in seen[0] + + +def test_below_the_declared_floor_fails(tmp_path): + run = _runner("[0] Accept rate: 1.05 (tokens=42, iterations=40)\n") + result = gates.run_acceptance_gate( + _env(), task_data=_task(tmp_path, 1.2), runner=run, unit="dspark", features=["dspark"] + ) + assert not result.passed + # Long enough to be a measurement, so the floor is what failed it. + assert "too short" not in result.detail + assert result.measured == pytest.approx(1.05) + + +# --------------------------------------------------------------- quickstart + + +def _function_source(name: str) -> str: + source = QUICKSTART.read_text(encoding="utf-8") + node = next( + n for n in ast.parse(source).body if isinstance(n, ast.FunctionDef) and n.name == name + ) + return ast.get_source_segment(source, node) or "" + + +def test_quickstart_builds_a_dspark_config(): + """The gates emit `--spec_decode_algo=DSPARK`; quickstart has to honour it.""" + setup = _function_source("setup_llm") + branch = setup[setup.index('spec_decode_algo == "DSPARK"') :] + branch = branch[: branch.index("elif")] + assert "DSparkDecodingConfig(" in branch + assert "max_draft_len=args.spec_decode_max_draft_len" in branch + + +def test_the_accept_rate_line_does_not_depend_on_the_algorithm(): + """The acceptance gate parses `Accept rate: ... (tokens=T, iterations=I)`. + + It is printed from the streamed iterations in `main`, which never looks at + which speculator is configured -- so DSpark feeds it exactly as MTP does. + """ + main = _function_source("main") + assert "Accept rate:" in main + assert "spec_decode_algo" not in main diff --git a/agent-flow/tests/workflows/staircase_onboard/test_environment.py b/agent-flow/tests/workflows/staircase_onboard/test_environment.py new file mode 100644 index 000000000000..dae27282d825 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_environment.py @@ -0,0 +1,654 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the resolved launch environment. + +This is the module that replaces four hops of agent-side environment +discovery with one string, so the tests care about exactly one thing: that +the string is right, and that the prompt block quotes the same string the +orchestrator will run. +""" + +from __future__ import annotations + +import pytest + +from agent_flow.workflows.staircase_onboard import environment +from agent_flow.workflows.staircase_onboard.environment import ( + MULTI_TASK_SLURM_SCRUB, + MULTI_TASK_SLURM_VARS, + RunEnvironment, + resolve_environment, +) + +_LOCAL_TASK = { + "environment": { + "kind": "local", + "python": "/opt/venv/bin/python", + "env": {"TRTLLM_MODELING_V2": "require", "TMPDIR": "/scratch/tmp"}, + } +} + +_SLURM_TASK = { + "environment": { + "kind": "slurm", + "python": "/opt/venv/bin/python", + "env": {"TRTLLM_MODELING_V2": "require"}, + "partition": "batch", + "account": "coreai", + "container_image": "/img/trtllm.sqsh", + "mounts": ["/lustre:/lustre", "/scratch:/scratch"], + "time_limit": "04:00:00", + } +} + + +def test_local_prefix_carries_devices_and_env(): + env = resolve_environment(_LOCAL_TASK, devices=[2, 3], world_size=2) + prefix = env.prefix + assert prefix.startswith("CUDA_VISIBLE_DEVICES=2,3 ") + assert "TRTLLM_MODELING_V2=require" in prefix + assert "TMPDIR=/scratch/tmp" in prefix + + +def test_local_env_order_is_stable(): + """A prefix that reorders between calls makes two identical runs look different.""" + env = resolve_environment(_LOCAL_TASK, devices=[0], world_size=1) + assert env.prefix == env.prefix + assert env.prefix.index("TMPDIR") < env.prefix.index("TRTLLM_MODELING_V2") + + +def test_local_prefix_without_devices_omits_the_mask(): + env = resolve_environment(_LOCAL_TASK, devices=[], world_size=1) + assert "CUDA_VISIBLE_DEVICES" not in env.prefix + + +def test_slurm_prefix_is_an_srun_line(): + env = resolve_environment(_SLURM_TASK, devices=[], world_size=4) + prefix = env.prefix + assert prefix.startswith("srun ") + assert "--partition=batch" in prefix + assert "--account=coreai" in prefix + assert "--time=04:00:00" in prefix + assert "--gres=gpu:4" in prefix + assert "--container-image=/img/trtllm.sqsh" in prefix + assert "--container-mounts=/lustre:/lustre,/scratch:/scratch" in prefix + assert "--export=ALL,TRTLLM_MODELING_V2=require" in prefix + + +def test_slurm_prefix_never_sets_cuda_visible_devices(): + """Under Slurm the allocation is the claim; masking it hides our own devices.""" + env = resolve_environment(_SLURM_TASK, devices=[0, 1, 2, 3], world_size=4) + assert "CUDA_VISIBLE_DEVICES" not in env.prefix + + +def test_command_prefixes_a_string(): + env = resolve_environment(_LOCAL_TASK, devices=[1], world_size=1) + assert env.command("trtllm-eval --model x gsm8k").endswith("trtllm-eval --model x gsm8k") + assert env.command("trtllm-eval").startswith("CUDA_VISIBLE_DEVICES=1") + + +def test_command_quotes_a_sequence(): + env = resolve_environment(_LOCAL_TASK, devices=[], world_size=1) + assert "'a b'" in env.command(["echo", "a b"]) + + +def test_python_command_uses_the_declared_interpreter(): + env = resolve_environment(_LOCAL_TASK, devices=[0], world_size=1) + assert "/opt/venv/bin/python script.py" in env.python_command("script.py") + + +def test_env_values_with_spaces_are_quoted(): + task = { + "environment": { + "kind": "local", + "python": "/p", + "env": {"SOME_FLAGS": "-a -b c"}, + } + } + env = resolve_environment(task, devices=[], world_size=1) + assert "SOME_FLAGS='-a -b c'" in env.prefix + + +def test_describe_quotes_the_same_prefix_the_orchestrator_runs(): + """The prompt block and the gate invocation must not be able to drift.""" + env = resolve_environment(_LOCAL_TASK, devices=[0, 1], world_size=2) + described = env.describe() + assert env.prefix in described + assert env.python in described + assert "world size 2" in described + assert "symptom to report" in described + + +def test_describe_tells_slurm_runs_the_devices_are_allocated(): + env = resolve_environment(_SLURM_TASK, devices=[], world_size=4) + assert "allocated by Slurm" in env.describe() + + +def test_world_size_is_at_least_one(): + env = resolve_environment(_LOCAL_TASK, devices=[], world_size=0) + assert env.world_size == 1 + + +def test_environment_is_frozen(): + env = RunEnvironment(kind="local", python="/p") + try: + env.kind = "slurm" # type: ignore[misc] + except Exception: + return + raise AssertionError("RunEnvironment must be immutable once resolved") + + +# ------------------------------------------------- the multi-rank launch rule + +# These pin TensorRT-LLM's own launch rules, carried from +# `.claude/skills/trtllm-test-script-builder/`. They assert exact strings on +# purpose: the whole failure mode is a plausible-looking srun line that is +# missing one flag, and a loose assertion would pass on exactly that. + + +def _slurm(world_size, gpus_per_node=4, **over): + kwargs = dict( + kind="slurm", + python="/venv/bin/python", + world_size=world_size, + gpus_per_node=gpus_per_node, + partition="batch", + account="acct", + container_image="/img.sqsh", + mounts=("/lustre:/lustre",), + time_limit="04:00:00", + env={"TRTLLM_MODELING_V2": "require"}, + ) + kwargs.update(over) + return RunEnvironment(**kwargs) + + +def test_one_task_per_rank_is_emitted(): + """The flag whose absence causes MPI_ERR_SPAWN. + + Without `--ntasks`, srun launches one task, the MPI world size is 1, and + `need_spawn_mpi_workers` (`llmapi/mpi_session.py`) is + `get_mpi_world_size() == 1 and model_world_size > 1` — true — so the LLM + API tries to create the other ranks dynamically, which pyxis forbids. + """ + command = _slurm(4).command("python x.py") + assert "--ntasks=4" in command + assert "--ntasks-per-node=4" in command + assert "--nodes=1" in command + + +def test_gres_is_per_node_not_per_world(): + """They differ the moment a run spans nodes, and `--gres` is per node.""" + env = _slurm(8, gpus_per_node=4) + assert env.nodes == 2 + command = env.command("python x.py") + assert "--gres=gpu:4" in command + assert "--gres=gpu:8" not in command + assert "--ntasks=8" in command + + +def test_multi_rank_wraps_the_payload_in_the_launcher(): + """`trtllm-llmapi-launch` makes rank 0 the leader and the rest workers.""" + command = _slurm(4).command("/venv/bin/python x.py") + assert "bash -c 'trtllm-llmapi-launch /venv/bin/python x.py'" in command + assert "--mpi=pmix" in command + + +def test_a_single_task_scrubs_mpi_instead_of_wrapping(): + """One task is the special case, and for a specific reason. + + mpi4py auto-initializes MPI whenever PMI variables are present, and inside + a PMIx-managed step they always are — so a single-task job initializes MPI + for no reason and then dies spawning a worker. Clearing the variables lets + it skip MPI init entirely. + """ + command = _slurm(1).command("/venv/bin/python x.py") + assert "trtllm-llmapi-launch" not in command + assert "--mpi=" not in command # `none` is expressed by omitting the flag + assert "${!PMIX@}" in command + assert "exec /venv/bin/python x.py" in command + + +def test_the_scrub_keeps_the_rank_identifiers(): + """They identify the rank rather than bootstrap MPI, and are read later.""" + command = _slurm(1).command("python x.py") + assert "SLURM_PROCID" in command + assert "SLURM_LOCALID" in command + + +def test_an_explicit_pmix_variant_is_honoured(): + """Clusters report `pmix`, `pmix_v3`, `pmix_v5`; `srun --mpi=list` says.""" + assert "--mpi=pmix_v5" in _slurm(4, mpi="pmix_v5").command("python x.py") + + +def test_a_container_name_is_emitted_so_pyxis_can_reuse_the_import(): + """A 26 GB image costs 4-8 minutes per job without it.""" + assert "--container-name=trtllm-run" in _slurm(4, container_name="trtllm-run").command("x") + assert "--container-name" not in _slurm(4).command("x") + + +def test_the_emitted_line_matches_the_deployment_guide_shape(): + """Checked against `docs/source/deployment-guide/` two-node example.""" + command = _slurm(8, gpus_per_node=4).command("python3 -m tensorrt_llm.commands.serve M") + for fragment in ( + "--nodes=2", + "--ntasks=8", + "--ntasks-per-node=4", + "--mpi=pmix", + "--gres=gpu:4", + "trtllm-llmapi-launch python3 -m tensorrt_llm.commands.serve M", + ): + assert fragment in command, fragment + + +def test_spool_and_local_apply_the_same_rule_through_mpirun(): + """There is no srun inside a spool script, but the engine rule is the same.""" + env = RunEnvironment(kind="session_spool", python="/venv/bin/python", world_size=4, spool="/s") + command = env.command("/venv/bin/python x.py") + assert "mpirun -n 4" in command + assert "trtllm-llmapi-launch" in command + + +def test_a_single_rank_spool_command_is_left_alone(): + env = RunEnvironment(kind="session_spool", python="/venv/bin/python", world_size=1, spool="/s") + assert "mpirun" not in env.command("/venv/bin/python x.py") + + +# ------------------------- mpirun inside someone else's Slurm allocation +# +# Two defects, one cause: a `session_spool` job is ONE Slurm task that owns +# the node and runs `mpirun` itself, so both Open MPI's slot accounting and +# `trtllm-llmapi-launch`'s rank derivation read the allocation and get the +# wrong answer. Neither is a property of the model, the target or the run. +# The `slurm` branch must keep neither workaround. + + +def test_mpirun_oversubscribes_because_the_allocation_is_one_task(): + """`-n 4` against a one-task allocation dies before any Python runs. + + Open MPI takes its slot count from the resource manager and refuses with + *"There are not enough slots available in the system"*. The allocation + holds the whole node, so the slot count is bookkeeping. + """ + env = RunEnvironment(kind="session_spool", python="/venv/bin/python", world_size=4, spool="/s") + assert "mpirun -n 4 --oversubscribe --allow-run-as-root" in env.command("/venv/bin/python x.py") + + +def test_the_slurm_branch_never_oversubscribes(): + """There a slot shortage is a mis-sized allocation reporting itself. + + `srun --ntasks` genuinely sizes the world under `slurm`, so the flag + would hide the one failure that is worth seeing. + """ + assert "--oversubscribe" not in _slurm(4).command("/venv/bin/python x.py") + assert "--oversubscribe" not in _slurm(8, gpus_per_node=4).command("python3 -m x") + + +def test_mpirun_scrubs_the_slurm_task_vars_the_launcher_reads_first(): + """Otherwise every rank is rank 0 of a world of 1, and the run hangs. + + `trtllm-llmapi-launch` reads `SLURM_PROCID` before + `OMPI_COMM_WORLD_RANK` and `SLURM_NTASKS` before `OMPI_COMM_WORLD_SIZE`. + Inside a one-task allocation those are `0` and `1` for every rank, so all + four became `mgmn_leader_node` and waited forever for workers. + """ + env = RunEnvironment(kind="session_spool", python="/venv/bin/python", world_size=4, spool="/s") + command = env.command("/venv/bin/python x.py") + for name in ("SLURM_PROCID", "SLURM_NTASKS", "SLURM_LOCALID", "SLURM_NODEID"): + assert f"-u {name}" in command, name + assert command.index(MULTI_TASK_SLURM_SCRUB) < command.index("mpirun") + + +def test_the_scrub_names_every_var_the_step_sets_for_its_task(): + """A partial scrub still leaves the launcher a Slurm answer to prefer.""" + assert set(MULTI_TASK_SLURM_VARS) == { + "SLURM_PROCID", + "SLURM_NTASKS", + "SLURM_LOCALID", + "SLURM_NODEID", + "SLURM_STEP_ID", + "SLURM_TASK_PID", + "SLURM_STEP_NUM_TASKS", + } + + +def test_the_slurm_branch_never_scrubs_the_task_vars(): + """Under `srun` the same scrub breaks a launch that is already correct. + + `SLURM_PROCID` really is the rank there and `SLURM_NTASKS` really is the + world size; the one-task branch's own scrub deliberately spares them. + """ + for command in (_slurm(4).command("/venv/bin/python x.py"), _slurm(1).command("/p x.py")): + assert "env -u SLURM_PROCID" not in command + assert "-u SLURM_NTASKS" not in command + + +def test_a_single_rank_spool_command_gets_neither_workaround(): + """Both are properties of `mpirun`, and at one task there is no `mpirun`.""" + env = RunEnvironment(kind="session_spool", python="/p", world_size=1, spool="/s") + command = env.command("/p x.py", single_process=True) + assert "--oversubscribe" not in command + assert "env -u SLURM_PROCID" not in command + + +def test_a_single_process_command_at_world_size_4_gets_neither(): + """The probe is re-asked at one task, so it leaves the `mpirun` rules behind.""" + env = RunEnvironment(kind="session_spool", python="/p", world_size=4, spool="/s") + command = env.python_command("probe.py", single_process=True) + assert "--oversubscribe" not in command + assert "env -u SLURM_PROCID" not in command + + +def test_the_spool_prompt_quotes_the_line_the_orchestrator_will_run(): + """The prompt is the agent's only instruction, so a stale copy is a defect. + + An agent handed the bare `mpirun` line reproduces both failures by hand, + and the second of them presents as an engine that never boots. + """ + env = RunEnvironment(kind="session_spool", python="/p", world_size=4, spool="/s") + block = env.describe() + assert "mpirun -n 4 --oversubscribe --allow-run-as-root" in block + assert MULTI_TASK_SLURM_SCRUB in block + assert "Copy that line whole" in block + + +def test_a_multi_rank_spool_prompt_tells_the_agent_how_to_launch(): + """The prompt is the agent's only instruction; it must carry the rule.""" + env = RunEnvironment(kind="session_spool", python="/p", world_size=4, spool="/s") + block = env.describe() + assert "mpirun -n 4" in block + assert "trtllm-llmapi-launch" in block + assert "MPI_ERR_SPAWN" in block + assert "Do not work around that by reducing the world size" in block + + +def test_a_single_rank_spool_prompt_stays_quiet_about_it(): + """A rule a run cannot break is noise in that run's prompt.""" + env = RunEnvironment(kind="session_spool", python="/p", world_size=1, spool="/s") + assert "mpirun" not in env.describe() + + +# ------------------------------------------- work that is not engine work + +# A read-only probe is launched by the same module and by a different rule. +# The rule is the one that already exists for one task -- no launcher, MPI +# scrubbed -- asked for by name rather than re-derived by the caller. + + +def test_a_single_process_command_drops_the_multi_rank_launcher(): + """The inventory probe wants one process at any world size. + + It imports the build, enumerates it and exits. Sent through the engine's + launch rules at world size 4 it became + `mpirun -n 4 ... trtllm-llmapi-launch`, which needs an MPI the run's + `PATH` had deleted -- so the probe died with rc=127 and the plan phase + was handed an empty survey. + """ + env = RunEnvironment(kind="session_spool", python="/venv/bin/python", world_size=4, spool="/s") + command = env.python_command("probe.py", single_process=True) + assert "mpirun" not in command + assert "trtllm-llmapi-launch" not in command + assert "/venv/bin/python probe.py" in command + + +def test_a_single_process_command_still_scrubs_the_mpi_environment(): + """mpi4py auto-initializes whenever PMI variables are present. + + A spool script runs inside the session's own allocation, so they are, + and a lone probe that joins a world it never asked for is the same + failure the one-task Slurm branch exists to prevent. + """ + env = RunEnvironment(kind="session_spool", python="/p", world_size=4, spool="/s") + command = env.command("/p probe.py", single_process=True) + assert "${!PMIX@}" in command + assert "exec /p probe.py" in command + + +def test_a_single_process_slurm_command_sizes_the_step_for_one_task(): + """Shrinking the payload without shrinking the step runs N copies of it.""" + command = _slurm(4).python_command("probe.py", single_process=True) + assert "--ntasks=1" in command + assert "--ntasks-per-node=1" in command + assert "--gres=gpu:1" in command + assert "--mpi=" not in command + assert "trtllm-llmapi-launch" not in command + assert "${!PMIX@}" in command + + +def test_engine_work_is_untouched_by_the_new_distinction(): + """The gates genuinely do want a rank per device; only the probe does not.""" + command = _slurm(4).python_command("x.py") + assert "--ntasks=4" in command + assert "trtllm-llmapi-launch /venv/bin/python x.py" in command + + +# --------------------------------------------------------------- PATH merging + + +def test_a_declared_path_extends_the_inherited_one_rather_than_deleting_it(): + """The line that produced `mpirun: command not found`. + + A spec pins its own interpreter by putting its venv first. Written as a + plain assignment that also deletes `/usr/local/mpi/bin`, and the failure + arrives as rc=127 inside a launcher line several layers away. + """ + task = { + "environment": { + "kind": "local", + "python": "/venv/bin/python", + "env": {"PATH": "/venv/bin:/usr/local/bin:/usr/bin:/bin"}, + } + } + prefix = resolve_environment(task, devices=[], world_size=1).prefix + assert prefix == 'PATH=/venv/bin:/usr/local/bin:/usr/bin:/bin:"$PATH"' + + +def test_the_declared_entries_still_win_the_lookup(): + """Prepending, not appending: the whole point of declaring a PATH.""" + task = {"environment": {"kind": "local", "python": "/p", "env": {"PATH": "/venv/bin"}}} + prefix = resolve_environment(task, devices=[], world_size=1).prefix + assert prefix.index("/venv/bin") < prefix.index("$PATH") + + +def test_a_path_that_already_names_the_inherited_one_is_not_doubled(): + """Both spellings of the suffix an author may write by hand.""" + for declared in ("/venv/bin:$PATH", "/venv/bin:${PATH}"): + task = {"environment": {"kind": "local", "python": "/p", "env": {"PATH": declared}}} + prefix = resolve_environment(task, devices=[], world_size=1).prefix + assert prefix == 'PATH=/venv/bin:"$PATH"' + + +def test_the_merge_reaches_the_slurm_export_list_too(): + """`--export` carries the same assignments and the same deletion.""" + command = _slurm(4, env={"PATH": "/venv/bin"}).command("python x.py") + assert '--export=ALL,PATH=/venv/bin:"$PATH"' in command + + +def test_other_variables_are_still_replaced_outright(): + """Only PATH decides *which program runs*; nothing else gets the exception.""" + task = { + "environment": { + "kind": "local", + "python": "/p", + "env": {"TRTLLM_MODELING_V2": "require", "TMPDIR": "/scratch/tmp"}, + } + } + prefix = resolve_environment(task, devices=[], world_size=1).prefix + assert "$" not in prefix + + +# ------------------------- the orchestrator must outlive the compute session + + +def _spool_env(tmp_path, kind="session_spool"): + return RunEnvironment(kind=kind, python="/p", spool=str(tmp_path)) + + +def test_sharing_an_allocation_with_the_watcher_warns(tmp_path, monkeypatch): + """Five session losses across two runs, each discarding answered work. + + The orchestrator is GPU-free and blocks for hours; the compute session is + GPU-heavy and dies — preempted, torn down by a worker teardown, or killed + outright when an agent's `pkill -9` catches the watcher, which IS the + allocation's main process. Putting the first inside the second couples + their lifetimes for no benefit at all. + """ + (tmp_path / "watcher.jobid").write_text("559470\n", encoding="utf-8") + monkeypatch.setenv("SLURM_JOB_ID", "559470") + message = environment.shares_allocation_with_spool(_spool_env(tmp_path)) + assert message and "559470" in message + assert "--gpus-per-node=0" in message + + +def test_a_different_allocation_is_silent(tmp_path, monkeypatch): + """The correct configuration must not be warned at, or the warning is noise.""" + (tmp_path / "watcher.jobid").write_text("559470\n", encoding="utf-8") + monkeypatch.setenv("SLURM_JOB_ID", "561807") + assert environment.shares_allocation_with_spool(_spool_env(tmp_path)) is None + + +@pytest.mark.parametrize("contents", ["", "unknown\n"]) +def test_an_unknown_watcher_job_is_silent(tmp_path, monkeypatch, contents): + """Unknown is not a warning. Absent evidence is absent, not damning.""" + (tmp_path / "watcher.jobid").write_text(contents, encoding="utf-8") + monkeypatch.setenv("SLURM_JOB_ID", "559470") + assert environment.shares_allocation_with_spool(_spool_env(tmp_path)) is None + + +def test_a_missing_marker_is_silent(tmp_path, monkeypatch): + """A spool whose watcher predates this file, or is not running yet.""" + monkeypatch.setenv("SLURM_JOB_ID", "559470") + assert environment.shares_allocation_with_spool(_spool_env(tmp_path)) is None + + +def test_outside_slurm_and_outside_a_spool_it_is_silent(tmp_path, monkeypatch): + monkeypatch.delenv("SLURM_JOB_ID", raising=False) + (tmp_path / "watcher.jobid").write_text("559470\n", encoding="utf-8") + assert environment.shares_allocation_with_spool(_spool_env(tmp_path)) is None + + monkeypatch.setenv("SLURM_JOB_ID", "559470") + assert environment.shares_allocation_with_spool(_spool_env(tmp_path, kind="slurm")) is None + + +def test_the_watcher_writes_the_marker_as_its_own_file(): + """`watcher.jobid` is the evidence, and it must not share the heartbeat line. + + `scripts/spool_submit.sh` pipes `watcher.heartbeat` straight into `date -d` + after stripping `alive=`, so anything extra on that line breaks every + submitter. A check nothing writes checks nothing, so this pins both halves: + the marker is written, and it is written somewhere else. + """ + from pathlib import Path + + watcher = Path(environment.__file__).with_name("spool_watcher.sh").read_text(encoding="utf-8") + assert '"$SPOOL/watcher.jobid"' in watcher + assert "${SLURM_JOB_ID:-unknown}" in watcher + for line in watcher.splitlines(): + if "watcher.heartbeat" in line and "echo" in line: + assert "SLURM_JOB_ID" not in line + + +# ------------------------------- per-workload substrate + + +_WORKLOAD_BASE = { + "environment": { + "kind": "session_spool", + "python": "/p", + "spool": "/s", + "env": {"A": "1", "B": "2"}, + "partition": "batch", + } +} + + +def test_without_an_override_nothing_changes(): + """Absent a `workloads` block, asking as a gate resolves to the base.""" + plain = resolve_environment(_WORKLOAD_BASE) + gated = resolve_environment(_WORKLOAD_BASE, workload="gate") + assert plain == gated + + +def test_a_gate_override_replaces_scalars_and_merges_env(): + """Short work on the shared session; a 10-15 minute gate in its own job.""" + spec = { + "environment": { + **_WORKLOAD_BASE["environment"], + "workloads": {"gate": {"kind": "slurm", "env": {"B": "override", "C": "3"}}}, + } + } + gate = resolve_environment(spec, workload="gate") + assert gate.kind == "slurm" + # Scalars the override does not mention are inherited, so a gate that only + # needs a different substrate does not restate mounts and accounts. + assert gate.partition == "batch" + # `env` merges key-by-key rather than replacing wholesale. + assert gate.env == {"A": "1", "B": "override", "C": "3"} + + +def test_an_override_for_another_workload_does_not_leak(): + """A gate override must not change the substrate of everything else.""" + spec = { + "environment": { + **_WORKLOAD_BASE["environment"], + "workloads": {"gate": {"kind": "slurm"}}, + } + } + assert resolve_environment(spec).kind == "session_spool" + assert resolve_environment(spec, workload=None).kind == "session_spool" + + +def test_the_workload_set_is_closed(): + """An override the orchestrator never dispatches under is a typo. + + It would silently never apply, and nothing downstream would say so — + `task_schema._validate_workloads` rejects it for exactly that reason. + """ + assert environment.WORKLOAD_KINDS == ("gate",) + + +# ------------------------------- image import is the parallelism lever + + +def _img_env(image, kind="session_spool"): + return RunEnvironment(kind=kind, python="/p", container_image=image, spool="/s") + + +def test_a_registry_reference_warns_with_the_command_to_fix_it(): + """Every allocation re-imports a 26 GB image: 5-7 minutes, 25-35 under load. + + That cost is what forces work onto one shared session, and a shared session + has one failure domain -- which is how a leaked server, a process-group + kill and a PMIx collision each took down everything else in one run. + """ + msg = environment.image_import_warning(_img_env("registry.example/ns/trtllm:tag")) + assert msg and "enroot import" in msg + assert "5-7 minutes" in msg + assert "25-35 minutes" in msg + assert "trtllm:tag" in msg + + +def test_an_existing_squashfs_is_silent(tmp_path): + """A correctly configured run must not be warned at.""" + sqsh = tmp_path / "trtllm.sqsh" + sqsh.write_bytes(b"x") + assert environment.image_import_warning(_img_env(str(sqsh))) is None + + +def test_a_squashfs_path_that_does_not_exist_is_a_different_warning(tmp_path): + """Pyxis fails to start rather than falling back to a registry.""" + msg = environment.image_import_warning(_img_env(str(tmp_path / "absent.sqsh"))) + assert msg and "does not exist" in msg + assert "enroot import" not in msg + + +def test_a_local_run_needs_no_container_and_is_silent(): + assert environment.image_import_warning(_img_env("registry/x:y", kind="local")) is None + assert environment.image_import_warning(_img_env("")) is None + + +def test_every_container_bearing_kind_is_checked(): + """`slurm`, `slurm_broker` and `session_spool` all hand the image to pyxis.""" + for kind in ("slurm", "slurm_broker", "session_spool"): + assert environment.image_import_warning(_img_env("registry/x:y", kind=kind)) diff --git a/agent-flow/tests/workflows/staircase_onboard/test_execution.py b/agent-flow/tests/workflows/staircase_onboard/test_execution.py new file mode 100644 index 000000000000..bedfce93dd64 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_execution.py @@ -0,0 +1,421 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Where the orchestrator's turns and commands actually go. + +Two things are being protected here and they are different in kind. + +The first is that the mapping from a ``task.yaml`` environment kind to a +mechanism lives in one place. Before this it was inline in +``_run_preflight``, twice, in the file that also owns resume and gate +sequencing — so adding the next mechanism meant editing the state machine. + +The second is that the abstraction is *load-bearing* rather than decorative. +An interface nothing calls is worse than no interface, because it reads as +structure. So several of these assert the real call path end to end: a run +configured for a spool really does put its gate command into that spool's +`pending` directory, through `CommandExecutor`, with no branch in the +workflow. +""" + +from __future__ import annotations + +import inspect +import json +from pathlib import Path + +import pytest +import yaml + +from agent_flow.runtime import ( + CommandSpec, + DelegatedCommandExecutor, + InProcessTurnExecutor, + LocalCommandExecutor, + SlurmStepExecutor, + SpoolCommandExecutor, + TurnOutcome, + TurnRequest, +) +from agent_flow.workflows.staircase_onboard import gates as gates_mod +from agent_flow.workflows.staircase_onboard import workflow as workflow_mod +from agent_flow.workflows.staircase_onboard.environment import RunEnvironment +from agent_flow.workflows.staircase_onboard.execution import ( + command_runner_for, + select_command_executor, + select_turn_executor, + shell_runner_for, +) +from agent_flow.workflows.staircase_onboard.preflight import CommandResult +from agent_flow.workflows.staircase_onboard.state import SCHEMA_VERSION, load_state +from agent_flow.workflows.staircase_onboard.survey import MARKER, Inventory +from agent_flow.workflows.staircase_onboard.transport import ExternalAgentTransport +from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + +from .test_workflow import Script, _install, healthy_smi + +# ------------------------------------------------------- executor selection + + +def _environment(kind: str, **kwargs) -> RunEnvironment: + return RunEnvironment(kind=kind, python="/opt/venv/bin/python", **kwargs) + + +def test_local_resolves_to_a_subprocess_here(): + assert isinstance(select_command_executor(_environment("local")), LocalCommandExecutor) + + +def test_slurm_resolves_to_a_step_executor_that_declares_itself_distributed(): + executor = select_command_executor(_environment("slurm", world_size=8)) + assert isinstance(executor, SlurmStepExecutor) + assert executor.capabilities.distributed + assert executor.capabilities.max_world_size == 8 + + +def test_a_session_spool_resolves_to_the_spool_it_names(tmp_path): + executor = select_command_executor(_environment("session_spool", spool=str(tmp_path))) + assert isinstance(executor, SpoolCommandExecutor) + assert executor.root == tmp_path + + +def test_a_session_spool_with_no_spool_falls_back_rather_than_crashing(): + """The schema already refuses this; resolution must not be the check. + + Two places enforcing one rule means the second one eventually disagrees + with the first. + """ + assert isinstance(select_command_executor(_environment("session_spool")), LocalCommandExecutor) + + +def test_the_broker_resolves_to_a_delegated_executor(tmp_path): + """The orchestrator cannot submit here, so a turn does — and only that. + + The verdict still comes from files the job writes; see + `tests/runtime/test_delegated_executor.py` for the property that an + agent's report cannot move it. + """ + executor = select_command_executor( + _environment("slurm_broker", container_image="img", partition="p"), + turns=InProcessTurnExecutor(lambda r: None), + workspace=tmp_path, + ) + assert isinstance(executor, DelegatedCommandExecutor) + assert executor.root == tmp_path / "broker" + + +def test_the_broker_submission_is_a_script_path_not_an_interpolated_command(tmp_path): + """The template quotes its payload; a gate command is full of quotes.""" + executor = select_command_executor( + _environment("slurm_broker", container_image="img", partition="p"), + turns=InProcessTurnExecutor(lambda r: None), + workspace=tmp_path, + ) + call = executor.render_submission("/w/broker/gate-1.sh") + assert 'command=["bash", "-lc", "bash /w/broker/gate-1.sh"])' in call + assert 'image="img"' in call + + +def test_without_a_turn_executor_the_broker_falls_back_rather_than_crashing(): + """A `--dry-run` has no turn executor to hand over and must still resolve.""" + executor = select_command_executor(_environment("slurm_broker")) + assert type(executor) is LocalCommandExecutor + + +# ------------------------------------------------ the legacy runner adapters + + +def test_the_shell_adapter_produces_the_shape_the_gates_take(): + run = shell_runner_for(LocalCommandExecutor()) + result = run("echo hi", None, 30) + assert isinstance(result, gates_mod.ShellResult) + assert result.ok and result.stdout.strip() == "hi" + + +def test_the_shell_adapter_keeps_a_timeout_distinct_from_a_failure(): + run = shell_runner_for(LocalCommandExecutor()) + assert run("exit 5", None, 30).returncode == 5 + timed = run("sleep 30", None, 1) + assert timed.timed_out and not timed.ok + + +def test_the_shell_adapter_passes_the_environment_through(): + run = shell_runner_for(LocalCommandExecutor()) + assert run("echo $STAIRCASE_ADAPTER", {"STAIRCASE_ADAPTER": "yes"}, 30).stdout.strip() == "yes" + + +def test_the_probe_adapter_produces_the_shape_preflight_takes(): + run = command_runner_for(LocalCommandExecutor()) + result = run(["echo", "a b"], None, 30) + assert isinstance(result, CommandResult) + assert result.stdout == "a b\n" + + +def test_the_probe_adapter_reports_a_missing_binary_as_127(): + run = command_runner_for(LocalCommandExecutor()) + assert run(["staircase-no-such-binary-xyz"], None, 30).returncode == 127 + + +def test_the_gate_default_runner_goes_through_the_local_executor(monkeypatch): + """Not merely equivalent to it -- actually it. + + The point of Phase 2 is that there is one implementation of "run a + command here". A second one living on in `gates.py` would be the thing + the refactor was for. + """ + seen: list[CommandSpec] = [] + real = LocalCommandExecutor.execute + + def spy(self, spec): + seen.append(spec) + return real(self, spec) + + monkeypatch.setattr(LocalCommandExecutor, "execute", spy) + gates_mod._default_shell_runner("echo hi", {"A": "b"}, 42) + + assert len(seen) == 1 + assert seen[0].argv == "echo hi" + assert seen[0].env == {"A": "b"} + assert seen[0].timeout_s == 42 + + +def test_the_spool_gate_runner_goes_through_the_spool_executor(tmp_path): + """`spool_runner` keeps its name and signature and owns none of the protocol.""" + run = gates_mod.spool_runner(tmp_path, poll_s=0.01, name_prefix="probe") + result = run("echo hi", None, 0) # nothing is watching, so it times out + assert result.timed_out + script = sorted((tmp_path / "pending").glob("probe-*.sh"))[-1] + assert script.read_text(encoding="utf-8").rstrip().endswith("echo hi") + + +# -------------------------------------------------------- turn executor choice + + +def test_no_transport_means_this_process_calls_the_layer(): + executor = select_turn_executor(None, lambda request: (lambda prompt: "hi")) + assert isinstance(executor, InProcessTurnExecutor) + assert executor.run(TurnRequest(role="smith", what="x", prompt="p")).text == "hi" + + +def test_a_turn_executor_is_used_as_given(tmp_path): + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + assert select_turn_executor(transport, lambda request: None) is transport + + +def test_a_pre_turn_request_callable_is_still_accepted(tmp_path): + """A keyword-callable transport still works. + + That is how the external transport was wired before the turn types + existed, and a caller outside this package has no way to know it was + supposed to change. + """ + seen: dict = {} + + def legacy(*, role, what, prompt, system_prompt="", decision_file=None): + seen.update(role=role, what=what, prompt=prompt, system_prompt=system_prompt) + return "legacy answer" + + executor = select_turn_executor(legacy, lambda request: None) + outcome = executor.run(TurnRequest(role="auditor", what="audit", prompt="p", system_prompt="S")) + assert isinstance(outcome, TurnOutcome) + assert outcome.text == "legacy answer" + assert seen == {"role": "auditor", "what": "audit", "prompt": "p", "system_prompt": "S"} + + +# ------------------------------------- the workflow no longer knows the names + + +def test_the_workflow_neither_names_nor_builds_a_platform_mechanism(): + """A mechanism should be removable without reading the state machine.""" + source = inspect.getsource(workflow_mod) + for name in ( + "session_spool", + "slurm_broker", + "spool_runner", + "SpoolCommandExecutor", + "LocalCommandExecutor", + "SlurmStepExecutor", + "srun", + ): + assert name not in source, f"workflow.py still names the mechanism {name!r}" + assert "environment.kind" not in source + assert ".kind ==" not in source + + +def test_the_one_remaining_platform_read_is_the_auditor_log_hint(): + """Pinned rather than hidden, because it is real and it is not Phase 2's. + + After a resume the gate log is empty (it is in memory), so the auditor is + told where the dumps are on disk instead of being handed an empty + section. That location is a property of where commands ran, so it does + belong to the executor — `ExecutionRecord.log_path` already carries it + per command. Routing the *glob* through the executor needs the evidence + plane that Phase 5 introduces, and doing it here would mean inventing + half of one. + + So: exactly one read, in exactly one method. If a second appears, this + fails and the question gets asked again. + """ + source = inspect.getsource(workflow_mod) + assert source.count('"spool"') == 1 + assert '"spool"' in inspect.getsource(StaircaseOnboardWorkflow._auditor_dump) + + +def test_run_turn_has_no_transport_branch(): + """Both transports answer `TurnExecutor.run`, so there is nothing to branch on. + + This is the Phase 1 acceptance criterion stated as a test: the file that + owns stage sequencing must not be able to tell the two apart. + """ + source = inspect.getsource(StaircaseOnboardWorkflow._run_turn) + assert "_agent_transport" not in source + assert "TurnRequest(" in source + assert ".run(request)" in source + + +def test_the_workflow_resolves_its_turn_executor_once(tmp_path): + workflow = StaircaseOnboardWorkflow(workspace=tmp_path / "ws") + assert isinstance(workflow._turns, InProcessTurnExecutor) + + transport = ExternalAgentTransport(tmp_path / "ws2", poll_s=0.01) + assert ( + StaircaseOnboardWorkflow(workspace=tmp_path / "ws2", agent_transport=transport)._turns + is transport + ) + + +# ------------------------------------------------- the real call path, wired + + +@pytest.fixture +def spool_task_file(tmp_path, repo, checkpoint): # noqa: F811 -- fixtures by name + spool = tmp_path / "spool" + spool.mkdir() + 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": "session_spool", + "python": "/opt/venv/bin/python", + "spool": str(spool), + "env": { + "TMPDIR": "/scratch/staircase-run-tmp", + "CLAUDE_CODE_TMPDIR": "/scratch/staircase-run-tmp", + }, + }, + "accuracy": { + "benchmark": "gsm8k", + "reference_yaml": "tests/integration/defs/accuracy/references/gsm8k.yaml", + "model_key": "openai/gpt-oss-120b", + "tolerance": 5.0, + }, + } + path = tmp_path / "task.yaml" + path.write_text(yaml.safe_dump(spec), encoding="utf-8") + return path + + +def test_a_spool_run_puts_its_commands_in_the_spool(tmp_path, spool_task_file, monkeypatch): + """End to end: preflight resolves the executor, and gates use it. + + `--dry-run` stops before the first agent turn, which is exactly far + enough to have resolved the executor. It does still take the inventory + survey, and that survey goes through the same spool — so it is stubbed + out here rather than left to block for its full deadline on a spool + nothing is watching. What this test is about is the mapping from an + environment kind to a mechanism, not the survey. + """ + monkeypatch.setattr(workflow_mod, "run_inventory_probe", lambda *a, **k: Inventory()) + workflow = StaircaseOnboardWorkflow( + workspace=tmp_path / "ws", + command_runner=healthy_smi(), + dry_run=True, + ) + _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(spool_task_file) + + assert isinstance(workflow._command_executor, SpoolCommandExecutor) + assert workflow._command_executor.root == tmp_path / "spool" + + # And the gate runner really goes there: nothing is watching, so it + # times out, but the work landed as a file. + result = workflow._shell_runner("echo gate-command", None, 0) + assert result.timed_out + scripts = sorted((tmp_path / "spool" / "pending").glob("*.sh")) + assert scripts and "echo gate-command" in scripts[-1].read_text(encoding="utf-8") + + +def test_an_injected_runner_still_wins_over_the_resolved_one( + tmp_path, spool_task_file, monkeypatch +): + """An operator debugging a gate, and every test in this suite.""" + calls: list[str] = [] + + def injected(command, env=None, timeout=0): + calls.append(command) + return gates_mod.ShellResult(returncode=0, stdout="ok") + + workflow = StaircaseOnboardWorkflow( + workspace=tmp_path / "ws", + command_runner=healthy_smi(), + shell_runner=injected, + dry_run=True, + ) + _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(spool_task_file) + + assert workflow._shell_runner is injected + assert workflow._shell_runner("x", None, 1).stdout == "ok" + assert calls[-1] == "x" + # The dry run's inventory probe went through it as well, which is the + # same claim stated twice: nothing resolved a second runner behind the + # injected one's back. Everything the run itself issued is that probe. + assert calls[:-1] and all(MARKER in command for command in calls[:-1]) + + +# ------------------------------------------------------- checkpoint compatibility + + +def test_the_checkpoint_schema_did_not_move(): + """Phase 1 and 2 must be resumable from a checkpoint written before them.""" + assert SCHEMA_VERSION == 1 + + +def test_a_checkpoint_written_before_the_refactor_still_loads(tmp_path): + """Byte-for-byte what the pre-refactor `save_state` wrote.""" + legacy = { + "version": 1, + "task_path": str(tmp_path / "task.yaml"), + "stage": "assemble", + "done": False, + "verdict": "", + "slug": "gpt-oss-120b-tp1", + "branch": "staircase/gpt-oss-120b-tp1", + "devices": [0], + "plan_round": 1, + "plan_approved": True, + "queue": [{"n": 1, "mode": "onboard", "target": "norm/rms", "status": "done"}], + "unit_index": 0, + "unit_turns": 2, + "units_done": [], + "audit_clean": False, + "human_stops": 0, + } + path = tmp_path / ".staircase_onboard_state.json" + path.write_text(json.dumps(legacy), encoding="utf-8") + + state = load_state(path) + assert state.stage == "assemble" + assert state.plan_approved is True + assert state.queue[0]["target"] == "norm/rms" + assert state.unit_turns == 2 + + +def test_the_dispatch_directory_layout_did_not_move(tmp_path): + """A resume shares its mailbox with the driver that was already running.""" + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + assert transport.root == Path(tmp_path) / "dispatch" diff --git a/agent-flow/tests/workflows/staircase_onboard/test_gate_wiring.py b/agent-flow/tests/workflows/staircase_onboard/test_gate_wiring.py new file mode 100644 index 000000000000..f61b4ec36518 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_gate_wiring.py @@ -0,0 +1,818 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""That the gates are REACHED — not that they work, which is test_gates.py's job. + +This file exists because `run_acceptance_gate` had three passing unit tests +and no call site anywhere in the orchestrator. Full coverage of a function's +own behaviour says nothing about whether anything calls it, and a gate nobody +calls gates nothing. Everything here asserts reachability and argument +routing. +""" + +from __future__ import annotations + +import inspect +from types import SimpleNamespace + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import gates, workflow +from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + +def _call_block(source: str, symbol: str, lines: int = 14) -> str: + """The `lines` source lines starting at the call to `symbol`. + + Sliced by line count rather than by a literal closing-paren sequence. + The gate calls moved one indentation level in when they were looped over + the member set, and an anchor keyed on indentation asserts formatting + where what is being checked is argument routing. + """ + rows = source.splitlines() + start = next(i for i, row in enumerate(rows) if symbol in row) + return "\n".join(rows[start : start + lines]) + + +# ------------------------------------------------------- reachability + + +def test_the_acceptance_gate_is_called_by_the_unit_loop(): + """The hole this file was written for. + + `grep -rn run_acceptance_gate` used to find a definition, an export and + three tests — and no caller. + """ + source = inspect.getsource(workflow.StaircaseOnboardWorkflow._assemble_one_unit) + assert "run_acceptance_gate" in source + + +def test_core_is_not_acceptance_gated(): + """Core has no draft path, so the comparison would be 1.0 against 1.0. + + A gate that passes whatever the code did is worse than no gate: it + appears in the evidence. + """ + assert "core" not in workflow.ACCEPTANCE_GATED_UNITS + + +def test_every_speculative_feature_is_acceptance_gated(): + """The acceptance gate's whole membership is the distribution-preserving features. + + If a new speculator lands in FEATURE_SKILLS and not here, its unit gets + boot and accuracy only — both of which are blind to a miscomputed draft + path by construction. + """ + from agent_flow.workflows.staircase_onboard.task_schema import FEATURE_SKILLS + + speculative = { + name for name, skill in FEATURE_SKILLS.items() if skill == "staircase-modeling-speculator" + } + assert speculative <= workflow.ACCEPTANCE_GATED_UNITS + + +def test_the_accuracy_gate_is_called_with_the_unit(): + """Without it the mtp unit is measured under core's LLM-API YAML. + + `resolve_llm_api_options` exists precisely to prevent that and was being + handed `unit=None` by the only caller in the package. + """ + source = inspect.getsource(workflow.StaircaseOnboardWorkflow._assemble_one_unit) + assert "unit=unit" in _call_block(source, "run_accuracy_gate") + + +def test_the_anchor_is_resolved_per_unit(): + source = inspect.getsource(workflow.StaircaseOnboardWorkflow._assemble_one_unit) + assert "self._reference_score(unit)" in source + + +# ------------------------------------------------------- per-unit resolution + + +_ACC = { + "benchmark": {"default": "gsm8k", "multimodal": "mmmu"}, + "reference_yaml": {"default": "refs/gsm8k.yaml", "multimodal": "refs/mmmu.yaml"}, + "llm_api_options": {"default": "/core.yaml", "mtp": "/mtp.yaml"}, + "discriminators": { + "default": {"quant_algo": "MIXED_PRECISION"}, + "mtp": {"quant_algo": "MIXED_PRECISION", "spec_dec_algo": "MTP"}, + }, +} + + +@pytest.mark.parametrize( + "unit,expected", + [("core", "gsm8k"), ("mtp", "gsm8k"), ("multimodal", "mmmu"), (None, "gsm8k")], +) +def test_benchmark_resolves_per_unit(unit, expected): + assert gates.resolve_per_unit(_ACC, "benchmark", unit) == expected + + +def test_reference_yaml_resolves_per_unit(): + """A multimodal unit anchors in a different FILE, not just a different row.""" + assert gates.resolve_per_unit(_ACC, "reference_yaml", "multimodal") == "refs/mmmu.yaml" + assert gates.resolve_per_unit(_ACC, "reference_yaml", "core") == "refs/gsm8k.yaml" + + +def test_a_scalar_still_means_the_whole_run(): + """Nothing that already worked may change shape.""" + assert gates.resolve_per_unit({"benchmark": "mmlu"}, "benchmark", "mtp") == "mmlu" + + +def test_discriminators_resolve_per_unit(): + assert gates.resolve_discriminators(_ACC, "mtp") == { + "quant_algo": "MIXED_PRECISION", + "spec_dec_algo": "MTP", + } + assert gates.resolve_discriminators(_ACC, "core") == {"quant_algo": "MIXED_PRECISION"} + + +def test_a_flat_discriminator_mapping_applies_to_every_unit(): + flat = {"discriminators": {"quant_algo": "NVFP4"}} + assert gates.resolve_discriminators(flat, "mtp") == {"quant_algo": "NVFP4"} + assert gates.resolve_discriminators(flat, None) == {"quant_algo": "NVFP4"} + + +def test_no_discriminators_means_none_rather_than_empty(): + """`{}` and `None` are different questions for load_reference_score. + + `{}` asks for the row with no tags at all; `None` means the caller did + not express an opinion. Conflating them is how a MIXED_PRECISION + checkpoint anchors to the bf16 number. + """ + assert gates.resolve_discriminators({}, "core") is None + + +# ------------------------------------------------------- the anchor itself + + +def _reference_file(tmp_path): + path = tmp_path / "gsm8k.yaml" + path.write_text( + yaml.safe_dump( + { + "Qwen/Model": [ + {"accuracy": 95.72}, + {"quant_algo": "MIXED_PRECISION", "accuracy": 96.36}, + {"quant_algo": "MIXED_PRECISION", "spec_dec_algo": "MTP", "accuracy": 95.57}, + ] + } + ), + encoding="utf-8", + ) + return path + + +def test_an_untagged_lookup_picks_the_bf16_row_which_is_why_unit_matters(tmp_path): + """Documents the defect rather than the fix: this is what used to happen.""" + assert gates.load_reference_score(_reference_file(tmp_path), "Qwen/Model") == 95.72 + + +def test_resolve_reference_uses_the_units_discriminators(tmp_path): + task = { + "trtllm_repo_path": str(tmp_path), + "accuracy": { + "reference_yaml": "gsm8k.yaml", + "discriminators": _ACC["discriminators"], + }, + } + _reference_file(tmp_path) + assert gates.resolve_reference(task, model_key="Qwen/Model", unit="core") == 96.36 + assert gates.resolve_reference(task, model_key="Qwen/Model", unit="mtp") == 95.57 + + +def test_a_model_with_no_untagged_row_needs_discriminators(tmp_path): + """The other half of the defect, and the more expensive half. + + `Qwen/Qwen3.8-27B` has only tagged rows, so an untagged lookup does not + pick the wrong number — it raises, at the END of an assembly. + """ + path = tmp_path / "r.yaml" + path.write_text( + yaml.safe_dump({"M": [{"quant_algo": "MIXED_PRECISION", "accuracy": 95.53}]}), + encoding="utf-8", + ) + with pytest.raises(KeyError, match="no entry matching"): + gates.load_reference_score(path, "M") + assert gates.load_reference_score(path, "M", discriminators={"quant_algo": "MIXED_PRECISION"}) + + +# ------------------------------------------- the re-draft has to say what was wrong + + +def test_a_schema_rejection_puts_the_validator_output_in_the_redraft_prompt(tmp_path): + """Otherwise the drafter is told to read a review that never happened. + + A malformed `vocabulary.yaml` is rejected BEFORE the reviewer runs, so + `progress.yaml` holds no reviewer entry. The re-draft prompt used to say + "fetch the PlanReviewer's REJECT and address every item it names" — and + there was nothing there. The drafter would re-emit the same file and, at + the default one-round cap, end the run with `plan-failed`. + """ + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + workspace = tmp_path / "ws" + workspace.mkdir() + (workspace / "vocabulary.yaml").write_text( + "verdict: assemblable\nsteps:\n- step: norm\n state: GAP\n", encoding="utf-8" + ) + flow = StaircaseOnboardWorkflow(workspace=workspace) + prompt = flow._plan_drafter_prompt(2) + + assert "schema check" in prompt + assert "PlanReviewer's REJECT" not in prompt + # The validator's own words, not a paraphrase. + assert "tier" in prompt + + +def test_a_first_round_prompt_is_not_a_redraft(tmp_path): + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + workspace = tmp_path / "ws" + workspace.mkdir() + flow = StaircaseOnboardWorkflow(workspace=workspace) + assert "re-draft" not in flow._plan_drafter_prompt(1) + + +# ------------------------------------------------- concurrent smiths and devices + + +def _smith_workflow(tmp_path, gpus_per_node): + import yaml as _yaml + + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + workspace = tmp_path / "ws" + workspace.mkdir() + flow = StaircaseOnboardWorkflow(workspace=workspace) + flow.task_data = { + "trtllm_repo_path": str(tmp_path), + "arch": "sm_103", + "environment": {"gpus_per_node": gpus_per_node}, + } + flow._kernel_sources_block = "" + del _yaml + return flow + + +def _item(n): + from agent_flow.workflows.staircase_onboard.vocabulary import QueueItem + + return QueueItem(n=n, mode="onboard", target="norm/x") + + +def test_concurrent_smiths_are_spread_across_the_nodes_devices(tmp_path): + """The launch protocol hands every role `CUDA_VISIBLE_DEVICES=0`. + + At ten parallel smiths that is ten processes on one device, and the + symptom is `Executor creation failed due to insufficient GPU memory` on a + 284 GB card — which reads as a target that allocates too much. Measured + on this host with four concurrent loads. + """ + flow = _smith_workflow(tmp_path, 4) + assigned = { + n: [ + line + for line in flow._smith_prompt(_item(n)).splitlines() + if "CUDA_VISIBLE_DEVICES" in line + ] + for n in range(1, 9) + } + got = [row[0].split("CUDA_VISIBLE_DEVICES=")[1].split("`")[0] for row in assigned.values()] + assert got == ["1", "2", "3", "0", "1", "2", "3", "0"] + + +def test_a_single_device_node_says_nothing_extra(tmp_path): + """No assignment to make; a contradictory instruction is worse than none.""" + flow = _smith_workflow(tmp_path, 1) + assert "CUDA_VISIBLE_DEVICES" not in flow._smith_prompt(_item(3)) + + +# ------------------------------------------- assembly turns and their decisions + + +def test_the_assembler_decision_path_tracks_the_turn_not_the_previous_phase(tmp_path): + """`current_iteration` was last set by the SMITH phase, to a batch index. + + Turn 2 of an assembly was therefore told to write `assembler-1.yaml` — + the path turn 1 had already used. Either the re-draft overwrites its + predecessor's decision, or the role notices the contradiction and writes + the right file anyway, which is what happened here. + """ + import inspect + + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + source = inspect.getsource(StaircaseOnboardWorkflow._assemble_one_unit) + head = source[: source.index("self._dispatch(")] + assert "self._progress_ctx.current_iteration = state.unit_turns" in head + + +def test_a_blocked_decision_from_a_previous_process_is_handled_before_redispatching(tmp_path): + """A resume re-enters the loop at the top and never acts on it otherwise. + + The gap handler runs immediately after a dispatch, so a BLOCKED decision + written by a process that then died is invisible: the resumed run + dispatches the assembler again and pays a full turn to re-derive an + answer already on disk. Measured — this run paid exactly that. + """ + import inspect + + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + source = inspect.getsource(StaircaseOnboardWorkflow._assemble_one_unit) + head = source[: source.index("self._dispatch(")] + assert "_handle_blocked_assembler" in head + + +def test_several_vocabulary_gaps_dispatch_several_smiths(tmp_path): + """`missing_entry` is one string and an assembly turn finds several gaps. + + All three of this run's gaps were the same class, found in one pass — + which is how assembly gaps come, because the turn sees the whole forward + before it stops. Dispatching the concatenation sends one smith after a + target that does not exist. + """ + import inspect + + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + source = inspect.getsource(StaircaseOnboardWorkflow._handle_blocked_assembler) + assert "re.split" in source + # Concurrently: gaps found in one turn have no dependency on each other. + assert "ThreadPoolExecutor" in source + + +# ------------------------------------------------- the runtime block, per unit + + +def _task(runtime): + return {"runtime": runtime, "trtllm_repo_path": "/repo", "parallel": "tp1"} + + +_PER_UNIT_RUNTIME = { + "default": { + "cuda_graph": True, + "chunked_prefill": True, + "overlap_scheduler": True, + "kv_cache_reuse": True, + "extra_flags": ["--max_seq_len=8192"], + }, + "multimodal": {"kv_cache_reuse": False}, +} + + +def test_a_flat_runtime_block_still_means_the_whole_run(): + flat = {"cuda_graph": True, "kv_cache_reuse": True, "extra_flags": []} + assert gates.resolve_runtime(_task(flat), "multimodal") == flat + + +def test_a_per_unit_block_is_layered_over_default(): + """A unit states only what differs, so `default` supplies the rest.""" + merged = gates.resolve_runtime(_task(_PER_UNIT_RUNTIME), "multimodal") + assert merged["kv_cache_reuse"] is False + assert merged["cuda_graph"] is True + assert merged["extra_flags"] == ["--max_seq_len=8192"] + + +def test_a_unit_with_no_block_of_its_own_gets_default(): + assert gates.resolve_runtime(_task(_PER_UNIT_RUNTIME), "core")["kv_cache_reuse"] is True + + +def test_the_units_really_do_disagree_in_the_emitted_flags(): + """This is the whole point, and it is not hypothetical. + + `core` and `mtp` were gated with block reuse ON and passed at 96.25 and + 96.10; the multimodal path requires it OFF. Block reuse is one of the two + switches spelled as a NEGATION, so "on" emits nothing and "off" emits + `--disable_kv_cache_reuse` — which is exactly the translation a reader + gets backwards. + """ + core = gates.runtime_flags(gates.resolve_runtime(_task(_PER_UNIT_RUNTIME), "core")) + mm = gates.runtime_flags(gates.resolve_runtime(_task(_PER_UNIT_RUNTIME), "multimodal")) + assert "--disable_kv_cache_reuse" not in core + assert "--disable_kv_cache_reuse" in mm + + +def test_the_boot_gate_takes_a_unit_because_it_cannot_read_the_options_file(): + """`quickstart_advanced.py` has no `--extra_llm_api_options`. + + So the gate that runs first and most often is the only one that cannot + see a unit's declared configuration through `llm_api_options`, and the + runtime block is its sole channel. + """ + import inspect + + assert "unit" in inspect.signature(gates.run_boot_gate).parameters + source = inspect.getsource(gates.run_boot_gate) + assert "resolve_runtime(task_data, unit)" in source + + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + loop = inspect.getsource(StaircaseOnboardWorkflow._assemble_one_unit) + assert "unit=unit" in _call_block(loop, "run_boot_gate") + + +def test_a_resume_only_auto_handles_a_vocabulary_gap(): + """Any other blocked reason already ended the run once. + + Handling it again just ends it again — measured: a `knowledge_gap` block + written at a session cut-off re-closed the run within seconds of every + resume, before the newly-extended session could dispatch anything. An + operator who resumed deliberately has said the condition should be + revisited, so the other reasons fall through to a fresh turn. + """ + import inspect + + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + source = inspect.getsource(StaircaseOnboardWorkflow._assemble_one_unit) + head = source[: source.index("self._dispatch(")] + assert 'prior_reason == "vocabulary_gap"' in head + + +# ------------------------------------------- the auditor's runtime-value dump + + +class _FakeGateLog: + def __init__(self, results): + self.results = results + + def render(self): + return "" + + +class _Res: + def __init__(self, name, stdout_tail): + self.name = name + self.stdout_tail = stdout_tail + + +class _DumpWorkflow: + """Just enough of the orchestrator to exercise `_auditor_dump`.""" + + RUNTIME_VALUE_MARKER = StaircaseOnboardWorkflow.RUNTIME_VALUE_MARKER + _auditor_dump = StaircaseOnboardWorkflow._auditor_dump + + def __init__(self, results, spool=None): + self.gate_log = _FakeGateLog(results) + self._spool = spool + + def _gate_env(self): + return SimpleNamespace(spool=self._spool) + + +_MARKER = StaircaseOnboardWorkflow.RUNTIME_VALUE_MARKER + + +def test_every_gate_with_a_dump_is_included_not_just_boot(): + """The boot gate drives two short prompts and records warm-up shapes. + + Both of this project's auditor findings were visible only in an accuracy + gate's dump, so passing `boot.stdout_tail` alone -- which is what this did + for its whole life -- hands the auditor the one dump that cannot show them. + """ + wf = _DumpWorkflow( + [ + _Res("boot", f"{_MARKER} ===\n RV ('batch', 8192)\n"), + _Res("accuracy", f"{_MARKER} ===\n RV ('vision', 12544)\n"), + ] + ) + dump = wf._auditor_dump() + assert "`boot` gate" in dump and "`accuracy` gate" in dump + assert "('vision', 12544)" in dump + + +def test_a_gate_with_no_dump_contributes_nothing(): + """It falls through to the directions, which legitimately mention `boot`. + + So this asserts on the section header, not on the word: a gate whose log + never carried a dump must not be rendered as though it had one. + """ + wf = _DumpWorkflow([_Res("boot", "just some log output\n")]) + dump = wf._auditor_dump() + assert "`boot` gate" not in dump + assert "just some log output" not in dump + + +def test_with_no_dump_at_all_the_auditor_is_told_where_to_look(): + """`gate_log` is in-memory, so a run RESUMED into deliver has no history. + + That is the normal way to reach this stage after a repair, and it is + exactly when the auditor most needs the dump. Handing it an empty section + is what made two auditors recover the logs by hand. + """ + wf = _DumpWorkflow([], spool="/spool/run") + dump = wf._auditor_dump() + assert "/spool/run/logs/*.log" in dump + assert _MARKER in dump + assert "union" in dump + # and it must not read as "the gates did not run" + assert "not a sign the gates did not run" in dump + + +def test_without_a_spool_it_still_names_the_marker(): + wf = _DumpWorkflow([], spool=None) + dump = wf._auditor_dump() + assert _MARKER in dump + assert "gate logs" in dump + + +def test_the_directions_prefer_the_accuracy_gates(): + """A boot dump is not a substitute, and the prompt has to say why.""" + wf = _DumpWorkflow([], spool="/spool/run") + assert "accuracy gates" in wf._auditor_dump() + + +# -------------------------------------------- adding a unit to a delivered run + + +class _GrowthWorkflow: + """Just enough of the orchestrator to exercise `_grown_units`.""" + + _grown_units = StaircaseOnboardWorkflow._grown_units + + def __init__(self, features): + self.task_data = {"features": list(features)} + + +def test_a_spec_that_grew_names_the_new_units_in_order(): + wf = _GrowthWorkflow(["mtp", "multimodal", "host_offload", "disagg"]) + state = SimpleNamespace(units_done=["core", "mtp", "multimodal"]) + assert wf._grown_units(state) == ["host_offload", "disagg"] + + +def test_a_spec_that_did_not_grow_names_nothing(): + wf = _GrowthWorkflow(["mtp"]) + state = SimpleNamespace(units_done=["core", "mtp"]) + assert wf._grown_units(state) == [] + + +def test_only_a_delivered_run_re_enters_on_growth(): + """A failed run is ALSO missing units from `units_done`. + + Re-entering there would re-dispatch paid agent turns with nothing having + changed -- exactly what REMEDIABLE_VERDICTS exists to prevent. So the + growth branch is guarded on the verdict being `delivered`, and this + asserts the guard is in the source rather than implied by it. + """ + src = inspect.getsource(StaircaseOnboardWorkflow._init_state) + assert 'state.verdict == "delivered" and self._grown_units(state)' in src + + +def test_growth_re_enters_at_assemble_and_clears_the_audit(): + """The new unit has to be gated, and then re-delivered and re-audited. + + Leaving `audit_clean` set would carry the previous unit's audit over a + target that has since changed. + """ + src = inspect.getsource(StaircaseOnboardWorkflow._init_state) + growth = src[src.index('state.verdict == "delivered"') :] + assert "state.stage = STAGE_ASSEMBLE" in growth + assert "state.audit_clean = False" in growth + + +# ------------------------------- decisions must travel on every transport + + +def test_an_external_transport_always_uses_decision_files(): + """The MCP tools need a session; an external transport has none. + + Without this the orchestrator neither tells the role where to write its + decision nor collects one, and `_run_deliver_phase` reads `progress.yaml` + and gets whatever was there before -- the PREVIOUS run's verdict. + + Measured: three audit passes in a row had their verdicts discarded this + way, including a `CLEAN` that was read back as the `UNCOVERED` from two + passes earlier. A loop like that cannot converge and looks exactly like a + gate that keeps failing. + """ + src = inspect.getsource(StaircaseOnboardWorkflow.__init__) + line = next(ln for ln in src.splitlines() if "self._decision_files =" in ln) + assert "agent_transport is not None" in line + + +def test_the_enterprise_mcp_reason_is_not_dropped(): + """The two reasons are independent; either alone requires decision files.""" + src = inspect.getsource(StaircaseOnboardWorkflow.__init__) + line = next(ln for ln in src.splitlines() if "self._decision_files =" in ln) + assert "ENTERPRISE_MCP_CONFIG.exists()" in line + assert " or " in line + + +# ------------------------------------------- witnesses for the new feature kinds + + +def test_every_feature_whose_standing_gates_are_blind_has_a_witness(): + """A new witnessed feature in FEATURE_SKILLS and not here gets boot and accuracy only.""" + from agent_flow.workflows.staircase_onboard.task_schema import FEATURE_SKILLS + + witnessed_skills = { + "staircase-modeling-host-offload", + "staircase-modeling-disagg", + "staircase-modeling-bounded-replay", + "staircase-modeling-kv-quant", + } + needing = {name for name, skill in FEATURE_SKILLS.items() if skill in witnessed_skills} + assert needing == set(workflow.WITNESS_GATED_UNITS) + assert workflow.WITNESS_GATED_UNITS["decoder_replay"] == "replay" + assert workflow.WITNESS_GATED_UNITS["encoder_replay"] == "replay" + assert workflow.WITNESS_GATED_UNITS["kv_quant"] == "kv_quant" + + +def test_the_witness_dispatch_reaches_the_new_witnesses(): + source = inspect.getsource(StaircaseOnboardWorkflow._run_witness) + assert "run_replay_witness" in source + assert "run_kv_quant_witness" in source + # The replay workload is declared, so it must travel from the block. + for key in ("window", "prompt_tokens", "prefix_tokens", "reuse_alignment"): + assert f'"{key}"' in source + + +@pytest.mark.parametrize( + "unit,kind,block,gap", + [ + ("decoder_replay", "replay", {"contrast_env": {"X": "0"}, "window": 128}, "prompt_tokens"), + ("encoder_replay", "replay", {"contrast_env": {"X": "0"}, "prefix_tokens": 9}, "window"), + ("kv_quant", "kv_quant", {"contrast_env": {"X": "bf16"}}, "expected_bytes_per_token"), + ("kv_quant", "kv_quant", {"expected_bytes_per_token": 890}, "contrast_env"), + ("host_offload", "residency", {"expected_mib": 1}, "contrast_env"), + ("disagg", "equivalence", None, "a block at all"), + # Without its own launcher leg B is leg A again (F021a). + ("disagg", "equivalence", {"contrast_env": {"TRTLLM_DISAGG": "1"}}, "contrast_command"), + ("disagg", "equivalence", {"contrast_command": "/leg_b.py"}, "contrast_env"), + ], +) +def test_an_unconfigured_witness_names_what_is_missing(unit, kind, block, gap): + assert gap in StaircaseOnboardWorkflow._witness_block_gap(unit, kind, block) + + +@pytest.mark.parametrize( + "unit,kind,block", + [ + ("decoder_replay", "replay", {"contrast_env": {"X": "0"}, "window": 1, "prompt_tokens": 9}), + ("kv_quant", "kv_quant", {"expected_bytes_per_token": 890, "exact": False}), + ("host_offload", "residency", {"contrast_env": {"X": "0"}}), + ("disagg", "equivalence", {"contrast_env": {"X": "1"}, "contrast_command": "/b.py"}), + ], +) +def test_a_configured_witness_has_no_gap(unit, kind, block): + """A lossy kv_quant drops the identity leg, and with it the contrast.""" + assert StaircaseOnboardWorkflow._witness_block_gap(unit, kind, block) is None + + +def test_the_disagg_block_reaches_the_equivalence_witness(tmp_path, monkeypatch): + """The declared workload, the per-leg timeout and a workdir leg B can read.""" + from agent_flow.workflows.staircase_onboard import witness + from agent_flow.workflows.staircase_onboard.environment import RunEnvironment + + from .test_workflow import _workflow + + seen: dict = {} + + def fake(environment, **kwargs): + seen.update(kwargs) + return gates.GateResult(name="equivalence", passed=True, detail="ok") + + monkeypatch.setattr(witness, "run_equivalence_witness", fake) + wf = _workflow(tmp_path) + wf.task_data = { + "features": ["disagg"], + "accuracy": { + "witness": { + "disagg": { + "contrast_env": {"TRTLLM_DISAGG": 1}, + "contrast_command": "/onboard/leg_b.py", + "prompt_tokens": [131, 2050], + "timeout": 2700, + "legs_env": {"TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY": 0}, + } + } + }, + } + wf.environment = RunEnvironment(kind="local", python="/py", world_size=4) + wf._gate_environment = wf.environment + assert wf._run_witness("disagg", "tep4", None).passed + assert seen["contrast_command"] == "/onboard/leg_b.py" + assert seen["contrast_env"] == {"TRTLLM_DISAGG": "1"} + assert seen["legs_env"] == {"TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY": "0"} + assert seen["prompt_tokens"] == [131, 2050] + assert seen["timeout"] == 2700 + assert seen["member"] == "tep4" + assert seen["workdir"] == wf.workspace / "witness" / "disagg" + + +def test_a_witness_without_a_timeout_keeps_the_boot_budget(tmp_path, monkeypatch): + from agent_flow.workflows.staircase_onboard import witness + from agent_flow.workflows.staircase_onboard.environment import RunEnvironment + + from .test_workflow import _workflow + + seen: dict = {} + + def fake(environment, **kwargs): + seen.update(kwargs) + return gates.GateResult(name="residency", passed=True, detail="ok") + + monkeypatch.setattr(witness, "run_residency_witness", fake) + wf = _workflow(tmp_path) + wf.task_data = { + "features": ["host_offload"], + "accuracy": {"witness": {"host_offload": {"contrast_env": {"X": "0"}, "expected_mib": 1}}}, + } + wf.environment = RunEnvironment(kind="local", python="/py", world_size=1) + wf._gate_environment = wf.environment + wf._run_witness("host_offload", "tp1", None) + assert seen["timeout"] == gates.BOOT_TIMEOUT_S + + +# ------------------------------- the batch limit reaches every token witness + + +@pytest.mark.parametrize( + "unit,function,block", + [ + ( + "kv_quant", + "run_kv_quant_witness", + {"contrast_env": {"X": "bf16"}, "expected_bytes_per_token": 890}, + ), + ( + "disagg", + "run_equivalence_witness", + {"contrast_env": {"X": "1"}, "contrast_command": "/b.py"}, + ), + ("host_offload", "run_residency_witness", {"contrast_env": {"X": "0"}, "expected_mib": 1}), + ], +) +@pytest.mark.parametrize("declared", [None, 1]) +def test_the_batch_limit_travels_from_the_block( + tmp_path, monkeypatch, unit, function, block, declared +): + """F036: absent it is None, so a block written before the key runs as it did.""" + from agent_flow.workflows.staircase_onboard import witness + from agent_flow.workflows.staircase_onboard.environment import RunEnvironment + + from .test_workflow import _workflow + + seen: dict = {} + + def fake(environment, **kwargs): + seen.update(kwargs) + return gates.GateResult(name=unit, passed=True, detail="ok") + + monkeypatch.setattr(witness, function, fake) + wf = _workflow(tmp_path) + if declared is not None: + block = block | {"max_batch_size": declared} + wf.task_data = {"features": [unit], "accuracy": {"witness": {unit: block}}} + wf.environment = RunEnvironment(kind="local", python="/py", world_size=1) + wf._gate_environment = wf.environment + wf._run_witness(unit, "tp1", None) + assert seen["max_batch_size"] == declared + + +def test_a_replay_witness_is_not_handed_a_batch_limit(): + """It counts work, not tokens, and takes no such argument: passing one would raise.""" + from agent_flow.workflows.staircase_onboard import witness + + source = inspect.getsource(StaircaseOnboardWorkflow._run_witness) + assert "schedule" not in _call_block(source, "run_replay_witness", lines=12) + assert "max_batch_size" not in inspect.signature(witness.run_replay_witness).parameters + + +def test_a_failed_witness_leaves_both_legs_in_the_journal(tmp_path, monkeypatch): + """The one-line summary cannot carry the excerpt; the journal and console do. + + Generated text can hold what rich reads as markup (`[/INST]`), which + must print as text rather than raise. + """ + from agent_flow.workflows.staircase_onboard import witness + from agent_flow.workflows.staircase_onboard.environment import RunEnvironment + + from .test_workflow import _workflow + + tail = "continuation #0: first differs at token 3\n native: '[/INST] yes'\n bf16 : 'no'" + + def fake(environment, **kwargs): + return gates.GateResult( + name="kv_quant", passed=False, detail="native storage changed '[/x]'", stdout_tail=tail + ) + + monkeypatch.setattr(witness, "run_kv_quant_witness", fake) + wf = _workflow(tmp_path) + wf.task_data = { + "features": ["kv_quant"], + "accuracy": { + "witness": { + "kv_quant": {"contrast_env": {"X": "bf16"}, "expected_bytes_per_token": 890} + } + }, + } + wf.environment = RunEnvironment(kind="local", python="/py", world_size=1) + wf._gate_environment = wf.environment + assert not wf._run_witness("kv_quant", "tp1", None).passed + journal = wf.journal.path.read_text(encoding="utf-8") + assert "witness `kv_quant` failed: evidence" in journal + assert "native: '[/INST] yes'" in journal diff --git a/agent-flow/tests/workflows/staircase_onboard/test_gates.py b/agent-flow/tests/workflows/staircase_onboard/test_gates.py new file mode 100644 index 000000000000..edf46e14069f --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_gates.py @@ -0,0 +1,1032 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the gates — the commands whose exit codes are the verdicts.""" + +from __future__ import annotations + +import itertools +import json +import shlex +import subprocess + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import gates +from agent_flow.workflows.staircase_onboard.environment import resolve_environment +from agent_flow.workflows.staircase_onboard.gates import ShellResult + +_TASK = { + "checkpoint_path": "/ckpt/Qwen3-30B-A3B", + "parallel": "tp1", + "trtllm_repo_path": "/repo", + "accuracy": { + "benchmark": "gsm8k", + "reference_yaml": "tests/integration/defs/accuracy/references/gsm8k.yaml", + "tolerance": 5.0, + }, +} + + +def _env(world_size: int = 1, devices=(0,)): + return resolve_environment( + {"environment": {"kind": "local", "python": "/opt/venv/bin/python", "env": {}}}, + devices=devices, + world_size=world_size, + ) + + +def recorder(result: ShellResult | list[ShellResult]): + """A runner that returns canned results and records the commands it saw.""" + results = result if isinstance(result, list) else [result] + seen: list[str] = [] + + def run(command, env=None, timeout=0): + seen.append(command) + return results[min(len(seen) - 1, len(results) - 1)] + + run.seen = seen # type: ignore[attr-defined] + return run + + +# --------------------------------------------------------------- topology + + +@pytest.mark.parametrize( + "segment,expected", + [ + ("tp1", ["--tp_size=1"]), + ("tp8", ["--tp_size=8"]), + ("tep4", ["--tp_size=4", "--moe_ep_size=4"]), + ("dep4", ["--tp_size=4", "--moe_ep_size=4", "--enable_attention_dp"]), + ], +) +def test_topology_flags_follow_the_segment(segment, expected): + """Every gate is topology-blind; deriving the flags is what keeps the name honest.""" + assert gates.topology_flags(segment, world_size=1) == expected + + +def test_unified_takes_its_size_from_the_world_size(): + """And, with no member, only ``--tp_size`` — which is the whole defect. + + Kept as the documented fallback for a caller that predates the member + set. The orchestrator always passes one; this is what it looks like when + nobody does, and it is why a unified target used to be gated at pure + tensor parallelism while claiming three configurations. + """ + assert gates.topology_flags("unified", world_size=4) == ["--tp_size=4"] + + +@pytest.mark.parametrize( + "member,expected", + [ + ("tp4", ["--tp_size=4"]), + ("tep4", ["--tp_size=4", "--moe_ep_size=4"]), + ("dep4", ["--tp_size=4", "--moe_ep_size=4", "--enable_attention_dp"]), + ], +) +def test_a_member_is_what_gets_translated(member, expected): + """`unified` names a family; the member names the configuration served. + + Without this, `--moe_ep_size` and `--enable_attention_dp` are emitted by + no gate for the one segment whose entire purpose is to carry both + branches. + """ + assert gates.topology_flags("unified", world_size=4, member=member) == expected + + +def test_the_member_overrides_the_declared_segment(): + """It is the configuration actually being served, not a hint.""" + assert gates.topology_flags("tp8", world_size=8, member="dep4") == [ + "--tp_size=4", + "--moe_ep_size=4", + "--enable_attention_dp", + ] + + +def test_the_member_reaches_the_verdict(): + """A gate line that does not name the topology cannot be told apart from the members it did not run.""" + result = gates.GateResult(name="accuracy[gsm8k]", passed=True, detail="x", member="dep4") + assert "accuracy[gsm8k]@dep4" in result.summary() + assert "@" not in gates.GateResult(name="boot", passed=True, detail="x").summary() + + +# ------------------------------------------------------------ score parsing + + +def test_lm_eval_score_is_parsed(): + out = "lm-eval gsm8k exact_match,flexible-extract accuracy: 90.67\n" + assert gates.parse_eval_score(out) == pytest.approx(90.67) + + +def test_lm_eval_average_form_is_parsed(): + assert gates.parse_eval_score("lm-eval gsm8k average accuracy: 88.10\n") == pytest.approx(88.10) + + +def test_mmlu_form_is_parsed(): + out = "MMLU weighted average accuracy: 74.98 (14042)\n" + assert gates.parse_eval_score(out) == pytest.approx(74.98) + + +def test_the_final_score_wins_over_partial_estimates(): + """A long run prints running estimates; only the last line is the result.""" + out = ( + "Partial scores after 100/1319 responses (estimate, 0~100): gsm8k,acc ~ 12.00\n" + "Partial scores after 900/1319 responses (estimate, 0~100): gsm8k,acc ~ 80.00\n" + "lm-eval gsm8k exact_match,flexible-extract accuracy: 90.67\n" + ) + assert gates.parse_eval_score(out) == pytest.approx(90.67) + + +def test_no_score_returns_none_rather_than_a_guess(): + assert gates.parse_eval_score("engine built, nothing else\n") is None + + +def test_acceptance_length_is_parsed(): + out = ( + "Spec-dec stats: AL (decoded tokens/step, iteration-weighted) 3.351 " + "(per-request min 1.000, max 4.000, n=1319)\n" + ) + assert gates.parse_acceptance_length(out) == pytest.approx(3.351) + + +def test_acceptance_length_alternate_form_is_parsed(): + assert gates.parse_acceptance_length("acceptance_length: 3.28\n") == pytest.approx(3.28) + + +# --------------------------------------------------------------- reference + + +def _reference_file(tmp_path, entries): + path = tmp_path / "gsm8k.yaml" + path.write_text(yaml.safe_dump({"org/Model": entries}), encoding="utf-8") + return path + + +def test_reference_matches_on_the_exact_discriminator_set(tmp_path): + path = _reference_file( + tmp_path, + [ + {"accuracy": 74.20}, + {"quant_algo": "FP8", "accuracy": 74.30}, + {"quant_algo": "NVFP4", "kv_cache_quant_algo": "FP8", "accuracy": 66.03}, + ], + ) + assert gates.load_reference_score(path, "org/Model") == pytest.approx(74.20) + assert gates.load_reference_score( + path, "org/Model", discriminators={"quant_algo": "FP8"} + ) == pytest.approx(74.30) + assert gates.load_reference_score( + path, + "org/Model", + discriminators={"quant_algo": "NVFP4", "kv_cache_quant_algo": "FP8"}, + ) == pytest.approx(66.03) + + +def test_an_fp8_anchor_is_not_silently_used_for_bf16(tmp_path): + path = _reference_file(tmp_path, [{"quant_algo": "FP8", "accuracy": 74.30}]) + with pytest.raises(KeyError, match="no entry matching"): + gates.load_reference_score(path, "org/Model") + + +def test_a_missing_model_says_the_anchor_is_never_invented(tmp_path): + path = _reference_file(tmp_path, [{"accuracy": 1.0}]) + with pytest.raises(KeyError, match="never invented"): + gates.load_reference_score(path, "org/Other") + + +def test_an_ambiguous_anchor_is_refused(tmp_path): + path = _reference_file(tmp_path, [{"accuracy": 74.20}, {"accuracy": 75.00}]) + with pytest.raises(KeyError, match="unambiguous"): + gates.load_reference_score(path, "org/Model") + + +def test_resolve_reference_joins_a_relative_path_to_the_repo(tmp_path): + repo = tmp_path / "repo" + (repo / "refs").mkdir(parents=True) + (repo / "refs" / "gsm8k.yaml").write_text( + yaml.safe_dump({"org/Model": [{"accuracy": 89.99}]}), encoding="utf-8" + ) + task = {**_TASK, "trtllm_repo_path": str(repo)} + task["accuracy"] = {**_TASK["accuracy"], "reference_yaml": "refs/gsm8k.yaml"} + assert gates.resolve_reference(task, model_key="org/Model") == pytest.approx(89.99) + + +# -------------------------------------------------------------- boot gate + + +def test_boot_gate_passes_on_clean_exit_with_output(): + run = recorder(ShellResult(returncode=0, stdout="Paris is the capital\n")) + result = gates.run_boot_gate(_env(), task_data=_TASK, runner=run) + assert result.passed + assert "quickstart_advanced.py" in result.command + assert "TRTLLM_MODELING_V2=require" in result.command + + +def test_every_gate_asks_the_target_for_its_runtime_value_dump(): + """A target prints its dump only under the switch; the auditor reads it from gate logs.""" + run = recorder(ShellResult(returncode=0, stdout="ok")) + gates.run_boot_gate(_env(), task_data=_TASK, runner=run) + gates.run_accuracy_gate(_env(), task_data=_TASK, reference=90.0, runner=run) + assert run.seen # type: ignore[attr-defined] + for command in run.seen: # type: ignore[attr-defined] + assert f"{gates.RUNTIME_VALUE_DUMP_ENV}=1" in command + + +def test_boot_gate_always_forces_require_mode(): + """Under `auto` the built-in implementation answers and the gate lies.""" + run = recorder(ShellResult(returncode=0, stdout="ok")) + gates.run_boot_gate(_env(), task_data={**_TASK, "parallel": "dep4"}, runner=run) + command = run.seen[0] # type: ignore[attr-defined] + assert "TRTLLM_MODELING_V2=require" in command + assert "--enable_attention_dp" in command + + +def test_boot_gate_fails_on_nonzero_exit(): + run = recorder(ShellResult(returncode=1, stderr="CUDA error")) + result = gates.run_boot_gate(_env(), task_data=_TASK, runner=run) + assert not result.passed and "exited 1" in result.detail + + +def test_boot_gate_fails_on_empty_output(): + run = recorder(ShellResult(returncode=0, stdout=" \n")) + result = gates.run_boot_gate(_env(), task_data=_TASK, runner=run) + assert not result.passed and "no output" in result.detail + + +def test_boot_gate_checks_keywords_when_the_spec_has_them(): + spec = gates.BootSpec(prompts=("x",), keywords=("Paris", "Berlin")) + run = recorder(ShellResult(returncode=0, stdout="Paris only\n")) + result = gates.run_boot_gate(_env(), task_data=_TASK, spec=spec, runner=run) + assert not result.passed and "Berlin" in result.detail + + +def test_boot_spec_defaults_to_no_keywords(): + """Keywords are frozen only after observing what this model actually says.""" + assert gates.BootSpec().keywords == () + + +def test_boot_spec_loads_from_disk(tmp_path): + path = tmp_path / "boot.yaml" + path.write_text(yaml.safe_dump({"prompts": ["a"], "keywords": ["k"], "max_tokens": 8})) + spec = gates.BootSpec.load(path) + assert spec.prompts == ("a",) and spec.keywords == ("k",) and spec.max_tokens == 8 + + +def test_boot_spec_missing_file_is_the_default(tmp_path): + assert gates.BootSpec.load(tmp_path / "absent.yaml") == gates.BootSpec() + + +def test_boot_gate_timeout_says_it_is_no_longer_a_boot_gate(): + run = recorder(ShellResult(returncode=124, timed_out=True)) + result = gates.run_boot_gate(_env(), task_data=_TASK, runner=run, timeout=5) + assert not result.passed and "not catching catastrophes" in result.detail + + +# ----------------------------------------------------------- accuracy gate + + +def test_accuracy_gate_is_one_sided(): + run = recorder(ShellResult(returncode=0, stdout="lm-eval gsm8k average accuracy: 85.50\n")) + result = gates.run_accuracy_gate(_env(), task_data=_TASK, reference=89.99, runner=run) + assert result.passed # 85.50 >= 89.99 - 5.0 + assert result.measured == pytest.approx(85.50) + + +def test_accuracy_gate_fails_below_the_floor(): + run = recorder(ShellResult(returncode=0, stdout="lm-eval gsm8k average accuracy: 80.00\n")) + result = gates.run_accuracy_gate(_env(), task_data=_TASK, reference=89.99, runner=run) + assert not result.passed + assert "80.0000 < 84.9900" in result.detail + + +def test_accuracy_gate_with_no_parsable_score_is_not_a_pass(): + """A gate with no number is not a pass, however clean the exit.""" + run = recorder(ShellResult(returncode=0, stdout="done\n")) + result = gates.run_accuracy_gate(_env(), task_data=_TASK, reference=89.99, runner=run) + assert not result.passed and "no score" in result.detail + + +def test_accuracy_gate_passes_the_dataset_path_when_declared(): + task = {**_TASK, "accuracy": {**_TASK["accuracy"], "dataset_path": "/data/mmlu"}} + run = recorder(ShellResult(returncode=0, stdout="lm-eval gsm8k average accuracy: 99\n")) + gates.run_accuracy_gate(_env(), task_data=task, reference=1.0, runner=run) + assert '--dataset_path="/data/mmlu"' in run.seen[0] # type: ignore[attr-defined] + + +def test_gate_summary_carries_the_numbers_verbatim(): + run = recorder(ShellResult(returncode=0, stdout="lm-eval gsm8k average accuracy: 95.0720\n")) + result = gates.run_accuracy_gate(_env(), task_data=_TASK, reference=89.9962, runner=run) + summary = result.summary() + assert "95.0720" in summary and "89.9962" in summary and "PASS" in summary + + +# --------------------------------------------------------- acceptance gate + + +def test_acceptance_gate_compares_require_against_off(): + run = recorder( + [ + ShellResult(returncode=0, stdout="Spec-dec stats: AL (x) 3.3514 (y)\n"), + ShellResult(returncode=0, stdout="Spec-dec stats: AL (x) 3.2752 (y)\n"), + ] + ) + result = gates.run_acceptance_gate(_env(), task_data=_TASK, runner=run) + assert result.passed + assert result.measured == pytest.approx(3.3514) + assert result.reference == pytest.approx(3.2752) + seen = run.seen # type: ignore[attr-defined] + assert "TRTLLM_MODELING_V2=require" in seen[0] + assert "TRTLLM_MODELING_V2=off" in seen[1] + + +def test_acceptance_of_one_means_every_draft_was_rejected(): + run = recorder( + [ + ShellResult(returncode=0, stdout="Spec-dec stats: AL (x) 1.0000 (y)\n"), + ShellResult(returncode=0, stdout="Spec-dec stats: AL (x) 3.2752 (y)\n"), + ] + ) + result = gates.run_acceptance_gate(_env(), task_data=_TASK, runner=run) + assert not result.passed and "every draft was rejected" in result.detail + + +def test_acceptance_below_the_reference_fails(): + """Clearly above 1.0 but below the reference means subtly wrong, not slower.""" + run = recorder( + [ + ShellResult(returncode=0, stdout="Spec-dec stats: AL (x) 2.0000 (y)\n"), + ShellResult(returncode=0, stdout="Spec-dec stats: AL (x) 3.2752 (y)\n"), + ] + ) + result = gates.run_acceptance_gate(_env(), task_data=_TASK, runner=run) + assert not result.passed + + +# ------------------------------------------------ the acceptance workload + +# Measured on DeepSeek-V4.1 with DSpark on, tep4, the BOOT prompts at 256 +# tokens: every answer is 2-3 tokens then EOS, so the rate is 1.00 for any +# drafter (05-units/20260923-164756-dspark-preassembly.md, §4c). +_BOOT_WORKLOAD_OUTPUT = ( + "[0] Accept rate: 1.00 (tokens=2, iterations=2)\n" + "[1] Accept rate: 1.00 (tokens=2, iterations=2)\n" + "[2] Accept rate: 1.00 (tokens=3, iterations=3)\n" +) +# The same command on three long-answer prompts, after the drafter fix. +_LONG_WORKLOAD_OUTPUT = ( + "[0] Accept rate: 3.45 (tokens=152, iterations=44)\n" + "[1] Accept rate: 2.25 (tokens=256, iterations=114)\n" + "[2] Accept rate: 5.69 (tokens=239, iterations=42)\n" +) +_LONG_PROMPTS = [ + "Explain step by step how a refrigerator moves heat.", + 'Write a Python module that parses "key = $value" lines, with tests.', +] + + +def _no_stock_runner(output): + """`=require` prints `output`; `=off` cannot load a new architecture.""" + seen: list[str] = [] + + def run(command, env=None, timeout=0): + seen.append(command) + if "TRTLLM_MODELING_V2=off" in command: + return ShellResult( + returncode=1, + stderr="RuntimeError: Unknown model architecture: NewForCausalLM", + ) + return ShellResult(returncode=0, stdout=output) + + run.seen = seen # type: ignore[attr-defined] + return run + + +def _acceptance_task(**acceptance): + return { + **_TASK, + "accuracy": { + **_TASK["accuracy"], + "acceptance": {"dspark": {"min_absolute": 2.0}, **acceptance}, + }, + } + + +def test_a_short_answer_workload_fails_as_too_short_not_as_a_low_number(): + """The F029 case: the boot prompts read 1.00 whatever the drafter does. + + Reported as a floor miss, 1.00 against 2.0 reads as a broken drafter and + sends the assembler after a defect the target does not have. + """ + result = gates.run_acceptance_gate( + _env(), + task_data=_acceptance_task(), + runner=_no_stock_runner(_BOOT_WORKLOAD_OUTPUT), + unit="dspark", + ) + assert not result.passed + assert "workload too short to measure acceptance" in result.detail + assert "4 generation iteration(s) over 3 prompt(s)" in result.detail + assert "the boot prompts" in result.detail + # No number is reported as the target's: there was no measurement. + assert result.measured is None + assert "[2] 1.00 (3 tok/3 it)" in result.summary() + + +def test_the_two_leg_comparison_is_held_to_the_same_minimum(): + run = recorder( + [ + ShellResult(returncode=0, stdout=_LONG_WORKLOAD_OUTPUT), + ShellResult(returncode=0, stdout=_BOOT_WORKLOAD_OUTPUT), + ] + ) + result = gates.run_acceptance_gate(_env(), task_data=_TASK, runner=run) + assert not result.passed + assert "workload too short" in result.detail + assert "TRTLLM_MODELING_V2=off leg" in result.detail + + +def test_per_prompt_acceptance_is_reported_beside_the_mean(): + """The verdict is on the iteration-weighted rate; the journal shows every prompt.""" + result = gates.run_acceptance_gate( + _env(), + task_data=_acceptance_task(prompts=_LONG_PROMPTS), + runner=_no_stock_runner(_LONG_WORKLOAD_OUTPUT), + unit="dspark", + ) + assert result.passed, result.detail + assert result.measured == pytest.approx(647 / 200) + summary = result.summary() + for expected in ("[0] 3.45 (152 tok/44 it)", "[1] 2.25", "[2] 5.69", "prompt mean 3.80"): + assert expected in summary, summary + assert "197 generation iterations" in summary + + +def test_a_declared_workload_replaces_the_boot_prompts(): + run = _no_stock_runner(_LONG_WORKLOAD_OUTPUT) + gates.run_acceptance_gate( + _env(), + task_data=_acceptance_task(prompts=_LONG_PROMPTS, max_tokens=300), + runner=run, + unit="dspark", + spec=gates.BootSpec(prompts=("The capital of France is",)), + ) + command = run.seen[0] # type: ignore[attr-defined] + assert "--max_tokens=300" in command + assert "capital of France" not in command + # Shell-quoted, so a prompt's own quotes and `$` reach quickstart intact. + words = shlex.split(command) + assert words[words.index("--prompt") + 1 :] == _LONG_PROMPTS + + +def test_no_declared_workload_keeps_the_boot_prompts(): + """Backward compatible: a task that declares nothing measures what it always did.""" + spec = gates.BootSpec(prompts=("The capital of France is",)) + workload = gates.resolve_acceptance_workload( + {"acceptance": {"mtp": {"min_absolute": 1.5}}}, "mtp", fallback=spec + ) + assert workload.prompts == spec.prompts + assert workload.max_tokens == gates.DEFAULT_ACCEPTANCE_MAX_TOKENS == 256 + assert workload.min_draft_iterations == gates.DEFAULT_MIN_DRAFT_ITERATIONS + assert "boot prompts" in workload.source + assert gates.resolve_acceptance_workload({}, None, fallback=spec).prompts == spec.prompts + + +def test_a_unit_entry_overrides_the_run_wide_workload(): + accuracy = { + "acceptance": { + "prompts": ["run-wide"], + "max_tokens": 256, + "min_draft_iterations": 64, + "dspark": {"min_absolute": 2.0}, + "multimodal": {"min_absolute": 2.0, "prompts": ["image-side"], "max_tokens": 128}, + } + } + dspark = gates.resolve_acceptance_workload(accuracy, "dspark") + assert dspark.prompts == ("run-wide",) and dspark.max_tokens == 256 + assert dspark.min_draft_iterations == 64 + assert dspark.source == "accuracy.acceptance.prompts" + multimodal = gates.resolve_acceptance_workload(accuracy, "multimodal") + assert multimodal.prompts == ("image-side",) and multimodal.max_tokens == 128 + # A key the unit does not state still comes from the run-wide level. + assert multimodal.min_draft_iterations == 64 + + +def test_a_prompts_file_is_read_in_either_shape(tmp_path): + as_list = tmp_path / "prompts.json" + as_list.write_text(json.dumps(["a long question", "another"]), encoding="utf-8") + as_boot = tmp_path / "boot.yaml" + as_boot.write_text(yaml.safe_dump({"prompts": ["from boot"], "keywords": ["x"]})) + for path, expected in ((as_list, ("a long question", "another")), (as_boot, ("from boot",))): + workload = gates.resolve_acceptance_workload( + {"acceptance": {"prompts_file": str(path)}}, "dspark" + ) + assert workload.prompts == expected + assert str(path) in workload.source + + +def test_an_unreadable_prompts_file_fails_the_gate_before_it_runs(tmp_path): + run = _no_stock_runner(_LONG_WORKLOAD_OUTPUT) + result = gates.run_acceptance_gate( + _env(), + task_data=_acceptance_task(prompts_file=str(tmp_path / "gone.json")), + runner=run, + unit="dspark", + ) + assert not result.passed + assert "declared acceptance workload cannot be read" in result.detail + assert run.seen == [] # type: ignore[attr-defined] + + +def test_a_declared_minimum_is_the_one_applied(): + result = gates.run_acceptance_gate( + _env(), + task_data=_acceptance_task(prompts=_LONG_PROMPTS, min_draft_iterations=500), + runner=_no_stock_runner(_LONG_WORKLOAD_OUTPUT), + unit="dspark", + ) + assert not result.passed + assert "197 generation iteration(s)" in result.detail and "500" in result.detail + + +def test_a_long_workload_with_a_dead_drafter_still_fails_the_floor(): + """Long answers are what let the floor see the drafter at all.""" + dead = ( + "[0] Accept rate: 1.00 (tokens=256, iterations=256)\n" + "[1] Accept rate: 1.00 (tokens=256, iterations=256)\n" + ) + result = gates.run_acceptance_gate( + _env(), + task_data=_acceptance_task(prompts=_LONG_PROMPTS), + runner=_no_stock_runner(dead), + unit="dspark", + ) + assert not result.passed + assert "too short" not in result.detail + assert result.measured == pytest.approx(1.0) and result.reference == pytest.approx(2.0) + + +# ------------------------------------------------------------------- log + + +def test_gate_log_keeps_order_and_finds_the_latest(): + log = gates.GateLog() + log.record(gates.GateResult(name="boot", passed=True, detail="a")) + log.record(gates.GateResult(name="accuracy[gsm8k]", passed=False, detail="b")) + log.record(gates.GateResult(name="boot", passed=False, detail="c")) + assert log.last("boot").detail == "c" + assert log.last("accuracy").passed is False + assert log.last("acceptance") is None + assert len(log.render().splitlines()) == 3 + + +def test_the_two_gate_commands_spell_expert_parallelism_differently(): + """`--moe_ep_size` for quickstart, `--ep_size` for trtllm-eval. + + Measured on the first `tep4` accuracy gate of a real run:: + + Error: No such option: --moe_ep_size + (Possible options: --ep_size, --pp_size, --tp_size) + + `tp4` had passed the same gate minutes earlier, because it emits only + `--tp_size`, which both commands accept — so the defect was invisible at + the one member that does not exercise it. + """ + from agent_flow.workflows.staircase_onboard.gates import EVAL_DIALECT, topology_flags + + assert topology_flags("unified", 4, "tep4") == ["--tp_size=4", "--moe_ep_size=4"] + assert topology_flags("unified", 4, "tep4", dialect=EVAL_DIALECT) == [ + "--tp_size=4", + "--ep_size=4", + ] + + +def test_attention_dp_is_not_a_trtllm_eval_flag(): + """`trtllm-eval --help` lists tp/pp/ep and nothing for attention DP. + + So a `dep` member cannot carry it on the argv, and emitting it there + reintroduces the failure above. It travels in the options overlay + instead — and the overlay must not be empty, or the gate would serve + `tep` while the verdict claimed `dep`. + """ + from agent_flow.workflows.staircase_onboard.gates import ( + EVAL_DIALECT, + member_llm_api_overlay, + topology_flags, + ) + + eval_flags = topology_flags("unified", 4, "dep4", dialect=EVAL_DIALECT) + assert "--enable_attention_dp" not in eval_flags + assert eval_flags == ["--tp_size=4", "--ep_size=4"] + assert member_llm_api_overlay("dep4") == {"enable_attention_dp": True} + # The quickstart dialect keeps the flag, because that command takes it. + assert "--enable_attention_dp" in topology_flags("unified", 4, "dep4") + + +def test_only_dep_members_need_an_overlay(): + """An overlay that fires for every member would be a different bug.""" + from agent_flow.workflows.staircase_onboard.gates import member_llm_api_overlay + + assert member_llm_api_overlay("tp4") == {} + assert member_llm_api_overlay("tep4") == {} + assert member_llm_api_overlay(None) == {} + + +def test_the_draft_depth_comes_from_the_options_yaml(tmp_path): + """All three gates must ask for the depth the unit was built for. + + `trtllm-eval` takes no speculation flags, so the accuracy gate reads + `speculative_config.num_nextn_predict_layers` from the unit's options + YAML. The boot and acceptance gates drive `quickstart_advanced.py`, which + cannot read that file and takes `--spec_decode_max_draft_len` instead -- + hard-defaulted to 1 before this. + + Measured on a real unit: the YAML said 3 and the flag said 1, so + acceptance would have measured MTP-1 against a ceiling of 2.0 while the + target's real acceptance is 2.44-2.59 against a ceiling of 4.0. It fails + toward green, because a shallower speculator accepts a larger fraction of + what it drafts. + """ + import yaml as _yaml + + from agent_flow.workflows.staircase_onboard.gates import ( + resolve_max_draft_len, + speculation_flags, + ) + + options = tmp_path / "mtp_llm_api.yaml" + options.write_text( + _yaml.safe_dump( + {"speculative_config": {"decoding_type": "MTP", "num_nextn_predict_layers": 3}} + ) + ) + accuracy = {"llm_api_options": {"mtp": str(options)}} + + assert resolve_max_draft_len(accuracy, "mtp") == 3 + assert speculation_flags(["mtp"], max_draft_len=3) == [ + "--spec_decode_algo=MTP", + "--spec_decode_max_draft_len=3", + ] + + +def test_a_unit_with_no_speculator_still_answers_one(tmp_path): + """The old answer has to survive for every unit that is not a speculator.""" + import yaml as _yaml + + from agent_flow.workflows.staircase_onboard.gates import resolve_max_draft_len + + assert resolve_max_draft_len({}, "core") == 1 + assert resolve_max_draft_len({"llm_api_options": "/nonexistent.yaml"}, "core") == 1 + + plain = tmp_path / "core_llm_api.yaml" + plain.write_text(_yaml.safe_dump({"enable_chunked_prefill": True})) + assert resolve_max_draft_len({"llm_api_options": str(plain)}, "core") == 1 + + +def test_the_off_leg_can_be_impossible_rather_than_broken(tmp_path): + """A new architecture has no built-in modeling to compare a speculator to. + + The acceptance gate's `=off` leg runs the same checkpoint under stock + in-tree modeling. For a genuinely new architecture that path cannot load + the checkpoint at all — so the gate is unavailable exactly when onboarding + is doing its job, which is the substrate gate's problem one gate down. + + Skipping it would be worse than failing: it is the ONLY gate that can see + a miscomputed draft path, because rejection sampling makes a wrong draft + produce correct text more slowly. So a declared absolute floor stands in, + and the weaker claim is recorded. + """ + import yaml as _yaml + + from agent_flow.workflows.staircase_onboard.gates import resolve_acceptance_floor + + options = tmp_path / "mtp.yaml" + options.write_text(_yaml.safe_dump({"speculative_config": {"num_nextn_predict_layers": 3}})) + accuracy = { + "llm_api_options": {"mtp": str(options)}, + "acceptance": {"mtp": {"min_absolute": 1.5}}, + } + assert resolve_acceptance_floor(accuracy, "mtp") == 1.5 + # A unit that declares none keeps the two-leg comparison and its hard fail. + assert resolve_acceptance_floor(accuracy, "core") is None + assert resolve_acceptance_floor({}, "mtp") is None + + +def test_the_unsupported_architecture_markers_are_shared_with_preflight(): + """One definition of "the stock path does not know this model". + + The substrate gate and the acceptance gate ask the same question for the + same reason, so they must not drift apart. + """ + from agent_flow.workflows.staircase_onboard.preflight import reports_unsupported_architecture + + assert reports_unsupported_architecture( + "exited 1", "Unknown architecture for AutoModelForCausalLM" + ) + assert reports_unsupported_architecture( + "exited 1", + "ValueError: The repository /ckpt contains custom code which must be " + "executed to correctly load the model.", + ) + # A real fault must still be a hard stop, not a fallback. + assert not reports_unsupported_architecture( + "exited 1", "Impossible to fit in any sequence in kvCache" + ) + + +def test_the_registry_miss_is_also_an_unsupported_architecture(): + """`get_model_architecture` raises its own message, by its own path. + + `Unknown model architecture:` comes from the registry lookup in + `_torch/models/modeling_utils.py`, not from the `AutoModelFor*` config + loaders. A checkpoint can fail there without matching any other marker -- + and for the acceptance gate that is the difference between falling back to + a declared floor and hard-failing a working unit. + """ + from agent_flow.workflows.staircase_onboard.preflight import reports_unsupported_architecture + + assert reports_unsupported_architecture( + "exited 1", "RuntimeError: Unknown model architecture: MiMoV2ForCausalLM" + ) + + +# ------------------------------------------------- the tree a gate measured + + +def _git(repo, *args): + """Run one git command in ``repo``, with an identity the host may not have.""" + return subprocess.run( + [ + "git", + "-c", + "user.name=staircase", + "-c", + "user.email=staircase@example.invalid", + "-c", + "commit.gpgsign=false", + "-C", + str(repo), + *args, + ], + capture_output=True, + text=True, + check=True, + ).stdout.strip() + + +def _repo_with_target(tmp_path, *, family="gpt_oss"): + """A real one-commit repo holding a modeling_v2 family directory. + + Real git rather than a mocked one on purpose: what this guard has to get + right is the exact output of `git status --porcelain`, which is the kind + of thing a stub agrees with and the tool does not. + """ + repo = tmp_path / "repo" + family_dir = repo / gates.MODELING_V2_PACKAGE / "models" / family + family_dir.mkdir(parents=True) + (family_dir / "routing.py").write_text("# committed\n", encoding="utf-8") + (repo / "docs").mkdir() + (repo / "docs" / "note.md").write_text("committed\n", encoding="utf-8") + _git(repo, "init", "-q", ".") + _git(repo, "add", "-A") + _git(repo, "commit", "-q", "-m", "target") + return repo + + +def _task_for(repo, family="gpt_oss"): + return {**_TASK, "trtllm_repo_path": str(repo), "target": {"family": family}} + + +def test_the_tree_state_reaches_the_summary(): + """`progress.yaml` and the journal read this line; attribution rides on it.""" + plain = gates.GateResult(name="boot", passed=True, detail="ok", tree_state="v1-3-gabc1234") + assert "[tree v1-3-gabc1234]" in plain.summary() + numeric = gates.GateResult( + name="accuracy[gsm8k]", + passed=True, + detail="ok", + measured=90.0, + reference=89.0, + tolerance=5.0, + tree_state="abc1234", + ) + assert "[tree abc1234]" in numeric.summary() + + +def test_a_gate_records_the_commit_it_measured(tmp_path): + repo = _repo_with_target(tmp_path) + run = recorder(ShellResult(returncode=0, stdout="Paris\n")) + result = gates.run_boot_gate(_env(), task_data=_task_for(repo), runner=run) + assert result.passed + assert result.tree_state == _git(repo, "describe", "--always") + assert result.tree_state in result.summary() + + +def test_the_accuracy_and_acceptance_gates_record_it_too(tmp_path): + """A number is attributable or it is not; which gate produced it is not the question.""" + repo = _repo_with_target(tmp_path) + expected = _git(repo, "describe", "--always") + accuracy = gates.run_accuracy_gate( + _env(), + task_data=_task_for(repo), + reference=89.99, + runner=recorder(ShellResult(returncode=0, stdout="lm-eval gsm8k average accuracy: 90\n")), + ) + acceptance = gates.run_acceptance_gate( + _env(), + task_data=_task_for(repo), + runner=recorder( + [ + ShellResult(returncode=0, stdout="Spec-dec stats: AL (x) 3.3514 (y)\n"), + ShellResult(returncode=0, stdout="Spec-dec stats: AL (x) 3.2752 (y)\n"), + ] + ), + ) + assert accuracy.tree_state == expected + assert acceptance.tree_state == expected + + +def test_a_gate_still_runs_where_git_cannot_answer(tmp_path, monkeypatch): + """No repository, no commits, no git binary — none of them is a modeling defect.""" + # A TMPDIR inside a checkout (a short one is needed for the RPC socket + # path) would otherwise let git walk up into that checkout and answer. + monkeypatch.setenv("GIT_CEILING_DIRECTORIES", str(tmp_path.parent)) + run = recorder(ShellResult(returncode=0, stdout="Paris\n")) + result = gates.run_boot_gate( + _env(), task_data={**_TASK, "trtllm_repo_path": str(tmp_path)}, runner=run + ) + assert result.passed + assert result.tree_state == gates.UNKNOWN_TREE_STATE + + +def test_git_missing_entirely_does_not_fail_a_gate(tmp_path, monkeypatch): + """The host with no `git` on PATH. A gate adjudicates a model, not a toolchain.""" + repo = _repo_with_target(tmp_path) + + def no_git(*args, **kwargs): + raise FileNotFoundError("git") + + monkeypatch.setattr(gates.subprocess, "run", no_git) + run = recorder(ShellResult(returncode=0, stdout="Paris\n")) + result = gates.run_boot_gate(_env(), task_data=_task_for(repo), runner=run) + assert result.passed and result.tree_state == gates.UNKNOWN_TREE_STATE + + +# --------------------------------------------------- the dirty-target guard + + +def test_a_dirty_target_path_stops_the_gate_before_it_spends_anything(tmp_path): + """At gate time the assembler has already committed; dirty means another writer.""" + repo = _repo_with_target(tmp_path) + modeling = repo / gates.MODELING_V2_PACKAGE / "models" / "gpt_oss" / "modeling.py" + modeling.write_text("# half-written by somebody else\n", encoding="utf-8") + + run = recorder(ShellResult(returncode=0, stdout="lm-eval gsm8k average accuracy: 90\n")) + result = gates.run_accuracy_gate(_env(), task_data=_task_for(repo), reference=89.99, runner=run) + assert not result.passed + assert "modeling.py" in result.detail + # And it cost nothing: the guard is in front of the command, not after it. + assert run.seen == [] # type: ignore[attr-defined] + + +def test_every_gate_refuses_a_dirty_target_path(tmp_path): + repo = _repo_with_target(tmp_path) + (repo / gates.MODELING_V2_PACKAGE / "models" / "gpt_oss" / "weights.py").write_text( + "# in flight\n", encoding="utf-8" + ) + task = _task_for(repo) + boot = gates.run_boot_gate(_env(), task_data=task, runner=recorder([])) + acceptance = gates.run_acceptance_gate(_env(), task_data=task, runner=recorder([])) + assert not boot.passed and "weights.py" in boot.detail + assert not acceptance.passed and "weights.py" in acceptance.detail + + +def test_the_guard_is_scoped_to_the_target_not_to_the_repo(tmp_path): + """A run has records and workflow fixes in flight constantly. + + A repo-wide check would refuse every gate of every run, which is a guard + nobody can leave switched on. + """ + repo = _repo_with_target(tmp_path) + (repo / "docs" / "note.md").write_text("edited while the gate runs\n", encoding="utf-8") + (repo / "07-workflow-feedback.md").write_text("new record\n", encoding="utf-8") + run = recorder(ShellResult(returncode=0, stdout="Paris\n")) + result = gates.run_boot_gate(_env(), task_data=_task_for(repo), runner=run) + assert result.passed and not result.dirty_override + + +def test_the_escape_hatch_runs_the_gate_and_marks_the_number(tmp_path): + repo = _repo_with_target(tmp_path) + (repo / gates.MODELING_V2_PACKAGE / "models" / "gpt_oss" / "modeling.py").write_text( + "# debugging by hand\n", encoding="utf-8" + ) + run = recorder(ShellResult(returncode=0, stdout="Paris\n")) + result = gates.run_boot_gate(_env(), task_data=_task_for(repo), runner=run, allow_dirty=True) + assert result.passed and result.dirty_override + assert "DIRTY TARGET PATH" in result.summary() + assert result.tree_state.endswith("-dirty") + + +def test_the_escape_hatch_can_come_from_the_environment(tmp_path, monkeypatch): + repo = _repo_with_target(tmp_path) + (repo / gates.MODELING_V2_PACKAGE / "models" / "gpt_oss" / "modeling.py").write_text( + "# debugging by hand\n", encoding="utf-8" + ) + monkeypatch.setenv(gates.ALLOW_DIRTY_ENV, "1") + result = gates.run_boot_gate( + _env(), task_data=_task_for(repo), runner=recorder(ShellResult(returncode=0, stdout="x\n")) + ) + assert result.passed and result.dirty_override + + +def test_a_clean_target_is_not_marked_as_overridden(tmp_path, monkeypatch): + """The override marks a number measured over someone's edits, not every number.""" + monkeypatch.setenv(gates.ALLOW_DIRTY_ENV, "1") + repo = _repo_with_target(tmp_path) + result = gates.run_boot_gate( + _env(), task_data=_task_for(repo), runner=recorder(ShellResult(returncode=0, stdout="x\n")) + ) + assert result.passed and not result.dirty_override + assert "DIRTY" not in result.summary() + + +def test_the_family_path_follows_the_checkpoint_when_nothing_overrides_it(tmp_path): + """Derived the way preflight derives it, so the guard watches what is written.""" + checkpoint = tmp_path / "ckpt" + checkpoint.mkdir() + (checkpoint / "config.json").write_text( + json.dumps({"architectures": ["GptOssForCausalLM"]}), encoding="utf-8" + ) + task = {**_TASK, "checkpoint_path": str(checkpoint)} + assert gates.target_family_rel(task) == f"{gates.MODELING_V2_PACKAGE}/models/gpt_oss" + assert ( + gates.target_family_rel({**task, "target": {"family": "mimo_v2"}}) + == f"{gates.MODELING_V2_PACKAGE}/models/mimo_v2" + ) + + +def test_an_underivable_family_guards_nothing_rather_than_guessing(tmp_path): + assert gates.target_family_rel({"checkpoint_path": str(tmp_path / "absent")}) is None + + +# ---------------------------------------------------------- member ordering + + +def test_members_are_ordered_cheapest_first(): + """A unit is gated at every member and a broken unit fails at all of them.""" + assert gates.order_members_by_cost(["tp4", "dep4", "tep4"]) == ["tep4", "dep4", "tp4"] + + +def test_the_order_is_a_function_of_the_member_set_alone(): + """Two runs of the same unit must gate in the same sequence, or the journals stop comparing line by line.""" + orderings = { + tuple(gates.order_members_by_cost(list(permutation))) + for permutation in itertools.permutations(["tp4", "tep4", "dep4"]) + } + assert orderings == {("tep4", "dep4", "tp4")} + + +def test_a_larger_world_size_breaks_the_tie(): + assert gates.order_members_by_cost(["tep8", "tep2", "tep4"]) == ["tep2", "tep4", "tep8"] + + +def test_an_unrecognized_member_sorts_last_and_does_not_raise(): + """Including `unified`, which names a family rather than a configuration. + + This function decides a running order; a gate loop that died while + choosing one would have failed worse than one that ran in a poor order. + """ + ordered = gates.order_members_by_cost(["unified", "nonsense", "tp4", "tep4"]) + assert ordered[:2] == ["tep4", "tp4"] + assert set(ordered[2:]) == {"unified", "nonsense"} + + +def test_ordering_keeps_every_member(): + members = ["tp4", "tep4", "dep4", "tep4"] + assert sorted(gates.order_members_by_cost(members)) == sorted(members) + + +def test_the_degenerate_member_set_is_returned_unchanged(): + """A concrete segment is gated at itself, once.""" + assert gates.order_members_by_cost(["tp8"]) == ["tp8"] + assert gates.order_members_by_cost([]) == [] + + +def test_a_member_overlay_never_lands_in_the_source_tree(tmp_path, monkeypatch): + """A generated file in the repo is indistinguishable from a checked-in one. + + With no unit options file the derived copy has nowhere natural to go, and + the cwd is the one place it must not be: the cwd of a test run is the + repository, and this wrote `core.dep1.llm_api.yaml` into it. + """ + from pathlib import Path + + from agent_flow.workflows.staircase_onboard.gates import _write_member_options + + monkeypatch.chdir(tmp_path) + written = Path( + _write_member_options(None, {"enable_attention_dp": True}, unit="core", member="dep4") + ) + assert written.is_file() + assert written.parent != tmp_path, "the overlay landed in the working directory" diff --git a/agent-flow/tests/workflows/staircase_onboard/test_kernel_sources.py b/agent-flow/tests/workflows/staircase_onboard/test_kernel_sources.py new file mode 100644 index 000000000000..eae0880b5f70 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_kernel_sources.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The kernel-source registry, held to the tree it describes. + +Same contract as the runtime map: pointers resolve, no transcribed code, and +the claims that matter are checked rather than trusted. The one addition is a +coverage check in both directions --- a registration file the registry does +not know about is a kernel source an agent will meet and cannot classify. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + +from agent_flow.workflows.staircase_onboard.runtime import ( + KERNEL_SOURCES_PATH, + RuntimeIndexError, + format_kernel_sources, + load_kernel_sources, + source_for_stem, +) + +REPO_ROOT = Path(__file__).resolve().parents[4] +CUSTOM_OPS = REPO_ROOT / "tensorrt_llm/_torch/custom_ops" +in_checkout = pytest.mark.skipif( + not CUSTOM_OPS.is_dir(), reason="not inside a TensorRT-LLM checkout" +) + +#: Registration mechanisms rather than kernel sources. An op attributed to +#: one of these has a real provenance elsewhere, and the registry says so in +#: prose rather than inventing an entry for it. +NOT_A_SOURCE = {"__init__", "fast_custom_op", "cutedsl_matmul_heuristics"} + + +@pytest.fixture(scope="module") +def sources(): + return load_kernel_sources() + + +def test_every_source_declares_what_certifying_one_requires(sources): + """Provenance decides the shape of a receipt; that is the point of the file.""" + for source in sources.values(): + assert source.tier in (1, 2, 3, 4), f"{source.name}: tier={source.tier}" + assert len(source.availability.split()) >= 5, f"{source.name}: availability too thin" + assert len(source.certifying.split()) >= 5, f"{source.name}: certifying too thin" + assert source.pointers, f"{source.name}: points nowhere" + + +@in_checkout +def test_every_path_resolves(sources): + missing = [ + f"{source.name}: {path}" + for source in sources.values() + for path in source.pointers + if not (REPO_ROOT / path).exists() + ] + assert not missing, "kernel_sources.yaml points at paths that do not exist:\n " + "\n ".join( + missing + ) + + +@in_checkout +def test_every_registration_file_in_the_tree_is_classified(sources): + """The direction that actually rots. + + A new `custom_ops/_custom_ops.py` lands and nothing tells an + agent what kind of kernel it holds. This fails when that happens. + """ + known = {Path(path).stem for source in sources.values() for path in source.registers} + unclassified = sorted( + py.stem + for py in CUSTOM_OPS.glob("*.py") + if py.stem not in known and py.stem not in NOT_A_SOURCE + ) + assert not unclassified, ( + "custom_ops files no kernel source claims --- an agent meeting an op " + f"from one of these cannot classify it: {unclassified}" + ) + + +@in_checkout +def test_the_two_registration_planes_are_both_covered(sources): + """Missing either plane makes attribution silently wrong for ~a third of ops.""" + registers = {p for s in sources.values() for p in s.registers} + assert any("custom_ops" in p for p in registers), "the Python plane is unrepresented" + assert any("cpp/tensorrt_llm/thop" in p for p in registers), "the C++ plane is unrepresented" + + +@in_checkout +def test_attribution_resolves_real_catalog_ops_to_real_sources(sources): + """End to end against the ops the catalog actually wraps. + + This is the join the smith phase depends on: an op name, discovered on + the host, resolved to the provenance that decides its receipt. + """ + import yaml + + index = REPO_ROOT / "tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml" + if not index.is_file(): + pytest.skip("catalog index not present") + + registrations = {} + for py in sorted(CUSTOM_OPS.glob("*.py")): + for name in re.findall(r'"trtllm::([a-zA-Z0-9_]+)', py.read_text(encoding="utf-8")): + registrations.setdefault(name, py.stem) + + entries = yaml.safe_load(index.read_text(encoding="utf-8"))["entries"] + ops = [ + e["impl"].rsplit(".", 1)[-1] + for e in entries + if str(e.get("impl", "")).startswith("torch.ops.trtllm.") + ] + assert ops, "no trtllm ops in the catalog index" + + unresolved = [] + for op in ops: + stem = registrations.get(op, "cpp_thop") + if stem in NOT_A_SOURCE: + continue # a registration mechanism; provenance is elsewhere + if source_for_stem(stem, sources) is None: + unresolved.append(f"{op} (registered by {stem})") + assert not unresolved, "catalog ops whose provenance cannot be resolved:\n " + "\n ".join( + unresolved + ) + + +def test_the_cpp_plane_is_the_conclusion_for_an_unclaimed_op(sources): + """Not a fallback: an op no Python file registers IS a C++ op.""" + assert source_for_stem("cpp_thop", sources) is not None + assert source_for_stem("cpp_custom_ops", sources).name == "cpp_thop" + + +def test_flashinfer_is_marked_as_the_conditional_one(sources): + """The only source whose absence is an environment fact, not a build one.""" + availability = sources["flashinfer"].availability + assert "IS_FLASHINFER_AVAILABLE" in availability + assert "version" in sources["flashinfer"].certifying # no version pins + + +def test_vendored_triton_is_outside_the_op_plane(sources): + """It is called as Python, so an op inventory cannot see it at all.""" + triton = sources["vendored_triton"] + assert not triton.registers + assert triton.tier == 2 + assert "invisible to an op" in " ".join(triton.summary.split()) + + +def test_the_registry_carries_no_code(): + text = KERNEL_SOURCES_PATH.read_text(encoding="utf-8") + assert "```" not in text + from .test_skills import BARE_LINE_REFERENCE, LINE_CITATION + + assert not LINE_CITATION.findall(text) + BARE_LINE_REFERENCE.findall(text) + + +def test_schema_rejects_a_source_that_points_nowhere(tmp_path): + bad = tmp_path / "k.yaml" + bad.write_text( + "x:\n summary: s\n tier: 1\n availability: a b c d e\n certifying: a b c d e\n", + encoding="utf-8", + ) + with pytest.raises(RuntimeIndexError, match="at least one of"): + load_kernel_sources(bad) + + +def test_the_rendered_block_explains_why_provenance_matters(sources): + block = format_kernel_sources(sources) + assert "registration planes" in block + assert "certifying one" in block + assert "tier 2" in block # vendored_triton renders its tier + + +def test_the_registry_reaches_the_smith(tmp_path): + """The role that decides where a computation comes from. + + A smith is dispatched with a tier; the registry turns "tier 1" into "and + here is what certifying a CuTe DSL kernel actually requires". + """ + from agent_flow.workflows.staircase_onboard.vocabulary import QueueItem + + from .test_workflow import _workflow + + workflow = _workflow(tmp_path) + workflow.task_data.update({"trtllm_repo_path": "/repo", "arch": "sm_103"}) + workflow._kernel_sources_block = format_kernel_sources(load_kernel_sources()) + prompt = workflow._smith_prompt(QueueItem(n=1, mode="onboard", target="x", tier=1)) + + assert "Where kernels come from" in prompt + assert "cute_dsl" in prompt and "flashinfer" in prompt + assert "registration planes" in prompt + + +def test_a_run_without_a_survey_still_dispatches(tmp_path): + """The block defaults empty; a smith must not depend on it existing.""" + from agent_flow.workflows.staircase_onboard.vocabulary import QueueItem + + from .test_workflow import _workflow + + workflow = _workflow(tmp_path) + workflow.task_data.update({"trtllm_repo_path": "/repo", "arch": "sm_103"}) + assert workflow._smith_prompt(QueueItem(n=1, mode="onboard", target="x", tier=1)) diff --git a/agent-flow/tests/workflows/staircase_onboard/test_macro_expansion.py b/agent-flow/tests/workflows/staircase_onboard/test_macro_expansion.py new file mode 100644 index 000000000000..18c3d1c19db4 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_macro_expansion.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for `MACRO` — the outcome that says what the next work is. + +`SMITH_DECISIONS` has always included `MACRO`, and the orchestrator has +always scored it as a failure. On the first real onboarding the largest +dispatch returned MACRO with an eleven-item decomposition, and the run +closed `smith-blocked` while holding exactly the queue extension it needed. + +The fix is two-sided: the decomposition has to be machine-readable, and the +orchestrator has to act on it. These tests hold both sides. +""" + +from __future__ import annotations + +import yaml + +from agent_flow.workflows.staircase_onboard.progress import ( + ProgressContext, + collect_decision_file, + decision_path, + init_progress_file, + latest_entry, + parse_constituents, +) +from agent_flow.workflows.staircase_onboard.state import WorkflowState +from agent_flow.workflows.staircase_onboard.vocabulary import QueueItem +from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + +# ------------------------------------------------------- parsing the list + + +def test_a_constituent_needs_a_target_and_gets_a_default_tier(): + rows = parse_constituents([{"target": "causal_conv1d_fwd"}]) + assert rows == [{"target": "causal_conv1d_fwd", "tier": 1}] + + +def test_a_constituent_without_a_target_is_dropped(): + """It cannot be dispatched, and guessing a name is worse than dropping it.""" + assert parse_constituents([{"tier": 2, "note": "something"}]) == [] + + +def test_a_non_list_decomposition_is_not_a_crash(): + """A smith writing YAML by hand should not take the run down with it.""" + assert parse_constituents("conv, then the delta rule") == [] + assert parse_constituents(None) == [] + + +def test_optional_fields_survive_and_unknown_ones_do_not(): + rows = parse_constituents( + [{"target": "x", "tier": 3, "unit": "mtp", "surface": "P>1", "colour": "red"}] + ) + assert rows[0] == {"target": "x", "tier": 3, "unit": "mtp", "surface": "P>1"} + + +def test_a_bad_tier_falls_back_rather_than_raising(): + assert parse_constituents([{"target": "x", "tier": "two"}])[0]["tier"] == 1 + + +def test_constituents_round_trip_through_a_decision_file(tmp_path): + progress = tmp_path / "progress.yaml" + init_progress_file(progress) + ctx = ProgressContext(path=progress) + path = decision_path(tmp_path, "smith", 1) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + yaml.safe_dump( + { + "summary": "registry lookup, not an op", + "decision": "MACRO", + "constituents": [ + {"target": "causal_conv1d_fwd", "tier": 1}, + {"target": "chunk_gated_delta_rule", "tier": 2}, + ], + } + ), + encoding="utf-8", + ) + assert collect_decision_file(ctx, tmp_path, "smith", 1) == "MACRO" + stored = latest_entry(progress, "smith")["constituents"] + assert [c["target"] for c in stored] == ["causal_conv1d_fwd", "chunk_gated_delta_rule"] + + +# ------------------------------------------------ the orchestrator's action + + +def _workflow(tmp_path, *, features): + workflow = StaircaseOnboardWorkflow(tmp_path / "ws") + workflow.task_data = {"features": features} + init_progress_file(workflow.progress_path) + return workflow + + +def _record_macro(workflow, children): + ctx = ProgressContext(path=workflow.progress_path) + ctx.append( + "smith", + {"iteration": 1, "agent": "smith", "decision": "MACRO", "constituents": children}, + ) + + +def test_constituents_are_appended_to_the_queue(tmp_path): + workflow = _workflow(tmp_path, features=[]) + state = WorkflowState(task_path="t", queue=[{"n": 1, "target": "gdn", "status": "running"}]) + _record_macro(workflow, [{"target": "conv_fwd", "tier": 1}, {"target": "delta", "tier": 2}]) + + added = workflow._enqueue_constituents( + state, QueueItem(n=1, mode="onboard", target="gdn"), None + ) + + assert added == 2 + assert [row["target"] for row in state.queue[1:]] == ["conv_fwd", "delta"] + assert [row["n"] for row in state.queue[1:]] == [2, 3] + assert state.queue[2]["tier"] == 2 + # Provenance, so a later reader knows why this dispatch exists. + assert "constituent of dispatch 1" in state.queue[1]["note"] + + +def test_constituents_for_a_later_unit_are_deferred(tmp_path): + """Core must be green before the first feature increment. + + Measured: of one macro's eleven constituents, four were MTP-only. + Pulling them into the core batch would gate core against vocabulary it + does not use. + """ + workflow = _workflow(tmp_path, features=["mtp"]) + state = WorkflowState(task_path="t", unit_index=0, queue=[{"n": 1, "target": "gdn"}]) + _record_macro( + workflow, + [ + {"target": "conv_fwd", "tier": 1, "unit": "core"}, + {"target": "recurrent_update", "tier": 2, "unit": "mtp"}, + ], + ) + + added = workflow._enqueue_constituents( + state, QueueItem(n=1, mode="onboard", target="gdn"), None + ) + + assert added == 1 + assert [row["target"] for row in state.queue[1:]] == ["conv_fwd"] + + +def test_a_later_unit_accepts_its_own_constituents(tmp_path): + workflow = _workflow(tmp_path, features=["mtp"]) + state = WorkflowState(task_path="t", unit_index=1, queue=[{"n": 1, "target": "gdn"}]) + _record_macro(workflow, [{"target": "recurrent_update", "tier": 2, "unit": "mtp"}]) + + assert ( + workflow._enqueue_constituents(state, QueueItem(n=1, mode="onboard", target="gdn"), None) + == 1 + ) + + +def test_constituents_of_a_feature_only_row_carry_its_unit(tmp_path): + """A feature-only row is dispatched in the smith phase, so its pieces are too. + + Under its tag, though: a piece is needed no earlier than the row it + decomposes, and inheriting the tag is what makes a failure among the + pieces defer to that unit rather than block core. + """ + workflow = _workflow(tmp_path, features=["mtp"]) + state = WorkflowState( + task_path="t", + unit_index=0, + queue=[{"n": 1, "mode": "onboard", "target": "gdn", "unit": "mtp"}], + ) + _record_macro( + workflow, + [{"target": "conv_fwd", "tier": 1}, {"target": "delta", "tier": 2, "unit": "core"}], + ) + + added = workflow._enqueue_constituents( + state, QueueItem(n=1, mode="onboard", target="gdn", unit="mtp"), None + ) + + assert added == 2 + assert [row["unit"] for row in state.queue[1:]] == ["mtp", "mtp"] + + +def test_a_target_already_queued_is_not_dispatched_twice(tmp_path): + """Two macros naming a shared piece, or a resume, must not duplicate it.""" + workflow = _workflow(tmp_path, features=[]) + state = WorkflowState(task_path="t", queue=[{"n": 1, "target": "conv_fwd"}]) + _record_macro(workflow, [{"target": "conv_fwd", "tier": 1}]) + + assert ( + workflow._enqueue_constituents(state, QueueItem(n=1, mode="onboard", target="gdn"), None) + == 0 + ) + + +def test_a_macro_with_no_usable_constituents_adds_nothing(tmp_path): + """Then the caller marks it failed — which is the old behaviour, correctly.""" + workflow = _workflow(tmp_path, features=[]) + state = WorkflowState(task_path="t", queue=[{"n": 1, "target": "gdn"}]) + _record_macro(workflow, [{"tier": 1, "note": "no target"}]) + + assert ( + workflow._enqueue_constituents(state, QueueItem(n=1, mode="onboard", target="gdn"), None) + == 0 + ) + + +# ------------------------------------------------- the tier reaches the agent + + +def test_a_writing_dispatch_is_told_it_is_writing(tmp_path): + """The tier must reach the agent. + + Silently dropping this was a real bug: the variable was built and never + interpolated, so every dispatch read as though it were wrapping. + """ + from agent_flow.workflows.staircase_onboard.vocabulary import TIER_TRITON + + workflow = _workflow(tmp_path, features=[]) + workflow.task_data.update({"trtllm_repo_path": "/repo", "arch": "sm_103"}) + prompt = workflow._smith_prompt( + QueueItem(n=1, mode="onboard", target="exotic_norm", tier=TIER_TRITON) + ) + assert "tier 4" in prompt + assert "you are writing it, not wrapping one" in prompt + + +def test_a_wrapping_dispatch_is_not_told_to_write(tmp_path): + workflow = _workflow(tmp_path, features=[]) + workflow.task_data.update({"trtllm_repo_path": "/repo", "arch": "sm_103"}) + prompt = workflow._smith_prompt(QueueItem(n=1, mode="onboard", target="x", tier=1)) + assert "tier 1" in prompt + assert "you are writing it" not in prompt diff --git a/agent-flow/tests/workflows/staircase_onboard/test_package_hygiene.py b/agent-flow/tests/workflows/staircase_onboard/test_package_hygiene.py new file mode 100644 index 000000000000..1f2281b2a91a --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_package_hygiene.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""What the staircase package tracks: no session placeholders, no one's home directory. + +A sandboxed session leaves empty `.mcp.json` and `.claude/*` files in its +working directory, and a commit that stages by directory picks them up. A path +under one operator's home or project scratch works on exactly one machine. +""" + +from __future__ import annotations + +import re +import shutil +import subprocess +from pathlib import Path + +import pytest + +PACKAGE = Path(__file__).resolve().parents[3] / "agent_flow" / "workflows" / "staircase_onboard" +TESTS = Path(__file__).resolve().parent + +#: A path under a user's home, or under a named cluster's per-user storage. +PERSONAL_PATH = re.compile(r"/home/[a-z][a-z0-9_-]*/|/(?:scratch|lustre)/fsw/portfolios/[a-z]") + + +def _tracked(directory: Path) -> list[Path]: + if shutil.which("git") is None: + pytest.skip("git is not available") + result = subprocess.run( + ["git", "ls-files", "-z", "--", "."], + cwd=directory, + capture_output=True, + check=False, + ) + if result.returncode != 0: + pytest.skip("not inside a git checkout") + return [directory / name for name in result.stdout.decode().split("\0") if name] + + +def test_no_session_placeholder_is_tracked(): + tracked = _tracked(PACKAGE) + placeholders = [ + path.relative_to(PACKAGE) + for path in tracked + if path.name == ".mcp.json" or ".claude" in path.relative_to(PACKAGE).parts + ] + assert not placeholders, f"session placeholders tracked: {placeholders}" + empty = [ + path.relative_to(PACKAGE) + for path in tracked + if path.is_file() and path.stat().st_size == 0 and path.name != "__init__.py" + ] + assert not empty, f"empty files tracked: {empty}" + + +@pytest.mark.parametrize("root", [PACKAGE, TESTS], ids=["package", "tests"]) +def test_no_tracked_file_names_a_personal_path(root): + offences = [] + for path in _tracked(root): + if path.suffix not in {".py", ".md", ".yaml", ".sh", ".json"} or not path.is_file(): + continue + for number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + if PERSONAL_PATH.search(line): + offences.append(f"{path.relative_to(root)}:{number}: {line.strip()[:100]}") + assert not offences, "personal paths:\n " + "\n ".join(offences) diff --git a/agent-flow/tests/workflows/staircase_onboard/test_preflight.py b/agent-flow/tests/workflows/staircase_onboard/test_preflight.py new file mode 100644 index 000000000000..472ef34aa802 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_preflight.py @@ -0,0 +1,578 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for Phase A — the checks that decide whether the run happens. + +Every probe is injected, so none of this needs a GPU. That is the point: +the checks most worth getting right are the ones about broken hardware, and +those are exactly the ones a test machine cannot reproduce. +""" + +from __future__ import annotations + +import json +import os +import socket +import uuid +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import preflight +from agent_flow.workflows.staircase_onboard.preflight import ( + CommandResult, + DeviceReport, + PreflightError, +) + +_SMI_ARGS = "--query-gpu=index,memory.free,memory.total,utilization.gpu,compute_cap" + + +def _smi_rows(*rows: str) -> str: + return "\n".join(rows) + "\n" + + +def _healthy_row(index: int, cap: str = "10.3") -> str: + return f"{index}, 180000, 183359, 0, {cap}" + + +def make_runner(*, smi: str, probe_ok=lambda index: True, multirank_ok: bool = True): + """Build a runner that answers nvidia-smi, the lone probe, and torchrun.""" + + def runner(argv, env=None, timeout=0): + joined = " ".join(argv) + if argv and argv[0] == "nvidia-smi": + return CommandResult(returncode=0, stdout=smi) + if argv and argv[0] == "torchrun": + if multirank_ok: + return CommandResult(returncode=0, stdout="[preflight] OK\n") + return CommandResult(returncode=1, stderr="NCCL init failed") + if "torch.cuda.set_device" in joined: + index = int((env or {}).get("CUDA_VISIBLE_DEVICES", "0")) + if probe_ok(index): + return CommandResult(returncode=0, stdout="ok\n") + return CommandResult(returncode=1, stderr="cudaErrorDevicesUnavailable") + raise AssertionError(f"unexpected command: {argv}") + + return runner + + +# A declared scratch location that is not under /tmp. It is never created: +# `check_tmpdir` reads the declared value, it does not stat it — and pytest's +# own `tmp_path` lives under /tmp on some hosts, which is exactly what the +# check rejects. +_SCRATCH = "/scratch/staircase-run-tmp" + + +def _task(tmp_path, **overrides): + checkpoint = tmp_path / "ckpt" + checkpoint.mkdir(exist_ok=True) + (checkpoint / "config.json").write_text(json.dumps({"architectures": ["Qwen3MoeForCausalLM"]})) + task = { + "checkpoint_path": str(checkpoint), + "arch": "sm_103", + "parallel": "tp1", + "trtllm_repo_path": str(tmp_path), + "environment": { + "kind": "local", + "python": "/opt/venv/bin/python", + "env": {"TMPDIR": _SCRATCH, "CLAUDE_CODE_TMPDIR": _SCRATCH}, + }, + } + task.update(overrides) + return task + + +# ---------------------------------------------------------------- checkpoint + + +def test_checkpoint_must_parse(tmp_path): + checkpoint = tmp_path / "ckpt" + checkpoint.mkdir() + (checkpoint / "config.json").write_text("{not json") + with pytest.raises(PreflightError, match="not valid JSON"): + preflight.check_checkpoint(checkpoint) + + +def test_missing_config_json_is_named(tmp_path): + with pytest.raises(PreflightError, match="not a HF checkpoint dir"): + preflight.check_checkpoint(tmp_path) + + +# ------------------------------------------------------------------- tmpdir + + +def test_tmpdir_must_be_set(): + with pytest.raises(PreflightError, match="TMPDIR is not set"): + preflight.check_tmpdir({"CLAUDE_CODE_TMPDIR": "/scratch/x"}) + + +def test_tmpdir_does_not_fall_back_to_the_launching_shell(monkeypatch): + """Agents run under the declared prefix, not under the orchestrator's env. + + A run whose scratch location is inherited rather than declared is one + nobody can reproduce, so an omitted value is an error even when the + launching shell happens to have a perfectly good one. + """ + monkeypatch.setenv("TMPDIR", "/scratch/inherited") + monkeypatch.setenv("CLAUDE_CODE_TMPDIR", "/scratch/inherited") + with pytest.raises(PreflightError, match="TMPDIR is not set"): + preflight.check_tmpdir({}) + + +@pytest.mark.parametrize("value", ["/tmp", "/tmp/mine", "/tmp/a/b"]) +def test_tmpdir_under_slash_tmp_is_rejected(value): + with pytest.raises(PreflightError, match="points into /tmp"): + preflight.check_tmpdir({"TMPDIR": value, "CLAUDE_CODE_TMPDIR": value}) + + +def test_tmpdir_elsewhere_is_accepted(): + preflight.check_tmpdir({"TMPDIR": _SCRATCH, "CLAUDE_CODE_TMPDIR": _SCRATCH}) + + +# ------------------------------------------------------ tmpdir, by length +# +# The property that breaks first on a cluster whose storage paths are deep, +# and the one an author is least likely to suspect: the directory is on the +# right filesystem, is writable, has space, and still fails because its NAME +# is long. Observed as a ZeroMQ traceback naming neither the workflow nor the +# spec line, after four ranks had started and loaded weights. + + +def test_the_bound_is_derived_from_the_two_fixed_costs(): + """Not a literal. `107 - 45 - len("/")`, and the separator is the easy miss. + + An `AF_UNIX` bind succeeds at 107 characters and fails at 108; the + basename `rpc_test_` is `len("rpc_test_") + 36` == 45; and + `os.path.join` puts a separator between them. Dropping the separator + yields 62 and admits one `TMPDIR` that does not work. + """ + assert preflight.MAX_TMPDIR_LEN == 61 + assert ( + preflight.MAX_TMPDIR_LEN + == preflight.UNIX_SOCKET_PATH_MAX - preflight.RPC_SOCKET_BASENAME_LEN - 1 + ) + + +def test_a_tmpdir_too_long_for_a_unix_socket_is_rejected(): + """A 92-character value, the length a real run declared, and what it cost.""" + value = "/scratch/projects/example_org/example_project/inference/users/example_user/run/staircase-tmp" + assert len(value) == 92 + with pytest.raises(PreflightError) as excinfo: + preflight.check_tmpdir({"TMPDIR": value, "CLAUDE_CODE_TMPDIR": value}) + message = str(excinfo.value) + # The arithmetic belongs in the message: a bare "too long" sends the + # reader to the filesystem, which is not where the problem is. + assert "92 characters" in message + assert "rpc_test_" in message + assert "sizeof(sockaddr_un.sun_path)" in message + assert "at most 61" in message + + +def test_the_boundary_is_inclusive_on_the_side_that_binds(): + """61 works and 62 does not — verified against a real bind, not asserted.""" + preflight.check_tmpdir({"TMPDIR": "/s" + "a" * 59, "CLAUDE_CODE_TMPDIR": _SCRATCH}) + with pytest.raises(PreflightError, match="at most 61"): + preflight.check_tmpdir({"TMPDIR": "/s" + "a" * 60, "CLAUDE_CODE_TMPDIR": _SCRATCH}) + + +def test_the_derived_bound_agrees_with_an_actual_af_unix_bind(): + """The arithmetic checked against the kernel rather than against itself. + + This is the one measurement behind the bound that can be taken locally + with no GPU, so there is no reason to take it on trust: bind a socket at + exactly the budget and one character past it, under the real basename. + + ``/tmp`` on purpose — it is the shortest directory that certainly exists, + and the length of the prefix is the whole variable here. + """ + root = Path("/tmp") + if not os.access(root, os.W_OK): + pytest.skip("no writable /tmp to measure the socket path limit against") + + basename = f"rpc_test_{uuid.uuid4()}" + assert len(basename) == preflight.RPC_SOCKET_BASENAME_LEN + + def binds(tmpdir_len: int) -> bool: + directory = root / ("a" * (tmpdir_len - len("/tmp/"))) + directory.mkdir(exist_ok=True) + path = directory / basename + assert len(str(path)) == tmpdir_len + 1 + len(basename) + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + sock.bind(str(path)) + return True + except OSError: + return False + finally: + sock.close() + path.unlink(missing_ok=True) + directory.rmdir() + + assert binds(preflight.MAX_TMPDIR_LEN) + assert not binds(preflight.MAX_TMPDIR_LEN + 1) + + +def test_the_length_bound_does_not_reach_the_harness_scratch_dir(): + """`CLAUDE_CODE_TMPDIR` binds no socket, so a socket's budget is not its. + + Holding it to one would reject a value that works, for a reason that + would not be true of it. + """ + long_value = "/scratch/" + "a" * 100 + assert len(long_value) > preflight.MAX_TMPDIR_LEN + preflight.check_tmpdir({"TMPDIR": _SCRATCH, "CLAUDE_CODE_TMPDIR": long_value}) + + +# ------------------------------------------------------------------ devices + + +def test_compute_cap_maps_to_the_receipt_key(): + assert DeviceReport(0, 1, 1, 0, "10.3").arch == "sm_103" + assert DeviceReport(0, 1, 1, 0, "10.0").arch == "sm_100" + assert DeviceReport(0, 1, 1, 0, "9.0").arch == "sm_90" + + +def test_a_spinning_card_is_broken_not_busy(): + """100% utilization invalidates every number taken there.""" + report = DeviceReport(0, 180000, 183359, 100, "10.3") + assert report.looks_idle is False + assert "broken for our purposes" in report.why_not_idle() + + +def test_a_card_holding_memory_is_not_claimable(): + report = DeviceReport(0, 20000, 183359, 0, "10.3") + assert report.looks_idle is False + assert "free" in report.why_not_idle() + + +def test_claim_picks_healthy_devices(tmp_path): + runner = make_runner(smi=_smi_rows(_healthy_row(0), _healthy_row(1), _healthy_row(2))) + assert preflight.claim_devices(2, "sm_103", "/p", runner=runner) == (0, 1) + + +def test_claim_skips_a_card_that_fails_the_lone_probe(): + """A card that fails alone is unusable, not contended — waiting never helps.""" + runner = make_runner( + smi=_smi_rows(_healthy_row(0), _healthy_row(1), _healthy_row(2)), + probe_ok=lambda index: index != 0, + ) + notes: list[str] = [] + assert preflight.claim_devices(2, "sm_103", "/p", runner=runner, notes=notes) == (1, 2) + assert any("lone-process probe" in n for n in notes) + + +def test_claim_reports_why_when_it_cannot_fill_the_world_size(): + runner = make_runner( + smi=_smi_rows(_healthy_row(0), "1, 20000, 183359, 0, 10.3", "2, 180000, 183359, 97, 10.3") + ) + with pytest.raises(PreflightError) as excinfo: + preflight.claim_devices(3, "sm_103", "/p", runner=runner) + message = str(excinfo.value) + assert "need 3 healthy sm_103" in message + assert "device 1" in message and "device 2" in message + + +def test_wrong_architecture_is_refused_with_the_reason(): + runner = make_runner(smi=_smi_rows(_healthy_row(0, cap="10.0"), _healthy_row(1, cap="10.0"))) + with pytest.raises(PreflightError) as excinfo: + preflight.claim_devices(1, "sm_103", "/p", runner=runner) + assert "sm_100" in str(excinfo.value) + assert "another machine's certification" in str(excinfo.value) + + +def test_nvidia_smi_failure_is_not_silently_an_empty_claim(): + def runner(argv, env=None, timeout=0): + return CommandResult(returncode=127, stderr="no nvidia-smi") + + with pytest.raises(PreflightError, match="nvidia-smi did not answer"): + preflight.query_devices(runner) + + +def test_unparsable_smi_row_is_rejected(): + def runner(argv, env=None, timeout=0): + return CommandResult(returncode=0, stdout="0, 1, 2\n") + + with pytest.raises(PreflightError, match="unparsable"): + preflight.query_devices(runner) + + +# ---------------------------------------------------------------- multirank + + +def test_multirank_probe_is_skipped_at_world_size_one(): + def runner(argv, env=None, timeout=0): + raise AssertionError("should not run") + + preflight.check_multirank_substrate(1, None, runner) # type: ignore[arg-type] + + +def test_multirank_timeout_explains_the_hang(tmp_path): + from agent_flow.workflows.staircase_onboard.environment import resolve_environment + + env = resolve_environment(_task(tmp_path), devices=[0, 1], world_size=2) + + def runner(argv, env_=None, timeout=0): + return CommandResult(returncode=124, timed_out=True) + + with pytest.raises(PreflightError) as excinfo: + preflight.check_multirank_substrate(2, env, runner) + assert "hangs rather than fails" in str(excinfo.value) + + +# ---------------------------------------------------------- accuracy anchors +# +# The failure these cover happened: the `multimodal` unit's gate named +# `mmmu.yaml`, that file had no row for the checkpoint, and nobody noticed +# until someone looked by luck while an earlier unit was still gating. The +# whole unit would otherwise have been assembled, booted on three members and +# killed at its accuracy gate hours later, on a spec error readable at +# startup. So what these tests actually assert is *when* it is noticed. + + +_REFERENCES_DIR = "tests/integration/defs/accuracy/references" + + +def _reference(tmp_path, name, payload): + """Write a reference file where `trtllm_repo_path` will resolve it.""" + path = tmp_path / _REFERENCES_DIR / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(yaml.safe_dump(payload), encoding="utf-8") + return f"{_REFERENCES_DIR}/{name}" + + +def test_a_missing_model_key_is_named_before_anything_is_built(tmp_path): + reference = _reference(tmp_path, "gsm8k.yaml", {"openai/gpt-oss-120b": [{"accuracy": 90.0}]}) + task = _task( + tmp_path, + accuracy={"benchmark": "gsm8k", "reference_yaml": reference, "model_key": "XiaomiMiMo/x"}, + ) + with pytest.raises(PreflightError) as excinfo: + preflight.check_accuracy_anchors(task) + message = str(excinfo.value) + assert "unit `core`" in message + assert "XiaomiMiMo/x" in message + assert reference in message + # `load_reference_score`'s own words, surfaced rather than flattened to + # "not found" -- including the instruction that the anchor is never + # invented, which is the part that says what to do next. + assert "has no entry for 'XiaomiMiMo/x'" in message + assert "never invented" in message + + +def test_a_model_with_no_matching_row_lists_the_rows_that_are_there(tmp_path): + """The discriminator case, which is the one a spec gets subtly wrong. + + A model key that is present reads as "the anchor exists", and the run + then measures against a row that does not describe this checkpoint -- + or, here, against nothing at all. + """ + reference = _reference( + tmp_path, + "gsm8k.yaml", + {"openai/gpt-oss-120b": [{"quant_algo": "FP8", "accuracy": 90.0}]}, + ) + task = _task( + tmp_path, + accuracy={ + "benchmark": "gsm8k", + "reference_yaml": reference, + "model_key": "openai/gpt-oss-120b", + "discriminators": {"quant_algo": "MIXED_PRECISION"}, + }, + ) + with pytest.raises(PreflightError) as excinfo: + preflight.check_accuracy_anchors(task) + message = str(excinfo.value) + assert "discriminators: {'quant_algo': 'MIXED_PRECISION'}" in message + # What was searched for AND what is available, because the gap between + # the two is the edit the author has to make. + assert "has no entry matching {'quant_algo': 'MIXED_PRECISION'}" in message + assert "Available discriminator sets: [{'quant_algo': 'FP8'}]" in message + + +def test_every_unit_is_checked_not_just_core(tmp_path): + """A valid multi-unit spec resolves, and each unit is reported. + + `benchmark`, `reference_yaml` and `discriminators` are all per-unit + here, which is the shape the run that produced this check used. + """ + gsm8k = _reference( + tmp_path, + "gsm8k.yaml", + {"XiaomiMiMo/MiMo-V2.6-Flash": [{"quant_algo": "MIXED_PRECISION", "accuracy": 91.2}]}, + ) + mmmu = _reference( + tmp_path, + "mmmu.yaml", + {"XiaomiMiMo/MiMo-V2.6-Flash": [{"quant_algo": "MIXED_PRECISION", "accuracy": 64.8}]}, + ) + task = _task( + tmp_path, + features=["multimodal"], + accuracy={ + "benchmark": {"default": "gsm8k", "multimodal": "mmmu"}, + "reference_yaml": {"default": gsm8k, "multimodal": mmmu}, + "model_key": "XiaomiMiMo/MiMo-V2.6-Flash", + "discriminators": {"quant_algo": "MIXED_PRECISION"}, + }, + ) + notes = preflight.check_accuracy_anchors(task) + assert len(notes) == 2 + assert "accuracy anchor [core]: gsm8k XiaomiMiMo/MiMo-V2.6-Flash" in notes[0] + assert "= 91.2" in notes[0] + assert "accuracy anchor [multimodal]: mmmu XiaomiMiMo/MiMo-V2.6-Flash" in notes[1] + assert "= 64.8" in notes[1] + + +def test_the_unit_that_differs_is_the_one_that_fails(tmp_path): + """Core resolving says nothing about the increment, and vice versa. + + This is the run's actual failure: the text anchor was fine, and the + multimodal one -- a different benchmark, in a different file -- was not. + """ + gsm8k = _reference(tmp_path, "gsm8k.yaml", {"XiaomiMiMo/M": [{"accuracy": 91.2}]}) + mmmu = _reference(tmp_path, "mmmu.yaml", {"some/other-model": [{"accuracy": 60.0}]}) + task = _task( + tmp_path, + features=["multimodal"], + accuracy={ + "benchmark": {"default": "gsm8k", "multimodal": "mmmu"}, + "reference_yaml": {"default": gsm8k, "multimodal": mmmu}, + "model_key": "XiaomiMiMo/M", + }, + ) + with pytest.raises(PreflightError, match="unit `multimodal`"): + preflight.check_accuracy_anchors(task) + + +def test_a_unit_with_no_accuracy_block_is_skipped(tmp_path): + """Two spellings of "there is nothing to check", neither a failure.""" + assert preflight.check_accuracy_anchors(_task(tmp_path)) == [] + + task = _task( + tmp_path, + features=["multimodal"], + accuracy={"benchmark": "gsm8k", "reference_yaml": {"core": "refs/gsm8k.yaml"}}, + ) + notes = preflight.check_accuracy_anchors(task) + assert "no reference_yaml for this unit" in notes[1] + + +def test_a_reference_file_this_checkout_does_not_have_is_not_a_failure(tmp_path): + """The check owns the ROW. A file it cannot read tells it nothing. + + A sparse checkout, or a path that resolves outside the repo entirely, + must not stop a run -- but the note has to state the weaker claim rather + than read as a pass. + """ + task = _task( + tmp_path, + accuracy={ + "benchmark": "gsm8k", + "reference_yaml": "/elsewhere/gsm8k.yaml", + "model_key": "openai/gpt-oss-120b", + }, + ) + notes = preflight.check_accuracy_anchors(task) + assert "accuracy reference file not found" in notes[0] + assert "unchecked here" in notes[0] + + +def test_the_model_key_falls_back_to_the_checkpoint_repo_name(tmp_path): + """Because `accuracy.model_key` is optional, and the gate derives it too. + + Deriving a different key here than the gate derives would make this + check pass on a row nothing ever reads. + """ + task = _task(tmp_path) + assert preflight.resolve_model_key(task) == f"{tmp_path.name}/ckpt" + + +# -------------------------------------------------------------- end to end + + +def test_run_preflight_resolves_identity(tmp_path): + task = _task(tmp_path, parallel="dep4") + (tmp_path / "ckpt").rename(tmp_path / "Qwen3-30B-A3B") + task["checkpoint_path"] = str(tmp_path / "Qwen3-30B-A3B") + runner = make_runner(smi=_smi_rows(*[_healthy_row(i) for i in range(4)])) + result = preflight.run_preflight(task, runner=runner) + assert result.slug == "qwen3-30b-a3b-sm_103-dep4" + assert result.branch == "staircase/qwen3-30b-a3b-sm_103-dep4" + assert result.devices == (0, 1, 2, 3) + assert result.world_size == 4 + assert result.checkpoint_config["architectures"] == ["Qwen3MoeForCausalLM"] + + +def test_run_preflight_skips_the_claim_under_slurm(tmp_path): + task = _task(tmp_path) + task["environment"]["kind"] = "slurm" + task["environment"]["partition"] = "batch" + task["environment"]["container_image"] = "/img.sqsh" + + def runner(argv, env=None, timeout=0): + if argv and argv[0] == "torchrun": + return CommandResult(returncode=0, stdout="[preflight] OK\n") + raise AssertionError(f"login-node cards must not be probed: {argv}") + + result = preflight.run_preflight(task, runner=runner) + assert result.devices == () + assert result.environment.prefix.startswith("srun ") + + +def test_unified_segment_defaults_its_world_size(tmp_path): + task = _task(tmp_path, parallel="unified") + runner = make_runner(smi=_smi_rows(*[_healthy_row(i) for i in range(4)])) + result = preflight.run_preflight(task, runner=runner) + assert result.world_size == 4 + assert any("unified segment" in n for n in result.notes) + + +def test_a_bad_anchor_stops_the_run_before_a_single_device_is_probed(tmp_path): + """Where this check sits in the order, asserted rather than described. + + It is pure file reading, so everything below it -- the device claim, the + multi-rank probe, the substrate gate -- is work a failure here would + have wasted. + """ + reference = _reference(tmp_path, "gsm8k.yaml", {"someone/else": [{"accuracy": 90.0}]}) + task = _task( + tmp_path, + accuracy={ + "benchmark": "gsm8k", + "reference_yaml": reference, + "model_key": "openai/gpt-oss-120b", + }, + ) + + def runner(argv, env=None, timeout=0): # pragma: no cover - must not run + raise AssertionError(f"probed the host before checking the spec: {argv}") + + with pytest.raises(PreflightError, match="unit `core`"): + preflight.run_preflight(task, runner=runner) + + +def test_the_resolved_anchor_reaches_the_report(tmp_path): + reference = _reference(tmp_path, "gsm8k.yaml", {"openai/gpt-oss-120b": [{"accuracy": 90.5}]}) + task = _task( + tmp_path, + accuracy={ + "benchmark": "gsm8k", + "reference_yaml": reference, + "model_key": "openai/gpt-oss-120b", + }, + ) + runner = make_runner(smi=_smi_rows(_healthy_row(0))) + result = preflight.run_preflight(task, runner=runner) + rendered = preflight.format_preflight_report(result) + assert "accuracy anchor [core]: gsm8k openai/gpt-oss-120b {} = 90.5" in rendered + + +def test_format_report_carries_the_prefix(tmp_path): + runner = make_runner(smi=_smi_rows(_healthy_row(0))) + result = preflight.run_preflight(_task(tmp_path), runner=runner) + rendered = preflight.format_preflight_report(result) + assert result.environment.prefix in rendered + assert result.slug in rendered diff --git a/agent-flow/tests/workflows/staircase_onboard/test_prompts.py b/agent-flow/tests/workflows/staircase_onboard/test_prompts.py new file mode 100644 index 000000000000..e8f30597e46b --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_prompts.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the six role prompts and the bundle that scopes them to a run.""" + +from __future__ import annotations + +import dataclasses + +import pytest + +from agent_flow.workflows.staircase_onboard import prompts as prompts_mod +from agent_flow.workflows.staircase_onboard.environment import resolve_environment +from agent_flow.workflows.staircase_onboard.prompts import ( + DEFAULT_PROMPTS, + PromptBundle, + build_staircase_onboard_prompts, +) + +ROLES = tuple(f.name for f in dataclasses.fields(DEFAULT_PROMPTS)) + +# The one role that reports no typed decision, because it makes none: it +# relays a submission to a scheduler the orchestrator cannot reach, and the +# verdict comes from files the job wrote. Every rule below that is about +# reporting a verdict, or about reading the in-tree layout, is therefore not +# about this role. +RUNNER = "runner" +DECIDING_ROLES = tuple(r for r in ROLES if r != RUNNER) + +# Things the port retired. A prompt naming one of them is telling an agent to +# use a product that no longer exists, which is worse than saying nothing. +RETIRED = ( + "uv run", + "ty check", + "TARGET.md", + "smoke.py", + "llm_args.yaml", + "model_dir", + "1.3.0rc25", + ".allow-trtllm", + "isolation hook", + "bench/accuracy.py", + "bench/perf.py", +) + + +def _env(): + return resolve_environment( + {"environment": {"kind": "local", "python": "/opt/venv/bin/python", "env": {}}}, + devices=[0, 1], + world_size=2, + ) + + +def test_the_bundle_has_exactly_the_roles_the_workflow_dispatches(): + assert ROLES == ( + "plan_drafter", + "plan_reviewer", + "smith", + "assembler", + "reviewer", + "auditor", + "runner", + ) + + +def test_every_prompt_is_substantial(): + for role in ROLES: + assert len(getattr(DEFAULT_PROMPTS, role)) > 2000, role + + +def test_every_deciding_prompt_names_its_own_progress_tool(): + """The typed decision is the only channel a verdict reaches the orchestrator.""" + for role in DECIDING_ROLES: + assert f"append_{role}_progress" in getattr(DEFAULT_PROMPTS, role), role + + +def test_the_runner_is_told_that_nothing_reads_its_answer(): + """It has no progress tool because it has no verdict. + + Stating that in the prompt is what keeps it from inventing one: a role + that believes it is reporting a result will report a result. + """ + prompt = DEFAULT_PROMPTS.runner + assert "append_runner_progress" not in prompt + assert "nothing reads your answer" in prompt + assert "trigger, not a witness" in prompt + + +def test_no_prompt_names_a_retired_product(): + for role in ROLES: + text = getattr(DEFAULT_PROMPTS, role) + for retired in RETIRED: + assert retired not in text, f"{role} still names {retired!r}" + + +def test_every_prompt_that_reads_the_tree_carries_its_layout(): + """The runner is excluded because it never opens the repository. + + It is handed a script path and a submission call. Telling it where the + catalog lives would be telling it about work that is not its own, which + is how a relay starts having opinions. + """ + for role in DECIDING_ROLES: + assert "_experimental/modeling_v2" in getattr(DEFAULT_PROMPTS, role), role + assert "_experimental/modeling_v2" not in DEFAULT_PROMPTS.runner + + +def test_every_role_has_a_red_flags_table(): + """Each row has an incident behind it; they are the expensive part.""" + for role in ROLES: + assert "## Red flags" in getattr(DEFAULT_PROMPTS, role), role + assert "| Impulse | Reality |" in getattr(DEFAULT_PROMPTS, role), role + + +def test_the_test_naming_convention_reaches_the_roles_that_write_tests(): + for role in ("smith", "assembler", "reviewer", "auditor"): + assert "test_modeling_v2_" in getattr(DEFAULT_PROMPTS, role), role + + +def test_the_modeling_v2_switch_reaches_the_roles_that_run_the_engine(): + for role in ("plan_drafter", "assembler"): + assert "TRTLLM_MODELING_V2" in getattr(DEFAULT_PROMPTS, role), role + assert "require" in getattr(DEFAULT_PROMPTS, role), role + + +def test_the_anchor_is_never_invented_reaches_the_planner(): + assert "never invented" in DEFAULT_PROMPTS.plan_drafter + assert "never invented" in DEFAULT_PROMPTS.plan_reviewer.lower() or ( + "invented" in DEFAULT_PROMPTS.plan_reviewer + ) + + +def test_receipts_are_three_state_in_the_smith_prompt(): + assert "absent — unknown" in DEFAULT_PROMPTS.smith + assert "Tolerance inflation manufactures green" in DEFAULT_PROMPTS.smith + + +def test_the_runtime_dump_incident_reaches_the_roles_that_can_catch_it(): + """34 of 35 batch sizes outside the contract — the failure the auditor exists for.""" + for role in ("assembler", "reviewer", "auditor"): + assert "34 of the 35" in getattr(DEFAULT_PROMPTS, role), role + + +def test_the_source_boundary_survives_the_hook_removal(): + """The hook is gone; the rule it defended is prose in every writing role.""" + for role in ("plan_drafter", "smith", "assembler"): + text = getattr(DEFAULT_PROMPTS, role) + assert "never to copy a structure" in text, role + assert "ordinary read" in DEFAULT_PROMPTS.smith + + +def test_the_reviewer_is_told_not_to_re_run_the_gates(): + assert "Do not re-run the gates" in DEFAULT_PROMPTS.reviewer + assert "already exists" in DEFAULT_PROMPTS.reviewer + + +def test_the_plan_phase_is_told_it_does_not_block_the_build(): + assert "not** a failure" in DEFAULT_PROMPTS.plan_drafter + assert "the build is not blocked on you" in DEFAULT_PROMPTS.plan_reviewer.lower() or ( + "not blocked on you" in DEFAULT_PROMPTS.plan_reviewer + ) + + +def test_the_assembler_is_told_not_to_reformat(): + assert "lints once" in DEFAULT_PROMPTS.assembler + + +# ------------------------------------------------------------------- bundling + + +def test_empty_extensions_leave_the_bundle_unchanged(): + extended = DEFAULT_PROMPTS.with_extensions(smith="", assembler=" \n ") + for role in ROLES: + assert getattr(extended, role) == getattr(DEFAULT_PROMPTS, role) + + +def test_a_non_empty_extension_is_appended_after_one_blank_line(): + extended = DEFAULT_PROMPTS.with_extensions(smith="## Extra\nSomething.\n") + assert extended.smith.startswith(DEFAULT_PROMPTS.smith.rstrip()) + assert extended.smith.endswith("\n\n## Extra\nSomething.\n") + assert "\n\n\n" not in extended.smith + assert extended.auditor == DEFAULT_PROMPTS.auditor + + +def test_the_bundle_is_frozen(): + with pytest.raises(Exception): + DEFAULT_PROMPTS.smith = "mutated" # type: ignore[misc] + + +def test_replace_swaps_one_prompt_outright(): + swapped = DEFAULT_PROMPTS.replace(auditor="tiny") + assert swapped.auditor == "tiny" + assert swapped.smith == DEFAULT_PROMPTS.smith + + +def test_the_environment_block_reaches_every_role(): + """A role that was not told the rule is a role that invents a launcher command.""" + bundle = build_staircase_onboard_prompts(environment=_env()) + for role in ROLES: + text = getattr(bundle, role) + assert "CUDA_VISIBLE_DEVICES=0,1" in text, role + assert "symptom to report" in text, role + + +def test_without_an_environment_the_bundle_is_the_default(): + assert build_staircase_onboard_prompts() == DEFAULT_PROMPTS + + +def test_full_reviewer_mode_adds_its_extension_only_to_the_reviewer(): + bundle = build_staircase_onboard_prompts(reviewer_mode="full") + assert "full` reviewer mode" in bundle.reviewer + assert bundle.smith == DEFAULT_PROMPTS.smith + + +def test_light_reviewer_mode_needs_no_extension(): + assert build_staircase_onboard_prompts(reviewer_mode="light") == DEFAULT_PROMPTS + + +def test_full_mode_still_refuses_to_re_run_the_accuracy_gate(): + """An hour of wall-clock to reproduce a number inside its own noise is delay.""" + bundle = build_staircase_onboard_prompts(reviewer_mode="full") + assert "not re-run the accuracy gate" in bundle.reviewer + + +def test_default_prompts_match_the_module_constants(): + assert DEFAULT_PROMPTS.plan_drafter == prompts_mod.PLAN_DRAFTER_SYSTEM_PROMPT + assert DEFAULT_PROMPTS.smith == prompts_mod.SMITH_SYSTEM_PROMPT + assert DEFAULT_PROMPTS.auditor == prompts_mod.AUDITOR_SYSTEM_PROMPT + + +def test_bundle_equality_is_by_value(): + assert PromptBundle(*(getattr(DEFAULT_PROMPTS, r) for r in ROLES)) == DEFAULT_PROMPTS + + +def test_every_role_is_told_how_to_write_down_what_it_read(): + """The rule has to reach the role, not just the contract-writing skill. + + A smith writes contracts, but an assembler writes target comments and a + reviewer writes findings — all of it read back later against an engine + that has moved. Whoever records a fact needs the rule. + """ + prompts = build_staircase_onboard_prompts() + for role in ( + prompts.plan_drafter, + prompts.plan_reviewer, + prompts.smith, + prompts.assembler, + prompts.reviewer, + prompts.auditor, + ): + assert "Describe the behaviour and name the symbol — never the line" in role + assert "quoted diagnostic" in role diff --git a/agent-flow/tests/workflows/staircase_onboard/test_remediation.py b/agent-flow/tests/workflows/staircase_onboard/test_remediation.py new file mode 100644 index 000000000000..cbe30ca0ecb3 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_remediation.py @@ -0,0 +1,550 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The convergence rules of the UNCOVERED remediation cycle. + +These are the tests the whole module exists for. The loop spends GPU time, +so "does it stop, and does it stop for a nameable reason" is the property +that matters, and it is decidable on a laptop because every rule here is a +pure function over plain data. + +The four failure modes being ruled out, each of which is a real way a +naive `while auditor == UNCOVERED` loop burns a night: + + * one smith per value instead of one per surface; + * re-dispatching a finding that is already queued; + * re-dispatching a surface that was already certified; + * running forever because the auditor keeps saying the same thing. +""" + +from __future__ import annotations + +from agent_flow.workflows.staircase_onboard import remediation as R +from agent_flow.workflows.staircase_onboard.progress import parse_uncovered_records +from agent_flow.workflows.staircase_onboard.state import is_remediable + +# ------------------------------------------------------------------- parsing + + +def test_a_row_without_entry_arch_or_axis_is_dropped_not_guessed(): + """A guessed axis merges unrelated surfaces and certifies neither.""" + rows, dropped = R.parse_uncovered_rows( + [ + {"entry": "attention/x", "arch": "sm_103", "axis": "batch_size", "values": [1]}, + {"entry": "attention/x", "arch": "sm_103"}, # no axis + {"arch": "sm_103", "axis": "batch_size"}, # no entry + "not a mapping", + ] + ) + assert len(rows) == 1 + assert dropped == 3 + + +def test_a_single_value_is_accepted_in_either_spelling(): + rows, _ = R.parse_uncovered_rows([{"entry": "e", "arch": "a", "axis": "x", "value": 7}]) + assert rows[0].values == ("7",) + + +def test_values_are_canonicalized_so_two_spellings_hash_alike(): + """``[12, 3, 5]`` and ``[5, 3, 12]`` are the same gap.""" + one, _ = R.parse_uncovered_rows( + [{"entry": "e", "arch": "a", "axis": "x", "values": [12, 3, 5]}] + ) + two, _ = R.parse_uncovered_rows( + [{"entry": "e", "arch": "a", "axis": "x", "values": [5, 12, 3]}] + ) + assert one[0].values == two[0].values == ("3", "5", "12") + + +def test_numeric_values_sort_numerically_not_lexically(): + """Lexical order would put 12 before 3 and change the identity.""" + rows, _ = R.parse_uncovered_rows( + [{"entry": "e", "arch": "a", "axis": "x", "values": [2, 31, 9]}] + ) + assert rows[0].values == ("2", "9", "31") + + +def test_the_decision_file_parser_agrees_with_the_typed_one(): + """`progress` normalizes for persistence; `remediation` types it. Same rule.""" + records = parse_uncovered_records( + [ + {"entry": "e", "arch": "a", "axis": "x", "values": [3], "contract_line": "L"}, + {"entry": "e", "arch": "a"}, + ] + ) + assert len(records) == 1 + rows, dropped = R.parse_uncovered_rows(records) + assert dropped == 0 and rows[0].contract_line == "L" + + +# -------------------------------------------------------------- canonicalize + + +def _rows(*specs): + return [ + R.UncoveredRow(entry=e, arch=a, axis=x, values=tuple(str(v) for v in vals)) + for e, a, x, vals in specs + ] + + +def test_seven_uncovered_values_on_one_entry_are_one_dispatch(): + """The whole reason the queue is made of surfaces rather than values.""" + rows = _rows( + ("attention/paged", "sm_103", "batch_size", [3, 5, 7]), + ("attention/paged", "sm_103", "batch_size", [9, 12, 15, 31]), + ) + items = R.canonicalize(rows) + assert len(items) == 1 + assert items[0].values == ("3", "5", "7", "9", "12", "15", "31") + + +def test_different_axes_on_one_entry_stay_separate(): + """They are different tests and different contract lines.""" + items = R.canonicalize( + _rows( + ("attention/paged", "sm_103", "batch_size", [3]), + ("attention/paged", "sm_103", "num_pages", [64]), + ) + ) + assert len(items) == 2 + + +def test_the_same_surface_on_two_architectures_is_two_items(): + """Receipts are per-architecture; absent means unknown.""" + items = R.canonicalize( + _rows( + ("norm/rms", "sm_103", "hidden", [4096]), + ("norm/rms", "sm_100", "hidden", [4096]), + ) + ) + assert len(items) == 2 + + +def test_identity_is_stable_across_processes(): + a = R.canonicalize(_rows(("e", "a", "x", [1, 2])), target_commit="c", capture_id="k") + b = R.canonicalize(_rows(("e", "a", "x", [2, 1])), target_commit="c", capture_id="k") + assert a[0].identity == b[0].identity + + +def test_a_new_commit_or_capture_makes_it_new_work(): + """A wrapper change can open a runtime branch that did not exist before.""" + base = R.canonicalize(_rows(("e", "a", "x", [1])), target_commit="c1", capture_id="k") + newer = R.canonicalize(_rows(("e", "a", "x", [1])), target_commit="c2", capture_id="k") + other = R.canonicalize(_rows(("e", "a", "x", [1])), target_commit="c1", capture_id="k2") + assert base[0].identity != newer[0].identity != other[0].identity + + +def test_the_brief_carries_the_whole_surface_and_the_contract_line(): + rows = [ + R.UncoveredRow( + entry="e", + arch="sm_103", + axis="batch_size", + values=("3", "5"), + contract_line="certified at 1, 2, 4", + ) + ] + brief = R.canonicalize(rows)[0].brief() + assert "batch_size" in brief and "3, 5" in brief + assert "certified at 1, 2, 4" in brief + assert "sm_103 receipt" in brief + + +# --------------------------------------------------------------------- merge + + +def test_genuinely_new_work_is_appended(): + merge = R.merge_into_queue([], R.canonicalize(_rows(("e", "a", "x", [1])))) + assert len(merge.appended) == 1 and merge.has_new_work + + +def test_a_finding_already_queued_is_not_dispatched_twice(): + first = R.canonicalize(_rows(("e", "a", "x", [1, 2]))) + merge = R.merge_into_queue(first, first) + assert merge.duplicates and not merge.appended and not merge.has_new_work + + +def test_a_wider_surface_widens_the_open_item_rather_than_adding_one(): + """Two dispatches on one entry would collide on the same file.""" + queue = R.canonicalize(_rows(("e", "a", "x", [1]))) + merge = R.merge_into_queue(queue, R.canonicalize(_rows(("e", "a", "x", [1, 5])))) + assert len(merge.queue) == 1 + assert merge.widened and merge.queue[0].values == ("1", "5") + + +def test_a_certified_surface_coming_back_is_a_regression_not_a_redispatch(): + """Re-running the same smith cannot fix a stale receipt or a bad read.""" + queue = [ + R.canonicalize(_rows(("e", "a", "x", [1])))[0].__class__( + entry="e", arch="a", axis="x", values=("1",), status=R.QUEUE_CERTIFIED + ) + ] + merge = R.merge_into_queue(queue, R.canonicalize(_rows(("e", "a", "x", [1])))) + assert merge.regressions and not merge.appended + + +def test_a_surface_whose_attempt_failed_is_not_retried(): + queue = [ + R.RemediationItem( + entry="e", arch="a", axis="x", values=("1",), status=R.QUEUE_FAILED_RECEIPT + ) + ] + merge = R.merge_into_queue(queue, R.canonicalize(_rows(("e", "a", "x", [1, 2])))) + assert merge.exhausted and not merge.appended and not merge.widened + + +def test_the_new_item_cap_reports_rather_than_silently_dropping(): + """Proceeding on a partial set would converge on an incomplete audit.""" + incoming = R.canonicalize(_rows(*[(f"e{i}", "a", "x", [1]) for i in range(5)])) + merge = R.merge_into_queue([], incoming, budget=R.RemediationBudget(max_new_items=2)) + assert len(merge.appended) == 2 + assert len(merge.exhausted) == 3 + + +# ----------------------------------------------------------------- signature + + +def test_the_signature_ignores_closed_items(): + """Only unresolved work decides whether the cycle is shrinking.""" + open_only = [R.RemediationItem(entry="e", arch="a", axis="x", values=("1",))] + plus_done = open_only + [ + R.RemediationItem(entry="z", arch="a", axis="x", values=("9",), status=R.QUEUE_CERTIFIED) + ] + assert R.unresolved_signature(open_only) == R.unresolved_signature(plus_done) + + +def test_an_empty_queue_has_no_signature(): + assert R.unresolved_signature([]) == "" + + +# --------------------------------------------------------------- round plan + + +def _merge_of(*specs, queue=()): + return R.merge_into_queue(list(queue), R.canonicalize(_rows(*specs))) + + +def test_nothing_open_means_the_audit_is_closed(): + decision = R.decide_round(R.merge_into_queue([], []), budget=R.RemediationBudget(), round_no=1) + assert decision.action == "clean" + + +def test_new_work_dispatches_capped_at_the_parallel_budget(): + merge = _merge_of(*[(f"e{i}", "a", "x", [1]) for i in range(6)]) + decision = R.decide_round( + merge, budget=R.RemediationBudget(max_parallel_verifiers=2), round_no=1 + ) + assert decision.action == "dispatch" + assert len(decision.dispatch) == 6 + assert decision.workers == 2 + + +def test_a_regression_stops_the_cycle_whatever_else_is_true(): + """Evidence that cannot be trusted is not repaired by more rounds.""" + queue = [ + R.RemediationItem(entry="e", arch="a", axis="x", values=("1",), status=R.QUEUE_CERTIFIED) + ] + merge = _merge_of(("e", "a", "x", [1]), ("new", "a", "x", [2]), queue=queue) + decision = R.decide_round(merge, budget=R.RemediationBudget(), round_no=1) + assert decision.stopping + assert decision.verdict == R.STOP_NO_PROGRESS + assert "certified in an earlier round" in decision.detail + + +def test_an_auditor_repeating_itself_stops_the_cycle(): + queue = R.canonicalize(_rows(("e", "a", "x", [1]))) + merge = R.merge_into_queue(queue, queue) + decision = R.decide_round(merge, budget=R.RemediationBudget(), round_no=2) + assert decision.stopping and decision.verdict == R.STOP_NO_PROGRESS + + +def test_spending_the_round_budget_stops_with_its_own_verdict(): + merge = _merge_of(("e", "a", "x", [1])) + decision = R.decide_round(merge, budget=R.RemediationBudget(max_rounds=1), round_no=2) + assert decision.stopping and decision.verdict == R.STOP_BUDGET + + +def test_an_unchanged_unresolved_set_stops_even_when_rows_look_new(): + """Monotone convergence is the contract, not a heuristic.""" + merge = _merge_of(("e", "a", "x", [1])) + signature = R.unresolved_signature(merge.queue) + decision = R.decide_round( + merge, budget=R.RemediationBudget(), round_no=1, previous_signature=signature + ) + assert decision.stopping and decision.verdict == R.STOP_NO_PROGRESS + assert "byte-identical" in decision.detail + + +# ------------------------------------------------------------------ outcomes + + +def test_a_landed_verification_certifies_its_surface(): + queue = R.canonicalize(_rows(("e", "a", "x", [1]))) + updated, failed = R.apply_outcomes(queue, {queue[0].identity: "LANDED"}) + assert updated[0].status == R.QUEUE_CERTIFIED + assert updated[0].attempts == 1 + assert not failed + + +def test_a_failed_receipt_is_terminal_and_reported(): + queue = R.canonicalize(_rows(("e", "a", "x", [1]))) + updated, failed = R.apply_outcomes(queue, {queue[0].identity: "FAILED_RECEIPT"}) + assert updated[0].status == R.QUEUE_FAILED_RECEIPT and len(failed) == 1 + + +def test_a_verification_that_said_nothing_counts_as_blocked_not_as_success(): + """Treating silence as success is how an uncertified surface reaches a merge.""" + queue = R.canonicalize(_rows(("e", "a", "x", [1]))) + updated, failed = R.apply_outcomes(queue, {}) + assert updated[0].status == R.QUEUE_BLOCKED and len(failed) == 1 + + +def test_already_closed_items_are_left_alone(): + queue = [ + R.RemediationItem( + entry="e", arch="a", axis="x", values=("1",), status=R.QUEUE_CERTIFIED, attempts=1 + ) + ] + updated, failed = R.apply_outcomes(queue, {}) + assert updated[0].status == R.QUEUE_CERTIFIED and updated[0].attempts == 1 + assert not failed + + +# ----------------------------------------------------------------- checkpoint + + +def test_the_cycle_survives_a_resume_byte_for_byte(): + """A resume that forgets the queue re-dispatches work it already did.""" + queue = list(R.canonicalize(_rows(("e", "a", "x", [3, 1])), target_commit="c")) + state = R.RemediationState(queue=queue, round_no=2, signature="abc123") + restored = R.RemediationState.from_dict(state.as_dict()) + assert restored.round_no == 2 and restored.signature == "abc123" + assert [i.identity for i in restored.queue] == [i.identity for i in queue] + + +def test_a_checkpoint_from_before_the_cycle_existed_loads_as_empty(): + assert R.RemediationState.from_dict(None).queue == [] + assert R.RemediationState.from_dict({}).round_no == 0 + + +def test_every_stop_verdict_is_re_enterable(): + """Each names a repair made outside the run, and the queue is persisted.""" + for verdict in (R.STOP_NO_PROGRESS, R.STOP_BUDGET, R.STOP_VERIFICATION_FAILED): + assert is_remediable(verdict), verdict + + +# ------------------------------------------------------ the cycle, end to end +# +# Everything above is a pure function. These drive the real orchestrator, so +# they are what catches the wiring: a queue that is never checkpointed, a +# verify dispatch that never runs, a loop that re-audits without re-reading +# delivery. + +import json # noqa: E402 -- the end-to-end half needs the harness below + +import pytest # noqa: E402 + +from agent_flow.workflows.staircase_onboard.state import load_state # noqa: E402 + +from .test_workflow import Script, _install, _workflow, _write_vocabulary # noqa: E402 + +_ROW = { + "entry": "attention/thop_attention", + "arch": "sm_103", + "axis": "batch_size", + "values": ["3", "5", "7", "9", "12", "15", "31"], + "contract_line": "certified at batch sizes 1, 2, 4, 8 on sm_103", +} + + +def _drive(tmp_path, task_file, monkeypatch, *, auditor, smith=None): + """Run the whole workflow with a scripted auditor and smith.""" + workflow = _workflow(tmp_path) + script = Script(None, auditor=auditor, **({"smith": smith} if smith else {})) + calls = _install(monkeypatch, workflow, script) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + return workflow, calls + + +def _state(tmp_path): + return load_state(tmp_path / "ws" / ".staircase_onboard_state.json") + + +def test_a_converging_cycle_certifies_and_then_delivers(tmp_path, task_file, monkeypatch): + """The loop this module exists to create: find, fix, re-check, ship. + + Before this, round one ended the run and a human did the rest. + """ + workflow, calls = _drive( + tmp_path, + task_file, + monkeypatch, + auditor=lambda turn, prompt: ( + {"decision": "UNCOVERED", "summary": "7 values outside", "uncovered": [_ROW]} + if turn == 1 + else {"decision": "CLEAN", "summary": "all covered"} + ), + smith=lambda turn, prompt: { + "decision": "LANDED", + "summary": "certified attention/thop_attention", + "entry": "attention/thop_attention", + }, + ) + state = _state(tmp_path) + assert state.done and state.verdict == "delivered" + assert state.audit_clean + + # Two auditor turns and exactly one verify dispatch. + assert [n for n, _ in calls].count("auditor") == 2 + assert [n for n, _ in calls].count("smith") == 1 + + cycle = R.RemediationState.from_dict(state.remediation) + assert [i.status for i in cycle.queue] == [R.QUEUE_CERTIFIED] + + +def test_seven_uncovered_values_produce_one_smith_not_seven(tmp_path, task_file, monkeypatch): + """The cost saving that motivated the whole design.""" + _, calls = _drive( + tmp_path, + task_file, + monkeypatch, + auditor=lambda turn, prompt: ( + {"decision": "UNCOVERED", "summary": "s", "uncovered": [_ROW]} + if turn == 1 + else {"decision": "CLEAN", "summary": "s"} + ), + smith=lambda turn, prompt: { + "decision": "LANDED", + "summary": "ok", + "entry": "attention/thop_attention", + }, + ) + smith_prompts = [p for n, p in calls if n == "smith"] + assert len(smith_prompts) == 1 + # ...and that one smith was handed the whole surface. + assert "3, 5, 7, 9, 12, 15, 31" in smith_prompts[0] + assert "certified at batch sizes 1, 2, 4, 8" in smith_prompts[0] + + +def test_an_auditor_that_repeats_itself_stops_instead_of_looping(tmp_path, task_file, monkeypatch): + """The failure the bound exists for: a cycle that cannot converge. + + The auditor reports the identical surface forever and the smith keeps + claiming success. Without the convergence rule this runs until someone + notices. + """ + _, calls = _drive( + tmp_path, + task_file, + monkeypatch, + auditor=lambda turn, prompt: { + "decision": "UNCOVERED", + "summary": "same as last time", + "uncovered": [_ROW], + }, + smith=lambda turn, prompt: { + "decision": "LANDED", + "summary": "ok", + "entry": "attention/thop_attention", + }, + ) + state = _state(tmp_path) + assert state.verdict == R.STOP_NO_PROGRESS + # It stopped early rather than spending the round budget. + assert [n for n, _ in calls].count("smith") == 1 + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "certified in an earlier round" in journal + + +def test_a_verification_that_cannot_land_stops_with_its_own_verdict( + tmp_path, task_file, monkeypatch +): + _, calls = _drive( + tmp_path, + task_file, + monkeypatch, + auditor=lambda turn, prompt: {"decision": "UNCOVERED", "summary": "s", "uncovered": [_ROW]}, + smith=lambda turn, prompt: { + "decision": "FAILED_RECEIPT", + "summary": "the test does not pass at batch 31", + "entry": "attention/thop_attention", + }, + ) + state = _state(tmp_path) + assert state.verdict == R.STOP_VERIFICATION_FAILED + assert [n for n, _ in calls].count("smith") == 1 + cycle = R.RemediationState.from_dict(state.remediation) + assert cycle.queue[0].status == R.QUEUE_FAILED_RECEIPT + + +def test_an_uncovered_with_no_structured_rows_keeps_the_old_behaviour( + tmp_path, task_file, monkeypatch +): + """Backward compatible on purpose. + + An auditor that reports prose only has said something real; the run must + not treat it as clean. It ends exactly as it did before the cycle + existed. + """ + _drive( + tmp_path, + task_file, + monkeypatch, + auditor=lambda turn, prompt: { + "decision": "UNCOVERED", + "summary": "batch 257 outside the column", + }, + ) + assert _state(tmp_path).verdict == "audit-uncovered" + + +def test_the_queue_is_checkpointed_so_a_resume_does_not_redo_it(tmp_path, task_file, monkeypatch): + """The manual loop's real cost was re-deriving what was already known.""" + _drive( + tmp_path, + task_file, + monkeypatch, + auditor=lambda turn, prompt: {"decision": "UNCOVERED", "summary": "s", "uncovered": [_ROW]}, + smith=lambda turn, prompt: { + "decision": "FAILED_RECEIPT", + "summary": "no", + "entry": "attention/thop_attention", + }, + ) + raw = json.loads((tmp_path / "ws" / ".staircase_onboard_state.json").read_text()) + assert raw["remediation"]["queue"], "the queue never reached the checkpoint" + assert raw["remediation"]["round_no"] >= 1 + restored = R.RemediationState.from_dict(raw["remediation"]) + assert restored.queue[0].entry == "attention/thop_attention" + + +@pytest.mark.parametrize("budget", [{"max_rounds": 0}, {"max_new_items": -1}]) +def test_a_budget_of_zero_or_less_is_refused_by_the_schema(budget, tmp_path): + """There is no spelling for 'unlimited', by design.""" + import yaml + + from agent_flow.workflows.staircase_onboard.task_schema import ( + TaskSchemaError, + load_and_validate_task_yaml, + ) + + spec = { + "checkpoint_path": str(tmp_path), + "arch": "sm_103", + "parallel": "tp1", + "trtllm_repo_path": str(tmp_path), + "catalog_remediation": budget, + } + path = tmp_path / "t.yaml" + path.write_text(yaml.safe_dump(spec), encoding="utf-8") + with pytest.raises(TaskSchemaError, match="catalog_remediation"): + load_and_validate_task_yaml(path) diff --git a/agent-flow/tests/workflows/staircase_onboard/test_runtime_and_scope.py b/agent-flow/tests/workflows/staircase_onboard/test_runtime_and_scope.py new file mode 100644 index 000000000000..60736ed6aadb --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_runtime_and_scope.py @@ -0,0 +1,669 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the run-scoped declarations: runtime switches, scope, embargo. + +These three arrived together because one real onboarding needed all three, +and each of them is a place where the obvious implementation is wrong in the +same way: it trusts that saying a thing is the same as its being true. A +runtime switch is not a flag (two of them are spelled as negations), an +embargo on a path that does not exist protects nothing, and a text-only +scope that nobody tells the weight loader about reads as a defect. +""" + +from __future__ import annotations + +import textwrap +from pathlib import Path + +import pytest + +from agent_flow.workflows.staircase_onboard.environment import resolve_environment +from agent_flow.workflows.staircase_onboard.gates import runtime_flags, speculation_flags +from agent_flow.workflows.staircase_onboard.prompts import build_staircase_onboard_prompts +from agent_flow.workflows.staircase_onboard.prompts._common import scope_block, source_embargo_block +from agent_flow.workflows.staircase_onboard.task_schema import ( + RUNTIME_SWITCHES, + TaskSchemaError, + load_and_validate_task_yaml, +) + +# Never created. The TMPDIR check reads the declared value rather than +# stat-ing it, and pytest's own tmp_path lives under /tmp, which the check +# rejects by design. +_SCRATCH = "/scratch/staircase-run-tmp" + + +# ----------------------------------------------------------- runtime flags + + +def test_enabling_a_positive_switch_adds_its_flag(): + flags = runtime_flags({"cuda_graph": True, "chunked_prefill": True}) + assert "--use_cuda_graph" in flags + assert "--enable_chunked_prefill" in flags + + +def test_wanting_a_negative_switch_on_emits_nothing(): + """The trap this function exists for. + + Overlap scheduling and block reuse are on by default in the engine, so + the only flags that exist for them turn them *off*. Asked to "enable + block reuse", the obvious move adds ``--disable_kv_cache_reuse`` and + silently gates the opposite of what was requested. + """ + flags = runtime_flags({"overlap_scheduler": True, "kv_cache_reuse": True}) + assert flags == [] + + +def test_turning_a_negative_switch_off_emits_its_disable_flag(): + flags = runtime_flags({"overlap_scheduler": False, "kv_cache_reuse": False}) + assert flags == ["--disable_overlap_scheduler", "--disable_kv_cache_reuse"] + + +def test_an_absent_switch_takes_the_engine_default_not_false(): + """``{}`` must mean "the engine's defaults", not "everything off".""" + assert runtime_flags({}) == [] + + +def test_extra_flags_pass_through_in_order(): + flags = runtime_flags({"extra_flags": ["--mamba_ssm_cache_dtype=float32", "--x=1"]}) + assert flags == ["--mamba_ssm_cache_dtype=float32", "--x=1"] + + +@pytest.mark.parametrize( + ("features", "expected_algo"), + [(["mtp"], "MTP"), (["eagle3"], "Eagle3"), (["MTP"], "MTP"), (["dspark"], "DSPARK")], +) +def test_a_speculator_feature_becomes_spec_decode_flags(features, expected_algo): + flags = speculation_flags(features) + assert flags[0] == f"--spec_decode_algo={expected_algo}" + assert flags[1] == "--spec_decode_max_draft_len=1" + + +def test_a_non_speculative_feature_produces_no_flags(): + """`eplb` changes the target, not the command that runs it.""" + assert speculation_flags(["eplb", "dwdp"]) == [] + assert speculation_flags(None) == [] + + +# --------------------------------------------------------- boot gate mode + + +def _boot_task(tmp_path): + return { + "checkpoint_path": str(tmp_path), + "trtllm_repo_path": str(tmp_path), + "parallel": "tp1", + "runtime": {"cuda_graph": True}, + "features": [], + } + + +def _capture_runner(sink): + def run(command, env=None, timeout=0): + sink.append(command) + from agent_flow.workflows.staircase_onboard.gates import ShellResult + + return ShellResult(returncode=0, stdout="Generated text: Paris") + + return run + + +def test_the_boot_gate_runs_under_require_by_default(tmp_path): + from agent_flow.workflows.staircase_onboard.gates import run_boot_gate + + commands = [] + environment = resolve_environment( + {"environment": {"kind": "local", "python": "/py"}}, world_size=1 + ) + result = run_boot_gate( + environment, task_data=_boot_task(tmp_path), runner=_capture_runner(commands) + ) + assert "TRTLLM_MODELING_V2=require" in commands[0] + assert result.name == "boot" + + +def test_a_control_run_is_named_so_it_cannot_be_mistaken_for_the_gate(tmp_path): + """Two rows both called "boot" would be a journal that lies. + + The control exists to answer "is this the target or the host?", and it + measures the *stock* implementation. Recording it under the real gate's + name is the same class of error as reporting a curve for 34 of 35 + captured batch sizes as though it covered all of them. + """ + from agent_flow.workflows.staircase_onboard.gates import MODE_OFF, run_boot_gate + + commands = [] + environment = resolve_environment( + {"environment": {"kind": "local", "python": "/py"}}, world_size=1 + ) + result = run_boot_gate( + environment, + task_data=_boot_task(tmp_path), + runner=_capture_runner(commands), + mode=MODE_OFF, + ) + assert "TRTLLM_MODELING_V2=off" in commands[0] + assert result.name == "boot[off]" + + +def test_the_control_run_is_otherwise_the_same_command(tmp_path): + """One variable, or it is not a control.""" + from agent_flow.workflows.staircase_onboard.gates import MODE_OFF, run_boot_gate + + require_cmds, off_cmds = [], [] + environment = resolve_environment( + {"environment": {"kind": "local", "python": "/py"}}, world_size=1 + ) + task = _boot_task(tmp_path) + run_boot_gate(environment, task_data=task, runner=_capture_runner(require_cmds)) + run_boot_gate(environment, task_data=task, runner=_capture_runner(off_cmds), mode=MODE_OFF) + + normalize = lambda c: c.replace("TRTLLM_MODELING_V2=off", "TRTLLM_MODELING_V2=require") # noqa: E731 + assert normalize(off_cmds[0]) == require_cmds[0] + + +# ------------------------------------------------------------------- scope + + +def test_text_only_scope_tells_the_loader_to_ignore_vision_tensors(): + block = scope_block("text_only") + assert "ignore the vision tensors" in block + # The specific failure worth naming: a strict loader on a correct target. + assert "asserts every tensor was consumed" in block + + +def test_full_scope_adds_nothing(): + assert scope_block("full") == "" + + +# ---------------------------------------------------------------- embargo + + +def test_an_embargo_names_every_path_and_permits_running_them(): + block = source_embargo_block(["tensorrt_llm/_torch/models/modeling_qwen3_5.py"]) + assert "`tensorrt_llm/_torch/models/modeling_qwen3_5.py`" in block + # Reading is forbidden; running is the point. A block that forbade both + # would throw away the run's only accuracy oracle. + assert "off limits to read" in block + assert "TRTLLM_MODELING_V2=off" in block + + +def test_no_embargo_adds_nothing(): + assert source_embargo_block([]) == "" + + +def test_an_out_of_repo_entry_does_not_claim_this_repository_implements_it(): + """The opening sentence has to split on where the paths point. + + "This checkpoint already has an implementation in this repository" is + false for a sibling vLLM checkout, and a rule whose premise is visibly + wrong is a rule an agent reasons its way around. + """ + block = source_embargo_block(["/lustre/users/x/vllm"]) + assert "in this repository" not in block + assert "another engine's checkout on this filesystem" in block + # Running it is still the point — as a server rather than under + # TRTLLM_MODELING_V2=off, which means nothing to another engine. + assert "black-box oracle is expected" in block + assert "TRTLLM_MODELING_V2=off" not in block + + +def test_both_kinds_of_entry_each_get_their_own_premise_and_permission(): + block = source_embargo_block( + ["tensorrt_llm/_torch/models/modeling_qwen3_5.py", "/lustre/users/x/sglang"] + ) + assert "already has an implementation in this repository" in block + assert "another engine's checkout on this filesystem" in block + assert "TRTLLM_MODELING_V2=off" in block + assert "black-box oracle is expected" in block + + +def test_a_directory_entry_is_said_to_embargo_the_tree_beneath_it(): + """A sibling engine checkout is a tree. + + Enumerating its model files in a spec is asking the spec to go stale. + """ + assert "the whole tree beneath it" in source_embargo_block(["/lustre/users/x/vllm"]) + + +def test_every_role_gets_the_embargo_and_the_scope(): + """A role not told the rule is the role that breaks it. + + The reviewer is the specific worry: its whole job is comparison, so an + un-briefed reviewer will helpfully open the embargoed file to check the + target against it. + """ + bundle = build_staircase_onboard_prompts( + source_embargo=["tensorrt_llm/_torch/models/modeling_qwen3_5.py"], + scope="text_only", + ) + for role in ("plan_drafter", "plan_reviewer", "smith", "assembler", "reviewer", "auditor"): + prompt = getattr(bundle, role) + assert "Embargoed sources" in prompt, role + assert "Scope: text only" in prompt, role + + +# ------------------------------------------------------- slurm_broker kind + + +def _broker_environment(): + return resolve_environment( + { + "environment": { + "kind": "slurm_broker", + "python": "/repo/.venv/bin/python", + "partition": "batch", + "account": "some-account", + "container_image": "registry/image:tag", + "gpus_per_node": 4, + "time_limit": "02:00:00", + "mounts": ["/lustre:/lustre"], + "env": {"TRTLLM_MODELING_V2": "require"}, + } + }, + world_size=1, + ) + + +def test_the_broker_template_is_fully_resolved(): + """No image to pick, no partition to discover — that is the whole point.""" + call = _broker_environment().broker_submit("echo hi") + for fragment in ( + 'image="registry/image:tag"', + 'partition="batch"', + 'account="some-account"', + "gpus_per_node=4", + 'time_limit="02:00:00"', + '"/lustre:/lustre"', + "echo hi", + ): + assert fragment in call, fragment + + +def test_the_broker_prefix_is_in_container_env_not_an_srun_line(): + """Under the broker the transport is the call; the prefix is what runs inside.""" + environment = _broker_environment() + assert environment.prefix == "TRTLLM_MODELING_V2=require" + assert "srun" not in environment.prefix + + +def test_the_broker_block_forbids_the_blocking_call(): + """`slurm_run` reports failure for jobs that are running fine. + + The MCP layer aborts a tool call after 300s of silence and every job + here is longer, so the agent has to be told before it tries. + """ + described = _broker_environment().describe() + assert "Never use the blocking `slurm_run`" in described + assert "300 seconds" in described + # And it must preempt the wrong conclusion from a dead CLI. + assert "do not conclude from their failure that the cluster is down" in described + + +# --------------------------------------------- preflight under the broker + + +def test_preflight_claims_no_devices_under_slurm_broker(tmp_path): + """A login node has no GPUs at all, and that is not a failure. + + The claim check was written for `kind != "slurm"`, which lets + `slurm_broker` through to probe a host where `nvidia-smi` cannot reach + a driver — failing a run that is perfectly fine. + """ + import json + + from agent_flow.workflows.staircase_onboard import preflight + + checkpoint = tmp_path / "ckpt" + checkpoint.mkdir() + (checkpoint / "config.json").write_text( + json.dumps({"architectures": ["Qwen3_5ForConditionalGeneration"]}), encoding="utf-8" + ) + + def _explode(*args, **kwargs): # pragma: no cover - must never be called + raise AssertionError("preflight probed devices under a scheduler-owned allocation") + + result = preflight.run_preflight( + { + "checkpoint_path": str(checkpoint), + "arch": "sm_103", + "parallel": "tp1", + "trtllm_repo_path": str(tmp_path), + "environment": { + "kind": "slurm_broker", + "python": "/opt/venv/bin/python", + "env": {"TMPDIR": _SCRATCH, "CLAUDE_CODE_TMPDIR": _SCRATCH}, + }, + }, + runner=_explode, + ) + assert result.devices == () + # The trust being extended has to be visible in the report, because + # nothing local verified it. + assert any("taken on trust" in note for note in result.notes) + + +# ----------------------------------------------------------- schema wiring + + +def _task_yaml(tmp_path: Path, body: str) -> Path: + path = tmp_path / "task.yaml" + path.write_text(textwrap.dedent(body), encoding="utf-8") + return path + + +def _base(repo: Path, checkpoint: Path) -> str: + return f""" + checkpoint_path: {checkpoint} + trtllm_repo_path: {repo} + arch: sm_103 + parallel: tp1 + accuracy: + benchmark: gsm8k + reference_yaml: tests/x/gsm8k.yaml + environment: + kind: local + python: /usr/bin/python3 + env: + TMPDIR: {_SCRATCH} + """ + + +def test_runtime_defaults_are_filled_in_when_the_block_is_absent(tmp_path): + repo = tmp_path / "repo" + repo.mkdir() + task = _task_yaml(tmp_path, _base(repo, repo)) + data = load_and_validate_task_yaml(task) + assert data["runtime"] == dict(RUNTIME_SWITCHES) | {"extra_flags": []} + assert data["scope"] == "full" + assert data["source_embargo"] == [] + + +def test_a_misspelled_runtime_switch_is_an_error_not_a_shrug(tmp_path): + """Silently ignoring it gates a feature set nobody asked for.""" + repo = tmp_path / "repo" + repo.mkdir() + task = _task_yaml( + tmp_path, + _base(repo, repo) + + """ + runtime: + cuda_graphs: true + """, + ) + with pytest.raises(TaskSchemaError, match="unknown switch"): + load_and_validate_task_yaml(task) + + +def test_an_embargo_on_a_missing_path_is_rejected(tmp_path): + repo = tmp_path / "repo" + repo.mkdir() + task = _task_yaml( + tmp_path, + _base(repo, repo) + + """ + source_embargo: + - tensorrt_llm/_torch/models/modeling_nothing.py + """, + ) + with pytest.raises(TaskSchemaError, match="protects nothing"): + load_and_validate_task_yaml(task) + + +def test_an_embargo_on_a_real_path_is_accepted(tmp_path): + repo = tmp_path / "repo" + (repo / "pkg").mkdir(parents=True) + (repo / "pkg" / "modeling_x.py").write_text("", encoding="utf-8") + task = _task_yaml( + tmp_path, + _base(repo, repo) + + """ + source_embargo: + - pkg/modeling_x.py + """, + ) + assert load_and_validate_task_yaml(task)["source_embargo"] == ["pkg/modeling_x.py"] + + +def test_an_embargo_may_name_a_tree_outside_the_repo(tmp_path): + """The trees most likely to be copied from are not in this repository. + + TensorRT-LLM ships no implementation of a genuinely new checkpoint — that + is the premise of the run — while a sibling vLLM or SGLang checkout on the + same filesystem does, is one `Read` away from every agent, and is in the + run's working set anyway because the accuracy anchor comes from it. + """ + repo = tmp_path / "repo" + repo.mkdir() + sibling = tmp_path / "vllm" + (sibling / "model_executor").mkdir(parents=True) + task = _task_yaml( + tmp_path, + _base(repo, repo) + + f""" + source_embargo: + - {sibling} + """, + ) + assert load_and_validate_task_yaml(task)["source_embargo"] == [str(sibling)] + + +def test_an_absolute_path_inside_the_repo_is_still_rejected(tmp_path): + """The case the original error was actually about, and the fix is spelling. + + A repo-relative list means the same thing from any working directory; an + absolute one that happens to point inside the repo does not. + """ + repo = tmp_path / "repo" + (repo / "pkg").mkdir(parents=True) + (repo / "pkg" / "modeling_x.py").write_text("", encoding="utf-8") + task = _task_yaml( + tmp_path, + _base(repo, repo) + + f""" + source_embargo: + - {repo}/pkg/modeling_x.py + """, + ) + with pytest.raises(TaskSchemaError, match="Spell an in-repo path repo-relative"): + load_and_validate_task_yaml(task) + + +def test_an_out_of_repo_embargo_must_still_exist(tmp_path): + """The must-exist argument is path-shape-independent. + + An embargo naming a path that is not there reads as protection while + protecting nothing. + """ + repo = tmp_path / "repo" + repo.mkdir() + task = _task_yaml( + tmp_path, + _base(repo, repo) + + """ + source_embargo: + - /definitely/not/here/vllm + """, + ) + with pytest.raises(TaskSchemaError, match="protects nothing"): + load_and_validate_task_yaml(task) + + +def test_slurm_broker_requires_an_account(tmp_path): + """The broker rejects every submit that omits it — cheaper to catch here.""" + repo = tmp_path / "repo" + repo.mkdir() + task = _task_yaml( + tmp_path, + f""" + checkpoint_path: {repo} + trtllm_repo_path: {repo} + arch: sm_103 + parallel: tp1 + accuracy: + benchmark: gsm8k + reference_yaml: tests/x/gsm8k.yaml + environment: + kind: slurm_broker + python: /usr/bin/python3 + partition: batch + container_image: registry/image:tag + gpus_per_node: 4 + time_limit: "02:00:00" + """, + ) + with pytest.raises(TaskSchemaError, match="account"): + load_and_validate_task_yaml(task) + + +def test_slurm_broker_requires_gpus_per_node_not_gpus(tmp_path): + repo = tmp_path / "repo" + repo.mkdir() + task = _task_yaml( + tmp_path, + f""" + checkpoint_path: {repo} + trtllm_repo_path: {repo} + arch: sm_103 + parallel: tp1 + accuracy: + benchmark: gsm8k + reference_yaml: tests/x/gsm8k.yaml + environment: + kind: slurm_broker + python: /usr/bin/python3 + partition: batch + account: some-account + container_image: registry/image:tag + time_limit: "02:00:00" + """, + ) + with pytest.raises(TaskSchemaError, match="gpus_per_node"): + load_and_validate_task_yaml(task) + + +# ------------------------------------------- per-unit gating of features + + +def test_core_is_gated_without_the_speculator_flags(tmp_path): + """Core is green *before* the first feature increment — so must its gate be. + + Measured on the real run: the core target refused + `--spec_decode_algo=MTP` through its own fail-fast assert, and the + assembler escalated rather than deleting the assert to make the gate + pass. Deleting it would have been the silent wrong — one-engine MTP is + drafted by the model, so a base model under that config emits no drafts + and the run would have gated a speculator that was not there. + """ + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + workflow = StaircaseOnboardWorkflow(tmp_path / "ws") + workflow.task_data = {"features": ["mtp", "eagle3"]} + assert workflow._features_through("core") == [] + assert workflow._features_through("mtp") == ["mtp"] + assert workflow._features_through("eagle3") == ["mtp", "eagle3"] + + +def test_an_unknown_unit_falls_back_to_the_whole_feature_list(tmp_path): + """Conservative: gate against everything rather than silently against nothing.""" + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + workflow = StaircaseOnboardWorkflow(tmp_path / "ws") + workflow.task_data = {"features": ["mtp"]} + assert workflow._features_through("not-a-unit") == ["mtp"] + + +def test_the_boot_gate_honours_the_features_override(tmp_path): + from agent_flow.workflows.staircase_onboard.gates import run_boot_gate + + commands = [] + environment = resolve_environment( + {"environment": {"kind": "local", "python": "/py"}}, world_size=1 + ) + task = _boot_task(tmp_path) | {"features": ["mtp"]} + + run_boot_gate(environment, task_data=task, runner=_capture_runner(commands)) + assert "--spec_decode_algo=MTP" in commands[0] + + run_boot_gate(environment, task_data=task, runner=_capture_runner(commands), features=[]) + assert "--spec_decode_algo=MTP" not in commands[1] + + +# ------------------------------------------- the eval gates' configuration + + +def test_the_eval_gates_carry_the_declared_runtime_configuration(tmp_path): + """`trtllm-eval` takes none of the quickstart flags, so it needs the YAML. + + Without it the accuracy gate measured the engine's defaults — including + `max_seq_len=262144`, the checkpoint's own maximum — while the anchor had + been measured under the requested configuration. That is not a loose + comparison; it is a different question. + """ + from agent_flow.workflows.staircase_onboard.gates import eval_command + + options = tmp_path / "llm_api.yaml" + options.write_text("kv_cache_config: {dtype: fp8}\n", encoding="utf-8") + environment = resolve_environment( + {"environment": {"kind": "local", "python": "/py"}}, world_size=1 + ) + task = { + "checkpoint_path": "/ckpt", + "parallel": "tp1", + "accuracy": { + "llm_api_options": str(options), + "max_seq_len": 8192, + "max_batch_size": 32, + }, + } + parts = eval_command(task, environment, "gsm8k") + joined = " ".join(parts) + assert f'--extra_llm_api_options="{options}"' in joined + assert "--max_seq_len=8192" in joined + assert "--max_batch_size=32" in joined + # The benchmark is a positional and must stay last of the eval args. + assert parts[-1] == "gsm8k" + + +def test_the_benchmark_stays_last_when_a_dataset_path_is_present(): + """MMLU needs --dataset_path, and it follows the positional benchmark.""" + from agent_flow.workflows.staircase_onboard.gates import eval_command + + environment = resolve_environment( + {"environment": {"kind": "local", "python": "/py"}}, world_size=1 + ) + task = { + "checkpoint_path": "/ckpt", + "parallel": "tp1", + "accuracy": {"dataset_path": "/data/mmlu"}, + } + parts = eval_command(task, environment, "mmlu") + assert parts[-2] == "mmlu" + assert parts[-1] == '--dataset_path="/data/mmlu"' + + +def test_an_llm_api_options_file_that_does_not_exist_is_refused(tmp_path): + repo = tmp_path / "repo" + repo.mkdir() + task = tmp_path / "task.yaml" + task.write_text( + textwrap.dedent(f""" + checkpoint_path: {repo} + trtllm_repo_path: {repo} + arch: sm_103 + parallel: tp1 + accuracy: + benchmark: gsm8k + reference_yaml: tests/x/gsm8k.yaml + llm_api_options: {tmp_path}/nope.yaml + environment: + kind: local + python: /usr/bin/python3 + env: + TMPDIR: {_SCRATCH} + """), + encoding="utf-8", + ) + with pytest.raises(TaskSchemaError, match="measures the engine's defaults"): + load_and_validate_task_yaml(task) diff --git a/agent-flow/tests/workflows/staircase_onboard/test_runtime_index.py b/agent-flow/tests/workflows/staircase_onboard/test_runtime_index.py new file mode 100644 index 000000000000..1b546578d543 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_runtime_index.py @@ -0,0 +1,575 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The runtime map, held to the tree it maps. + +The index is only worth consulting if a wrong entry is a loud failure. These +checks are what make that true, and they are in two groups: + +* **the pointers resolve** — every path and every symbol, against the actual + checkout. This is the whole premise: a map of paths survives a refactor + only while someone checks that it did. +* **the map stays a map** — no code, no signatures, no line numbers. That + restraint is the design, and a rule nothing enforces is a rule that decays. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + +from agent_flow.workflows.staircase_onboard.runtime import ( + ALWAYS_RELEVANT, + INDEX_PATH, + RuntimeArea, + RuntimeIndexError, + areas_for_demands, + format_areas, + load_index, + unmatched_demands, +) +from agent_flow.workflows.staircase_onboard.survey import read_checkpoint_surface + +REPO_ROOT = Path(__file__).resolve().parents[4] +in_checkout = pytest.mark.skipif( + not (REPO_ROOT / "tensorrt_llm" / "_torch").is_dir(), + reason="not inside a TensorRT-LLM checkout; nothing to resolve against", +) + + +@pytest.fixture(scope="module") +def index(): + return load_index() + + +# ------------------------------------------------------------- the pointers + + +@in_checkout +def test_every_path_resolves(index): + """A pointer that has moved is a failure, not a stale sentence.""" + missing = [ + f"{area.name}: {path}" + for area in index + for path in area.pointers + if not (REPO_ROOT / path).exists() + ] + assert not missing, "runtime_index.yaml points at paths that do not exist:\n " + "\n ".join( + missing + ) + + +@in_checkout +def test_every_symbol_resolves(index): + """The other half of "name the symbol, not the line". + + Searched inside the area's own `code` paths rather than the whole tree: + a symbol that exists somewhere else entirely is not evidence that this + entry points anywhere useful. + """ + missing = [] + for area in index: + haystack = "" + for path in area.code: + target = REPO_ROOT / path + if target.is_dir(): + for source in target.rglob("*.py"): + haystack += source.read_text(encoding="utf-8", errors="ignore") + elif target.is_file(): + haystack += target.read_text(encoding="utf-8", errors="ignore") + for symbol in area.symbols: + if symbol not in haystack: + missing.append(f"{area.name}: {symbol} not found under {list(area.code)}") + assert not missing, "runtime_index.yaml names symbols that do not resolve:\n " + "\n ".join( + missing + ) + + +@in_checkout +def test_each_area_anchors_on_at_least_one_defined_symbol(index): + """Substring presence is not enough, and an audit proved it. + + `test_every_symbol_resolves` concatenates the area's files and asks + whether the name appears anywhere. Three entries passed that while + pointing nowhere useful: `MambaHybridCacheManager` matched only inside + *comments* (and as a substring of a different class), `Mapping` matched + only import lines, and `MultimodalPlaceholderMetadata` matched only + because model files import it from elsewhere. + + So every area must anchor on at least one symbol actually *defined* + under its own `code` paths. Not all of them --- a config field is a + legitimate symbol and is defined elsewhere --- but at least one, or the + entry is not pointing at an implementation. + """ + import re as _re + + unanchored = [] + for area in index: + if not area.symbols: + continue + anchored = False + for path in area.code: + target = REPO_ROOT / path + sources = ( + target.rglob("*.py") if target.is_dir() else [target] if target.is_file() else [] + ) + for source in sources: + text = source.read_text(encoding="utf-8", errors="ignore") + for symbol in area.symbols: + pattern = ( + rf"^\s*(?:class|def)\s+{_re.escape(symbol)}\b|^{_re.escape(symbol)}\s*[:=]" + ) + if _re.search(pattern, text, _re.M): + anchored = True + break + if anchored: + break + if anchored: + break + if not anchored: + unanchored.append( + f"{area.name}: none of {list(area.symbols)} is DEFINED under {list(area.code)}" + ) + assert not unanchored, ( + "an area whose symbols only appear in comments or imports points nowhere:\n " + + "\n ".join(unanchored) + ) + + +# --------------------------------------------------------- it stays a map + + +def test_the_index_carries_no_code(index): + """No fences, no signatures. It says where to look, not what is there.""" + text = INDEX_PATH.read_text(encoding="utf-8") + assert "```" not in text, "runtime_index.yaml contains a code block" + # `def foo(` / `class Foo(` are the shapes that creep in first. + offenders = re.findall(r"\b(?:def|class)\s+\w+\s*\(", text) + assert not offenders, f"runtime_index.yaml is transcribing code: {offenders}" + + +def test_the_index_cites_no_line_numbers(index): + """Same rule the catalog and the skills are held to.""" + from .test_skills import BARE_LINE_REFERENCE, LINE_CITATION + + text = INDEX_PATH.read_text(encoding="utf-8") + offenders = LINE_CITATION.findall(text) + BARE_LINE_REFERENCE.findall(text) + assert not offenders, f"runtime_index.yaml cites code by line: {offenders}" + + +@pytest.mark.parametrize("area_name", [a for a in load_index().areas]) +def test_each_summary_stays_high_level(index, area_name): + """A long summary is a summary that has started transcribing. + + The cap is arbitrary and the point is that there is one: an entry that + needs more than a short paragraph is describing the code instead of + locating it. + """ + words = len(index.get(area_name).summary.split()) + assert 10 <= words <= 70, f"{area_name}: summary is {words} words" + + +# ----------------------------------------------------------------- schema + + +def test_an_unknown_field_is_rejected(tmp_path): + bad = tmp_path / "i.yaml" + bad.write_text("x:\n summary: s\n code: [tensorrt_llm]\n wat: 1\n", encoding="utf-8") + with pytest.raises(RuntimeIndexError, match="unknown field"): + load_index(bad) + + +def test_a_missing_required_field_is_rejected(tmp_path): + bad = tmp_path / "i.yaml" + bad.write_text("x:\n summary: s\n", encoding="utf-8") + with pytest.raises(RuntimeIndexError, match="missing required field 'code'"): + load_index(bad) + + +def test_an_empty_index_is_rejected(tmp_path): + bad = tmp_path / "i.yaml" + bad.write_text("{}\n", encoding="utf-8") + with pytest.raises(RuntimeIndexError): + load_index(bad) + + +# ------------------------------------------------- the join to the survey + + +def _surface(**config): + base = dict( + architectures=["XForCausalLM"], + num_attention_heads=32, + num_key_value_heads=8, + hidden_act="silu", + rms_norm_eps=1e-5, + ) + return read_checkpoint_surface(base | config).demand_names + + +def test_a_dense_model_is_not_handed_the_exotic_areas(index): + """The reason this is a query and not a document. + + Loading mamba, KDA, mHC and disaggregation for a dense Llama costs + context and offers the agent detail it has no business using. + """ + names = {a.name for a in areas_for_demands(_surface(), index=index)} + assert {"attention", "mlp", "normalization", "quantization"} <= names + assert not names & {"mamba", "gated_deltanet", "kda", "moe", "mla", "speculation"} + + +def test_a_hybrid_recurrent_model_gets_the_recurrent_areas(index): + names = { + a.name + for a in areas_for_demands( + _surface( + mamba_d_state=128, conv_kernel=4, layer_types=["linear_attention", "full_attention"] + ), + index=index, + ) + } + assert {"mamba", "gated_deltanet"} <= names + assert "moe" not in names + + +def test_an_mla_moe_model_gets_those_and_not_the_recurrent_ones(index): + names = { + a.name + for a in areas_for_demands( + _surface( + kv_lora_rank=512, q_lora_rank=1536, n_routed_experts=256, num_experts_per_tok=8 + ), + index=index, + ) + } + assert {"mla", "moe"} <= names + assert not names & {"mamba", "gated_deltanet"} + + +def test_the_always_relevant_areas_are_always_there(index): + """A target always sits in a topology and is always captured or not.""" + names = {a.name for a in areas_for_demands([], index=index)} + assert set(ALWAYS_RELEVANT) <= names + for name in ALWAYS_RELEVANT: + assert index.get(name) is not None, name + + +def test_areas_come_back_in_index_order(index): + """Two runs on one checkpoint must read the same document.""" + demands = _surface(n_routed_experts=8, num_experts_per_tok=2) + first = [a.name for a in areas_for_demands(demands, index=index)] + assert first == [a.name for a in areas_for_demands(reversed(list(demands)), index=index)] + + +def test_every_demand_the_survey_can_emit_is_accounted_for(index): + """The drift guard between `survey.py` and this index. + + A demand no area claims is not automatically wrong — some are genuinely + informational. But it must be *visible*, because the alternative is a + checkpoint feature that silently maps to no engine knowledge at all. + """ + emitted = set() + for config in ( + dict(kv_lora_rank=512, q_lora_rank=1536), + dict(n_routed_experts=8, num_experts_per_tok=2, mlp_only_layers=[0]), + dict(mamba_d_state=128, conv_kernel=4, layer_types=["a", "b"]), + dict(sliding_window=4096, attn_logit_softcapping=50.0, sinks=4), + dict(num_nextn_predict_layers=1, head_dim=256, use_qk_norm=True), + dict(rope_scaling={"rope_type": "yarn"}, partial_rotary_factor=0.25), + dict(quantization_config={"quant_algo": "NVFP4"}), + ): + emitted.update(_surface(**config)) + + unclaimed = set(unmatched_demands(emitted, index=index)) + # Known-informational: these describe the stack rather than name an area. + expected_unclaimed = {"stack.hybrid"} + surprising = unclaimed - expected_unclaimed + assert not surprising, ( + "survey.py emits demands no runtime area claims — either add the area " + f"or mark the demand informational: {sorted(surprising)}" + ) + + +# --------------------------------------------------------------- rendering + + +def test_the_rendered_block_hands_over_pointers_not_content(index): + block = format_areas(areas_for_demands(_surface(), index=index)) + assert "Read the source" in block + assert "no line numbers" in block + assert "tests/unittest/_torch" in block # the test pointer is half the value + + +def test_unmatched_demands_are_surfaced_as_a_signal(): + block = format_areas( + (RuntimeArea(name="x", summary="s", code=("tensorrt_llm",)),), + unmatched=("something.nobody.claims",), + ) + assert "Demands no area claims" in block + assert "something.nobody.claims" in block + + +def test_an_empty_selection_renders_nothing(): + assert format_areas(()) == "" + + +# ------------------------------------------------- reaching the plan phase + + +def test_the_map_reaches_the_plan_drafter(tmp_path, monkeypatch): + """A map nobody reads is a catalogue. + + And it must be *narrowed*: the drafter for a dense checkpoint should not + be handed the recurrent areas. + """ + import json + + import yaml + + from .test_delivery import _build_repo + from .test_workflow import Script, _install, _workflow, passing_shell + + repo = _build_repo(tmp_path) + checkpoint = tmp_path / "ckpt" + checkpoint.mkdir() + checkpoint.joinpath("config.json").write_text( + json.dumps( + { + "architectures": ["LlamaForCausalLM"], + "num_attention_heads": 32, + "num_key_value_heads": 8, + "hidden_act": "silu", + "rms_norm_eps": 1e-5, + } + ) + ) + task = tmp_path / "task.yaml" + task.write_text( + yaml.safe_dump( + { + "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/t", + "CLAUDE_CODE_TMPDIR": "/scratch/t", + }, + }, + "accuracy": { + "benchmark": "gsm8k", + "reference_yaml": "r.yaml", + "model_key": "openai/gpt-oss-120b", + "tolerance": 5.0, + }, + } + ), + encoding="utf-8", + ) + + workflow = _workflow(tmp_path, shell_runner=passing_shell()) + calls = _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task) + + prompt = next(p for name, p in calls if name == "plan_drafter") + assert "## Runtime map" in prompt + assert "tensorrt_llm/_torch/attention" in prompt + assert "tests/unittest/_torch/attention" in prompt + # Narrowed, not a catalogue. + assert "kimi_kda" not in prompt + assert "modules/mamba" not in prompt + + +def test_the_three_writing_roles_are_told_the_map_protocol(): + """Only the roles that read the engine to write code need it. + + `SOURCE_BOUNDARY` binds plan_drafter, smith and assembler — the three + that derive a target. A reviewer handed a "go read the engine" protocol + is a reviewer that will go read the engine instead of the diff. + """ + from agent_flow.workflows.staircase_onboard.prompts import DEFAULT_PROMPTS + + for role in ("plan_drafter", "smith", "assembler"): + text = getattr(DEFAULT_PROMPTS, role) + assert "runtime map" in text, role + assert "Take the location from" in text, role + for role in ("reviewer", "auditor", "plan_reviewer"): + assert "runtime map" not in getattr(DEFAULT_PROMPTS, role), role + + +# ---------------------------------------------------------------- gating + + +def test_every_area_declares_how_it_is_proven(index): + """Required, because "configured" once shipped as "verified". + + Three of seven features on the first real run were enabled and never + witnessed. An area that cannot say how it is proven is an area that will + be claimed on the strength of a gate that cannot see it. + """ + from agent_flow.workflows.staircase_onboard.runtime import GATES + + for area in index: + assert area.gating is not None, f"{area.name} declares no gating" + assert area.gating.gate in GATES, f"{area.name}: gate={area.gating.gate!r}" + assert len(area.gating.signal.split()) >= 5, f"{area.name}: signal is too thin" + assert len(area.gating.blind.split()) >= 3, f"{area.name}: blind is too thin" + + +def test_the_areas_needing_a_witness_are_the_ones_we_know_are_unverified(index): + """A cross-check against measured reality, not a tautology. + + overlap scheduler, chunked prefill and block reuse shipped unverified. + They live in `scheduling` and `kv_cache`, and both must say `none` --- if + either ever claims a standard gate covers it, that claim is the bug. + """ + assert index.get("scheduling").gating.needs_a_witness + assert index.get("kv_cache").gating.needs_a_witness + # And the areas the gates genuinely do cover must not claim otherwise. + assert index.get("attention").gating.gate == "accuracy" + assert index.get("speculation").gating.gate == "acceptance" + + +def test_speculation_records_that_both_ordinary_gates_are_blind(index): + """The single most expensive blind spot in the workflow.""" + blind = index.get("speculation").gating.blind + assert "Boot and accuracy" in blind + assert "distribution" in blind + + +def test_gating_reaches_the_rendered_block(index): + """A blind spot nobody is shown is a blind spot nobody avoids.""" + block = format_areas(areas_for_demands(_surface(num_nextn_predict_layers=1), index=index)) + assert "**proven by**" in block + assert "**blind spot**" in block + assert "acceptance gate" in block + + +def test_a_bad_gate_name_is_rejected(tmp_path): + bad = tmp_path / "i.yaml" + bad.write_text( + "x:\n summary: s\n code: [tensorrt_llm]\n" + " gating: {gate: sometimes, signal: a b c d e, blind: a b c}\n", + encoding="utf-8", + ) + with pytest.raises(RuntimeIndexError, match="must be one of"): + load_index(bad) + + +def test_incomplete_gating_is_rejected(tmp_path): + bad = tmp_path / "i.yaml" + bad.write_text( + "x:\n summary: s\n code: [tensorrt_llm]\n gating: {gate: boot}\n", encoding="utf-8" + ) + with pytest.raises(RuntimeIndexError, match="missing 'signal'"): + load_index(bad) + + +# ------------------------------------------- the features must have a home + + +def test_every_feature_the_schema_accepts_has_a_runtime_area(index): + """`features: [eplb]` must resolve to somewhere an agent can read. + + This caught two real gaps: `eplb` and `dwdp` had no area at all, and the + `multimodal` feature could never reach its area because the area was + spelled `multimodality`. + """ + from agent_flow.workflows.staircase_onboard.task_schema import FEATURE_SKILLS + + # The speculative family all live in one area, and `disagg` is the + # feature-list spelling of the `disaggregation` area. Native KV storage + # and encoder replay are pool and prefix-reuse questions, so `kv_cache`; + # decoder replay truncates a sliding window, so `sparse_attention`. + aliases = { + "mtp": "speculation", + "eagle3": "speculation", + "dflash": "speculation", + "dspark": "speculation", + "disagg": "disaggregation", + "kv_quant": "kv_cache", + "encoder_replay": "kv_cache", + "decoder_replay": "sparse_attention", + } + missing = [ + feature for feature in FEATURE_SKILLS if index.get(aliases.get(feature, feature)) is None + ] + assert not missing, f"features with no runtime area: {missing}" + + +# --------------------------------------------------------------- overview + + +def test_the_overview_names_only_areas_that_exist(index): + """A flow stage pointing at a phantom area is a map with a dead end.""" + unknown = [ + f"{stage.stage}: {name}" + for stage in index.overview.flow + for name in stage.areas + if index.get(name) is None + ] + assert not unknown, f"_overview.flow names areas that do not exist: {unknown}" + + +def test_every_area_is_placed_somewhere_in_the_flow(index): + """An area in no stage is an area nobody knows who owns.""" + placed = {name for stage in index.overview.flow for name in stage.areas} + orphans = sorted(set(index.areas) - placed) + assert not orphans, f"areas absent from _overview.flow: {orphans}" + + +def test_the_forward_stage_is_where_a_target_lives(index): + """The one stage a modeling_v2 target replaces.""" + forward = next(s for s in index.overview.flow if "forward" in s.stage) + assert {"attention", "mlp", "normalization"} <= set(forward.areas) + # The things a target must NOT try to do. + assert "scheduling" not in forward.areas + assert "kv_cache" not in forward.areas + assert "speculation" not in forward.areas + + +def test_the_overview_leads_the_rendered_block(index): + block = format_areas(areas_for_demands(_surface(), index=index), overview=index.overview) + assert block.index("How the engine fits together") < block.index("## Runtime map") + assert "THE TARGET" in block + + +def test_a_reserved_key_that_is_not_the_overview_is_rejected(tmp_path): + bad = tmp_path / "i.yaml" + bad.write_text("_wat:\n a: 1\n", encoding="utf-8") + with pytest.raises(RuntimeIndexError, match="unknown reserved key"): + load_index(bad) + + +# ------------------------------------------------------------ jit/autotune + + +def test_jit_is_always_relevant(index): + """Nothing in a config.json says "affected by tactic selection". + + Every model is, so it cannot hang off a demand. + """ + from agent_flow.workflows.staircase_onboard.runtime import ALWAYS_RELEVANT + + assert "jit_autotune" in ALWAYS_RELEVANT + assert "jit_autotune" in {a.name for a in areas_for_demands([], index=index)} + + +def test_jit_records_that_a_cold_receipt_describes_a_different_path(index): + """The systemic consequence, stated where a smith taking a receipt sees it. + + A receipt taken in a cold process certifies tactics that serving --- which + warms the autotuner during startup --- will not select. + """ + blind = index.get("jit_autotune").gating.blind + assert "cold" in blind and "receipt" in blind + signal = index.get("jit_autotune").gating.signal + assert "enable_autotuner" in signal + assert "_run_autotuner_warmup" in signal diff --git a/agent-flow/tests/workflows/staircase_onboard/test_skills.py b/agent-flow/tests/workflows/staircase_onboard/test_skills.py new file mode 100644 index 000000000000..8d524def62ee --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_skills.py @@ -0,0 +1,326 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the packaged agent skills. + +The skills are ~7,000 lines of prose ported from a repository with a +different layout. Prose has no compiler, so the failure mode is a skill that +confidently tells an agent to write a file into a directory that no longer +exists — and the agent then spends a turn discovering that. These checks are +the compiler that prose does not have. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest +import yaml + +import agent_flow.workflows.staircase_onboard as pkg +from agent_flow.workflows.staircase_onboard.task_schema import FEATURE_SKILLS + +SKILLS_DIR = Path(pkg.__file__).parent / "skills" +REPO_ROOT = Path(__file__).resolve().parents[4] + +# Retired by the port. A skill naming one is telling an agent to use a +# product that no longer exists. +RETIRED = ( + "uv run", + "ty check", + "1.3.0rc25", + ".allow-trtllm", + "bench/accuracy.py", + "bench/perf.py", +) + +# Narrow, per-file exceptions to the retired-product check. The op inventory +# is *generated*, and its provenance line records the engine version the +# snapshot was taken at — that is a stamp on an observation, not a pin the +# targets assert against, and the file's header says so. Deleting it would +# make the snapshot undateable. +RETIRED_EXCEPTIONS = { + ("staircase-kernel-feasibility/references/trtllm-ops.md", "1.3.0rc25"), +} + +# Repo-relative path prefixes worth resolving. Anything quoted in backticks +# that starts with one of these should exist in the checkout. +CHECKED_PREFIXES = ("tensorrt_llm/", "tests/", "examples/", "jenkins/") + +# Paths a skill may name that do not exist yet, each with a reason. Keeping +# this list short is the point: it is the inventory of what the port still +# owes, not a place to silence the check. +KNOWN_ABSENT = { + # Ported separately; referenced by the kernel skills. + "agent_flow/workflows/staircase_onboard/tools/list_trtllm_ops.py", +} + + +def _skill_dirs() -> list[Path]: + """Every packaged skill, and nothing else that happens to be a directory. + + Dot-directories are excluded because tooling leaves them here. A Claude + Code session writes an empty `.claude/.cc-writes/` into whatever tree it + touches, and without this every check below reported the packaged skills + as malformed — in a working tree, on a file git does not even track, + which is the most confusing possible way for a suite to go red. + """ + if not SKILLS_DIR.is_dir(): + return [] + return sorted(p for p in SKILLS_DIR.iterdir() if p.is_dir() and not p.name.startswith(".")) + + +def _markdown_files() -> list[Path]: + return sorted(p for skill in _skill_dirs() for p in skill.rglob("*.md")) + + +pytestmark = pytest.mark.skipif( + not _skill_dirs(), reason="no skills packaged yet (the knowledge port has not landed)" +) + + +def _front_matter(path: Path) -> dict: + text = path.read_text(encoding="utf-8") + match = re.match(r"^---\n(.*?)\n---\n", text, re.DOTALL) + assert match, f"{path} has no YAML front matter" + return yaml.safe_load(match.group(1)) or {} + + +def test_every_skill_directory_has_a_skill_md(): + for skill in _skill_dirs(): + assert (skill / "SKILL.md").is_file(), skill.name + + +def test_every_skill_is_prefixed(): + """The prefix is what keeps these from colliding with the neighbouring roots.""" + for skill in _skill_dirs(): + assert skill.name.startswith("staircase-"), skill.name + + +def test_front_matter_name_matches_the_directory(): + for skill in _skill_dirs(): + data = _front_matter(skill / "SKILL.md") + assert data.get("name") == skill.name, skill.name + + +def test_every_skill_has_a_description(): + """The description is what the model matches on when deciding to load it.""" + for skill in _skill_dirs(): + description = _front_matter(skill / "SKILL.md").get("description", "") + assert len(description) > 60, skill.name + + +def test_no_skill_names_a_retired_product(): + offences = [] + for path in _markdown_files(): + relative = str(path.relative_to(SKILLS_DIR)) + text = path.read_text(encoding="utf-8") + for retired in RETIRED: + if (relative, retired) in RETIRED_EXCEPTIONS: + continue + # Word-bounded on the left, or "ty check" matches "identity + # check" and a porting agent spends a turn rewriting a sentence + # that was already correct. + pattern = re.compile(rf"(?{}") or token.endswith("/"): + continue + candidates = [REPO_ROOT / token] + if not Path(token).suffix: + # A catalog entry is habitually named without an extension + # (`catalog/comm/allgather`), because the entry is the + # contract *and* the wrapper *and* the test. Resolve either. + candidates += [REPO_ROOT / f"{token}.py", REPO_ROOT / f"{token}.md"] + if not any(candidate.exists() for candidate in candidates): + missing.append(f"{path.relative_to(SKILLS_DIR)}: {token}") + assert not missing, "skills name paths that do not exist:\n " + "\n ".join( + sorted(set(missing)) + ) + + +def test_cross_skill_references_resolve(): + """A skill pointing at a sibling must point at one that is packaged.""" + names = {skill.name for skill in _skill_dirs()} + dangling: list[str] = [] + for path in _markdown_files(): + text = path.read_text(encoding="utf-8") + for referenced in re.findall(r"staircase-[a-z0-9-]+", text): + if referenced not in names: + dangling.append(f"{path.relative_to(SKILLS_DIR)}: {referenced}") + assert not dangling, "dangling skill references:\n " + "\n ".join(sorted(set(dangling))) + + +def test_every_feature_has_the_skill_the_schema_promises(): + """`features: [mtp]` must resolve to a skill the assembler can actually load.""" + names = {skill.name for skill in _skill_dirs()} + for feature, skill in FEATURE_SKILLS.items(): + assert skill in names, f"feature {feature!r} maps to {skill!r}, which is not packaged" + + +def test_the_assembler_always_loads_modeling_core(): + assert (SKILLS_DIR / "staircase-modeling-core" / "SKILL.md").is_file() + + +def test_the_two_smith_modes_are_packaged(): + for name in ("staircase-kernel-onboard", "staircase-kernel-verify"): + assert (SKILLS_DIR / name / "SKILL.md").is_file(), name + + +def test_port_markers_are_visible_rather_than_silent(): + """`` is a deliberate to-do, and it should be findable. + + This does not fail on them — leaving one is the correct action when a fact + could not be confirmed. It asserts they carry a note, so the next pass + knows what was being questioned. + """ + for path in _markdown_files(): + text = path.read_text(encoding="utf-8") + for match in re.finditer(r"", text, re.DOTALL): + assert match.group(1).strip(), f"{path.relative_to(SKILLS_DIR)}: bare PORT marker" + + +# --------------------------------------------- how engine knowledge is written + +#: A skill's job is to tell an agent what the engine does and where to look — +#: not to reproduce what the engine's code says. Reproducing it is the trap: +#: engine code moves every day, prose does not, and prose that has quietly +#: stopped being true is worse than prose that was never written, because the +#: agent believes it and spends a turn on it. +#: +#: So: **describe the behaviour, name the symbol, let the agent read the +#: source.** A file path and a function, class or field name survive a +#: refactor. A line number does not — measured, not assumed: the sibling +#: catalog held 7 line citations and all 7 pointed at the wrong line within +#: days of being written, while every symbol and path beside them still +#: resolved. `tests/unittest/_torch/modeling_v2/test_modeling_v2_no_stale_claims.py` +#: enforces the same rule on that tree; this is the skills half. +LINE_CITATION = re.compile(r"\b[A-Za-z0-9_./-]+\.(?:py|h|hpp|cpp|cu|cuh):\d+(?:[-,]\d+)*") +BARE_LINE_REFERENCE = re.compile(r"`:\d+`") + +#: The one exception, and not a loophole: a line number inside quoted engine +#: output is part of the message the engine prints, so it is what an agent +#: greps for on hitting it rather than a direction to navigate. +DIAGNOSTIC_MARKERS = ("Assertion failed", "RuntimeError", "Error:", "error:", "[ERROR]") +DIAGNOSTIC_WINDOW = 3 + + +#: A citation inside a ``double-backtick`` span is being *shown*, not made -- +#: markdown's idiom for displaying a literal. The rule's own documentation has +#: to be able to print the bad form next to the good one, and a rule that +#: cannot state its own counter-example is a rule nobody can teach. +_SHOWN_LITERAL = re.compile(r"``.*?``") + + +def _is_shown_as_an_example(line: str, start: int, end: int) -> bool: + return any(m.start() <= start and end <= m.end() for m in _SHOWN_LITERAL.finditer(line)) + + +def _quotes_a_diagnostic(lines: list[str], index: int) -> bool: + lo = max(0, index - DIAGNOSTIC_WINDOW) + hi = min(len(lines), index + DIAGNOSTIC_WINDOW + 1) + return any(marker in line for line in lines[lo:hi] for marker in DIAGNOSTIC_MARKERS) + + +def test_no_skill_cites_engine_code_by_line_number(): + """Knowledge with a shelf life, kept out of the knowledge layer.""" + offences: list[str] = [] + for path in _markdown_files(): + lines = path.read_text(encoding="utf-8").splitlines() + for index, line in enumerate(lines): + bare = BARE_LINE_REFERENCE.search(line) + if bare and not _is_shown_as_an_example(line, bare.start(), bare.end()): + offences.append(f"{path.relative_to(SKILLS_DIR)}:{index + 1}: bare `:NNN`") + hit = LINE_CITATION.search(line) + if ( + hit is None + or _quotes_a_diagnostic(lines, index) + or _is_shown_as_an_example(line, hit.start(), hit.end()) + ): + continue + offences.append(f"{path.relative_to(SKILLS_DIR)}:{index + 1}: {hit.group(0)}") + assert not offences, ( + "skills must name the file and the symbol, never the line — a line number " + "is wrong by the next merge and sends the agent to whatever now occupies " + "it:\n " + "\n ".join(sorted(set(offences))) + ) + + +def test_the_line_citation_rule_actually_bites(): + """A guard nobody has seen fail is a guard nobody knows works. + + Both halves matter: it must flag navigation, and it must leave a quoted + diagnostic alone. The second is the one that would quietly over-fire and + get the whole rule switched off. + """ + navigation = ["warmup is at `_torch/pyexecutor/model_engine.py:1103`"] + assert LINE_CITATION.search(navigation[0]) and not _quotes_a_diagnostic(navigation, 0) + + symptom = ["raises `RuntimeError: bad dtype (../tensorrt_llm/runtime/torchUtils.h:113)`"] + assert LINE_CITATION.search(symptom[0]) and _quotes_a_diagnostic(symptom, 0) + + # A pasted traceback wraps, putting the marker a line above the citation. + wrapped = [ + "```", + "RuntimeError: [TensorRT-LLM][ERROR] Assertion failed:", + "Failed: MPI error", + "(../tensorrt_llm/runtime/utils/mpiUtils.cpp:277)", + ] + assert _quotes_a_diagnostic(wrapped, 3) + + good = ["warmup calls `_run_autotuner_warmup` (`_torch/pyexecutor/model_engine.py`)"] + assert LINE_CITATION.search(good[0]) is None + + +def test_the_entry_spec_separates_pool_geometry_from_batch_geometry(): + """The one-line omission that cost nine entries. + + Eight different smiths, over two days, each sized a state pool to their + test's batch -- the natural reading of "test it at the shape it is called + with", and the one arrangement the runtime never produces. An audit found + all nine outside the engine's real width of 461 at once. + + The spec has to say this explicitly because the correct instinct produces + the wrong answer here: the batch comes from the executor, the pool's slot + count from the cache manager's capacity planning, and the two are + unrelated numbers. + """ + spec = (SKILLS_DIR / "staircase-kernel-onboard" / "references" / "entry-spec.md").read_text( + encoding="utf-8" + ) + assert "never from your batch" in spec + # and the reason, not just the rule -- a rule with no failure attached is + # the first thing a hurried reader skips + assert "pool_slots == batch_size" in spec + assert "certified column" in spec diff --git a/agent-flow/tests/workflows/staircase_onboard/test_smith_deferral.py b/agent-flow/tests/workflows/staircase_onboard/test_smith_deferral.py new file mode 100644 index 000000000000..59a2897f9395 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_smith_deferral.py @@ -0,0 +1,487 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A smith failure only a later unit pays for must not stop core. + +Measured on the DeepSeek-V4.1 run: the plan queued 33 dispatches, six of them +needed only by later feature units, and the smith phase closed the run +`smith-blocked` on ANY failure -- before core, which needed none of those six, +had started. Queuing them up front is still right, since they run alongside +core's for free; what changes is what their failure blocks. + +Same harness as `test_workflow.py`: no agent is constructed, and the fake +writes the typed decision the real role would. The smith here is scripted per +target and per attempt, so a test can say "fails once, lands on the retry". +""" + +from __future__ import annotations + +import json +import re + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard.state import ( + ITEM_DEFERRED, + ITEM_DONE, + ITEM_FAILED, + STAGE_ASSEMBLE, + STAGE_SMITH, + load_state, +) + +from .test_workflow import Script, _install, _plan_then_vocabulary, _workflow + +_STATE = ".staircase_onboard_state.json" + + +def _features(task_file, features): + spec = yaml.safe_load(task_file.read_text()) + spec["features"] = list(features) + task_file.write_text(yaml.safe_dump(spec), encoding="utf-8") + + +def _smith(outcomes: dict[str, list[str]] | None = None): + """A smith scripted per target and per attempt; ``LANDED`` when unscripted. + + It stamps the dispatch number it was given as its ``iteration``, which is + what lets the orchestrator tell its entry apart from a concurrent + sibling's -- the real collector does the same. + """ + outcomes = outcomes or {} + attempts: dict[str, int] = {} + + def smith(turn, prompt): + n = int(re.search(r"Dispatch (\d+):", prompt).group(1)) + target = re.search(r"the catalog entry `([^`]+)`", prompt).group(1) + attempts[target] = attempts.get(target, 0) + 1 + script = outcomes.get(target, ["LANDED"]) + decision = script[min(attempts[target], len(script)) - 1] + return {"iteration": n, "decision": decision, "summary": f"{decision} {target}"} + + return smith + + +def _smith_calls(calls, target=None): + return [ + (index, prompt) + for index, (role, prompt) in enumerate(calls) + if role == "smith" and (target is None or f"catalog entry `{target}`" in prompt) + ] + + +def _assembler_index(calls, unit): + return next( + index + for index, (role, prompt) in enumerate(calls) + if role == "assembler" and f"Unit: **{unit}**" in prompt + ) + + +def _run(tmp_path, task_file, monkeypatch, queue, smith, **roles): + workflow = _workflow(tmp_path) + calls = _install(monkeypatch, workflow, Script(None, smith=smith, **roles)) + _plan_then_vocabulary(workflow, tmp_path, monkeypatch, queue) + with workflow: + workflow.run(task_file) + return calls, load_state(tmp_path / "ws" / _STATE) + + +def _row(state, n): + return next(row for row in state.queue if row["n"] == n) + + +# ----------------------------------------------------------- core still blocks + + +def test_a_core_failure_still_blocks_before_assembly(tmp_path, task_file, monkeypatch): + """An untagged row is core's, and core cannot assemble without it.""" + _features(task_file, ["mtp"]) + queue = [ + {"n": 1, "mode": "onboard", "target": "norm/core_only"}, + {"n": 2, "mode": "onboard", "target": "moe/drafter_moe", "unit": "mtp"}, + ] + calls, state = _run( + tmp_path, task_file, monkeypatch, queue, _smith({"norm/core_only": ["BLOCKED"]}) + ) + + assert "assembler" not in [role for role, _ in calls] + assert state.verdict == "smith-blocked" + assert _row(state, 1)["status"] == ITEM_FAILED + + +def test_a_dependency_of_a_core_row_is_core_whatever_its_tag(tmp_path, task_file, monkeypatch): + """Core's own entries were built on it, so deferring it would start core without it.""" + _features(task_file, ["mtp"]) + queue = [ + {"n": 1, "mode": "onboard", "target": "quant/shared_packing", "unit": "mtp"}, + {"n": 2, "mode": "onboard", "target": "attn/core_logits", "depends_on": [1]}, + ] + calls, state = _run( + tmp_path, task_file, monkeypatch, queue, _smith({"quant/shared_packing": ["BLOCKED"]}) + ) + + assert state.verdict == "smith-blocked" + assert "assembler" not in [role for role, _ in calls] + # The core row that waits on it is never dispatched: there is nothing to + # build it on. + assert not _smith_calls(calls, "attn/core_logits") + + +# ------------------------------------------------- feature-only rows defer + + +def test_a_feature_only_failure_is_deferred_and_core_proceeds(tmp_path, task_file, monkeypatch): + _features(task_file, ["mtp"]) + queue = [ + {"n": 1, "mode": "onboard", "target": "norm/core_only"}, + {"n": 2, "mode": "onboard", "target": "moe/drafter_moe", "unit": "mtp"}, + ] + calls, state = _run( + tmp_path, + task_file, + monkeypatch, + queue, + _smith({"moe/drafter_moe": ["BLOCKED", "LANDED"]}), + ) + + assert "Unit: **core**" in calls[_assembler_index(calls, "core")][1] + assert state.units_done == ["core", "mtp"] + assert state.verdict == "delivered" + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "smith dispatches deferred" in journal + assert "2:moe/drafter_moe → mtp" in journal + + +def test_a_deferred_row_is_retried_once_before_its_unit(tmp_path, task_file, monkeypatch): + """After core is green, before the first mtp turn -- and told what failed.""" + _features(task_file, ["mtp"]) + queue = [ + {"n": 1, "mode": "onboard", "target": "norm/core_only"}, + {"n": 2, "mode": "onboard", "target": "moe/drafter_moe", "unit": "mtp"}, + ] + calls, state = _run( + tmp_path, + task_file, + monkeypatch, + queue, + _smith({"moe/drafter_moe": ["BLOCKED", "LANDED"]}), + ) + + attempts = _smith_calls(calls, "moe/drafter_moe") + assert len(attempts) == 2 + retry_at, retry_prompt = attempts[1] + assert _assembler_index(calls, "core") < retry_at < _assembler_index(calls, "mtp") + assert "This is a retry" in retry_prompt + assert "smith decided BLOCKED" in retry_prompt + # Pointed at the report the first attempt left, by its real path. + report = re.search(r"previous attempt's report is `([^`]+)`", retry_prompt).group(1) + assert "smith-onboard-moe-drafter-moe" in report + first_report = open(report, encoding="utf-8").read() + assert "Dispatch 2: **onboard** the catalog entry `moe/drafter_moe`" in first_report + assert "This is a retry" not in first_report + # The first attempt was an ordinary dispatch: no retry note. + assert "This is a retry" not in attempts[0][1] + assert _row(state, 2)["status"] == ITEM_DONE + + +def test_a_retry_that_fails_again_closes_the_run_naming_its_unit(tmp_path, task_file, monkeypatch): + _features(task_file, ["mtp"]) + queue = [ + {"n": 1, "mode": "onboard", "target": "norm/core_only"}, + {"n": 2, "mode": "onboard", "target": "moe/drafter_moe", "unit": "mtp"}, + ] + calls, state = _run( + tmp_path, task_file, monkeypatch, queue, _smith({"moe/drafter_moe": ["BLOCKED"]}) + ) + + assert state.verdict == "smith-blocked:mtp" + # Core was built and gated; the stop is at mtp, before its first turn. + assert state.units_done == ["core"] + assert [p for r, p in calls if r == "assembler" and "Unit: **mtp**" in p] == [] + assert len(_smith_calls(calls, "moe/drafter_moe")) == 2, "one retry, not more" + assert _row(state, 2)["status"] == ITEM_FAILED + + +def test_a_row_waiting_on_a_deferred_row_is_deferred_undispatched(tmp_path, task_file, monkeypatch): + """It is dispatched for the first time at its unit, after what it waits on lands.""" + _features(task_file, ["mtp", "eplb"]) + queue = [ + {"n": 1, "mode": "onboard", "target": "norm/core_only"}, + {"n": 2, "mode": "onboard", "target": "moe/drafter_moe", "unit": "mtp"}, + { + "n": 3, + "mode": "onboard", + "target": "moe/balancer_slots", + "unit": "eplb", + "depends_on": [2], + }, + ] + calls, state = _run( + tmp_path, + task_file, + monkeypatch, + queue, + _smith({"moe/drafter_moe": ["BLOCKED", "LANDED"]}), + ) + + waiting = _smith_calls(calls, "moe/balancer_slots") + assert len(waiting) == 1, "held back in the smith phase, then dispatched once" + # Dispatched at eplb's start: after mtp was assembled, before eplb's turn. + assert _assembler_index(calls, "mtp") < waiting[0][0] < _assembler_index(calls, "eplb") + # Its first attempt, so there is no earlier report to point it at. + assert "This is a retry" not in waiting[0][1] + assert state.verdict == "delivered" + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "3:moe/balancer_slots → eplb (not dispatched: waits on [2])" in journal + + +# ---------------------------------------------------------------------- resume + + +def test_a_resume_keeps_deferrals_and_does_not_redispatch_landed_rows( + tmp_path, task_file, monkeypatch +): + _features(task_file, ["mtp"]) + queue = [ + {"n": 1, "mode": "onboard", "target": "norm/core_only"}, + {"n": 2, "mode": "onboard", "target": "moe/drafter_moe", "unit": "mtp"}, + ] + + def die(turn, prompt): + raise KeyboardInterrupt + + first = _workflow(tmp_path) + _install( + monkeypatch, + first, + Script(None, smith=_smith({"moe/drafter_moe": ["BLOCKED"]}), reviewer=die), + ) + _plan_then_vocabulary(first, tmp_path, monkeypatch, queue) + with pytest.raises(KeyboardInterrupt), first: + first.run(task_file) + + saved = load_state(tmp_path / "ws" / _STATE) + assert saved.stage == STAGE_ASSEMBLE + assert _row(saved, 1)["status"] == ITEM_DONE + assert _row(saved, 2)["status"] == ITEM_DEFERRED + assert _row(saved, 2)["deferred_to"] == "mtp" + + second = _workflow(tmp_path) + calls = _install(monkeypatch, second, Script(None, smith=_smith())) + with second: + second.run(task_file) + + assert not _smith_calls(calls, "norm/core_only"), "a landed row is never re-dispatched" + retried = _smith_calls(calls, "moe/drafter_moe") + assert len(retried) == 1 + assert "This is a retry" in retried[0][1] + assert retried[0][0] < _assembler_index(calls, "mtp") + assert load_state(tmp_path / "ws" / _STATE).verdict == "delivered" + + +def test_a_resume_after_the_retry_landed_dispatches_no_smith(tmp_path, task_file, monkeypatch): + """The retry is owed once. A restart inside the unit it served pays nothing again.""" + _features(task_file, ["mtp"]) + queue = [ + {"n": 1, "mode": "onboard", "target": "norm/core_only"}, + {"n": 2, "mode": "onboard", "target": "moe/drafter_moe", "unit": "mtp"}, + ] + + def assembler(turn, prompt): + if "Unit: **mtp**" in prompt: + raise KeyboardInterrupt + return {"decision": "UNIT_READY", "summary": "assembled"} + + first = _workflow(tmp_path) + _install( + monkeypatch, + first, + Script( + None, + smith=_smith({"moe/drafter_moe": ["BLOCKED", "LANDED"]}), + assembler=assembler, + ), + ) + _plan_then_vocabulary(first, tmp_path, monkeypatch, queue) + with pytest.raises(KeyboardInterrupt), first: + first.run(task_file) + assert _row(load_state(tmp_path / "ws" / _STATE), 2)["status"] == ITEM_DONE + + second = _workflow(tmp_path) + calls = _install(monkeypatch, second, Script(None, smith=_smith())) + with second: + second.run(task_file) + + assert _smith_calls(calls) == [] + assert load_state(tmp_path / "ws" / _STATE).verdict == "delivered" + + +def test_a_retry_interrupted_in_flight_is_still_owed(tmp_path, task_file, monkeypatch): + """A deferred row checkpointed `running` died mid-retry; its sibling that landed did not. + + The walk marks a batch `running` and checkpoints before it dispatches, so + that is the state a crash inside the retry batch leaves behind. + """ + _features(task_file, ["mtp"]) + queue = [ + {"n": 1, "mode": "onboard", "target": "norm/core_only"}, + {"n": 2, "mode": "onboard", "target": "moe/drafter_moe", "unit": "mtp"}, + {"n": 3, "mode": "onboard", "target": "moe/drafter_router", "unit": "mtp"}, + ] + + def die(turn, prompt): + raise KeyboardInterrupt + + first = _workflow(tmp_path) + blocked = {"moe/drafter_moe": ["BLOCKED"], "moe/drafter_router": ["BLOCKED"]} + _install(monkeypatch, first, Script(None, smith=_smith(blocked), reviewer=die)) + _plan_then_vocabulary(first, tmp_path, monkeypatch, queue) + with pytest.raises(KeyboardInterrupt), first: + first.run(task_file) + + state_path = tmp_path / "ws" / _STATE + data = json.loads(state_path.read_text(encoding="utf-8")) + for row in data["queue"]: + if row["n"] == 2: + row["status"] = ITEM_DONE + if row["n"] == 3: + assert row["status"] == ITEM_DEFERRED + row["status"] = "running" + state_path.write_text(json.dumps(data), encoding="utf-8") + + second = _workflow(tmp_path) + calls = _install(monkeypatch, second, Script(None, smith=_smith())) + with second: + second.run(task_file) + + assert not _smith_calls(calls, "norm/core_only") + assert not _smith_calls(calls, "moe/drafter_moe") + router = _smith_calls(calls, "moe/drafter_router") + assert len(router) == 1 + assert router[0][0] < _assembler_index(calls, "mtp") + assert load_state(state_path).verdict == "delivered" + + +# ------------------------------- applying the change to a run already closed + + +def _strip_units(state_path): + """Make the checkpoint look like one written before `unit` existed.""" + data = json.loads(state_path.read_text(encoding="utf-8")) + for row in data["queue"]: + row.pop("unit", None) + state_path.write_text(json.dumps(data), encoding="utf-8") + + +def _close_smith_blocked(tmp_path, task_file, monkeypatch, queue): + workflow = _workflow(tmp_path) + _install(monkeypatch, workflow, Script(None, smith=_smith({"moe/drafter_moe": ["BLOCKED"]}))) + _plan_then_vocabulary(workflow, tmp_path, monkeypatch, queue) + with workflow: + workflow.run(task_file) + state_path = tmp_path / "ws" / _STATE + closed = load_state(state_path) + assert closed.verdict == "smith-blocked" and closed.stage == STAGE_SMITH + _strip_units(state_path) + return state_path + + +def _tag_vocabulary(tmp_path, tags): + path = tmp_path / "ws" / "vocabulary.yaml" + doc = yaml.safe_load(path.read_text(encoding="utf-8")) + for row in doc["queue"]: + if row["n"] in tags: + row["unit"] = tags[row["n"]] + path.write_text(yaml.safe_dump(doc), encoding="utf-8") + + +def test_tagging_a_closed_runs_vocabulary_re_enters_it(tmp_path, task_file, monkeypatch): + """The operator's path for a run whose checkpoint predates units. + + The run closed `smith-blocked` over a feature-only row. Tagging that row in + `vocabulary.yaml` and re-running is the whole repair: the failed row is + deferred without a second dispatch now, the row the close left unstarted + is dispatched, core assembles, and the retry happens at mtp. + """ + _features(task_file, ["mtp"]) + queue = [ + {"n": 1, "mode": "onboard", "target": "norm/core_only"}, + {"n": 2, "mode": "onboard", "target": "moe/drafter_moe"}, + {"n": 3, "mode": "onboard", "target": "attn/core_after", "depends_on": [1]}, + ] + state_path = _close_smith_blocked(tmp_path, task_file, monkeypatch, queue) + assert _row(load_state(state_path), 3)["status"] == "pending" + _tag_vocabulary(tmp_path, {2: "mtp"}) + + second = _workflow(tmp_path) + calls = _install(monkeypatch, second, Script(None, smith=_smith())) + with second: + second.run(task_file) + + assert not _smith_calls(calls, "norm/core_only") + assert len(_smith_calls(calls, "attn/core_after")) == 1 + retried = _smith_calls(calls, "moe/drafter_moe") + assert len(retried) == 1, "deferred without a second smith-phase dispatch" + assert _assembler_index(calls, "core") < retried[0][0] < _assembler_index(calls, "mtp") + final = load_state(state_path) + assert final.verdict == "delivered" + assert _row(final, 2)["unit"] == "mtp" + + +def test_a_closed_run_stays_closed_when_core_needed_the_failure(tmp_path, task_file, monkeypatch): + """No tag, or a tag that does not cover the failure: re-running still refuses.""" + _features(task_file, ["mtp"]) + queue = [ + {"n": 1, "mode": "onboard", "target": "norm/core_only"}, + {"n": 2, "mode": "onboard", "target": "moe/drafter_moe"}, + ] + state_path = _close_smith_blocked(tmp_path, task_file, monkeypatch, queue) + + second = _workflow(tmp_path) + calls = _install(monkeypatch, second, Script(None, smith=_smith())) + with second: + second.run(task_file) + + assert calls == [] + assert load_state(state_path).verdict == "smith-blocked" + + +def test_a_tag_naming_a_unit_the_run_does_not_build_is_not_applied( + tmp_path, task_file, monkeypatch +): + """Untagged stays core -- the blocking direction -- rather than a guess.""" + _features(task_file, ["mtp"]) + queue = [ + {"n": 1, "mode": "onboard", "target": "norm/core_only"}, + {"n": 2, "mode": "onboard", "target": "moe/drafter_moe"}, + ] + state_path = _close_smith_blocked(tmp_path, task_file, monkeypatch, queue) + _tag_vocabulary(tmp_path, {2: "dspark"}) + + second = _workflow(tmp_path) + calls = _install(monkeypatch, second, Script(None, smith=_smith())) + with second: + second.run(task_file) + + assert calls == [] + assert "unit" not in _row(load_state(state_path), 2) + + +# ------------------------------------------------------------------ plan phase + + +def test_the_plan_phase_rejects_a_unit_the_run_does_not_build(tmp_path, task_file, monkeypatch): + """Caught as a re-draft, where the drafter can fix it, not at the first failure.""" + _features(task_file, ["mtp"]) + queue = [{"n": 1, "mode": "onboard", "target": "moe/drafter_moe", "unit": "dspark"}] + workflow = _workflow(tmp_path, max_plan_rounds=0) + calls = _install(monkeypatch, workflow, Script(None)) + _plan_then_vocabulary(workflow, tmp_path, monkeypatch, queue) + with workflow: + workflow.run(task_file) + + assert "smith" not in [role for role, _ in calls] + assert load_state(tmp_path / "ws" / _STATE).verdict == "plan-failed" + assert "not one of this run's assembly units" in workflow._vocabulary_errors diff --git a/agent-flow/tests/workflows/staircase_onboard/test_smith_resume_adoption.py b/agent-flow/tests/workflows/staircase_onboard/test_smith_resume_adoption.py new file mode 100644 index 000000000000..cbb731e4860c --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_smith_resume_adoption.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A resumed smith phase collects answers the dead process never read. + +Measured on the DeepSeek-V4.1 run: the compute session was killed with 30 +smith turns in flight and 16 LANDED decisions already on disk; the state file +had every item `running`, so a restart re-dispatched all 30. +""" + +from __future__ import annotations + +import yaml + +from agent_flow.workflows.staircase_onboard.progress import decision_path, mark_decision_dispatch +from agent_flow.workflows.staircase_onboard.vocabulary import QueueItem + +from .test_workflow import Script, _install, _workflow + + +def _decision(workspace, n, decision="LANDED"): + path = decision_path(workspace, "smith", n) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + yaml.safe_dump({"decision": decision, "summary": "landed", "entry": "norm/x"}), + encoding="utf-8", + ) + + +def test_a_decision_written_after_the_last_dispatch_is_adopted(tmp_path, monkeypatch): + wf = _workflow(tmp_path) + calls = _install(monkeypatch, wf, Script(None)) + wf._decision_files = True + mark_decision_dispatch(wf.workspace, "smith", 7) # the dead process's dispatch + _decision(wf.workspace, 7) # answered after it + + result = wf._run_one_smith(QueueItem(n=7, mode="onboard", target="norm/x")) + + assert result.startswith("adopted LANDED") + assert not [c for c in calls if c[0] == "smith"] + assert wf._smith_decision_for(QueueItem(n=7, mode="onboard", target="norm/x")) == "LANDED" + + +def test_a_retry_never_adopts_the_failure_it_is_retrying(tmp_path, monkeypatch): + wf = _workflow(tmp_path) + calls = _install(monkeypatch, wf, Script(None)) + wf._decision_files = True + monkeypatch.setattr(type(wf), "_smith_prompt", lambda self, item, note="": "Dispatch 7") + mark_decision_dispatch(wf.workspace, "smith", 7) + _decision(wf.workspace, 7, decision="FAILED") + + wf._run_one_smith(QueueItem(n=7, mode="onboard", target="norm/x"), note="retry") + + assert [c for c in calls if c[0] == "smith"] + + +def test_no_decision_on_disk_dispatches_as_before(tmp_path, monkeypatch): + wf = _workflow(tmp_path) + calls = _install(monkeypatch, wf, Script(None)) + wf._decision_files = True + monkeypatch.setattr(type(wf), "_smith_prompt", lambda self, item, note="": "Dispatch 3") + + wf._run_one_smith(QueueItem(n=3, mode="onboard", target="norm/y")) + + assert [c for c in calls if c[0] == "smith"] diff --git a/agent-flow/tests/workflows/staircase_onboard/test_state.py b/agent-flow/tests/workflows/staircase_onboard/test_state.py new file mode 100644 index 000000000000..6aea10354e3b --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_state.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the staircase-onboard checkpoint state.""" + +from __future__ import annotations + +import json + +import pytest + +from agent_flow.workflows.staircase_onboard import state as state_mod +from agent_flow.workflows.staircase_onboard.state import ( + ITEM_DONE, + ITEM_PENDING, + SCHEMA_VERSION, + STAGE_ASSEMBLE, + STAGE_PREFLIGHT, + WorkflowState, + load_state, + save_state, +) + + +def test_round_trip_preserves_every_field(tmp_path): + path = tmp_path / "state.json" + original = WorkflowState( + task_path="/x/task.yaml", + stage=STAGE_ASSEMBLE, + done=False, + slug="qwen3-30b-a3b-sm_103-tp1", + branch="staircase/qwen3-30b-a3b-sm_103-tp1", + devices=[2, 3], + plan_round=1, + plan_approved=True, + queue=[{"n": 1, "mode": "onboard", "target": "moe_gelu", "status": ITEM_DONE}], + unit_index=1, + unit_turns=3, + units_done=["core"], + audit_clean=False, + human_stops=2, + ) + save_state(path, original) + assert load_state(path) == original + + +def test_default_state_starts_at_preflight(): + assert WorkflowState(task_path="t").stage == STAGE_PREFLIGHT + + +def test_save_is_atomic_and_leaves_no_temp_files(tmp_path): + path = tmp_path / "state.json" + save_state(path, WorkflowState(task_path="t")) + save_state(path, WorkflowState(task_path="t", stage=STAGE_ASSEMBLE)) + assert [p.name for p in tmp_path.iterdir()] == ["state.json"] + + +def test_unknown_version_is_rejected(tmp_path): + path = tmp_path / "state.json" + path.write_text(json.dumps({"version": 99, "task_path": "t", "stage": "plan"})) + with pytest.raises(ValueError, match="Unsupported checkpoint version"): + load_state(path) + + +def test_unknown_stage_is_rejected(tmp_path): + path = tmp_path / "state.json" + path.write_text(json.dumps({"version": SCHEMA_VERSION, "task_path": "t", "stage": "tuning"})) + with pytest.raises(ValueError, match="Unsupported stage"): + load_state(path) + + +def test_queue_with_unknown_status_is_rejected(tmp_path): + """A silently-dropped queue row is a catalog entry that never gets built.""" + path = tmp_path / "state.json" + path.write_text( + json.dumps( + { + "version": SCHEMA_VERSION, + "task_path": "t", + "stage": "smith", + "queue": [{"n": 1, "status": "probably-fine"}], + } + ) + ) + with pytest.raises(ValueError, match="probably-fine"): + load_state(path) + + +def test_queue_must_be_a_list_of_mappings(tmp_path): + path = tmp_path / "state.json" + path.write_text( + json.dumps( + { + "version": SCHEMA_VERSION, + "task_path": "t", + "stage": "smith", + "queue": ["not-a-mapping"], + } + ) + ) + with pytest.raises(ValueError, match="must be a mapping"): + load_state(path) + + +def test_missing_optional_fields_default(tmp_path): + path = tmp_path / "state.json" + path.write_text(json.dumps({"version": SCHEMA_VERSION, "task_path": "t", "stage": "plan"})) + loaded = load_state(path) + assert loaded.queue == [] + assert loaded.human_stops == 0 + assert loaded.unit_index == 0 + assert loaded.plan_approved is False + + +def test_item_statuses_are_distinct(): + statuses = { + ITEM_PENDING, + ITEM_DONE, + state_mod.ITEM_RUNNING, + state_mod.ITEM_FAILED, + state_mod.ITEM_DEFERRED, + } + assert len(statuses) == 5 + + +# ------------------------------------------------------- remediable verdicts + + +def test_the_verdict_survives_a_round_trip(): + """A closed run that cannot say why it closed has to be started over.""" + path_state = WorkflowState(task_path="t", done=True, verdict="delivery-failed") + assert path_state.verdict == "delivery-failed" + + +def test_a_state_written_before_verdicts_existed_still_loads(tmp_path): + path = tmp_path / "state.json" + path.write_text( + json.dumps({"version": SCHEMA_VERSION, "task_path": "t", "stage": STAGE_PREFLIGHT}), + encoding="utf-8", + ) + assert load_state(path).verdict == "" + + +@pytest.mark.parametrize( + "verdict", + [ + "delivery-failed", + "audit-uncovered", + "witness-unconfigured:host_offload", + "assembler-blocked:knowledge_gap", + "assembler-blocked:environment", + ], +) +def test_verdicts_naming_a_repair_outside_the_run_are_remediable(verdict): + """These close a run whose agent work is done. + + The repair is an edit in the repo or in task.yaml -- a missing CI test + file, an entry needing certification, a witness block never configured -- + and the run cannot observe it except by running its own check again. + """ + assert state_mod.is_remediable(verdict) + + +@pytest.mark.parametrize( + "verdict", + [ + "ceiling", + "plan-failed", + "smith-blocked", + "assembly-exhausted:mtp", + "assembler-blocked:vocabulary_gap", + "delivered", + "", + ], +) +def test_budgets_and_agent_judgments_are_not_remediable(verdict): + """Re-running these re-dispatches paid turns with nothing changed. + + That is how a workflow burns a night discovering the same thing four + times. They need `--clean` and a human deciding what to do differently. + """ + assert not state_mod.is_remediable(verdict) + + +def test_a_prefix_is_not_enough_to_count_as_remediable(): + """`startswith` alone would match a verdict that merely begins the same.""" + assert not state_mod.is_remediable("delivery-failed-somehow-else") + + +def test_a_deferred_row_survives_a_round_trip(tmp_path): + """A resume that forgot a deferral would assemble that unit without the entry.""" + path = tmp_path / "state.json" + row = { + "n": 2, + "mode": "onboard", + "target": "moe/drafter_moe", + "unit": "dspark", + "status": state_mod.ITEM_DEFERRED, + "deferred_to": "dspark", + "deferred_why": "smith decided BLOCKED", + } + save_state(path, WorkflowState(task_path="t", queue=[row])) + assert load_state(path).queue == [row] + + +def test_re_entering_an_assembler_block_restores_the_unit_turn_budget(): + """The budget bounds the unit's own self-heal loop. + + The operator's repair happened outside it, so the resumed unit gets its + turns back. + """ + import inspect + + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + src = inspect.getsource(StaircaseOnboardWorkflow._init_state) + branch = src[src.index("re-entered a closed run") :] + branch = branch[: branch.index("state.done = False")] + assert 'state.verdict.startswith("assembler-blocked")' in branch + assert "state.unit_turns = 0" in branch diff --git a/agent-flow/tests/workflows/staircase_onboard/test_substrate_gate.py b/agent-flow/tests/workflows/staircase_onboard/test_substrate_gate.py new file mode 100644 index 000000000000..3b61b01153b2 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_substrate_gate.py @@ -0,0 +1,377 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the substrate gates — the checks that come before any agent. + +Both exist because one real onboarding spent roughly half its wall clock on +failures that had nothing to do with the model, and every one of them +surfaced mid-run inside an agent's turn. +""" + +from __future__ import annotations + +import pytest + +from agent_flow.workflows.staircase_onboard.environment import resolve_environment +from agent_flow.workflows.staircase_onboard.gates import ShellResult +from agent_flow.workflows.staircase_onboard.preflight import ( + PreflightError, + agent_plane_gate, + substrate_gate, +) + +_TASK = { + "checkpoint_path": "/ckpt", + "trtllm_repo_path": "/repo", + "parallel": "tp1", + "runtime": {"cuda_graph": True, "extra_flags": ["--kv_cache_dtype fp8"]}, + "features": ["mtp"], +} + + +def _environment(): + return resolve_environment( + {"environment": {"kind": "local", "python": "/py", "env": {"PATH": "/venv/bin"}}}, + world_size=1, + ) + + +def _runner(result, sink=None): + def run(command, env=None, timeout=0): + if sink is not None: + sink.append((command, timeout)) + return result + + return run + + +def _scripted(*results, sink=None): + """A runner that answers each call with the next result in turn. + + The fallback path runs the gate twice — the run's own checkpoint, then + the proxy — so a single fixed answer cannot tell the two apart. + """ + remaining = list(results) + + def run(command, env=None, timeout=0): + if sink is not None: + sink.append((command, timeout)) + return remaining.pop(0) if remaining else results[-1] + + return run + + +#: What a stock path that has never heard of this architecture actually +#: prints. Quoted from the raise sites named in +#: `preflight.UNSUPPORTED_ARCHITECTURE_MARKERS`. +_UNKNOWN_ARCH = "ValueError: Unknown architecture for AutoModelForCausalLM: MiMoV2ForCausalLM" + + +# ------------------------------------------------------------ GPU plane + + +def test_a_host_that_generates_passes(): + notes: list[str] = [] + substrate_gate( + _TASK, + _environment(), + shell_runner=_runner(ShellResult(returncode=0, stdout="Generated text: Paris")), + notes=notes, + ) + assert any("substrate gate" in n for n in notes) + + +def test_a_host_that_cannot_generate_fails_before_any_agent(): + runner = _runner(ShellResult(returncode=1, stdout="MPI_ERR_SPAWN: could not spawn processes")) + with pytest.raises(PreflightError, match="cannot generate a token"): + substrate_gate(_TASK, _environment(), shell_runner=runner) + + +def test_the_failure_says_the_target_is_not_at_fault(): + """The whole point: these errors read as model defects and are not.""" + runner = _runner(ShellResult(returncode=1, stdout='could not open source file "cuda.h"')) + with pytest.raises(PreflightError) as excinfo: + substrate_gate(_TASK, _environment(), shell_runner=runner) + message = str(excinfo.value) + assert "no target exists yet" in message + # And it must carry the command, or the operator cannot reproduce it. + assert "quickstart_advanced.py" in message + + +def test_it_runs_the_stock_path_not_the_target(): + """There is no target at preflight time, so `require` would always fail.""" + sink: list[tuple[str, int]] = [] + substrate_gate( + _TASK, + _environment(), + shell_runner=_runner(ShellResult(returncode=0, stdout="ok"), sink), + ) + command = sink[0][0] + assert "TRTLLM_MODELING_V2=off" in command + assert "TRTLLM_MODELING_V2=require" not in command + + +def test_it_carries_the_runs_own_runtime_block(): + """Each of the three motivating failures depended on these flags.""" + sink: list[tuple[str, int]] = [] + substrate_gate( + _TASK, + _environment(), + shell_runner=_runner(ShellResult(returncode=0, stdout="ok"), sink), + ) + command = sink[0][0] + assert "--use_cuda_graph" in command + assert "--kv_cache_dtype fp8" in command + assert "PATH=/venv/bin" in command + + +def test_it_does_not_carry_any_feature_flags(): + """No unit is assembled yet, so a speculator flag would be a lie.""" + sink: list[tuple[str, int]] = [] + substrate_gate( + _TASK, + _environment(), + shell_runner=_runner(ShellResult(returncode=0, stdout="ok"), sink), + ) + assert "--spec_decode_algo" not in sink[0][0] + + +# ------------------------------------------- the architecture that has no stock path + + +def _with_proxy(proxy="/proxy-ckpt"): + return _TASK | { + "environment": {"kind": "local", "python": "/py", "substrate_checkpoint": proxy} + } + + +@pytest.mark.parametrize( + "message", + [ + "ValueError: Unknown architecture for AutoModelForCausalLM: MiMoV2ForCausalLM", + "ValueError: Unknown architecture for AutoModelForMultimodalEncoder: MiMoV2ForCausalLM", + "NotImplementedError: The given huggingface model architecture MiMoV2ForCausalLM " + "is not supported in TRT-LLM yet", + ], +) +def test_every_way_the_stock_path_says_it_does_not_know_the_architecture(message): + """Matched on the message text, because the gate reads a log and not an exception. + + All three come from the two architecture registries — `modeling_auto.py`'s + `AutoModelForCausalLM.from_config` and `automodel.py`'s `MODEL_MAP` + lookup — and all three mean the same thing about this run. + """ + runner = _scripted(ShellResult(returncode=1, stdout=message)) + with pytest.raises(PreflightError, match="does not know this checkpoint's architecture"): + substrate_gate(_TASK, _environment(), shell_runner=runner) + + +def test_an_unknown_architecture_with_no_proxy_says_to_declare_one(): + """Not the raw architecture error, which reads as a target defect. + + There is no target at preflight time, so a message naming the + checkpoint's architecture as unsupported points the operator at exactly + the thing that is supposed to be unsupported — the failure mode this + gate exists to prevent, reproduced by the gate itself. + """ + runner = _scripted(ShellResult(returncode=1, stdout=_UNKNOWN_ARCH)) + with pytest.raises(PreflightError) as excinfo: + substrate_gate(_TASK, _environment(), shell_runner=runner) + message = str(excinfo.value) + assert "environment.substrate_checkpoint" in message + assert "not a fault in the checkpoint or the host" in message + # And it must steer away from the blunter escape hatch, which also gives + # up the two checks that do not depend on the checkpoint at all. + assert "--skip-substrate-gate" in message + + +def test_the_proxy_is_run_under_the_same_environment_and_runtime_block(): + sink: list[tuple[str, int]] = [] + notes: list[str] = [] + substrate_gate( + _with_proxy(), + _environment(), + shell_runner=_scripted( + ShellResult(returncode=1, stdout=_UNKNOWN_ARCH), + ShellResult(returncode=0, stdout="Generated text: Paris"), + sink=sink, + ), + notes=notes, + ) + assert len(sink) == 2 + first, second = sink[0][0], sink[1][0] + assert '--model_dir="/ckpt"' in first + assert '--model_dir="/proxy-ckpt"' in second + # Same environment, same runtime block, same stock path — one variable. + for command in (first, second): + assert "TRTLLM_MODELING_V2=off" in command + assert "--use_cuda_graph" in command + assert "--kv_cache_dtype fp8" in command + + +def test_the_weaker_claim_is_recorded_verbatim(): + """The weaker claim is recorded verbatim. + + A note that said only "substrate gate passed" would be read later as + the claim the unqualified gate makes, and the difference is precisely + which failure this run has been protected from. + """ + notes: list[str] = [] + substrate_gate( + _with_proxy(), + _environment(), + shell_runner=_scripted( + ShellResult(returncode=1, stdout=_UNKNOWN_ARCH), + ShellResult(returncode=0, stdout="Generated text: Paris"), + ), + notes=notes, + ) + assert notes == [ + "substrate proven on `/proxy-ckpt`; this checkpoint's own geometry is " + "unverified, so a kvCache sizing fault will surface at the first boot " + "gate rather than here." + ] + + +def test_a_real_host_failure_is_never_answered_by_measuring_another_model(): + """The whole reason the discriminator is narrow. + + Widening the fallback to "any failure" would let a genuine kvCache + sizing fault on the target checkpoint be answered by silently booting a + different model and reporting green — the same class of mistake as + gating under `TRTLLM_MODELING_V2=auto`. + """ + sink: list[tuple[str, int]] = [] + runner = _scripted( + ShellResult(returncode=1, stdout="Impossible to fit in any sequence in kvCache"), + ShellResult(returncode=0, stdout="Generated text: Paris"), + sink=sink, + ) + with pytest.raises(PreflightError, match="cannot generate a token"): + substrate_gate(_with_proxy(), _environment(), shell_runner=runner) + assert len(sink) == 1, "the proxy was run for a failure that was not about the architecture" + + +def test_a_proxy_that_also_fails_fails_the_run_and_names_the_proxy(): + """The proxy is the control, so its failure is a host failure.""" + runner = _scripted( + ShellResult(returncode=1, stdout=_UNKNOWN_ARCH), + ShellResult(returncode=1, stdout="MPI_ERR_SPAWN: could not spawn processes"), + ) + with pytest.raises(PreflightError) as excinfo: + substrate_gate(_with_proxy(), _environment(), shell_runner=runner) + message = str(excinfo.value) + assert "failed on the declared proxy checkpoint /proxy-ckpt" in message + assert "no target exists yet" in message + + +def test_a_checkpoint_the_stock_path_knows_never_reaches_the_proxy(): + """Nothing changes for the run that was already working.""" + sink: list[tuple[str, int]] = [] + notes: list[str] = [] + substrate_gate( + _with_proxy(), + _environment(), + shell_runner=_scripted(ShellResult(returncode=0, stdout="Paris"), sink=sink), + notes=notes, + ) + assert len(sink) == 1 + assert any("stock path generated" in note for note in notes) + + +# ---------------------------------------------------------- agent plane + + +def test_one_good_agent_turn_passes(): + notes: list[str] = [] + agent_plane_gate(lambda: "OK", notes) + assert any("agent plane" in n for n in notes) + + +def test_an_agent_that_cannot_start_fails_loudly(): + def probe(): + raise RuntimeError("EROFS: read-only file system, mkdir '.../credentials.lock'") + + with pytest.raises(PreflightError, match="not a model or target"): + agent_plane_gate(probe) + + +def test_an_empty_reply_is_a_failure(): + """A role that cannot answer a trivial prompt will not write a decision.""" + with pytest.raises(PreflightError, match="empty reply"): + agent_plane_gate(lambda: " ") + + +def test_trust_remote_code_refusal_is_an_unsupported_architecture(): + """A checkpoint whose implementation travels with its weights. + + `transformers` raises this instead of any registry message when the + checkpoint carries an `auto_map`, and that is the same condition: the + stock path has no built-in architecture for it. Keying the discriminator + only on TensorRT-LLM's registry missed it, and the run stopped on a + checkpoint nobody expected the stock path to load. + """ + from agent_flow.workflows.staircase_onboard.preflight import _reports_unsupported_architecture + + class _Result: + detail = "exited 1" + stdout_tail = ( + "ValueError: The repository /path/to/ckpt contains custom code which " + "must be executed to correctly load the model. ... Please pass the " + "argument `trust_remote_code=True` to allow custom code to be run." + ) + + assert _reports_unsupported_architecture(_Result()) + + +def test_an_unknown_model_type_is_an_unsupported_architecture(): + """A checkpoint newer than the installed `transformers`, with no remote code. + + Measured on DeepSeek-V4.1-Flash-NVFP4 (`model_type: deepseek_v41`, + transformers 5.5.4): `AutoConfig.from_pretrained` refuses before any + registry is consulted, and the run stopped instead of falling back. + """ + from agent_flow.workflows.staircase_onboard.preflight import _reports_unsupported_architecture + + class _Result: + detail = "exited 1" + stdout_tail = ( + "ValueError('The checkpoint you are trying to load has model type " + "`deepseek_v41` but Transformers does not recognize this architecture. " + "This could be because of an issue with the checkpoint, or because your " + "version of Transformers is out of date.')" + ) + + assert _reports_unsupported_architecture(_Result()) + + +def test_a_registration_shim_without_a_target_is_an_unsupported_architecture(): + """A registration shim without a target is an unsupported architecture. + + The stock path finds a class registered under the architecture name + (so a multimodal input processor can be found) that raises when built + without a modeling_v2 target: the acceptance gate's `=off` leg must fall + back to the declared floor rather than hard-fail. + """ + from agent_flow.workflows.staircase_onboard.preflight import _reports_unsupported_architecture + + class _Result: + detail = "exited 1" + stdout_tail = ( + "ValueError: DeepseekV41ForCausalLM has no built-in implementation: DeepSeek-V4.1 " + "runs only as a modeling_v2 target (tensorrt_llm/_torch/_experimental/modeling_v2/" + "models/deepseek_v41/), and this class only registers its multimodal input processor." + ) + + assert _reports_unsupported_architecture(_Result()) + + +def test_a_kv_cache_failure_is_still_a_hard_stop(): + """The widening must not turn the fallback into "any failure".""" + from agent_flow.workflows.staircase_onboard.preflight import _reports_unsupported_architecture + + class _Result: + detail = "exited 1" + stdout_tail = "Impossible to fit in any sequence in kvCache" + + assert not _reports_unsupported_architecture(_Result()) diff --git a/agent-flow/tests/workflows/staircase_onboard/test_survey.py b/agent-flow/tests/workflows/staircase_onboard/test_survey.py new file mode 100644 index 000000000000..86211df849cd --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_survey.py @@ -0,0 +1,702 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The surface survey: what the checkpoint demands, what the build supplies. + +The demand half is tested against five *different families* on purpose. A +reader that only ever sees one checkpoint converges on that checkpoint, and +the whole reason this module exists is the next model — the one with a layer +nobody has a word for. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from agent_flow.workflows.staircase_onboard.environment import RunEnvironment +from agent_flow.workflows.staircase_onboard.survey import ( + MARKER, + RECEIPT_ABSENT, + RECEIPT_PASSED, + CatalogEntry, + CheckpointSurface, + Inventory, + format_surface, + parse_inventory, + probe_command, + read_checkpoint_surface, + run_inventory_probe, + write_surface_file, +) + +# --------------------------------------------------------------- the demands + + +def _dense(): + """An ordinary dense decoder — the boring baseline.""" + return dict( + architectures=["LlamaForCausalLM"], + model_type="llama", + num_attention_heads=32, + num_key_value_heads=8, + hidden_act="silu", + rope_theta=500000.0, + rms_norm_eps=1e-5, + hidden_size=4096, + num_hidden_layers=32, + vocab_size=128256, + ) + + +def _names(config): + return read_checkpoint_surface(config).demand_names + + +def test_a_dense_decoder_reports_only_what_it_has(): + names = _names(_dense()) + assert "attention.shape" in names + assert "mlp.activation" in names + assert "norm.style" in names + # Nothing it does not have. + assert not any(n.startswith(("moe.", "recurrent.", "stack.", "speculation.")) for n in names) + + +@pytest.mark.parametrize( + "kv_heads,shape", + [(32, "MHA"), (8, "GQA"), (1, "MQA")], +) +def test_attention_shape_is_read_from_the_head_counts(kv_heads, shape): + config = _dense() | {"num_key_value_heads": kv_heads} + demand = next(d for d in read_checkpoint_surface(config).demands if d.name == "attention.shape") + assert shape in demand.detail + + +def test_mla_is_not_confused_with_ordinary_attention(): + config = _dense() | {"kv_lora_rank": 512, "q_lora_rank": 1536} + assert "attention.mla" in _names(config) + + +def test_a_hybrid_stack_is_reported_as_hybrid_not_as_a_window(): + """The bug this test exists for: `layer_types` read as sliding-window. + + A checkpoint that interleaves linear and full attention has no window + anywhere in it, and reporting one sends the plan looking for a windowing + op that the model never calls. + """ + config = _dense() | {"layer_types": ["linear_attention"] * 3 + ["full_attention"]} + names = _names(config) + assert "stack.hybrid" in names + assert "attention.windowing" not in names + + +def test_a_real_window_is_still_reported(): + config = _dense() | {"sliding_window": 4096} + assert "attention.windowing" in _names(config) + + +def test_a_uniform_layer_types_list_is_not_a_hybrid(): + config = _dense() | {"layer_types": ["full_attention"] * 8} + assert "stack.hybrid" not in _names(config) + + +def test_moe_is_read_through_any_of_its_spellings(): + for key in ("num_experts", "num_local_experts", "n_routed_experts", "moe_num_experts"): + config = _dense() | {key: 128, "num_experts_per_tok": 8} + demand = next(d for d in read_checkpoint_surface(config).demands if d.name == "moe.routing") + assert "128 experts" in demand.detail, key + assert "top-8" in demand.detail, key + + +def test_a_recurrent_plane_is_called_out_as_a_different_cache(): + config = _dense() | {"mamba_d_state": 128, "conv_kernel": 4} + demand = next(d for d in read_checkpoint_surface(config).demands if d.name == "recurrent.state") + assert "state cache manager" in demand.detail + + +def test_quantization_absence_is_stated_rather_than_omitted(): + """An absent demand and an unmeasured one read the same; say which.""" + demand = next(d for d in read_checkpoint_surface(_dense()).demands if d.name == "quantization") + assert "none declared" in demand.detail + + +def test_quantization_reports_the_held_out_modules(): + config = _dense() | { + "quantization_config": { + "quant_algo": "NVFP4", + "kv_cache_quant_algo": "FP8", + "exclude_modules": ["lm_head", "vision"], + } + } + demand = next(d for d in read_checkpoint_surface(config).demands if d.name == "quantization") + assert "NVFP4" in demand.detail + assert "kv=FP8" in demand.detail + assert "2 module(s) held out" in demand.detail + + +def test_a_multimodal_checkpoint_is_read_through_its_text_tower(): + """Reading the outer config would describe the vision encoder.""" + config = { + "architectures": ["SomeVLForConditionalGeneration"], + "model_type": "some_vl", + "vision_config": {"num_attention_heads": 16, "hidden_size": 1408}, + "text_config": _dense() | {"num_attention_heads": 40, "num_key_value_heads": 8}, + } + surface = read_checkpoint_surface(config) + assert surface.submodels == ("vision_config", "text_config") + demand = next(d for d in surface.demands if d.name == "attention.shape") + assert "40 query heads" in demand.detail + + +# ------------------------------------------------------- the unknown column + + +def test_an_unclaimed_structural_key_is_surfaced_not_dropped(): + """The whole point, for a model this reader has never seen. + + `query_pre_attn_scalar` changes the attention scale. A reader that drops + what it does not recognize produces a target that is quietly wrong, and + the discovery moves to a smith at tier 4 with no budget left. + """ + config = _dense() | {"query_pre_attn_scalar": 224, "no_rope_layers": [0, 1, 1]} + surface = read_checkpoint_surface(config) + assert "query_pre_attn_scalar" in surface.unrecognized + assert "no_rope_layers" in surface.unrecognized + + +def test_bookkeeping_keys_are_not_reported_as_discoveries(): + config = _dense() | {"bos_token_id": 1, "transformers_version": "4.51", "_name_or_path": "x"} + assert read_checkpoint_surface(config).unrecognized == () + + +def test_a_key_a_probe_already_read_is_not_also_unrecognized(): + config = _dense() | {"rope_scaling": {"rope_type": "yarn"}, "num_experts": 8} + surface = read_checkpoint_surface(config) + assert "rope_scaling" not in surface.unrecognized + assert "num_experts" not in surface.unrecognized + + +# ------------------------------------------------------------- the inventory + + +def _inventory_payload(**overrides): + payload = { + "trtllm_ops": ["fused_qk_norm_rope", "flashinfer_silu_and_mul", "brand_new_op"], + "catalog": [ + { + "path": "attention/fused_qk_norm_rope.py", + "impl": "torch.ops.trtllm.fused_qk_norm_rope", + "summary": "s", + "receipt": "passed", + "impl_resolves": True, + }, + { + "path": "activation/flashinfer_silu_and_mul.py", + "impl": "torch.ops.trtllm.flashinfer_silu_and_mul", + "summary": "s", + "receipt": "absent", + "impl_resolves": True, + }, + { + "path": "gemm/gone.py", + "impl": "torch.ops.trtllm.gone", + "summary": "s", + "receipt": "passed", + "impl_resolves": False, + }, + ], + "modules": ["rms_norm.RMSNorm"], + "versions": {"torch": "2.12"}, + "error": "", + } + payload.update(overrides) + return payload + + +def test_the_payload_is_found_in_a_log_full_of_engine_chatter(): + stdout = ( + "[TRT-LLM] loading plugins\nrank 0 initialized\n" + + MARKER + + json.dumps(_inventory_payload()) + + "\n[TRT-LLM] shutting down\n" + ) + inventory = parse_inventory(stdout, arch="sm_103") + assert inventory.ok + assert len(inventory.catalog) == 3 + + +def test_a_silent_probe_is_an_error_carrying_the_tail(): + inventory = parse_inventory("import failed\nTraceback\nImportError: no module", arch="sm_103") + assert not inventory.ok + assert "ImportError" in inventory.error + + +def test_entries_are_sorted_into_spendable_uncertified_and_rotted(): + inventory = parse_inventory(MARKER + json.dumps(_inventory_payload()), arch="sm_103") + assert [e.name for e in inventory.spendable] == ["fused_qk_norm_rope"] + assert [e.name for e in inventory.uncertified] == ["flashinfer_silu_and_mul"] + assert [e.name for e in inventory.rotted] == ["gone"] + + +def test_a_rotted_entry_is_never_spendable_however_good_its_receipt(): + """It carries a passing receipt and still cannot be called.""" + entry = CatalogEntry(path="gemm/gone.py", receipt=RECEIPT_PASSED, impl_resolves=False) + assert not entry.spendable + + +def test_an_op_no_entry_wraps_is_reported_as_the_tier_one_rung(): + inventory = parse_inventory(MARKER + json.dumps(_inventory_payload()), arch="sm_103") + assert inventory.unwrapped_ops == ("brand_new_op",) + + +def test_the_probe_command_carries_the_runs_launch_prefix(): + env = RunEnvironment(kind="local", python="/venv/bin/python", env={"TRTLLM_MODELING_V2": "off"}) + command = probe_command(env, repo="/repo", arch="sm_103") + assert "TRTLLM_MODELING_V2=off" in command + assert "/venv/bin/python" in command + assert "/repo" in command and "sm_103" in command + + +def test_the_probe_is_never_launched_as_engine_work(): + """Read-only introspection, so one process at any world size. + + At world size 4 it went out as `mpirun -n 4 ... trtllm-llmapi-launch`, + which is the rule for a gate. The run's own `PATH` carried no `mpirun`, + the probe exited 127, and the plan phase drafted against an empty + inventory -- every certified entry read as a gap. + """ + env = RunEnvironment(kind="session_spool", python="/venv/bin/python", world_size=4, spool="/s") + command = probe_command(env, repo="/repo", arch="sm_103") + assert "mpirun" not in command + assert "trtllm-llmapi-launch" not in command + assert "__SURVEY_JSON__" in command + + +def test_a_failed_probe_is_reported_rather_than_raised(): + """A run survives a survey failure; it just plans on assumption.""" + + class _Result: + ok, returncode, timed_out = False, 1, False + stdout, stderr = "", "ModuleNotFoundError: tensorrt_llm" + + inventory = run_inventory_probe( + {"arch": "sm_103", "trtllm_repo_path": "/repo"}, + RunEnvironment(kind="local", python="python3"), + shell_runner=lambda *a, **k: _Result(), + ) + assert not inventory.ok + assert "exited 1" in inventory.error + + +def test_a_timeout_says_so(): + class _Result: + ok, returncode, timed_out = False, None, True + stdout, stderr = "", "" + + inventory = run_inventory_probe( + {"arch": "sm_103", "trtllm_repo_path": "/repo"}, + RunEnvironment(kind="local", python="python3"), + shell_runner=lambda *a, **k: _Result(), + timeout=7, + ) + assert "timed out after 7s" in inventory.error + + +# ---------------------------------------------------------------- rendering + + +def test_the_block_lays_both_columns_out_without_joining_them(): + surface = read_checkpoint_surface(_dense() | {"query_pre_attn_scalar": 224}) + inventory = parse_inventory(MARKER + json.dumps(_inventory_payload()), arch="sm_103") + block = format_surface(surface, inventory) + assert "Joining them is your job" in block + assert "query_pre_attn_scalar" in block + assert "fused_qk_norm_rope" in block + assert "brand_new_op" in block + assert "gone" in block + + +def test_a_failed_inventory_tells_the_drafter_it_is_planning_blind(): + block = format_surface( + read_checkpoint_surface(_dense()), Inventory(arch="sm_103", error="probe died") + ) + assert "probe died" in block + assert "you do not know what exists" in block + + +def test_a_failed_probe_is_not_presented_as_a_measurement(): + """The block opens by saying which of its two columns is a reading. + + Without that, a supply column that failed to land renders with the same + shape as one that landed empty, and the drafter -- which is told it is + reading a measurement -- plans every entry as a gap. + """ + surface = read_checkpoint_surface(_dense()) + measured = parse_inventory(MARKER + json.dumps(_inventory_payload()), arch="sm_103") + assert format_surface(surface, measured).startswith("## Surface survey\n\nMeasured") + + failed = format_surface(surface, Inventory(arch="sm_103", error="probe exited 127")) + assert "The supply column was NOT" in failed + assert "an assumption and not a reading" in failed + assert "evidence of absence" in failed + + +def test_the_workspace_file_is_unabridged(tmp_path): + """The prompt block is capped; the audit record is not.""" + payload = _inventory_payload( + trtllm_ops=[f"op_{i}" for i in range(200)], + modules=[f"mod.Class{i}" for i in range(90)], + ) + inventory = parse_inventory(MARKER + json.dumps(payload), arch="sm_103") + path = write_surface_file(tmp_path / "surface.md", read_checkpoint_surface(_dense()), inventory) + text = path.read_text() + assert "op_199" in text + assert "mod.Class89" in text + assert "torch: 2.12" in text + assert "op_199" not in format_surface(read_checkpoint_surface(_dense()), inventory) + + +def test_an_empty_surface_still_renders(): + block = format_surface(CheckpointSurface(architecture="", model_type=""), Inventory()) + assert "Surface survey" in block + + +def test_receipt_absent_is_the_default_not_a_pass(): + assert CatalogEntry(path="x/y.py").receipt == RECEIPT_ABSENT + assert not CatalogEntry(path="x/y.py").spendable + + +# ------------------------------------------------- the probe, actually run + + +def _fake_repo(tmp_path, entries, indent=""): + """A repo tree with a catalog index and contracts, and nothing else. + + ``indent`` nests the list under ``entries:`` the way the in-tree index + does; the default keeps it at column 0. + """ + catalog = tmp_path / "tensorrt_llm/_torch/_experimental/modeling_v2/catalog" + catalog.mkdir(parents=True) + lines = ["entries:"] + for path, impl, frontmatter in entries: + lines += [f"{indent}- path: {path}", f"{indent} impl: {impl}", f"{indent} summary: s"] + md = catalog / path.replace(".py", ".md") + md.parent.mkdir(parents=True, exist_ok=True) + md.write_text(f"---\n{frontmatter}---\n\n# {Path(path).stem}\n", encoding="utf-8") + (catalog / "index.yaml").write_text("\n".join(lines) + "\n", encoding="utf-8") + modules = tmp_path / "tensorrt_llm/_torch/modules" + modules.mkdir(parents=True) + (modules / "rms_norm.py").write_text("class RMSNorm:\n pass\n\ndef _hidden():\n pass\n") + return tmp_path + + +#: Makes `import yaml` raise, so the probe's hand-rolled fallback parser is +#: the thing under test. Without this the fallback is dead code in CI — the +#: interpreter that runs these tests has yaml, and so does every real run's — +#: and dead code is exactly where a parser bug hides. +_BLOCK_YAML = ( + "import sys\n" + "class _NoYaml:\n" + " def find_module(self, name, path=None):\n" + " return self if name == 'yaml' else None\n" + " def load_module(self, name):\n" + " raise ImportError('blocked for test')\n" + "sys.meta_path.insert(0, _NoYaml())\n" +) + + +def _run_probe(repo, arch="sm_103", *, block_yaml=False): + import subprocess + import sys + + from agent_flow.workflows.staircase_onboard.survey import PROBE_SOURCE + + source = (_BLOCK_YAML + PROBE_SOURCE) if block_yaml else PROBE_SOURCE + proc = subprocess.run( + [sys.executable, "-c", source, str(repo), arch], + capture_output=True, + text=True, + timeout=120, + ) + return parse_inventory(proc.stdout + proc.stderr, arch=arch) + + +#: The two shapes actually present in the tree, and the two the old +#: hand-rolled reader got wrong in the same direction — inline FAILED read as +#: passed, block FAILED read as absent. Both are the dangerous direction: +#: they make a known-broken entry look spendable or merely un-run. +_INLINE_PASS = "receipts:\n sm_103: {status: passed, tests: 9}\n" +_INLINE_FAIL = "receipts:\n sm_103: {status: failed, tests: 57}\n" +_BLOCK_FAIL = "receipts:\n sm_103:\n status: failed\n tests: 57\n note: >-\n long\n" +_BLOCK_PASS = "receipts:\n sm_103:\n status: passed\n tests: 12\n" +_OTHER_ARCH = "receipts:\n sm_100: {status: passed, tests: 9}\n" + + +@pytest.mark.parametrize( + "frontmatter,expected", + [ + (_INLINE_PASS, RECEIPT_PASSED), + (_INLINE_FAIL, "failed"), + (_BLOCK_FAIL, "failed"), + (_BLOCK_PASS, RECEIPT_PASSED), + (_OTHER_ARCH, RECEIPT_ABSENT), + ("", RECEIPT_ABSENT), + ], +) +@pytest.mark.parametrize("block_yaml", [False, True], ids=["yaml", "fallback"]) +def test_the_probe_reads_both_frontmatter_shapes(tmp_path, frontmatter, expected, block_yaml): + repo = _fake_repo(tmp_path, [("norm/x.py", "torch.ops.trtllm.x", frontmatter)]) + inventory = _run_probe(repo, block_yaml=block_yaml) + assert inventory.catalog[0].receipt == expected + + +@pytest.mark.parametrize("indent", ["", " "], ids=["column0", "indented"]) +def test_the_probe_reads_an_index_nested_under_entries(tmp_path, indent): + """The in-tree index indents its list under `entries:`. + + Measured on the DeepSeek-V4.1 dry run: the column-0 match read the real + index as empty, and the survey reported `0/0 entries certified` over 19 + entries that all carry sm_103 receipts -- the exact reading this survey + exists to prevent. + """ + repo = _fake_repo( + tmp_path, + [ + ("norm/x.py", "torch.ops.trtllm.x", _INLINE_PASS), + ("gemm/y.py", "torch.ops.trtllm.y", _INLINE_PASS), + ], + indent=indent, + ) + inventory = _run_probe(repo) + assert [e.path for e in inventory.catalog] == ["norm/x.py", "gemm/y.py"] + assert [e.impl for e in inventory.catalog] == ["torch.ops.trtllm.x", "torch.ops.trtllm.y"] + assert all(e.receipt == RECEIPT_PASSED for e in inventory.catalog) + + +def test_a_receipt_for_another_arch_is_not_this_arch(tmp_path): + """The architecture is the whole key; a pass elsewhere says nothing here.""" + repo = _fake_repo(tmp_path, [("norm/x.py", "torch.ops.trtllm.x", _OTHER_ARCH)]) + assert not _run_probe(repo, arch="sm_103").catalog[0].spendable + + +def test_the_probe_reads_the_kernel_plane_without_importing_it(tmp_path): + repo = _fake_repo(tmp_path, [("norm/x.py", "torch.ops.trtllm.x", _INLINE_PASS)]) + inventory = _run_probe(repo) + assert "rms_norm.RMSNorm" in inventory.modules + assert not any(m.endswith("._hidden") for m in inventory.modules) + + +def test_an_engine_that_will_not_import_does_not_condemn_the_whole_catalog(tmp_path): + """Measured: the probe runs on an interpreter with no torch. + + Without the guard this reported all 48 real entries as rotted, which is + the most alarming possible way to say "I could not measure this". + """ + repo = _fake_repo(tmp_path, [("norm/x.py", "torch.ops.trtllm.x", _INLINE_PASS)]) + inventory = _run_probe(repo) + # This interpreter has no engine, whether or not it has torch — so the + # op inventory cannot be taken either way, which is the condition. + assert inventory.error + assert inventory.rotted == () + assert inventory.catalog[0].receipt == RECEIPT_PASSED + + +# ------------------------------------------------- torch mirrors carry none + + +def test_a_torch_mirror_is_spendable_without_a_receipt(): + """Absence is by design there, not the three-state unknown.""" + entry = CatalogEntry(path="torch/add.py", impl="torch.add") + assert not entry.certifiable + assert entry.spendable + + +def test_a_torch_mirror_is_not_listed_as_needing_a_receipt_run(): + payload = _inventory_payload( + catalog=[ + { + "path": "torch/add.py", + "impl": "torch.add", + "receipt": "absent", + "impl_resolves": True, + }, + { + "path": "norm/y.py", + "impl": "torch.ops.trtllm.y", + "receipt": "absent", + "impl_resolves": True, + }, + ] + ) + inventory = parse_inventory(MARKER + json.dumps(payload), arch="sm_103") + assert [e.name for e in inventory.uncertified] == ["y"] + assert "add" in [e.name for e in inventory.spendable] + + +def test_a_failed_receipt_is_listed_loudly_rather_than_as_unrun(): + """A known-failing entry that reads as "not yet run" gets spent.""" + payload = _inventory_payload( + catalog=[ + { + "path": "attention/thop_attention.py", + "impl": "torch.ops.trtllm.thop_attention", + "receipt": "failed", + "impl_resolves": True, + } + ] + ) + inventory = parse_inventory(MARKER + json.dumps(payload), arch="sm_103") + assert inventory.spendable == () + assert [(e.name, e.receipt) for e in inventory.uncertified] == [("thop_attention", "failed")] + assert "receipt: failed" in format_surface(read_checkpoint_surface(_dense()), inventory) + + +def test_a_partial_inventory_is_still_rendered(): + """An import failure still leaves the catalog, which is most of the value.""" + payload = _inventory_payload(trtllm_ops=[], error="ModuleNotFoundError: torch") + inventory = parse_inventory(MARKER + json.dumps(payload), arch="sm_103") + block = format_surface(read_checkpoint_surface(_dense()), inventory) + assert "ModuleNotFoundError" in block + assert "fused_qk_norm_rope" in block # measured, so still shown + assert "nothing can be called rotted" in block + + +def test_failed_and_unrun_are_not_lumped_into_one_heading(): + """They mean opposite things about how much work is left.""" + payload = _inventory_payload( + catalog=[ + { + "path": "ssm/broken.py", + "impl": "torch.ops.trtllm.broken", + "receipt": "failed", + "impl_resolves": True, + }, + { + "path": "norm/unrun.py", + "impl": "torch.ops.trtllm.unrun", + "receipt": "absent", + "impl_resolves": True, + }, + ] + ) + block = format_surface( + read_checkpoint_surface(_dense()), + parse_inventory(MARKER + json.dumps(payload), arch="sm_103"), + ) + assert "Known failing on this arch" in block + assert "Present but un-run on this arch" in block + # The failing one must not be described as merely awaiting a run. + failing_section = block.split("Known failing")[1].split("**Present but un-run")[0] + assert "broken" in failing_section + assert "unrun" not in failing_section + + +# ------------------------------------------- the linter cannot see this code + + +def test_the_probe_source_compiles(): + """`ruff` and `ruff format` do not look inside a string literal. + + Both embedded programs in this package have already shipped a bug that + a linter would have caught on a real file. This one: a `\"\"\"` docstring + written inside the `r\"\"\"...\"\"\"` literal terminated it, and the module + stopped importing. `compile()` is the cheapest possible stand-in for + the check ruff would run. + """ + from agent_flow.workflows.staircase_onboard.survey import PROBE_SOURCE + + compile(PROBE_SOURCE, "", "exec") + + +def test_the_probe_declares_the_marker_the_parser_looks_for(): + """Two literals, in two modules, that must agree.""" + from agent_flow.workflows.staircase_onboard.survey import MARKER, PROBE_SOURCE + + assert MARKER in PROBE_SOURCE + + +# ------------------------------------------------- how the op plane is enumerated + +#: A stand-in `torch` whose `ops.trtllm` is LAZY, exactly as the real one is: +#: `dir()` lists only what something has already resolved. The schema registry +#: knows the rest. Measured on a real GB300 build: `dir()` -> 91, +#: `_jit_get_all_schemas()` -> 273. +_FAKE_TORCH = ''' +import sys, types + +class _Schema: + def __init__(self, name): + self.name = name + +class _LazyNamespace: + """dir() reports only the two names somebody happened to touch.""" + def __dir__(self): + return ["already_resolved_a", "already_resolved_b"] + +torch = types.ModuleType("torch") +torch.__version__ = "9.9.9-fake" +torch.ops = types.SimpleNamespace(trtllm=_LazyNamespace()) +torch._C = types.SimpleNamespace( + _jit_get_all_schemas=lambda: [ + _Schema("trtllm::allgather"), + _Schema("trtllm::fp4_quantize"), + _Schema("trtllm::fused_qk_norm_rope"), + _Schema("aten::add"), + ] +) +sys.modules["torch"] = torch + +trtllm = types.ModuleType("tensorrt_llm") +trtllm.__version__ = "0.0.0-fake" +sys.modules["tensorrt_llm"] = trtllm +''' + + +def _run_probe_with_fake_torch(repo, arch="sm_103"): + import subprocess + import sys + + from agent_flow.workflows.staircase_onboard.survey import PROBE_SOURCE + + proc = subprocess.run( + [sys.executable, "-c", _FAKE_TORCH + PROBE_SOURCE, str(repo), arch], + capture_output=True, + text=True, + timeout=120, + ) + return parse_inventory(proc.stdout + proc.stderr, arch=arch) + + +def test_the_op_plane_is_read_from_the_schema_registry_not_from_dir(tmp_path): + """`torch.ops.` is a lazy proxy and `dir()` under-reports it badly. + + This was live: the survey reported 91 of 274 ops, and therefore reported + thirteen perfectly good catalog entries as "Rotted --- do not spend + these", including `allgather`, `fp4_quantize`, `fused_qk_norm_rope` and + `causal_conv1d_fwd`. A plan drafted against that would route around the + exact vocabulary a quantized hybrid MoE target needs, and dispatch a + smith to re-onboard each one. + """ + repo = _fake_repo(tmp_path, [("comm/allgather.py", "torch.ops.trtllm.allgather", _INLINE_PASS)]) + inventory = _run_probe_with_fake_torch(repo) + + assert "allgather" in inventory.trtllm_ops + assert "fp4_quantize" in inventory.trtllm_ops + # And the lazily-resolved names are still there, so the registry is a + # superset rather than a replacement. + assert "already_resolved_a" in inventory.trtllm_ops + + +def test_ops_from_other_namespaces_are_not_claimed_as_trtllms(tmp_path): + repo = _fake_repo(tmp_path, [("comm/allgather.py", "torch.ops.trtllm.allgather", _INLINE_PASS)]) + assert "add" not in _run_probe_with_fake_torch(repo).trtllm_ops + + +def test_an_entry_whose_op_is_registered_is_not_reported_as_rotted(tmp_path): + """The consequence, stated as the thing a plan reads.""" + repo = _fake_repo(tmp_path, [("comm/allgather.py", "torch.ops.trtllm.allgather", _INLINE_PASS)]) + entry = _run_probe_with_fake_torch(repo).catalog[0] + assert entry.impl_resolves + assert entry.spendable diff --git a/agent-flow/tests/workflows/staircase_onboard/test_survey_probes.py b/agent-flow/tests/workflows/staircase_onboard/test_survey_probes.py new file mode 100644 index 000000000000..b9a37652fa0e --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_survey_probes.py @@ -0,0 +1,504 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The checkpoint probes added for architectures the first two runs did not have. + +Each of these was a MISS on a real checkpoint, and the shape of the miss is +the same every time: the config says something structural, no probe claims it, +the key lands in `unrecognized`, and the runtime area that knows about it is +never surfaced to the plan phase. The key showing up in `unrecognized` is the +safety net working — but a net is not a map. + +Measured on Qwen3.8-Flash-Next-NVFP4: before these probes the survey produced +10 demands and reached 14 runtime areas; after, 16 demands and 20 areas. The +six it had been silent about were the sparse-attention indexer, mrope, plain +RoPE, hyper-connections, the vision tower, and a 20-million-row embedding +table. +""" + +from __future__ import annotations + +import pytest + +from agent_flow.workflows.staircase_onboard.runtime import areas_for_demands, load_index +from agent_flow.workflows.staircase_onboard.survey import read_checkpoint_surface + + +def _names(config): + return set(read_checkpoint_surface(config).demand_names) + + +def _detail(config, name): + for demand in read_checkpoint_surface(config).demands: + if demand.name == name: + return demand.detail + raise AssertionError(f"{name} not produced; got {sorted(_names(config))}") + + +def _areas(config): + return { + a.name + for a in areas_for_demands(read_checkpoint_surface(config).demand_names, index=load_index()) + } + + +# ------------------------------------------------- the sparse-attention indexer + + +def test_an_indexer_is_a_sparse_attention_demand(): + """A learned selector, not a mask. + + A window leaves cache addressing alone; an indexer decides per query + which past positions are attended to, so the top-k selection is part of + the forward and part of what a receipt certifies. A target that computes + full attention here boots and scores close on short prompts. + """ + config = {"indexer_budget": 2048, "indexer_n_heads": 4, "indexer_head_dim": 128} + assert "attention.indexer" in _names(config) + assert "budget=2048" in _detail(config, "attention.indexer") + assert "sparse_attention" in _areas(config) + + +def test_the_deepseek_spelling_is_also_an_indexer(): + assert "attention.indexer" in _names({"index_topk": 2048, "index_n_heads": 64}) + + +def test_a_plain_attention_config_produces_no_indexer_demand(): + assert "attention.indexer" not in _names({"num_attention_heads": 32, "head_dim": 128}) + + +# ----------------------------------------------------------------------- rope + + +def test_rope_is_read_from_rope_parameters_as_well_as_rope_scaling(): + """Transformers 5.x renamed the block and the probe only knew the old name. + + The consequence on a real checkpoint was that the ONLY rope demand + produced was `rope.partial`, so mrope — the thing a multimodal target is + most likely to get wrong — was never mentioned. + """ + config = {"rope_parameters": {"rope_type": "default", "rope_theta": 10000000}} + assert "rope.plain" in _names(config) + assert "theta=10000000" in _detail(config, "rope.plain") + + +def test_a_scaled_rope_under_the_new_name_is_still_scaling(): + config = {"rope_parameters": {"rope_type": "yarn", "factor": 4}} + assert "rope.scaling" in _names(config) + assert "yarn" in _detail(config, "rope.scaling") + + +def test_mrope_is_detected_under_the_new_name_with_its_sections(): + config = { + "rope_parameters": { + "rope_type": "default", + "rope_theta": 10000000, + "mrope_section": [11, 11, 10], + "mrope_interleaved": True, + } + } + detail = _detail(config, "rope.mrope") + assert "[11, 11, 10]" in detail + assert "interleaved" in detail + + +def test_the_old_rope_scaling_spelling_still_works(): + """Nothing that already worked may change shape.""" + config = {"rope_scaling": {"rope_type": "linear", "factor": 2}} + assert "rope.scaling" in _names(config) + + +def test_partial_rotary_is_read_from_either_place(): + assert "rope.partial" in _names({"partial_rotary_factor": 0.25}) + assert "rope.partial" in _names({"rope_parameters": {"partial_rotary_factor": 0.25}}) + + +# --------------------------------------------------------- hyper-connections + + +def test_hyper_connections_reach_their_runtime_area(): + """The area existed and no demand routed to it, so it was never read. + + The residual stream is not one tensor: `hc_count` parallel streams are + mixed at every block boundary, which changes the shape of every residual + add. A target that assumes one stream is a different model. + """ + config = {"hc_count": 4, "hc_lowrank": 320} + assert "residual.hyper_connection" in _names(config) + detail = _detail(config, "residual.hyper_connection") + assert "4 parallel residual streams" in detail + assert "320" in detail + assert "hyper_connection" in _areas(config) + + +def test_the_generic_spelling_is_also_hyper_connections(): + assert "residual.hyper_connection" in _names({"num_residual_streams": 2}) + + +# ------------------------------------------------------------- vision towers + + +def test_a_vision_tower_reaches_the_multimodal_area(): + """A vision tower reaches the multimodal area. + + Same shape of miss: the survey PRINTED "Multimodal checkpoint" and + produced no demand, so the multimodal runtime area stayed unread. + """ + config = {"vision_config": {"depth": 27}, "image_token_id": 248056} + assert "multimodal.tower" in _names(config) + assert "multimodal" in _areas(config) + + +def test_the_tower_is_read_from_the_outer_config_not_the_text_tower(): + """`scoped` is the text tower by then; the vision config is its sibling.""" + config = {"text_config": {"num_attention_heads": 24}, "vision_config": {"depth": 27}} + assert "multimodal.tower" in _names(config) + + +def test_placeholder_tokens_alone_are_enough(): + """Some checkpoints declare the tokens and inline the tower.""" + assert "multimodal.tower" in _names({"image_token_id": 1, "video_token_id": 2}) + + +def test_a_text_only_checkpoint_produces_no_tower_demand(): + assert "multimodal.tower" not in _names({"num_attention_heads": 32, "vocab_size": 1000}) + + +# ------------------------------------------------------- huge auxiliary tables + + +def test_a_secondary_table_far_larger_than_the_vocabulary_is_a_residency_decision(): + """20 M rows against a 248 K vocabulary was 53.7 GB of a 124 GB checkpoint. + + Whether it is resident decides how much of the device is left for KV, so + it is a decision rather than a default — which is what the `host_offload` + feature exists to make explicit. + """ + config = {"vocab_size": 248320, "ngram_vocab_size_base": 20000000} + assert "embedding.huge_table" in _names(config) + detail = _detail(config, "embedding.huge_table") + assert "20,000,000" in detail + assert "80x" in detail + assert {"host_offload", "embedding"} <= _areas(config) + + +def test_a_merely_larger_table_is_not_huge(): + """The threshold is an order of magnitude, not "different". + + A padded or slightly-extended vocabulary is ordinary and must not + surface a residency decision nobody has to make. + """ + assert "embedding.huge_table" not in _names({"vocab_size": 100000, "extra_vocab_size": 120000}) + + +def test_the_primary_vocabulary_is_never_its_own_secondary_table(): + assert "embedding.huge_table" not in _names({"vocab_size": 248320}) + + +@pytest.mark.parametrize("value", ["20000000", None, True]) +def test_a_non_integer_table_size_is_ignored_rather_than_crashing(value): + """Configs in the wild carry strings and nulls. + + A probe that raises on one costs the whole survey, which is never worth it. + """ + assert "embedding.huge_table" not in _names( + {"vocab_size": 248320, "ngram_vocab_size_base": value} + ) + + +# ------------------------------------------------------------ the hybrid stack + + +def _unrecognized(config): + return set(read_checkpoint_surface(config).unrecognized) + + +_MIMO_HYBRID = { + # 39 of one kind and 9 of the other, in the real checkpoint's own order. + "hybrid_layer_pattern": [0, 1, 1, 1, 1] + ([0] + [1] * 5) * 7 + [0], + "num_hidden_layers": 48, + "num_attention_heads": 64, + "num_key_value_heads": 4, + "head_dim": 192, + "v_head_dim": 128, + "rope_theta": 1e7, + "swa_num_attention_heads": 64, + "swa_num_key_value_heads": 8, + "swa_head_dim": 192, + "swa_v_head_dim": 128, + "swa_rope_theta": 1e4, +} + + +def test_a_per_layer_pattern_is_a_hybrid_stack_not_a_window(): + """`attention.windowing` said `sliding_window=128` and nothing else. + + That reads as one window over a uniform stack. It is 48 layers of two + different kinds, and the difference decides whether the cache manager is + asked for one geometry or two. + """ + assert "attention.hybrid" in _names(_MIMO_HYBRID) + detail = _detail(_MIMO_HYBRID, "attention.hybrid") + assert "48 layers" in detail + assert "39 of `1`" in detail and "9 of `0`" in detail + assert {"sparse_attention", "kv_cache"} <= _areas(_MIMO_HYBRID) + + +def test_the_survey_reports_the_pattern_and_refuses_to_interpret_it(): + """Which value means sliding-window is the config author's convention. + + On this checkpoint `1` is sliding-window and `0` is full attention, and + nothing in `config.json` says so. A survey that guessed the polarity + would silently window the wrong 39 layers, which is worse than handing + the plan phase the counts and the question. + """ + detail = _detail(_MIMO_HYBRID, "attention.hybrid") + assert "question for the plan phase" in detail + for leak in ("sliding-window layers:", "39 sliding", "9 full"): + assert leak not in detail + + +def test_different_head_geometry_between_the_two_kinds_is_the_headline(): + """Not the same shape under two masks: two shapes in one stack. + + The KV head counts differ (8 against 4) and the RoPE bases differ by + three orders of magnitude, which is the weight manifest's shard math and + the attention entry's positional scheme, not a mask. + """ + detail = _detail(_MIMO_HYBRID, "attention.hybrid") + assert "DIFFERENT geometry" in detail + assert "swa_num_key_value_heads=8 against num_key_value_heads=4" in detail + assert "swa_rope_theta" in detail + # The ones that agree must not be reported as differences. + assert "swa_head_dim=192 against head_dim=192" not in detail + + +def test_per_kind_geometry_that_agrees_is_said_to_agree(): + """A hybrid stack whose kinds share a shape is a different, easier target.""" + config = { + "hybrid_layer_pattern": [0, 1] * 8, + "num_key_value_heads": 8, + "swa_num_key_value_heads": 8, + } + detail = _detail(config, "attention.hybrid") + assert "declared and agrees" in detail + assert "DIFFERENT" not in detail + + +def test_an_exact_period_is_reported_and_an_inexact_one_is_not_invented(): + """An exact period is reported, and an inexact one is not invented. + + A stack whose global layers land every sixth position and one whose + first block is short are different stacks. + """ + assert "repeating every 2 layer(s)" in _detail( + {"hybrid_layer_pattern": [0, 1] * 12}, "attention.hybrid" + ) + assert "no fixed period" in _detail(_MIMO_HYBRID, "attention.hybrid") + + +def test_a_period_that_occurs_only_once_is_not_a_period(): + """A period that occurs only once is not a period. + + A 48-entry list trivially satisfies p=47 when its last entry equals its + first, and reporting that as a repeat is worse than reporting none. + """ + from agent_flow.workflows.staircase_onboard.survey import _repeating_period + + assert _repeating_period([0, 1, 1, 1, 0]) is None + assert _repeating_period([0, 1, 0, 1]) == 2 + + +@pytest.mark.parametrize("key", ["hybrid_layer_pattern", "layer_types", "moe_layer_pattern"]) +def test_every_per_layer_pattern_spelling_fires(key): + """Config authors spell the same concept a dozen ways. + + `*_layer_pattern` catches the ones nobody has written down yet. + """ + assert "attention.hybrid" in _names({key: ["full", "swa", "full", "swa"]}) + + +def test_an_index_list_names_one_kind_and_implies_the_rest(): + """`full_attention_idx` says which layers are full and leaves the rest.""" + detail = _detail( + {"full_attention_idx": [0, 9, 18, 27], "num_hidden_layers": 36}, "attention.hybrid" + ) + assert "36 layers, 4 of them named by `full_attention_idx` and 32 not" in detail + + +def test_a_per_layer_list_that_is_not_attention_does_not_fire(): + """`moe_layer_freq` is the same shape and a different question.""" + assert "attention.hybrid" not in _names({"moe_layer_freq": [0] + [1] * 47}) + + +def test_a_uniform_stack_produces_no_hybrid_demand(): + assert "attention.hybrid" not in _names({"num_attention_heads": 32, "head_dim": 128}) + + +def test_the_per_kind_geometry_keys_stop_being_reported_as_unrecognized(): + """They are the evidence for a demand now, not leftovers.""" + left = _unrecognized(_MIMO_HYBRID) + assert not {k for k in left if k.startswith("swa_")} + assert "hybrid_layer_pattern" not in left + + +# ------------------------------------------------------------ attention sinks + + +_SINKS = {"add_swa_attention_sink_bias": True, "add_full_attention_sink_bias": False} + + +def test_a_sink_bias_names_the_layer_kind_that_carries_it(): + """A sink on every layer and a sink on the local layers only are different models. + + The layer kind is in the key name and nowhere else. + """ + detail = _detail(_SINKS, "attention.sink") + assert "swa layers=True" in detail + assert "full layers=False" in detail + assert "sparse_attention" in _areas(_SINKS) + + +def test_an_unprefixed_sink_flag_is_read_as_applying_to_every_layer(): + assert "all layers=True" in _detail({"attention_sink": True}, "attention.sink") + + +def test_the_sink_keys_stop_being_reported_as_unrecognized(): + assert not _unrecognized(_SINKS) + + +def test_the_count_of_sink_tokens_is_still_its_own_demand(): + """`attention.sinks` reads the count spelling; this probe reads the per-layer-kind boolean. + + Neither may swallow the other. + """ + assert "attention.sinks" in _names({"sinks": 4}) + assert "attention.sink" not in _names({"sinks": 4}) + + +# ------------------------------------------------- asymmetric qk / v head dims + + +def test_a_v_head_dim_that_differs_is_stated_as_its_own_clause(): + """The single hardest kernel-feasibility question this family asks. + + It decides whether any existing paged-attention entry can be reused at + all, or whether the target needs a padded-V fallback — so a drafter that + reads only `head_dim` sizes the queue without the hardest item in it. + """ + detail = _detail({"head_dim": 192, "v_head_dim": 128}, "attention.head_dim") + assert "v_head_dim=128 differs from head_dim=192" in detail + assert "asymmetric qk/v" in detail + assert "v_head_dim" not in _unrecognized({"head_dim": 192, "v_head_dim": 128}) + + +def test_a_v_head_dim_equal_to_head_dim_adds_no_clause(): + """The ordinary case must stay one line.""" + detail = _detail({"head_dim": 128, "v_head_dim": 128}, "attention.head_dim") + assert detail == "128, declared not derived" + + +def test_a_checkpoint_with_no_v_head_dim_is_unchanged(): + assert _detail({"head_dim": 256}, "attention.head_dim") == "256, declared not derived" + + +# --------------------------------------------------------------- moe routing + + +def test_the_routing_method_is_reported_beside_the_scoring_function(): + """The scoring function says how an expert is SCORED; these say how the top-k is SELECTED. + + They pick a different routing entry rather than a different constant. + """ + detail = _detail( + { + "n_routed_experts": 256, + "num_experts_per_tok": 8, + "scoring_func": "sigmoid", + "topk_method": "noaux_tc", + "n_group": 1, + "topk_group": 1, + }, + "moe.routing", + ) + assert "scoring=sigmoid" in detail + assert "topk_method=noaux_tc" in detail + assert "n_group=1" in detail + assert "topk_group=1" in detail + + +def test_a_plain_moe_reports_no_routing_method_it_does_not_have(): + detail = _detail({"n_routed_experts": 8, "num_experts_per_tok": 2}, "moe.routing") + assert "topk_method" not in detail + assert "topk_group" not in detail + + +# -------------------------------------------------- the invariant behind all of it + + +def test_every_demand_a_runtime_area_declares_is_one_some_probe_can_produce(): + """A dangling demand name is an area that can never be surfaced. + + This is the cheap, permanent version of the audit that found the six + misses above: an area may only claim demands the survey knows how to + emit, and a probe may only emit demands some area claims. + """ + import pathlib + import re + + import yaml + + from agent_flow.workflows.staircase_onboard import survey as survey_mod + + index = yaml.safe_load( + (pathlib.Path(survey_mod.__file__).parent / "runtime_index.yaml").read_text() + ) + source = pathlib.Path(survey_mod.__file__).read_text() + produced = set(re.findall(r'Demand\(\s*"([a-z_.]+)"', source)) + + declared: dict[str, list[str]] = {} + for name, area in index.items(): + if name == "_overview": + continue + for demand in area.get("demands") or []: + declared.setdefault(demand, []).append(name) + + dangling = {d: owners for d, owners in declared.items() if d not in produced} + assert not dangling, f"areas declare demands no probe emits: {dangling}" + + orphan = sorted(d for d in produced if d not in declared) + assert not orphan, f"probes emit demands no area claims: {orphan}" + + +# ------------------------------------------- a DSpark drafter under mtp.* + +# DeepSeek-V4.1-Flash's own `text_config`, trimmed to the keys that matter. +_DSPARK = { + "architectures": ["DeepseekV41ForCausalLM"], + "text_config": { + "num_nextn_predict_layers": 3, + "dspark_block_size": 5, + "dspark_markov_rank": 256, + "dspark_noise_token_id": 128799, + "dspark_target_layer_ids": [37, 38, 39], + }, +} + + +def test_a_dspark_drafter_is_not_reported_as_mtp(): + """The survey is what the plan drafter reads first. + + Before this it said "MTP is the checkpoint's own" for DeepSeek-V4.1, whose + `mtp.*` stages are a DSpark block drafter -- which is how a plan ends up + requesting the feature that builds the wrong drafter. + """ + detail = _detail(_DSPARK, "speculation.heads") + assert "`dspark`, NOT `mtp`" in detail + assert "block 5" in detail + assert "MTP is the checkpoint's own" not in detail + assert not {k for k in _unrecognized(_DSPARK) if k.startswith("dspark_")} + + +def test_a_plain_mtp_checkpoint_still_reads_as_mtp(): + detail = _detail({"num_nextn_predict_layers": 1}, "speculation.heads") + assert "MTP is the checkpoint's own" in detail diff --git a/agent-flow/tests/workflows/staircase_onboard/test_task_schema.py b/agent-flow/tests/workflows/staircase_onboard/test_task_schema.py new file mode 100644 index 000000000000..32aedd4d5545 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_task_schema.py @@ -0,0 +1,1116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for staircase-onboard's ``task.yaml`` validation.""" + +from __future__ import annotations + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard 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 reintroduces the TP-only gate under a new 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", "eplb", "dwdp", "multimodal"} + for feature, skill in task_schema.FEATURE_SKILLS.items(): + assert skill.startswith("staircase-modeling-"), feature + + +def test_slurm_kind_requires_partition_and_image(tmp_path): + spec = _valid_spec(tmp_path) + spec["environment"]["kind"] = "slurm" + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert "partition" in str(excinfo.value) + assert "container_image" in str(excinfo.value) + + +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", "eplb"] + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert task_schema.assembly_units(data) == ["core", "mtp", "eplb"] + + +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.staircase_onboard 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. + + `feature:` instead of `features:` used to validate cleanly and resolve + to `[]`. The run then onboarded core, gated it green, and reported + success for a feature set it never built — a green run that delivered + 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): + """`runtme:` reproduced `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"], + } + spec["environment"] = { + "kind": "slurm", + "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__).parent / "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. + Before this test it documented 6 of 11 top-level fields — and the five + it omitted (`runtime`, `target`, `scope`, `source_embargo`, and the + accuracy protocol options) were precisely the ones added because a run + had failed without them. A user copying the example would have hit the + same failures in the same order. + """ + import re + from pathlib import Path + + example = Path(task_schema.__file__).parent / "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__).parent / "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. + + `skip_substrate_gate` was exactly that: an escape hatch from the + longest preflight step, reachable only by editing Python. + """ + import inspect + + from agent_flow.workflows.staircase_onboard.cli import _parse_args + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + args = vars(_parse_args(["--task", "t.yaml"])) + parameters = inspect.signature(StaircaseOnboardWorkflow.__init__).parameters + # Injection seams for the tests, not operator knobs. + injected = {"self", "workspace", "prompts", "command_runner", "shell_runner"} + unreachable = sorted( + name + for name in parameters + if name not in injected + and name.replace("-", "_") not in args + # `--reviewer` is spelled shorter than `reviewer_mode`. + and name != "reviewer_mode" + ) + assert not unreachable, f"no CLI flag reaches: {unreachable}" + assert args["reviewer"] in ("off", "light", "full") + + +def test_the_substrate_gate_flag_reaches_the_workflow(): + from agent_flow.workflows.staircase_onboard.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 + + +# ------------------------------- per-workload launch substrate + + +def _with_workloads(tmp_path, workloads): + spec = _valid_spec(tmp_path) + spec.setdefault("environment", {})["workloads"] = workloads + return _write(tmp_path, spec) + + +def test_a_gate_may_declare_its_own_substrate(tmp_path): + """The right substrate depends on how long the job runs. + + A shared session imports a 26 GB container once and costs nothing per job + after that -- correct for a 1-minute entry test, and one onboarding ran 443 + of them. It is wrong for a 10-15 minute gate, where the import is 30-50% + overhead and isolation is the point: a shared session has one failure + domain, so a gate that leaks a server, kills a process group or collides on + PMIx takes everything else with it. + """ + path = _with_workloads(tmp_path, {"gate": {"kind": "slurm", "time_limit": "01:00:00"}}) + data = task_schema.load_and_validate_task_yaml(path) + assert data["environment"]["workloads"]["gate"]["kind"] == "slurm" + + +def test_a_workload_the_orchestrator_never_dispatches_under_is_rejected(tmp_path): + """An override that never applies is invisible -- the run just uses the base.""" + path = _with_workloads(tmp_path, {"smithy": {"kind": "slurm"}}) + with pytest.raises(task_schema.TaskSchemaError, match="smithy"): + task_schema.load_and_validate_task_yaml(path) + + +def test_a_workload_with_an_unknown_kind_is_rejected(tmp_path): + path = _with_workloads(tmp_path, {"gate": {"kind": "kubernetes"}}) + with pytest.raises(task_schema.TaskSchemaError, match="kind"): + task_schema.load_and_validate_task_yaml(path) + + +def test_a_workload_with_an_unknown_key_is_rejected(tmp_path): + """The same closed-block rule as `environment` itself, one level down.""" + path = _with_workloads(tmp_path, {"gate": {"partitionn": "batch"}}) + with pytest.raises(task_schema.TaskSchemaError, match="partitionn"): + task_schema.load_and_validate_task_yaml(path) + + +def test_a_workload_may_not_carry_its_own_workloads(tmp_path): + """A nesting nothing resolves is a substrate that silently never applies.""" + path = _with_workloads(tmp_path, {"gate": {"workloads": {"gate": {"kind": "slurm"}}}}) + with pytest.raises(task_schema.TaskSchemaError, match="workloads"): + task_schema.load_and_validate_task_yaml(path) + + +def test_a_workload_block_that_is_not_a_mapping_is_rejected(tmp_path): + path = _with_workloads(tmp_path, {"gate": "slurm"}) + with pytest.raises(task_schema.TaskSchemaError, match="must be a mapping"): + task_schema.load_and_validate_task_yaml(path) + + +def test_workloads_that_is_not_a_mapping_is_rejected(tmp_path): + path = _with_workloads(tmp_path, ["gate"]) + with pytest.raises(task_schema.TaskSchemaError, match="workloads"): + task_schema.load_and_validate_task_yaml(path) + + +def test_omitting_workloads_is_still_valid(tmp_path): + """Nothing that worked may change shape.""" + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, _valid_spec(tmp_path))) + assert "workloads" not in data["environment"] + + +# ------------------------------- 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, as they had to be, the spec stops working on any other + machine -- one run produced a 220-line task.yaml with six absolute paths, + and its most reusable artifact was 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 shape this run actually wrote: 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_a_witness_boot_spec_resolves_the_same_way(tmp_path): + """`witness..boot_spec` is written beside the spec for the same reason. + + Exercised against `_resolve_spec_relative_paths` directly, because the + `witness` block in this checkout is closed over `contrast_env`, + `expected_mib` and `contrast_command` -- `boot_spec` is resolved here so + that whichever change admits the key inherits the portability, rather than + having to remember it. + """ + data = { + "accuracy": { + "witness": { + "host_offload": { + "contrast_env": {"X": "0"}, + "boot_spec": "boot.residency.yaml", + } + } + } + } + task_schema._resolve_spec_relative_paths(data, tmp_path) + resolved = data["accuracy"]["witness"]["host_offload"]["boot_spec"] + assert resolved == str((tmp_path / "boot.residency.yaml").resolve()) + + +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 new feature kinds + + +def test_the_new_feature_kinds_map_to_their_skills(): + skills = task_schema.FEATURE_SKILLS + assert skills["dspark"] == "staircase-modeling-speculator" + assert skills["decoder_replay"] == "staircase-modeling-bounded-replay" + assert skills["encoder_replay"] == "staircase-modeling-bounded-replay" + assert skills["kv_quant"] == "staircase-modeling-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_its_own_launcher_is_refused(tmp_path): + """Nothing reads TRTLLM_DISAGG: leg B would be leg A, and the witness vacuous.""" + block = {"contrast_env": {"TRTLLM_DISAGG": "1"}} + with pytest.raises(task_schema.TaskSchemaError, match="contrast_command' is required"): + task_schema.load_and_validate_task_yaml( + _with_witness(tmp_path, ["disagg"], {"disagg": block}) + ) + + +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. F039.""" + 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): + """Every task written before the workload could be declared looks like this.""" + 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 (F036) + + +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): + """Every task.yaml written before the key keeps validating, and 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}}) + ) diff --git a/agent-flow/tests/workflows/staircase_onboard/test_transport.py b/agent-flow/tests/workflows/staircase_onboard/test_transport.py new file mode 100644 index 000000000000..7b73c3be9675 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_transport.py @@ -0,0 +1,376 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The external agent transport: the orchestrator's turn, delivered as a file. + +What these protect is the hand-off contract, because both sides are written +independently and a mismatch presents as a hang rather than an error. The +ordering rule (`.done` last) and the system prompt travelling with the turn +are the two that would be silently wrong. +""" + +from __future__ import annotations + +import json +import threading +import time +from pathlib import Path + +import pytest + +from agent_flow.workflows.staircase_onboard.transport import ( + DISPATCH_DIRNAME, + ExternalAgentTransport, + ExternalTransportTimeout, +) + + +def _answer_after(transport: ExternalAgentTransport, request_id: str, result: str, delay=0.05): + """Play the driver from another thread.""" + + def run(): + time.sleep(delay) + transport.answer(request_id, result) + + thread = threading.Thread(target=run, daemon=True) + thread.start() + return thread + + +def test_a_turn_round_trips(tmp_path): + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + _answer_after(transport, "001-smith", "LANDED, receipt written") + assert transport(role="smith", what="onboard foo", prompt="do it") == ( + "LANDED, receipt written" + ) + + +def test_the_system_prompt_travels_with_the_turn(tmp_path): + """A driver given only the user prompt runs a DIFFERENT agent. + + No role, no invariants, no citation discipline, no embargo -- and the + answer still looks plausible, which is what makes the omission dangerous + rather than obvious. + """ + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + _answer_after(transport, "001-assembler", "ok") + transport(role="assembler", what="core", prompt="p", system_prompt="YOU ARE THE ASSEMBLER") + + root = tmp_path / DISPATCH_DIRNAME + # `startswith`, not equality: the transport appends its own note about the + # one instruction it cannot honour. See `_transport_note`. + assert ( + (root / "001-assembler.system.md") + .read_text(encoding="utf-8") + .startswith("YOU ARE THE ASSEMBLER") + ) + request = json.loads((root / "001-assembler.request.json").read_text(encoding="utf-8")) + assert request["system_prompt_file"].endswith("001-assembler.system.md") + + +def test_the_request_names_the_decision_file_the_orchestrator_will_read(tmp_path): + """The driver has to know where the typed decision goes. + + Otherwise the turn completes and the orchestrator finds nothing to branch on. + """ + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + _answer_after(transport, "001-reviewer", "ok") + transport(role="reviewer", what="core", prompt="p", decision_file="/ws/decisions/r-1.yaml") + request = json.loads( + (tmp_path / DISPATCH_DIRNAME / "001-reviewer.request.json").read_text(encoding="utf-8") + ) + assert request["decision_file"] == "/ws/decisions/r-1.yaml" + + +def test_done_is_written_after_the_result(tmp_path): + """A reader that sees `.done` must be able to trust the result beside it. + + Same rule as the compute session's `.rc`, and for the same reason: the + two are read by a different process than the one that writes them. + """ + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + root = tmp_path / DISPATCH_DIRNAME + transport.answer("007-smith", "the answer") + assert (root / "007-smith.done").stat().st_mtime_ns >= ( + root / "007-smith.result.md" + ).stat().st_mtime_ns + + +def test_pending_lists_only_unanswered_requests_in_order(tmp_path): + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + for n, role in ((1, "smith"), (2, "smith"), (3, "assembler")): + (tmp_path / DISPATCH_DIRNAME / f"{n:03d}-{role}.request.json").write_text( + json.dumps({"id": f"{n:03d}-{role}"}), encoding="utf-8" + ) + transport.answer("002-smith", "done") + assert [p.name for p in transport.pending()] == [ + "001-smith.request.json", + "003-assembler.request.json", + ] + + +def test_ids_continue_across_a_resume(tmp_path): + """Renumbering from 1 would overwrite the previous attempt's request files. + + Those files are the evidence for what the run asked for. + """ + root = tmp_path / DISPATCH_DIRNAME + root.mkdir(parents=True) + (root / "004-smith.request.json").write_text("{}", encoding="utf-8") + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + _answer_after(transport, "005-smith", "ok") + transport(role="smith", what="x", prompt="p") + assert (root / "005-smith.result.md").is_file() + + +def test_concurrent_turns_get_distinct_ids(tmp_path): + """Ten smiths run in one batch; two sharing an id would cross their answers.""" + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + seen: list[str] = [] + lock = threading.Lock() + + def one(): + with lock: + n = len(seen) + 1 + request_id = f"{n:03d}-smith" + _answer_after(transport, request_id, request_id, delay=0.2) + with lock: + seen.append(request_id) + + # Drive the transport itself concurrently and check the ids it minted. + results: list[str] = [] + threads = [] + for _ in range(5): + one() + for _ in range(5): + t = threading.Thread( + target=lambda: results.append(transport(role="smith", what="x", prompt="p")) + ) + t.start() + threads.append(t) + for t in threads: + t.join(timeout=10) + assert sorted(results) == sorted(seen) + + +def test_a_stale_driver_heartbeat_fails_fast_instead_of_waiting_hours(tmp_path): + """The three-hour turn timeout bounds a slow turn, not a dead driver.""" + transport = ExternalAgentTransport(tmp_path, poll_s=0.01, timeout_s=60) + beat = tmp_path / DISPATCH_DIRNAME / "driver.heartbeat" + beat.write_text("x", encoding="utf-8") + import os + + old = time.time() - 10_000 + os.utime(beat, (old, old)) + with pytest.raises(ExternalTransportTimeout, match="heartbeat is stale"): + transport(role="smith", what="x", prompt="p") + + +def test_no_heartbeat_at_all_is_not_treated_as_a_dead_driver(tmp_path): + """A driver that does not bother writing one must not be penalized.""" + transport = ExternalAgentTransport(tmp_path, poll_s=0.01, timeout_s=60) + _answer_after(transport, "001-smith", "ok") + assert transport(role="smith", what="x", prompt="p") == "ok" + + +def test_the_timeout_names_the_prompt_file_so_the_turn_can_be_run_by_hand(tmp_path): + transport = ExternalAgentTransport(tmp_path, poll_s=0.01, timeout_s=0.05) + with pytest.raises(ExternalTransportTimeout, match="001-smith.prompt.md"): + transport(role="smith", what="x", prompt="p") + + +def test_an_answer_with_no_result_file_is_empty_not_an_error(tmp_path): + """A driver may legitimately have nothing to say. + + The decision file is what the orchestrator branches on. + """ + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + root = tmp_path / DISPATCH_DIRNAME + + def run(): + time.sleep(0.05) + (root / "001-smith.done").write_text("1\n", encoding="utf-8") + + threading.Thread(target=run, daemon=True).start() + assert transport(role="smith", what="x", prompt="p") == "" + + +def test_files_are_renamed_into_place_never_written_in_place(tmp_path): + """A driver polling the directory must never see a half-written request.""" + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + _answer_after(transport, "001-smith", "ok") + transport(role="smith", what="x", prompt="p" * 10000) + leftovers = list((tmp_path / DISPATCH_DIRNAME).glob("*.tmp")) + assert not leftovers, f"staging files left behind: {leftovers}" + + +def test_the_workflow_accepts_a_transport_and_uses_it_for_every_role(tmp_path): + """Reachability, in the same spirit as test_gate_wiring.py.""" + import inspect + + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + assert "agent_transport" in inspect.signature(StaircaseOnboardWorkflow.__init__).parameters + for method in ("_dispatch", "_run_one_smith"): + source = inspect.getsource(getattr(StaircaseOnboardWorkflow, method)) + assert "_run_turn" in source, f"{method} bypasses the transport" + turn = inspect.getsource(StaircaseOnboardWorkflow._run_turn) + assert "system_prompt=" in turn + + +# ---------------------------------------- persistence, for the roles that need it + + +def test_prior_turns_restore_what_a_persistent_role_would_have_remembered(tmp_path): + """plan_drafter, plan_reviewer and assembler are persistent by design. + + The transport starts a fresh session per turn, so their prompts' claim + ("You are in the same session as your first draft") is false under it. A + drafter that believes it will revise a file it has never read. The + artifacts are on disk either way; naming them is what makes them usable. + """ + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + _answer_after(transport, "001-plan_drafter", "draft one done") + transport(role="plan_drafter", what="round 1", prompt="draft it") + + _answer_after(transport, "002-plan_drafter", "draft two done") + transport(role="plan_drafter", what="round 2", prompt="redraft it") + + request = json.loads( + (tmp_path / DISPATCH_DIRNAME / "002-plan_drafter.request.json").read_text(encoding="utf-8") + ) + assert [t["id"] for t in request["prior_turns"]] == ["001-plan_drafter"] + assert request["prior_turns"][0]["prompt_file"].endswith("001-plan_drafter.prompt.md") + + +def test_prior_turns_are_scoped_to_one_role(tmp_path): + """A smith must not be handed the assembler's history.""" + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + _answer_after(transport, "001-assembler", "ok") + transport(role="assembler", what="core", prompt="p") + _answer_after(transport, "002-smith", "ok") + transport(role="smith", what="x", prompt="p") + assert [t["id"] for t in transport.prior_turns("smith")] == ["002-smith"] + assert [t["id"] for t in transport.prior_turns("assembler")] == ["001-assembler"] + + +def test_an_unanswered_turn_is_not_prior_history(tmp_path): + """Concurrent smiths would otherwise hand each other half-finished turns.""" + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + (tmp_path / DISPATCH_DIRNAME / "001-smith.request.json").write_text("{}", encoding="utf-8") + assert transport.prior_turns("smith") == [] + + +# --------------------------------------------------- one orchestrator per workspace + + +def test_a_second_orchestrator_on_one_workspace_is_refused(tmp_path): + """Two on one workspace is not a race anybody notices. + + They share the checkpoint file, the dispatch directory and the request-id + counter, so they mint the SAME id, overwrite each other's request JSON, + and both unblock on one answer. Measured: a stale process a `pkill` had + failed to kill raced a fresh one for three seconds, and the survivor's + request was the older code's — missing the system-prompt pointer the + driver needed, which is a turn run with no role. + """ + import subprocess + import sys + import textwrap + + workspace = tmp_path / "ws" + script = textwrap.dedent( + f""" + import sys, time + sys.path.insert(0, {str(Path(__file__).resolve().parents[3])!r}) + from agent_flow.workflows.staircase_onboard.cli import _claim_workspace + from pathlib import Path + _claim_workspace(Path({str(workspace)!r})) + print("HELD", flush=True) + time.sleep(30) + """ + ) + holder = subprocess.Popen([sys.executable, "-c", script], stdout=subprocess.PIPE, text=True) + try: + assert holder.stdout.readline().strip() == "HELD" + second = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=60 + ) + assert second.returncode == 3 + assert "already holds" in second.stderr + finally: + holder.kill() + + +def test_the_lock_is_released_when_the_process_exits(tmp_path): + """A crash must not leave a workspace unusable; flock is held by the fd.""" + import subprocess + import sys + import textwrap + + workspace = tmp_path / "ws" + script = textwrap.dedent( + f""" + import sys + sys.path.insert(0, {str(Path(__file__).resolve().parents[3])!r}) + from agent_flow.workflows.staircase_onboard.cli import _claim_workspace + from pathlib import Path + _claim_workspace(Path({str(workspace)!r})) + """ + ) + for _ in range(2): + done = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True) + assert done.returncode == 0, done.stderr + + +# -------------------------------- what this transport cannot honour, said once + + +def test_the_system_prompt_says_the_progress_tool_is_unavailable(tmp_path): + """Every role prompt ends by telling the agent to call `append__progress`. + + The tool is real -- the in-process transport hands it to the session it + spawns -- but this transport has no session. The shared decision-file + block that follows says to write a file instead, so the prompt arrives + carrying two different endings and only the second one works here. + + Three auditor turns in a row noticed and picked correctly. The next role + might not: calling a missing tool looks from the inside like a turn that + ended, so the run would block for its full timeout on a turn that believed + it was finished. + """ + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + _answer_after(transport, "001-auditor", "ok") + transport(role="auditor", what="w", prompt="p", system_prompt="ROLE PROMPT") + + system = (tmp_path / DISPATCH_DIRNAME / "001-auditor.system.md").read_text(encoding="utf-8") + assert system.startswith("ROLE PROMPT") + assert "append_auditor_progress" in system + assert "not available here" in system + assert "`.done`" in system + + +def test_the_note_names_the_role_that_was_dispatched(tmp_path): + """A generic note would send a smith looking for the auditor's tool.""" + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + _answer_after(transport, "001-smith", "ok") + transport(role="smith", what="w", prompt="p", system_prompt="ROLE") + + system = (tmp_path / DISPATCH_DIRNAME / "001-smith.system.md").read_text(encoding="utf-8") + assert "append_smith_progress" in system + assert "append_auditor_progress" not in system + + +def test_the_role_prompt_is_never_edited_only_appended_to(tmp_path): + """The instruction is correct for the transport it was written for. + + Rewriting the role prompts to suit this one would make them wrong in + process, so the difference stays where it belongs. + """ + transport = ExternalAgentTransport(tmp_path, poll_s=0.01) + _answer_after(transport, "001-reviewer", "ok") + original = "ROLE\n\n## Ending your turn\n\nCall `append_reviewer_progress`.\n" + transport(role="reviewer", what="w", prompt="p", system_prompt=original) + + system = (tmp_path / DISPATCH_DIRNAME / "001-reviewer.system.md").read_text(encoding="utf-8") + assert system.startswith(original) diff --git a/agent-flow/tests/workflows/staircase_onboard/test_troubleshooting.py b/agent-flow/tests/workflows/staircase_onboard/test_troubleshooting.py new file mode 100644 index 000000000000..63b82337b66e --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_troubleshooting.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The troubleshooting doc, kept honest. + +Infrastructure knowledge rots the same way engine knowledge does, and in the +same direction: the symptom stays true while everything around it moves. So +this doc is keyed on **verbatim output** — the string you actually saw — and +these checks hold it to that. + +Two failure modes are worth spending a test on. A doc that names a check +which no longer exists sends the reader looking for code that is not there; +and a doc nobody is pointed at when they need it is a doc nobody reads, which +is the more common way a troubleshooting file dies. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import agent_flow.workflows.staircase_onboard as pkg +from agent_flow.workflows.staircase_onboard.workflow import TROUBLESHOOTING_HINT + +PACKAGE = Path(pkg.__file__).parent +DOC = PACKAGE / "TROUBLESHOOTING.md" + + +def _text() -> str: + return DOC.read_text(encoding="utf-8") + + +def test_the_doc_is_packaged_next_to_the_workflow(): + """Not in a docs tree somewhere: the failure message names this path.""" + assert DOC.is_file() + + +def test_the_failure_message_points_at_a_file_that_exists(): + """The hint is built from a path; a typo there is silent until it fires.""" + quoted = re.search(r"`([^`]+TROUBLESHOOTING\.md)`", TROUBLESHOOTING_HINT) + assert quoted, TROUBLESHOOTING_HINT + assert Path(quoted.group(1)).is_file() + + +def test_the_readme_sends_the_reader_here(): + assert "TROUBLESHOOTING.md" in (PACKAGE / "README.md").read_text(encoding="utf-8") + + +def test_every_check_the_doc_credits_actually_exists(): + """A `*Caught by:* `x`` line is a claim about code, so resolve it. + + A doc that credits a check which has been renamed or deleted is worse + than one that credits nothing: the reader concludes the case is handled + and stops looking. + """ + sources = "\n".join( + path.read_text(encoding="utf-8") + for path in PACKAGE.rglob("*.py") + if "__pycache__" not in path.parts + ) + # The test suite is the other legitimate home for a named check. + sources += "\n".join( + path.read_text(encoding="utf-8") for path in (Path(__file__).parent).glob("test_*.py") + ) + + # Only the *first* backticked token names the check. What follows is + # prose that legitimately quotes config keys and env vars — `runtime`, + # `TMPDIR` — and resolving those as symbols is how this check would + # become noisy enough to switch off. + missing: list[str] = [] + for line in _text().splitlines(): + if not line.startswith("*Caught by:*"): + continue + for name in re.findall(r"`([A-Za-z_][A-Za-z0-9_]*)\(?\)?`", line)[:1]: + if re.search( + rf"\b(def|class)\s+{re.escape(name)}\b|^{re.escape(name)}\s*=", sources, re.M + ): + continue + missing.append(f"{name} (in: {line.strip()})") + assert not missing, "TROUBLESHOOTING.md credits checks that do not exist:\n " + "\n ".join( + missing + ) + + +def test_the_preflight_table_matches_the_checks_preflight_runs(): + """The table is a promise about what runs before an agent exists.""" + from agent_flow.workflows.staircase_onboard import preflight + + text = _text() + for name in ( + "check_checkpoint", + "check_tmpdir", + "claim_devices", + "check_multirank_substrate", + "substrate_gate", + "agent_plane_gate", + ): + assert hasattr(preflight, name), name + assert f"`{name}`" in text, f"{name} is a preflight check the doc does not list" + + +def test_the_doc_obeys_the_citation_rule_it_relies_on(): + """It is keyed on verbatim output, which is the durable half. + + A line number would undercut the whole premise: the reader greps for the + symptom, and the symptom is what survives. Quoted diagnostics keep their + line numbers — those are part of the message the engine prints. + """ + from .test_skills import ( + BARE_LINE_REFERENCE, + LINE_CITATION, + _is_shown_as_an_example, + _quotes_a_diagnostic, + ) + + lines = _text().splitlines() + offences: list[str] = [] + for index, line in enumerate(lines): + if BARE_LINE_REFERENCE.search(line): + offences.append(f"{index + 1}: bare `:NNN`") + hit = LINE_CITATION.search(line) + if hit is None or _quotes_a_diagnostic(lines, index): + continue + if _is_shown_as_an_example(line, hit.start(), hit.end()): + continue + offences.append(f"{index + 1}: {hit.group(0)}") + assert not offences, "TROUBLESHOOTING.md cites code by line:\n " + "\n ".join(offences) + + +def test_the_open_problem_is_still_labelled_open(): + """`get_stats()` is why three features are unverified. Do not lose that.""" + text = _text() + assert "get_stats()" in text + assert "*Open.*" in text + + +def test_the_doc_is_declared_as_package_data(): + """Otherwise a non-editable install has no file behind the hint. + + This is the failure that cannot be caught by running from a source tree, + which is the only way these tests ever run — so it is asserted against + the packaging declaration instead. + """ + import tomllib + + pyproject = Path(__file__).resolve().parents[3] / "pyproject.toml" + data = tomllib.loads(pyproject.read_text(encoding="utf-8")) + declared = data["tool"]["setuptools"]["package-data"]["agent_flow.workflows.staircase_onboard"] + assert "TROUBLESHOOTING.md" in declared diff --git a/agent-flow/tests/workflows/staircase_onboard/test_vocabulary.py b/agent-flow/tests/workflows/staircase_onboard/test_vocabulary.py new file mode 100644 index 000000000000..28298926e621 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_vocabulary.py @@ -0,0 +1,420 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for ``vocabulary.yaml`` — the typed dispatch queue.""" + +from __future__ import annotations + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import vocabulary as vocab + + +def _ok_step(**overrides): + step = { + "step": "rmsnorm_q", + "values": {"dtype": "bf16", "head_dim": 128}, + "state": vocab.STATE_OK, + "entry": "catalog/norm/flashinfer_rmsnorm.py", + "contract_line": "certified: sm_103, bf16, head_dim in {64, 128}", + } + step.update(overrides) + return step + + +def _doc(**overrides): + doc = { + "verdict": vocab.VERDICT_ASSEMBLABLE, + "steps": [_ok_step()], + "queue": [], + } + doc.update(overrides) + return doc + + +def test_minimal_assemblable_document_parses(): + parsed = vocab.parse_vocabulary(_doc()) + assert parsed.verdict == vocab.VERDICT_ASSEMBLABLE + assert parsed.queue == () + assert parsed.is_ceiling is False + assert parsed.counts()[vocab.STATE_OK] == 1 + + +def test_ok_row_without_contract_line_is_rejected(): + """The quoted line is the evidence; without it nobody can audit the claim.""" + doc = _doc(steps=[_ok_step(contract_line="")]) + with pytest.raises(vocab.VocabularyError, match="contract_line"): + vocab.parse_vocabulary(doc) + + +def test_gap_row_without_a_tier_is_rejected(): + """A GAP has to say which rung supplies it before anything else.""" + doc = _doc( + verdict=vocab.VERDICT_AFTER_ADDITIONS, + steps=[{"step": "moe_gelu", "state": vocab.STATE_GAP}], + queue=[{"n": 1, "mode": vocab.MODE_ONBOARD, "target": "moe_gelu"}], + ) + with pytest.raises(vocab.VocabularyError, match="tier"): + vocab.parse_vocabulary(doc) + + +def test_a_wrapping_tier_must_name_what_it_wraps(): + doc = _doc( + verdict=vocab.VERDICT_AFTER_ADDITIONS, + steps=[{"step": "moe_gelu", "state": vocab.STATE_GAP, "tier": vocab.TIER_TRTLLM_OP}], + queue=[{"n": 1, "mode": vocab.MODE_ONBOARD, "target": "moe_gelu"}], + ) + with pytest.raises(vocab.VocabularyError, match="candidate"): + vocab.parse_vocabulary(doc) + + +def test_a_writing_tier_must_say_what_it_rejected(): + """Otherwise the ladder is a way to skip straight to writing a kernel.""" + doc = _doc( + verdict=vocab.VERDICT_AFTER_ADDITIONS, + steps=[{"step": "exotic_norm", "state": vocab.STATE_GAP, "tier": vocab.TIER_TRITON}], + queue=[{"n": 1, "mode": vocab.MODE_ONBOARD, "target": "exotic_norm"}], + ) + with pytest.raises(vocab.VocabularyError, match="tried"): + vocab.parse_vocabulary(doc) + + +def test_a_writing_tier_with_its_search_recorded_parses(): + """The case that matters for a model the engine does not implement.""" + parsed = vocab.parse_vocabulary( + _doc( + verdict=vocab.VERDICT_AFTER_ADDITIONS, + steps=[ + { + "step": "exotic_norm", + "state": vocab.STATE_GAP, + "tier": vocab.TIER_NATIVE_TORCH, + "tried": "no trtllm op; nothing under modules/** computes it", + } + ], + queue=[ + { + "n": 1, + "mode": vocab.MODE_ONBOARD, + "target": "exotic_norm", + "tier": vocab.TIER_NATIVE_TORCH, + } + ], + ) + ) + item = parsed.queue[0] + assert item.tier == vocab.TIER_NATIVE_TORCH + # The orchestrator needs to know this dispatch owes an implementation, + # not just a wrapper, so it can budget turns differently. + assert item.writes_implementation is True + assert "native torch" not in item.tier_brief.lower() or "compose" in item.tier_brief + + +def test_a_wrapping_dispatch_is_not_marked_as_writing(): + parsed = vocab.parse_vocabulary( + _doc( + verdict=vocab.VERDICT_AFTER_ADDITIONS, + steps=[ + { + "step": "moe_gelu", + "state": vocab.STATE_GAP, + "tier": vocab.TIER_INTREE_CALLABLE, + "candidate": "tensorrt_llm._torch.modules.x.y", + } + ], + queue=[{"n": 1, "mode": vocab.MODE_ONBOARD, "target": "moe_gelu", "tier": 2}], + ) + ) + assert parsed.queue[0].writes_implementation is False + + +def test_the_tier_survives_a_round_trip_through_run_state(): + """A resume must not silently demote a writing dispatch to a wrapping one.""" + item = vocab.QueueItem(n=1, mode=vocab.MODE_ONBOARD, target="x", tier=vocab.TIER_TRITON) + assert vocab.QueueItem(**item.as_dict()).tier == vocab.TIER_TRITON + + +def test_verify_dispatch_without_surface_is_rejected(): + """The excluded value is what smith is dispatched with; vagueness costs a run.""" + doc = _doc( + verdict=vocab.VERDICT_AFTER_ADDITIONS, + queue=[{"n": 1, "mode": vocab.MODE_VERIFY, "target": "thop_attention"}], + ) + with pytest.raises(vocab.VocabularyError, match="surface"): + vocab.parse_vocabulary(doc) + + +def test_the_two_chapters_may_not_overlap(): + doc = _doc( + verdict=vocab.VERDICT_AFTER_ADDITIONS, + queue=[ + {"n": 1, "mode": vocab.MODE_ONBOARD, "target": "fused_moe"}, + {"n": 2, "mode": vocab.MODE_VERIFY, "target": "fused_moe", "surface": "h=72"}, + ], + ) + with pytest.raises(vocab.VocabularyError, match="both chapters"): + vocab.parse_vocabulary(doc) + + +def test_duplicate_dispatch_numbers_are_rejected(): + doc = _doc( + verdict=vocab.VERDICT_AFTER_ADDITIONS, + queue=[ + {"n": 1, "mode": vocab.MODE_ONBOARD, "target": "a"}, + {"n": 1, "mode": vocab.MODE_ONBOARD, "target": "b"}, + ], + ) + with pytest.raises(vocab.VocabularyError, match="duplicate"): + vocab.parse_vocabulary(doc) + + +def test_dependency_on_a_missing_item_is_rejected(): + doc = _doc( + verdict=vocab.VERDICT_AFTER_ADDITIONS, + queue=[{"n": 1, "mode": vocab.MODE_ONBOARD, "target": "a", "depends_on": [7]}], + ) + with pytest.raises(vocab.VocabularyError, match="depends on 7"): + vocab.parse_vocabulary(doc) + + +def test_ceiling_verdict_requires_its_evidence(): + doc = _doc(verdict=vocab.VERDICT_CEILING, ceiling={}) + with pytest.raises(vocab.VocabularyError, match="ceiling"): + vocab.parse_vocabulary(doc) + + +def test_ceiling_with_evidence_parses(): + parsed = vocab.parse_vocabulary( + _doc( + verdict=vocab.VERDICT_CEILING, + ceiling={ + "step": "exotic_norm", + "searched": ["inventory", "catalog"], + "why": "none", + "exhausted": { + 1: "no torch.ops.trtllm candidate", + 2: "nothing under modules/** computes it", + 3: "not expressible in torch ops without a scan", + 4: "a Triton kernel needs a spec nobody has", + }, + }, + ) + ) + assert parsed.is_ceiling + assert parsed.ceiling is not None and parsed.ceiling["step"] == "exotic_norm" + + +def test_a_ceiling_that_ignores_the_writing_tiers_is_rejected(): + """The expensive mistake: reporting "nobody wrote it yet" as "it cannot exist".""" + doc = _doc( + verdict=vocab.VERDICT_CEILING, + ceiling={ + "step": "exotic_norm", + "searched": ["inventory"], + "exhausted": {1: "no trtllm op", 2: "nothing in modules/**"}, + }, + ) + with pytest.raises(vocab.VocabularyError, match=r"tier\(s\) \[3, 4\]"): + vocab.parse_vocabulary(doc) + + +def test_ceiling_block_without_a_ceiling_verdict_is_rejected(): + doc = _doc(ceiling={"step": "x"}) + with pytest.raises(vocab.VocabularyError, match="must agree"): + vocab.parse_vocabulary(doc) + + +def test_assemblable_with_a_queue_is_contradictory(): + doc = _doc(queue=[{"n": 1, "mode": vocab.MODE_ONBOARD, "target": "a"}]) + with pytest.raises(vocab.VocabularyError, match="assemblable_after_additions"): + vocab.parse_vocabulary(doc) + + +def test_after_additions_with_an_empty_queue_is_contradictory(): + doc = _doc(verdict=vocab.VERDICT_AFTER_ADDITIONS, queue=[]) + with pytest.raises(vocab.VocabularyError, match="names none"): + vocab.parse_vocabulary(doc) + + +def test_every_problem_is_batched(): + doc = { + "verdict": "maybe", + "steps": [{"step": "", "state": "NOPE"}], + "queue": [{"n": 1, "mode": "guess", "target": ""}], + } + with pytest.raises(vocab.VocabularyError) as excinfo: + vocab.parse_vocabulary(doc) + message = str(excinfo.value) + assert "verdict" in message and "state" in message and "mode" in message + + +def test_queue_item_maps_mode_to_skill(): + onboard = vocab.QueueItem(n=1, mode=vocab.MODE_ONBOARD, target="a") + verify = vocab.QueueItem(n=2, mode=vocab.MODE_VERIFY, target="b", surface="s") + assert onboard.skill == "staircase-kernel-onboard" + assert verify.skill == "staircase-kernel-verify" + + +def test_independent_items_land_in_one_batch(): + """Everything without a dependency runs at once — this is the wall-clock win.""" + items = [ + vocab.QueueItem(n=1, mode=vocab.MODE_ONBOARD, target="a"), + vocab.QueueItem(n=2, mode=vocab.MODE_ONBOARD, target="b"), + vocab.QueueItem(n=3, mode=vocab.MODE_ONBOARD, target="c"), + ] + batches = vocab.dependency_batches(items) + assert len(batches) == 1 + assert [i.n for i in batches[0]] == [1, 2, 3] + + +def test_dependencies_layer_into_ordered_batches(): + items = [ + vocab.QueueItem(n=1, mode=vocab.MODE_ONBOARD, target="a"), + vocab.QueueItem(n=2, mode=vocab.MODE_VERIFY, target="b", surface="s", depends_on=(1,)), + vocab.QueueItem(n=3, mode=vocab.MODE_ONBOARD, target="c"), + vocab.QueueItem(n=4, mode=vocab.MODE_VERIFY, target="d", surface="s", depends_on=(2, 3)), + ] + batches = vocab.dependency_batches(items) + assert [[i.n for i in batch] for batch in batches] == [[1, 3], [2], [4]] + + +def test_a_dependency_cycle_is_named_not_broken_arbitrarily(): + items = [ + vocab.QueueItem(n=1, mode=vocab.MODE_ONBOARD, target="a", depends_on=(2,)), + vocab.QueueItem(n=2, mode=vocab.MODE_ONBOARD, target="b", depends_on=(1,)), + ] + with pytest.raises(vocab.VocabularyError, match="cycle"): + vocab.dependency_batches(items) + + +def test_empty_queue_yields_no_batches(): + assert vocab.dependency_batches([]) == [] + + +def test_load_from_disk_round_trips(tmp_path): + doc = _doc( + verdict=vocab.VERDICT_AFTER_ADDITIONS, + steps=[_ok_step(state=vocab.STATE_CERT)], + queue=[ + { + "n": 1, + "mode": vocab.MODE_VERIFY, + "target": "flashinfer_rmsnorm", + "surface": "head_dim=72", + } + ], + ) + path = tmp_path / "vocabulary.yaml" + path.write_text(yaml.safe_dump(doc), encoding="utf-8") + parsed = vocab.load_vocabulary(path) + assert parsed.queue[0].surface == "head_dim=72" + reparsed = vocab.parse_vocabulary(yaml.safe_load(vocab.dump_vocabulary(parsed))) + assert reparsed.queue[0].target == "flashinfer_rmsnorm" + + +def test_missing_file_names_the_producer(tmp_path): + with pytest.raises(vocab.VocabularyError, match="plan phase"): + vocab.load_vocabulary(tmp_path / "absent.yaml") + + +def test_empty_file_is_rejected(tmp_path): + path = tmp_path / "vocabulary.yaml" + path.write_text(" \n", encoding="utf-8") + with pytest.raises(vocab.VocabularyError, match="empty"): + vocab.load_vocabulary(path) + + +# ------------------------------------------------------------------ unit + + +_UNITS = ("core", "dspark", "multimodal", "kv_quant") + + +def _queue_doc(*rows): + return _doc(verdict=vocab.VERDICT_AFTER_ADDITIONS, queue=list(rows)) + + +def test_a_queue_row_may_name_the_unit_that_first_needs_it(): + parsed = vocab.parse_vocabulary( + _queue_doc( + {"n": 1, "mode": vocab.MODE_ONBOARD, "target": "moe/drafter_moe", "unit": "dspark"}, + {"n": 2, "mode": vocab.MODE_ONBOARD, "target": "norm/core_only"}, + ), + units=_UNITS, + ) + assert [item.unit for item in parsed.queue] == ["dspark", vocab.UNIT_CORE] + + +def test_a_unit_the_run_does_not_build_is_rejected(): + """A misspelt tag could only be guessed at, and guessing late lets a core failure through.""" + doc = _queue_doc({"n": 1, "mode": vocab.MODE_ONBOARD, "target": "x", "unit": "dspakr"}) + with pytest.raises(vocab.VocabularyError, match="not one of this run's assembly units"): + vocab.parse_vocabulary(doc, units=_UNITS) + + +@pytest.mark.parametrize("bad", ["", " ", 3, ["dspark"]]) +def test_a_unit_must_be_a_name(bad): + doc = _queue_doc({"n": 1, "mode": vocab.MODE_ONBOARD, "target": "x", "unit": bad}) + with pytest.raises(vocab.VocabularyError, match=r"queue\[0\]\.unit"): + vocab.parse_vocabulary(doc) + + +def test_without_the_run_units_only_the_shape_is_checked(): + """The validator cannot see task.yaml on its own; the orchestrator passes the units.""" + doc = _queue_doc({"n": 1, "mode": vocab.MODE_ONBOARD, "target": "x", "unit": "anything"}) + assert vocab.parse_vocabulary(doc).queue[0].unit == "anything" + + +def test_the_unit_survives_a_round_trip_through_run_state(): + """A checkpointed row also carries the orchestrator's bookkeeping, which is not the item.""" + item = vocab.QueueItem(n=4, mode=vocab.MODE_ONBOARD, target="x", depends_on=(1,), unit="dspark") + row = { + **item.as_dict(), + "status": "deferred", + "deferred_to": "dspark", + "deferred_why": "smith decided BLOCKED", + } + assert vocab.QueueItem.from_row(row) == item + + +def test_a_row_checkpointed_before_units_existed_reads_as_core(): + row = {"n": 1, "mode": vocab.MODE_ONBOARD, "target": "x", "depends_on": [], "status": "failed"} + assert vocab.QueueItem.from_row(row).unit == vocab.UNIT_CORE + + +def _item(n, unit="core", depends_on=()): + return vocab.QueueItem( + n=n, mode=vocab.MODE_ONBOARD, target=f"t{n}", unit=unit, depends_on=depends_on + ) + + +def test_a_row_is_needed_by_its_own_unit(): + needed = vocab.needed_by([_item(1), _item(2, "dspark")], _UNITS) + assert needed == {1: "core", 2: "dspark"} + + +def test_a_dependency_of_a_core_row_is_core_transitively(): + """Chain: core row 3 -> row 2 (kv_quant) -> row 1 (multimodal). Core needs all three.""" + queue = [ + _item(1, "multimodal"), + _item(2, "kv_quant", depends_on=(1,)), + _item(3, "core", depends_on=(2,)), + ] + assert vocab.needed_by(queue, _UNITS) == {1: "core", 2: "core", 3: "core"} + + +def test_a_dependency_is_needed_by_its_earliest_dependent(): + """Tagged kv_quant, but a dspark row waits on it: it is owed before dspark, not kv_quant.""" + queue = [_item(1, "kv_quant"), _item(2, "dspark", depends_on=(1,)), _item(3, "kv_quant")] + assert vocab.needed_by(queue, _UNITS) == {1: "dspark", 2: "dspark", 3: "kv_quant"} + + +def test_a_dependent_is_not_pulled_earlier_by_what_it_waits_on(): + queue = [_item(1, "core"), _item(2, "kv_quant", depends_on=(1,))] + assert vocab.needed_by(queue, _UNITS) == {1: "core", 2: "kv_quant"} + + +def test_a_unit_outside_the_run_is_needed_by_core(): + """The blocking direction: an unrecognised tag never lets a failure slip past assembly.""" + assert vocab.needed_by([_item(1, "eagle3")], _UNITS) == {1: "core"} diff --git a/agent-flow/tests/workflows/staircase_onboard/test_witness.py b/agent-flow/tests/workflows/staircase_onboard/test_witness.py new file mode 100644 index 000000000000..2aa9dd49d07c --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_witness.py @@ -0,0 +1,1140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The two-leg gates for features that are output-identical by construction. + +Written against the real tools' output formats, not against invented strings. +That rule is the whole reason this file exists in the shape it does: the +acceptance gate spent its entire life parsing `Spec-dec stats: AL`, a string +that appears nowhere in this tree outside its own tests' fixtures, and three +passing unit tests never noticed. +""" + +from __future__ import annotations + +from pathlib import Path + +from agent_flow.workflows.staircase_onboard.environment import RunEnvironment +from agent_flow.workflows.staircase_onboard.gates import ShellResult +from agent_flow.workflows.staircase_onboard.witness import ( + DEVICE_MEMORY_MARKER, + LEG_B_ARGV_ENV, + LEG_B_DRAFT_LEN_ENV, + LEG_B_LEGS_ENV_ENV, + LEG_B_MAX_BATCH_SIZE_ENV, + LEG_B_MAX_TOKENS_ENV, + LEG_B_MEMBER_ENV, + LEG_B_PROMPTS_FILE_ENV, + LEG_B_TIMEOUT_ENV, + PREFIX_OBSERVABLE, + WITNESS_COUNTER_ENV, + WITNESS_COUNTER_MARKER, + WITNESS_OUTPUT_ENV, + compare_continuations, + filler_words, + first_difference, + parse_device_memory_mib, + parse_generated_texts, + parse_witness_counters, + parse_witness_texts, + parse_witness_token_ids, + run_equivalence_witness, + run_kv_quant_witness, + run_replay_witness, + run_residency_witness, + with_max_batch_size, + witness_command, +) + +_TASK = { + "trtllm_repo_path": "/repo", + "checkpoint_path": "/ckpt", + "parallel": "tp1", + "features": [], + "runtime": {"cuda_graph": True, "kv_cache_reuse": True, "extra_flags": []}, +} + + +def _env(): + return RunEnvironment(kind="local", python="/opt/venv/bin/python", world_size=1) + + +#: Verbatim from a real `quickstart_advanced.py` run in this project's logs. +def _quickstart(text_a: str, text_b: str, mib: float | None = None) -> str: + out = ( + f"[0] Accept rate: 1.00 (tokens=32, iterations=32)\n" + f"[0] Prompt: 'The capital of France is'..., Generated text: {text_a!r}...\n" + f"[1] Accept rate: 1.00 (tokens=32, iterations=32)\n" + f"[1] Prompt: 'In one sentence, photosynthesis is'..., Generated text: {text_b!r}...\n" + ) + if mib is not None: + out += f"{DEVICE_MEMORY_MARKER}: {mib}\n" + return out + + +# ------------------------------------------------------------------- parsing + + +def test_continuations_are_read_in_order(): + texts = parse_generated_texts(_quickstart(" Paris.", " the process")) + assert texts == ("' Paris.'", "' the process'") + + +def test_no_continuations_is_empty_not_an_error(): + assert parse_generated_texts("nothing here") == () + + +def test_device_memory_is_read_from_a_marker_not_from_nvidia_smis_format(): + """The workload decides WHEN to sample, because the gate cannot know. + + The number that matters is taken after the engine reports ready, and a + gate sampling from outside has no way to tell when that was. + """ + assert parse_device_memory_mib(f"{DEVICE_MEMORY_MARKER}: 183542\n") == 183542.0 + assert parse_device_memory_mib(f"{DEVICE_MEMORY_MARKER}=1024.5") == 1024.5 + assert parse_device_memory_mib("no marker") is None + + +def test_the_last_reading_wins(): + """A two-phase engine start prints one per phase; the real one is last.""" + both = f"{DEVICE_MEMORY_MARKER}: 100\n...\n{DEVICE_MEMORY_MARKER}: 200\n" + assert parse_device_memory_mib(both) == 200.0 + + +# ------------------------------------------------------------------- command + + +def test_the_witness_runs_quickstart_not_the_eval_harness(): + """`trtllm-eval` reports neither device memory nor the continuations. + + It also costs nine minutes a leg where this costs four. + """ + argv = witness_command(_TASK, _env(), unit="core") + assert "quickstart_advanced.py" in argv + assert "trtllm-eval" not in argv + + +def test_the_witness_carries_the_units_runtime_block(): + task = dict(_TASK) | { + "runtime": {"default": {"kv_cache_reuse": True}, "mm": {"kv_cache_reuse": False}} + } + assert "--disable_kv_cache_reuse" in witness_command(task, _env(), unit="mm") + assert "--disable_kv_cache_reuse" not in witness_command(task, _env(), unit="core") + + +# ----------------------------------------------------------------- residency + + +def _residency(on_mib, off_mib, on_text=(" A", " B"), off_text=(" A", " B"), expected=100.0): + calls: list[str] = [] + + def runner(command, env=None, timeout=0): + calls.append(command) + first = len(calls) == 1 + mib, text = (on_mib, on_text) if first else (off_mib, off_text) + return ShellResult(returncode=0, stdout=_quickstart(text[0], text[1], mib)) + + return run_residency_witness( + _env(), + task_data=_TASK, + expected_mib=expected, + contrast_env={"TARGET_OFFLOAD": "0"}, + runner=runner, + ), calls + + +def test_a_real_offload_shows_the_predicted_memory_delta(): + result, calls = _residency(on_mib=1000, off_mib=1200, expected=100.0) + assert result.passed + assert result.measured == 200.0 + assert len(calls) == 2 + + +def test_an_offload_that_did_nothing_fails_even_though_the_text_is_right(): + """The failure this gate exists for. + + A target that quietly loaded the table onto the device anyway produces + identical output and passes boot and accuracy. Only the memory delta + distinguishes it. + """ + result, _ = _residency(on_mib=1200, off_mib=1205, expected=100.0) + assert not result.passed + assert "delta 5 MiB < predicted 100 MiB" in result.detail + + +def test_legs_that_disagree_on_text_fail_regardless_of_the_memory(): + """Offloading is output-identical by construction. + + So a disagreement means the gather is wrong, and the accuracy tolerance + would hide it. + """ + result, _ = _residency(on_mib=1000, off_mib=1200, on_text=(" A", " B"), off_text=(" A", " C")) + assert not result.passed + assert "identical text" in result.detail + + +def test_a_missing_memory_reading_is_not_a_pass(): + result, _ = _residency(on_mib=None, off_mib=1200) + assert not result.passed + assert "is not a pass" in result.detail + + +def test_the_control_leg_runs_with_the_offload_switched_off(): + """The changed variable has to be the feature, and only the feature.""" + seen: list[dict] = [] + + def runner(command, env=None, timeout=0): + seen.append({"cmd": command}) + return ShellResult(returncode=0, stdout=_quickstart(" A", " B", 1000 + 200 * len(seen))) + + run_residency_witness( + _env(), + task_data=_TASK, + expected_mib=1.0, + contrast_env={"TARGET_OFFLOAD": "0"}, + runner=runner, + ) + assert "TARGET_OFFLOAD=0" not in seen[0]["cmd"] + assert "TARGET_OFFLOAD=0" in seen[1]["cmd"] + + +# --------------------------------------------------------------- equivalence + + +def _equivalence(agg_text, dis_text, dis_ok=True): + calls: list[str] = [] + + def runner(command, env=None, timeout=0): + calls.append(command) + first = len(calls) == 1 + if not first and not dis_ok: + return ShellResult(returncode=1, stdout="", stderr="hang") + text = agg_text if first else dis_text + return ShellResult(returncode=0, stdout=_quickstart(*text)) + + return run_equivalence_witness( + _env(), + task_data=_TASK, + contrast_env={"TRTLLM_DISAGG": "1"}, + runner=runner, + ), calls + + +def test_token_identical_output_across_the_hop_passes(): + result, calls = _equivalence((" Paris.", " the process"), (" Paris.", " the process")) + assert result.passed + assert "token-identical" in result.detail + assert len(calls) == 2 + + +def test_fluent_but_different_output_fails(): + """Not 'similar' — identical. + + Missing recurrent state is recomputed from zero and the model recovers + into text that is fluent and wrong, which a short benchmark under a + five-point tolerance does not catch. + """ + result, _ = _equivalence((" Paris.", " the process"), (" Paris.", " a process")) + assert not result.passed + assert "fluent, wrong text" in result.detail + + +def test_a_disagg_leg_that_fails_names_the_usual_causes(): + """A disagg launch fails by hanging more often than by raising.""" + result, _ = _equivalence((" A", " B"), (" A", " B"), dis_ok=False) + assert not result.passed + assert "cache transceiver" in result.detail + + +def test_the_second_leg_runs_disaggregated_and_the_first_does_not(): + result, calls = _equivalence((" A", " B"), (" A", " B")) + assert "TRTLLM_DISAGG=1" not in calls[0] + assert "TRTLLM_DISAGG=1" in calls[1] + + +def test_both_legs_ask_the_target_for_its_runtime_value_dump(): + _, calls = _equivalence((" A", " B"), (" A", " B")) + assert len(calls) == 2 + assert all("TRTLLM_MODELING_V2_VALUE_DUMP=1" in call for call in calls) + + +def test_the_disagg_leg_can_be_a_different_command_entirely(): + """Disaggregation is not an environment variable over one process. + + It is a context server, a generation server and a client. An env overlay + alone would run leg B **aggregated**, and the witness would compare a run + to itself and pass trivially — the worst possible outcome for a gate whose + entire job is to catch a silent equality. + """ + calls: list[str] = [] + + def runner(command, env=None, timeout=0): + calls.append(command) + return ShellResult(returncode=0, stdout=_quickstart(" Paris.", " the process")) + + result = run_equivalence_witness( + _env(), + task_data=_TASK, + contrast_env={"TRTLLM_DISAGG": "1"}, + contrast_command="/run/launch_disagg.sh --prompts boot", + runner=runner, + ) + assert result.passed + assert "quickstart_advanced.py" in calls[0] + assert "launch_disagg.sh" in calls[1] + # Leg A's argv is HANDED to leg B in its environment; it is not what runs. + assert "quickstart_advanced.py" not in calls[1].split("bash -c", 1)[1] + + +def test_without_a_contrast_command_both_legs_are_the_same_argv(): + """Recorded so the trivial-pass hazard above is visible in the tests. + + This is the correct shape for the residency witness, where the feature + really is one switch, and the wrong one for disagg. + """ + calls: list[str] = [] + + def runner(command, env=None, timeout=0): + calls.append(command) + return ShellResult(returncode=0, stdout=_quickstart(" A", " B")) + + run_equivalence_witness(_env(), task_data=_TASK, contrast_env={"X": "1"}, runner=runner) + assert calls[0].split("python")[-1] in calls[1] + + +# ------------------------------------------- equivalence: leg B as a launcher + + +def _whole(texts, ids=None) -> str: + """Quickstart's output under TRTLLM_WITNESS_FULL_OUTPUT=1 (see the example script).""" + out = "" + for i, text in enumerate(texts): + out += f"[{i}] Prompt: 'p'..., Generated text: {text[:200]!r}...\n" + if ids is not None: + out += f"[{i}] Witness token ids: {list(ids[i])}\n" + out += f"[{i}] Witness text: {text!r}\n" + return out + + +def _run_equivalence(tmp_path, agg_out, dis_out, **kwargs): + calls: list[str] = [] + + def runner(command, env=None, timeout=0): + calls.append(command) + return ShellResult(returncode=0, stdout=agg_out if len(calls) == 1 else dis_out) + + kwargs.setdefault("contrast_command", "/onboard/scripts/disagg_leg_b.py") + kwargs.setdefault("workdir", tmp_path / "witness") + kwargs.setdefault("task_data", _TASK) + result = run_equivalence_witness( + kwargs.pop("environment", _env()), + contrast_env={"TRTLLM_DISAGG": "1"}, + runner=runner, + **kwargs, + ) + return result, calls + + +def _assignment(command: str, name: str) -> str: + import shlex + + for word in shlex.split(command): + if word.startswith(f"{name}="): + return word.split("=", 1)[1] + raise AssertionError(f"{name} not exported in {command}") + + +def test_the_contrast_command_runs_as_one_process_not_as_a_rank(tmp_path): + """Leg B starts engines; it is not one. F021b. + + As engine work at world size 4 it went out behind `mpirun -n 4 + trtllm-llmapi-launch`: three idle workers held ranks 1-3, and the command + inherited the wrapper's spawn-proxy address, which makes the engines it + starts dial the outer proxy and hang. + """ + four = RunEnvironment( + kind="local", python="/opt/venv/bin/python", world_size=4, devices=(0, 1, 2, 3) + ) + out = _whole([" A", " B"]) + result, calls = _run_equivalence(tmp_path, out, out, environment=four) + assert result.passed + assert "trtllm-llmapi-launch" in calls[0] and "mpirun" in calls[0] + leg_b = calls[1] + assert "trtllm-llmapi-launch" not in leg_b and "mpirun" not in leg_b + assert "CUDA_VISIBLE_DEVICES=0,1,2,3" in leg_b + for name in ("TLLM_SPAWN_PROXY_PROCESS", "TLLM_SPAWN_PROXY_PROCESS_IPC_ADDR"): + assert f"-u {name} " in leg_b + assert "/opt/venv/bin/python /onboard/scripts/disagg_leg_b.py" in leg_b + + +def test_a_contrast_command_that_is_not_python_runs_as_written(tmp_path): + out = _whole([" A"]) + _, calls = _run_equivalence( + tmp_path, out, out, contrast_command="/onboard/scripts/leg_b.sh --fast" + ) + assert "/onboard/scripts/leg_b.sh --fast" in calls[1] + assert "/opt/venv/bin/python /onboard/scripts/leg_b.sh" not in calls[1] + + +def test_leg_b_is_handed_what_leg_a_resolved(tmp_path): + """A restated workload desynchronises at the next edit to either side. F021c.""" + import json + + out = _whole([" A", " B"]) + _, calls = _run_equivalence(tmp_path, out, out, member="tep4", max_tokens=128, timeout=2700) + leg_b = calls[1] + prompts = json.loads(Path(_assignment(leg_b, LEG_B_PROMPTS_FILE_ENV)).read_text())["prompts"] + for prompt in prompts: + assert f'"{prompt}"' in calls[0] # the very prompts leg A ran + assert _assignment(leg_b, LEG_B_MAX_TOKENS_ENV) == "128" + assert _assignment(leg_b, LEG_B_MEMBER_ENV) == "tep4" + assert _assignment(leg_b, LEG_B_TIMEOUT_ENV) == "2700" + argv = _assignment(leg_b, LEG_B_ARGV_ENV) + assert "quickstart_advanced.py" in argv and "--max_tokens=128" in argv + assert "--prompt" not in argv # the prompts travel in the file + # Nothing speculates here, so there is no depth to hand over. + assert f"{LEG_B_DRAFT_LEN_ENV}=" not in leg_b + # Leg A is not the one being handed anything. + assert LEG_B_PROMPTS_FILE_ENV not in calls[0] + + +def test_leg_b_gets_the_draft_depth_when_leg_a_speculates(tmp_path): + import yaml + + options = tmp_path / "dspark.yaml" + options.write_text( + yaml.safe_dump({"speculative_config": {"decoding_type": "DSpark", "max_draft_len": 5}}) + ) + task = dict(_TASK) | {"accuracy": {"llm_api_options": {"disagg": str(options)}}} + out = _whole([" A"]) + _, calls = _run_equivalence( + tmp_path, out, out, task_data=task, unit="disagg", features=["dspark"] + ) + assert _assignment(calls[1], LEG_B_DRAFT_LEN_ENV) == "5" + + +def test_declared_prompt_lengths_reach_both_legs_beside_the_boot_prompts(tmp_path): + """The plan's disagg gate wants lengths hand-written prompts do not land on. F021d. + + Over one prefill chunk, odd (a ratio-2 compressor's pending token), and + one past a page edge. A leg-B-only prompt cannot be added: the witness + compares continuation lists and a count mismatch fails. + """ + import json + + from agent_flow.workflows.staircase_onboard.gates import BootSpec + + spec = BootSpec(prompts=("What is the capital of France?",)) + out = _whole([" A", " B", " C"]) + _, calls = _run_equivalence(tmp_path, out, out, spec=spec, prompt_tokens=[5, 7]) + expected = ["What is the capital of France?", filler_words(5), filler_words(7)] + for prompt in expected: + assert f'"{prompt}"' in calls[0] + handed = json.loads(Path(_assignment(calls[1], LEG_B_PROMPTS_FILE_ENV)).read_text()) + assert handed["prompts"] == expected + + +def test_both_legs_are_asked_for_whole_continuations(tmp_path): + out = _whole([" A"]) + _, calls = _run_equivalence(tmp_path, out, out) + assert all(f"{WITNESS_OUTPUT_ENV}=1" in call for call in calls) + + +def test_a_difference_past_the_printed_prefix_fails(tmp_path): + """The 200-character line is under a quarter of a 256-token continuation. F021e.""" + same = "x" * 200 + agg, dis = _whole([same + " tail one"]), _whole([same + " tail two"]) + assert parse_generated_texts(agg) == parse_generated_texts(dis) + result, _ = _run_equivalence(tmp_path, agg, dis) + assert not result.passed + assert "full text" in result.detail + assert "first differs at character 206" in result.detail + + +def test_token_ids_win_over_text_when_both_legs_print_them(tmp_path): + """Two token sequences can detokenize to one string; ids are the claim.""" + agg = _whole([" A B"], ids=[(1, 2, 3)]) + dis = _whole([" A B"], ids=[(1, 2, 4)]) + result, _ = _run_equivalence(tmp_path, agg, dis) + assert not result.passed + assert "token ids" in result.detail + assert "first differs at token 2" in result.detail + assert result.measured == 2 + + +def test_a_leg_without_whole_continuations_falls_back_to_the_prefix_and_says_so(tmp_path): + """A contrast command written before the full-output lines still compares.""" + result, _ = _run_equivalence(tmp_path, _whole([" A", " B"]), _quickstart(" A", " B")) + assert result.passed + assert PREFIX_OBSERVABLE in result.detail + assert "not compared" in result.detail + + +def test_a_failing_verdict_says_where_each_continuation_first_differs(tmp_path): + """Lost state shows at the first decode step; a late flip is numerics. F021f.""" + agg = _whole([" one two three", " same"], ids=[(1, 2, 3, 4, 5), (7, 7)]) + dis = _whole([" one two four", " same"], ids=[(1, 2, 3, 9, 5), (7, 7)]) + result, _ = _run_equivalence(tmp_path, agg, dis) + assert not result.passed + assert "continuation(s) [0]" in result.detail + assert "#0 first differs at token 3" in result.detail + assert "logit" in result.detail and "fluent, wrong text" in result.detail + assert "around token 3" in result.stdout_tail + + +def test_a_count_mismatch_is_named_as_one(tmp_path): + result, _ = _run_equivalence(tmp_path, _whole([" A", " B"]), _whole([" A"])) + assert not result.passed + assert "printed 2 continuation(s), disagg 1" in result.detail + + +def test_whole_continuations_are_parsed_verbatim(): + out = ( + "[0] Witness token ids: [128798, 17, 13]\n" + "0: [1] Witness token ids: []\n" + "[0] Witness text: 'it\\'s \"72\"\\nnext'\n" + ) + assert parse_witness_token_ids(out) == ((128798, 17, 13), ()) + assert parse_witness_texts(out) == ('it\'s "72"\nnext',) + + +def test_first_difference_of_a_truncated_continuation_is_where_it_stops(): + assert first_difference((1, 2, 3), (1, 2, 3, 4)) == 3 + assert first_difference("abc", "abx") == 2 + assert compare_continuations(_whole(["x"]), _whole(["x"])).identical + + +# ------------------------------------------------------ the draft depth, again + + +def test_a_witness_asks_for_the_depth_the_unit_was_built_for(tmp_path): + """The boot gate read the depth from the options YAML; the witness did not. + + Harmless for MTP at depth 1. Fatal for DSpark: `TorchLlmArgs` rejects a + DSpark config whose `block_size` differs from `max_draft_len`, so every + witness after a `dspark` unit would die at startup and report a leg that + "did not complete" -- a message that reads as the unit's own defect. + """ + import yaml + + options = tmp_path / "dspark.yaml" + options.write_text( + yaml.safe_dump({"speculative_config": {"decoding_type": "DSpark", "max_draft_len": 5}}) + ) + task = dict(_TASK) | {"accuracy": {"llm_api_options": {"dspark": str(options)}}} + argv = witness_command(task, _env(), unit="dspark", features=["dspark"]) + assert "--spec_decode_algo=DSPARK" in argv + assert "--spec_decode_max_draft_len=5" in argv + + +# -------------------------------------------------------- the counter contract + + +def _counter(**fields) -> str: + return f"{WITNESS_COUNTER_MARKER} " + " ".join(f"{k}={v}" for k, v in fields.items()) + + +def test_counter_lines_are_read_whatever_precedes_them(): + """A logger prefix or a rank tag in front of the marker changes nothing.""" + out = ( + "[TensorRT-LLM][INFO] " + _counter(forwarded_tokens=4100, decoder_tokens=128) + "\n" + "[rank1] " + _counter(kv_bytes_per_token=890) + "\n" + "an ordinary line\n" + ) + assert parse_witness_counters(out) == [ + {"forwarded_tokens": 4100, "decoder_tokens": 128}, + {"kv_bytes_per_token": 890}, + ] + + +def test_the_filler_is_one_short_word_per_token_at_least(): + """The long prompt's length is DECLARED, so the workload is built to it.""" + text = filler_words(4096) + assert len(text.split()) == 4096 + assert max(len(word) for word in text.split()) <= 5 + # 16384 words must stay under Linux's 128 KiB cap on one argv string. + assert len(filler_words(16384).encode()) < 128 * 1024 + + +# ------------------------------------------------------------ decoder replay + +_REPLAY_OFF = "REPLAY_OFF_FOR_TEST=1" + + +def _leg_runner(on_out: str, off_out: str, *, marker: str = _REPLAY_OFF, fail_leg=None): + """A runner whose output depends on whether the contrast env is in the command.""" + calls: list[str] = [] + + def runner(command, env=None, timeout=0): + calls.append(command) + off = marker in command + if fail_leg == ("OFF" if off else "ON"): + return ShellResult(returncode=1, stdout="", stderr="boom") + return ShellResult(returncode=0, stdout=(off_out if off else on_out)) + + return runner, calls + + +def _lines(*rows: dict) -> str: + return "".join(_counter(**row) + "\n" for row in rows) + _quickstart(" A", " B") + + +def _decoder(on_rows, off_rows, **kw): + runner, calls = _leg_runner(_lines(*on_rows), _lines(*off_rows), **kw) + result = run_replay_witness( + _env(), + task_data=_TASK, + feature="decoder_replay", + contrast_env={"REPLAY_OFF_FOR_TEST": "1"}, + window=128, + prompt_tokens=4096, + runner=runner, + unit="decoder_replay", + ) + return result, calls + + +# Quickstart runs every prompt twice -- a batched pass, then a streamed one -- +# so a real log carries a cold request and a warm (prefix-hit) one. +_COLD = {"forwarded_tokens": 4100} +_WARM = {"forwarded_tokens": 36} + + +def test_a_bounded_decoder_passes(): + result, calls = _decoder( + on_rows=[_COLD | {"decoder_tokens": 128}, _WARM | {"decoder_tokens": 36}], + off_rows=[_COLD | {"decoder_tokens": 4100}, _WARM | {"decoder_tokens": 36}], + ) + assert result.passed, result.detail + assert result.name == "replay[decoder_replay]" + assert len(calls) == 2 + + +def test_a_target_that_implements_nothing_fails_the_on_leg(): + """The case this witness exists for: full prefill, accuracy green.""" + result, _ = _decoder( + on_rows=[_COLD | {"decoder_tokens": 4100}], + off_rows=[_COLD | {"decoder_tokens": 4100}], + ) + assert not result.passed + assert "not bounded" in result.detail + + +def test_a_counter_that_prints_a_formula_is_caught_by_the_off_leg(): + """`decoder_tokens=min(prompt, window)` passes ON; OFF is why there are two legs.""" + result, _ = _decoder( + on_rows=[_COLD | {"decoder_tokens": 128}], + off_rows=[_COLD | {"decoder_tokens": 128}], + ) + assert not result.passed + assert "did not turn the feature off" in result.detail + + +def test_a_prompt_that_never_ran_cold_at_length_proves_nothing(): + """Under the window, bounded and unbounded decoders do the same work.""" + short = {"forwarded_tokens": 100} + result, _ = _decoder( + on_rows=[short | {"decoder_tokens": 100}], + off_rows=[short | {"decoder_tokens": 100}], + ) + assert not result.passed + assert "prompt_tokens=4096" in result.detail + + +def test_no_counter_line_is_a_named_failure_not_a_pass(): + runner, _ = _leg_runner(_quickstart(" A", " B"), _quickstart(" A", " B")) + result = run_replay_witness( + _env(), + task_data=_TASK, + feature="decoder_replay", + contrast_env={"REPLAY_OFF_FOR_TEST": "1"}, + window=128, + prompt_tokens=4096, + runner=runner, + ) + assert not result.passed + assert "printed no witness counter" in result.detail + + +def test_a_counter_line_missing_a_field_names_it(): + result, _ = _decoder(on_rows=[{"forwarded_tokens": 4100}], off_rows=[_COLD]) + assert not result.passed + assert "decoder_tokens" in result.detail + + +def test_a_storage_line_under_the_same_marker_is_not_a_replay_line(): + """One target may implement kv_quant and replay; each witness reads its own.""" + result, _ = _decoder( + on_rows=[{"kv_bytes_per_token": 890}, _COLD | {"decoder_tokens": 128}], + off_rows=[{"kv_bytes_per_token": 890}, _COLD | {"decoder_tokens": 4100}], + ) + assert result.passed, result.detail + + +def test_both_legs_count_and_only_the_second_is_switched_off(): + _, calls = _decoder( + on_rows=[_COLD | {"decoder_tokens": 128}], + off_rows=[_COLD | {"decoder_tokens": 4100}], + ) + assert all(f"{WITNESS_COUNTER_ENV}=1" in call for call in calls) + assert all("TRTLLM_MODELING_V2=require" in call for call in calls) + assert _REPLAY_OFF not in calls[0] and _REPLAY_OFF in calls[1] + # The workload is the declared one, not the boot spec's two short prompts. + assert "capital of France" not in calls[0] + assert calls[0].count(" one ") >= 400 + + +def test_a_leg_that_dies_is_named(): + runner, _ = _leg_runner(_lines(_COLD | {"decoder_tokens": 128}), "", fail_leg="OFF") + result = run_replay_witness( + _env(), + task_data=_TASK, + feature="decoder_replay", + contrast_env={"REPLAY_OFF_FOR_TEST": "1"}, + window=128, + prompt_tokens=4096, + runner=runner, + ) + assert not result.passed + assert "replay-OFF leg did not complete" in result.detail + + +# ------------------------------------------------------------ encoder replay + +_TASK_REUSE = dict(_TASK) | {"runtime": {"kv_cache_reuse": True, "extra_flags": []}} +_COLD_E = { + "prompt_tokens": 4105, + "reused_tokens": 0, + "replayed_tokens": 0, + "forwarded_tokens": 4105, +} + + +def _encoder(on_rows, off_rows, *, task=_TASK_REUSE, alignment=1): + runner, calls = _leg_runner(_lines(*on_rows), _lines(*off_rows)) + result = run_replay_witness( + _env(), + task_data=task, + feature="encoder_replay", + contrast_env={"REPLAY_OFF_FOR_TEST": "1"}, + window=128, + prefix_tokens=4096, + reuse_alignment=alignment, + runner=runner, + unit="encoder_replay", + ) + return result, calls + + +def _hit(reused: int, replayed: int, forwarded: int | None = None) -> dict: + prompt = 4105 + if forwarded is None: + forwarded = prompt - reused + replayed + return { + "prompt_tokens": prompt, + "reused_tokens": reused, + "replayed_tokens": replayed, + "forwarded_tokens": forwarded, + } + + +def test_a_prefix_hit_that_replays_the_window_passes(): + result, calls = _encoder( + on_rows=[_COLD_E, _hit(4096, 128)], + off_rows=[_COLD_E, _hit(4096, 0)], + ) + assert result.passed, result.detail + assert len(calls) == 2 + + +def test_the_encoder_workload_is_two_prompts_sharing_the_prefix(): + _, calls = _encoder(on_rows=[_COLD_E, _hit(4096, 128)], off_rows=[_COLD_E]) + assert "What comes next" in calls[0] and "Say it backwards" in calls[0] + + +def test_a_hit_that_recomputes_the_whole_prefix_breaks_the_accounting(): + """The shortcut "SWA KV is missing" invites: report the hit, forward everything.""" + result, _ = _encoder( + on_rows=[_COLD_E, _hit(4096, 128, forwarded=4105)], + off_rows=[_COLD_E, _hit(4096, 0)], + ) + assert not result.passed + assert "prompt - reused + replayed" in result.detail + + +def test_no_prefix_hit_means_nothing_was_witnessed(): + """No reuse at all satisfies every identity trivially; it must not pass.""" + result, _ = _encoder(on_rows=[_COLD_E, _COLD_E], off_rows=[_COLD_E, _COLD_E]) + assert not result.passed + assert "no request reused more than window=128" in result.detail + + +def test_a_replay_of_the_wrong_length_fails(): + result, _ = _encoder( + on_rows=[_COLD_E, _hit(4096, 64)], + off_rows=[_COLD_E, _hit(4096, 0)], + ) + assert not result.passed + assert "min(window=128, reused)" in result.detail + + +def test_the_off_leg_must_replay_nothing(): + result, _ = _encoder( + on_rows=[_COLD_E, _hit(4096, 128)], + off_rows=[_COLD_E, _hit(4096, 128)], + ) + assert not result.passed + assert "did not turn the" in result.detail + + +def test_an_odd_boundary_on_a_ratio_two_model_fails_when_declared(): + result, _ = _encoder( + on_rows=[_COLD_E, _hit(4095, 128)], + off_rows=[_COLD_E, _hit(4095, 0)], + alignment=2, + ) + assert not result.passed + assert "not a multiple of 2" in result.detail + # Undeclared, the same boundary is not this witness's business. + result, _ = _encoder( + on_rows=[_COLD_E, _hit(4095, 128)], + off_rows=[_COLD_E, _hit(4095, 0)], + ) + assert result.passed, result.detail + + +def test_block_reuse_off_refuses_before_launching_anything(): + task = dict(_TASK) | {"runtime": {"kv_cache_reuse": False, "extra_flags": []}} + result, calls = _encoder(on_rows=[_COLD_E], off_rows=[_COLD_E], task=task) + assert not result.passed + assert calls == [] + assert "kv_cache_reuse" in result.detail + + +# -------------------------------------------------------------------- kv_quant + +_BF16 = "KV_STORAGE_FOR_TEST=bf16" + + +def _kv(native_out, bf16_out, *, exact=True, tolerance=0.02): + runner, calls = _leg_runner(native_out, bf16_out, marker=_BF16) + result = run_kv_quant_witness( + _env(), + task_data=_TASK, + expected_bytes_per_token=890, + exact=exact, + contrast_env={"KV_STORAGE_FOR_TEST": "bf16"}, + bytes_tolerance=tolerance, + runner=runner, + unit="kv_quant", + ) + return result, calls + + +def _storage(bytes_per_token: int, a=" Paris.", b=" the process") -> str: + return _counter(kv_bytes_per_token=bytes_per_token) + "\n" + _quickstart(a, b) + + +def test_native_storage_at_the_declared_figure_and_identical_output_passes(): + result, calls = _kv(_storage(890), _storage(3200)) + assert result.passed, result.detail + assert len(calls) == 2 + assert all(f"{WITNESS_COUNTER_ENV}=1" in call for call in calls) + assert _BF16 not in calls[0] and _BF16 in calls[1] + + +def test_a_unit_still_on_bf16_pools_fails_on_bytes(): + """The case this witness exists for: output identical, storage unchanged.""" + result, calls = _kv(_storage(3200), _storage(3200)) + assert not result.passed + assert "3200 B/token against the declared 890" in result.detail + assert len(calls) == 1 + + +def test_padding_passes_and_a_layout_change_does_not(): + assert _kv(_storage(900), _storage(3200))[0].passed + # The indexer scale padded from 4 bytes to 16: 890 -> 920. + assert not _kv(_storage(920), _storage(3200))[0].passed + + +def test_a_contrast_that_selects_nothing_cannot_vouch_for_identity(): + result, _ = _kv(_storage(890), _storage(890)) + assert not result.passed + assert "did not select bf16 storage" in result.detail + + +def test_native_storage_that_moves_the_output_stored_other_values(): + result, _ = _kv(_storage(890), _storage(3200, b=" a process")) + assert not result.passed + assert "stored values are not the reference's" in result.detail + + +def test_no_storage_line_is_a_named_failure(): + result, _ = _kv(_quickstart(" A", " B"), _storage(3200)) + assert not result.passed + assert "printed no witness counter" in result.detail + + +def test_a_lossy_attention_path_drops_the_identity_leg_and_says_so(): + result, calls = _kv(_storage(890), _storage(3200, b=" different"), exact=False) + assert result.passed + assert len(calls) == 1 + assert "WEAKER CLAIM" in result.detail + + +# ---------------------------------------- one schedule on both legs (F036) + +_BATCHED = dict(_TASK) | { + "runtime": {"cuda_graph": True, "extra_flags": ["--max_num_tokens=2048", "--max_batch_size=8"]} +} + + +def test_a_declared_batch_limit_replaces_the_runtime_blocks(): + """Batch composition is a run-to-run source; the witness states ONE limit. + + argparse would take the last of two `--max_batch_size` flags, but leg B + reads leg A's argv for its limits, and a stale `=8` beside the `=1` is a + second answer to one question. + """ + argv = witness_command(_BATCHED, _env(), unit="kv_quant", max_batch_size=1) + assert argv.count("--max_batch_size") == 1 + assert "--max_batch_size=1" in argv + assert "--max_batch_size=8" not in argv + assert "--max_num_tokens=2048" in argv + + +def test_without_a_declared_limit_the_argv_is_unchanged(): + """A task.yaml written before the key keeps its exact commands.""" + argv = witness_command(_BATCHED, _env(), unit="kv_quant") + assert "--max_batch_size=8" in argv and argv.count("--max_batch_size") == 1 + + +def test_every_spelling_of_the_batch_flag_is_replaced(): + for flags in ( + ["--max_batch_size=8", "--x"], + ["--max_batch_size 8", "--x"], + ["--max_batch_size", "8", "--x"], + ): + assert with_max_batch_size(flags, 1) == ["--x", "--max_batch_size=1"] + assert with_max_batch_size(["--max_batch_size=8"], None) == ["--max_batch_size=8"] + + +def test_the_equivalence_witness_runs_both_legs_at_the_declared_limit(tmp_path): + out = _whole([" A", " B"]) + result, calls = _run_equivalence(tmp_path, out, out, task_data=_BATCHED, max_batch_size=1) + assert result.passed + assert "max_batch_size=1" in result.detail + assert "--max_batch_size=1" in calls[0] and "--max_batch_size=8" not in calls[0] + leg_b = calls[1] + assert _assignment(leg_b, LEG_B_MAX_BATCH_SIZE_ENV) == "1" + handed = _assignment(leg_b, LEG_B_ARGV_ENV) + assert "--max_batch_size=1" in handed and "--max_batch_size=8" not in handed + + +def test_leg_b_is_not_handed_a_limit_nobody_declared(tmp_path): + out = _whole([" A"]) + _, calls = _run_equivalence(tmp_path, out, out, task_data=_BATCHED) + assert f"{LEG_B_MAX_BATCH_SIZE_ENV}=" not in calls[1] + assert "--max_batch_size=8" in _assignment(calls[1], LEG_B_ARGV_ENV) + + +def test_an_env_only_equivalence_leg_b_gets_the_limit_in_its_argv(): + calls: list[str] = [] + + def runner(command, env=None, timeout=0): + calls.append(command) + return ShellResult(returncode=0, stdout=_quickstart(" A", " B")) + + run_equivalence_witness( + _env(), task_data=_BATCHED, contrast_env={"X": "1"}, runner=runner, max_batch_size=1 + ) + assert all("--max_batch_size=1" in call for call in calls) + + +# ---------------------------------- an environment both legs share (F039) + +_REPLAYS_OFF = { + "TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY": "0", + "TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY": "0", +} + + +def test_legs_env_reaches_leg_a_and_is_handed_to_leg_b(tmp_path): + """Leg B's engines guard the replays off, so leg A runs them off too. + + Otherwise the legs differ in the replays as well as in the hop, and token + identity no longer holds by construction. + """ + import json + + out = _whole([" A", " B"]) + result, calls = _run_equivalence(tmp_path, out, out, legs_env=_REPLAYS_OFF) + assert result.passed + agg, dis = calls + for name, value in _REPLAYS_OFF.items(): + assert _assignment(agg, name) == value + # In leg B's own environment, so a launcher that copies it passes it on... + assert _assignment(dis, name) == value + # ...and named, for one that sets engine variables of its own afterwards. + assert json.loads(_assignment(dis, LEG_B_LEGS_ENV_ENV)) == _REPLAYS_OFF + assert LEG_B_LEGS_ENV_ENV not in agg + # The disagg switch is still leg B's alone. + assert "TRTLLM_DISAGG" not in agg + assert _assignment(dis, "TRTLLM_DISAGG") == "1" + # A pass holds under these settings, so the verdict says which. + assert "both legs under TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY=0, " in result.detail + + +def test_without_legs_env_neither_leg_changes(tmp_path): + out = _whole([" A"]) + result, calls = _run_equivalence(tmp_path, out, out) + assert LEG_B_LEGS_ENV_ENV not in calls[1] + assert not any(name in call for name in _REPLAYS_OFF for call in calls) + assert "both legs under" not in result.detail + + +def test_legs_env_names_the_failing_verdict_too(tmp_path): + result, _ = _run_equivalence( + tmp_path, _whole([" A"]), _whole([" B"]), legs_env={"REPLAY_FOR_TEST": 0} + ) + assert not result.passed + assert "both legs under REPLAY_FOR_TEST=0" in result.detail + + +def test_an_env_only_equivalence_applies_legs_env_to_both_legs(): + calls: list[str] = [] + + def runner(command, env=None, timeout=0): + calls.append(command) + return ShellResult(returncode=0, stdout=_quickstart(" A", " B")) + + run_equivalence_witness( + _env(), task_data=_TASK, contrast_env={"X": "1"}, runner=runner, legs_env={"Y": "0"} + ) + assert all(_assignment(call, "Y") == "0" for call in calls) + assert "X=1" not in calls[0] and _assignment(calls[1], "X") == "1" + + +def test_the_kv_quant_legs_share_the_declared_limit_and_print_whole_continuations(): + runner, calls = _leg_runner(_storage(890), _storage(3200), marker=_BF16) + result = run_kv_quant_witness( + _env(), + task_data=_BATCHED, + expected_bytes_per_token=890, + contrast_env={"KV_STORAGE_FOR_TEST": "bf16"}, + runner=runner, + unit="kv_quant", + max_batch_size=1, + ) + assert result.passed, result.detail + assert "max_batch_size=1" in result.detail + for call in calls: + assert "--max_batch_size=1" in call and "--max_batch_size=8" not in call + assert f"{WITNESS_OUTPUT_ENV}=1" in call + + +def test_the_residency_legs_share_the_declared_limit(): + calls: list[str] = [] + + def runner(command, env=None, timeout=0): + calls.append(command) + return ShellResult(returncode=0, stdout=_quickstart(" A", " B", 1000 + 200 * len(calls))) + + result = run_residency_witness( + _env(), + task_data=_BATCHED, + expected_mib=100.0, + contrast_env={"TARGET_OFFLOAD": "0"}, + runner=runner, + max_batch_size=1, + ) + assert result.passed, result.detail + assert all("--max_batch_size=1" in call for call in calls) + assert all(f"{WITNESS_OUTPUT_ENV}=1" in call for call in calls) + + +# ------------------------------- a failed comparison shows both legs (F036) + +_AGG_TEXT = " The pass key is 41627. I found it by scanning for the only sentence." +_DIS_TEXT = " The pass key is 41627. I found it by scanning for the one sentence." + + +def test_a_failing_equivalence_verdict_quotes_both_legs_around_the_difference(tmp_path): + """The index says WHERE; the words say whether it is lost state or a near-tie.""" + agg = _whole([_AGG_TEXT, " same"], ids=[(1, 2, 3, 4, 5, 6), (7,)]) + dis = _whole([_DIS_TEXT, " same"], ids=[(1, 2, 3, 4, 9, 6), (7,)]) + result, _ = _run_equivalence(tmp_path, agg, dis) + assert not result.passed + assert "#0 first differs at token 4" in result.detail + # The one-line verdict quotes the words on both sides of the flip... + assert ( + "around it, #0 text [27:69] around character 57: " + "aggregated 'ound it by scanning for the only sentence.' vs " + "disagg 'ound it by scanning for the one sentence.'" + ) in result.detail + # ...and the evidence block shows ids and text, both legs, labelled. + tail = result.stdout_tail + assert "continuation #0: first differs at token 4" in tail + assert "token ids [0:6] around token 4:" in tail + assert "aggregated: [1, 2, 3, 4, 5, 6]" in tail + assert "disagg : [1, 2, 3, 4, 9, 6]" in tail + assert "text [" in tail and "the only sentence" in tail and "the one sentence" in tail + + +def test_ids_that_differ_under_identical_text_are_shown_as_ids(tmp_path): + agg = _whole([" A B"], ids=[(1, 2, 3)]) + dis = _whole([" A B"], ids=[(1, 2, 4)]) + result, _ = _run_equivalence(tmp_path, agg, dis) + assert ( + "around it, #0 token ids [0:3] around token 2: aggregated [1, 2, 3] vs disagg [1, 2, 4]" + ) in result.detail + assert "text: identical on both legs" in result.stdout_tail + + +def test_the_evidence_covers_several_differing_continuations_and_counts_the_rest(tmp_path): + texts_a = [f" answer {i} alpha" for i in range(6)] + texts_b = [f" answer {i} beta" for i in range(6)] + result, _ = _run_equivalence(tmp_path, _whole(texts_a), _whole(texts_b)) + tail = result.stdout_tail + for i in range(4): + assert f"continuation #{i}:" in tail + assert "continuation #4:" not in tail + assert "2 more differing continuation(s)" in tail + + +def test_a_failing_kv_quant_verdict_names_the_first_difference_and_both_legs(): + """It used to list the continuation and quote two 200-character prefixes.""" + ids_native, ids_bf16 = [(1,), (5, 6, 7)], [(1,), (5, 6, 8)] + native = _counter(kv_bytes_per_token=890) + "\n" + _whole([" same", _AGG_TEXT], ids_native) + bf16 = _counter(kv_bytes_per_token=3200) + "\n" + _whole([" same", _DIS_TEXT], ids_bf16) + runner, _ = _leg_runner(native, bf16, marker=_BF16) + result = run_kv_quant_witness( + _env(), + task_data=_TASK, + expected_bytes_per_token=890, + contrast_env={"KV_STORAGE_FOR_TEST": "bf16"}, + runner=runner, + unit="kv_quant", + max_batch_size=1, + ) + assert not result.passed + assert "continuation(s) [1]" in result.detail + assert "#1 first differs at token 2" in result.detail + assert "native 'ound it by scanning for the only sentence.'" in result.detail + assert "bf16 'ound it by scanning for the one sentence.'" in result.detail + assert "stored values are not the reference's" in result.detail + assert "max_batch_size=1" in result.detail + assert result.measured == 2 + assert "native: [5, 6, 7]" in result.stdout_tail + assert "bf16 : [5, 6, 8]" in result.stdout_tail + + +def test_a_failing_residency_verdict_names_the_first_difference(): + result, _ = _residency( + on_mib=1000, off_mib=1200, on_text=(" A", " B x"), off_text=(" A", " B y") + ) + assert not result.passed + assert "identical text" in result.detail + assert "#1 first differs at character 3" in result.detail + assert "offloaded ' B x'" in result.detail and "resident ' B y'" in result.detail + assert "around character 3" in result.stdout_tail diff --git a/agent-flow/tests/workflows/staircase_onboard/test_workflow.py b/agent-flow/tests/workflows/staircase_onboard/test_workflow.py new file mode 100644 index 000000000000..e5df7ddbbb82 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_workflow.py @@ -0,0 +1,1390 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the orchestrator. + +No agent is ever constructed: ``_make_agent`` is replaced by a fake that +appends the scripted progress entry the real agent would have written. That +is the whole contract between the orchestrator and its agents — a typed +decision in ``progress.yaml`` — so a fake that honours it exercises the real +control flow. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import progress as progress_mod +from agent_flow.workflows.staircase_onboard import vocabulary as vocab +from agent_flow.workflows.staircase_onboard.gates import ShellResult +from agent_flow.workflows.staircase_onboard.preflight import CommandResult +from agent_flow.workflows.staircase_onboard.state import ( + STAGE_ASSEMBLE, + STAGE_DELIVER, + STAGE_PLAN, + STAGE_PREFLIGHT, + STAGE_SMITH, + load_state, +) +from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + +from .test_delivery import TARGET_REL, _build_repo + +_SCRATCH = "/scratch/staircase-run-tmp" + + +# --------------------------------------------------------------------- fakes + + +class FakeLayer: + """Stands in for one ``AgentLayer``. Writes what the real role would write.""" + + def __init__(self, name, script, calls): + self.name = name + self.script = script + self.calls = calls + + def __call__(self, prompt: str) -> str: + self.calls.append((self.name, prompt)) + turn = sum(1 for n, _ in self.calls if n == self.name) + entry = self.script(self.name, turn, prompt) + if entry is not None: + self.script.ctx.append(self.name, {"agent": self.name, **entry}) + return f"{self.name} turn {turn} done" + + def __exit__(self, *args) -> None: + return None + + +class Script: + """Per-role scripted decisions, keyed by role and turn number.""" + + def __init__(self, ctx, plan=("DRAFT_READY",), review=("APPROVE",), **roles): + self.ctx = ctx + self.plan = list(plan) + self.review = list(review) + self.roles = roles + + def __call__(self, role, turn, prompt): + if role == "plan_drafter": + return {"decision": self.plan[min(turn, len(self.plan)) - 1], "summary": "drafted"} + if role == "plan_reviewer": + return {"decision": self.review[min(turn, len(self.review)) - 1], "summary": "reviewed"} + override = self.roles.get(role) + if callable(override): + return override(turn, prompt) + if role == "smith": + target = _target_from(prompt) + return {"decision": "LANDED", "summary": f"landed {target}", "entry": target} + if role == "assembler": + return {"decision": "UNIT_READY", "summary": "assembled"} + if role == "reviewer": + return {"decision": "APPROVE", "summary": "looks right"} + if role == "auditor": + return {"decision": "CLEAN", "summary": "8 values, 3 entries"} + raise AssertionError(f"unscripted role: {role}") + + +def _target_from(prompt: str) -> str: + for token in prompt.split("`"): + if token and " " not in token and not token.startswith("/"): + return token + return "unknown" + + +def _install(monkeypatch, workflow, script): + calls: list[tuple[str, str]] = [] + script.ctx = workflow._progress_ctx + + def fake_make_agent(self, name, system_prompt, *, required_tools, session_mode="stateless"): + return FakeLayer(name, script, calls) + + monkeypatch.setattr(StaircaseOnboardWorkflow, "_make_agent", fake_make_agent) + return calls + + +# ------------------------------------------------------------------ fixtures + + +@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, + }, + } + path = tmp_path / "task.yaml" + path.write_text(yaml.safe_dump(spec), encoding="utf-8") + return path + + +def git_probe(command: str, *, head: str = "c0ffee1", dirt: str = "") -> ShellResult | None: + """Answer the two git questions the resume check asks, or ``None``. + + `StaircaseOnboardWorkflow._target_marker` asks which commit last touched + the target and whether anything is uncommitted under it, through the same + shell runner the gates use. Every fake here therefore has to answer them, + and — this is the part worth stating — **none of them may count one as a + boot**: a probe is not a gate, and a counter that conflates the two turns + a cheap question into a phantom engine launch. + """ + if "git -C" not in command: + return None + if "log -1" in command: + return ShellResult(returncode=0, stdout=head + "\n") + if "status --porcelain" in command: + return ShellResult(returncode=0, stdout=dirt) + return ShellResult(returncode=0, stdout="") + + +def healthy_smi(count: int = 1): + rows = "\n".join(f"{i}, 180000, 183359, 0, 10.3" for i in range(count)) + + def runner(argv, env=None, timeout=0): + if argv and argv[0] == "nvidia-smi": + return CommandResult(returncode=0, stdout=rows + "\n") + return CommandResult(returncode=0, stdout="ok\n") + + return runner + + +def passing_shell(score: float = 90.0, accepted: int = 320, iterations: int = 100, dirt: str = ""): + def runner(command, env=None, timeout=0): + probe = git_probe(command, dirt=dirt) + if probe is not None: + return probe + if "trtllm-eval" in command: + return ShellResult(returncode=0, stdout=f"lm-eval gsm8k average accuracy: {score}\n") + # The acceptance gate runs the QUICKSTART command, because nothing in + # tree makes `trtllm-eval` report a speculative statistic. The fake + # has to print what the real script prints, or every speculator unit + # is unpassable for a reason that looks like a broken draft path. + return ShellResult( + returncode=0, + stdout=( + "Paris is the capital of France\n" + f"[0] Accept rate: {accepted / iterations:.2f} " + f"(tokens={accepted}, iterations={iterations})\n" + ), + ) + + return runner + + +def _workflow(tmp_path, **kwargs): + kwargs.setdefault("command_runner", healthy_smi()) + kwargs.setdefault("shell_runner", passing_shell()) + # These tests drive stages other than preflight, and several use a shell + # runner scripted to fail a specific boot. The substrate gate issues a + # boot-shaped command of its own, so it would consume those scripts. + # It has dedicated coverage in test_substrate_gate.py, plus the two + # end-to-end cases below. + kwargs.setdefault("skip_substrate_gate", True) + return StaircaseOnboardWorkflow(workspace=tmp_path / "ws", **kwargs) + + +def _write_vocabulary(workspace: Path, queue): + doc = { + "verdict": vocab.VERDICT_AFTER_ADDITIONS if queue else vocab.VERDICT_ASSEMBLABLE, + "steps": [ + { + "step": "norm", + "state": vocab.STATE_OK, + "entry": "catalog/norm/flashinfer_rmsnorm.py", + "contract_line": "certified sm_103 bf16", + } + ], + "queue": queue, + } + (workspace / "vocabulary.yaml").write_text(yaml.safe_dump(doc), encoding="utf-8") + + +# ------------------------------------------------------------------- preflight + + +def test_preflight_failure_ends_the_run_before_any_agent(tmp_path, task_file, monkeypatch): + def broken(argv, env=None, timeout=0): + return CommandResult(returncode=0, stdout="0, 20000, 183359, 99, 10.3\n") + + workflow = _workflow(tmp_path, command_runner=broken) + calls = _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task_file) + assert calls == [] + assert "preflight failed" in (tmp_path / "ws" / "journal.md").read_text() + + +def test_dry_run_stops_before_the_first_agent_turn(tmp_path, task_file, monkeypatch): + """And it leaves the checkpoint at preflight, so the real run re-runs it. + + A device claim is a property of the host *at that moment*. Letting a dry + run mark preflight complete would let a real run hours later start on a + claim that has since gone stale — which is the one thing preflight exists + to prevent. + """ + workflow = _workflow(tmp_path, dry_run=True) + calls = _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task_file) + assert calls == [] + assert load_state(tmp_path / "ws" / ".staircase_onboard_state.json").stage == STAGE_PREFLIGHT + + +def test_preflight_resolves_the_target_path(tmp_path, task_file, monkeypatch): + workflow = _workflow(tmp_path, dry_run=True) + _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task_file) + assert workflow._preflight.target_rel == TARGET_REL + + +# ------------------------------------------------------------------ plan phase + + +def test_happy_path_walks_every_stage(tmp_path, task_file, monkeypatch): + workflow = _workflow(tmp_path) + script = Script(None) + calls = _install(monkeypatch, workflow, script) + + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + roles = [name for name, _ in calls] + assert roles == ["plan_drafter", "plan_reviewer", "assembler", "reviewer", "auditor"] + assert load_state(tmp_path / "ws" / ".staircase_onboard_state.json").done + + +def test_a_ceiling_ends_the_run(tmp_path, task_file, monkeypatch): + workflow = _workflow(tmp_path) + calls = _install(monkeypatch, workflow, Script(None, plan=("CEILING",))) + with workflow: + workflow.run(task_file) + assert [n for n, _ in calls] == ["plan_drafter"] + assert "vocabulary ceiling" in (tmp_path / "ws" / "journal.md").read_text() + + +def test_plan_cap_records_open_risks_and_starts_the_build(tmp_path, task_file, monkeypatch): + """Reaching the cap is a budget, not a gate: the build starts anyway.""" + workflow = _workflow(tmp_path, max_plan_rounds=1) + script = Script(None, plan=("DRAFT_READY", "DRAFT_READY"), review=("REJECT", "REJECT")) + calls = _install(monkeypatch, workflow, script) + + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + plan_text = (tmp_path / "ws" / "plan.md").read_text() + assert "## Open risks" in plan_text + assert "reviewed" in plan_text + # Two plan rounds, then the build ran anyway. + assert [n for n, _ in calls].count("plan_drafter") == 2 + assert "assembler" in [n for n, _ in calls] + + +def test_a_supplied_plan_and_vocabulary_skip_the_plan_phase(tmp_path, task_file, monkeypatch): + workspace = tmp_path / "ws" + workspace.mkdir(parents=True) + plan = tmp_path / "plan.md" + plan.write_text("# supplied\n", encoding="utf-8") + vocabulary = tmp_path / "vocab.yaml" + vocabulary.write_text( + yaml.safe_dump({"verdict": vocab.VERDICT_ASSEMBLABLE, "steps": [], "queue": []}), + encoding="utf-8", + ) + workflow = _workflow(tmp_path, plan=plan, vocabulary=vocabulary) + calls = _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task_file) + roles = [n for n, _ in calls] + assert "plan_drafter" not in roles and "plan_reviewer" not in roles + assert (workspace / "plan.md").read_text() == "# supplied\n" + + +# ----------------------------------------------------------------- smith phase + + +def test_independent_queue_items_run_in_one_batch(tmp_path, task_file, monkeypatch): + workflow = _workflow(tmp_path) + script = Script(None) + calls = _install(monkeypatch, workflow, script) + + queue = [ + {"n": 1, "mode": "onboard", "target": "moe_gelu", "depends_on": []}, + {"n": 2, "mode": "onboard", "target": "rms_norm_gated", "depends_on": []}, + { + "n": 3, + "mode": "verify", + "target": "thop_attention", + "depends_on": [1], + "surface": "h=72", + }, + ] + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", queue) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + smith_calls = [p for n, p in calls if n == "smith"] + assert len(smith_calls) == 3 + state = load_state(tmp_path / "ws" / ".staircase_onboard_state.json") + assert all(row["status"] == "done" for row in state.queue) + + +def test_a_smith_that_cannot_land_stops_the_run(tmp_path, task_file, monkeypatch): + workflow = _workflow(tmp_path) + script = Script( + None, + smith=lambda turn, prompt: { + "decision": "BLOCKED", + "summary": "blocked on moe_gelu", + "entry": "moe_gelu", + }, + ) + calls = _install(monkeypatch, workflow, script) + + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", [{"n": 1, "mode": "onboard", "target": "moe_gelu"}]) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + assert "assembler" not in [n for n, _ in calls] + + +# -------------------------------------------------------------- assemble phase + + +def test_a_failing_boot_gate_loops_back_to_the_assembler(tmp_path, task_file, monkeypatch): + attempts = {"boot": 0} + + def flaky(command, env=None, timeout=0): + probe = git_probe(command) + if probe is not None: + return probe + if "trtllm-eval" in command: + return ShellResult(returncode=0, stdout="lm-eval gsm8k average accuracy: 90.0\n") + if "__SURVEY_JSON__" in command: + # The survey probe is a shell command like any other and must not + # be miscounted as a boot attempt. It is left failing on purpose: + # a run whose survey does not land still has to complete. + return ShellResult(returncode=1, stderr="no engine here") + attempts["boot"] += 1 + if attempts["boot"] == 1: + return ShellResult(returncode=1, stderr="CUDA error") + return ShellResult(returncode=0, stdout="Paris\n") + + workflow = _workflow(tmp_path, shell_runner=flaky) + calls = _install(monkeypatch, workflow, Script(None)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + assert [n for n, _ in calls].count("assembler") == 2 + + +def test_a_failing_accuracy_gate_is_not_papered_over(tmp_path, task_file, monkeypatch): + """The gate is the release criterion; exhausting the budget ends the run.""" + workflow = _workflow( + tmp_path, shell_runner=passing_shell(score=10.0), max_assembler_turns_per_unit=2 + ) + calls = _install(monkeypatch, workflow, Script(None)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + assert [n for n, _ in calls].count("assembler") == 2 + assert "auditor" not in [n for n, _ in calls] + assert "assembly-exhausted" in (tmp_path / "ws" / "journal.md").read_text() + + +def test_a_vocabulary_gap_dispatches_smith_and_retries(tmp_path, task_file, monkeypatch): + state_box = {"blocked_once": False} + + def assembler(turn, prompt): + if not state_box["blocked_once"]: + state_box["blocked_once"] = True + return { + "decision": "BLOCKED", + "summary": "need moe_gelu", + "blocked_reason": "vocabulary_gap", + "missing_entry": "moe_gelu", + } + return {"decision": "UNIT_READY", "summary": "assembled"} + + workflow = _workflow(tmp_path) + calls = _install(monkeypatch, workflow, Script(None, assembler=assembler)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + roles = [n for n, _ in calls] + assert roles.count("smith") == 1 + assert roles.count("assembler") == 2 + assert "auditor" in roles + + +def test_a_knowledge_gap_ends_the_run_with_the_question(tmp_path, task_file, monkeypatch): + def assembler(turn, prompt): + return { + "decision": "BLOCKED", + "summary": "which metadata field carries the sink offset?", + "blocked_reason": "knowledge_gap", + } + + workflow = _workflow(tmp_path) + calls = _install(monkeypatch, workflow, Script(None, assembler=assembler)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "knowledge_gap" in journal + assert "sink offset" in journal + assert "auditor" not in [n for n, _ in calls] + + +def test_reviewer_off_skips_the_review_turn(tmp_path, task_file, monkeypatch): + workflow = _workflow(tmp_path, reviewer_mode="off") + calls = _install(monkeypatch, workflow, Script(None)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + assert "reviewer" not in [n for n, _ in calls] + + +def test_features_are_assembled_after_core_one_at_a_time(tmp_path, task_file, monkeypatch): + spec = yaml.safe_load(task_file.read_text()) + spec["features"] = ["mtp"] + task_file.write_text(yaml.safe_dump(spec), encoding="utf-8") + + workflow = _workflow(tmp_path) + calls = _install(monkeypatch, workflow, Script(None)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + units = [p for n, p in calls if n == "assembler"] + assert "Unit: **core**" in units[0] + assert "Unit: **mtp**" in units[1] + state = load_state(tmp_path / "ws" / ".staircase_onboard_state.json") + assert state.units_done == ["core", "mtp"] + + +# ----------------------------------------------------- the member set, gated + + +def _unified_task(task_file, members): + """Rewrite the task spec as a unified run over `members`. + + World size 1 so the harness stays the one every other test here uses: the + device claim sees one card and the multi-rank substrate probe short- + circuits. What is under test is the LOOP, not the allocation. + """ + spec = yaml.safe_load(task_file.read_text()) + spec["parallel"] = "unified" + spec["members"] = members + task_file.write_text(yaml.safe_dump(spec), encoding="utf-8") + + +def _member_counting_shell(fail_on_boot=None): + """Records the topology flags of every boot and accuracy command.""" + boots: list[str] = [] + evals: list[str] = [] + + def runner(command, env=None, timeout=0): + probe = git_probe(command) + if probe is not None: + return probe + if "__SURVEY_JSON__" in command: + return ShellResult(returncode=1, stderr="no engine here") + if "trtllm-eval" in command: + evals.append(command) + return ShellResult(returncode=0, stdout="lm-eval gsm8k average accuracy: 90.0\n") + boots.append(command) + if fail_on_boot is not None and len(boots) == fail_on_boot: + return ShellResult(returncode=1, stderr="CUDA error") + return ShellResult(returncode=0, stdout="Paris\n") + + runner.boots = boots # type: ignore[attr-defined] + runner.evals = evals # type: ignore[attr-defined] + return runner + + +def _run_unified(tmp_path, task_file, monkeypatch, shell, **kwargs): + workflow = _workflow(tmp_path, shell_runner=shell, **kwargs) + calls = _install(monkeypatch, workflow, Script(None)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + return workflow, calls + + +def test_a_unified_unit_is_gated_once_per_member(tmp_path, task_file, monkeypatch): + """Three members, three boots and three accuracy runs — at three topologies. + + The finding this closes: a unified target was built once and gated once, + at pure tensor parallelism, so `--moe_ep_size` and `--enable_attention_dp` + were emitted by no gate while the target's Coverage claimed all three + configurations. + """ + _unified_task(task_file, ["tp1", "tep1", "dep1"]) + shell = _member_counting_shell() + workflow, _ = _run_unified(tmp_path, task_file, monkeypatch, shell) + + assert len(shell.boots) == 3, shell.boots + assert len(shell.evals) == 3, shell.evals + # And each ran at a DIFFERENT topology, which is the whole point. + # + # The two commands spell the same topology differently, so the boot and + # accuracy legs are asserted separately rather than together. `tep` and + # `dep` both carry expert parallelism; only `dep` carries attention DP, + # and `trtllm-eval` has no flag for it at all -- it arrives through the + # member's generated options file, so asserting it on the eval argv would + # be asserting the defect this split fixed. + assert sum("--moe_ep_size=1" in c for c in shell.boots) == 2 + assert sum("--enable_attention_dp" in c for c in shell.boots) == 1 + assert sum("--ep_size=1" in c for c in shell.evals) == 2 + assert not any("--moe_ep_size" in c for c in shell.evals) + assert not any("--enable_attention_dp" in c for c in shell.evals) + # One verdict per member, never collapsed, each naming its member. + rendered = workflow.gate_log.render() + for member in ("tp1", "tep1", "dep1"): + assert f"accuracy[gsm8k]@{member}" in rendered + + +def test_a_failure_on_the_second_member_stops_the_unit(tmp_path, task_file, monkeypatch): + """The third member would spend an hour measuring the same defect.""" + _unified_task(task_file, ["tp1", "tep1", "dep1"]) + shell = _member_counting_shell(fail_on_boot=2) + workflow, calls = _run_unified( + tmp_path, task_file, monkeypatch, shell, max_assembler_turns_per_unit=1 + ) + + assert len(shell.boots) == 2, shell.boots + assert shell.evals == [] # the unit never reached the accuracy gate + assert [n for n, _ in calls].count("assembler") == 1 + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "boot@tep1" in journal # the message names the failing member + assert "assembly-exhausted" in journal + + +def test_a_single_member_is_still_a_member(tmp_path, task_file, monkeypatch): + """`members:` is how a run buys a cheaper answer honestly — one, named.""" + _unified_task(task_file, ["dep1"]) + shell = _member_counting_shell() + workflow, _ = _run_unified(tmp_path, task_file, monkeypatch, shell) + + assert len(shell.boots) == 1 + assert len(shell.evals) == 1 + # `dep` reaches the boot gate as a flag and the accuracy gate as a + # generated options file, because `trtllm-eval` has no attention-DP flag. + assert "--enable_attention_dp" in shell.boots[0] + assert "--ep_size=1" in shell.evals[0] + assert "--enable_attention_dp" not in shell.evals[0] + assert "@dep1" in workflow.gate_log.render() + + +# ---------------------------------------------------------------- deliver phase + + +def test_an_uncovered_audit_blocks_delivery(tmp_path, task_file, monkeypatch): + workflow = _workflow(tmp_path) + script = Script( + None, + auditor=lambda turn, prompt: { + "decision": "UNCOVERED", + "summary": "batch 257 outside thop_attention's column", + }, + ) + _install(monkeypatch, workflow, script) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "audit-uncovered" in journal + assert "batch 257" in journal + + +# ----------------------------------------------------------------------- resume + + +def test_resume_picks_up_at_the_saved_stage(tmp_path, task_file, monkeypatch): + workflow = _workflow(tmp_path) + script = Script(None, plan=("DRAFT_READY",), review=("APPROVE",)) + _install(monkeypatch, workflow, script) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + # Second run over the same workspace: already done, so it no-ops rather + # than re-spending the plan phase. + second = _workflow(tmp_path) + calls = _install(monkeypatch, second, Script(None)) + with second: + second.run(task_file) + assert calls == [] + + +# --------------------------------------- a resume must not spend a turn + + +def _plan_then_vocabulary(workflow, tmp_path, monkeypatch, queue=()): + """The idiom every test here uses: write a vocabulary, then plan.""" + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", list(queue)) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + + +def _die_after_the_assembler(tmp_path, task_file, monkeypatch, **kwargs): + """One orchestrator that assembles `core`, then dies before the unit closes. + + This is the shape of every restart in the run that motivated the fix: the + assembler turn finished, its `UNIT_READY` and its work are on disk, and + the process that dispatched it is gone. A `KeyboardInterrupt` raised from + the reviewer's turn reproduces it exactly — `run` journals it and + re-raises, leaving the checkpoint at `assemble` with the unit still open. + """ + + def reviewer(turn, prompt): + raise KeyboardInterrupt + + workflow = _workflow(tmp_path, **kwargs) + _install(monkeypatch, workflow, Script(None, reviewer=reviewer)) + _plan_then_vocabulary(workflow, tmp_path, monkeypatch) + with pytest.raises(KeyboardInterrupt), workflow: + workflow.run(task_file) + state = load_state(tmp_path / "ws" / ".staircase_onboard_state.json") + assert state.stage == STAGE_ASSEMBLE and state.unit_turns == 1 + return state + + +def test_a_resume_onto_a_ready_unit_does_not_dispatch(tmp_path, task_file, monkeypatch): + """The single largest measurable waste in the run this closes. + + Ten assembler dispatches for four units, six of which built nothing, + because a spec fix needs a restart and every restart re-entered the unit + and dispatched. The prior `UNIT_READY` was on disk the whole time. + """ + before = _die_after_the_assembler(tmp_path, task_file, monkeypatch) + assert before.unit_marker["core"], "the marker is what the resume compares against" + + second = _workflow(tmp_path) + calls = _install(monkeypatch, second, Script(None)) + with second: + second.run(task_file) + + # The unit picks up at the gates: no assembler turn, and the run closes. + assert [n for n, _ in calls] == ["reviewer", "auditor"], calls + assert load_state(tmp_path / "ws" / ".staircase_onboard_state.json").verdict == "delivered" + assert "did not spend a turn" in (tmp_path / "ws" / "journal.md").read_text() + + +def test_a_resume_onto_a_modified_target_re_dispatches(tmp_path, task_file, monkeypatch): + """Whether the target is unchanged is checked, not assumed. + + A restart after somebody edited the target is a different situation from + a clean one: the `UNIT_READY` no longer describes what is on disk, so the + turn is worth spending. + """ + _die_after_the_assembler(tmp_path, task_file, monkeypatch) + + dirty = passing_shell(dirt=" M tensorrt_llm/.../modeling.py\n") + second = _workflow(tmp_path, shell_runner=dirty) + calls = _install(monkeypatch, second, Script(None)) + with second: + second.run(task_file) + + assert [n for n, _ in calls].count("assembler") == 1, calls + + +def test_a_marker_that_cannot_be_established_re_dispatches(tmp_path, task_file, monkeypatch): + """No git, no answer — and no answer means "assume changed". + + A wasted turn costs one dispatch; a skipped turn that was needed costs + the unit, so the unanswerable case takes the expensive side. + """ + _die_after_the_assembler(tmp_path, task_file, monkeypatch) + + def no_git(command, env=None, timeout=0): + if "git -C" in command: + return ShellResult(returncode=128, stderr="not a git repository") + return passing_shell()(command, env, timeout) + + second = _workflow(tmp_path, shell_runner=no_git) + calls = _install(monkeypatch, second, Script(None)) + with second: + second.run(task_file) + + assert [n for n, _ in calls].count("assembler") == 1, calls + + +def test_the_turn_a_resume_gives_back_is_still_spendable(tmp_path, task_file, monkeypatch): + """`unit_turns` counts work, so the skipped restart leaves the budget intact. + + Two turns, and the resumed run needs all of: the skip, a failing boot + gate, and a real assembler turn. Under the old accounting the restart + would have spent turn 2 and the gate failure would have ended the unit + `assembly-exhausted` with a budget it never used on work. + """ + _die_after_the_assembler(tmp_path, task_file, monkeypatch) + + attempts = {"boot": 0} + + def one_bad_boot(command, env=None, timeout=0): + probe = git_probe(command) + if probe is not None: + return probe + if "trtllm-eval" in command or "__SURVEY_JSON__" in command: + return passing_shell()(command, env, timeout) + attempts["boot"] += 1 + if attempts["boot"] == 1: + return ShellResult(returncode=1, stderr="CUDA error") + return ShellResult(returncode=0, stdout="Paris\n") + + second = _workflow(tmp_path, shell_runner=one_bad_boot, max_assembler_turns_per_unit=2) + calls = _install(monkeypatch, second, Script(None)) + with second: + second.run(task_file) + + assert [n for n, _ in calls].count("assembler") == 1, calls + assert load_state(tmp_path / "ws" / ".staircase_onboard_state.json").verdict == "delivered" + + +def test_a_resume_onto_a_blocked_assembler_still_dispatches_smith(tmp_path, task_file, monkeypatch): + """The pre-existing resume handling, unchanged by the UNIT_READY branch.""" + workflow = _workflow(tmp_path) + + def assembler(turn, prompt): + if turn == 1: + return { + "decision": "BLOCKED", + "summary": "need moe_gelu", + "blocked_reason": "vocabulary_gap", + "missing_entry": "moe_gelu", + } + return {"decision": "UNIT_READY", "summary": "assembled"} + + _install(monkeypatch, workflow, Script(None, assembler=assembler)) + _plan_then_vocabulary(workflow, tmp_path, monkeypatch) + + # Die between the decision landing and anything acting on it, which is + # the case `_resumed_into_unit` was introduced for. + def died(state, log): + raise KeyboardInterrupt + + monkeypatch.setattr(workflow, "_handle_blocked_assembler", died) + with pytest.raises(KeyboardInterrupt), workflow: + workflow.run(task_file) + + # The resumed process's assembler has the entry it was missing, so its + # turn lands — the default scripted `UNIT_READY`. + second = _workflow(tmp_path) + calls = _install(monkeypatch, second, Script(None)) + with second: + second.run(task_file) + + roles = [n for n, _ in calls] + assert roles.count("smith") == 1, roles + assert roles.count("assembler") == 1, roles + assert "auditor" in roles + + +# --------------------------------------------------- host facts, corrections + + +def test_host_facts_are_written_and_reach_a_role_prompt(tmp_path, task_file, monkeypatch): + """Six facts were hand-injected into ~15 briefs on the run this closes.""" + workflow = _workflow(tmp_path) + calls = _install(monkeypatch, workflow, Script(None)) + _plan_then_vocabulary(workflow, tmp_path, monkeypatch) + with workflow: + workflow.run(task_file) + + facts = (tmp_path / "ws" / "host_facts.md").read_text() + assert "TRTLLM_SKIP_SYS_PATH_CHECK=1" in facts + assert "UCX ERROR" in facts + # Seeded from what the run resolved, not hand-typed: this host's own + # TMPDIR and the bound preflight owns. + assert _SCRATCH in facts + assert "61 characters" in facts + # And it says which document it is not, so nobody merges the two. + assert "TROUBLESHOOTING.md" in facts + + for role in ("plan_drafter", "assembler", "auditor"): + prompt = next(p for name, p in calls if name == role) + assert str(tmp_path / "ws" / "host_facts.md") in prompt, role + assert "Learned during this run" in prompt, role + + +def test_a_fact_a_role_appended_survives_the_next_preflight(tmp_path, task_file, monkeypatch): + """Appendable, or it is just another list somebody else wrote.""" + _die_after_the_assembler(tmp_path, task_file, monkeypatch) + facts = tmp_path / "ws" / "host_facts.md" + with facts.open("a", encoding="utf-8") as handle: + handle.write("\n- **`nvidia-smi` lies about MIG here.** Read `/proc` instead.\n") + + second = _workflow(tmp_path) + _install(monkeypatch, second, Script(None)) + with second: + second.run(task_file) + + text = facts.read_text() + assert "lies about MIG here" in text + assert "TRTLLM_SKIP_SYS_PATH_CHECK=1" in text # the seeded half is still there + + +def test_no_corrections_means_no_reference_to_them(tmp_path, task_file, monkeypatch): + """A pointer at an empty file teaches the reader the pointer is worthless.""" + workflow = _workflow(tmp_path) + calls = _install(monkeypatch, workflow, Script(None)) + _plan_then_vocabulary(workflow, tmp_path, monkeypatch) + with workflow: + workflow.run(task_file) + + # The file exists — a role or an operator needs somewhere to write — but + # it is empty, so nothing points at it. + assert (tmp_path / "ws" / "corrections.md").is_file() + assert not any("corrections.md" in prompt for _, prompt in calls), calls + + +def test_a_late_review_becomes_a_correction_that_reaches_a_dispatch( + tmp_path, task_file, monkeypatch +): + """`SMITH_CORRECTIONS.md`, made first-class. + + The round-2 plan review landed after the smith queue was written. One of + its findings was that a rank's FP8 GEMM never sees `N=192`; without it a + smith would have certified a width no rank ever passes. + """ + workflow = _workflow(tmp_path, max_plan_rounds=1) + script = Script(None, plan=("DRAFT_READY", "DRAFT_READY"), review=("REJECT", "REJECT")) + calls = _install(monkeypatch, workflow, script) + _plan_then_vocabulary( + workflow, tmp_path, monkeypatch, queue=[{"n": 1, "mode": "onboard", "target": "moe_gelu"}] + ) + with workflow: + workflow.run(task_file) + + corrections = (tmp_path / "ws" / "corrections.md").read_text() + assert "PlanReviewer" in corrections + assert "reviewed" in corrections # the reviewer's own summary, verbatim + assert "this file wins" in corrections + + # And it reaches the dispatches the review was too late for, with the + # precedence rule stated rather than implied. + for role in ("smith", "assembler"): + prompt = next(p for name, p in calls if name == role) + assert str(tmp_path / "ws" / "corrections.md") in prompt, role + assert "the corrections file wins" in prompt, role + + +def test_journal_records_dispatch_prompts_before_they_are_sent(tmp_path, task_file, monkeypatch): + """After a session cut the journal is the only explanation of the disk.""" + workflow = _workflow(tmp_path) + _install(monkeypatch, workflow, Script(None)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + reports = sorted((tmp_path / "ws" / "reports").iterdir()) + assert reports, "every dispatch writes a report" + first = reports[0].read_text() + assert "## Dispatch prompt" in first + assert "## Result" in first + # Ordinals are sequential across the whole run, not per role. + assert reports[0].name.startswith("01-") + + +def test_progress_entries_are_typed_and_stage_scoped(tmp_path, task_file, monkeypatch): + workflow = _workflow(tmp_path) + _install(monkeypatch, workflow, Script(None)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + data = progress_mod.read_progress(tmp_path / "ws" / "progress.yaml") + assert [e["agent"] for e in data[progress_mod.PLAN_STAGE]] == ["plan_drafter", "plan_reviewer"] + assert [e["agent"] for e in data[progress_mod.BUILD_STAGE]] == ["assembler", "reviewer"] + assert [e["agent"] for e in data[progress_mod.DELIVER_STAGE]] == ["auditor"] + + +def test_stage_constants_advance_in_order(): + assert [STAGE_PLAN, STAGE_SMITH, STAGE_ASSEMBLE, STAGE_DELIVER] == [ + "plan", + "smith", + "assemble", + "deliver", + ] + + +# --------------------------------------------------- the substrate gate in run() + + +def test_run_invokes_the_substrate_gate_before_any_agent(tmp_path, task_file, monkeypatch): + """It must come first, or its whole purpose is defeated.""" + order: list[str] = [] + + def shell(command, env=None, timeout=0): + if "TRTLLM_MODELING_V2=off" in command and "quickstart" in command: + order.append("substrate") + return ShellResult(returncode=0, stdout="Paris\n") + + workflow = _workflow(tmp_path, shell_runner=shell, skip_substrate_gate=False) + + original = workflow._run_plan_phase + + def note_plan(state, log): + order.append("plan") + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", note_plan) + _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task_file) + + assert order and order[0] == "substrate", order + + +def test_a_failing_substrate_gate_stops_the_run_with_no_agent_turn( + tmp_path, task_file, monkeypatch +): + """A broken host must cost seconds, not a planning turn.""" + + def shell(command, env=None, timeout=0): + return ShellResult(returncode=1, stdout="MPI_ERR_SPAWN: could not spawn processes") + + workflow = _workflow(tmp_path, shell_runner=shell, skip_substrate_gate=False) + calls = _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task_file) + + assert calls == [], f"an agent was dispatched despite a dead host: {calls}" + + +# --------------------------------------------------------------- the survey + + +def test_the_survey_reaches_the_plan_drafter(tmp_path, task_file, checkpoint, monkeypatch): + """A measurement nobody reads is not a measurement.""" + checkpoint.joinpath("config.json").write_text( + json.dumps( + { + "architectures": ["GptOssForCausalLM"], + "num_attention_heads": 64, + "num_key_value_heads": 8, + "num_local_experts": 128, + "sinks": 64, + "hidden_act": "silu", + "rms_norm_eps": 1e-5, + "some_layer_nobody_supports": True, + } + ) + ) + + def surveying(command, env=None, timeout=0): + if "__SURVEY_JSON__" in command: + payload = { + "trtllm_ops": ["fused_qk_norm_rope", "an_op_with_no_entry"], + "catalog": [ + { + "path": "attention/fused_qk_norm_rope.py", + "impl": "torch.ops.trtllm.fused_qk_norm_rope", + "summary": "s", + "receipt": "passed", + "impl_resolves": True, + } + ], + "modules": ["rms_norm.RMSNorm"], + "versions": {"torch": "2.12"}, + "error": "", + } + return ShellResult(returncode=0, stdout="__SURVEY_JSON__" + json.dumps(payload)) + return passing_shell()(command, env, timeout) + + workflow = _workflow(tmp_path, shell_runner=surveying) + calls = _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task_file) + + prompt = next(p for name, p in calls if name == "plan_drafter") + assert "moe.routing" in prompt + assert "attention.sinks" in prompt + assert "some_layer_nobody_supports" in prompt # the unknown, surfaced at plan time + assert "fused_qk_norm_rope" in prompt # free vocabulary + assert "an_op_with_no_entry" in prompt # the tier-1 rung + assert (tmp_path / "ws" / "surface.md").is_file() + + +def test_a_survey_that_cannot_run_does_not_end_the_run(tmp_path, task_file, monkeypatch): + """Planning on assumption is how every plan before this one was drafted.""" + + def no_engine(command, env=None, timeout=0): + if "__SURVEY_JSON__" in command: + return ShellResult(returncode=1, stderr="ModuleNotFoundError: tensorrt_llm") + return passing_shell()(command, env, timeout) + + workflow = _workflow(tmp_path, shell_runner=no_engine) + calls = _install(monkeypatch, workflow, Script(None)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + prompt = next(p for name, p in calls if name == "plan_drafter") + assert "you do not know what exists" in prompt + assert [n for n, _ in calls].count("assembler") >= 1 # the run still ran + + +def test_a_survey_that_did_not_land_is_not_printed_as_a_check_mark( + tmp_path, task_file, monkeypatch, capsys +): + """`✔ survey ... 0/0 entries certified` was the whole defect. + + An inventory that failed to land and a catalog with nothing certified + rendered as the same green line, so an operator had no way to tell + "nothing is certified" from "nothing was measured". The run still + proceeds -- that is the documented policy -- it just proceeds loudly. + """ + + def no_mpirun(command, env=None, timeout=0): + if "__SURVEY_JSON__" in command: + return ShellResult(returncode=127, stderr="mpirun: command not found") + return passing_shell()(command, env, timeout) + + workflow = _workflow(tmp_path, shell_runner=no_mpirun, dry_run=True) + _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task_file) + + printed = " ".join(capsys.readouterr().out.split()) + assert "⚠ survey" in printed + assert "✔ survey" not in printed + assert "inventory probe failed" in printed + + +def test_a_dry_run_takes_the_survey(tmp_path, task_file, monkeypatch, capsys): + """It is a read-only, single-process measurement now, so it is affordable. + + A dry run that reports `0/0 entries certified` is the reading this + survey exists to prevent, and an operator checking whether a host can + host a run is exactly who needs the real numbers. + """ + + def surveying(command, env=None, timeout=0): + if "__SURVEY_JSON__" in command: + payload = { + "trtllm_ops": ["fused_qk_norm_rope"], + "catalog": [ + { + "path": "attention/fused_qk_norm_rope.py", + "impl": "torch.ops.trtllm.fused_qk_norm_rope", + "summary": "s", + "receipt": "passed", + "impl_resolves": True, + } + ], + "modules": [], + "versions": {}, + "error": "", + } + return ShellResult(returncode=0, stdout="__SURVEY_JSON__" + json.dumps(payload)) + return passing_shell()(command, env, timeout) + + workflow = _workflow(tmp_path, shell_runner=surveying, dry_run=True) + _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task_file) + + printed = " ".join(capsys.readouterr().out.split()) + assert "✔ survey" in printed + assert "1/1 entries certified on sm_103" in printed + + +# ------------------------------------------------------ smith concurrency + + +def _concurrency_probe(): + """Records the true peak of simultaneously-live smith dispatches.""" + import threading + + state = {"live": 0, "peak": 0} + lock = threading.Lock() + gate = threading.Event() + + def enter(): + with lock: + state["live"] += 1 + state["peak"] = max(state["peak"], state["live"]) + # Hold every dispatch open until the last one that will ever be + # admitted has arrived, so the peak is the pool's real limit and not + # an artefact of fast dispatches retiring before slow ones start. + gate.wait(timeout=5) + with lock: + state["live"] -= 1 + + return state, gate, enter + + +def _run_with_queue(tmp_path, task_file, monkeypatch, *, count, cap): + """Drive the smith phase over `count` independent items at `cap`.""" + state, gate, enter = _concurrency_probe() + + kwargs = {} if cap is None else {"max_parallel_smiths": cap} + workflow = _workflow(tmp_path, **kwargs) + calls = _install(monkeypatch, workflow, Script(None)) + + original = workflow._run_one_smith + + def counted(item): + enter() + return original(item) + + monkeypatch.setattr(workflow, "_run_one_smith", counted) + + queue = [{"n": i, "mode": "onboard", "target": f"entry_{i}"} for i in range(1, count + 1)] + + def plan_then_write(s, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", queue) + return workflow.__class__._run_plan_phase(workflow, s, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + + # Release the gate once the pool cannot admit anyone else. + import threading + + def release(): + import time + + deadline = time.monotonic() + 4 + while time.monotonic() < deadline: + if state["live"] >= min(cap or count, count): + break + time.sleep(0.01) + gate.set() + + releaser = threading.Thread(target=release, daemon=True) + releaser.start() + with workflow: + workflow.run(task_file) + releaser.join(timeout=6) + return state["peak"], calls + + +def test_the_default_caps_concurrent_smiths(tmp_path, task_file, monkeypatch): + """The default used to be 0, which reads as "no limit" and meant "all". + + On the queue this workflow was built against that was 19 agent sessions + opened at once, because every entry declared no dependencies and the + layering produced a single batch. + """ + from agent_flow.workflows.staircase_onboard.workflow import DEFAULT_MAX_PARALLEL_SMITHS + + assert DEFAULT_MAX_PARALLEL_SMITHS == 10 + peak, _ = _run_with_queue(tmp_path, task_file, monkeypatch, count=19, cap=None) + assert peak <= DEFAULT_MAX_PARALLEL_SMITHS, f"peak {peak} exceeded the default cap" + + +def test_an_explicit_cap_is_honoured(tmp_path, task_file, monkeypatch): + peak, _ = _run_with_queue(tmp_path, task_file, monkeypatch, count=12, cap=3) + assert peak <= 3 + + +def test_every_item_still_runs_under_a_cap(tmp_path, task_file, monkeypatch): + """Capping is a throttle, not a filter.""" + _, calls = _run_with_queue(tmp_path, task_file, monkeypatch, count=12, cap=3) + assert [n for n, _ in calls].count("smith") == 12 + + +def test_zero_still_means_uncapped_for_whoever_asks_for_it(tmp_path, task_file, monkeypatch): + """It stays available, but it is now an explicit choice rather than the default.""" + peak, _ = _run_with_queue(tmp_path, task_file, monkeypatch, count=14, cap=0) + assert peak > 10 + + +def test_both_dispatch_sites_stamp_the_decision_marker(): + """A turn that never answers must not leave the previous verdict fresh. + + `(agent, iteration)` is not unique across a run -- `iteration` is 1 for + every deliver turn and restarts per unit in assembly -- so two turns + routinely share a decision path. Without a marker stamped at dispatch, + `collect_decision_file` cannot tell this turn's answer from the last + one's. The predecessor run read three audit verdicts from two passes + earlier, which made the audit loop structurally unable to converge. + """ + import inspect + + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + for method in (StaircaseOnboardWorkflow._dispatch, StaircaseOnboardWorkflow._run_one_smith): + assert "mark_decision_dispatch" in inspect.getsource(method), ( + f"{method.__name__} dispatches a decision-file turn without stamping a marker" + ) + + +def test_preflight_surfaces_the_two_shaping_hazards(): + """Both are legal, both cost hours, and both are invisible until they bite.""" + import inspect + + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + source = inspect.getsource(StaircaseOnboardWorkflow._run_preflight) + assert "image_import_warning" in source + assert "shares_allocation_with_spool" in source + + +def test_a_gate_may_have_its_own_launch_substrate(): + """`environment.workloads.gate` exists so a gate can be its own failure domain.""" + import inspect + + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + source = inspect.getsource(StaircaseOnboardWorkflow._gate_env) + assert 'workload="gate"' in source, "the gate env ignores the workloads override" From 0d972897ea7dd018ab820d63b65e66ac0b74720e Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 12 Oct 2026 00:55:58 -0700 Subject: [PATCH 02/70] [TRTLLM-16460][feat] staircase: serving envelope, feature matrix and certification of what a target claims task.yaml declares a serving envelope and a feature-combination matrix; preflight derives the gate tuples and the certified set from them (replacing 'the tuples the gates happened to run'). Adds the capacity-constant audit of the target against the envelope, the per-call-site smith brief and mandatory integer-width probes in kernel-verify, a typed certification_gap decision fanned out to verify smiths by the orchestrator, early delivery checks run as warnings after planning and after every unit, and a preflight refusal of a speculation algorithm the gates' launch path cannot run. Squashed from user/williamj/staircase-polish @ c8109d6f90 (7 commits): - [feat] staircase: declare the serving envelope and feature matrix, derive the certified set from them - [feat] staircase: audit the target's capacity constants against the serving envelope - [feat] staircase: smith brief per call site, and mandatory integer-width probes in kernel-verify - [feat] staircase: typed certification_gap decision, fanned out to verify smiths by the orchestrator - [feat] staircase: run the delivery checks that need no finished target early, as warnings - [fix] staircase: envelope follow-ups found replaying the DeepSeek-V4.1 spec - [feat] staircase: preflight refuses a speculation algorithm the gates' launch path has no branch for Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../staircase_onboard/capacity_audit.py | 759 +++++++ .../staircase_onboard/early_checks.py | 608 ++++++ .../workflows/staircase_onboard/envelope.py | 1788 +++++++++++++++++ .../workflows/staircase_onboard/preflight.py | 241 +++ .../workflows/staircase_onboard/progress.py | 87 +- .../staircase_onboard/prompts/assembler.py | 16 + .../staircase_onboard/prompts/auditor.py | 14 + .../staircase_onboard/prompts/plan_drafter.py | 33 + .../prompts/plan_reviewer.py | 6 + .../staircase_onboard/prompts/reviewer.py | 32 +- .../staircase_onboard/prompts/smith.py | 10 + .../staircase_onboard/remediation.py | 173 ++ .../references/entry-spec.md | 6 + .../skills/staircase-kernel-verify/SKILL.md | 48 + .../references/integer-width-probes.md | 65 + .../skills/staircase-modeling-core/SKILL.md | 20 + .../staircase_onboard/smith_brief.py | 436 ++++ .../staircase_onboard/task.example.yaml | 64 + .../staircase_onboard/task_schema.py | 27 + .../workflows/staircase_onboard/workflow.py | 229 ++- .../staircase_onboard/test_capacity_audit.py | 240 +++ .../test_certification_gap.py | 300 +++ .../staircase_onboard/test_early_checks.py | 252 +++ .../staircase_onboard/test_envelope.py | 615 ++++++ .../staircase_onboard/test_gate_wiring.py | 9 +- .../staircase_onboard/test_smith_brief.py | 223 ++ .../test_speculation_support.py | 173 ++ 27 files changed, 6448 insertions(+), 26 deletions(-) create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/capacity_audit.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/early_checks.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/envelope.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/references/integer-width-probes.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/smith_brief.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_capacity_audit.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_certification_gap.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_early_checks.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_envelope.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_smith_brief.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_speculation_support.py diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/capacity_audit.py b/agent-flow/agent_flow/workflows/staircase_onboard/capacity_audit.py new file mode 100644 index 000000000000..b94d6e5961be --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/capacity_audit.py @@ -0,0 +1,759 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Every fixed size in a target must answer to an engine limit. + +A catalog entry is certified over the serving envelope, kernel by kernel. A +target can still set a ceiling none of those receipts sees: a module-level +``_TOPK_MAX_ROWS = 8192`` that sizes an offsets buffer at construction makes +every engine with ``max_num_tokens`` above 8192 die in warm-up, while every +kernel it feeds is certified to 16384. No gate launched above 8192, so no gate +saw it either. + +This module reads a target statically -- no import, no torch -- and lists: + +- **capacity constants**: module- and class-level integer constants whose + name says they bound something (``MAX``, ``ROWS``, ``TOKENS``, ``BATCH``, + ``CAPACITY``, ...), or that size an allocation, or that a comparison or a + ``min``/``max`` uses as a ceiling; +- **fixed-size buffers**: allocation calls (``torch.empty``/``zeros``/ + ``arange``/...) sized by such a constant or by an integer literal of 1024 or + more. + +Each one must be **tied** to the engine: sized from an engine argument +(``model_config.max_num_tokens``, ``max_batch_size``, ``max_seq_len``, ...), +which the audit recognizes and passes; or **bound** by a declaration, which +the audit then checks against the envelope. Two spellings of the declaration, +both read without importing the target:: + + _TOPK_MAX_ROWS = 16384 # capacity: max_num_tokens + _WINDOW = 128 # capacity: fixed: the checkpoint's sliding window + + CAPACITY_BOUNDS = { + "_GATHER_ROWS": "max_num_tokens * world_size", + "Indexer.SLOTS": "max_batch_size", + } + +A bound constant passes when its value is at least the envelope's limit for +that axis (times the factor, where ``world_size`` is the widest attention-DP +gather the envelope claims). ``fixed: `` records that the number is +not a capacity at all -- a tile width, a window the checkpoint fixes -- and the +reason is what the reviewer checks. +""" + +from __future__ import annotations + +import ast +import operator +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Iterable, Mapping + +from .envelope import ENVELOPE_AXES, ServingEnvelope + +#: Engine arguments a size can be derived from. A size expression that reads +#: one of these -- as a name or as an attribute of anything -- follows the +#: engine's configuration and is tied by construction. +ENGINE_ARGS: tuple[str, ...] = ( + *ENVELOPE_AXES, + "max_draft_len", + "max_beam_width", + "max_num_requests", + "max_num_sequences", + "tokens_per_block", + "max_total_draft_tokens", +) + +#: What a binding may name: an envelope axis, or the draft length. +BINDABLE_AXES: tuple[str, ...] = (*ENVELOPE_AXES, "max_draft_len") + +#: The module-level mapping a target may declare its bindings in. +BOUNDS_NAME = "CAPACITY_BOUNDS" + +#: ``# capacity: `` on the constant's or the allocation's line. +_INLINE = re.compile(r"#\s*capacity:\s*(?P.+?)\s*$") + +#: Name segments that say a constant bounds something. +CAPACITY_WORDS: tuple[str, ...] = ( + "MAX", + "ROWS", + "TOKENS", + "BATCH", + "CAPACITY", + "CAP", + "SEQ", + "SEQLEN", + "LEN", + "LIMIT", + "SLOTS", + "POOL", + "PAGES", + "BLOCKS", + "ENTRIES", + "REQUESTS", + "BUF", + "BUFFER", +) + +#: Calls that allocate a buffer whose size is their argument. +ALLOC_FUNCS: frozenset[str] = frozenset( + { + "empty", + "zeros", + "ones", + "full", + "arange", + "empty_strided", + "randn", + "rand", + "randint", + "new_empty", + "new_zeros", + "new_ones", + "new_full", + } +) + +#: An integer literal at least this large, sizing a buffer, is a capacity. +LITERAL_THRESHOLD = 1024 + +_CONSTANT_NAME = re.compile(r"^_?[A-Z][A-Z0-9_]*$") + +_BINARY = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.FloorDiv: operator.floordiv, + ast.LShift: operator.lshift, + ast.Pow: operator.pow, +} + +STATUS_OK = "ok" +STATUS_FAIL = "fail" + + +@dataclass(frozen=True) +class Binding: + """A parsed declaration: an axis times a factor, or ``fixed`` with a reason.""" + + raw: str + axis: str | None = None + factor: int = 1 + world: bool = False + reason: str = "" + error: str = "" + + @property + def fixed(self) -> bool: + return self.axis is None and not self.error + + def required(self, envelope: ServingEnvelope) -> int | None: + """The smallest value that covers the envelope, or None if unknowable.""" + if self.axis is None: + return None + if self.axis == "max_draft_len": + base = envelope.max_draft_len or None + else: + base = envelope.limit(self.axis) + if base is None: + return None + world = envelope.max_gather_world_size if self.world else 1 + return base * self.factor * world + + +def parse_binding(text: str) -> Binding: + """``max_num_tokens``, ``max_num_tokens * 4``, ``... * world_size``, ``fixed: why``.""" + raw = text.strip() + if raw.lower().startswith("fixed"): + reason = raw.partition(":")[2].strip() + if not reason: + return Binding( + raw=raw, error="`fixed` needs a reason: `fixed: `" + ) + return Binding(raw=raw, reason=reason) + parts = [part.strip() for part in raw.split("*")] + axis = parts[0] + if axis not in BINDABLE_AXES: + return Binding( + raw=raw, + error=f"binds to {axis!r}, which is not one of {list(BINDABLE_AXES)} or `fixed: `", + ) + factor = 1 + world = False + for part in parts[1:]: + if part == "world_size": + world = True + elif part.isdigit() and int(part) >= 1: + factor *= int(part) + else: + return Binding(raw=raw, error=f"factor {part!r} is not an integer or `world_size`") + return Binding(raw=raw, axis=axis, factor=factor, world=world) + + +@dataclass(frozen=True) +class Constant: + """One integer constant at module or class scope.""" + + name: str + value: int + file: str + line: int + inline: str = "" + + +@dataclass(frozen=True) +class Allocation: + """One allocation call sized by a constant or a large literal.""" + + call: str + file: str + line: int + scope: str + constants: tuple[str, ...] = () + literal: int | None = None + tied: bool = False + inline: str = "" + + +@dataclass(frozen=True) +class CapacityFinding: + """One audited capacity, and its verdict.""" + + kind: str + name: str + file: str + line: int + value: int | None + status: str + detail: str + binding: str = "" + + @property + def passed(self) -> bool: + return self.status == STATUS_OK + + +@dataclass(frozen=True) +class CapacityAudit: + """The audit of one target against one envelope.""" + + findings: tuple[CapacityFinding, ...] = () + files: tuple[str, ...] = () + tied: tuple[Allocation, ...] = () + envelope_declared: bool = True + errors: tuple[str, ...] = field(default_factory=tuple) + + @property + def failures(self) -> tuple[CapacityFinding, ...]: + return tuple(f for f in self.findings if not f.passed) + + @property + def passed(self) -> bool: + return not self.failures and not self.errors + + def checks(self) -> list[Any]: + """Delivery checks: one per failing capacity, or one summary when clean.""" + from .delivery import Check + + checks = [ + Check(name=f"capacity:{f.name}", passed=False, detail=f"{f.file}:{f.line}: {f.detail}") + for f in self.failures + ] + checks.extend(Check(name="capacity:parse", passed=False, detail=e) for e in self.errors) + if not checks: + checks.append( + Check( + name="capacity:audit", + passed=True, + detail=( + f"{len(self.findings)} capacity constant(s)/buffer(s) bound or fixed, " + f"{len(self.tied)} buffer(s) sized from engine arguments, over " + f"{len(self.files)} file(s)" + ), + ) + ) + return checks + + def render(self) -> str: + """The reviewer's table: every capacity, its binding and its verdict.""" + basis = ( + "the declared serving envelope" + if self.envelope_declared + else "a DERIVED envelope (the gate tuples' hull) -- the spec declares none, so a " + "constant that passes here is proven only up to what the gates ran" + ) + lines = [ + "## Capacity constants in the target", + "", + f"Audited statically against {basis}. Each row is a fixed size: a FAIL is " + "unbound or below the envelope; an ok row is a CLAIM -- check that the " + "binding is true of the code (the constant really caps that axis, the " + "`fixed:` reason really holds).", + "", + ] + if not self.findings and not self.errors: + lines.append( + f"No capacity constant or fixed-size buffer found in {len(self.files)} file(s); " + f"{len(self.tied)} allocation(s) are sized from engine arguments." + ) + return "\n".join(lines) + "\n" + lines += [ + "| verdict | kind | name | value | binding | where | detail |", + "|---|---|---|---|---|---|---|", + ] + for f in self.findings: + value = "" if f.value is None else str(f.value) + lines.append( + f"| {'ok' if f.passed else 'FAIL'} | {f.kind} | `{f.name}` | {value} | " + f"{f.binding or '(none)'} | {f.file}:{f.line} | {f.detail} |" + ) + for error in self.errors: + lines.append(f"| FAIL | parse | | | | | {error} |") + return "\n".join(lines) + "\n" + + +# ------------------------------------------------------------------ parsing + + +def _call_name(func: ast.expr) -> str: + if isinstance(func, ast.Attribute): + return func.attr + if isinstance(func, ast.Name): + return func.id + return "" + + +def _references(node: ast.AST) -> tuple[set[str], bool]: + """Constant-like names an expression reads, and whether it reads an engine arg.""" + names: set[str] = set() + engine = False + for child in ast.walk(node): + if isinstance(child, ast.Name): + if child.id in ENGINE_ARGS: + engine = True + elif _CONSTANT_NAME.match(child.id): + names.add(child.id) + elif isinstance(child, ast.Attribute): + if child.attr in ENGINE_ARGS: + engine = True + elif _CONSTANT_NAME.match(child.attr): + # `Class.CONST` or `self.CONST`: keep both spellings. + names.add(child.attr) + if isinstance(child.value, ast.Name): + names.add(f"{child.value.id}.{child.attr}") + return names, engine + + +def _evaluate(node: ast.AST, known: Mapping[str, int]) -> int | None: + """An integer constant expression's value, or None.""" + if isinstance(node, ast.Constant): + value = node.value + return value if isinstance(value, int) and not isinstance(value, bool) else None + if isinstance(node, ast.Name): + return known.get(node.id) + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub): + inner = _evaluate(node.operand, known) + return -inner if inner is not None else None + if isinstance(node, ast.BinOp) and type(node.op) in _BINARY: + left = _evaluate(node.left, known) + right = _evaluate(node.right, known) + if left is None or right is None: + return None + if isinstance(node.op, ast.Pow) and (right < 0 or right > 64): + return None + try: + return int(_BINARY[type(node.op)](left, right)) + except (ArithmeticError, ValueError): + return None + return None + + +def _inline(lines: list[str], node: ast.AST) -> str: + for number in {getattr(node, "lineno", 0), getattr(node, "end_lineno", 0)}: + if 1 <= number <= len(lines): + match = _INLINE.search(lines[number - 1]) + if match: + return match.group("binding") + return "" + + +def _size_arguments(call: ast.Call) -> list[ast.expr]: + """The arguments of an allocation call that give its size.""" + name = _call_name(call.func) + sizes = [kw.value for kw in call.keywords if kw.arg in ("size", "end", "high")] + if name in ("full", "new_full"): + sizes.extend(call.args[:1]) + elif name == "randint": + sizes.extend(call.args[1:3]) + elif name == "arange": + sizes.extend(call.args[:2]) + else: + sizes.extend(call.args) + return sizes + + +class _Collector(ast.NodeVisitor): + def __init__(self, file: str, lines: list[str]) -> None: + self.file = file + self.lines = lines + self.constants: dict[str, Constant] = {} + self.values: dict[str, int] = {} + self.allocations: list[Allocation] = [] + self.ceilings: set[str] = set() + self.bounds: dict[str, str] = {} + self.bounds_line = 0 + self.errors: list[str] = [] + self._scope: list[str] = [] + + # Constants live at module scope and directly in class bodies. + def _record_constant(self, target: ast.expr, value: ast.AST, node: ast.AST) -> None: + if not isinstance(target, ast.Name): + return + if target.id == BOUNDS_NAME and not self._scope: + self._record_bounds(value, node) + return + if not _CONSTANT_NAME.match(target.id): + return + number = _evaluate(value, self.values) + if number is None: + return + owner = ".".join(s for s in self._scope if not s.startswith("def ")) + qualified = f"{owner}.{target.id}" if owner else target.id + self.values[target.id] = number + self.constants[qualified] = Constant( + name=qualified, + value=number, + file=self.file, + line=getattr(node, "lineno", 0), + inline=_inline(self.lines, node), + ) + + def _record_bounds(self, value: ast.AST, node: ast.AST) -> None: + self.bounds_line = getattr(node, "lineno", 0) + try: + parsed = ast.literal_eval(value) + except ValueError: + self.errors.append( + f"{self.file}:{self.bounds_line}: {BOUNDS_NAME} must be a literal mapping of " + f"constant name to binding" + ) + return + if not isinstance(parsed, dict) or not all( + isinstance(k, str) and isinstance(v, str) for k, v in parsed.items() + ): + self.errors.append( + f"{self.file}:{self.bounds_line}: {BOUNDS_NAME} must map constant names to " + f"binding strings" + ) + return + self.bounds.update(parsed) + + def _in_function(self) -> bool: + return any(s.startswith("def ") for s in self._scope) + + def visit_Assign(self, node: ast.Assign) -> None: + if not self._in_function() and len(node.targets) == 1: + self._record_constant(node.targets[0], node.value, node) + self.generic_visit(node) + + def visit_AnnAssign(self, node: ast.AnnAssign) -> None: + if not self._in_function() and node.value is not None: + self._record_constant(node.target, node.value, node) + self.generic_visit(node) + + def visit_ClassDef(self, node: ast.ClassDef) -> None: + self._scope.append(node.name) + self.generic_visit(node) + self._scope.pop() + + def visit_FunctionDef(self, node: ast.FunctionDef) -> None: + self._scope.append(f"def {node.name}") + self.generic_visit(node) + self._scope.pop() + + visit_AsyncFunctionDef = visit_FunctionDef # type: ignore[assignment] + + def visit_Compare(self, node: ast.Compare) -> None: + for operand in (node.left, *node.comparators): + names, _ = _references(operand) + self.ceilings.update(names) + self.generic_visit(node) + + def visit_Call(self, node: ast.Call) -> None: + name = _call_name(node.func) + if name in ("min", "max", "clamp", "clamp_max", "clip"): + for arg in node.args: + names, _ = _references(arg) + self.ceilings.update(names) + if name in ALLOC_FUNCS: + self._record_allocation(node, name) + self.generic_visit(node) + + def _record_allocation(self, node: ast.Call, name: str) -> None: + names: set[str] = set() + tied = False + literal: int | None = None + for size in _size_arguments(node): + found, engine = _references(size) + names.update(found) + tied = tied or engine + # A literal expression (`1 << 20`) as a whole, then every literal + # inside a shape tuple (`(8192, 4)`). + for child in ast.walk(size): + value = _evaluate(child, {}) + if value is not None and value >= LITERAL_THRESHOLD: + literal = max(literal or 0, value) + if not names and literal is None and not tied: + return + functions = [s[4:] for s in self._scope if s.startswith("def ")] + scope = "import" if not functions else functions[-1] + self.allocations.append( + Allocation( + call=name, + file=self.file, + line=node.lineno, + scope=scope, + constants=tuple(sorted(names)), + literal=literal, + tied=tied, + inline=_inline(self.lines, node), + ) + ) + + +def _is_capacity_name(name: str) -> bool: + words = name.split(".")[-1].strip("_").split("_") + return any(word in CAPACITY_WORDS for word in words) + + +def target_files(target_dir: Path) -> list[Path]: + """Every Python file of a target, tests and caches excluded.""" + return sorted( + path + for path in Path(target_dir).rglob("*.py") + if "__pycache__" not in path.parts and not path.name.startswith("test_") + ) + + +def audit_target(target_dir: Path | str, envelope: ServingEnvelope) -> CapacityAudit: + """Audit every capacity constant and fixed-size buffer under ``target_dir``.""" + root = Path(target_dir) + findings: list[CapacityFinding] = [] + tied: list[Allocation] = [] + errors: list[str] = [] + files = target_files(root) if root.is_dir() else [] + for path in files: + rel = str(path.relative_to(root)) + source = path.read_text(encoding="utf-8") + try: + tree = ast.parse(source) + except SyntaxError as exc: + errors.append(f"{rel}: cannot parse: {exc}") + continue + collector = _Collector(rel, source.splitlines()) + collector.visit(tree) + errors.extend(collector.errors) + findings.extend(_judge_file(collector, envelope)) + tied.extend( + a for a in collector.allocations if a.tied and not a.constants and not a.literal + ) + return CapacityAudit( + findings=tuple(findings), + files=tuple(str(p.relative_to(root)) for p in files), + tied=tuple(tied), + envelope_declared=envelope.declared, + errors=tuple(errors), + ) + + +def _lookup(collector: _Collector, name: str) -> Constant | None: + if name in collector.constants: + return collector.constants[name] + matches = [c for key, c in collector.constants.items() if key.split(".")[-1] == name] + return matches[0] if len(matches) == 1 else None + + +def _judge_file(collector: _Collector, envelope: ServingEnvelope) -> list[CapacityFinding]: + sized = {name for alloc in collector.allocations for name in alloc.constants} + relevant: dict[str, Constant] = {} + for qualified, constant in collector.constants.items(): + short = qualified.split(".")[-1] + # A small constant that sizes a buffer is a dimension (a head width, a + # tile), not a capacity; it still counts when its name or its use as + # a ceiling says otherwise. + sizes_large = constant.value >= LITERAL_THRESHOLD and (short in sized or qualified in sized) + if ( + _is_capacity_name(qualified) + or sizes_large + or short in collector.ceilings + or qualified in collector.ceilings + or qualified in collector.bounds + or short in collector.bounds + or constant.inline + ): + relevant[qualified] = constant + + findings: list[CapacityFinding] = [] + for name in sorted(collector.bounds): + if _lookup(collector, name) is None: + findings.append( + CapacityFinding( + kind="binding", + name=name, + file=collector.file, + line=collector.bounds_line, + value=None, + status=STATUS_FAIL, + detail=f"{BOUNDS_NAME} binds {name!r}, which this file does not define", + binding=collector.bounds[name], + ) + ) + + for qualified, constant in sorted(relevant.items()): + short = qualified.split(".")[-1] + declared = constant.inline or collector.bounds.get(qualified) or collector.bounds.get(short) + uses = [] + if short in sized or qualified in sized: + uses.append("sizes a buffer") + if short in collector.ceilings or qualified in collector.ceilings: + uses.append("used as a ceiling") + findings.append( + _judge( + kind="constant", + name=qualified, + file=collector.file, + line=constant.line, + value=constant.value, + declared=declared, + uses=uses, + envelope=envelope, + ) + ) + + for alloc in collector.allocations: + if alloc.literal is None or alloc.tied: + # No large literal: any constant in the size carries the verdict. + # Sized from an engine argument: tied by construction. + continue + findings.append( + _judge( + kind="buffer", + name=f"{alloc.call}@{alloc.file}:{alloc.line}", + file=alloc.file, + line=alloc.line, + value=alloc.literal, + declared=alloc.inline, + uses=[f"literal size at {alloc.scope}"], + envelope=envelope, + ) + ) + return findings + + +def _judge( + *, + kind: str, + name: str, + file: str, + line: int, + value: int | None, + declared: str, + uses: Iterable[str], + envelope: ServingEnvelope, +) -> CapacityFinding: + context = f" ({', '.join(uses)})" if uses else "" + if not declared: + return CapacityFinding( + kind=kind, + name=name, + file=file, + line=line, + value=value, + status=STATUS_FAIL, + detail=( + f"unbound capacity{context}: derive it from an engine argument at " + f"construction, or declare `# capacity: ` / `# capacity: fixed: ` " + f"so it can be checked against the envelope" + ), + ) + binding = parse_binding(declared) + if binding.error: + return CapacityFinding( + kind=kind, + name=name, + file=file, + line=line, + value=value, + status=STATUS_FAIL, + detail=binding.error, + binding=binding.raw, + ) + if binding.fixed: + return CapacityFinding( + kind=kind, + name=name, + file=file, + line=line, + value=value, + status=STATUS_OK, + detail=f"not a capacity: {binding.reason}{context}", + binding=binding.raw, + ) + required = binding.required(envelope) + if required is None: + return CapacityFinding( + kind=kind, + name=name, + file=file, + line=line, + value=value, + status=STATUS_FAIL, + detail=( + f"bound to {binding.raw}, but the envelope states no {binding.axis}, so the " + f"bound cannot be proven" + ), + binding=binding.raw, + ) + if value is None or value < required: + return CapacityFinding( + kind=kind, + name=name, + file=file, + line=line, + value=value, + status=STATUS_FAIL, + detail=( + f"{value} < {binding.raw} = {required} at the envelope{context}: every engine " + f"configured above {value} exceeds this" + ), + binding=binding.raw, + ) + return CapacityFinding( + kind=kind, + name=name, + file=file, + line=line, + value=value, + status=STATUS_OK, + detail=f"{value} >= {binding.raw} = {required}{context}", + binding=binding.raw, + ) + + +__all__ = [ + "ALLOC_FUNCS", + "BINDABLE_AXES", + "BOUNDS_NAME", + "CAPACITY_WORDS", + "ENGINE_ARGS", + "LITERAL_THRESHOLD", + "Allocation", + "Binding", + "CapacityAudit", + "CapacityFinding", + "Constant", + "audit_target", + "parse_binding", + "target_files", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/early_checks.py b/agent-flow/agent_flow/workflows/staircase_onboard/early_checks.py new file mode 100644 index 000000000000..9e9df0406a2e --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/early_checks.py @@ -0,0 +1,608 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Delivery checks that need no finished target, run early, as warnings. + +A run whose every unit is green can still fail delivery on things no unit gate +looks at: a catalog test under a name the CI selector does not match, no +integration accuracy test for the model, no test-db line that runs it. Found +at delivery, each costs a turn and -- on the run this was written for -- hours +of an idle run. None of them needs the target to be finished, so this module +runs them at the end of the plan phase and after every unit, and reports them +as **warnings**: a run is never stopped here, it is told early. + +It also answers a question delivery never asked: **will CI actually run the +accuracy test, or skip it?** The test is read statically for what it needs -- +a checkpoint under ``LLM_MODELS_ROOT``, environment variables, a GPU count, an +architecture -- and compared with the test-db stages that list it (pre- or +post-merge, which GPUs, how many). What it needs is emitted as the +**required CI assets** list, because a test that skips on every CI machine +reads green forever. +""" + +from __future__ import annotations + +import ast +import os +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + +import yaml + +from .delivery import ( + ACCURACY_TESTS_DIR, + CATALOG_DIR, + CATALOG_TESTS_DIR, + MIRROR_CATEGORY, + TEST_DB_DIR, +) + +LEVEL_OK = "ok" +LEVEL_WARN = "warn" +LEVEL_INFO = "info" + +MODELS_ROOT_ENV = "LLM_MODELS_ROOT" +_MODELS_ROOT_CALL = "llm_models_root" + +#: Which GPU names in a test-db wildcard serve an architecture. Only used to +#: say "this stage can never satisfy the test's architecture skip". +ARCH_GPUS: Mapping[str, tuple[str, ...]] = { + "sm_80": ("a100", "a30", "a10"), + "sm_89": ("l40", "l4"), + "sm_90": ("h100", "h200", "gh200", "h20"), + "sm_100": ("b200", "gb200"), + "sm_103": ("b300", "gb300"), + "sm_120": ("rtx", "gb202", "gb203"), + "sm_121": ("gb10",), +} +_ARCH_FROM_SKIP = {"blackwell": "sm_100", "hopper": "sm_90", "ada": "sm_89"} + + +@dataclass(frozen=True) +class EarlyFinding: + name: str + level: str + detail: str + + def render(self) -> str: + return f"[{self.level}] {self.name}: {self.detail}" + + +@dataclass(frozen=True) +class CiAsset: + """One thing CI must provide, or the accuracy test skips.""" + + kind: str + what: str + source: str + note: str = "" + + def render(self) -> str: + note = f" -- {self.note}" if self.note else "" + return f"- **{self.kind}**: {self.what} (`{self.source}`){note}" + + +@dataclass(frozen=True) +class CiListing: + """One test-db block that lists the accuracy test.""" + + file: str + stage: str + gpus: tuple[str, ...] + gpu_count: tuple[int | None, int | None] + + +@dataclass(frozen=True) +class EarlyReport: + stage: str + findings: tuple[EarlyFinding, ...] = () + assets: tuple[CiAsset, ...] = () + listings: tuple[CiListing, ...] = field(default_factory=tuple) + + @property + def warnings(self) -> tuple[EarlyFinding, ...]: + return tuple(f for f in self.findings if f.level == LEVEL_WARN) + + def render(self) -> str: + lines = [ + f"# Early delivery checks ({self.stage})", + "", + "Nothing here stops the run. A `warn` line is something delivery fails on, or " + "something CI will skip without saying so; fix it now rather than at the end.", + "", + *(f"- {finding.render()}" for finding in self.findings), + "", + "## Required CI assets", + "", + "What CI must provide for the accuracy test to RUN rather than skip:", + "", + ] + lines += [asset.render() for asset in self.assets] or ["- (no accuracy test to read yet)"] + if self.listings: + lines += ["", "## Where CI runs it", ""] + for listing in self.listings: + low, high = listing.gpu_count + lines.append( + f"- `{listing.file}`: stage {listing.stage or '?'}, GPUs " + f"{', '.join(listing.gpus) or '*'}, count {low or '?'}..{high or '?'}" + ) + return "\n".join(lines) + "\n" + + +# ------------------------------------------------------------ catalog tests + + +def vocabulary_entries(vocabulary_path: Path | str, catalog_root: Path) -> list[tuple[str, str]]: + """``(category, name)`` for every entry the plan names, resolved against the catalog.""" + path = Path(vocabulary_path) + if not path.is_file(): + return [] + try: + data = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + except yaml.YAMLError: + return [] + if not isinstance(data, dict): + return [] + named: list[str] = [] + for step in data.get("steps") or (): + if isinstance(step, dict) and step.get("entry"): + named.append(str(step["entry"])) + for row in data.get("queue") or (): + if isinstance(row, dict) and row.get("target"): + named.append(str(row["target"])) + entries: set[tuple[str, str]] = set() + for raw in named: + text = raw.strip().removesuffix(".py") + if "catalog/" in text: + text = text.split("catalog/", 1)[1] + if "/" in text: + category, name = text.rsplit("/", 1) + entries.add((category.rsplit("/", 1)[-1], name)) + continue + # A bare name: find its category in the catalog, if it exists yet. + matches = sorted(catalog_root.glob(f"*/{text}.py")) if catalog_root.is_dir() else [] + if len(matches) == 1: + entries.add((matches[0].parent.name, text)) + return sorted(entries) + + +def check_test_naming(repo_root: Path, entries: Iterable[tuple[str, str]]) -> list[EarlyFinding]: + """Every planned entry's GPU test is (or will be) under the name CI matches.""" + findings: list[EarlyFinding] = [] + tests_root = repo_root / CATALOG_TESTS_DIR + for category, name in entries: + if category == MIRROR_CATEGORY: + continue + expected = tests_root / category / f"test_modeling_v2_{name}.py" + if expected.is_file(): + continue + others = ( + sorted( + p.name + for p in (tests_root / category).glob(f"test_*{name}*.py") + if p.name != expected.name + ) + if (tests_root / category).is_dir() + else [] + ) + if others: + findings.append( + EarlyFinding( + name=f"test-name:{category}/{name}", + level=LEVEL_WARN, + detail=( + f"tested as {', '.join(others)}; delivery and the CI selector expect " + f"{CATALOG_TESTS_DIR}/{category}/{expected.name}" + ), + ) + ) + elif (repo_root / CATALOG_DIR / category / f"{name}.py").is_file(): + findings.append( + EarlyFinding( + name=f"test-name:{category}/{name}", + level=LEVEL_WARN, + detail=f"the entry exists and has no test at {CATALOG_TESTS_DIR}/{category}/{expected.name}", + ) + ) + return findings + + +# ------------------------------------------------------------ accuracy test + + +def accuracy_test_path(repo_root: Path, family: str) -> Path: + return repo_root / ACCURACY_TESTS_DIR / f"test_modeling_v2_{family}.py" + + +def find_test_db_listings(repo_root: Path, stem: str) -> list[CiListing]: + """Every test-db block whose tests name ``stem``.""" + listings: list[CiListing] = [] + db = repo_root / TEST_DB_DIR + if not db.is_dir(): + return listings + for path in sorted(db.glob("*.yml")): + text = path.read_text(encoding="utf-8") + if stem not in text: + continue + try: + data = yaml.safe_load(text) or {} + except yaml.YAMLError: + listings.append( + CiListing(file=path.name, stage="unparsable", gpus=(), gpu_count=(None, None)) + ) + continue + # `{l0_x: [block, ...]}` is the house shape; a bare list of blocks is + # read the same way. + groups = list(data.values()) if isinstance(data, dict) else [data] + for blocks in groups: + for block in blocks if isinstance(blocks, list) else (): + if not isinstance(block, dict): + continue + tests = [str(t) for t in block.get("tests") or ()] + if not any(stem in t for t in tests): + continue + condition = block.get("condition") or {} + terms = condition.get("terms") or {} + wildcards = condition.get("wildcards") or {} + count = (condition.get("ranges") or {}).get("system_gpu_count") or {} + listings.append( + CiListing( + file=path.name, + stage=str(terms.get("stage", "")), + gpus=tuple(str(g) for g in wildcards.get("gpu") or ()), + gpu_count=(count.get("gte"), count.get("lte")), + ) + ) + return listings + + +def _models_root_suffix(node: ast.AST) -> str | None: + """The path a checkpoint expression appends to ``llm_models_root()``, if it is one.""" + + def is_root(expr: ast.AST) -> bool: + return isinstance(expr, ast.Call) and _call_name(expr.func) == _MODELS_ROOT_CALL + + if isinstance(node, ast.JoinedStr): + parts = node.values + for index, part in enumerate(parts): + if isinstance(part, ast.FormattedValue) and is_root(part.value): + rest = "".join(p.value for p in parts[index + 1 :] if isinstance(p, ast.Constant)) + return rest.lstrip("/") + if isinstance(node, ast.Call) and _call_name(node.func) == "join" and node.args: + if is_root(node.args[0]): + return "/".join(str(a.value) for a in node.args[1:] if isinstance(a, ast.Constant)) + if isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Div, ast.Add)): + left = node.left + if isinstance(left, ast.Call) and _call_name(left.func) in ("Path",) and left.args: + left = left.args[0] + if is_root(left) and isinstance(node.right, ast.Constant): + return str(node.right.value).lstrip("/") + return None + + +def _call_name(func: ast.AST) -> str: + if isinstance(func, ast.Attribute): + return func.attr + if isinstance(func, ast.Name): + return func.id + return "" + + +def _sm_requirement(node: ast.AST) -> str | None: + """``get_sm_version() != 103`` (a skip condition) -> ``sm_103``.""" + for child in ast.walk(node): + if not isinstance(child, ast.Compare) or len(child.comparators) != 1: + continue + left, right = child.left, child.comparators[0] + if isinstance(left, ast.Call) and _call_name(left.func) == "get_sm_version": + if isinstance(right, ast.Constant) and isinstance(right.value, int): + if isinstance(child.ops[0], ast.NotEq): + return f"sm_{right.value}" + if isinstance(child.ops[0], ast.Lt): + return f"sm_{right.value}+" + return None + + +def analyze_skips(test_file: Path) -> list[CiAsset]: + """What the accuracy test needs from CI, read without importing it.""" + rel = str(test_file) + try: + source = test_file.read_text(encoding="utf-8") + tree = ast.parse(source) + except (OSError, SyntaxError) as exc: + return [CiAsset(kind="unreadable", what=str(exc), source=rel)] + assets: list[CiAsset] = [] + seen: set[tuple[str, str]] = set() + + def add(kind: str, what: str, line: int, note: str = "") -> None: + if (kind, what) in seen: + return + seen.add((kind, what)) + assets.append(CiAsset(kind=kind, what=what, source=f"{test_file.name}:{line}", note=note)) + + for node in ast.walk(tree): + suffix = ( + _models_root_suffix(node) + if isinstance(node, (ast.JoinedStr, ast.Call, ast.BinOp)) + else None + ) + if suffix is not None: + add( + "checkpoint", + f"${MODELS_ROOT_ENV}/{suffix}", + node.lineno, + "the test's model path; absent on a CI machine, the load fails or the case skips", + ) + if isinstance(node, ast.Call): + name = _call_name(node.func) + if name == "skipif" and node.args: + reason = next( + ( + kw.value.value + for kw in node.keywords + if kw.arg == "reason" and isinstance(kw.value, ast.Constant) + ), + "", + ) + arch = _sm_requirement(node.args[0]) + if arch: + add("arch", arch, node.lineno, str(reason)) + else: + add("skip", ast.unparse(node.args[0]), node.lineno, str(reason)) + elif name in ("skip_less_device", "skip_less_mpi_world_size") and node.args: + count = node.args[0] + if isinstance(count, ast.Constant): + add("gpus", f">= {count.value}", node.lineno) + elif name == "skip_less_device_memory" and node.args: + memory = node.args[0] + if isinstance(memory, ast.Constant): + add("gpu memory", f">= {memory.value} MiB", node.lineno) + elif name == "skip" and isinstance(node.func, ast.Attribute): + reason = ( + node.args[0].value + if node.args and isinstance(node.args[0], ast.Constant) + else "" + ) + add("skip", "pytest.skip(...) at runtime", node.lineno, str(reason)) + elif name in ("getenv", "get") and node.args and isinstance(node.args[0], ast.Constant): + owner = node.func.value if isinstance(node.func, ast.Attribute) else None + if name == "getenv" or ( + isinstance(owner, ast.Attribute) and owner.attr == "environ" + ): + add("env", str(node.args[0].value), node.lineno) + elif isinstance(node, ast.Subscript): + if ( + isinstance(node.value, ast.Attribute) + and node.value.attr == "environ" + and isinstance(node.slice, ast.Constant) + ): + add("env", str(node.slice.value), node.lineno) + elif isinstance(node, ast.Name) and node.id.startswith( + ("skip_pre_", "skip_no_", "skip_post_") + ): + if isinstance(node.ctx, ast.Load): + generation = node.id.rsplit("_", 1)[-1] + arch = _ARCH_FROM_SKIP.get(generation) + if node.id.startswith("skip_pre_") and arch: + add("arch", f"{arch}+", node.lineno, node.id) + else: + add("skip", node.id, node.lineno) + return assets + + +def _arch_served(arch: str, gpus: Sequence[str]) -> bool | None: + """Whether a listing's GPU wildcards can satisfy an arch requirement; None = unknown.""" + base = arch.rstrip("+") + names = ARCH_GPUS.get(base) + if not names or not gpus: + return None + if arch.endswith("+"): + # At least this arch: anything listed at or above it. + newer = [a for a in ARCH_GPUS if _arch_number(a) >= _arch_number(base)] + names = tuple(n for a in newer for n in ARCH_GPUS[a]) + return any(any(name in g.lower() for name in names) for g in gpus) + + +def _arch_number(arch: str) -> int: + digits = re.sub(r"\D", "", arch) + return int(digits) if digits else 0 + + +def check_ci_skip( + repo_root: Path, + family: str, + *, + env: Mapping[str, str] | None = None, +) -> tuple[list[EarlyFinding], list[CiAsset], list[CiListing]]: + """Will CI run the accuracy test? Findings, the required assets, the listings.""" + test = accuracy_test_path(repo_root, family) + findings: list[EarlyFinding] = [] + if not test.is_file(): + findings.append( + EarlyFinding( + name="tests:accuracy-file", + level=LEVEL_WARN, + detail=( + f"no {ACCURACY_TESTS_DIR}/{test.name}: delivery requires the model's " + f"integration accuracy test; write it with the unit, not after the last one" + ), + ) + ) + return findings, [], [] + findings.append(EarlyFinding(name="tests:accuracy-file", level=LEVEL_OK, detail=str(test))) + + listings = find_test_db_listings(repo_root, test.stem) + if not listings: + findings.append( + EarlyFinding( + name="tests:test-db", + level=LEVEL_WARN, + detail=f"{test.stem} appears in no {TEST_DB_DIR}/*.yml; CI would never run it", + ) + ) + else: + stages = {listing.stage for listing in listings} + level = LEVEL_OK + detail = ", ".join(sorted({f"{item.file} ({item.stage or '?'})" for item in listings})) + if stages and stages <= {"post_merge"}: + level = LEVEL_WARN + detail += "; post-merge only, so no pre-merge run ever executes it" + findings.append(EarlyFinding(name="tests:test-db", level=level, detail=detail)) + + assets = analyze_skips(test) + for listing in listings: + assets.append( + CiAsset( + kind="stage", + what=f"{listing.file}: {listing.stage or '?'}", + source=f"{TEST_DB_DIR}/{listing.file}", + note=f"GPUs {', '.join(listing.gpus) or '*'}", + ) + ) + + environment = dict(os.environ) + environment.update({str(k): str(v) for k, v in (env or {}).items()}) + models_root = environment.get(MODELS_ROOT_ENV) + for asset in assets: + if asset.kind == "checkpoint": + if models_root: + local = Path(models_root) / asset.what.split("/", 1)[1] + level = LEVEL_OK if local.exists() else LEVEL_WARN + detail = ( + f"{local} exists here" + if local.exists() + else f"{local} does not exist under this host's {MODELS_ROOT_ENV}; " + f"CI's model store must hold it or every case skips or fails" + ) + else: + level = LEVEL_INFO + detail = ( + f"{MODELS_ROOT_ENV} is not set here, so whether CI's model store holds " + f"{asset.what} cannot be checked -- confirm it before merge" + ) + findings.append(EarlyFinding(name="ci:checkpoint", level=level, detail=detail)) + elif asset.kind == "arch": + for listing in listings: + served = _arch_served(asset.what, listing.gpus) + if served is False: + findings.append( + EarlyFinding( + name="ci:arch", + level=LEVEL_WARN, + detail=( + f"the test needs {asset.what} ({asset.source}) and " + f"{listing.file} lists GPUs {', '.join(listing.gpus)}: it " + f"always skips there" + ), + ) + ) + elif asset.kind == "gpus": + wanted = int(re.sub(r"\D", "", asset.what) or 0) + for listing in listings: + high = listing.gpu_count[1] + if isinstance(high, int) and high < wanted: + findings.append( + EarlyFinding( + name="ci:gpus", + level=LEVEL_WARN, + detail=( + f"the test needs {asset.what} GPUs and {listing.file} runs " + f"on at most {high}: it always skips there" + ), + ) + ) + elif asset.kind == "env" and asset.what not in environment: + findings.append( + EarlyFinding( + name=f"ci:env:{asset.what}", + level=LEVEL_WARN, + detail=( + f"the test reads ${asset.what} ({asset.source}); it is unset here, " + f"and CI must provide it" + ), + ) + ) + return findings, assets, listings + + +# -------------------------------------------------------------------- run + + +def run_early_checks( + *, + repo_root: Path | str, + family: str, + stage: str, + vocabulary_path: Path | str | None = None, + env: Mapping[str, str] | None = None, + capacity: Any = None, + certification: Any = None, + workspace: Path | str | None = None, +) -> EarlyReport: + """Every delivery check that needs no finished target, as warnings. + + ``capacity`` is a :class:`~.capacity_audit.CapacityAudit` and + ``certification`` a :class:`~.envelope.CertificationPlan`; when given, + their delivery failures are reported here early as well. + """ + root = Path(repo_root) + findings: list[EarlyFinding] = [] + entries = vocabulary_entries(vocabulary_path, root / CATALOG_DIR) if vocabulary_path else [] + findings.extend(check_test_naming(root, entries)) + + ci_findings, assets, listings = check_ci_skip(root, family, env=env) + findings.extend(ci_findings) + + if capacity is not None: + for finding in capacity.failures: + findings.append( + EarlyFinding( + name=f"capacity:{finding.name}", + level=LEVEL_WARN, + detail=f"{finding.file}:{finding.line}: {finding.detail}", + ) + ) + if certification is not None and workspace is not None: + from .envelope import check_external_evidence + + for check in check_external_evidence(certification, workspace): + if not check.passed: + findings.append( + EarlyFinding(name=check.name, level=LEVEL_WARN, detail=check.detail) + ) + if vocabulary_path: + from .smith_brief import load_steps, parse_call_sites + + for step in load_steps(vocabulary_path): + for error in parse_call_sites(step)[1]: + findings.append( + EarlyFinding(name="vocabulary:call_sites", level=LEVEL_WARN, detail=error) + ) + return EarlyReport( + stage=stage, + findings=tuple(findings), + assets=tuple(assets), + listings=tuple(listings), + ) + + +__all__ = [ + "ARCH_GPUS", + "CiAsset", + "EarlyFinding", + "EarlyReport", + "LEVEL_INFO", + "LEVEL_OK", + "LEVEL_WARN", + "MODELS_ROOT_ENV", + "CiListing", + "accuracy_test_path", + "analyze_skips", + "check_ci_skip", + "check_test_naming", + "run_early_checks", + "find_test_db_listings", + "vocabulary_entries", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/envelope.py b/agent-flow/agent_flow/workflows/staircase_onboard/envelope.py new file mode 100644 index 000000000000..b1c1268430e0 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/envelope.py @@ -0,0 +1,1788 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""What a target is certified FOR: the serving envelope and the feature matrix. + +A gate launches an engine at one configuration tuple -- a batch limit, a +token budget per step, a sequence length, a CUDA-graph capture list -- and +measures that tuple. The tempting reading of a green run is that the tuples +the gates launched ARE the certified set. That reading hides every defect that +only appears outside them: a run whose gates all used ``max_num_tokens`` 8192 +certifies nothing about 16384, and a target-side fixed-size buffer sized 8192, +or a kernel that overflows a 32-bit offset past 2**25 rows, passes every gate +and crashes the first deployment that asks for more. + +So the spec declares two things up front, and this module derives everything +else from them: + +``serving_envelope`` + What the target is claimed to serve: ``max_batch_size``, + ``max_num_tokens``, ``max_seq_len``, the CUDA-graph policy (an explicit + batch list, or the padding policy the engine generates one from) and the + MEMBERS -- every parallel configuration claimed, including multi-node and + pipeline-parallel ones a single allocation cannot run. Catalog receipts, + capacity constants and integer-width limits are judged against this box, + never against the gate tuples. + +``feature_matrix`` + The feature combinations that must have ENGINE evidence: which features + on, at which members, which kind of evidence (boot, accuracy, witness, + acceptance), at which point of the envelope (the gate tuple, the ceiling, + or a named point). + +From those, :func:`derive_certification` produces a +:class:`CertificationPlan`: + +- the **gate tuples** the run's existing gates launch, each checked to lie + inside the envelope (a gate measuring outside the envelope measures a + configuration nobody claimed); +- the **envelope probes** the run must launch in addition -- a boot at the + envelope's ceiling, per unit, so that warm-up at the real ``max_num_tokens`` + and capture at every claimed graph batch happen in the run, not in the first + deployment; +- the **matrix cells**, each classified as gated (an existing gate covers + it), probe (a derived probe covers it) or external (this allocation cannot + produce it, so delivery requires a recorded evidence file); +- the **certified set** every smith, reviewer and auditor certifies against. + +A spec with no ``serving_envelope`` is still accepted, as the migration path: +the envelope is then DERIVED from the gate tuples (the old semantics, stated +as such everywhere it is rendered) and ``python -m +agent_flow.workflows.staircase_onboard.envelope migrate `` prints a +declared block to start from. +""" + +from __future__ import annotations + +import argparse +import copy +import math +import re +import sys +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Callable, Iterable, Mapping, Sequence + +import yaml + +from .task_schema import ( + ACCURACY_FIELD, + FEATURE_MATRIX_FIELD, + FEATURES_FIELD, + MAX_WITNESS_PROMPT_WORDS, + SERVING_ENVELOPE_FIELD, + TaskSchemaError, + assembly_units, + members_for, + parse_parallel_segment, + reject_unknown_keys, +) + +# ------------------------------------------------------------------- schema + +#: The three engine limits an envelope bounds, by their engine-argument names. +#: Spelled as the engine spells them on purpose: a capacity constant is tied +#: to one of these by name (see :mod:`.capacity_audit`), and a probe passes +#: them as ``--=``. +ENVELOPE_AXES: tuple[str, ...] = ("max_batch_size", "max_num_tokens", "max_seq_len") + +ENVELOPE_KEYS: tuple[str, ...] = ( + *ENVELOPE_AXES, + "cuda_graph", + "members", + "max_draft_len", + "gpus_per_node", + "probe_prompt_tokens", +) +ENVELOPE_REQUIRED: tuple[str, ...] = (*ENVELOPE_AXES, "cuda_graph", "members") +CUDA_GRAPH_KEYS: tuple[str, ...] = ("batch_sizes", "padding", "max_batch_size") + +MATRIX_CELL_KEYS: tuple[str, ...] = ( + "name", + "features", + "members", + "evidence", + "at", + "prompt_tokens", + "note", +) +EVIDENCE_BOOT = "boot" +EVIDENCE_ACCURACY = "accuracy" +EVIDENCE_WITNESS = "witness" +EVIDENCE_ACCEPTANCE = "acceptance" +EVIDENCE_KINDS: tuple[str, ...] = ( + EVIDENCE_BOOT, + EVIDENCE_ACCURACY, + EVIDENCE_WITNESS, + EVIDENCE_ACCEPTANCE, +) +AT_GATE = "gate" +AT_CEILING = "ceiling" + +#: How a matrix cell gets its evidence. +CELL_GATED = "gated" +CELL_PROBE = "probe" +CELL_EXTERNAL = "external" + +#: Where an external cell's evidence is recorded, under the workspace. +EVIDENCE_DIRNAME = "evidence" + +#: The name a derived probe gate is recorded under. +PROBE_GATE = "envelope" + +#: ``tp4``, ``tep4``, ``dep8``, ``tep4pp2``. The gated member grammar +#: (:data:`~.task_schema.PARALLEL_PATTERN`) has no pipeline axis because one +#: allocation serves every gated member; an envelope member is a CLAIM and may +#: name a configuration this allocation cannot run. +ENVELOPE_MEMBER_PATTERN = re.compile(r"^(tp|tep|dep)(\d+)(?:pp(\d+))?$") + +#: Engine defaults where a gate command does not state a limit. The same for +#: ``quickstart_advanced.py`` and ``trtllm-eval``: both take them from the +#: LLM-API's own field defaults. ``max_seq_len`` defaults to the checkpoint's +#: ``max_position_embeddings``, which is why it is ``None`` here. +ENGINE_DEFAULTS: Mapping[str, int | None] = { + "max_batch_size": 2048, + "max_num_tokens": 8192, + "max_seq_len": None, +} + +#: CudaGraphConfig's own default bound when neither a list nor a maximum is +#: given. +DEFAULT_GRAPH_MAX_BATCH_SIZE = 128 + +_FLAG_AXES = {f"--{axis}": axis for axis in ENVELOPE_AXES} +_GRAPH_LIST_FLAG = "--cuda_graph_batch_sizes" +_GRAPH_PADDING_FLAG = "--cuda_graph_padding_enabled" + + +def _positive_int(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value >= 1 + + +def _non_negative_int(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value >= 0 + + +# ------------------------------------------------------------------ members + + +@dataclass(frozen=True) +class EnvelopeMember: + """One parallel configuration the target is claimed to serve. + + ``size`` is the width of one pipeline stage (the trailing number of the + gated grammar); ``pp`` multiplies it. ``tep4pp2`` is eight devices, two + stages of four. + """ + + name: str + kind: str + size: int + pp: int = 1 + + @property + def world_size(self) -> int: + return self.size * self.pp + + @property + def attention_dp(self) -> bool: + """Whether ops after the attention-DP gather see every rank's rows.""" + return self.kind == "dep" + + def nodes(self, gpus_per_node: int | None) -> int: + if not gpus_per_node: + return 1 + return max(1, math.ceil(self.world_size / gpus_per_node)) + + def describe(self, gpus_per_node: int | None) -> str: + parts = [f"{self.world_size} GPU(s)"] + if self.pp > 1: + parts.insert(0, f"PP{self.pp}") + nodes = self.nodes(gpus_per_node) + if nodes > 1: + parts.append(f"{nodes} nodes") + return f"{self.name} ({', '.join(parts)})" + + +def parse_envelope_member(text: str) -> EnvelopeMember: + """Parse ``tep4`` / ``dep8`` / ``tp4pp2``. Raises :class:`ValueError`.""" + match = ENVELOPE_MEMBER_PATTERN.match(str(text)) + if match is None: + raise ValueError( + f"{text!r} is not an envelope member; expected tp, tep or dep, " + f"optionally followed by pp (e.g. 'tep4pp2')" + ) + size = int(match.group(2)) + pp = int(match.group(3) or 1) + if size < 1 or pp < 1: + raise ValueError(f"{text!r} declares a zero-width member") + return EnvelopeMember(name=str(text), kind=match.group(1), size=size, pp=pp) + + +# -------------------------------------------------------------- graph policy + + +def generate_graph_batch_sizes(max_batch_size: int, padding: bool) -> list[int]: + """The capture list the engine generates from a bound and a padding policy. + + Mirrors ``CudaGraphConfig._generate_cuda_graph_batch_sizes`` in + ``tensorrt_llm/llmapi/llm_args.py``: with padding, ``1, 2, 4`` then + multiples of 8 to 128 then steps of 64; without, every size to 31, then + 32, 64, 128 and powers of two. Restated rather than imported so this + module stays importable on a login node with no engine. + """ + if max_batch_size < 1: + return [] + if padding: + sizes = [1, 2, 4] + [i * 8 for i in range(1, 17)] + while sizes[-1] + 64 <= max_batch_size: + sizes.append(sizes[-1] + 64) + else: + sizes = list(range(1, 32)) + [32, 64, 128] + if max_batch_size > 1: + sizes += [2**i for i in range(8, math.ceil(math.log(max_batch_size, 2)))] + sizes = sorted({size for size in sizes if size <= max_batch_size}) + if not sizes or sizes[-1] != max_batch_size: + sizes.append(max_batch_size) + return sizes + + +@dataclass(frozen=True) +class GraphPolicy: + """Which batch sizes are captured. Empty ``batch_sizes`` means eager.""" + + batch_sizes: tuple[int, ...] = () + padding: bool = False + source: str = "eager" + + @property + def enabled(self) -> bool: + return bool(self.batch_sizes) + + def clipped(self, max_batch_size: int | None) -> tuple[int, ...]: + """The sizes an engine with this batch limit actually captures.""" + if max_batch_size is None: + return self.batch_sizes + return tuple(size for size in self.batch_sizes if size <= max_batch_size) + + +def resolve_graph_policy(block: Any, max_batch_size: int | None) -> GraphPolicy: + """Read ``serving_envelope.cuda_graph`` (already validated).""" + if block is False or block is None: + return GraphPolicy() + block = dict(block) + explicit = block.get("batch_sizes") + padding = bool(block.get("padding", False)) + if explicit: + return GraphPolicy( + batch_sizes=tuple(sorted({int(v) for v in explicit})), + padding=padding, + source="explicit list", + ) + bound = int(block.get("max_batch_size") or max_batch_size or DEFAULT_GRAPH_MAX_BATCH_SIZE) + return GraphPolicy( + batch_sizes=tuple(generate_graph_batch_sizes(bound, padding)), + padding=padding, + source=f"{'padding' if padding else 'no-padding'} policy up to {bound}", + ) + + +# ------------------------------------------------------------------ envelope + + +@dataclass(frozen=True) +class ServingEnvelope: + """The box every certification is judged against. + + ``declared`` is False when the spec stated no envelope and this one was + derived from the gate tuples. That is the legacy reading -- the certified + set IS the gate tuples -- and every rendering says so. + """ + + max_batch_size: int | None + max_num_tokens: int | None + max_seq_len: int | None + graph: GraphPolicy = field(default_factory=GraphPolicy) + members: tuple[EnvelopeMember, ...] = () + max_draft_len: int = 0 + gpus_per_node: int | None = None + probe_prompt_tokens: int | None = None + declared: bool = True + + def limit(self, axis: str) -> int | None: + return getattr(self, axis) + + @property + def capture_batches(self) -> tuple[int, ...]: + return self.graph.clipped(self.max_batch_size) + + def member(self, name: str) -> EnvelopeMember | None: + return next((m for m in self.members if m.name == name), None) + + @property + def max_world_size(self) -> int: + return max((m.world_size for m in self.members), default=1) + + @property + def max_gather_world_size(self) -> int: + """The widest attention-DP gather any member performs, or 1.""" + return max((m.world_size for m in self.members if m.attention_dp), default=1) + + +def envelope_from_spec( + block: Mapping[str, Any], + *, + gpus_per_node: int | None = None, +) -> ServingEnvelope: + """Build the declared envelope from a validated ``serving_envelope`` block.""" + max_batch_size = int(block["max_batch_size"]) + return ServingEnvelope( + max_batch_size=max_batch_size, + max_num_tokens=int(block["max_num_tokens"]), + max_seq_len=int(block["max_seq_len"]), + graph=resolve_graph_policy(block.get("cuda_graph"), max_batch_size), + members=tuple(parse_envelope_member(m) for m in block.get("members") or ()), + max_draft_len=int(block.get("max_draft_len") or 0), + gpus_per_node=int(block.get("gpus_per_node") or gpus_per_node or 0) or None, + probe_prompt_tokens=( + int(block["probe_prompt_tokens"]) if block.get("probe_prompt_tokens") else None + ), + declared=True, + ) + + +# ------------------------------------------------------------- engine tuples + + +@dataclass(frozen=True) +class EngineTuple: + """One configuration an engine is launched at, and who launches it.""" + + gate: str + unit: str + member: str + max_batch_size: int | None + max_num_tokens: int | None + max_seq_len: int | None + graph_batch_sizes: tuple[int, ...] = () + source: str = "" + #: A generated prompt of at least this many words, sent beside the boot + #: prompts. Probes only. + prompt_tokens: int | None = None + #: Whether batches are padded up to the next captured size. Probes only: + #: a probe serves the envelope's own padding policy. + graph_padding: bool = False + + def limit(self, axis: str) -> int | None: + return getattr(self, axis) + + def point(self) -> str: + """``B=8 T=2048 L=16384 G=1..8``: the tuple, compactly.""" + + def show(value: int | None) -> str: + return "default" if value is None else str(value) + + graphs = compact_values(self.graph_batch_sizes) if self.graph_batch_sizes else "eager" + text = ( + f"B={show(self.max_batch_size)} T={show(self.max_num_tokens)} " + f"L={show(self.max_seq_len)} G={graphs}" + ) + if self.prompt_tokens: + text += f" prompt>={self.prompt_tokens}" + return text + + def label(self) -> str: + return f"{self.gate} {self.unit}@{self.member} ({self.point()})" + + def outside(self, envelope: ServingEnvelope, *, checkpoint_max: int | None) -> list[str]: + """Every way this tuple exceeds the envelope, as sentences.""" + problems: list[str] = [] + for axis in ENVELOPE_AXES: + bound = envelope.limit(axis) + value = self.limit(axis) + if bound is None: + continue + if value is None and axis == "max_seq_len": + # The engine then sizes for the checkpoint's own maximum. + value = checkpoint_max + if value is None: + continue + if value is None: + continue + if value > bound: + problems.append(f"{axis}={value} > serving_envelope.{axis}={bound}") + return problems + + +def compact_values(values: Sequence[int]) -> str: + """``1,2,3,4,8,16`` -> ``1..4,8,16``: a capture list on one line.""" + ordered = sorted(set(int(v) for v in values)) + runs: list[str] = [] + start = prev = None + for value in ordered: + if start is None: + start = prev = value + continue + if value == prev + 1: + prev = value + continue + runs.append(str(start) if start == prev else f"{start}..{prev}") + start = prev = value + if start is not None: + runs.append(str(start) if start == prev else f"{start}..{prev}") + return ",".join(runs) + + +def _flag_values(flags: Iterable[str]) -> dict[str, list[str]]: + """``{"--max_batch_size": ["8"], ...}`` from runtime ``extra_flags``. + + Three spellings survive in real specs -- ``--flag=v``, ``--flag v`` as one + string, and the flag and its value(s) as separate elements -- and a + ``nargs='+'`` flag (the graph list) takes every following bare value. + """ + words: list[str] = [] + for flag in flags: + words.extend(str(flag).split()) + found: dict[str, list[str]] = {} + current: str | None = None + for word in words: + if word.startswith("--"): + name, _, value = word.partition("=") + found[name] = [value] if value else [] + current = None if value else name + continue + if current is not None: + found[current].append(word) + return found + + +def _int_or_none(values: Sequence[str] | None) -> int | None: + if not values: + return None + try: + return int(values[-1]) + except ValueError: + return None + + +def runtime_tuple( + runtime: Mapping[str, Any], + *, + gate: str, + unit: str, + member: str, + max_batch_size: int | None = None, +) -> EngineTuple: + """The tuple a quickstart gate (boot, witness, acceptance) launches. + + ``max_batch_size`` is a witness's own override (both legs at one limit). + """ + flags = _flag_values(runtime.get("extra_flags") or ()) + values: dict[str, int | None] = {} + for spelled, axis in _FLAG_AXES.items(): + values[axis] = _int_or_none(flags.get(spelled)) + if values[axis] is None: + values[axis] = ENGINE_DEFAULTS[axis] + if max_batch_size is not None: + values["max_batch_size"] = int(max_batch_size) + graphs: tuple[int, ...] = () + if runtime.get("cuda_graph"): + explicit = [int(v) for v in flags.get(_GRAPH_LIST_FLAG, []) if v.isdigit()] + if explicit: + graphs = tuple(sorted(set(explicit))) + else: + graphs = tuple( + generate_graph_batch_sizes( + DEFAULT_GRAPH_MAX_BATCH_SIZE, _GRAPH_PADDING_FLAG in flags + ) + ) + graphs = tuple( + g for g in graphs if values["max_batch_size"] is None or g <= values["max_batch_size"] + ) + return EngineTuple( + gate=gate, + unit=unit, + member=member, + graph_batch_sizes=graphs, + source="runtime block", + **values, + ) + + +def options_tuple( + accuracy: Mapping[str, Any], + options_path: str | None, + *, + unit: str, + member: str, +) -> EngineTuple: + """The tuple ``trtllm-eval`` launches for the accuracy gate. + + The options YAML is read first because ``trtllm-eval`` lets it override + its own flags; ``accuracy.max_seq_len`` / ``max_batch_size`` fill what the + YAML leaves out; the engine defaults fill the rest. + """ + options: dict[str, Any] = {} + if options_path and Path(options_path).is_file(): + loaded = yaml.safe_load(Path(options_path).read_text(encoding="utf-8")) or {} + options = loaded if isinstance(loaded, dict) else {} + values: dict[str, int | None] = {} + for axis in ENVELOPE_AXES: + value = options.get(axis) + if value is None: + value = accuracy.get(axis) + values[axis] = int(value) if value is not None else ENGINE_DEFAULTS[axis] + graphs: tuple[int, ...] = () + graph_config = options.get("cuda_graph_config") + if isinstance(graph_config, dict): + policy = resolve_graph_policy( + { + "batch_sizes": graph_config.get("batch_sizes"), + "padding": bool(graph_config.get("enable_padding", False)), + "max_batch_size": graph_config.get("max_batch_size") + or DEFAULT_GRAPH_MAX_BATCH_SIZE, + }, + values["max_batch_size"], + ) + graphs = policy.clipped(values["max_batch_size"]) + source = f"options {Path(options_path).name}" if options_path else "trtllm-eval defaults" + return EngineTuple( + gate=EVIDENCE_ACCURACY, + unit=unit, + member=member, + graph_batch_sizes=graphs, + source=source, + **values, + ) + + +# -------------------------------------------------------------- unit gates + + +@dataclass(frozen=True) +class UnitGates: + """What the run's own gates measure for one unit.""" + + unit: str + features: tuple[str, ...] + members: tuple[str, ...] + kinds: frozenset[str] + #: The witness runs at one member only; see the workflow's call site. + witness_member: str | None = None + + +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 default_unit_kinds(task_data: Mapping[str, Any], unit: str) -> frozenset[str]: + """The gate kinds a unit runs, read from what the spec declares for it. + + The orchestrator passes its own exact answer; this is the reading a + caller without it (preflight in a test, the migration command) uses: a + witness where ``accuracy.witness`` has a block for the unit, an + acceptance gate where ``accuracy.acceptance`` names it. + """ + accuracy = task_data.get(ACCURACY_FIELD) or {} + kinds = {EVIDENCE_BOOT, EVIDENCE_ACCURACY} + witness = accuracy.get("witness") if isinstance(accuracy, dict) else None + if isinstance(witness, dict) and unit in witness: + kinds.add(EVIDENCE_WITNESS) + acceptance = accuracy.get("acceptance") if isinstance(accuracy, dict) else None + if isinstance(acceptance, dict) and unit in acceptance: + kinds.add(EVIDENCE_ACCEPTANCE) + return frozenset(kinds) + + +def resolve_unit_gates( + task_data: Mapping[str, Any], + world_size: int, + *, + unit_kinds: Callable[[str], Iterable[str]] | None = None, +) -> list[UnitGates]: + """Every unit's gated members and gate kinds, cheapest member first.""" + from .gates import order_members_by_cost + + gates: list[UnitGates] = [] + for unit in assembly_units(task_data): + members = tuple(order_members_by_cost(members_for(task_data, unit, world_size))) + kinds = frozenset( + unit_kinds(unit) if unit_kinds is not None else default_unit_kinds(task_data, unit) + ) + gates.append( + UnitGates( + unit=unit, + features=features_through(task_data, unit), + members=members, + kinds=kinds, + witness_member=members[0] if EVIDENCE_WITNESS in kinds and members else None, + ) + ) + return gates + + +def gate_tuples(task_data: Mapping[str, Any], gates: Sequence[UnitGates]) -> list[EngineTuple]: + """Every tuple the run's existing gates launch.""" + from .gates import resolve_llm_api_options, resolve_runtime + + accuracy = task_data.get(ACCURACY_FIELD) or {} + witness_blocks = accuracy.get("witness") if isinstance(accuracy, dict) else None + witness_blocks = witness_blocks if isinstance(witness_blocks, dict) else {} + tuples: list[EngineTuple] = [] + for unit_gates in gates: + unit = unit_gates.unit + runtime = resolve_runtime(task_data, unit) + options = resolve_llm_api_options(accuracy, unit) if isinstance(accuracy, dict) else None + for member in unit_gates.members: + if EVIDENCE_BOOT in unit_gates.kinds: + tuples.append(runtime_tuple(runtime, gate=EVIDENCE_BOOT, unit=unit, member=member)) + if EVIDENCE_ACCURACY in unit_gates.kinds: + tuples.append(options_tuple(accuracy, options, unit=unit, member=member)) + if EVIDENCE_ACCEPTANCE in unit_gates.kinds: + tuples.append( + runtime_tuple(runtime, gate=EVIDENCE_ACCEPTANCE, unit=unit, member=member) + ) + if unit_gates.witness_member is not None: + block = witness_blocks.get(unit) or {} + batch = block.get("max_batch_size") if isinstance(block, dict) else None + tuples.append( + runtime_tuple( + runtime, + gate=EVIDENCE_WITNESS, + unit=unit, + member=unit_gates.witness_member, + max_batch_size=int(batch) if batch else None, + ) + ) + return tuples + + +def derived_envelope(tuples: Sequence[EngineTuple], members: Sequence[str]) -> ServingEnvelope: + """The legacy envelope: the hull of the gate tuples. + + This is what a run without ``serving_envelope`` has always certified -- + the gate tuples and nothing more -- made explicit so it can be rendered + as exactly that. + """ + + def hull(axis: str) -> int | None: + known = [t.limit(axis) for t in tuples if t.limit(axis) is not None] + return max(known) if known else None + + graphs = sorted({g for t in tuples for g in t.graph_batch_sizes}) + parsed: list[EnvelopeMember] = [] + for name in members: + try: + parsed.append(parse_envelope_member(name)) + except ValueError: + continue + return ServingEnvelope( + max_batch_size=hull("max_batch_size"), + max_num_tokens=hull("max_num_tokens"), + max_seq_len=hull("max_seq_len"), + graph=GraphPolicy( + batch_sizes=tuple(graphs), + source="union of the gates' capture lists" if graphs else "eager", + ), + members=tuple(parsed), + declared=False, + ) + + +# ------------------------------------------------------------------- matrix + + +@dataclass(frozen=True) +class MatrixCell: + """One required piece of engine evidence, and how it will be obtained.""" + + cell: str + features: tuple[str, ...] + member: str + evidence: str + at: str + status: str + reason: str + unit: str | None = None + + @property + def slug(self) -> str: + """A file-name-safe identity: the evidence file for an external cell. + + Keyed by the cell's NAME rather than its feature list, which for a + late unit is every feature of the run; the schema keeps names unique. + """ + raw = f"{self.cell}__{self.member}__{self.evidence}__{self.at}" + return re.sub(r"[^A-Za-z0-9_.+-]+", "-", raw).strip("-") + + def render(self) -> str: + feats = "+".join(self.features) or "core" + where = f" via unit `{self.unit}`" if self.unit else "" + return ( + f"[{self.status:<8}] {feats} @ {self.member}: {self.evidence} at {self.at}" + f"{where} -- {self.reason}" + ) + + +def _at_label(at: Any) -> str: + if isinstance(at, dict): + return ",".join(f"{k}={v}" for k, v in sorted(at.items())) + return str(at or AT_GATE) + + +def _point_tuple( + at: Any, + envelope: ServingEnvelope, + base: EngineTuple | None, + *, + unit: str, + member: str, + prompt_tokens: int | None, +) -> EngineTuple: + """The engine tuple a probe at ``at`` launches.""" + if at == AT_CEILING: + values = {axis: envelope.limit(axis) for axis in ENVELOPE_AXES} + label = AT_CEILING + else: + values = {axis: (base.limit(axis) if base else None) for axis in ENVELOPE_AXES} + values.update({axis: int(v) for axis, v in dict(at).items()}) + label = _at_label(at) + graphs = envelope.graph.clipped(values["max_batch_size"]) + return EngineTuple( + gate=PROBE_GATE, + unit=unit, + member=member, + graph_batch_sizes=graphs, + source=f"serving envelope ({label})", + prompt_tokens=prompt_tokens or envelope.probe_prompt_tokens, + graph_padding=envelope.graph.padding and bool(graphs), + **values, + ) + + +def _reaches(tuple_: EngineTuple, target: EngineTuple) -> bool: + """Whether an existing gate tuple already measures ``target``'s point.""" + for axis in ENVELOPE_AXES: + want = target.limit(axis) + have = tuple_.limit(axis) + if want is None: + continue + if have is None or have < want: + return False + return set(target.graph_batch_sizes) <= set(tuple_.graph_batch_sizes) + + +# --------------------------------------------------------------------- plan + + +@dataclass(frozen=True) +class CertificationPlan: + """Everything derived from the envelope and the matrix, for one run.""" + + envelope: ServingEnvelope + tuples: tuple[EngineTuple, ...] = () + probes: tuple[EngineTuple, ...] = () + cells: tuple[MatrixCell, ...] = () + errors: tuple[str, ...] = () + warnings: tuple[str, ...] = () + #: The world size of the run's own allocation. + world_size: int = 1 + + @property + def declared(self) -> bool: + return self.envelope.declared + + def probes_for(self, unit: str) -> list[EngineTuple]: + return [probe for probe in self.probes if probe.unit == unit] + + @property + def external_cells(self) -> list[MatrixCell]: + return [cell for cell in self.cells if cell.status == CELL_EXTERNAL] + + @property + def certified_capture_batches(self) -> tuple[int, ...]: + """Every batch size a graph is captured at: the envelope's and the gates'. + + The union, because a gate that captures a size the envelope does not + list still runs the catalog at that size, and an uncertified value in + a gate is an uncertified value however it got there. + """ + sizes = set(self.envelope.capture_batches) + for tuple_ in (*self.tuples, *self.probes): + sizes.update(tuple_.graph_batch_sizes) + return tuple(sorted(sizes)) + + def standard(self) -> list[str]: + """The certification standard, one line per axis.""" + env = self.envelope + gather = env.max_gather_world_size + lines = [ + f"tokens per step per rank: T <= {env.max_num_tokens}" + + ( + f" (x{gather} = {env.max_num_tokens * gather} after an attention-DP gather)" + if gather > 1 and env.max_num_tokens + else "" + ), + f"requests per step: <= {env.max_batch_size}", + f"per-request length: <= {env.max_seq_len}", + "CUDA-graph capture batches G: " + + (compact_values(self.certified_capture_batches) or "none (eager only)") + + "; a call site captured at row multiple m certifies m*G for every G", + "pools: any size up to the largest the GPU allocates at the envelope's " + "settings, with rows addressed at the far end", + ] + if env.max_draft_len: + lines.append( + f"speculation: up to {env.max_draft_len} draft token(s) per request per step " + f"({env.max_draft_len + 1} rows per request at the verify step)" + ) + return lines + + def notes(self) -> list[str]: + """Preflight's notes: what is certified, what is probed, what is owed.""" + env = self.envelope + notes: list[str] = [] + if env.declared: + members = ", ".join(m.describe(env.gpus_per_node) for m in env.members) + notes.append( + f"serving envelope: B<={env.max_batch_size} T<={env.max_num_tokens} " + f"L<={env.max_seq_len} G={compact_values(env.capture_batches) or 'eager'}; " + f"members {members}" + ) + else: + notes.append( + "no serving_envelope declared: the certified set is the hull of the gate " + f"tuples (B<={env.max_batch_size} T<={env.max_num_tokens} " + f"L<={env.max_seq_len}). Nothing outside those tuples is certified. " + "`python -m agent_flow.workflows.staircase_onboard.envelope migrate " + "` prints a block to declare." + ) + if self.probes: + notes.append( + f"{len(self.probes)} envelope probe(s): " + + "; ".join(probe.label() for probe in self.probes) + ) + external = self.external_cells + if external: + notes.append( + f"{len(external)} feature-matrix cell(s) need evidence this allocation cannot " + f"produce; delivery requires `{EVIDENCE_DIRNAME}/.yaml` for each: " + + ", ".join(cell.slug for cell in external) + ) + notes.extend(self.warnings) + return notes + + def render(self) -> str: + """``certification.md``: the whole plan, for every role to read.""" + env = self.envelope + head = ( + "the DECLARED serving envelope" + if env.declared + else "a DERIVED envelope -- the hull of the gate tuples, because the spec " + "declares no `serving_envelope`. Nothing outside these tuples is certified" + ) + lines = [ + "# Certification plan", + "", + f"The certified set is {head}. Catalog receipts, capacity constants and " + "integer-width limits are judged against it. **A gate tuple is evidence " + "inside the set, not the set.**", + "", + "## Certification standard", + "", + *(f"- {line}" for line in self.standard()), + "", + "## Members", + "", + ] + gated = {t.member for t in self.tuples} + for member in env.members: + where = ( + "gated by this run" if member.name in gated else "not runnable in this allocation" + ) + lines.append(f"- {member.describe(env.gpus_per_node)}: {where}") + lines += ["", "## Gate tuples (what the run's gates launch)", ""] + lines += [f"- {t.label()} [{t.source}]" for t in self.tuples] or ["- (none)"] + lines += ["", "## Envelope probes (launched in addition, per unit)", ""] + lines += [f"- {p.label()}" for p in self.probes] or ["- (none)"] + lines += ["", "## Feature matrix", ""] + lines += [f"- {cell.render()}" for cell in self.cells] or ["- (none)"] + external = self.external_cells + if external: + lines += [ + "", + "## Evidence owed from outside this allocation", + "", + f"Delivery fails until each of these has `{EVIDENCE_DIRNAME}/.yaml` " + "in the workspace, with `verdict: pass`, the `tree` it was measured on and " + "the `log` that shows it:", + "", + *(f"- `{cell.slug}`" for cell in external), + ] + if self.warnings: + lines += ["", "## Warnings", "", *(f"- {w}" for w in self.warnings)] + return "\n".join(lines) + "\n" + + def prompt_block(self, path: Path | str | None = None) -> str: + """The block a role's prompt carries: the standard, and where the rest is.""" + env = self.envelope + kind = "declared" if env.declared else "DERIVED from the gate tuples (legacy)" + where = f" The full plan is `{path}`." if path else "" + return ( + f"## Serving envelope ({kind})\n\n" + "Certify against the envelope, never against the tuples a gate happened to " + f"run.{where}\n\n" + "\n".join(f"- {line}" for line in self.standard()) + "\n" + ) + + +def _matrix_cells_from_spec(task_data: Mapping[str, Any]) -> list[dict[str, Any]]: + block = task_data.get(FEATURE_MATRIX_FIELD) + return [dict(cell) for cell in block] if isinstance(block, list) else [] + + +def _unit_for(features: Sequence[str], gates: Sequence[UnitGates]) -> UnitGates | None: + """The earliest unit whose target implements every feature in ``features``.""" + wanted = set(features) + return next((g for g in gates if wanted <= set(g.features)), None) + + +def derive_certification( + task_data: Mapping[str, Any], + *, + world_size: int, + gpus_per_node: int | None = None, + checkpoint_config: Mapping[str, Any] | None = None, + unit_kinds: Callable[[str], Iterable[str]] | None = None, +) -> CertificationPlan: + """Derive the gate tuples, probes, matrix cells and certified set for a run. + + Pure apart from reading the accuracy gate's options YAMLs. ``errors`` are + inconsistencies the run must not start with; preflight raises on them. + """ + gates = resolve_unit_gates(task_data, world_size, unit_kinds=unit_kinds) + tuples = gate_tuples(task_data, gates) + gated_members = sorted({m for g in gates for m in g.members}) + checkpoint_max = None + if checkpoint_config: + text_config = checkpoint_config.get("text_config") + source = text_config if isinstance(text_config, dict) else checkpoint_config + raw = source.get("max_position_embeddings") or checkpoint_config.get( + "max_position_embeddings" + ) + checkpoint_max = int(raw) if isinstance(raw, int) else None + + block = task_data.get(SERVING_ENVELOPE_FIELD) + environment = task_data.get("environment") or {} + node_width = gpus_per_node or environment.get("gpus_per_node") + errors: list[str] = [] + warnings: list[str] = [] + + if not isinstance(block, dict): + envelope = derived_envelope(tuples, gated_members) + cells = [ + MatrixCell( + cell=f"{g.unit}-gates", + features=g.features, + member=member, + evidence=kind, + at=AT_GATE, + status=CELL_GATED, + reason="the unit's own gate", + unit=g.unit, + ) + for g in gates + for member in g.members + for kind in sorted(g.kinds) + if kind != EVIDENCE_WITNESS or member == g.witness_member + ] + return CertificationPlan( + envelope=envelope, + tuples=tuple(tuples), + cells=tuple(cells), + world_size=world_size, + ) + + envelope = envelope_from_spec(block, gpus_per_node=node_width) + claimed = {m.name for m in envelope.members} + + for member in gated_members: + if member not in claimed: + errors.append( + f"member {member} is gated by this run but is not in " + f"serving_envelope.members {sorted(claimed)}: a gate would measure a " + f"configuration the target does not claim" + ) + for tuple_ in tuples: + problems = tuple_.outside(envelope, checkpoint_max=checkpoint_max) + if problems: + hint = ( + " (unset max_seq_len means the checkpoint's max_position_embeddings; pass " + "--max_seq_len)" + if any(p.startswith("max_seq_len") for p in problems) and tuple_.max_seq_len is None + else "" + ) + errors.append( + f"the {tuple_.label()} gate launches outside the envelope: " + f"{'; '.join(problems)}{hint}" + ) + if ( + checkpoint_max is not None + and envelope.max_seq_len + and envelope.max_seq_len > checkpoint_max + ): + errors.append( + f"serving_envelope.max_seq_len={envelope.max_seq_len} exceeds the checkpoint's " + f"max_position_embeddings={checkpoint_max}" + ) + if ( + envelope.max_batch_size + and envelope.max_num_tokens + and envelope.max_batch_size * (1 + envelope.max_draft_len) > envelope.max_num_tokens + ): + warnings.append( + f"a full batch of decode steps needs {envelope.max_batch_size} x " + f"{1 + envelope.max_draft_len} = " + f"{envelope.max_batch_size * (1 + envelope.max_draft_len)} tokens, over " + f"max_num_tokens={envelope.max_num_tokens}: the scheduler can never fill the " + f"batch limit in one step" + ) + if node_width: + for member in envelope.members: + if member.nodes(int(node_width)) > 1: + warnings.append( + f"{member.describe(int(node_width))} spans nodes; its evidence comes " + f"from a multi-node session, not from this run's gates" + ) + + cells: list[MatrixCell] = [] + probes: dict[tuple[str, str, str], EngineTuple] = {} + + def gate_base(unit: str, member: str, kind: str = EVIDENCE_BOOT) -> EngineTuple | None: + """The tuple ``kind``'s gate launches for (unit, member): a point's other axes.""" + return next( + (t for t in tuples if t.unit == unit and t.member == member and t.gate == kind), + None, + ) + + def add_probe(probe: EngineTuple) -> None: + key = (probe.unit, probe.member, probe.point()) + probes.setdefault(key, probe) + + # The gates the run already runs are cells by construction. + for g in gates: + for member in g.members: + for kind in sorted(g.kinds): + if kind == EVIDENCE_WITNESS and member != g.witness_member: + continue + cells.append( + MatrixCell( + cell=f"{g.unit}-gates", + features=g.features, + member=member, + evidence=kind, + at=AT_GATE, + status=CELL_GATED, + reason="the unit's own gate", + unit=g.unit, + ) + ) + # The minimum envelope evidence, per unit: one boot at the ceiling, at + # the cheapest member. Warm-up at the real token budget and capture at + # every claimed batch size are what expose a fixed-size buffer. + if g.members: + probe = _point_tuple( + AT_CEILING, + envelope, + gate_base(g.unit, g.members[0]), + unit=g.unit, + member=g.members[0], + prompt_tokens=None, + ) + add_probe(probe) + cells.append( + MatrixCell( + cell=f"{g.unit}-ceiling", + features=g.features, + member=g.members[0], + evidence=EVIDENCE_BOOT, + at=AT_CEILING, + status=CELL_PROBE, + reason=probe.label(), + unit=g.unit, + ) + ) + + all_features = tuple(str(f) for f in task_data.get(FEATURES_FIELD) or ()) + for member in envelope.members: + if member.name in gated_members: + continue + cells.append( + MatrixCell( + cell="envelope-member", + features=all_features, + member=member.name, + evidence=EVIDENCE_BOOT, + at=AT_CEILING, + status=CELL_EXTERNAL, + reason=( + f"{member.describe(envelope.gpus_per_node)} is claimed by the envelope " + f"and gated by no unit of this run" + ), + ) + ) + + for index, raw in enumerate(_matrix_cells_from_spec(task_data)): + name = str(raw.get("name") or f"cell{index}") + features = tuple(str(f) for f in raw.get("features") or ()) + evidence = raw.get("evidence") or [EVIDENCE_BOOT, EVIDENCE_ACCURACY] + evidence = [evidence] if isinstance(evidence, str) else list(evidence) + at = raw.get("at") or AT_GATE + members = raw.get("members") or [m.name for m in envelope.members] + unit_gates = _unit_for(features, gates) + for member in members: + for kind in evidence: + status, reason, unit = _classify( + unit_gates, + member=str(member), + kind=str(kind), + at=at, + envelope=envelope, + tuples=tuples, + prompt_tokens=raw.get("prompt_tokens"), + gate_base=gate_base, + add_probe=add_probe, + ) + cells.append( + MatrixCell( + cell=name, + features=features, + member=str(member), + evidence=str(kind), + at=_at_label(at), + status=status, + reason=reason, + unit=unit, + ) + ) + + return CertificationPlan( + envelope=envelope, + tuples=tuple(tuples), + probes=tuple(probes.values()), + cells=tuple(_dedupe_cells(cells)), + errors=tuple(errors), + warnings=tuple(warnings), + world_size=world_size, + ) + + +def _classify( + unit_gates: UnitGates | None, + *, + member: str, + kind: str, + at: Any, + envelope: ServingEnvelope, + tuples: Sequence[EngineTuple], + prompt_tokens: Any, + gate_base: Callable[..., EngineTuple | None], + add_probe: Callable[[EngineTuple], None], +) -> tuple[str, str, str | None]: + """Classify one declared cell as gated, probe or external.""" + if unit_gates is None: + return CELL_EXTERNAL, "no unit of this run implements this feature set", None + unit = unit_gates.unit + if member not in unit_gates.members: + return ( + CELL_EXTERNAL, + f"{member} is not gated by the `{unit}` unit (gated: " + f"{', '.join(unit_gates.members) or 'none'})", + unit, + ) + if kind not in unit_gates.kinds: + return CELL_EXTERNAL, f"the `{unit}` unit runs no {kind} gate", unit + if kind == EVIDENCE_WITNESS and member != unit_gates.witness_member: + return ( + CELL_EXTERNAL, + f"the `{unit}` witness runs at {unit_gates.witness_member} only", + unit, + ) + if at == AT_GATE: + return CELL_GATED, f"the `{unit}` unit's {kind} gate", unit + # The point moves only the axes it names; the rest stay at what this + # kind's own gate launches. A long prompt is a boot probe's workload. + target = _point_tuple( + at, + envelope, + gate_base(unit, member, kind), + unit=unit, + member=member, + prompt_tokens=int(prompt_tokens) if prompt_tokens else None, + ) + if kind != EVIDENCE_BOOT: + target = replace(target, prompt_tokens=None) + existing = [t for t in tuples if t.unit == unit and t.member == member and t.gate == kind] + if any(_reaches(t, target) for t in existing): + return CELL_GATED, f"the `{unit}` {kind} gate already runs at {target.point()}", unit + if kind == EVIDENCE_BOOT: + add_probe(target) + return CELL_PROBE, target.label(), unit + return ( + CELL_EXTERNAL, + f"the `{unit}` {kind} gate runs below {target.point()}; record engine evidence " + f"at that point", + unit, + ) + + +def _dedupe_cells(cells: Iterable[MatrixCell]) -> list[MatrixCell]: + seen: dict[tuple[str, str, str, str], MatrixCell] = {} + rank = {CELL_GATED: 0, CELL_PROBE: 1, CELL_EXTERNAL: 2} + for cell in cells: + key = (",".join(cell.features), cell.member, cell.evidence, cell.at) + held = seen.get(key) + # The cheapest way to get the evidence wins: a declared cell that an + # existing gate already covers is not owed twice. On a tie the later + # cell -- the declared one, over the derived default -- keeps its name. + if held is None or rank[cell.status] <= rank[held.status]: + seen[key] = cell + return list(seen.values()) + + +# ----------------------------------------------------------------- probes + + +def probe_flags(probe: EngineTuple) -> list[str]: + """The engine flags that put a quickstart boot at ``probe``'s point.""" + flags = [f"--{axis}={probe.limit(axis)}" for axis in ENVELOPE_AXES if probe.limit(axis)] + if probe.graph_batch_sizes: + flags.append(f"{_GRAPH_LIST_FLAG} " + " ".join(str(g) for g in probe.graph_batch_sizes)) + if probe.graph_padding: + flags.append(_GRAPH_PADDING_FLAG) + return flags + + +def _without_axis_flags(flags: Iterable[str]) -> list[str]: + """``flags`` with every envelope-axis and graph-list spelling removed.""" + owned = (*_FLAG_AXES, _GRAPH_LIST_FLAG, _GRAPH_PADDING_FLAG) + kept: list[str] = [] + skip_values = False + for flag in flags: + word = str(flag) + head = word.split()[0] if word.split() else word + name = head.partition("=")[0] + if name in owned: + # A bare flag's value(s) may follow as separate elements. + skip_values = "=" not in head and len(word.split()) == 1 + continue + if skip_values and not word.startswith("--"): + continue + skip_values = False + kept.append(word) + return kept + + +def probe_task_data( + task_data: Mapping[str, Any], + unit: str, + probe: EngineTuple, +) -> dict[str, Any]: + """A copy of the spec whose runtime block puts ``unit``'s boot at the probe point. + + The probe IS the boot gate -- same launcher, same prompts, same keywords, + same runtime switches -- with only the envelope's limits and capture list + changed. A probe built any other way would measure two variables at once. + """ + from .gates import resolve_runtime + + data = copy.deepcopy(dict(task_data)) + runtime = dict(resolve_runtime(task_data, unit)) + runtime["extra_flags"] = [ + *_without_axis_flags(runtime.get("extra_flags") or ()), + *probe_flags(probe), + ] + runtime["cuda_graph"] = bool(probe.graph_batch_sizes) + data["runtime"] = {"default": runtime} + return data + + +def probe_prompts(base: Sequence[str], probe: EngineTuple) -> tuple[str, ...]: + """The boot prompts, plus one generated long prompt when the probe asks for one.""" + if not probe.prompt_tokens: + return tuple(base) + from .witness import filler_words + + words = min(int(probe.prompt_tokens), MAX_WITNESS_PROMPT_WORDS) + return (*base, filler_words(words)) + + +def run_envelope_probe( + environment: Any, + *, + task_data: Mapping[str, Any], + probe: EngineTuple, + spec: Any, + runner: Any, + features: Any, +) -> Any: + """Boot the unit at ``probe``'s point. Returns a ``GateResult``. + + Recorded as the ``envelope`` gate with the point in its detail, so the + gate history reads as a third kind of evidence rather than as a second + boot that happened to pass. + """ + from .gates import BootSpec, run_boot_gate + + probe_spec = BootSpec( + prompts=probe_prompts(spec.prompts, probe), + keywords=tuple(spec.keywords), + max_tokens=int(spec.max_tokens), + ) + result = run_boot_gate( + environment, + task_data=probe_task_data(task_data, probe.unit, probe), + spec=probe_spec, + runner=runner, + features=features, + unit=probe.unit, + member=probe.member, + ) + return replace(result, name=PROBE_GATE, detail=f"at {probe.point()}: {result.detail}") + + +# ----------------------------------------------------------------- evidence + + +def check_external_evidence(plan: CertificationPlan, workspace: Path | str) -> list[Any]: + """Delivery checks: every external matrix cell has a recorded verdict. + + An evidence file is YAML with ``verdict: pass``, the ``tree`` it was + measured on and the ``log`` that shows it. A missing or failing one is a + failed check, because the envelope claims that cell and nothing in this + run measured it. + """ + from .delivery import Check + + root = Path(workspace) / EVIDENCE_DIRNAME + checks: list[Check] = [] + for cell in plan.external_cells: + path = root / f"{cell.slug}.yaml" + name = f"matrix:{cell.slug}" + if not path.is_file(): + checks.append( + Check( + name=name, + passed=False, + detail=f"no evidence at {path}; {cell.reason}", + ) + ) + continue + try: + data = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + except yaml.YAMLError as exc: + checks.append(Check(name=name, passed=False, detail=f"{path} is not YAML: {exc}")) + continue + if not isinstance(data, dict): + checks.append(Check(name=name, passed=False, detail=f"{path} is not a mapping")) + continue + missing = [key for key in ("verdict", "tree", "log") if not data.get(key)] + verdict = str(data.get("verdict", "")).strip().lower() + if missing: + detail = f"{path} lacks {', '.join(missing)}" + elif verdict != "pass": + detail = f"{path} records verdict {verdict!r}" + else: + detail = f"{path}: pass on {data['tree']} ({data['log']})" + checks.append(Check(name=name, passed=not missing and verdict == "pass", detail=detail)) + return checks + + +# --------------------------------------------------------------- validation + + +def validate_serving_envelope(block: Any, errors: list[str]) -> None: + """Schema-validate ``serving_envelope``. Appends to ``errors``.""" + where = SERVING_ENVELOPE_FIELD + if not isinstance(block, dict): + errors.append(f"'{where}' must be a mapping, got {type(block).__name__}") + return + reject_unknown_keys(block, ENVELOPE_KEYS, errors, where=where) + for key in ENVELOPE_REQUIRED: + if block.get(key) is None: + errors.append( + f"'{where}.{key}' is required: the envelope is what every certification is " + f"judged against, so none of its limits may be left to an engine default" + ) + for key in ENVELOPE_AXES: + value = block.get(key) + if value is not None and not _positive_int(value): + errors.append(f"'{where}.{key}' must be an integer >= 1, got {value!r}") + for key in ("gpus_per_node", "probe_prompt_tokens"): + value = block.get(key) + if value is not None and not _positive_int(value): + errors.append(f"'{where}.{key}' must be an integer >= 1, got {value!r}") + probe_words = block.get("probe_prompt_tokens") + if _positive_int(probe_words) and probe_words > MAX_WITNESS_PROMPT_WORDS: + errors.append( + f"'{where}.probe_prompt_tokens' is {probe_words}; the probe prompt travels as one " + f"argv string, which Linux caps at 128 KiB, so the limit is " + f"{MAX_WITNESS_PROMPT_WORDS}" + ) + draft = block.get("max_draft_len") + if draft is not None and not _non_negative_int(draft): + errors.append(f"'{where}.max_draft_len' must be an integer >= 0, got {draft!r}") + + _validate_graph_policy(block.get("cuda_graph"), block.get("max_batch_size"), errors) + + members = block.get("members") + if members is not None: + if not isinstance(members, list) or not members: + errors.append(f"'{where}.members' must be a non-empty list of parallel members") + else: + seen: set[str] = set() + for index, member in enumerate(members): + try: + parse_envelope_member(member) + except ValueError as exc: + errors.append(f"'{where}.members[{index}]': {exc}") + continue + if member in seen: + errors.append(f"'{where}.members' lists {member!r} twice") + seen.add(member) + + +def _validate_graph_policy(block: Any, max_batch_size: Any, errors: list[str]) -> None: + where = f"{SERVING_ENVELOPE_FIELD}.cuda_graph" + if block is None or block is False: + return + if not isinstance(block, dict): + errors.append( + f"'{where}' must be false (eager only) or a mapping with 'batch_sizes' or " + f"'padding'; got {block!r}" + ) + return + reject_unknown_keys(block, CUDA_GRAPH_KEYS, errors, where=where) + sizes = block.get("batch_sizes") + if sizes is None and "padding" not in block: + errors.append( + f"'{where}' needs 'batch_sizes' (the capture list) or 'padding' (the policy the " + f"engine generates one from): the capture list decides which row counts every " + f"captured call site is certified at" + ) + if sizes is not None: + if not isinstance(sizes, list) or not sizes or not all(_positive_int(s) for s in sizes): + errors.append(f"'{where}.batch_sizes' must be a non-empty list of integers >= 1") + elif _positive_int(max_batch_size) and max(sizes) > max_batch_size: + errors.append( + f"'{where}.batch_sizes' captures {max(sizes)}, over " + f"{SERVING_ENVELOPE_FIELD}.max_batch_size={max_batch_size}" + ) + padding = block.get("padding") + if padding is not None and not isinstance(padding, bool): + errors.append(f"'{where}.padding' must be true or false") + bound = block.get("max_batch_size") + if bound is not None: + if not _positive_int(bound): + errors.append(f"'{where}.max_batch_size' must be an integer >= 1") + elif _positive_int(max_batch_size) and bound > max_batch_size: + errors.append( + f"'{where}.max_batch_size'={bound} exceeds " + f"{SERVING_ENVELOPE_FIELD}.max_batch_size={max_batch_size}" + ) + + +def validate_feature_matrix(data: Mapping[str, Any], errors: list[str]) -> None: + """Schema-validate ``feature_matrix`` against the spec's features and envelope.""" + where = FEATURE_MATRIX_FIELD + block = data.get(FEATURE_MATRIX_FIELD) + envelope = data.get(SERVING_ENVELOPE_FIELD) + if not isinstance(envelope, dict): + errors.append( + f"'{where}' needs '{SERVING_ENVELOPE_FIELD}': a matrix cell names members and " + f"points of the envelope, so there has to be one" + ) + return + if not isinstance(block, list): + errors.append(f"'{where}' must be a list of cells, got {type(block).__name__}") + return + features = {str(f) for f in data.get(FEATURES_FIELD) or ()} + claimed = {str(m) for m in envelope.get("members") or ()} + names: set[str] = set() + for index, cell in enumerate(block): + cell_where = f"{where}[{index}]" + if not isinstance(cell, dict): + errors.append(f"'{cell_where}' must be a mapping") + continue + reject_unknown_keys(cell, MATRIX_CELL_KEYS, errors, where=cell_where) + name = cell.get("name") + if name is not None: + if not isinstance(name, str) or not re.fullmatch(r"[A-Za-z0-9_.+-]+", name): + errors.append( + f"'{cell_where}.name' must be a file-name-safe string (letters, digits, " + f"'_', '.', '+', '-'): it names the cell's evidence file" + ) + elif name in names: + errors.append(f"'{cell_where}.name' {name!r} is used by another cell") + names.add(str(name)) + cell_features = cell.get("features") or [] + if not isinstance(cell_features, list) or not all( + isinstance(f, str) for f in cell_features + ): + errors.append(f"'{cell_where}.features' must be a list of feature names ([] = core)") + else: + unknown = sorted(set(cell_features) - features) + if unknown: + errors.append( + f"'{cell_where}.features' names {unknown}, which this run does not build " + f"(features: {sorted(features)})" + ) + members = cell.get("members") + if members is not None: + if not isinstance(members, list) or not members: + errors.append(f"'{cell_where}.members' must be a non-empty list") + else: + stray = sorted(str(m) for m in members if str(m) not in claimed) + if stray: + errors.append( + f"'{cell_where}.members' names {stray}, which " + f"'{SERVING_ENVELOPE_FIELD}.members' does not claim" + ) + evidence = cell.get("evidence") + kinds = [evidence] if isinstance(evidence, str) else evidence + if kinds is not None and ( + not isinstance(kinds, list) or not kinds or any(k not in EVIDENCE_KINDS for k in kinds) + ): + errors.append( + f"'{cell_where}.evidence' must be one of {list(EVIDENCE_KINDS)} or a list of " + f"them, got {evidence!r}" + ) + _validate_point(cell.get("at"), envelope, errors, f"{cell_where}.at") + prompt_tokens = cell.get("prompt_tokens") + if prompt_tokens is not None and ( + not _positive_int(prompt_tokens) or prompt_tokens > MAX_WITNESS_PROMPT_WORDS + ): + errors.append( + f"'{cell_where}.prompt_tokens' must be an integer in 1..{MAX_WITNESS_PROMPT_WORDS}" + ) + + +def _validate_point(at: Any, envelope: Mapping[str, Any], errors: list[str], where: str) -> None: + if at is None or at in (AT_GATE, AT_CEILING): + return + if not isinstance(at, dict) or not at: + errors.append( + f"'{where}' must be '{AT_GATE}', '{AT_CEILING}' or a mapping of " + f"{list(ENVELOPE_AXES)} to values; got {at!r}" + ) + return + reject_unknown_keys(at, ENVELOPE_AXES, errors, where=where) + for axis, value in at.items(): + if axis not in ENVELOPE_AXES: + continue + if not _positive_int(value): + errors.append(f"'{where}.{axis}' must be an integer >= 1") + continue + bound = envelope.get(axis) + if _positive_int(bound) and value > bound: + errors.append( + f"'{where}.{axis}'={value} is outside " + f"{SERVING_ENVELOPE_FIELD}.{axis}={bound}; a point outside the envelope " + f"is a plan change, not a matrix cell" + ) + + +def validate_certification_blocks(data: Mapping[str, Any], errors: list[str]) -> None: + """The schema's entry point for both blocks.""" + if data.get(SERVING_ENVELOPE_FIELD) is not None: + validate_serving_envelope(data[SERVING_ENVELOPE_FIELD], errors) + if data.get(FEATURE_MATRIX_FIELD) is not None: + validate_feature_matrix(data, errors) + + +# ---------------------------------------------------------------- migration + + +def migration_block( + task_data: Mapping[str, Any], + *, + world_size: int, + unit_kinds: Callable[[str], Iterable[str]] | None = None, +) -> str: + """A ``serving_envelope`` + ``feature_matrix`` block for a spec that has none. + + Derived from the gate tuples, so adopting it unchanged keeps the old + certified set -- now stated rather than implied. The comments say what to + widen: the point of declaring an envelope is to claim what the target is + deployed at, which is usually more than the gates ran. + """ + gates = resolve_unit_gates(task_data, world_size, unit_kinds=unit_kinds) + tuples = gate_tuples(task_data, gates) + members = sorted({m for g in gates for m in g.members}) + hull = derived_envelope(tuples, members) + graphs = list(hull.graph.batch_sizes) + envelope: dict[str, Any] = { + "max_batch_size": hull.max_batch_size or ENGINE_DEFAULTS["max_batch_size"], + "max_num_tokens": hull.max_num_tokens or ENGINE_DEFAULTS["max_num_tokens"], + "max_seq_len": hull.max_seq_len or 0, + "cuda_graph": {"batch_sizes": graphs} if graphs else False, + "members": members, + } + matrix: list[dict[str, Any]] = [] + for g in gates: + matrix.append( + { + "name": f"{g.unit}-gates", + "features": list(g.features), + "members": list(g.members), + "evidence": sorted(g.kinds - {EVIDENCE_WITNESS}), + "at": AT_GATE, + } + ) + # The witness runs at one member; a cell claiming it at the others + # would be owed as external evidence by the block meant to change + # nothing. + if g.witness_member is not None: + matrix.append( + { + "name": f"{g.unit}-witness", + "features": list(g.features), + "members": [g.witness_member], + "evidence": [EVIDENCE_WITNESS], + "at": AT_GATE, + } + ) + body = yaml.safe_dump( + {SERVING_ENVELOPE_FIELD: envelope, FEATURE_MATRIX_FIELD: matrix}, + sort_keys=False, + default_flow_style=None, + ) + header = ( + "# ---------------------------------------------------- serving envelope\n" + "# Generated from this spec's gate tuples: adopting it unchanged certifies\n" + "# exactly what the gates ran -- the legacy certified set, now written down.\n" + "# WIDEN IT to what the target is deployed at (max_num_tokens, max_seq_len,\n" + "# max_batch_size, the capture list, every member including multi-node and\n" + "# PP ones). Every limit raised here is probed at the ceiling by the run.\n" + ) + if not hull.max_seq_len: + header += ( + "# max_seq_len: no gate states one, so the engine used the checkpoint's\n" + "# max_position_embeddings. Fill it in before using this block.\n" + ) + return header + body + + +def _main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser( + prog="python -m agent_flow.workflows.staircase_onboard.envelope", + description="Serving-envelope tools for staircase-onboard task specs.", + ) + sub = parser.add_subparsers(dest="command", required=True) + migrate = sub.add_parser( + "migrate", + help="print a serving_envelope + feature_matrix block derived from the gate tuples", + ) + migrate.add_argument("task", help="path to task.yaml") + migrate.add_argument( + "--world-size", + type=int, + default=None, + help="world size of a unified run (default: from the member set)", + ) + migrate.add_argument( + "--write", + action="store_true", + help="append the block to the task file instead of printing it", + ) + plan = sub.add_parser("plan", help="print the certification plan for a task.yaml") + plan.add_argument("task", help="path to task.yaml") + plan.add_argument("--world-size", type=int, default=None) + args = parser.parse_args(argv) + + from .task_schema import load_and_validate_task_yaml + + try: + data = load_and_validate_task_yaml(args.task) + except TaskSchemaError as exc: + sys.stderr.write(f"{exc}\n") + return 2 + world_size = args.world_size + if world_size is None: + segment = str(data.get("parallel")) + _, size = parse_parallel_segment(segment) + if size is None: + _, size = parse_parallel_segment(members_for(data, None, 4)[0]) + world_size = size or 4 + + if args.command == "plan": + result = derive_certification(data, world_size=world_size) + sys.stdout.write(result.render()) + for error in result.errors: + sys.stderr.write(f"ERROR: {error}\n") + return 1 if result.errors else 0 + + if data.get(SERVING_ENVELOPE_FIELD) is not None: + sys.stderr.write(f"{args.task} already declares '{SERVING_ENVELOPE_FIELD}'\n") + return 1 + block = migration_block(data, world_size=world_size) + if args.write: + path = Path(args.task) + text = path.read_text(encoding="utf-8") + path.write_text(text.rstrip("\n") + "\n\n" + block, encoding="utf-8") + sys.stdout.write(f"appended a serving_envelope block to {path}\n") + else: + sys.stdout.write(block) + return 0 + + +__all__ = [ + "AT_CEILING", + "AT_GATE", + "CELL_EXTERNAL", + "CELL_GATED", + "CELL_PROBE", + "CUDA_GRAPH_KEYS", + "ENGINE_DEFAULTS", + "ENVELOPE_AXES", + "ENVELOPE_KEYS", + "EVIDENCE_DIRNAME", + "EVIDENCE_KINDS", + "CertificationPlan", + "EngineTuple", + "EnvelopeMember", + "GraphPolicy", + "MatrixCell", + "PROBE_GATE", + "ServingEnvelope", + "UnitGates", + "check_external_evidence", + "compact_values", + "default_unit_kinds", + "derive_certification", + "derived_envelope", + "envelope_from_spec", + "features_through", + "gate_tuples", + "generate_graph_batch_sizes", + "migration_block", + "options_tuple", + "parse_envelope_member", + "probe_flags", + "probe_prompts", + "probe_task_data", + "resolve_graph_policy", + "resolve_unit_gates", + "run_envelope_probe", + "runtime_tuple", + "validate_certification_blocks", + "validate_feature_matrix", + "validate_serving_envelope", +] + + +if __name__ == "__main__": + sys.exit(_main()) diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/preflight.py b/agent-flow/agent_flow/workflows/staircase_onboard/preflight.py index 56b1820f0451..fc0050d0b541 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/preflight.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/preflight.py @@ -36,15 +36,19 @@ from __future__ import annotations +import ast import json import shlex from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable, Mapping, Sequence +import yaml + from agent_flow.runtime import CommandSpec, LocalCommandExecutor from .delivery import MODELING_V2_PACKAGE +from .envelope import CertificationPlan, derive_certification from .environment import RunEnvironment, resolve_environment from .task_schema import ( assembly_units, @@ -180,12 +184,30 @@ class PreflightResult: target_rel: str = "" checkpoint_config: dict[str, Any] = field(default_factory=dict) notes: list[str] = field(default_factory=list) + #: What the run certifies: the serving envelope, the gate tuples inside + #: it, the envelope probes and the feature-matrix cells. See `envelope.py`. + certification: CertificationPlan | None = None class PreflightError(RuntimeError): """A precondition failed. The run ends here, before the branch exists.""" +class SpecInconsistencyError(PreflightError): + """The spec cannot run as written; the host is not at fault. + + Raised where preflight finds the spec's own blocks disagreeing -- a gate + launched outside the declared serving envelope -- or asking the tree for + something it does not have, so the operator is pointed at the spec or the + tree rather than at the infrastructure troubleshooting guide every other + preflight failure belongs in. + """ + + +class UnsupportedSpeculationError(SpecInconsistencyError): + """A gate would ask a launch path for a speculation algorithm it has no branch for.""" + + #: The one process-wide local executor for probes. Stateless. _LOCAL_EXECUTOR = LocalCommandExecutor() @@ -476,6 +498,184 @@ def check_accuracy_anchors( return notes +#: The launcher the boot, acceptance, witness and envelope-probe gates run +#: (see ``gates.run_boot_gate``), and the module whose ``SpeculativeConfig`` +#: decides what ``trtllm-eval`` and the LLM API accept, both repo-relative. +QUICKSTART_REL = "examples/llm-api/quickstart_advanced.py" +LLM_ARGS_REL = "tensorrt_llm/llmapi/llm_args.py" + +#: The variable ``quickstart_advanced.py`` dispatches its algorithms on. +_QUICKSTART_ALGO_NAME = "spec_decode_algo" +_SPECULATIVE_UNION = "SpeculativeConfig" + + +def quickstart_spec_algorithms(path: Path | str) -> set[str] | None: + """The ``--spec_decode_algo`` values the quickstart script has a branch for. + + Read from the script's own ``spec_decode_algo == ""`` comparisons, + upper-cased the way the script upper-cases its argument. None when the + file is absent or does not parse, which means "cannot tell", not "none". + """ + try: + tree = ast.parse(Path(path).read_text(encoding="utf-8")) + except (OSError, SyntaxError): + return None + found: set[str] = set() + for node in ast.walk(tree): + if not isinstance(node, ast.Compare): + continue + operands = [node.left, *node.comparators] + if not any(isinstance(o, ast.Name) and o.id == _QUICKSTART_ALGO_NAME for o in operands): + continue + for operand in operands: + constants = ( + operand.elts if isinstance(operand, (ast.Tuple, ast.List, ast.Set)) else [operand] + ) + for constant in constants: + if isinstance(constant, ast.Constant) and isinstance(constant.value, str): + found.add(constant.value.upper()) + return found + + +def llm_api_decoding_types(path: Path | str) -> set[str] | None: + """The ``speculative_config.decoding_type`` values the LLM API accepts. + + The ``decoding_type`` literals of the classes in ``SpeculativeConfig``'s + union; every decoding config's literal when the union cannot be found. + None when the file is absent or does not parse. + """ + try: + tree = ast.parse(Path(path).read_text(encoding="utf-8")) + except (OSError, SyntaxError): + return None + literals: dict[str, str] = {} + members: set[str] = set() + for node in tree.body: + if isinstance(node, ast.ClassDef): + for item in node.body: + if ( + isinstance(item, ast.AnnAssign) + and isinstance(item.target, ast.Name) + and item.target.id == "decoding_type" + and isinstance(item.annotation, ast.Subscript) + ): + value = item.annotation.slice + if isinstance(value, ast.Constant) and isinstance(value.value, str): + literals[node.name] = value.value + target = None + if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): + target, value = node.target.id, node.value + elif isinstance(node, ast.Assign) and len(node.targets) == 1: + first = node.targets[0] + target, value = (first.id if isinstance(first, ast.Name) else None), node.value + if target == _SPECULATIVE_UNION and value is not None: + members = {n.id for n in ast.walk(value) if isinstance(n, ast.Name)} + if members: + return {literal for name, literal in literals.items() if name in members} + return set(literals.values()) + + +def check_speculation_support(task_data: Mapping[str, Any]) -> list[str]: + """Every speculation algorithm the gates will emit has a branch in the launcher. + + Two launch paths, each asked differently: + + - the quickstart gates (boot, acceptance, witness, envelope probe) pass + ``--spec_decode_algo=`` from ``gates.speculation_flags``. A value + the script has no branch for falls through to ``spec_config = None``: + the gate runs the unit WITHOUT speculation, its acceptance statistic is + missing or 1.0, and the failure reads as a broken draft path; + - the accuracy gate's ``trtllm-eval`` reads the unit's options YAML, whose + ``speculative_config.decoding_type`` must be one the LLM API's + ``SpeculativeConfig`` union accepts. + + Raises :class:`UnsupportedSpeculationError` naming the algorithm, the + unit, the file and the branch it lacks. Returns notes, including one for a + launcher this preflight could not read. + """ + from .envelope import features_through + from .gates import resolve_llm_api_options, speculation_flags + + repo = Path(str(task_data.get("trtllm_repo_path") or ".")) + accuracy = task_data.get("accuracy") or {} + emitted: dict[str, list[str]] = {} + asked: dict[str, list[tuple[str, str]]] = {} + for unit in assembly_units(task_data): + for flag in speculation_flags(features_through(task_data, unit)): + if flag.startswith("--spec_decode_algo="): + emitted.setdefault(flag.split("=", 1)[1].upper(), []).append(unit) + options = resolve_llm_api_options(accuracy, unit) if isinstance(accuracy, dict) else None + if options and Path(options).is_file(): + try: + loaded = yaml.safe_load(Path(options).read_text(encoding="utf-8")) or {} + except yaml.YAMLError: + continue + config = loaded.get("speculative_config") if isinstance(loaded, dict) else None + decoding = config.get("decoding_type") if isinstance(config, dict) else None + if decoding: + asked.setdefault(str(decoding), []).append((unit, str(options))) + if not emitted and not asked: + return [] + + notes: list[str] = [] + errors: list[str] = [] + verified: set[str] = set() + if emitted: + quickstart = repo / QUICKSTART_REL + supported = quickstart_spec_algorithms(quickstart) + if supported is None: + notes.append( + f"speculation: cannot read {quickstart}, so whether it handles " + f"{', '.join(sorted(emitted))} is unchecked; the boot gate will say" + ) + else: + for algo, units in sorted(emitted.items()): + if algo in supported: + verified.add(algo) + continue + errors.append( + f"unit(s) {', '.join(units)} make the boot/acceptance/witness/probe gates " + f"pass `--spec_decode_algo={algo}`, but {quickstart} has no " + f"`{_QUICKSTART_ALGO_NAME} == '{algo}'` branch (it handles " + f"{', '.join(sorted(supported)) or 'none'}); the value falls through to " + f"`spec_config = None` and every quickstart gate would run the unit " + f"without speculation. Add the {algo} branch to {QUICKSTART_REL} on this " + f"tree before the run" + ) + if asked: + llm_args = repo / LLM_ARGS_REL + accepted = llm_api_decoding_types(llm_args) + if accepted is None: + notes.append( + f"speculation: cannot read {llm_args}, so the decoding types " + f"{', '.join(sorted(asked))} are unchecked" + ) + else: + lowered = {a.lower() for a in accepted} + for decoding, uses in sorted(asked.items()): + if decoding.lower() in lowered: + verified.add(decoding) + continue + where = "; ".join(f"{unit} ({options})" for unit, options in uses) + errors.append( + f"the accuracy gate's options ask trtllm-eval for " + f"`speculative_config.decoding_type: {decoding}` ({where}), but " + f"`{_SPECULATIVE_UNION}` in {llm_args} accepts only " + f"{', '.join(sorted(accepted)) or 'none'}. Add the {decoding} config to " + f"that union in {LLM_ARGS_REL}, or fix the options file" + ) + if errors: + raise UnsupportedSpeculationError( + "a speculation algorithm the gates will ask for has no branch on this tree:\n - " + + "\n - ".join(errors) + ) + if verified: + notes.append( + f"speculation: {', '.join(sorted(verified))} supported by the gates' launch paths" + ) + return notes + + def check_tmpdir(env: Mapping[str, str]) -> None: """Assert ``TMPDIR`` and ``CLAUDE_CODE_TMPDIR`` are set, short, and not ``/tmp``. @@ -831,11 +1031,16 @@ def run_preflight( *, runner: CommandRunner = _default_runner, skip_device_claim: bool = False, + unit_kinds: Callable[[str], Any] | None = None, ) -> PreflightResult: """Assert every precondition in order and resolve the run's identity. ``skip_device_claim`` exists for the CPU-only tests and for a ``--dry-run`` on a machine with no GPUs; it is never set by a real run. + + ``unit_kinds`` answers which gate kinds (boot, accuracy, witness, + acceptance) the orchestrator runs for a unit; without it they are read + from what the spec declares. See :func:`~.envelope.derive_certification`. """ notes: list[str] = [] @@ -890,6 +1095,29 @@ def run_preflight( # the unit list is decided here. notes.extend(check_accuracy_anchors(task_data, units)) + # The certified set, derived before anything touches the host: it is file + # reading, and a gate that launches outside the declared envelope -- or a + # gated member the envelope does not claim -- is a spec inconsistency the + # run must not start with. Everything it derives (probes, external matrix + # cells) is printed here so the bill is readable on a dry run. + certification = derive_certification( + task_data, + world_size=world_size, + gpus_per_node=_int_or_none((task_data.get("environment") or {}).get("gpus_per_node")), + checkpoint_config=checkpoint_config, + unit_kinds=unit_kinds, + ) + if certification.errors: + raise SpecInconsistencyError( + "the serving envelope and the gates disagree:\n - " + + "\n - ".join(certification.errors) + ) + notes.extend(certification.notes()) + # File reading too, and a spec that asks a launcher for a speculation + # algorithm it cannot build would otherwise surface at the first boot + # gate as a unit that "drafts nothing". + notes.extend(check_speculation_support(task_data)) + env_block = dict(task_data.get("environment") or {}) declared_env = {str(k): str(v) for k, v in (env_block.get("env") or {}).items()} check_tmpdir(declared_env) @@ -934,9 +1162,17 @@ def run_preflight( target_rel=target_rel, checkpoint_config=checkpoint_config, notes=notes, + certification=certification, ) +def _int_or_none(value: Any) -> int | None: + try: + return int(value) if value is not None else None + except (TypeError, ValueError): + return None + + def resolve_target_identity( task_data: Mapping[str, Any], checkpoint_config: Mapping[str, Any], @@ -1000,6 +1236,11 @@ def format_preflight_report(result: PreflightResult) -> str: "UNIX_SOCKET_PATH_MAX", "PreflightError", "PreflightResult", + "SpecInconsistencyError", + "UnsupportedSpeculationError", + "check_speculation_support", + "llm_api_decoding_types", + "quickstart_spec_algorithms", "UNSUPPORTED_ARCHITECTURE_MARKERS", "check_accuracy_anchors", "check_checkpoint", diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/progress.py b/agent-flow/agent_flow/workflows/staircase_onboard/progress.py index a47fa06c34bb..4fef63083d86 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/progress.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/progress.py @@ -46,6 +46,8 @@ from agent_flow.console import print_layer_panel from agent_flow.logger import get_logger +from .remediation import CERTIFICATION_GAP_KEY, normalize_certification_gap_records + PLAN_STAGE = "plan_stage" BUILD_STAGE = "build_stage" DELIVER_STAGE = "deliver_stage" @@ -72,7 +74,13 @@ # Staircase's triage table, as an enum. Each maps to a different # orchestrator action, which is the whole reason it is not free text. -BLOCKED_REASONS = ("vocabulary_gap", "knowledge_gap", "accuracy_shortfall", "environment") +BLOCKED_REASONS = ( + "vocabulary_gap", + "knowledge_gap", + "accuracy_shortfall", + "environment", + "certification_gap", +) def _empty() -> dict[str, list[dict[str, Any]]]: @@ -317,14 +325,19 @@ async def handler(args: dict[str, Any]) -> dict[str, Any]: # extending it. Measured on the first real onboarding, on its largest # item. "smith": ("entry", "receipt", "constituents"), - "assembler": ("blocked_reason", "missing_entry"), + "assembler": ("blocked_reason", "missing_entry", CERTIFICATION_GAP_KEY), + # `certification_gap`: "this entry is correct and must be certified at + # these values". The orchestrator fans each row out to a verify smith + # itself (remediation.py), so a certification-only finding no longer + # re-opens the assembler or waits on a human to relay it. + "reviewer": (CERTIFICATION_GAP_KEY,), # `uncovered` is to the auditor what `constituents` is to a MACRO: the # difference between a verdict the orchestrator can act on and one a # human has to read. Prose cannot be deduplicated, grouped by entry # surface, or compared across rounds -- so an UNCOVERED whose findings # live only in `summary` ends the run, which is exactly what it used to # do. See remediation.py. - "auditor": ("uncovered",), + "auditor": ("uncovered", CERTIFICATION_GAP_KEY), } #: Keys an uncovered row may carry. `entry`, `arch` and `axis` are required: @@ -540,12 +553,50 @@ def collect_decision_file( rows = parse_uncovered_records(value) if rows: entry[key] = rows + elif key == CERTIFICATION_GAP_KEY: + gaps = normalize_certification_gap_records(value) + if gaps: + entry[key] = gaps elif str(value).strip(): entry[key] = str(value).strip() ctx.append(agent, entry) return decision +def _certification_gap_property(who: str) -> dict[str, Any]: + """The ``certification_gap`` tool property, worded for one role.""" + return { + "type": "array", + "description": ( + f"{who} One row per catalog entry that is CORRECT but not certified at " + "values the target reaches inside the serving envelope. The orchestrator " + "dispatches a verify smith per row itself -- no human relays it -- so a " + "certification-only finding never needs a REJECT or a BLOCKED." + ), + "items": { + "type": "object", + "properties": { + "entry": {"type": "string", "description": "Catalog entry, e.g. `moe/runner`."}, + "values": { + "type": "array", + "items": {"type": "string"}, + "description": "The values to certify (row counts, sizes, ...).", + }, + "axis": { + "type": "string", + "description": "The certification axis; default `envelope`.", + }, + "arch": {"type": "string", "description": "Default: this run's architecture."}, + "reason": { + "type": "string", + "description": "Why the target needs it: the call site, the contract line.", + }, + }, + "required": ["entry", "values"], + }, + } + + def build_progress_tools(ctx: ProgressContext) -> dict[str, list[Any]]: """Build the per-role tool lists for ``BackendConfig(tools=...)``.""" plan_drafter = _decision_tool( @@ -633,7 +684,10 @@ def build_progress_tools(ctx: ProgressContext) -> dict[str, list[Any]]: "shortfall that survived your self-heal — the run ends with the " "audit trail, because the gate is the release criterion and it is " "not yours to relax. `environment`: the launch prefix or the host " - "cannot do what the task needs." + "cannot do what the task needs. `certification_gap`: you cannot " + "go on until entries are certified at values the target needs — " + "list them in `certification_gap`; the orchestrator verifies them " + "and resumes you." ), }, "missing_entry": { @@ -643,8 +697,12 @@ def build_progress_tools(ctx: ProgressContext) -> dict[str, list[Any]]: "op if you found one. This becomes smith's dispatch brief." ), }, + CERTIFICATION_GAP_KEY: _certification_gap_property( + "Optional with UNIT_READY (the verifies run while the unit is gated), " + "required with BLOCKED/certification_gap." + ), }, - extra_keys=("blocked_reason", "missing_entry"), + extra_keys=("blocked_reason", "missing_entry", CERTIFICATION_GAP_KEY), ) reviewer = _decision_tool( ctx, @@ -657,9 +715,17 @@ def build_progress_tools(ctx: ProgressContext) -> dict[str, list[Any]]: decision_help=( "APPROVE if the diff is consistent with the gate evidence and the " "invariants hold. REJECT for a defect the gates cannot see and a reader " - "can — an unmapped call, a stale op list, an uncovered runtime value, a " - "sync in the forward, a topology that does not match the segment." + "can — an unmapped call, a stale op list, a sync in the forward, a " + "topology that does not match the segment. A value outside an entry's " + "certified column on a CORRECT target is not a REJECT: APPROVE and list " + "it in `certification_gap`." ), + extra_properties={ + CERTIFICATION_GAP_KEY: _certification_gap_property( + "Certification-only findings, with APPROVE or REJECT." + ) + }, + extra_keys=(CERTIFICATION_GAP_KEY,), ) auditor = _decision_tool( ctx, @@ -725,9 +791,12 @@ def build_progress_tools(ctx: ProgressContext) -> dict[str, list[Any]]: }, "required": ["entry", "arch", "axis", "values"], }, - } + }, + CERTIFICATION_GAP_KEY: _certification_gap_property( + "Optional, and dispatched exactly like `uncovered`." + ), }, - extra_keys=("uncovered",), + extra_keys=("uncovered", CERTIFICATION_GAP_KEY), ) return { diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/assembler.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/assembler.py index 68a2384f9a41..e8dc8f25de52 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/assembler.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/assembler.py @@ -85,6 +85,12 @@ is the difference between an unsupported input and a wrong answer. 9. **Self-containment.** No shared helpers across targets. Duplication between targets is the architecture, not a smell. +10. **Every fixed size answers to an engine limit.** A constant that bounds + rows, tokens, batch or slots, or that sizes a buffer, is derived from an + engine argument at construction or declared `# capacity: ` / + `# capacity: fixed: ` on its line. The orchestrator audits each one + against the serving envelope; unbound, or below the envelope, fails review + and delivery. ## The runtime-value dump @@ -158,6 +164,14 @@ `accuracy_shortfall` ends the run with the audit trail, because guessing at the first or relaxing the second is exactly what must not happen quietly. +**A correct call outside its entry's certified column is a +`certification_gap`, not a judgement call and not a reason to stop.** List +each one -- `{{entry, values, axis, reason}}`, the reason naming the call site +-- in `certification_gap`. With `UNIT_READY` the orchestrator dispatches the +verify smiths and gates the unit; with `BLOCKED` and `blocked_reason: +certification_gap` (only when you genuinely cannot go on without the answer) +it verifies first and resumes you. Either way nobody has to relay it. + ## Your report The skill specifies it in full. At minimum: the target path; the **derivation @@ -188,7 +202,9 @@ | "A shared helper for future targets would be cleaner" | Self-containment is the architecture. Everything lives in the target dir. | | "Skip the boot gate, the accuracy gate is the real one" | The ladder exists to make iteration cheap. Boot first, always. | | "The fused kernel is missing; I can compose it from torch mirrors" | Composing math from mirrors is always a vocabulary gap. A *single* torch callable may be a mirror, but only once a closed search shows no engine op serves it. | +| "This value is outside the contract, so I'll block on it" | It is a `certification_gap` row. The orchestrator verifies it while the unit is gated; BLOCKED is for when you cannot write the unit without the answer. | | "I'm not sure this value is certified — flag it for verify" | Read the contract. It is written to be sufficient alone. Flag it only once you have established it is not covered. | +| "8192 rows is plenty for this buffer" | Plenty for the gates is not the envelope. Size it from `max_num_tokens` or declare `# capacity: max_num_tokens` and let the audit prove it. | | "Let me reformat before I report" | One lint pass, at delivery. Reformatting every turn is how a run spends its budget on whitespace. | | "Three failed fixes, but the fourth idea is good" | Stop and report. Three hypothesis-backed fixes without movement means the hypothesis space is wrong, and the next person needs the audit trail more than the fourth attempt. | """ diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/auditor.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/auditor.py index f109e7d5dd67..308739d6e3c1 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/auditor.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/auditor.py @@ -44,6 +44,16 @@ because every *constant* it passed was inside. That is the failure this turn exists to prevent, and it is the one that has actually happened. +**Captured rows come per call site.** An entry reached from several call sites +is reached with several row multiples: `G` at a decode step, `(1 + draft)·G` at +a verify step, a fraction of that on a compressed path, times the gather width +after an attention-DP all-gather. Derive each site's multiple from the call +chain and check `multiple x every captured G` against the certified column; +`certification.md` lists the captured batch sizes and the envelope. A site on +a path that is never captured (context rows only) is eager: check it against +the eager rows up to the envelope's token budget, not against capture. Name +the call site and its multiple in `contract_line` when you report one. + Two things that are **not** coverage: - a receipt for a different architecture. Receipts are keyed by GPU arch and @@ -100,6 +110,10 @@ A row missing `entry`, `arch` or `axis` cannot be dispatched and is dropped. +`certification_gap` (`{{entry, values, axis, reason}}`, arch and axis +defaulted) is accepted too and dispatched exactly like `uncovered`; use +`uncovered` when you have the contract line, which is the better brief. + ## What is not yours - performance: measured, never gated, and not part of onboarding diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_drafter.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_drafter.py index 55747de9dcd9..f3e3b57f7778 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_drafter.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_drafter.py @@ -116,6 +116,24 @@ assembled from several `catalog/torch/` mirrors is always a vocabulary gap, never a mirror to add. +## Judge coverage against the serving envelope + +Your dispatch prompt carries the run's **serving envelope** and the full plan is +in `certification.md`: the token budget per step, the batch limit, the sequence +length, the CUDA-graph capture batches and every member the target is claimed to +serve. A contract decides a row against **that box**, not against the +configuration some gate will happen to launch. A row whose certified column +covers the gate tuples and stops below the envelope is `CERT`, not `OK` — the +values above the gate tuples are exactly the ones a deployment reaches and no +gate does. When the envelope was DERIVED (the spec declares none), say in +`plan.md`'s open risks that nothing above the gate tuples is certified. + +`plan.md` carries a short **envelope** section: the envelope as you read it, +the members a single allocation cannot run (their evidence is owed from +outside), and any limit the checkpoint or the runtime cannot honour — a limit +that cannot be served is a plan change to raise now, not a finding for the +last gate. + ## The accuracy anchor The plan names the benchmark **and its protocol**. A reasoning model scored as @@ -158,6 +176,10 @@ entry: catalog//.py # OK and CERT rows contract_line: "" candidate: torch.ops.trtllm. # GAP rows + call_sites: # where the target calls the entry + - {{site: "", row_multiple: 1}} + - {{site: "<...>", row_multiple: 6, gathered: true, row_width: 512}} + - {{site: "<...>", eager_only: true}} queue: - {{n: 1, mode: onboard, target: , depends_on: []}} - {{n: 2, mode: verify, target: , depends_on: [], @@ -178,6 +200,16 @@ a whole run - **the two chapters do not overlap.** An entry smith onboards is tested at this checkpoint's values by construction, so it never also needs verifying +- `call_sites` lists every place the target will call the entry, with the + rows per request it passes at a captured step: 1 for a decode step, + `1 + draft` at a speculative verify step, a fraction of that on a + compressed path, the drafter's own width for a drafter call. `gathered: + true` marks a call after an attention-DP all-gather (it sees every rank's + rows); `eager_only: true` a call on a path that is never captured (context + rows only); `row_width` the width the entry's offsets are computed over. + Each smith brief multiplies these by every captured batch size, so a site + missing here is a row count no smith certifies and the delivery auditor + finds at the end of the run - `depends_on` is for real ordering only (a verify row whose test state must be built by an onboard row first). Everything without a dependency runs **concurrently**, so a spurious dependency directly costs wall-clock @@ -221,6 +253,7 @@ | "I'll mark everything depends_on the previous item to be safe" | Every spurious dependency serializes a batch that would have run concurrently. Order only what genuinely must be ordered. | | "The row's `note` says it is dspark-only" | Nothing reads notes. Put it in `unit:`, or one failed dispatch the speculator alone needed stops the run before core has started. Measured: six such rows in a 33-dispatch queue, each able to close the run. | | "The reviewer will catch it, ship the draft" | The reviewer gets one round. What it does not catch, the gates catch — expensively. | +| "The contract covers what the gates run, so the row is OK" | The certified set is the serving envelope. A column that stops at the gate tuples is CERT: the values above them are the ones a deployment reaches and no gate does. | | "The parallel segment doesn't change the math, so the mapping is the same" | Collectives appear nowhere in the HF reference. Walk the segment separately or the assembler finds the gap mid-run. | """ ) diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_reviewer.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_reviewer.py index 227544b97a4a..5bd92c0f7475 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_reviewer.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/plan_reviewer.py @@ -53,6 +53,11 @@ - `depends_on` orders only what genuinely must be ordered. **Every spurious dependency serializes a batch that would have run concurrently**, so this is a real cost, not a nicety. +- coverage was judged against the **serving envelope** (`certification.md`), + not against the gate tuples. An `OK` row whose quoted contract line certifies + only values at or below what the gates launch, while the envelope reaches + higher, is a wrong `OK` — REJECT it: it is the class of row that let a + fixed-size buffer and a 32-bit offset overflow through every gate once. - the parallel segment was walked separately. Collectives appear nowhere in the HF reference, so a mapping that is all-green for a `tepN`/`depN` target and never mentions a reduce or a gather was not derived, it was assumed. @@ -97,6 +102,7 @@ | "The plan is thin on X — REJECT" | Would that waste an assembly turn or produce a wrong gate? If not, it is a summary note. You have one round. | | "The ceiling verdict looks reasonable" | It is the one verdict that ends the run. Check the catalog *and* the inventory *and* the uncertified surfaces before accepting it. | | "Dependencies are harmless, they just order things" | Each one serializes a batch. Spurious ordering is measured in wall-clock. | +| "The row is OK — its contract covers the gate configuration" | The certified set is the envelope. Coverage up to the gate tuples and not beyond is a CERT row mislabelled. | | "The anchor number looks plausible" | Plausible is not sourced. Ask which file it came from. | | "APPROVE, and I'll note my concerns nowhere" | The summary reaches the assembler. Silence throws away the only cheap part of this review. | """ diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/reviewer.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/reviewer.py index 733e26446a46..70d22039f695 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/reviewer.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/reviewer.py @@ -32,10 +32,11 @@ 2. **`REQUIRED_TRTLLM_OPS` equals the forward's call set** — not a superset, not a subset. A stale name is as wrong as a missing one. 3. **The runtime-value dump against the contracts.** Any value the engine - varied at runtime that sits outside its entry's certified column is a - blocker, not a note. This is the check that has been missed before: a - target merged with 34 of the 35 batch sizes its engine captured sitting - outside the contract, because every *constant* it passed was inside. + varied at runtime that sits outside its entry's certified column must be + certified before delivery -- see "Rejecting" for how to report it. This is + the check that has been missed before: a target merged with 34 of the 35 + batch sizes its engine captured sitting outside the contract, because + every *constant* it passed was inside. 4. **Graph safety and sync-freedom in the forward.** No `.item()`, no `.cpu()`, no data-dependent host branching. These are writing rules — a forward that violates them is wrong the first time it is captured, and the @@ -48,6 +49,13 @@ path declares. 7. **For a feature increment: the base path is unchanged.** With the feature off, the forward must be what the recorded gate measured, bit for bit. +8. **Capacity constants against the serving envelope.** Your prompt carries + the orchestrator's static audit of every fixed size in the target. A FAIL + row is a blocker. An ok row is a CLAIM: read the code and confirm the + constant really caps the axis it is bound to (a row ceiling bound to + `max_batch_size` is wrong), and that a `fixed:` reason is true. No gate + launches at the envelope's ceiling except the envelope probe, so this + reading is where a too-small buffer is caught. ## What not to do @@ -63,8 +71,18 @@ ## Rejecting REJECT for a defect the gates cannot see and a reader can: an unmapped call, a -stale op list, an uncovered runtime value, a sync in the forward, a topology -that does not match the segment, a feature that changed the base path. +stale op list, a sync in the forward, a topology that does not match the +segment, a feature that changed the base path, a capacity constant below the +envelope. + +**A runtime value outside an entry's certified column, on a target that is +otherwise correct, is a certification gap, not a REJECT.** List it in +`certification_gap` -- `{{entry, values, axis, reason}}`, one row per entry +surface -- and APPROVE. The orchestrator dispatches the verify smiths itself +and the unit moves on; a REJECT for it re-opens the assembler for hours over +work that is not the assembler's. Judge "outside" against the serving +envelope in your prompt: a value inside the envelope that the contract's +certified column covers by the envelope standard is not a gap at all. APPROVE when the diff is consistent with the gate evidence and the invariants hold. A gate that passed plus a clean audit is sufficient — you are not @@ -94,7 +112,9 @@ | "The whole file should be scanned for non-catalog calls" | That false-positives on init-time code and buries the real finding. Scope it to `forward` plus what it reaches. | | "The constants are all certified, so we're clean" | The values that escape are the ones the engine varies at runtime. Read the dump. | | "Both gates passed, so the topology is what the path says" | Every gate is topology-blind. Only reading the assembly against the segment catches it. | +| "The constant is bound, so it is fine" | A binding is a claim the audit takes at its word. Check it names the axis the constant actually caps. | | "It would read better formatted" | One lint pass, at delivery. Not yours. | +| "REJECT — this value is not certified" | If the code is right, it is a `certification_gap` row with an APPROVE. Verifies run without the assembler. | | "REJECT — I'd have written this differently" | Name the invariant it violates, or APPROVE. | """ ) diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/smith.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/smith.py index 51d646410a94..03fac5a2a883 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/smith.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/smith.py @@ -83,6 +83,15 @@ rather than by recollection. A receipt that predated its own test file has shipped here before, and it read as green. +## The smith brief + +Your dispatch prompt carries a **smith brief**: the run's serving envelope, +every call site of this entry with the rows per request it passes and the +captured row counts that follow from them, and the integer-width boundaries to +probe. The brief is the surface. Certify all of it, record the probes as hard +limits in the contract, and say in your report which call sites you confirmed +against the call chain when the brief says they were assumed. + ## Tolerance Never inflate one to manufacture a pass. If a comparison needs a looser @@ -154,6 +163,7 @@ | "The contract can say 'see source'" | A caller must be able to use the op from the contract alone. Sending them to the source makes the contract a pointer instead of an interface. | | "I'll onboard the two ops this macro needs myself" | One run, one entry. Decompose, land the report, then dispatch one child per constituent. | | "The test passed, so the receipt is good" | Only if it post-dates every file of the entry — in both trees. Lint and index edits happen after the test ran. | +| "The envelope never reaches 2^31 / width rows, so no probe" | The probe records a hard limit for the day a memory cap is lifted. Run it whatever today's envelope is. | | "I'll name the test after the op, like upstream does" | `test_modeling_v2_.py`, always. The name is what keeps it from colliding with upstream's test of the same op and what the CI selector matches. | """ ) diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/remediation.py b/agent-flow/agent_flow/workflows/staircase_onboard/remediation.py index 4435bdafa036..aaeab51c2207 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/remediation.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/remediation.py @@ -208,6 +208,116 @@ def parse_uncovered_rows(raw: Any) -> tuple[tuple[UncoveredRow, ...], int]: return tuple(rows), dropped +# --------------------------------------------------------- certification gaps + +#: The typed decision field an assembler, a reviewer or an auditor uses to say +#: "this entry is correct, but it must be certified at these values". Before +#: it existed the only typed outcomes were a target REJECT (which re-opens the +#: assembler for hours over what is a catalog item) or a run-ending BLOCKED, +#: and every verify was relayed by a human. The orchestrator now fans each +#: gap out to verify smiths itself, through the same queue the audit uses. +CERTIFICATION_GAP_KEY = "certification_gap" + +#: Roles whose decision may carry a gap. +CERTIFICATION_GAP_ROLES: tuple[str, ...] = ("assembler", "reviewer", "auditor") + +#: The axis a gap row gets when it names none: "the serving envelope", which is +#: what an entry is certified over. +DEFAULT_GAP_AXIS = "envelope" + +#: Keys a gap row may carry. Only `entry` is required: the arch defaults to the +#: run's, and the axis to the envelope. +CERTIFICATION_GAP_FIELDS: tuple[str, ...] = ( + "entry", + "values", + "value", + "axis", + "arch", + "reason", + "contract_line", +) + + +def normalize_certification_gap_records(raw: Any) -> list[dict[str, Any]]: + """Normalize a decision's gap list for persistence, dropping rows with no entry. + + Kept as plain dicts: this is the persistence boundary, and + :func:`parse_certification_gaps` turns them into typed rows. + """ + if not isinstance(raw, list): + return [] + rows: list[dict[str, Any]] = [] + for row in raw: + if not isinstance(row, Mapping): + continue + entry = str(row.get("entry", "")).strip() + if not entry: + continue + record: dict[str, Any] = {"entry": entry} + values = row.get("values") + if values is None: + values = row.get("value") + if isinstance(values, (str, int, float)): + values = [values] + if isinstance(values, Sequence) and not isinstance(values, (str, bytes)): + record["values"] = [str(v).strip() for v in values if str(v).strip()] + for key in ("axis", "arch", "reason", "contract_line"): + text = str(row.get(key, "") or "").strip() + if text: + record[key] = text + rows.append(record) + return rows + + +def parse_certification_gaps( + raw: Any, *, default_arch: str +) -> tuple[tuple[UncoveredRow, ...], int]: + """Turn a decision's ``certification_gap`` list into dispatchable rows. + + Returns the rows and how many were dropped (no entry), so the caller can + say so. The reason travels as the contract line: it is what the verify + smith reads to know why the surface is needed. + """ + if not isinstance(raw, Sequence) or isinstance(raw, (str, bytes)): + return (), 0 + rows: list[UncoveredRow] = [] + dropped = 0 + for candidate in raw: + if not isinstance(candidate, Mapping) or not str(candidate.get("entry", "")).strip(): + dropped += 1 + continue + reason = str(candidate.get("reason", "") or candidate.get("contract_line", "")).strip() + row = UncoveredRow.parse( + { + "entry": candidate.get("entry"), + "arch": str(candidate.get("arch", "") or default_arch), + "axis": str(candidate.get("axis", "") or DEFAULT_GAP_AXIS), + "values": candidate.get("values", candidate.get("value")), + "contract_line": reason, + } + ) + if row is None: + dropped += 1 + else: + rows.append(row) + return tuple(rows), dropped + + +def parse_audit_rows( + entry: Mapping[str, Any], *, default_arch: str +) -> tuple[tuple[UncoveredRow, ...], int]: + """An auditor's ``uncovered`` rows and its ``certification_gap`` rows, together. + + The two say the same thing -- a surface outside its certified column -- + and are dispatched by the same cycle. + """ + uncovered, dropped = parse_uncovered_rows(entry.get("uncovered")) + gaps, gap_dropped = parse_certification_gaps( + entry.get(CERTIFICATION_GAP_KEY), default_arch=default_arch + ) + return (*uncovered, *gaps), dropped + gap_dropped + + # ---------------------------------------------------------------------- items @@ -578,6 +688,61 @@ def decide_round( ) +def decide_gap_dispatch( + merge: MergeOutcome, + *, + budget: RemediationBudget, + incoming: Sequence[RemediationItem], +) -> RoundDecision: + """What to do with the certification gaps one decision reported. + + Not :func:`decide_round`: a gap arrives mid-assembly, not from an audit + round, so it spends no round budget and there is no previous signature to + converge against. The rules that remain: + + - a surface already certified is not re-dispatched -- the reporter misread + a contract, and the detail says so -- and does not stop anything; + - a surface whose verify already failed, or that is over the item budget, + cannot be dispatched again: ``stop``, because the unit needs a surface + the catalog cannot give it and that is a plan change; + - everything new or widened is dispatched, capped at + ``max_parallel_verifiers`` concurrent smiths. + """ + wanted = {item.group_key for item in incoming} + dispatch = tuple( + item for item in merge.queue if item.group_key in wanted and item.status in OPEN_STATES + ) + notes: list[str] = [] + if merge.regressions: + names = ", ".join(f"{i.entry}[{i.axis}]" for i in merge.regressions[:4]) + notes.append( + f"{len(merge.regressions)} already certified by an earlier verify, not re-dispatched " + f"({names}): re-read the contract" + ) + if merge.duplicates: + notes.append(f"{len(merge.duplicates)} already queued") + if merge.exhausted: + names = ", ".join(f"{i.entry}[{i.axis}]" for i in merge.exhausted[:4]) + return RoundDecision( + action="stop", + verdict=STOP_VERIFICATION_FAILED, + detail=( + f"{len(merge.exhausted)} gap(s) cannot be verified again -- a previous verify " + f"failed or the item budget is spent: {names}" + + (f"; {'; '.join(notes)}" if notes else "") + ), + ) + if not dispatch: + return RoundDecision(action="clean", detail="; ".join(notes) or "nothing to verify") + notes.insert(0, f"{len(dispatch)} verify dispatch(es)") + return RoundDecision( + action="dispatch", + dispatch=dispatch, + workers=max(1, min(budget.max_parallel_verifiers, len(dispatch))), + detail="; ".join(notes), + ) + + # ------------------------------------------------------------------- outcomes #: Smith decisions that certify a surface, and the queue state each maps to. @@ -660,6 +825,10 @@ def render(self) -> str: __all__ = [ + "CERTIFICATION_GAP_FIELDS", + "CERTIFICATION_GAP_KEY", + "CERTIFICATION_GAP_ROLES", + "DEFAULT_GAP_AXIS", "OPEN_STATES", "QUEUE_BLOCKED", "QUEUE_CERTIFIED", @@ -680,8 +849,12 @@ def render(self) -> str: "UncoveredRow", "apply_outcomes", "canonicalize", + "decide_gap_dispatch", "decide_round", "merge_into_queue", + "normalize_certification_gap_records", + "parse_audit_rows", + "parse_certification_gaps", "parse_uncovered_rows", "unresolved_signature", ] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/entry-spec.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/entry-spec.md index 87590f4fed99..1b36a006a780 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/entry-spec.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/entry-spec.md @@ -58,6 +58,12 @@ violating none of these must get a correct result. ## Notes Arch quirks observed in source or during testing. + +## Hard limits +The integer-width boundary of every index-typed argument (2^31 elements, +2^31 / row width rows, 2^24 + 1 through fp32), probed and recorded whether +or not today's envelope reaches it. Table format and probing recipe: +staircase-kernel-verify `references/integer-width-probes.md`. ``` ### `.py` — wrapper diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/SKILL.md index b329832d929e..82db861b6d27 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/SKILL.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/SKILL.md @@ -74,6 +74,53 @@ every existing caller is affected. constraint on every target that would have passed that value, and finding it here is far cheaper than finding it in an accuracy gate. +## The surface is the envelope, at every call site's row multiple + +Your dispatch prompt carries a **smith brief**: the run's serving envelope and, +for each place the target calls this entry, the rows per request it passes +(`G`, `3·G`, `6·G`, ...; times the gather width after an attention-DP +all-gather; `eager-only` for a site that is never captured). Certify what the +brief lists, not what one dump showed: + +- every captured row count it lists, replay on new data bitwise equal to eager; +- every eager row count up to the envelope's token budget, by the envelope + standard (a run at the maximum plus prefix-bitwise for grid-sized rows; + every bucket's edges for a kernel that picks by size); +- pools at the largest size the GPU allocates at the envelope's settings, with + rows addressed at the far end. + +The certified column you write says the envelope and the call-site multiples +it covers, so the auditor reads one sentence instead of re-deriving rows. A +brief that says its call sites were ASSUMED is a prompt to check them against +the call chain and say in your report which you found. + +## Integer-width boundary probes (mandatory) + +Every verify probes each **index-typed** argument -- offsets, indices, +positions, page and slot tables, row counts -- at the integer-width +boundaries, and records the outcome in the contract as a **hard limit**, +**whether or not today's envelope reaches it**: + +- **2^31 elements**: an int32 element count or offset wraps; +- **2^31 / row width rows**, for every row width the entry addresses: a row + index multiplied by the width in 32 bits wraps there (2^31 / 512 = 2^22 + rows), and the kernel writes out of place without raising; +- **2^24 + 1** for any index that passes through fp32: above 2^24 an fp32 + integer is no longer exact. + +The brief lists the row counts for this entry's widths and marks any that fall +inside the envelope. How to probe without the full allocation, and the table +format, are in `references/integer-width-probes.md`. + +The reason this is mandatory rather than envelope-driven: a contract certified +"up to the envelope of the day" stops being true the day another change +removes the memory cap that kept the envelope small. One kernel computed its +output offset in 32 bits and silently wrote rows at and above 2^25 out of +place; its contract was correct for the envelope it was certified at, and a +prefill-tiling fix that lifted a memory cap moved the reachable rows past it. +A hard limit in the contract is what makes that change visible: **any change +that removes a memory cap re-reads the hard limits of every entry it feeds.** + ## When the surface neighbours a known intermittent defect `entry-spec.md` states the rule and the case that produced it: prove the @@ -86,5 +133,6 @@ where a blind harness returns a confident, meaningless green. - entry path, the surface requested, and the verdict: already covered / newly certified / failed - the receipt as it now reads, and what the new coverage pinned +- the hard limits the integer-width probes established - **what remains uncertified** on this entry — the next caller reads this - every claim matched against disk before sending diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/references/integer-width-probes.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/references/integer-width-probes.md new file mode 100644 index 000000000000..31dd71d5e483 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/references/integer-width-probes.md @@ -0,0 +1,65 @@ +# Integer-width boundary probes + +Every verify and every onboard records, for each index-typed argument of the +entry, where 32-bit (and fp32) arithmetic stops being exact -- as a hard limit +in the contract, independent of the serving envelope. + +## Which arguments + +Anything the kernel turns into an address or an index: row offsets, cumulative +sequence lengths, token positions, page tables and slot indices, gather and +scatter indices, top-k indices, the row count itself. Read the wrapper's +signature and the kernel's launch code for the arithmetic: `row * width + col` +in `int`, an offset accumulated in `int32_t`, a `tl.int32` program-id product, +an index stored as `float` on its way through a reduction. + +## The three boundaries + +| boundary | what overflows | probe at | +|---|---|---| +| total elements 2^31 | an int32 element count or flat offset | 2^31 elements (or 2^31 + 1) | +| rows x width 2^31 | `row * width` in 32 bits, for each row width the entry addresses | `2^31 / width + 1` rows | +| fp32-exact 2^24 | an integer index carried through fp32 | index `2^24 + 1` | + +The smith brief computes the row counts for the widths the call sites declare +and says whether each lies inside the envelope. + +## Probing without the full footprint + +A 2^31-element probe of a bf16 tensor is 4 GiB, which most GPUs hold; where the +op's footprint does not fit, drive the boundary through the cheapest operand: + +- the smallest dtype the op accepts for the operand that carries the count; +- a page or slot table that points a small batch at far rows of a large pool, + rather than a large batch; +- strided or expanded views that give a large logical row count over little + memory, where the op accepts them; +- for the rows x width boundary, the narrowest width the op supports with the + real width's arithmetic, and the real width at the largest row count that + fits, both stated. + +Compare against a reference built from native torch ops on the rows at and +beyond the boundary, not only the first rows: the failure this exists for +wrote the rows past 2^25 out of place and left the rest correct. + +When a boundary genuinely cannot be reached on this machine, record it as +`unprobed` with the reason and the largest value that was probed: the hard +limit is then that value, not the boundary. + +## The contract table + +```markdown +## Hard limits + +| argument | arithmetic | boundary | probed at | result | receipt | +|---|---|---|---|---|---| +| `out` offset | `row * 512` in int32 | 2^31 / 512 = 4194304 rows | 4194305 rows | failed: rows >= 2^22 written out of place | sm_103 | +| `num_tokens` | element count, int64 | 2^31 elements | 2^31 + 1 | passed | sm_103 | +| `topk_idx` | through fp32 in the reduction | 2^24 | 16777217 | passed (int32 path, never fp32) | sm_103 | +``` + +A `failed` row is a hard limit every target must stay below, and a value of the +serving envelope above it is a blocker for the target, not for the entry. An +`unprobed` row limits the entry to what was probed. Rewrite the table only from +a run you watched, and re-read it whenever a change removes a memory cap that +kept the reachable rows small. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/SKILL.md index 272ad3abd48a..5620d908b77f 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/SKILL.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-core/SKILL.md @@ -73,6 +73,26 @@ These hold for every edit to a target's modeling code, by anyone, forever. is the difference between an unsupported input and a wrong answer. 9. **Self-containment.** No shared helpers across targets. Duplication between targets is the architecture, not a smell. +10. **Every fixed size answers to an engine limit.** A module- or class-level + integer that bounds something (rows, tokens, batch, slots, pages) or sizes + a buffer, and every buffer sized by a literal of 1024 or more, is either + derived from an engine argument at construction + (`model_config.max_num_tokens`, `max_batch_size`, `max_seq_len`) or + declared on its own line: + + ```python + _TOPK_MAX_ROWS = 16384 # capacity: max_num_tokens + _GATHER_ROWS = 65536 # capacity: max_num_tokens * world_size + _WINDOW = 128 # capacity: fixed: the checkpoint's sliding window + ``` + + (or in a module-level `CAPACITY_BOUNDS = {"NAME": ""}`). The + orchestrator audits every one statically against the serving envelope at + review and at delivery: unbound, or bound and below the envelope, fails. + The case that made this an invariant: a `_TOPK_MAX_ROWS = 8192` offsets + buffer killed every engine above `max_num_tokens` 8192 in warm-up while + every kernel it fed was certified to 16384, and no gate launched above + 8192 to see it. ## Products diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/smith_brief.py b/agent-flow/agent_flow/workflows/staircase_onboard/smith_brief.py new file mode 100644 index 000000000000..3defbeb76e8b --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/smith_brief.py @@ -0,0 +1,436 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The brief a smith dispatch carries: the entry's surface, per call site, over the envelope. + +A catalog entry is called from several places in a target, and each call site +reaches it with a different number of rows per request: a decode step passes +one row per request, a speculative verify step ``1 + draft`` rows, a +ratio-compressed path half of those, and an op after an attention-DP gather +sees every rank's rows. Under CUDA-graph capture each of those is multiplied +by every captured batch size ``G``. A smith certifying the entry at ``{G, +6G}`` -- "the decode rows and the verify rows" -- misses the call site that +passes ``3G`` and the one that passes ``5G``, and the delivery auditor, which +derives the multiple from the call chain, finds them uncovered at the end of +the run: seven verify dispatches to learn what the plan could have said. + +So the plan states each entry's call sites in ``vocabulary.yaml``:: + + steps: + - step: compressed-KV latent + entry: catalog/attention/compressor.py + call_sites: + - {site: "Compressor.forward, verify step", row_multiple: 3} + - {site: "MoE.forward after the DP gather", row_multiple: 6, gathered: true} + - {site: "_TailSlots, context rows only", eager_only: true} + +and this brief turns them into the rows to certify: ``m*G`` for every +captured ``G`` (times the gather width under attention DP), the eager rows up +to the envelope's token budget, ``eager-only`` sites marked so neither the +smith nor the auditor reasons about capture for them again -- plus the +integer-width probes every verify owes, whatever today's envelope reaches. + +``python -m agent_flow.workflows.staircase_onboard.smith_brief --workspace W +`` prints the brief for a dispatch already written, the way an +operator relaying a turn needs it. +""" + +from __future__ import annotations + +import argparse +import re +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + +import yaml + +from .envelope import CertificationPlan, compact_values + +#: The keys one call site may carry. +CALL_SITE_KEYS: tuple[str, ...] = ( + "site", + "row_multiple", + "gathered", + "eager_only", + "row_width", + "note", +) + +#: 32-bit signed element offsets overflow here. +INT32_LIMIT = 2**31 +#: Integers above this are not exactly representable in fp32. +FP32_EXACT_LIMIT = 2**24 + +_DISPATCH = re.compile(r"Dispatch (\d+): \*\*(\w+)\*\* the catalog entry `([^`]+)`") + + +@dataclass(frozen=True) +class CallSite: + """One place the target calls an entry, and how many rows per request it passes.""" + + site: str + row_multiple: int = 1 + gathered: bool = False + eager_only: bool = False + row_widths: tuple[int, ...] = () + note: str = "" + defaulted: bool = False + + +def entry_key(path: str) -> str: + """``catalog/attention/x.py`` / ``attention/x`` / ``.../catalog/attention/x`` -> ``attention/x``.""" + text = str(path).strip().removesuffix(".py").strip("/") + if "catalog/" in text: + text = text.split("catalog/", 1)[1] + return text + + +def _widths(value: Any) -> tuple[int, ...]: + if value is None: + return () + values = value if isinstance(value, list) else [value] + return tuple(int(v) for v in values if isinstance(v, int) and not isinstance(v, bool) and v > 0) + + +def parse_call_sites(step: Mapping[str, Any]) -> tuple[list[CallSite], list[str]]: + """The call sites one vocabulary step declares, and what is wrong with them.""" + raw = step.get("call_sites") + if raw is None: + return [], [] + where = f"step {step.get('step', '?')!r}" + if not isinstance(raw, list): + return [], [f"{where}: call_sites must be a list"] + sites: list[CallSite] = [] + errors: list[str] = [] + step_widths = _widths(step.get("row_width")) + for index, row in enumerate(raw): + if not isinstance(row, dict): + errors.append(f"{where}: call_sites[{index}] must be a mapping") + continue + unknown = sorted(set(row) - set(CALL_SITE_KEYS)) + if unknown: + errors.append(f"{where}: call_sites[{index}] has unknown key(s) {unknown}") + site = str(row.get("site", "")).strip() + if not site: + errors.append(f"{where}: call_sites[{index}] names no `site`") + continue + multiple = row.get("row_multiple", 1) + if not isinstance(multiple, int) or isinstance(multiple, bool) or multiple < 1: + errors.append(f"{where}: call_sites[{index}].row_multiple must be an integer >= 1") + continue + sites.append( + CallSite( + site=site, + row_multiple=multiple, + gathered=bool(row.get("gathered", False)), + eager_only=bool(row.get("eager_only", False)), + row_widths=_widths(row.get("row_width")) or step_widths, + note=str(row.get("note", "")).strip(), + ) + ) + return sites, errors + + +def steps_for(steps: Iterable[Mapping[str, Any]], target: str) -> list[Mapping[str, Any]]: + """The vocabulary steps an entry serves.""" + key = entry_key(target) + name = key.rsplit("/", 1)[-1] + found = [] + for step in steps: + entry = step.get("entry") + if not entry: + continue + candidate = entry_key(str(entry)) + if candidate == key or ("/" not in key and candidate.rsplit("/", 1)[-1] == name): + found.append(step) + return found + + +def default_call_sites(plan: CertificationPlan | None) -> list[CallSite]: + """What a brief assumes when the plan declared no call sites for the entry.""" + sites = [CallSite(site="decode step", row_multiple=1, defaulted=True)] + draft = plan.envelope.max_draft_len if plan is not None else 0 + if draft: + sites.append( + CallSite(site="speculative verify step", row_multiple=1 + draft, defaulted=True) + ) + return sites + + +def gather_widths(plan: CertificationPlan | None) -> list[tuple[str, int]]: + """``[(member, world), ...]`` for every attention-DP member the envelope claims.""" + if plan is None: + return [] + return sorted( + {(m.name, m.world_size) for m in plan.envelope.members if m.attention_dp}, + key=lambda item: (item[1], item[0]), + ) + + +def captured_rows(site: CallSite, batches: Sequence[int]) -> list[int]: + """Rows a captured call at this site sees, per rank, for every captured batch.""" + if site.eager_only: + return [] + return sorted({site.row_multiple * g for g in batches}) + + +def _multiple(m: int) -> str: + return "G" if m == 1 else f"{m}·G" + + +def integer_width_probes(widths: Iterable[int], *, reachable_rows: int | None) -> list[str]: + """The mandatory boundary probes, with the row counts they imply.""" + lines = [ + f"total elements at {INT32_LIMIT} (2^31): an int32 element offset or count wraps here", + f"any index that passes through fp32 at {FP32_EXACT_LIMIT + 1} (2^24 + 1): above " + f"2^24 an fp32 index is no longer exact", + ] + for width in sorted(set(widths)): + rows = INT32_LIMIT // width + reach = "" + if reachable_rows is not None: + reach = ( + " -- INSIDE the envelope's reachable rows" + if reachable_rows >= rows + else f" (the envelope reaches {reachable_rows} rows)" + ) + lines.append( + f"rows x width {width}: 2^31 / {width} = {rows} rows, probe at {rows + 1}{reach}" + ) + if not widths: + lines.append( + "rows x row width: 2^31 / rows, for every row width this entry " + "addresses (the call sites declare none -- read them from the contract)" + ) + return lines + + +def render_smith_brief( + item: Any, + *, + steps: Sequence[Mapping[str, Any]] = (), + plan: CertificationPlan | None = None, + queue_row: Mapping[str, Any] | None = None, +) -> str: + """The brief for one smith dispatch.""" + target = str(item.target) + lines = [f"## Smith brief: dispatch {item.n}, {item.mode} `{target}`", ""] + row = dict(queue_row or {}) + for key in ("tier", "unit", "depends_on", "surface", "note"): + value = row.get(key, getattr(item, key, None)) + if value not in (None, "", (), []): + lines.append(f"- {key}: {value}") + + batches = plan.certified_capture_batches if plan is not None else () + env = plan.envelope if plan is not None else None + if env is not None: + kind = "declared" if env.declared else "DERIVED from the gate tuples (legacy)" + lines += [ + "", + f"### Envelope ({kind})", + "", + f"- tokens per step per rank T <= {env.max_num_tokens}; requests per step <= " + f"{env.max_batch_size}; per-request length <= {env.max_seq_len}", + f"- captured batch sizes G: {compact_values(batches) or 'none (eager only)'}", + ] + + sites: list[CallSite] = [] + errors: list[str] = [] + for step in steps_for(steps, target): + found, problems = parse_call_sites(step) + sites.extend(found) + errors.extend(problems) + if not sites: + sites = default_call_sites(plan) + gathers = gather_widths(plan) + + lines += ["", "### Call sites and the rows each one passes", ""] + if any(site.defaulted for site in sites): + lines += [ + "`vocabulary.yaml` declares no `call_sites` for this entry, so these are " + "ASSUMED (decode, and the verify step when the envelope speculates). Confirm " + "them against the call chain; a site this list misses is what the delivery " + "auditor will find uncovered.", + "", + ] + lines += [ + "| call site | row multiple | captured rows per rank | after an attention-DP gather |", + "|---|---|---|---|", + ] + union: set[int] = set() + eager_rows: set[int] = set() + for site in sites: + if site.eager_only: + lines.append( + f"| {site.site} | eager-only | never captured | " + f"{'gathered' if site.gathered else '-'} |" + ) + else: + rows = captured_rows(site, batches) + union.update(rows) + gathered = "-" + if site.gathered and gathers: + parts = [] + for member, world in gathers: + wide = sorted({world * r for r in rows}) + union.update(wide) + parts.append( + f"{member}: {_multiple(world * site.row_multiple)} = " + f"{compact_values(wide) or '-'}" + ) + gathered = "; ".join(parts) + lines.append( + f"| {site.site} | {_multiple(site.row_multiple)} | {compact_values(rows) or '-'} | " + f"{gathered} |" + ) + if env is not None and env.max_num_tokens: + widest = max((w for _, w in gathers), default=1) if site.gathered else 1 + eager_rows.add(env.max_num_tokens * widest) + if errors: + lines += ["", "Malformed `call_sites` (fix them in `vocabulary.yaml`):", ""] + lines += [f"- {error}" for error in errors] + lines += [ + "", + f"**Certify every captured row count:** {compact_values(sorted(union)) or '(none: eager only)'}" + " -- replay on new data bitwise equal to eager at each.", + ] + if eager_rows: + lines.append( + f"**And every eager row count 1..{max(eager_rows)}** by the envelope standard " + "(grid-sized rows: a run at the maximum plus prefix-bitwise; bucketed kernels: " + "every bucket's edges)." + ) + + widths = sorted({w for site in sites for w in site.row_widths}) + reachable = max(eager_rows | union) if (eager_rows or union) else None + lines += [ + "", + "### Integer-width boundary probes (mandatory)", + "", + "Probe every index-typed argument at these boundaries and record each outcome in " + "the contract's `## Hard limits` table, whether or not today's envelope reaches it:", + "", + *(f"- {line}" for line in integer_width_probes(widths, reachable_rows=reachable)), + ] + return "\n".join(lines) + "\n" + + +def load_steps(vocabulary_path: Path | str) -> list[Mapping[str, Any]]: + """The vocabulary's steps, or nothing if the file is absent or unreadable.""" + path = Path(vocabulary_path) + if not path.is_file(): + return [] + try: + data = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + except yaml.YAMLError: + return [] + steps = data.get("steps") if isinstance(data, dict) else None + return [s for s in steps if isinstance(s, dict)] if isinstance(steps, list) else [] + + +def load_queue_row(vocabulary_path: Path | str, n: int) -> dict[str, Any]: + path = Path(vocabulary_path) + if not path.is_file(): + return {} + try: + data = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + except yaml.YAMLError: + return {} + queue = data.get("queue") if isinstance(data, dict) else None + for row in queue or (): + if isinstance(row, dict) and str(row.get("n")) == str(n): + return dict(row) + return {} + + +@dataclass(frozen=True) +class _Dispatch: + n: int + mode: str + target: str + + +def read_dispatch(workspace: Path, request: str) -> _Dispatch: + """The dispatch number, mode and target of a smith request already written.""" + rid = request if request.endswith("-smith") else f"{request}-smith" + for candidate in (rid, request): + path = workspace / "dispatch" / f"{candidate}.prompt.md" + if path.is_file(): + match = _DISPATCH.search(path.read_text(encoding="utf-8")) + if match is None: + raise ValueError(f"{path} names no catalog-entry dispatch") + return _Dispatch(n=int(match.group(1)), mode=match.group(2), target=match.group(3)) + raise FileNotFoundError(f"no dispatch prompt for {request!r} under {workspace / 'dispatch'}") + + +def _plan_for_workspace(workspace: Path) -> CertificationPlan | None: + from .envelope import derive_certification + from .task_schema import ( + TaskSchemaError, + load_and_validate_task_yaml, + members_for, + parse_parallel_segment, + ) + + task = workspace / "task.yaml" + if not task.is_file(): + return None + try: + data = load_and_validate_task_yaml(task) + except TaskSchemaError: + return None + _, size = parse_parallel_segment(str(data.get("parallel"))) + if size is None: + _, size = parse_parallel_segment(members_for(data, None, 4)[0]) + return derive_certification(data, world_size=size or 4) + + +def _main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser( + prog="python -m agent_flow.workflows.staircase_onboard.smith_brief", + description="Print the smith brief for a dispatch: call sites, rows, probes.", + ) + parser.add_argument("--workspace", required=True, help="the run's workspace directory") + parser.add_argument("request", help="a request id (e.g. 075-smith) or a dispatch number") + args = parser.parse_args(argv) + workspace = Path(args.workspace) + try: + dispatch = read_dispatch(workspace, args.request) + except (OSError, ValueError) as exc: + sys.stderr.write(f"{exc}\n") + return 2 + vocabulary = workspace / "vocabulary.yaml" + sys.stdout.write( + render_smith_brief( + dispatch, + steps=load_steps(vocabulary), + plan=_plan_for_workspace(workspace), + queue_row=load_queue_row(vocabulary, dispatch.n), + ) + ) + return 0 + + +__all__ = [ + "CALL_SITE_KEYS", + "FP32_EXACT_LIMIT", + "INT32_LIMIT", + "CallSite", + "captured_rows", + "default_call_sites", + "entry_key", + "gather_widths", + "integer_width_probes", + "load_queue_row", + "load_steps", + "parse_call_sites", + "read_dispatch", + "render_smith_brief", + "steps_for", +] + + +if __name__ == "__main__": + sys.exit(_main()) diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml b/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml index ef9a592137b7..697e69e25bb1 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml +++ b/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml @@ -171,6 +171,70 @@ trtllm_repo_path: /path/to/TensorRT-LLM # - --max_batch_size=8 # - --kv_cache_fraction=0.6 +# --------------------------------------------------------- serving envelope +# What the target is CLAIMED to serve. Every certification -- catalog +# receipts, capacity constants in the target, integer-width limits of the +# kernels -- is judged against this box, never against the tuples the gates +# happen to launch. +# +# The failure this prevents: when the certified set is "whatever the gates +# ran", a run whose gates all used max_num_tokens 8192 and max_seq_len 16384 +# certifies nothing above them, and every defect that lives above them -- a +# fixed-size buffer sized 8192 in the target, a recompute path that crashes +# past 16K, a 32-bit offset that overflows at 2^25 rows -- passes every gate +# and ships. +# +# From this block the run derives (preflight prints all of it, and writes it +# to /certification.md): +# - a check that every gate tuple lies INSIDE the envelope, and that every +# gated member is one the envelope claims; +# - one ENVELOPE PROBE per unit: the boot gate at the ceiling (these limits, +# this capture list), at the unit's cheapest member; +# - a boot-at-the-ceiling evidence requirement for every member no unit +# gates (multi-node, PP): delivery needs +# /evidence/.yaml {verdict: pass, tree: ..., log: ...}. +# +# Omitted: the envelope is DERIVED from the gate tuples (the legacy reading) +# and every rendering says so. To migrate a spec, print a block derived from +# its gates and then widen it: +# python -m agent_flow.workflows.staircase_onboard.envelope migrate task.yaml +# +# serving_envelope: +# max_batch_size: 256 # requests per step +# max_num_tokens: 16384 # tokens per step, per rank +# max_seq_len: 131072 # per-request length +# max_draft_len: 3 # optional: speculative tokens per request per step +# cuda_graph: # false = eager only +# padding: true # the engine's padding policy ... +# max_batch_size: 256 # ... generated up to here (default: max_batch_size) +# # batch_sizes: [1, 2, 4, 8, 16, 32, 64, 128, 256] # or an explicit list +# # Every configuration claimed, including ones this allocation cannot run: +# # [pp], e.g. tep8 (two 4-GPU nodes) or tep4pp2. +# members: [tp4, tep4, dep4, dep8] +# gpus_per_node: 4 # optional: default environment.gpus_per_node +# # Optional: a generated prompt of at least this many words every probe also +# # sends, so the probe prefills a long request. Keep it under max_num_tokens +# # unless the runtime block turns chunked_prefill on. +# probe_prompt_tokens: 12000 +# +# The feature combinations that must have ENGINE evidence beyond what the +# units' own gates produce. Each cell expands over its members and evidence +# kinds, and is classified as gated (an existing gate covers it), probe (the +# run launches a boot at that point) or external (evidence file owed at +# delivery, as above). +# +# feature_matrix: +# - name: mtp-long-context +# features: [mtp] # [] = core +# members: [tep4] # default: every envelope member +# evidence: [boot] # boot | accuracy | witness | acceptance +# at: {max_seq_len: 131072} # gate (default) | ceiling | a point +# prompt_tokens: 12000 +# - name: multi-node-accuracy +# features: [mtp] +# members: [dep8] +# evidence: accuracy + # ------------------------------------------------------------- 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. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py b/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py index 8840df865fb2..e01014e2092f 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py @@ -279,6 +279,16 @@ # only states one when this target is known to need a wider one. REMEDIATION_FIELD = "catalog_remediation" +# ------------------------------------------------ envelope / matrix / certification +# What the target is certified FOR, declared before the run: the serving +# envelope (engine limits, capture policy, every member claimed) and the +# feature combinations that must have engine evidence. Both are optional so an +# existing spec still validates -- its certified set is then the hull of its +# gate tuples, which every rendering states. The validators, the derivation and +# the migration command live in `envelope.py`. +SERVING_ENVELOPE_FIELD = "serving_envelope" +FEATURE_MATRIX_FIELD = "feature_matrix" + # Runtime switches the gates must exercise, and which therefore have to be # declared rather than left to whatever the engine defaults to. # @@ -318,6 +328,8 @@ SCOPE_FIELD, SOURCE_EMBARGO_FIELD, REMEDIATION_FIELD, + SERVING_ENVELOPE_FIELD, + FEATURE_MATRIX_FIELD, ) REMEDIATION_KEYS: tuple[str, ...] = ( @@ -1445,6 +1457,17 @@ def _require_declared_witnesses(data: Mapping[str, Any], errors: list[str]) -> N ) +def _validate_certification(data: Mapping[str, Any], errors: list[str]) -> None: + """Validate ``serving_envelope`` and ``feature_matrix``; see `envelope.py`. + + Imported here rather than at module scope: `envelope` consumes this + schema's constants and helpers, so a top-level import would be circular. + """ + from .envelope import validate_certification_blocks + + validate_certification_blocks(data, errors) + + 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. @@ -1664,6 +1687,8 @@ def load_and_validate_task_yaml(path: str | Path) -> dict[str, Any]: errors.append(str(exc)) _require_declared_witnesses(data, errors) + _validate_certification(data, errors) + if errors: bullet = "\n - " raise TaskSchemaError( @@ -1712,6 +1737,8 @@ def assembly_units(data: Mapping[str, Any]) -> list[str]: "REPLAY_WITNESS_REQUIRED", "EQUIVALENCE_WITNESS_FEATURES", "EQUIVALENCE_WITNESS_KEYS", + "FEATURE_MATRIX_FIELD", + "SERVING_ENVELOPE_FIELD", "WITNESS_BATCH_SIZE_KEY", "WITNESS_KEYS", "WITNESS_LEGS_ENV_KEY", diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py b/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py index 2e6cee7d9c43..4d7ed7d85310 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py @@ -52,10 +52,11 @@ from agent_flow.logger import get_logger from agent_flow.runtime import CommandExecutor, TurnExecutor, TurnRequest +from . import envelope as envelope_mod from . import gates as gates_mod from . import remediation from .committer import install_committer -from .delivery import summarize_for_auditor, verify_delivery +from .delivery import DeliveryReport, summarize_for_auditor, verify_delivery from .environment import ( RunEnvironment, image_import_warning, @@ -67,6 +68,7 @@ from .preflight import ( MAX_TMPDIR_LEN, PreflightError, + SpecInconsistencyError, format_preflight_report, run_preflight, substrate_gate, @@ -217,6 +219,12 @@ "kv_quant": "kv_quant", } +# The BLOCKED reasons a resume handles instead of re-dispatching: the two the +# orchestrator can act on by itself, dispatching smiths -- an onboard for a +# vocabulary gap, verifies for a certification gap. Any other reason already +# ended the run once, and handling it again just ends it again. +RESUMABLE_BLOCKED_REASONS: tuple[str, ...] = ("vocabulary_gap", "certification_gap") + # Shown on every infrastructure failure, and pointed at the file rather than # summarized, because the summary is the part that goes stale. A run stops # here perhaps once a week; the operator reading that message has no reason to @@ -442,6 +450,9 @@ def __init__( ) self._runtime_block: str = "" self._kernel_sources_block: str = "" + #: The certification standard every planning and certifying role + #: carries; written with `certification.md` at preflight. + self._certification_block: str = "" # ------------------------------------------------------------- lifecycle @@ -617,6 +628,7 @@ def run(self, task: str | Path) -> None: if state.stage == STAGE_PLAN: if not self._run_plan_phase(state, log): return + self._run_early_checks("end of the plan phase", log) state.stage = STAGE_SMITH self._checkpoint(state) @@ -657,12 +669,16 @@ def _run_preflight(self, state: WorkflowState, log) -> bool: if self._command_runner is not None: kwargs["runner"] = self._command_runner try: - result = run_preflight(self.task_data, **kwargs) + result = run_preflight(self.task_data, unit_kinds=self._unit_gate_kinds, **kwargs) except PreflightError as exc: + # A spec that contradicts itself is fixed in the spec; the + # infrastructure guide would send the operator the wrong way. + hint = "" if isinstance(exc, SpecInconsistencyError) else TROUBLESHOOTING_HINT print_message(f"[bold red]✘ preflight: {exc}[/bold red]", log) - print_message(TROUBLESHOOTING_HINT, log) + if hint: + print_message(hint, log) self.journal.open(slug="(preflight failed)") - self.journal.section("preflight failed", f"{exc}\n\n{TROUBLESHOOTING_HINT}") + self.journal.section("preflight failed", f"{exc}\n\n{hint}".rstrip()) return False state.slug = result.slug @@ -747,6 +763,7 @@ def _run_preflight(self, state: WorkflowState, log) -> bool: return False self._take_survey(result, log) + self._write_certification(result) # Prompts can only be built now: every role's prompt carries the # resolved launch prefix, which did not exist until the device claim @@ -1007,6 +1024,76 @@ def _take_survey(self, result, log) -> None: if surface.unrecognized: print_message(f" unrecognized structural keys: {', '.join(surface.unrecognized)}", log) + def _unit_gate_kinds(self, unit: str) -> set[str]: + """The gate kinds `_assemble_one_unit` runs for `unit`; see `envelope.py`.""" + kinds = {envelope_mod.EVIDENCE_BOOT, envelope_mod.EVIDENCE_ACCURACY} + if unit in WITNESS_GATED_UNITS: + kinds.add(envelope_mod.EVIDENCE_WITNESS) + if unit in ACCEPTANCE_GATED_UNITS: + kinds.add(envelope_mod.EVIDENCE_ACCEPTANCE) + return kinds + + def _write_certification(self, result) -> None: + """Write `certification.md` and the block the certifying roles carry.""" + plan = getattr(result, "certification", None) + if plan is None: + return + path = self.workspace / "certification.md" + path.write_text(plan.render(), encoding="utf-8") + self._certification_block = plan.prompt_block(path) + self.journal.section("certification plan", f"`{path}`\n\n" + "\n".join(plan.notes())) + + def _run_early_checks(self, stage: str, log) -> None: + """Run the delivery checks that need no finished target, as warnings only.""" + from .early_checks import run_early_checks + + if self._preflight is None: + return + environment = self.task_data.get("environment") or {} + report = run_early_checks( + repo_root=self.task_data["trtllm_repo_path"], + family=self._preflight.family, + stage=stage, + vocabulary_path=self.vocabulary_path, + env=environment.get("env") or {}, + capacity=self._capacity_audit(), + certification=getattr(self._preflight, "certification", None), + workspace=self.workspace, + ) + path = self.workspace / "early_checks.md" + path.write_text(report.render(), encoding="utf-8") + warnings = report.warnings + self.journal.section( + f"early delivery checks ({stage})", + f"{len(warnings)} warning(s); required CI assets and details in `{path}`\n\n" + + "\n".join(f"- {w.render()}" for w in warnings), + ) + for warning in warnings: + print_message(f"[bold yellow]⚠ {escape_markup(warning.render())}[/bold yellow]", log) + + def _run_envelope_probes(self, unit: str, log) -> bool: + """Boot `unit` at every envelope point its certification plan derived.""" + plan = getattr(self._preflight, "certification", None) + for probe in plan.probes_for(unit) if plan is not None else (): + result = self.gate_log.record( + envelope_mod.run_envelope_probe( + self._gate_env(), + task_data=self.task_data, + probe=probe, + spec=gates_mod.BootSpec.load(self.boot_spec_path), + runner=self._shell_runner, + features=self._features_through(unit), + ) + ) + self.journal.event("gate", result.summary()) + if not result.passed: + print_message( + f"[bold yellow]↻ {escape_markup(result.summary())}[/bold yellow]", log + ) + return False + print_message(f"[bold green]✔ {escape_markup(result.summary())}[/bold green]", log) + return True + def _run_plan_phase(self, state: WorkflowState, log) -> bool: """Drafter <-> reviewer, capped, and never blocking the build.""" if self.preset_plan is not None and self.preset_vocabulary is not None: @@ -1509,6 +1596,7 @@ def _run_assemble_phase(self, state: WorkflowState, log) -> bool: state.units_done.append(unit) state.unit_index += 1 state.unit_turns = 0 + self._run_early_checks(f"after the {unit} unit", log) # A unit is the assembler's unit of work, so its context ends # with it. The next unit starts clean, which is also what makes # "core is green before the first feature" mean something. @@ -1700,7 +1788,7 @@ def _assemble_one_unit(self, state: WorkflowState, unit: str, log) -> bool: # the blocking condition is meant to be revisited. So the # other reasons fall through to a fresh assembler turn, which # is free to reach the same conclusion and say so. - elif prior == "BLOCKED" and prior_reason == "vocabulary_gap": + elif prior == "BLOCKED" and prior_reason in RESUMABLE_BLOCKED_REASONS: print_message( "[bold cyan]→ resuming onto a BLOCKED assembler decision from a " "previous process; handling it rather than re-dispatching" @@ -1739,6 +1827,8 @@ def _assemble_one_unit(self, state: WorkflowState, unit: str, log) -> bool: if decision == "UNIT_READY": state.unit_marker[unit] = self._target_marker() self._checkpoint(state) + if not self._fan_out_certification_gaps(state, "assembler", unit, log): + continue # The configurations this unit is gated over, resolved once. A # `unified` segment is ONE codebase over a FAMILY of parallel @@ -1776,12 +1866,21 @@ def _assemble_one_unit(self, state: WorkflowState, unit: str, log) -> bool: ) if not boots: continue + # The envelope's ceiling, before review: warm-up at the real token + # budget and capture at every claimed batch are what a fixed-size + # buffer fails, and no gate tuple reaches them. + if not self._run_envelope_probes(unit, log): + continue if self.reviewer_mode != "off": # The last member's evidence: the reviewer reads one dump and # the summary line names which member produced it, so the # freshest is the one whose log is still meaningful to quote. self._dispatch("reviewer", what=unit, prompt=self._reviewer_prompt(unit, boots[-1])) + # Certification-only findings go to verify smiths, not back to + # the assembler: the target is correct, the catalog is short. + if not self._fan_out_certification_gaps(state, "reviewer", unit, log): + continue if latest_decision(self.progress_path, "reviewer", REVIEWER_DECISIONS) == "REJECT": print_message("[bold yellow]↻ reviewer REJECT[/bold yellow]", log) continue @@ -2102,7 +2201,9 @@ def _run_remediation_cycle(self, state: WorkflowState, report, boot, log) -> boo print_message("[bold green]✔ audit CLEAN[/bold green]", log) return True - rows, dropped = remediation.parse_uncovered_rows(entry.get("uncovered")) + rows, dropped = remediation.parse_audit_rows( + entry, default_arch=str(self.task_data.get("arch", "")) + ) if dropped: # Never silent: a dropped row is a finding nothing will act # on, and the run would otherwise converge on an incomplete @@ -2222,10 +2323,14 @@ def _dispatch_verifiers(self, cycle: "remediation.RemediationState", decision, l # One smith per surface, carrying the whole surface. The dispatch # numbers start well above the build queue's so a decision file from - # remediation is never mistaken for one the smith phase wrote. + # remediation is never mistaken for one the smith phase wrote, and are + # the item's position in the queue, which only grows: a later round or + # a certification-gap fan-out never reuses another surface's number, + # and so never reads another surface's decision. + position = {queued.group_key: index for index, queued in enumerate(cycle.queue)} dispatches = { item.identity: QueueItem( - n=9000 + index, + n=9000 + position.get(item.group_key, index), mode=MODE_VERIFY, target=item.entry, surface=item.brief(), @@ -2261,6 +2366,58 @@ def _dispatch_verifiers(self, cycle: "remediation.RemediationState", decision, l cycle.queue = list(updated) return failed + def _fan_out_certification_gaps(self, state: WorkflowState, role: str, unit: str, log) -> bool: + """Verify every `certification_gap` row `role`'s latest decision carries. + + True when there was nothing to verify or every verify landed; False + when one failed or cannot be retried, which sends the unit back to the + assembler with a correction naming the surfaces. The rows join the + remediation queue, so the delivery audit never re-dispatches them. + """ + entry = latest_entry(self.progress_path, role) or {} + rows, dropped = remediation.parse_certification_gaps( + entry.get(remediation.CERTIFICATION_GAP_KEY), + default_arch=str(self.task_data.get("arch", "")), + ) + if dropped: + print_message( + f"[bold yellow]! {dropped} certification_gap row(s) from the {role} named no " + f"entry and cannot be dispatched[/bold yellow]", + log, + ) + if not rows: + return True + budget = self._remediation_budget() + cycle = remediation.RemediationState.from_dict(state.remediation) + incoming = remediation.canonicalize( + rows, target_commit=self._target_commit(), capture_id=f"{role}:{unit}" + ) + merge = remediation.merge_into_queue(cycle.queue, incoming, budget=budget) + decision = remediation.decide_gap_dispatch(merge, budget=budget, incoming=incoming) + cycle.queue = list(merge.queue) + state.remediation = cycle.as_dict() + self._checkpoint(state) + self.journal.section(f"certification gap ({role}, {unit})", decision.detail) + print_message( + f"[bold cyan]→ certification gap from the {role}: {decision.detail}[/bold cyan]", log + ) + failed: tuple = () + if decision.action == "dispatch": + failed = self._dispatch_verifiers(cycle, decision, log) + state.remediation = cycle.as_dict() + self._checkpoint(state) + if decision.action != "stop" and not failed: + return True + names = ", ".join(f"{i.entry}[{i.axis}]" for i in failed) or decision.detail + self._record_correction( + f"certification gap, {role} on {unit}", + f"These surfaces did not certify: {names}. The entry fails there or could not be " + f"verified, so the target must not reach those values with it: use another entry, " + f"bound the call, or narrow the serving envelope (a plan change).", + ) + print_message(f"[bold red]✘ certification gap not closed: {names}[/bold red]", log) + return False + def _target_marker(self) -> str: """A cheap identity for the work under the target path, or ``""``. @@ -2339,7 +2496,7 @@ def _capture_id(self) -> str: def _delivery_report(self, state: WorkflowState): """The mechanical half of delivery. Cheap, and safe to repeat.""" - return verify_delivery( + report = verify_delivery( repo_root=self.task_data["trtllm_repo_path"], target_rel=self._preflight.target_rel, family=self._preflight.family, @@ -2351,6 +2508,16 @@ def _delivery_report(self, state: WorkflowState): model_key=self._model_key(), accuracy_test=f"test_modeling_v2_{self._preflight.family}.py", ) + return DeliveryReport(checks=(*report.checks, *self._certification_checks())) + + def _certification_checks(self) -> list: + """Delivery checks owed to the envelope: matrix evidence and capacity constants.""" + plan = getattr(self._preflight, "certification", None) + if plan is None: + return [] + audit = self._capacity_audit() + capacity = audit.checks() if audit is not None else [] + return [*envelope_mod.check_external_evidence(plan, self.workspace), *capacity] # ------------------------------------------------------------------ helpers @@ -2809,6 +2976,21 @@ def _handle_blocked_assembler(self, state: WorkflowState, log) -> bool: """ entry = latest_entry(self.progress_path, "assembler") or {} reason = str(entry.get("blocked_reason", "")).strip() + if reason == "certification_gap": + units = assembly_units(self.task_data) + unit = units[state.unit_index] if state.unit_index < len(units) else UNIT_CORE + if not entry.get(remediation.CERTIFICATION_GAP_KEY): + print_message( + "[bold red]✘ assembler blocked on a certification gap without listing " + "`certification_gap`, so there is nothing to verify.[/bold red]", + log, + ) + self._close_run(state, log, verdict="certification-gap-unnamed") + return False + if self._fan_out_certification_gaps(state, "assembler", unit, log): + return True + self._close_run(state, log, verdict=remediation.STOP_VERIFICATION_FAILED) + return False if reason == "vocabulary_gap": missing = str(entry.get("missing_entry", "")).strip() if not missing: @@ -2983,6 +3165,7 @@ def _plan_drafter_prompt(self, round_no: int) -> str: f"architecture, the parallel segment, the features and the accuracy anchor.\n\n" f"{self._surface_block}\n" f"{self._runtime_block}\n" + f"{self._certification_block}\n" f"Write `{self.plan_path}` and `{self.vocabulary_path}` as your system " f"prompt specifies. The orchestrator parses `vocabulary.yaml` and will " f"refuse a malformed one, so follow the schema exactly.\n\n" @@ -3061,6 +3244,7 @@ def _smith_prompt(self, item: QueueItem, note: str = "") -> str: f"Load the `{item.skill}` skill before anything else.\n" f"{surface}" f"{tier}\n" + f"{self._smith_brief(item)}\n" f"{self._kernel_sources_block}\n" f"Read `{self.vocabulary_path}` for the computation step this entry serves " f"and the values this checkpoint passes it.\n\n" @@ -3069,6 +3253,17 @@ def _smith_prompt(self, item: QueueItem, note: str = "") -> str: + (f"\n\n{context}\n" if context else "") ) + def _smith_brief(self, item: QueueItem) -> str: + """The entry's call sites, captured rows and integer-width probes; see `smith_brief.py`.""" + from .smith_brief import load_queue_row, load_steps, render_smith_brief + + return render_smith_brief( + item, + steps=load_steps(self.vocabulary_path), + plan=getattr(self._preflight, "certification", None), + queue_row=load_queue_row(self.vocabulary_path, item.n), + ) + def _assembler_prompt(self, unit: str, turn: int) -> str: feature_skills = "" if unit != "core": @@ -3116,9 +3311,24 @@ def _reviewer_prompt(self, unit: str, boot: gates_mod.GateResult) -> str: "table, vocabulary audit and dump review are what you check against the " "diff.\n\n" "Read the diff and decide. Do not re-run the gates.\n\n" + + self._capacity_block() + self._report_instruction("reviewer") ) + def _capacity_audit(self): + """The static capacity audit of the target, or None before preflight.""" + from .capacity_audit import audit_target + + plan = getattr(self._preflight, "certification", None) + rel = str(getattr(self._preflight, "target_rel", "") or "") + if plan is None or not rel: + return None + return audit_target(Path(self.task_data["trtllm_repo_path"]) / rel, plan.envelope) + + def _capacity_block(self) -> str: + audit = self._capacity_audit() + return f"{self._certification_block}\n{audit.render()}\n" if audit is not None else "" + #: The line every modeling_v2 target's `_emit_runtime_values` writes before #: its dump. Named here because the auditor has to find dumps this process #: never saw — see `_auditor_dump`. @@ -3173,6 +3383,7 @@ def _auditor_prompt(self, report, boot: gates_mod.GateResult | None) -> str: f"{summarize_for_auditor(report, self._auditor_dump())}\n" f"## Gates\n\n```\n{self.gate_log.render() or '(none in this process — see above)'}" f"\n```\n\n" + f"{self._certification_block}\n" f"Do the one check that is left: every value in the runtime-value dump " f"against its catalog entry's certified column, on this architecture.\n\n" + self._report_instruction("auditor") diff --git a/agent-flow/tests/workflows/staircase_onboard/test_capacity_audit.py b/agent-flow/tests/workflows/staircase_onboard/test_capacity_audit.py new file mode 100644 index 000000000000..40f071467903 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_capacity_audit.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the static capacity-constant audit (``capacity_audit.py``). + +The synthetic target reproduces the shape of the DeepSeek-V4.1 defect: a +module-level ``_TOPK_MAX_ROWS = 8192`` sizing a constant offsets buffer at +construction, which killed every engine above ``max_num_tokens`` 8192 in +warm-up. Under the gates' hull (8192) it passes, which is how it shipped; +under the declared envelope (16384) it must not. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import capacity_audit as audit +from agent_flow.workflows.staircase_onboard import envelope +from agent_flow.workflows.staircase_onboard.state import load_state + +from .test_delivery import TARGET_REL +from .test_workflow import Script, _install, _workflow, _write_vocabulary + +_TOPK = '''\ +"""An indexer with a fixed-size offsets buffer.""" + +import torch + +_TOPK_MAX_ROWS = 8192{binding} +_HEAD_DIM = 128 +_BLOCK = 64 + + +class TopK: + def __init__(self, model_config): + self._offsets = torch.arange(_TOPK_MAX_ROWS + 1, dtype=torch.int32) + self._rows = torch.empty(model_config.max_num_tokens, _HEAD_DIM) + + def forward(self, scores, n): + assert n <= _TOPK_MAX_ROWS + return scores +''' + +DECLARED = envelope.ServingEnvelope( + max_batch_size=256, max_num_tokens=16384, max_seq_len=1048576, members=() +) +GATE_HULL = envelope.ServingEnvelope( + max_batch_size=32, max_num_tokens=8192, max_seq_len=16384, declared=False +) + + +def _target(tmp_path: Path, body: str, name: str = "topk.py") -> Path: + target = tmp_path / "target" + target.mkdir(exist_ok=True) + (target / name).write_text(body, encoding="utf-8") + return target + + +def _by_name(result: audit.CapacityAudit) -> dict[str, audit.CapacityFinding]: + return {f.name: f for f in result.findings} + + +def test_an_unbound_capacity_constant_fails_whatever_the_envelope(tmp_path): + target = _target(tmp_path, _TOPK.format(binding="")) + result = audit.audit_target(target, DECLARED) + + finding = _by_name(result)["_TOPK_MAX_ROWS"] + assert not finding.passed + assert "unbound capacity" in finding.detail + assert "sizes a buffer" in finding.detail and "used as a ceiling" in finding.detail + # A small dimension is not a capacity; an engine-sized buffer is tied. + assert "_HEAD_DIM" not in _by_name(result) and "_BLOCK" not in _by_name(result) + assert len(result.tied) == 0 # the engine-sized buffer also reads _HEAD_DIM + assert not result.passed + + +def test_the_bound_constant_passes_the_gate_hull_and_fails_the_envelope(tmp_path): + """Exactly how it shipped: inside what the gates ran, outside what is served.""" + target = _target(tmp_path, _TOPK.format(binding=" # capacity: max_num_tokens")) + + hull = audit.audit_target(target, GATE_HULL) + assert _by_name(hull)["_TOPK_MAX_ROWS"].passed + assert "DERIVED envelope" in hull.render() + + declared = audit.audit_target(target, DECLARED) + finding = _by_name(declared)["_TOPK_MAX_ROWS"] + assert not finding.passed + assert "8192 < max_num_tokens = 16384" in finding.detail + checks = declared.checks() + assert [c.name for c in checks] == ["capacity:_TOPK_MAX_ROWS"] + assert "topk.py:5" in checks[0].detail + + +def test_a_capacity_sized_from_the_engine_needs_no_binding(tmp_path): + body = """\ +import torch + + +class TopK: + def __init__(self, model_config): + self._offsets = torch.arange(model_config.max_num_tokens + 1) + self._slots = torch.zeros(self.max_batch_size, 4) +""" + result = audit.audit_target(_target(tmp_path, body), DECLARED) + assert result.findings == () and len(result.tied) == 2 + assert result.passed + assert [c.name for c in result.checks()] == ["capacity:audit"] + + +def test_bindings_in_the_bounds_mapping_and_their_factors(tmp_path): + body = """\ +import torch + +CAPACITY_BOUNDS = { + "_GATHER_ROWS": "max_num_tokens * world_size", + "Pool.SLOTS": "max_batch_size", + "_GONE": "max_seq_len", +} +_GATHER_ROWS = 65536 +_WINDOW_LEN = 128 # capacity: fixed: the checkpoint's sliding window + + +class Pool: + SLOTS = 128 +""" + dep = envelope.ServingEnvelope( + max_batch_size=256, + max_num_tokens=16384, + max_seq_len=131072, + members=(envelope.parse_envelope_member("dep4"),), + ) + findings = _by_name(audit.audit_target(_target(tmp_path, body), dep)) + + assert findings["_GATHER_ROWS"].passed + assert "65536 >= max_num_tokens * world_size = 65536" in findings["_GATHER_ROWS"].detail + assert not findings["Pool.SLOTS"].passed # 128 < 256 + assert findings["_WINDOW_LEN"].passed and "sliding window" in findings["_WINDOW_LEN"].detail + assert "does not define" in findings["_GONE"].detail + + wider = envelope.ServingEnvelope( + max_batch_size=256, + max_num_tokens=16384, + max_seq_len=131072, + members=(envelope.parse_envelope_member("dep8"),), + ) + assert not _by_name(audit.audit_target(_target(tmp_path, body), wider))["_GATHER_ROWS"].passed + + +def test_a_large_literal_buffer_needs_a_binding(tmp_path): + body = """\ +import torch + +_WORKSPACE = torch.empty(1 << 20, dtype=torch.uint8) +_SCRATCH = torch.empty(2048) # capacity: fixed: one page of host scratch +_SMALL = torch.empty(16) +""" + result = audit.audit_target(_target(tmp_path, body), DECLARED) + buffers = [f for f in result.findings if f.kind == "buffer"] + assert [(f.value, f.passed) for f in buffers] == [(1 << 20, False), (2048, True)] + assert "literal size at import" in buffers[0].detail + + +@pytest.mark.parametrize( + ("text", "error"), + [ + ("fixed", "needs a reason"), + ("max_tokens", "not one of"), + ("max_num_tokens * two", "is not an integer"), + ], +) +def test_malformed_bindings_fail_with_the_reason(text, error): + binding = audit.parse_binding(text) + assert error in binding.error + + +def test_an_envelope_without_the_axis_cannot_prove_the_bound(tmp_path): + target = _target(tmp_path, "_MAX_LEN = 4096 # capacity: max_seq_len\n") + hull = envelope.ServingEnvelope(max_batch_size=8, max_num_tokens=2048, max_seq_len=None) + finding = _by_name(audit.audit_target(target, hull))["_MAX_LEN"] + assert not finding.passed and "cannot be proven" in finding.detail + + +def test_unparsable_files_and_bounds_are_reported(tmp_path): + target = _target(tmp_path, "CAPACITY_BOUNDS = dict(a='b')\n") + (target / "broken.py").write_text("def (:\n", encoding="utf-8") + result = audit.audit_target(target, DECLARED) + assert not result.passed + assert any("cannot parse" in e for e in result.errors) + assert any("literal mapping" in e for e in result.errors) + + +# ----------------------------------------------------------------- workflow + + +def test_the_reviewer_reads_the_audit_and_delivery_fails_on_it( + tmp_path, task_file, repo, monkeypatch +): + (repo / TARGET_REL / "topk.py").write_text(_TOPK.format(binding=""), encoding="utf-8") + spec = yaml.safe_load(task_file.read_text()) + spec["runtime"] = {"extra_flags": ["--max_seq_len=4096", "--max_batch_size=4"]} + spec["accuracy"].update(max_seq_len=4096, max_batch_size=4) + spec["serving_envelope"] = { + "max_batch_size": 64, + "max_num_tokens": 16384, + "max_seq_len": 32768, + "cuda_graph": False, + "members": ["tp1"], + } + task_file.write_text(yaml.safe_dump(spec), encoding="utf-8") + + workflow = _workflow(tmp_path) + calls = _install(monkeypatch, workflow, Script(None)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + reviewer = next(prompt for role, prompt in calls if role == "reviewer") + assert "## Capacity constants in the target" in reviewer + assert "| FAIL | constant | `_TOPK_MAX_ROWS` | 8192 |" in reviewer + assert "## Serving envelope (declared)" in reviewer + state = load_state(tmp_path / "ws" / ".staircase_onboard_state.json") + assert state.verdict == "delivery-failed" + assert "capacity:_TOPK_MAX_ROWS" in (tmp_path / "ws" / "journal.md").read_text() + + +def test_the_reviewer_and_assembler_are_told_about_capacity_constants(): + from agent_flow.workflows.staircase_onboard.prompts import DEFAULT_PROMPTS + + assert "Capacity constants against the serving envelope" in DEFAULT_PROMPTS.reviewer + assert "# capacity: " in DEFAULT_PROMPTS.assembler diff --git a/agent-flow/tests/workflows/staircase_onboard/test_certification_gap.py b/agent-flow/tests/workflows/staircase_onboard/test_certification_gap.py new file mode 100644 index 000000000000..832970624985 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_certification_gap.py @@ -0,0 +1,300 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the typed ``certification_gap`` decision and its automatic fan-out. + +The history behind it, from the DeepSeek-V4.1 run: three reviewer REJECTs of a +sound core target, each on certification alone ("the target code is sound and +needs no change"), each costing a reviewer turn plus an assembler turn; an +assembler that could only say ``knowledge_gap`` for four verifies it needed, +which closed the run; and four verify smiths dispatched by hand. Each case +below is one of those, now handled by the orchestrator. +""" + +from __future__ import annotations + +from pathlib import Path + +import yaml + +from agent_flow.workflows.staircase_onboard import progress as progress_mod +from agent_flow.workflows.staircase_onboard import remediation as R +from agent_flow.workflows.staircase_onboard.state import load_state + +from .test_workflow import Script, _install, _workflow, _write_vocabulary + +_GAP = { + "entry": "activation/silu_and_mul", + "values": ["2", "3", "6", "8192"], + "axis": "rows", + "reason": "MoE expert MLP at the verify step: 6·G rows", +} + + +# ---------------------------------------------------------------- parsing + + +def test_a_gap_row_needs_only_an_entry(): + rows, dropped = R.parse_certification_gaps( + [ + {"entry": "moe/runner", "values": [16384]}, + {"entry": "norm/rms", "value": "3", "axis": "T", "arch": "sm_100", "reason": "why"}, + {"values": ["1"]}, + "not a row", + ], + default_arch="sm_103", + ) + assert dropped == 2 + assert (rows[0].arch, rows[0].axis, rows[0].values) == ( + "sm_103", + R.DEFAULT_GAP_AXIS, + ("16384",), + ) + assert (rows[1].arch, rows[1].axis, rows[1].contract_line) == ("sm_100", "T", "why") + + +def test_the_auditor_may_mix_uncovered_and_gap_rows(): + entry = { + "uncovered": [{"entry": "a/x", "arch": "sm_103", "axis": "rows", "values": ["5"]}], + R.CERTIFICATION_GAP_KEY: [{"entry": "b/y", "values": ["7"]}, {"values": []}], + } + rows, dropped = R.parse_audit_rows(entry, default_arch="sm_103") + assert [r.entry for r in rows] == ["a/x", "b/y"] and dropped == 1 + + +def test_records_are_normalized_at_the_persistence_boundary(): + records = R.normalize_certification_gap_records( + [{"entry": " moe/runner ", "values": [1, " 2 ", ""], "reason": " r "}, {"entry": ""}, 3] + ) + assert records == [{"entry": "moe/runner", "values": ["1", "2"], "reason": "r"}] + + +def _item(entry: str, status: str = R.QUEUE_PENDING, values=("1",)) -> R.RemediationItem: + return R.RemediationItem(entry=entry, arch="sm_103", axis="rows", values=values, status=status) + + +def test_new_gaps_are_dispatched_and_certified_ones_are_not(): + budget = R.RemediationBudget() + queue = [_item("a/done", R.QUEUE_CERTIFIED)] + incoming = (_item("a/done"), _item("b/new")) + merge = R.merge_into_queue(queue, incoming, budget=budget) + decision = R.decide_gap_dispatch(merge, budget=budget, incoming=incoming) + assert decision.action == "dispatch" + assert [i.entry for i in decision.dispatch] == ["b/new"] + assert "already certified by an earlier verify" in decision.detail + + +def test_a_gap_whose_verify_already_failed_stops(): + budget = R.RemediationBudget() + queue = [_item("a/x", R.QUEUE_FAILED_RECEIPT)] + incoming = (_item("a/x", values=("2",)),) + merge = R.merge_into_queue(queue, incoming, budget=budget) + decision = R.decide_gap_dispatch(merge, budget=budget, incoming=incoming) + assert decision.stopping and decision.verdict == R.STOP_VERIFICATION_FAILED + + +def test_nothing_new_is_clean(): + budget = R.RemediationBudget() + queue = [_item("a/x", R.QUEUE_CERTIFIED)] + incoming = (_item("a/x"),) + merge = R.merge_into_queue(queue, incoming, budget=budget) + assert R.decide_gap_dispatch(merge, budget=budget, incoming=incoming).action == "clean" + + +# --------------------------------------------------------------- decisions + + +def test_the_three_roles_carry_the_gap_through_a_decision_file(tmp_path): + ctx = progress_mod.ProgressContext(path=tmp_path / "progress.yaml") + progress_mod.init_progress_file(ctx.path) + for role, decision in ( + ("assembler", "UNIT_READY"), + ("reviewer", "APPROVE"), + ("auditor", "CLEAN"), + ): + path = progress_mod.decision_path(tmp_path, role, 1) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + yaml.safe_dump( + {"decision": decision, "summary": "s", "certification_gap": [_GAP, {"x": 1}]} + ), + encoding="utf-8", + ) + assert progress_mod.collect_decision_file(ctx, tmp_path, role, 1) == decision + entry = progress_mod.latest_entry(ctx.path, role) + assert entry["certification_gap"] == [ + { + "entry": _GAP["entry"], + "values": _GAP["values"], + "axis": "rows", + "reason": _GAP["reason"], + } + ] + assert "certification_gap" in progress_mod.BLOCKED_REASONS + + +def test_the_tools_offer_the_gap_to_exactly_the_three_roles(): + for role in R.CERTIFICATION_GAP_ROLES: + assert R.CERTIFICATION_GAP_KEY in progress_mod.EXTRA_KEYS_BY_AGENT[role] + for role in ("smith", "plan_drafter", "plan_reviewer"): + assert R.CERTIFICATION_GAP_KEY not in progress_mod.EXTRA_KEYS_BY_AGENT.get(role, ()) + + +# ---------------------------------------------------------------- workflow + + +def _landing_smith(turn, prompt): + target = prompt.split("the catalog entry `", 1)[1].split("`", 1)[0] + return {"decision": "LANDED", "summary": f"certified {target}", "entry": target} + + +def _run(tmp_path, task_file, monkeypatch, **roles): + workflow = _workflow(tmp_path) + roles.setdefault("smith", _landing_smith) + calls = _install(monkeypatch, workflow, Script(None, **roles)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + return calls, load_state(tmp_path / "ws" / ".staircase_onboard_state.json") + + +def _roles(calls): + return [role for role, _ in calls] + + +def test_an_assembler_gap_is_verified_while_the_unit_moves_on(tmp_path, task_file, monkeypatch): + calls, state = _run( + tmp_path, + task_file, + monkeypatch, + assembler=lambda turn, prompt: { + "decision": "UNIT_READY", + "summary": "core assembled; silu rows uncovered", + "certification_gap": [_GAP], + }, + ) + assert _roles(calls) == [ + "plan_drafter", + "plan_reviewer", + "assembler", + "smith", + "reviewer", + "auditor", + ] + smith = next(prompt for role, prompt in calls if role == "smith") + assert "**verify** the catalog entry `activation/silu_and_mul`" in smith + assert "2, 3, 6, 8192" in smith and "6·G rows" in smith + assert state.verdict == "delivered" + cycle = R.RemediationState.from_dict(state.remediation) + assert [(i.entry, i.status, i.capture_id) for i in cycle.queue] == [ + ("activation/silu_and_mul", R.QUEUE_CERTIFIED, "assembler:core") + ] + + +def test_a_certification_only_review_does_not_reopen_the_assembler( + tmp_path, task_file, monkeypatch +): + calls, state = _run( + tmp_path, + task_file, + monkeypatch, + reviewer=lambda turn, prompt: { + "decision": "APPROVE", + "summary": "the target code is sound; one entry needs certifying", + "certification_gap": [_GAP], + }, + ) + assert _roles(calls).count("assembler") == 1 + assert _roles(calls).count("smith") == 1 + assert state.verdict == "delivered" + + +def test_a_blocked_assembler_gets_its_verifies_and_is_resumed(tmp_path, task_file, monkeypatch): + def assembler(turn, prompt): + if turn == 1: + return { + "decision": "BLOCKED", + "summary": "cannot pick the chunk size without knowing", + "blocked_reason": "certification_gap", + "certification_gap": [_GAP, {"entry": "moe/runner", "values": ["16384"]}], + } + return {"decision": "UNIT_READY", "summary": "assembled"} + + calls, state = _run(tmp_path, task_file, monkeypatch, assembler=assembler) + roles = _roles(calls) + assert roles.index("smith") < len(roles) - roles[::-1].index("assembler") - 1 + assert roles.count("smith") == 2 and roles.count("assembler") == 2 + assert state.verdict == "delivered" + # Two surfaces, two distinct dispatch numbers: neither reads the other's decision. + smiths = [prompt for role, prompt in calls if role == "smith"] + numbers = sorted(int(p.split("Dispatch ", 1)[1].split(":", 1)[0]) for p in smiths) + assert numbers == [9000, 9001] + + +def test_a_gap_that_fails_to_certify_goes_back_to_the_assembler(tmp_path, task_file, monkeypatch): + def assembler(turn, prompt): + entry = {"decision": "UNIT_READY", "summary": "assembled"} + if turn == 1: + entry["certification_gap"] = [_GAP] + return entry + + calls, state = _run( + tmp_path, + task_file, + monkeypatch, + assembler=assembler, + smith=lambda turn, prompt: { + "decision": "FAILED_RECEIPT", + "summary": "fails at 8192 rows", + "entry": "activation/silu_and_mul", + }, + ) + assert _roles(calls).count("assembler") == 2 + corrections = (Path(tmp_path) / "ws" / "corrections.md").read_text() + assert "did not certify: activation/silu_and_mul[rows]" in corrections + assert state.verdict == "delivered" + + +def test_a_blocked_gap_with_no_rows_ends_the_run_with_its_reason(tmp_path, task_file, monkeypatch): + calls, state = _run( + tmp_path, + task_file, + monkeypatch, + assembler=lambda turn, prompt: { + "decision": "BLOCKED", + "summary": "needs verifies", + "blocked_reason": "certification_gap", + }, + ) + assert state.verdict == "certification-gap-unnamed" + assert "smith" not in _roles(calls) + + +def test_an_auditor_gap_runs_the_remediation_cycle(tmp_path, task_file, monkeypatch): + calls, state = _run( + tmp_path, + task_file, + monkeypatch, + auditor=lambda turn, prompt: ( + {"decision": "UNCOVERED", "summary": "one gap", "certification_gap": [_GAP]} + if turn == 1 + else {"decision": "CLEAN", "summary": "covered"} + ), + ) + assert _roles(calls).count("auditor") == 2 and _roles(calls).count("smith") == 1 + assert state.verdict == "delivered" + + +def test_the_roles_are_told_to_report_gaps_rather_than_reject(): + from agent_flow.workflows.staircase_onboard.prompts import DEFAULT_PROMPTS + + assert "is a certification gap, not a REJECT" in DEFAULT_PROMPTS.reviewer + assert "`certification_gap`" in DEFAULT_PROMPTS.assembler + assert "`certification_gap`" in DEFAULT_PROMPTS.auditor diff --git a/agent-flow/tests/workflows/staircase_onboard/test_early_checks.py b/agent-flow/tests/workflows/staircase_onboard/test_early_checks.py new file mode 100644 index 000000000000..4446088f5f95 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_early_checks.py @@ -0,0 +1,252 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the early delivery checks (``early_checks.py``). + +The fixtures are the DeepSeek-V4.1 run's delivery surprises: after every unit +was green, delivery failed on two comm entries whose tests were not named +``test_modeling_v2_.py``, a missing integration accuracy test and a +missing test-db line; and later its accuracy cases turned out to be listed +post-merge only, on a 527 GB checkpoint nobody had confirmed CI's model store +holds, with skip conditions that made several cases skip on CI. +""" + +from __future__ import annotations + +from pathlib import Path + +import yaml + +from agent_flow.workflows.staircase_onboard import delivery, early_checks, envelope +from agent_flow.workflows.staircase_onboard.capacity_audit import audit_target +from agent_flow.workflows.staircase_onboard.state import load_state + +from .test_workflow import Script, _install, _workflow, _write_vocabulary + +_ACCURACY_TEST = """\ +import os + +import pytest + +from tensorrt_llm._utils import get_sm_version + +from ..conftest import llm_models_root + +skip_not_sm103 = pytest.mark.skipif(get_sm_version() != 103, reason="certified on sm_103 only") + + +class TestModelingV2DeepseekV41(object): + MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4.1-Flash-NVFP4" + + @skip_not_sm103 + @pytest.mark.skip_less_device(4) + def test_gsm8k(self): + golden = os.environ.get("STAIRCASE_GOLDEN_DIR") + if golden is None: + pytest.skip("no golden reference") +""" + + +def _db(root: Path, name: str, stage: str, gpus, count, test: str) -> None: + db = root / delivery.TEST_DB_DIR + db.mkdir(parents=True, exist_ok=True) + (db / name).write_text( + yaml.safe_dump( + { + name.removesuffix(".yml"): [ + { + "condition": { + "ranges": {"system_gpu_count": {"gte": count, "lte": count}}, + "wildcards": {"gpu": gpus}, + "terms": {"stage": stage}, + }, + "tests": [test], + } + ] + } + ), + encoding="utf-8", + ) + + +def _repo(tmp_path: Path, *, accuracy: bool = True) -> Path: + root = tmp_path / "repo" + catalog = root / delivery.CATALOG_DIR + tests = root / delivery.CATALOG_TESTS_DIR + for category, name, test_name in ( + ("comm", "allreduce", "test_allreduce_op_matrix.py"), + ("norm", "rmsnorm", "test_modeling_v2_rmsnorm.py"), + ("torch", "embedding", None), + ): + (catalog / category).mkdir(parents=True, exist_ok=True) + (catalog / category / f"{name}.py").write_text("def f(): ...\n", encoding="utf-8") + if test_name: + (tests / category).mkdir(parents=True, exist_ok=True) + (tests / category / test_name).write_text("def test(): ...\n", encoding="utf-8") + (catalog / "gemm").mkdir(parents=True) + (catalog / "gemm" / "mm_mxfp8.py").write_text("def f(): ...\n", encoding="utf-8") + if accuracy: + acc = root / delivery.ACCURACY_TESTS_DIR + acc.mkdir(parents=True) + (acc / "test_modeling_v2_deepseek_v41.py").write_text(_ACCURACY_TEST, encoding="utf-8") + return root + + +def _vocabulary(tmp_path: Path) -> Path: + path = tmp_path / "vocabulary.yaml" + path.write_text( + yaml.safe_dump( + { + "steps": [ + {"step": "tp all-reduce", "entry": "catalog/comm/allreduce.py"}, + {"step": "norm", "entry": "catalog/norm/rmsnorm.py"}, + {"step": "embed", "entry": "catalog/torch/embedding.py"}, + { + "step": "gemm", + "entry": "catalog/gemm/mm_mxfp8.py", + "call_sites": [{"site": "x", "row_multiple": 0}], + }, + ], + "queue": [{"n": 1, "mode": "onboard", "target": "sampling/draft_sample"}], + } + ), + encoding="utf-8", + ) + return path + + +def _names(report, level=early_checks.LEVEL_WARN): + return {f.name for f in report.findings if f.level == level} + + +def test_test_names_ci_does_not_match_are_found_before_delivery(tmp_path): + report = early_checks.run_early_checks( + repo_root=_repo(tmp_path), + family="deepseek_v41", + stage="plan", + vocabulary_path=_vocabulary(tmp_path), + ) + warned = {f.name: f.detail for f in report.warnings} + assert "test_allreduce_op_matrix.py" in warned["test-name:comm/allreduce"] + assert "has no test" in warned["test-name:gemm/mm_mxfp8"] + # Correctly named, a mirror, and an entry smith has not written yet: no warning. + assert not {"test-name:norm/rmsnorm", "test-name:torch/embedding"} & set(warned) + assert not any("draft_sample" in name for name in warned) + assert "vocabulary:call_sites" in warned + + +def test_a_missing_accuracy_test_is_a_warning_and_no_assets_are_listed(tmp_path): + report = early_checks.run_early_checks( + repo_root=_repo(tmp_path, accuracy=False), family="deepseek_v41", stage="plan" + ) + assert "tests:accuracy-file" in _names(report) + assert report.assets == () + assert "(no accuracy test to read yet)" in report.render() + + +def test_the_required_ci_assets_are_read_from_the_test(tmp_path): + root = _repo(tmp_path) + assets = early_checks.analyze_skips( + root / delivery.ACCURACY_TESTS_DIR / "test_modeling_v2_deepseek_v41.py" + ) + found = {(a.kind, a.what) for a in assets} + assert ("checkpoint", "$LLM_MODELS_ROOT/DeepSeek-V4.1-Flash-NVFP4") in found + assert ("arch", "sm_103") in found + assert ("gpus", ">= 4") in found + assert ("env", "STAIRCASE_GOLDEN_DIR") in found + assert ("skip", "pytest.skip(...) at runtime") in found + + +def test_ci_skips_are_predicted_from_the_test_db(tmp_path): + root = _repo(tmp_path) + test = "accuracy/test_modeling_v2_deepseek_v41.py::TestModelingV2DeepseekV41::test_gsm8k" + _db(root, "l0_gb300_multi_gpus.yml", "post_merge", ["*gb300*"], 4, test) + report = early_checks.run_early_checks(repo_root=root, family="deepseek_v41", stage="unit") + warned = {f.name: f.detail for f in report.warnings} + assert "post-merge only" in warned["tests:test-db"] + assert "ci:env:STAIRCASE_GOLDEN_DIR" in warned + assert "ci:arch" not in warned and "ci:gpus" not in warned + + _db(root, "l0_b200.yml", "pre_merge", ["*b200*"], 1, test) + report = early_checks.run_early_checks(repo_root=root, family="deepseek_v41", stage="unit") + details = [f.detail for f in report.warnings] + assert any("needs sm_103" in d and "l0_b200.yml" in d and "always skips" in d for d in details) + assert any("needs >= 4 GPUs and l0_b200.yml runs on at most 1" in d for d in details) + assert "post-merge only" not in " ".join(details) + assert "## Where CI runs it" in report.render() + + +def test_the_checkpoint_is_checked_where_the_model_store_is_known(tmp_path): + root = _repo(tmp_path) + store = tmp_path / "models" + store.mkdir() + env = {early_checks.MODELS_ROOT_ENV: str(store)} + report = early_checks.run_early_checks( + repo_root=root, family="deepseek_v41", stage="u", env=env + ) + assert ( + "does not exist under this host's LLM_MODELS_ROOT" + in {f.name: f.detail for f in report.warnings}["ci:checkpoint"] + ) + + (store / "DeepSeek-V4.1-Flash-NVFP4").mkdir() + report = early_checks.run_early_checks( + repo_root=root, family="deepseek_v41", stage="u", env=env + ) + assert "ci:checkpoint" in _names(report, early_checks.LEVEL_OK) + + +def test_capacity_and_matrix_failures_surface_early(tmp_path): + target = tmp_path / "target" + target.mkdir() + (target / "m.py").write_text("_MAX_ROWS = 8192\n", encoding="utf-8") + env = envelope.ServingEnvelope(max_batch_size=8, max_num_tokens=16384, max_seq_len=4096) + plan = envelope.CertificationPlan( + envelope=env, + cells=( + envelope.MatrixCell( + cell="m", + features=(), + member="dep8", + evidence="boot", + at="ceiling", + status=envelope.CELL_EXTERNAL, + reason="two nodes", + ), + ), + ) + report = early_checks.run_early_checks( + repo_root=_repo(tmp_path), + family="deepseek_v41", + stage="unit", + capacity=audit_target(target, env), + certification=plan, + workspace=tmp_path / "ws", + ) + assert {"capacity:_MAX_ROWS", "matrix:m__dep8__boot__ceiling"} <= _names(report) + + +def test_the_workflow_reports_after_the_plan_and_every_unit_and_does_not_stop( + tmp_path, task_file, monkeypatch +): + workflow = _workflow(tmp_path) + _install(monkeypatch, workflow, Script(None)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + workspace = tmp_path / "ws" + journal = (workspace / "journal.md").read_text() + assert "early delivery checks (end of the plan phase)" in journal + assert "early delivery checks (after the core unit)" in journal + text = (workspace / "early_checks.md").read_text() + assert "# Early delivery checks (after the core unit)" in text + assert "## Required CI assets" in text + assert load_state(workspace / ".staircase_onboard_state.json").verdict == "delivered" diff --git a/agent-flow/tests/workflows/staircase_onboard/test_envelope.py b/agent-flow/tests/workflows/staircase_onboard/test_envelope.py new file mode 100644 index 000000000000..05b82fe9aa55 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_envelope.py @@ -0,0 +1,615 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the serving envelope and the feature matrix (``envelope.py``). + +The fixture spec is the shape of the DeepSeek-V4.1 onboarding's gates: boot at +``max_num_tokens`` 2048 / ``max_batch_size`` 8 / ``max_seq_len`` 16384 with +graphs and no padding, accuracy at 8192 / 32 / 16384 with padded graphs to 32, +a kv_quant witness at batch 1. That run took "the tuples the gates launch" as +its certified set, and three crashes lived just above those tuples. The tests +pin both readings: the legacy hull reproduces that set exactly, and a declared +envelope probes above it. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import envelope, preflight, task_schema +from agent_flow.workflows.staircase_onboard.gates import ShellResult, resolve_runtime +from agent_flow.workflows.staircase_onboard.state import load_state + +from .test_workflow import Script, _install, _workflow, _write_vocabulary, passing_shell + + +def _write(path: Path, data) -> Path: + path.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8") + return path + + +def _history_spec(tmp_path: Path, **overrides) -> dict: + """The DeepSeek-V4.1 run's gate configuration, as a unified spec.""" + checkpoint = tmp_path / "ckpt" + checkpoint.mkdir(exist_ok=True) + (checkpoint / "config.json").write_text( + json.dumps( + { + "architectures": ["DeepseekV41ForCausalLM"], + "text_config": {"max_position_embeddings": 1048576}, + } + ), + encoding="utf-8", + ) + repo = tmp_path / "repo" + repo.mkdir(exist_ok=True) + options = tmp_path / "scripts" + options.mkdir(exist_ok=True) + graph = {"enable_padding": True, "max_batch_size": 32} + _write( + options / "core_llm_api.yaml", + {"max_num_tokens": 8192, "cuda_graph_config": graph, "enable_chunked_prefill": True}, + ) + _write( + options / "dspark_llm_api.yaml", + { + "max_num_tokens": 8192, + "cuda_graph_config": graph, + "speculative_config": {"decoding_type": "DSpark", "max_draft_len": 5}, + }, + ) + spec = { + "checkpoint_path": str(checkpoint), + "arch": "sm_103", + "parallel": "unified", + "members": {"default": ["tep4", "dep4", "tp4"]}, + "features": ["dspark", "kv_quant"], + "trtllm_repo_path": str(repo), + "runtime": { + "default": { + "cuda_graph": True, + "chunked_prefill": True, + "extra_flags": [ + "--max_seq_len=16384", + "--max_num_tokens=2048", + "--max_batch_size=8", + "--kv_cache_fraction=0.5", + ], + } + }, + "environment": { + "kind": "local", + "python": "/usr/bin/python3", + "gpus_per_node": 4, + "env": {"TMPDIR": "/scratch/tmp", "CLAUDE_CODE_TMPDIR": "/scratch/tmp"}, + }, + "accuracy": { + "benchmark": "gsm8k", + "reference_yaml": "tests/integration/defs/accuracy/references/gsm8k.yaml", + "llm_api_options": { + "default": str(options / "dspark_llm_api.yaml"), + "core": str(options / "core_llm_api.yaml"), + }, + "max_seq_len": 16384, + "max_batch_size": 32, + "witness": { + "kv_quant": { + "contrast_env": {"KV_STORAGE": "bf16"}, + "expected_bytes_per_token": 890, + "max_batch_size": 1, + } + }, + }, + } + spec.update(overrides) + return spec + + +#: What the DeepSeek-V4.1 target is deployed at, which no gate reached. +DECLARED = { + "max_batch_size": 256, + "max_num_tokens": 16384, + "max_seq_len": 1048576, + "max_draft_len": 5, + "cuda_graph": {"padding": True}, + "members": ["tep4", "dep4", "tp4", "dep8", "tep4pp2"], +} + + +def _load(tmp_path: Path, spec: dict) -> dict: + return task_schema.load_and_validate_task_yaml(_write(tmp_path / "task.yaml", spec)) + + +def _kinds(unit: str) -> set[str]: + kinds = {"boot", "accuracy"} + if unit == "kv_quant": + kinds.add("witness") + if unit == "dspark": + kinds.add("acceptance") + return kinds + + +def _plan(tmp_path: Path, spec: dict) -> envelope.CertificationPlan: + data = _load(tmp_path, spec) + config = json.loads((Path(data["checkpoint_path"]) / "config.json").read_text()) + return envelope.derive_certification( + data, world_size=4, gpus_per_node=4, checkpoint_config=config, unit_kinds=_kinds + ) + + +# --------------------------------------------------------------- primitives + + +def test_the_graph_generator_mirrors_the_engine(): + assert envelope.generate_graph_batch_sizes(8, False) == list(range(1, 9)) + assert envelope.generate_graph_batch_sizes(32, True) == [1, 2, 4, 8, 16, 24, 32] + padded = envelope.generate_graph_batch_sizes(256, True) + assert padded[-3:] == [128, 192, 256] + assert envelope.generate_graph_batch_sizes(300, False)[-2:] == [256, 300] + + +def test_envelope_members_carry_pipeline_and_node_count(): + member = envelope.parse_envelope_member("tep4pp2") + assert (member.kind, member.size, member.pp, member.world_size) == ("tep", 4, 2, 8) + assert member.nodes(4) == 2 + assert "PP2" in member.describe(4) and "2 nodes" in member.describe(4) + assert envelope.parse_envelope_member("dep8").attention_dp + for bad in ("unified", "tp", "pp2", "tep4p2", "ep4"): + with pytest.raises(ValueError): + envelope.parse_envelope_member(bad) + + +# ------------------------------------------------------------------- legacy + + +def test_without_an_envelope_the_certified_set_is_the_gate_hull(tmp_path): + """The legacy reading, stated: exactly R30's set, and nothing above it.""" + plan = _plan(tmp_path, _history_spec(tmp_path)) + + assert not plan.declared + env = plan.envelope + assert (env.max_batch_size, env.max_num_tokens, env.max_seq_len) == (32, 8192, 16384) + # Boot captures 1..8 without padding; accuracy captures the padded list to 32. + assert plan.certified_capture_batches == (1, 2, 3, 4, 5, 6, 7, 8, 16, 24, 32) + assert plan.probes == () + assert plan.external_cells == [] + assert plan.errors == () + note = plan.notes()[0] + assert "no serving_envelope declared" in note and "migrate" in note + assert "DERIVED" in plan.render() + + +def test_gate_tuples_are_read_from_the_runtime_block_options_and_witness(tmp_path): + plan = _plan(tmp_path, _history_spec(tmp_path)) + by_gate = {(t.gate, t.unit, t.member): t for t in plan.tuples} + + boot = by_gate[("boot", "core", "tep4")] + assert (boot.max_batch_size, boot.max_num_tokens, boot.max_seq_len) == (8, 2048, 16384) + assert boot.graph_batch_sizes == tuple(range(1, 9)) + + accuracy = by_gate[("accuracy", "dspark", "dep4")] + assert (accuracy.max_batch_size, accuracy.max_num_tokens) == (32, 8192) + assert accuracy.graph_batch_sizes == (1, 2, 4, 8, 16, 24, 32) + + # The witness runs once, at the cheapest member, at its own batch limit. + witnesses = [t for t in plan.tuples if t.gate == "witness"] + assert [(t.unit, t.member, t.max_batch_size) for t in witnesses] == [("kv_quant", "tep4", 1)] + assert ("acceptance", "dspark", "tp4") in by_gate + + +# ----------------------------------------------------------------- declared + + +def test_a_declared_envelope_derives_a_ceiling_probe_per_unit(tmp_path): + plan = _plan(tmp_path, _history_spec(tmp_path, serving_envelope=DECLARED)) + + assert plan.declared and plan.errors == () + assert [p.unit for p in plan.probes] == ["core", "dspark", "kv_quant"] + probe = plan.probes[0] + # The cheapest member, at every limit of the envelope. + assert probe.member == "tep4" + assert (probe.max_batch_size, probe.max_num_tokens, probe.max_seq_len) == ( + 256, + 16384, + 1048576, + ) + assert probe.graph_batch_sizes[-2:] == (192, 256) + flags = envelope.probe_flags(probe) + assert "--max_num_tokens=16384" in flags + assert "--cuda_graph_padding_enabled" in flags # the envelope's padding policy + assert any(flag.startswith("--cuda_graph_batch_sizes 1 2 4 8") for flag in flags) + + +def test_members_no_unit_gates_owe_external_evidence(tmp_path): + """Multi-node and PP members are claims this allocation cannot check.""" + plan = _plan(tmp_path, _history_spec(tmp_path, serving_envelope=DECLARED)) + + external = {cell.member: cell for cell in plan.external_cells} + assert set(external) == {"dep8", "tep4pp2"} + assert external["dep8"].features == ("dspark", "kv_quant") + assert external["dep8"].evidence == "boot" and external["dep8"].at == "ceiling" + assert any("tep4pp2" in note and "evidence" in note for note in plan.notes()) + assert any("spans nodes" in w for w in plan.warnings) + + +def test_a_gate_launched_outside_the_envelope_is_an_error(tmp_path): + narrow = {**DECLARED, "max_num_tokens": 4096} + plan = _plan(tmp_path, _history_spec(tmp_path, serving_envelope=narrow)) + assert any( + "accuracy core@tep4" in e + and "max_num_tokens=8192 > serving_envelope.max_num_tokens=4096" in e + for e in plan.errors + ) + + +def test_a_gated_member_the_envelope_does_not_claim_is_an_error(tmp_path): + spec = _history_spec(tmp_path, serving_envelope={**DECLARED, "members": ["tep4", "dep4"]}) + plan = _plan(tmp_path, spec) + assert any("member tp4 is gated by this run but is not in" in e for e in plan.errors) + + +def test_an_unset_max_seq_len_reads_as_the_checkpoint_maximum(tmp_path): + """Leaving it to the engine sizes for 1M tokens, which the envelope may not claim.""" + spec = _history_spec(tmp_path, serving_envelope={**DECLARED, "max_seq_len": 131072}) + spec["runtime"]["default"]["extra_flags"] = ["--max_num_tokens=2048", "--max_batch_size=8"] + plan = _plan(tmp_path, spec) + errors = [e for e in plan.errors if "boot core" in e] + assert errors and "max_seq_len=1048576 > serving_envelope.max_seq_len=131072" in errors[0] + assert "pass --max_seq_len" in errors[0] + + +def test_an_envelope_beyond_the_checkpoint_is_an_error(tmp_path): + spec = _history_spec(tmp_path, serving_envelope={**DECLARED, "max_seq_len": 2097152}) + plan = _plan(tmp_path, spec) + assert any("exceeds the checkpoint's max_position_embeddings" in e for e in plan.errors) + + +def test_a_batch_limit_the_token_budget_cannot_fill_is_a_warning(tmp_path): + spec = _history_spec(tmp_path, serving_envelope={**DECLARED, "max_num_tokens": 1024}) + spec["runtime"]["default"]["extra_flags"] = [ + "--max_seq_len=16384", + "--max_num_tokens=1024", + "--max_batch_size=8", + ] + for name in ("core_llm_api.yaml", "dspark_llm_api.yaml"): + path = tmp_path / "scripts" / name + data = yaml.safe_load(path.read_text()) + data["max_num_tokens"] = 1024 + _write(path, data) + plan = _plan(tmp_path, spec) + assert any("256 x 6 = 1536 tokens" in w for w in plan.warnings) + + +# ------------------------------------------------------------------- matrix + + +def test_matrix_cells_are_classified_by_how_their_evidence_is_obtained(tmp_path): + matrix = [ + {"name": "spec-acc", "features": ["dspark"], "members": ["dep4"], "evidence": "accuracy"}, + { + "name": "spec-long", + "features": ["dspark"], + "members": ["tep4"], + "evidence": ["boot", "accuracy"], + "at": {"max_seq_len": 131072}, + "prompt_tokens": 12000, + }, + { + "name": "kvq-witness", + "features": ["kv_quant"], + "members": ["tp4"], + "evidence": "witness", + }, + {"name": "mn", "features": [], "members": ["dep8"], "evidence": "accuracy"}, + ] + spec = _history_spec(tmp_path, serving_envelope=DECLARED, feature_matrix=matrix) + plan = _plan(tmp_path, spec) + cells = { + (c.cell, c.member, c.evidence): c + for c in plan.cells + if c.cell in {m["name"] for m in matrix} + } + + assert cells[("spec-acc", "dep4", "accuracy")].status == envelope.CELL_GATED + long_boot = cells[("spec-long", "tep4", "boot")] + assert long_boot.status == envelope.CELL_PROBE and long_boot.unit == "dspark" + long_accuracy = cells[("spec-long", "tep4", "accuracy")] + assert long_accuracy.status == envelope.CELL_EXTERNAL + # Measured against the accuracy gate's own tuple, with only the named axis moved. + assert "runs below B=32 T=8192 L=131072" in long_accuracy.reason + assert "prompt>=" not in long_accuracy.reason + assert "witness runs at tep4 only" in cells[("kvq-witness", "tp4", "witness")].reason + assert cells[("mn", "dep8", "accuracy")].status == envelope.CELL_EXTERNAL + + point = [p for p in plan.probes if p.unit == "dspark" and p.max_seq_len == 131072] + assert len(point) == 1 + # Only the named axis moves; the rest stay at the unit's boot tuple. + assert (point[0].max_num_tokens, point[0].max_batch_size) == (2048, 8) + assert point[0].prompt_tokens == 12000 + + +def test_a_declared_cell_an_existing_gate_covers_is_not_owed_twice(tmp_path): + matrix = [{"features": [], "members": ["tep4"], "evidence": ["boot", "accuracy"]}] + plan = _plan( + tmp_path, _history_spec(tmp_path, serving_envelope=DECLARED, feature_matrix=matrix) + ) + core = [c for c in plan.cells if c.features == () and c.member == "tep4" and c.at == "gate"] + assert {c.evidence for c in core} == {"boot", "accuracy"} + assert all(c.status == envelope.CELL_GATED for c in core) + + +# ------------------------------------------------------------------- probes + + +def test_a_probe_is_the_boot_gate_with_only_the_limits_changed(tmp_path): + data = _load(tmp_path, _history_spec(tmp_path, serving_envelope=DECLARED)) + plan = envelope.derive_certification(data, world_size=4, unit_kinds=_kinds) + probe = plan.probes_for("core")[0] + + probed = envelope.probe_task_data(data, "core", probe) + runtime = resolve_runtime(probed, "dspark") + flags = runtime["extra_flags"] + assert "--kv_cache_fraction=0.5" in flags + assert not any(f.startswith("--max_num_tokens=2048") for f in flags) + assert runtime["chunked_prefill"] is True and runtime["cuda_graph"] is True + # Round trip: reading the probe's runtime block back gives the probe's point. + back = envelope.runtime_tuple(runtime, gate="boot", unit="core", member="tep4") + assert ( + back.point() + == envelope.EngineTuple( + gate="boot", + unit="core", + member="tep4", + max_batch_size=probe.max_batch_size, + max_num_tokens=probe.max_num_tokens, + max_seq_len=probe.max_seq_len, + graph_batch_sizes=probe.graph_batch_sizes, + ).point() + ) + # The spec itself is untouched. + assert "--max_num_tokens=2048" in data["runtime"]["default"]["extra_flags"] + + +def test_a_probe_with_a_prompt_length_adds_one_long_prompt(): + probe = envelope.EngineTuple( + gate="envelope", + unit="core", + member="tp1", + max_batch_size=1, + max_num_tokens=1, + max_seq_len=1, + prompt_tokens=300, + ) + prompts = envelope.probe_prompts(("a", "b"), probe) + assert prompts[:2] == ("a", "b") and len(prompts[2].split()) == 300 + + +# ----------------------------------------------------------------- evidence + + +def test_external_cells_need_a_passing_evidence_file(tmp_path): + plan = _plan(tmp_path, _history_spec(tmp_path, serving_envelope=DECLARED)) + workspace = tmp_path / "ws" + checks = envelope.check_external_evidence(plan, workspace) + assert len(checks) == 2 and not any(c.passed for c in checks) + + cell = plan.external_cells[0] + root = workspace / envelope.EVIDENCE_DIRNAME + root.mkdir(parents=True) + _write(root / f"{cell.slug}.yaml", {"verdict": "fail", "tree": "abc", "log": "x.log"}) + assert not envelope.check_external_evidence(plan, workspace)[0].passed + _write(root / f"{cell.slug}.yaml", {"verdict": "pass", "tree": "abc", "log": "x.log"}) + checks = {c.name: c for c in envelope.check_external_evidence(plan, workspace)} + assert checks[f"matrix:{cell.slug}"].passed + + +# ------------------------------------------------------------------- schema + + +@pytest.mark.parametrize( + ("change", "message"), + [ + ({"cuda_graph": None}, "'serving_envelope.cuda_graph' is required"), + ({"cuda_graph": {"batch_sizes": [1, 512]}}, "captures 512, over"), + ({"cuda_graph": {}}, "needs 'batch_sizes'"), + ({"members": ["tep4", "tep4"]}, "lists 'tep4' twice"), + ({"members": ["unified"]}, "is not an envelope member"), + ({"max_num_tokens": 0}, "'serving_envelope.max_num_tokens' must be an integer >= 1"), + ({"probe_prompt_tokens": 99999}, "caps at 128 KiB"), + ({"max_num_token": 1}, "did you mean 'serving_envelope.max_num_tokens'"), + ], +) +def test_the_envelope_block_is_validated(tmp_path, change, message): + block = {**DECLARED, **change} + block = {k: v for k, v in block.items() if v is not None} + with pytest.raises( + task_schema.TaskSchemaError, match=message.replace("(", r"\(").replace(")", r"\)") + ): + _load(tmp_path, _history_spec(tmp_path, serving_envelope=block)) + + +@pytest.mark.parametrize( + ("cell", "message"), + [ + ({"features": ["mtp"]}, "which this run does not build"), + ({"members": ["tep16"]}, "does not claim"), + ({"evidence": "perf"}, "evidence' must be one of"), + ({"at": {"max_seq_len": 2097152}}, "is outside serving_envelope.max_seq_len"), + ({"at": "somewhere"}, "must be 'gate', 'ceiling' or a mapping"), + ({"atx": "gate"}, "unknown key"), + ({"name": "a b"}, "file-name-safe"), + ], +) +def test_matrix_cells_are_validated(tmp_path, cell, message): + spec = _history_spec(tmp_path, serving_envelope=DECLARED, feature_matrix=[cell]) + with pytest.raises(task_schema.TaskSchemaError, match=message): + _load(tmp_path, spec) + + +def test_cell_names_are_unique_because_they_name_evidence_files(tmp_path): + cells = [{"name": "x", "features": []}, {"name": "x", "features": ["dspark"]}] + spec = _history_spec(tmp_path, serving_envelope=DECLARED, feature_matrix=cells) + with pytest.raises(task_schema.TaskSchemaError, match="is used by another cell"): + _load(tmp_path, spec) + + +def test_a_matrix_needs_an_envelope(tmp_path): + spec = _history_spec(tmp_path, feature_matrix=[{"features": []}]) + with pytest.raises(task_schema.TaskSchemaError, match="needs 'serving_envelope'"): + _load(tmp_path, spec) + + +# ---------------------------------------------------------------- migration + + +def test_the_migration_block_restates_the_legacy_set_as_a_declared_one(tmp_path): + data = _load(tmp_path, _history_spec(tmp_path)) + block = envelope.migration_block(data, world_size=4) + assert "WIDEN IT" in block + + path = tmp_path / "task.yaml" + path.write_text(path.read_text() + "\n" + block, encoding="utf-8") + migrated = task_schema.load_and_validate_task_yaml(path) + plan = envelope.derive_certification(migrated, world_size=4, unit_kinds=_kinds) + assert plan.declared and plan.errors == () + # Adopted unchanged, it owes nothing the gates did not already produce. + assert plan.external_cells == [] + env = plan.envelope + assert (env.max_batch_size, env.max_num_tokens, env.max_seq_len) == (32, 8192, 16384) + assert plan.certified_capture_batches == (1, 2, 3, 4, 5, 6, 7, 8, 16, 24, 32) + + +def test_the_migrate_command_appends_once(tmp_path): + _write(tmp_path / "task.yaml", _history_spec(tmp_path)) + path = str(tmp_path / "task.yaml") + assert envelope._main(["migrate", path, "--write"]) == 0 + assert "serving_envelope:" in Path(path).read_text() + assert envelope._main(["migrate", path]) == 1 + assert envelope._main(["plan", path]) == 0 + + +# ---------------------------------------------------------------- preflight + + +def _smi(argv, env=None, timeout=0): + return preflight.CommandResult(returncode=0, stdout="ok\n") + + +def test_preflight_refuses_an_envelope_the_gates_contradict(tmp_path): + spec = _history_spec(tmp_path, serving_envelope={**DECLARED, "max_num_tokens": 4096}) + data = _load(tmp_path, spec) + with pytest.raises( + preflight.SpecInconsistencyError, match="serving envelope and the gates disagree" + ): + preflight.run_preflight(data, runner=_smi, skip_device_claim=True) + + +def test_preflight_carries_the_plan_and_prints_the_bill(tmp_path): + data = _load(tmp_path, _history_spec(tmp_path, serving_envelope=DECLARED)) + result = preflight.run_preflight(data, runner=_smi, skip_device_claim=True) + assert result.certification is not None and result.certification.declared + joined = "\n".join(result.notes) + assert "serving envelope: B<=256 T<=16384 L<=1048576" in joined + assert "envelope probe(s)" in joined + + +# ----------------------------------------------------------------- workflow + + +@pytest.fixture +def enveloped_task(tmp_path, task_file): + spec = yaml.safe_load(task_file.read_text()) + spec["runtime"] = {"extra_flags": ["--max_seq_len=4096", "--max_batch_size=4"]} + spec["accuracy"].update(max_seq_len=4096, max_batch_size=4) + spec["serving_envelope"] = { + "max_batch_size": 64, + "max_num_tokens": 16384, + "max_seq_len": 32768, + "cuda_graph": {"batch_sizes": [1, 2, 4, 8, 16, 32, 64]}, + "members": ["tp1", "tp2"], + } + task_file.write_text(yaml.safe_dump(spec), encoding="utf-8") + return task_file + + +def _recording_shell(commands, *, fail_probe: int = 0): + inner = passing_shell() + failures = {"left": fail_probe} + + def runner(command, env=None, timeout=0): + commands.append(command) + if "--max_num_tokens=16384" in command and failures["left"]: + failures["left"] -= 1 + return ShellResult(returncode=1, stdout="", stderr="index out of bounds in warm-up") + return inner(command, env, timeout) + + return runner + + +def _with_vocabulary(monkeypatch, workflow, tmp_path): + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + + +def test_the_workflow_probes_the_ceiling_and_owes_the_unclaimed_member( + tmp_path, enveloped_task, monkeypatch +): + commands: list[str] = [] + workflow = _workflow(tmp_path, shell_runner=_recording_shell(commands)) + calls = _install(monkeypatch, workflow, Script(None)) + _with_vocabulary(monkeypatch, workflow, tmp_path) + with workflow: + workflow.run(enveloped_task) + + probes = [c for c in commands if "--max_num_tokens=16384" in c] + assert len(probes) == 1 + assert "--max_seq_len=32768" in probes[0] and "--max_batch_size=64" in probes[0] + assert "--cuda_graph_batch_sizes 1 2 4 8 16 32 64" in probes[0] + assert "--use_cuda_graph" in probes[0] + + workspace = tmp_path / "ws" + assert "tp2" in (workspace / "certification.md").read_text() + drafter = next(prompt for role, prompt in calls if role == "plan_drafter") + assert "## Serving envelope (declared)" in drafter + # tp2 is claimed and gated by nothing, so delivery stops on it. + state = load_state(workspace / ".staircase_onboard_state.json") + assert state.verdict == "delivery-failed" + assert "matrix:envelope-member__tp2__boot__ceiling" in (workspace / "journal.md").read_text() + + +def test_a_failed_probe_sends_the_unit_back_to_the_assembler(tmp_path, enveloped_task, monkeypatch): + commands: list[str] = [] + workflow = _workflow(tmp_path, shell_runner=_recording_shell(commands, fail_probe=1)) + calls = _install(monkeypatch, workflow, Script(None)) + _with_vocabulary(monkeypatch, workflow, tmp_path) + evidence = tmp_path / "ws" / envelope.EVIDENCE_DIRNAME + evidence.mkdir(parents=True) + _write( + evidence / "envelope-member__tp2__boot__ceiling.yaml", + {"verdict": "pass", "tree": "c0ffee1", "log": "/logs/tp2.log"}, + ) + with workflow: + workflow.run(enveloped_task) + + assert [role for role, _ in calls].count("assembler") == 2 + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "[FAIL] envelope@tp1: at B=64 T=16384 L=32768" in journal + assert load_state(tmp_path / "ws" / ".staircase_onboard_state.json").verdict == "delivered" + + +def test_the_plan_roles_judge_coverage_against_the_envelope(): + from agent_flow.workflows.staircase_onboard.prompts import DEFAULT_PROMPTS + + assert "serving envelope" in DEFAULT_PROMPTS.plan_drafter + assert "is `CERT`, not `OK`" in DEFAULT_PROMPTS.plan_drafter + assert "serving envelope" in DEFAULT_PROMPTS.plan_reviewer diff --git a/agent-flow/tests/workflows/staircase_onboard/test_gate_wiring.py b/agent-flow/tests/workflows/staircase_onboard/test_gate_wiring.py index f61b4ec36518..6fcd866c7b27 100644 --- a/agent-flow/tests/workflows/staircase_onboard/test_gate_wiring.py +++ b/agent-flow/tests/workflows/staircase_onboard/test_gate_wiring.py @@ -425,11 +425,16 @@ def test_a_resume_only_auto_handles_a_vocabulary_gap(): """ import inspect - from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + from agent_flow.workflows.staircase_onboard.workflow import ( + RESUMABLE_BLOCKED_REASONS, + StaircaseOnboardWorkflow, + ) source = inspect.getsource(StaircaseOnboardWorkflow._assemble_one_unit) head = source[: source.index("self._dispatch(")] - assert 'prior_reason == "vocabulary_gap"' in head + assert "prior_reason in RESUMABLE_BLOCKED_REASONS" in head + # The two reasons the orchestrator can act on itself, and no others. + assert RESUMABLE_BLOCKED_REASONS == ("vocabulary_gap", "certification_gap") # ------------------------------------------- the auditor's runtime-value dump diff --git a/agent-flow/tests/workflows/staircase_onboard/test_smith_brief.py b/agent-flow/tests/workflows/staircase_onboard/test_smith_brief.py new file mode 100644 index 000000000000..87ad7807cf2d --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_smith_brief.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the smith brief (``smith_brief.py``) and the kernel-verify skill's probes. + +The fixture is the DeepSeek-V4.1 delivery audit that came back UNCOVERED on +seven entries for one reason: capture lists held ``G`` and ``6·G`` while the +target also reached those ops at ``3·G`` (ratio-2 compressed rows) and ``5·G`` +(rope's q form), ``dep4`` saw ``4·G`` / ``24·G`` after the gather, and one call +site (the decode-replay tail) only ever ran on context rows. The capture +batches are that run's: boot ``1..8`` unpadded plus the accuracy gate's padded +list to 32. +""" + +from __future__ import annotations + +from pathlib import Path + +import yaml + +from agent_flow.workflows.staircase_onboard import envelope, smith_brief +from agent_flow.workflows.staircase_onboard.vocabulary import QueueItem + +from .test_envelope import _history_spec, _load +from .test_workflow import Script, _install, _workflow, _write_vocabulary + +STEPS = [ + { + "step": "compressed latent, ratio 2", + "state": "CERT", + "entry": "catalog/attention/compressor.py", + "call_sites": [ + {"site": "Compressor.forward, verify step", "row_multiple": 3, "row_width": 512}, + {"site": "_TailSlots, decode-replay tail", "eager_only": True}, + ], + }, + { + "step": "rope, q form", + "state": "OK", + "entry": "catalog/rope/rope_q.py", + "call_sites": [{"site": "drafter q", "row_multiple": 5}], + }, + { + "step": "routed experts", + "state": "CERT", + "entry": "catalog/moe/nvfp4_runner.py", + "call_sites": [ + {"site": "MoE after the DP gather, decode", "row_multiple": 1, "gathered": True}, + {"site": "MoE after the DP gather, verify", "row_multiple": 6, "gathered": True}, + ], + }, +] + + +def _plan(tmp_path: Path, **envelope_block) -> envelope.CertificationPlan: + spec = _history_spec(tmp_path) + if envelope_block: + spec["serving_envelope"] = envelope_block + return envelope.derive_certification(_load(tmp_path, spec), world_size=4) + + +def _item(target: str, n: int = 75, mode: str = "verify") -> QueueItem: + return QueueItem(n=n, mode=mode, target=target) + + +def test_each_call_site_states_its_row_multiple_over_every_captured_batch(tmp_path): + brief = smith_brief.render_smith_brief( + _item("attention/compressor"), steps=STEPS, plan=_plan(tmp_path) + ) + # The legacy run's capture batches: 1..8 and 16, 24, 32. + assert "captured batch sizes G: 1..8,16,24,32" in brief + assert "| Compressor.forward, verify step | 3·G | 3,6,9,12,15,18,21,24,48,72,96 | - |" in brief + assert "| _TailSlots, decode-replay tail | eager-only | never captured | - |" in brief + assert "every eager row count 1..8192" in brief + + +def test_a_gathered_site_is_multiplied_by_every_attention_dp_width(tmp_path): + brief = smith_brief.render_smith_brief( + _item("moe/nvfp4_runner"), steps=STEPS, plan=_plan(tmp_path) + ) + assert "dep4: 4·G = 4,8,12,16,20,24,28,32,64,96,128" in brief + assert "dep4: 24·G = 24,48,72,96,120,144,168,192,384,576,768" in brief + # Eager rows after the gather reach the gathered token budget. + assert "every eager row count 1..32768" in brief + + +def test_the_declared_envelope_widens_the_rows_and_the_probes_flag_reachable_ones(tmp_path): + plan = _plan( + tmp_path, + max_batch_size=256, + max_num_tokens=16384, + max_seq_len=1048576, + max_draft_len=5, + cuda_graph={"padding": True}, + members=["tep4", "dep4", "tp4", "dep8"], + ) + brief = smith_brief.render_smith_brief(_item("moe/nvfp4_runner"), steps=STEPS, plan=plan) + assert "dep8: 48·G" in brief and "12288" in brief # 48 x 256 + assert "every eager row count 1..131072" in brief # 16384 x 8 after the dep8 gather + + compressor = smith_brief.render_smith_brief(_item("compressor"), steps=STEPS, plan=plan) + assert ( + "2^31 / 512 = 4194304 rows, probe at 4194305 (the envelope reaches 16384 rows)" + in compressor + ) + assert "2^24 + 1" in compressor and "2^31" in compressor + + +def test_without_declared_call_sites_the_brief_assumes_decode_and_verify(tmp_path): + plan = _plan( + tmp_path, + max_batch_size=32, + max_num_tokens=8192, + max_seq_len=16384, + max_draft_len=5, + cuda_graph={"padding": True}, + members=["tep4", "dep4", "tp4"], + ) + brief = smith_brief.render_smith_brief(_item("norm/rmsnorm"), steps=STEPS, plan=plan) + assert "ASSUMED" in brief + assert "| decode step | G |" in brief and "| speculative verify step | 6·G |" in brief + assert "read them from the contract" in brief + + +def test_malformed_call_sites_are_named(tmp_path): + steps = [ + { + "step": "x", + "entry": "catalog/a/b.py", + "call_sites": [ + {"site": "s", "row_multiple": 0}, + {"row_multiple": 2}, + {"site": "t", "rows": 3}, + ], + } + ] + sites, errors = smith_brief.parse_call_sites(steps[0]) + assert [s.site for s in sites] == ["t"] + assert len(errors) == 3 + brief = smith_brief.render_smith_brief(_item("a/b"), steps=steps, plan=_plan(tmp_path)) + assert "Malformed `call_sites`" in brief + + +def test_entry_keys_match_every_spelling(): + for spelled in ( + "catalog/attention/compressor.py", + "attention/compressor", + "tensorrt_llm/_torch/_experimental/modeling_v2/catalog/attention/compressor.py", + ): + assert smith_brief.entry_key(spelled) == "attention/compressor" + assert len(smith_brief.steps_for(STEPS, "compressor")) == 1 + assert smith_brief.steps_for(STEPS, "other/compressor") == [] + + +def test_the_brief_carries_the_queue_row(tmp_path): + row = {"n": 75, "tier": 1, "unit": "dspark", "surface": "rows 3G", "note": "R20 union"} + brief = smith_brief.render_smith_brief( + _item("attention/compressor"), steps=STEPS, plan=None, queue_row=row + ) + for line in ("- tier: 1", "- unit: dspark", "- surface: rows 3G", "- note: R20 union"): + assert line in brief + # No plan: no envelope section, and the probes are still owed. + assert "### Envelope" not in brief and "Integer-width boundary probes" in brief + + +def test_the_command_prints_the_brief_for_a_written_dispatch(tmp_path, capsys): + workspace = tmp_path / "ws" + (workspace / "dispatch").mkdir(parents=True) + (workspace / "dispatch" / "075-smith.prompt.md").write_text( + "Dispatch 75: **verify** the catalog entry `attention/compressor`.\n", encoding="utf-8" + ) + (workspace / "vocabulary.yaml").write_text( + yaml.safe_dump({"steps": STEPS, "queue": [{"n": 75, "tier": 1, "unit": "core"}]}), + encoding="utf-8", + ) + task = tmp_path / "task.yaml" + task.write_text(yaml.safe_dump(_history_spec(tmp_path)), encoding="utf-8") + (workspace / "task.yaml").write_text(task.read_text(), encoding="utf-8") + + assert smith_brief._main(["--workspace", str(workspace), "075"]) == 0 + out = capsys.readouterr().out + assert "dispatch 75, verify `attention/compressor`" in out + assert "- unit: core" in out and "3·G" in out + assert smith_brief._main(["--workspace", str(workspace), "999"]) == 2 + + +def test_every_smith_prompt_carries_the_brief(tmp_path, task_file, monkeypatch): + workflow = _workflow(tmp_path) + calls = _install(monkeypatch, workflow, Script(None)) + queue = [{"n": 1, "mode": "verify", "target": "flashinfer_rmsnorm", "surface": "T=3"}] + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", queue) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + + smith = next(prompt for role, prompt in calls if role == "smith") + assert "## Smith brief: dispatch 1, verify `flashinfer_rmsnorm`" in smith + assert "### Envelope (DERIVED from the gate tuples (legacy))" in smith + assert "Integer-width boundary probes (mandatory)" in smith + + +def test_the_verify_skill_and_the_prompts_carry_the_probes(): + from agent_flow.workflows.staircase_onboard.prompts import DEFAULT_PROMPTS + + skills = Path(smith_brief.__file__).parent / "skills" + verify = (skills / "staircase-kernel-verify" / "SKILL.md").read_text() + assert "## Integer-width boundary probes (mandatory)" in verify + assert "2^31 / row width rows" in verify and "2^24 + 1" in verify + probes = ( + skills / "staircase-kernel-verify" / "references" / "integer-width-probes.md" + ).read_text() + assert "## Hard limits" in probes + entry_spec = (skills / "staircase-kernel-onboard" / "references" / "entry-spec.md").read_text() + assert "## Hard limits" in entry_spec + assert "## The smith brief" in DEFAULT_PROMPTS.smith + assert "call_sites" in DEFAULT_PROMPTS.plan_drafter + assert "Captured rows come per call site" in DEFAULT_PROMPTS.auditor diff --git a/agent-flow/tests/workflows/staircase_onboard/test_speculation_support.py b/agent-flow/tests/workflows/staircase_onboard/test_speculation_support.py new file mode 100644 index 000000000000..e7f9e1a69048 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_speculation_support.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Preflight refuses a speculation algorithm the gates' launch path cannot build. + +The case it was written for: on main, ``quickstart_advanced.py`` has no DSPARK +branch, so ``--spec_decode_algo=DSPARK`` falls through to ``spec_config = +None`` and every quickstart gate of a ``dspark`` unit runs without +speculation -- found at the first boot or acceptance gate, as a unit that +"drafts nothing". Preflight now names the branch and the file instead. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import preflight + +from .test_workflow import Script, _install, _workflow + +REPO_ROOT = Path(__file__).resolve().parents[4] + +_QUICKSTART = """\ +def setup_llm(args): + spec_decode_algo = args.spec_decode_algo.upper() if args.spec_decode_algo else None + if spec_decode_algo == "MTP": + spec_config = 1 + elif spec_decode_algo == "EAGLE3": + spec_config = 2 +{extra} else: + spec_config = None + return spec_config +""" +_DSPARK_BRANCH = """ elif spec_decode_algo in ("DSPARK",): + spec_config = 3 +""" + +_LLM_ARGS = """\ +class MTPDecodingConfig: + decoding_type: Literal["MTP"] = Field(default="MTP") + + +class DSparkDecodingConfig: + decoding_type: Literal["DSpark"] = Field(default="DSpark") + + +class PrivateDecodingConfig: + decoding_type: Literal["Private"] = Field(default="Private") + + +SpeculativeConfig: TypeAlias = Annotated[ + Union[MTPDecodingConfig, DSparkDecodingConfig], + Field(discriminator="decoding_type"), +] +""" + + +def _repo(tmp_path: Path, *, dspark_branch: bool) -> Path: + repo = tmp_path / "repo" + quickstart = repo / preflight.QUICKSTART_REL + quickstart.parent.mkdir(parents=True, exist_ok=True) + quickstart.write_text( + _QUICKSTART.format(extra=_DSPARK_BRANCH if dspark_branch else ""), encoding="utf-8" + ) + llm_args = repo / preflight.LLM_ARGS_REL + llm_args.parent.mkdir(parents=True, exist_ok=True) + llm_args.write_text(_LLM_ARGS, encoding="utf-8") + return repo + + +def _task(repo: Path, features, options: Path | None = None) -> dict: + accuracy: dict = {} + if options is not None: + accuracy["llm_api_options"] = {"default": str(options), "core": str(options)} + return {"trtllm_repo_path": str(repo), "features": list(features), "accuracy": accuracy} + + +def _options(tmp_path: Path, decoding_type: str) -> Path: + path = tmp_path / "spec_llm_api.yaml" + path.write_text( + yaml.safe_dump({"speculative_config": {"decoding_type": decoding_type}}), encoding="utf-8" + ) + return path + + +def test_the_launchers_are_read_for_the_algorithms_they_handle(tmp_path): + repo = _repo(tmp_path, dspark_branch=True) + assert preflight.quickstart_spec_algorithms(repo / preflight.QUICKSTART_REL) == { + "MTP", + "EAGLE3", + "DSPARK", + } + # Only the union's members: a config class outside it is not accepted. + assert preflight.llm_api_decoding_types(repo / preflight.LLM_ARGS_REL) == {"MTP", "DSpark"} + assert preflight.quickstart_spec_algorithms(tmp_path / "absent.py") is None + + +def test_a_dspark_unit_without_a_quickstart_branch_is_a_named_error(tmp_path): + repo = _repo(tmp_path, dspark_branch=False) + with pytest.raises(preflight.UnsupportedSpeculationError) as excinfo: + preflight.check_speculation_support(_task(repo, ["dspark", "multimodal"])) + message = str(excinfo.value) + assert "--spec_decode_algo=DSPARK" in message + assert str(repo / preflight.QUICKSTART_REL) in message + assert "no `spec_decode_algo == 'DSPARK'` branch (it handles EAGLE3, MTP)" in message + assert "spec_config = None" in message + # Both units that inherit the speculator are named. + assert "unit(s) dspark, multimodal" in message + # Not an infrastructure failure: the workflow skips the troubleshooting hint. + assert isinstance(excinfo.value, preflight.SpecInconsistencyError) + + +def test_a_tree_with_the_branch_passes_and_says_so(tmp_path): + repo = _repo(tmp_path, dspark_branch=True) + notes = preflight.check_speculation_support( + _task(repo, ["dspark"], _options(tmp_path, "DSpark")) + ) + assert notes == ["speculation: DSPARK, DSpark supported by the gates' launch paths"] + + +def test_a_decoding_type_the_llm_api_does_not_accept_is_a_named_error(tmp_path): + repo = _repo(tmp_path, dspark_branch=True) + options = _options(tmp_path, "Private") + with pytest.raises(preflight.UnsupportedSpeculationError) as excinfo: + preflight.check_speculation_support(_task(repo, ["mtp"], options)) + message = str(excinfo.value) + assert "`speculative_config.decoding_type: Private`" in message + assert str(options) in message + assert ( + f"`SpeculativeConfig` in {repo / preflight.LLM_ARGS_REL} accepts only DSpark, MTP" + in message + ) + + +def test_no_speculator_reads_nothing_and_a_missing_launcher_is_only_a_note(tmp_path): + empty = tmp_path / "empty" + empty.mkdir() + assert preflight.check_speculation_support(_task(empty, [])) == [] + notes = preflight.check_speculation_support(_task(empty, ["dspark"])) + assert len(notes) == 1 and "cannot read" in notes[0] and "DSPARK" in notes[0] + + +def test_this_tree_handles_the_speculators_upstream_ships(): + """Read the real launchers, so a renamed variable fails here and not in a run.""" + algorithms = preflight.quickstart_spec_algorithms(REPO_ROOT / preflight.QUICKSTART_REL) + assert algorithms is not None and {"MTP", "EAGLE3"} <= algorithms + decoding = preflight.llm_api_decoding_types(REPO_ROOT / preflight.LLM_ARGS_REL) + assert decoding is not None and {"MTP", "Eagle3"} <= decoding + + +def test_the_workflow_stops_at_preflight_with_the_named_branch( + tmp_path, task_file, repo, monkeypatch +): + quickstart = repo / preflight.QUICKSTART_REL + quickstart.parent.mkdir(parents=True, exist_ok=True) + quickstart.write_text(_QUICKSTART.format(extra=""), encoding="utf-8") + spec = yaml.safe_load(task_file.read_text()) + spec["features"] = ["dspark"] + task_file.write_text(yaml.safe_dump(spec), encoding="utf-8") + + workflow = _workflow(tmp_path, dry_run=True) + calls = _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task_file) + + assert calls == [] + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "preflight failed" in journal + assert "no `spec_decode_algo == 'DSPARK'` branch" in journal + assert "infrastructure failure" not in journal From 5fba5c4d1db49a851dc316c853dc3ea0aed6662a Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 12 Oct 2026 00:56:02 -0700 Subject: [PATCH 03/70] [TRTLLM-16460][feat] staircase: adopt upstream agent-flow routing and prompt snapshots; no model-specific data in the workflow Adapts the workflow to upstream agent-flow: per-role backend/model routing through an agents block (#19522), the backend/model named on externally dispatched turns, composed role prompts snapshotted into the workspace (#19715), the oversized-system-prompt spill (#19766). Fixes concurrent file-dispatch turns sharing a request id. Moves the DeepSeek-V4.1 specifics out of skills, registries and the example spec into a worked example that is never materialized, and replaces tests that read a model PR's files with fixtures. Squashed from user/williamj/staircase-polish @ c8109d6f90 (13 commits): - [fix] staircase: concurrent dispatch turns never share a request id - [feat] staircase: per-role backend and model routing through the agents block - [feat] staircase: a dispatched turn names the backend and model its role is configured for - [feat] staircase: snapshot the composed role prompts into the workspace - [test] staircase: a role's oversized system prompt reaches the CLI through the spill - [test] staircase: check the quickstart DSPARK contract against a fixture - [fix] staircase: claim tree-local Triton custom-op files by content, not by name - [chore] staircase: the example spec and the schema's messages name no checkpoint - [feat] staircase: move DeepSeek-V4.1 specifics out of the feature skills into a worked example - [fix] staircase: delivery accepts the op-matrix test a multi-rank catalog entry has on main - [test] staircase: the same-process retry test answers only once the retry holds the request - [doc] staircase: README sections for the envelope, the evidence ledger and commit hygiene - [doc] staircase: README sections for the gates and the execution substrate; no placeholder left Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- agent-flow/agent_flow/runtime/turns/base.py | 6 + .../agent_flow/runtime/turns/file_dispatch.py | 26 ++- .../workflows/staircase_onboard/cli.py | 5 +- .../workflows/staircase_onboard/delivery.py | 15 +- .../examples/deepseek-v4.1.md | 188 ++++++++++++++++++ .../workflows/staircase_onboard/gates.py | 2 +- .../staircase_onboard/kernel_sources.yaml | 21 +- .../staircase_onboard/prompts/__init__.py | 5 + .../workflows/staircase_onboard/roles.py | 103 ++++++++++ .../workflows/staircase_onboard/runtime.py | 34 +++- .../staircase-kernel-feasibility/SKILL.md | 8 +- .../SKILL.md | 126 ++++++------ .../staircase-modeling-host-offload/SKILL.md | 25 +-- .../staircase-modeling-kv-quant/SKILL.md | 93 ++++----- .../staircase-modeling-speculator/SKILL.md | 4 +- .../references/algorithms.md | 56 +++--- .../staircase_onboard/task.example.yaml | 57 ++++-- .../staircase_onboard/task_schema.py | 16 +- .../workflows/staircase_onboard/vocabulary.py | 7 +- .../workflows/staircase_onboard/workflow.py | 39 ++-- .../tests/runtime/test_turn_executors.py | 91 +++++++++ .../staircase_onboard/test_delivery.py | 41 +++- .../staircase_onboard/test_dspark.py | 85 +++++++- .../staircase_onboard/test_kernel_sources.py | 63 +++++- .../staircase_onboard/test_package_hygiene.py | 47 +++++ .../staircase_onboard/test_prompt_snapshot.py | 50 +++++ .../workflows/staircase_onboard/test_roles.py | 174 ++++++++++++++++ .../staircase_onboard/test_substrate_docs.py | 58 ++++++ .../test_system_prompt_spill.py | 77 +++++++ .../staircase_onboard/test_worked_examples.py | 74 +++++++ 30 files changed, 1358 insertions(+), 238 deletions(-) create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/examples/deepseek-v4.1.md create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/roles.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_prompt_snapshot.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_roles.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_substrate_docs.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_system_prompt_spill.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_worked_examples.py diff --git a/agent-flow/agent_flow/runtime/turns/base.py b/agent-flow/agent_flow/runtime/turns/base.py index 220bc383fddf..773916fd3621 100644 --- a/agent-flow/agent_flow/runtime/turns/base.py +++ b/agent-flow/agent_flow/runtime/turns/base.py @@ -60,6 +60,12 @@ class TurnRequest: #: Paths the role should read as input. Named rather than inlined, so a #: transport that cannot afford the prompt budget can hand over pointers. context_refs: tuple[str, ...] = () + #: The backend and model the caller configured for this role. An + #: in-process transport already built its session from them; a transport + #: that hands the turn to someone else passes them on, because otherwise + #: nothing tells that driver which model the role was meant to run on. + backend: str | None = None + model: str | None = None @dataclass(frozen=True) diff --git a/agent-flow/agent_flow/runtime/turns/file_dispatch.py b/agent-flow/agent_flow/runtime/turns/file_dispatch.py index 16afc9e815e4..ceede3473a59 100644 --- a/agent-flow/agent_flow/runtime/turns/file_dispatch.py +++ b/agent-flow/agent_flow/runtime/turns/file_dispatch.py @@ -109,6 +109,11 @@ def __init__( self._note = transport_note self._counter = itertools.count(self._next_seq()) self._lock = threading.Lock() + #: Request ids a caller of this executor is waiting on right now. Such + #: a request is never a twin: its caller will read the answer, and the + #: next caller's prompt may differ even when the role, ``what`` and + #: decision file match. Guarded by ``_lock``. + self._in_flight: set[str] = set() # ------------------------------------------------------------ helpers @@ -210,10 +215,17 @@ def _unanswered_twin(self, request: TurnRequest) -> str | None: writing a second request would duplicate the work -- or, worse, give two agents the same decision path. Waiting on the first one is the same answer to the same question. + + A request some caller of this executor is still waiting on is not a + twin. Without that rule, whether two concurrent callers asking alike + shared one id depended on whether the first caller's request file was + on disk yet when the second one looked, so they sometimes received one + answer between them. An orphan is likewise adopted by one caller at a + time. Call with ``_lock`` held. """ for path in sorted(self.root.glob(f"*-{request.role}.request.json")): name = path.name[: -len(".request.json")] - if (self.root / f"{name}.done").exists(): + if name in self._in_flight or (self.root / f"{name}.done").exists(): continue try: data = json.loads(path.read_text(encoding="utf-8")) @@ -229,7 +241,14 @@ def run(self, request: TurnRequest) -> TurnOutcome: with self._lock: adopted = self._unanswered_twin(request) request_id = adopted or f"{next(self._counter):03d}-{request.role}" + self._in_flight.add(request_id) + try: + return self._await(request, request_id, adopted) + finally: + with self._lock: + self._in_flight.discard(request_id) + def _await(self, request: TurnRequest, request_id: str, adopted: str | None) -> TurnOutcome: started = time.monotonic() # The role's system prompt travels WITH the turn. In-process it is # set on the session and is easy to forget here -- and a driver that @@ -285,6 +304,11 @@ def _write_request(self, request, request_id, system, prompt, result, done) -> N "done_file": str(done), "decision_file": request.decision_file, "context_refs": list(request.context_refs), + # Which backend and model the role is configured for. A + # driver that runs every turn on whatever model it has + # would otherwise turn a cross-model review into a + # same-model one without anything on disk saying so. + "agent": {"backend": request.backend, "model": request.model}, # Named rather than inlined: a driver that wants the # role's own history reads these, and one that does not # is not forced to carry them in the prompt. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/cli.py b/agent-flow/agent_flow/workflows/staircase_onboard/cli.py index af97b1ac1ae2..60acd4243299 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/cli.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/cli.py @@ -10,6 +10,7 @@ from pathlib import Path from .agent_home import ensure_agent_home +from .roles import shares_the_drafters_model from .state import STATE_FILENAME from .task_schema import TaskSchemaError, load_and_validate_task_yaml from .transport import DEFAULT_TURN_TIMEOUT_S, ExternalAgentTransport @@ -234,10 +235,12 @@ def main(argv: list[str] | None = None) -> None: # agent is constructed, and the message must list every problem at # once so the user fixes them in a single pass. try: - load_and_validate_task_yaml(args.task) + task_data = load_and_validate_task_yaml(args.task) except TaskSchemaError as exc: print(f"error: {exc}", file=sys.stderr) sys.exit(2) + if warning := shares_the_drafters_model(task_data): + print(f"warning: {warning}", file=sys.stderr) # A sandboxed session's $HOME is read-only, and the first thing a spawned # `claude` CLI does is take a lock beside its credential store. Without diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/delivery.py b/agent-flow/agent_flow/workflows/staircase_onboard/delivery.py index adb977543b05..e92aefc40819 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/delivery.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/delivery.py @@ -173,18 +173,25 @@ def check_catalog_tests(repo_root: Path, entries: Iterable[tuple[str, str]]) -> ``test_modeling_v2_*`` precisely so it cannot collide with upstream's test of the same op, and so the CI selector's pattern matches it exactly. A correctly written test under any other name is a test on no list. + + A multi-rank entry's test is an op matrix over world sizes and is named + ``test_modeling_v2__op_matrix.py``, as the in-tree ``comm`` tests + are; either name satisfies the check. """ checks: list[Check] = [] for category, name in entries: if category == MIRROR_CATEGORY: continue - rel = f"{CATALOG_TESTS_DIR}/{category}/test_modeling_v2_{name}.py" - path = repo_root / rel + candidates = [ + f"{CATALOG_TESTS_DIR}/{category}/test_modeling_v2_{name}{suffix}.py" + for suffix in ("", "_op_matrix") + ] + found = next((rel for rel in candidates if (repo_root / rel).is_file()), None) checks.append( Check( name=f"test:{category}/{name}", - passed=path.is_file(), - detail=str(path) if path.is_file() else f"missing: {rel}", + passed=found is not None, + detail=str(repo_root / found) if found else f"missing: {candidates[0]}", ) ) return checks diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/examples/deepseek-v4.1.md b/agent-flow/agent_flow/workflows/staircase_onboard/examples/deepseek-v4.1.md new file mode 100644 index 000000000000..3b3fbc4e42e9 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/examples/deepseek-v4.1.md @@ -0,0 +1,188 @@ + + +# Worked example: DeepSeek-V4.1-Flash NVFP4 + +What one checkpoint required of the workflow's generic features, kept here as +a worked example. The skills describe each feature without naming a model; +this file shows one model's values filled in. + +**This file is for operators writing a `task.yaml`.** The run does not +materialize `examples/` into its workspace, so no role is handed it. A run +that onboards this checkpoint, or a sibling of it, to measure the workflow +should add this file to `source_embargo`, because the facts below are what +the run is meant to derive. + +Sources: the DeepSeek-V4.1-Flash tech report, the checkpoint's own reference +code, and SGLang's implementation (branch `origin/dsv4.1`). + +## `dspark` (staircase-modeling-speculator) + +- The tech report (§2.4.3): "We omit the MTP module during backbone + pre-training and use DSpark for speculative decoding." The drafter ships + under `mtp.0..2` with `num_nextn_predict_layers: 3`, which is what makes + `mtp` the easy wrong choice. The `dspark_*` keys sit under `text_config`. +- Block size 5, `dspark_noise_token_id` 128799. Each of the three stages is a + full block: sliding-window attention plus a 128-expert top-3 MoE, with + hyper-connections. +- Each stage's ring has window 128 and is filled from the *input* of target + layers 37, 38 and 39 (`dspark_target_layer_ids`), averaged over the + hyper-connection copies. vLLM's DeepSeek-V4 path captures layer outputs, + which is a different checkpoint's contract. +- The Markov head is a rank-256 bigram bias. +- Acceptance ceiling 6 (`block_size + 1`). vLLM's DeepSeek-V4 DSpark recipe + uses 1.2 as its smoke floor. + +```yaml +features: [dspark] +accuracy: + llm_api_options: + dspark: scripts/dspark_llm_api.yaml # decoding_type: DSpark, max_draft_len: 5 + acceptance: + dspark: + min_absolute: 1.2 +``` + +## `kv_quant` (staircase-modeling-kv-quant) + +The reference fake-quantizes KV to these formats (tech report §2.4.4): + +| pool | who owns one | format | bytes per entry | +|---|---|---|---| +| SWA KV | every layer (40), the window only | FP8 E4M3 per 32 channels with a UE8M0 scale, over all 512 dims including RoPE | 512 + 16 = **528** | +| main (compressed global) KV | the KV-source layers only: 2, 8, 14 (compress ratio 2) and 20 (ratio 1) | FP4 E2M1 with an E4M3 scale per 16, **no global scale** | 256 + 32 = **288** | +| indexer K | the same KV-source layers | MXFP4: E2M1 with a UE8M0 scale per 32, 128 dims | 64 + 4 = **68** | + +Only layers 2, 8, 14 and 20 own compressed and indexer pools; the others read +a source layer's pool through `kv_source_layer_ids`. A pool allocated per +compressed layer (38 of them) comes to 10324 B/token, which is the sign the +sharing was never wired. + +``` +3 x (288 + 68) / 2 + (288 + 68) = 534 + 356 = 890 B/token + ratio-2 layers ratio-1 layer +``` + +The same formula gives 3200 B/token for bf16 containers of these values (the +core unit) and 1490 B/token with FP8 main KV. Padding the 4-byte indexer scale +to 16 moves 890 to 920. The SWA window is per request: 40 x 128 x 528 B, +about 2.7 MB. + +How this differs from the DeepSeek-V4 layouts in +`tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/`: V4's +`fp8_ds_mla` is 448 FP8 + 64 bf16 RoPE dims with block-64 scales, where V4.1 +quantizes all 512 dims, RoPE included, after RoPE. V4's `nvfp4` pool has a +second-level global scale and a 64-dim residual, where V4.1 has no global +scale and all 512 dims in FP4. + +The sm_103 kernels that serve this model (trtllm-gen's H512 +DynamicTokenSparse path through FlashInfer's +`trtllm_batch_decode_sparse_mla_dsv4`, and the catalog's `thop_attention`) are +exact only over bf16 KV, which is why core stores bf16 containers. SGLang +reports that this model's output is not bitwise stable across batch +composition, so compare at batch 1. + +```yaml +features: [kv_quant] +accuracy: + witness: + kv_quant: + contrast_env: {TRTLLM_DSV41_KV_STORAGE: bf16} # the leg WITH bf16 storage + expected_bytes_per_token: 890 + exact: true + max_batch_size: 1 +``` + +## `decoder_replay` / `encoder_replay` (staircase-modeling-bounded-replay) + +The model is a causal encoder-decoder: layers 0-19 are the "encoder", layers +20-39 the "decoder", and decoder global KV is projected from the encoder's +final hidden state. Sliding window 128. The last entry of +`kv_source_layer_ids` is 20; the tech report starts the decoder there, while +SGLang computes `late_layer_start = max(kv_source_layer_ids) + 1` and runs +layer 20 over every prompt token, because its compressor produces the +decoder's global KV. + +The tech report (§3.2.2, "SWA Bounded Replay"): + +> "SWA Bounded Replay instead replays only the most recent n_win tokens and +> truncates SWA to the replay segment, accepting approximate states: for a +> replay starting at position r, a query at position i attends to SWA keys in +> [max(r, i−n_win+1), i]." + +> "**Encoder SWA Bounded Replay** makes prefix caching depend only on global +> KV, allowing SWA KV to be removed from the persistent KV cache. When the +> encoder SWA KV is missing, we replay the last n_win tokens of the cached +> prefix and process them together with the uncached suffix. The replayed +> tokens regenerate only SWA KV, reusing the cached global KV without +> recomputation or overwriting, while the uncached suffix generates both +> global KV and SWA KV. By design, the replayed prefix state is approximate…" + +> "**Decoder SWA Bounded Replay** bounds the decoder forward pass to n_win +> tokens, nearly halving total prefill computation. Under CED, decoder global +> KV is projected from the final encoder hidden states. The only obstacle to +> ending prefill at the encoder is decoder SWA KV … at every prefill, we +> replay the last n_win tokens of the prompt, feed their encoder outputs +> through the decoder layers under the same SWA truncation, and use the +> resulting decoder SWA KV only for decoding, not for prefix caching. … not +> mathematically equivalent … we additionally simulate the same replay during +> post-training for train-aware adaptation." + +SGLang's reference mechanics are in +`python/sglang/srt/model_executor/encoder_swa_replay.py`. Its argument checks +(`python/sglang/srt/arg_groups/deepseek_v4_hook.py`) refuse `encoder_replay` +with prefill CUDA graphs, attention DP, PD disaggregation and LoRA, and +`decoder_replay` with prefill CUDA graphs and attention DP. Compress ratio 2 +on layers 2, 8 and 14 makes the reuse boundary even. + +```yaml +features: [decoder_replay, encoder_replay] +accuracy: + witness: + decoder_replay: + contrast_env: {TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY: "0"} + window: 128 + prompt_tokens: 4096 + encoder_replay: + contrast_env: {TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY: "0"} + window: 128 + prefix_tokens: 4096 + reuse_alignment: 2 +``` + +## `host_offload` (staircase-modeling-host-offload) + +The Engram table: two modules of 384M rows x 256 FP8 plus E8M0 scales, +**202.76 GB**, larger than the rest of the checkpoint. Per-rank copies on a +four-GPU node would be 811 GB of host memory; the shared single-copy layout +holds it once. SGLang ships that layout behind +`SGLANG_DSV41_ENGRAM_HOST_TABLE_LAYOUT=shared` (`python/sglang/srt/layers/engram.py`). +Row-sharded at TP4, the residency witness's `expected_mib` is +202.76 GB / 4 = **48,340 MiB** per rank. + +## `disagg` + +Both bounded replays are refused under disaggregation, so the disagg +equivalence witness pins them off on both legs: + +```yaml +accuracy: + witness: + disagg: + contrast_env: {TRTLLM_DISAGG: "1"} + contrast_command: scripts/disagg_leg_b.py + max_batch_size: 1 + legs_env: + TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY: "0" + TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY: "0" +``` + +## The smith queue + +The run queued 33 dispatches, six of them needed only by a feature unit: the +Engram host-gather ports, the ViT attention verify, a DSpark-only routing +surface and the MXFP4 drafter MoE. Untagged, any one of them failing would +have closed the run before core started, which is why queue rows carry +`unit:`. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/gates.py b/agent-flow/agent_flow/workflows/staircase_onboard/gates.py index 133851f45aa2..3c0a7c93a587 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/gates.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/gates.py @@ -645,7 +645,7 @@ def spool_runner( """Return a :data:`ShellRunner` that executes through a compute session. The session is a long-lived allocation whose own main process watches - ``/pending`` (see ``scripts/session_watcher.sh``). Work is a + ``/pending`` (see ``spool_watcher.sh``). Work is a file, which is the property that matters: **anything able to write a file can run GPU work**, including a plain Python process on a host with no GPU and no reachable scheduler. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/kernel_sources.yaml b/agent-flow/agent_flow/workflows/staircase_onboard/kernel_sources.yaml index 09dbad885f94..22d3c722b93f 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/kernel_sources.yaml +++ b/agent-flow/agent_flow/workflows/staircase_onboard/kernel_sources.yaml @@ -33,6 +33,9 @@ # summary: what this source is, in a sentence or two # registers: the file(s) where its ops enter `torch.ops.trtllm` # implementation: where the kernel body lives, when that differs +# claims_files_containing: markers; a `custom_ops/*.py` file no +# `registers` path names, but which contains all of them, +# belongs to this source # tier: which sourcing tier (see vocabulary.py) it satisfies # availability: what decides whether its ops exist in a given build # certifying: what a receipt for one of these must say, beyond passing @@ -149,20 +152,20 @@ triton_custom_op: them. registers: - tensorrt_llm/_torch/custom_ops/nvfp4_sfc_finalize.py - - tensorrt_llm/_torch/custom_ops/nvfp4_fake_quantize.py - - tensorrt_llm/_torch/custom_ops/mxfp8_fake_quantize.py - - tensorrt_llm/_torch/custom_ops/mxfp8_dequantize.py - - tensorrt_llm/_torch/custom_ops/mxfp8_to_e4m3.py - - tensorrt_llm/_torch/custom_ops/nvfp4_kv_to_bf16.py - - tensorrt_llm/_torch/custom_ops/nvfp4_kv_to_e4m3.py - - tensorrt_llm/_torch/custom_ops/e4m3_quantize_pow2_dynamic.py + # These files follow no naming convention, and a model's onboarding adds + # them one op at a time, so they are claimed by what they contain: a Triton + # kernel and the custom-op registration that exposes it. + claims_files_containing: + - "@triton.jit" + - "torch.library.custom_op" tier: 2 availability: >- Never gated --- Triton is a hard dependency of the wheel, and each module imports it at module scope rather than behind a try. But an op exists only once its module is imported: `nvfp4_sfc_finalize` is imported by the - `custom_ops` package, the others by the catalog wrapper that calls them, - so an inventory that imports only the package does not list them. + `custom_ops` package, while a module the package does not import is + imported by the catalog wrapper that calls it, so an inventory that + imports only the package does not list its ops. certifying: >- Read the autotune. A Triton op's numerics are fixed but its launch configuration is chosen at runtime, so a receipt that names only the op diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/__init__.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/__init__.py index c56575826678..2550d73cfa01 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/__init__.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/__init__.py @@ -136,6 +136,10 @@ def _for_every_role(prompts: PromptBundle, block: str) -> PromptBundle: ) +#: Where the orchestrator snapshots the composed prompts inside the workspace. +PROMPTS_DIRNAME = "prompts" + + def build_staircase_onboard_prompts( *, environment: "RunEnvironment | None" = None, @@ -193,6 +197,7 @@ def build_staircase_onboard_prompts( "DEFAULT_PROMPTS", "PLAN_DRAFTER_SYSTEM_PROMPT", "PLAN_REVIEWER_SYSTEM_PROMPT", + "PROMPTS_DIRNAME", "PromptBundle", "REVIEWER_FULL_EXTENSION", "REVIEWER_SYSTEM_PROMPT", diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/roles.py b/agent-flow/agent_flow/workflows/staircase_onboard/roles.py new file mode 100644 index 000000000000..62f783008100 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/roles.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The staircase roles, and which backend and model each one runs on. + +The selection is the shared per-role routing in :mod:`agent_flow.agent_runtime`, +which reads an optional ``agents`` block in ``task.yaml``. This module adds one +staircase-specific default: the ``plan_reviewer`` runs on Codex. The other +roles run on Claude Code, because they need its skill system. + +The plan reviewer runs on a different model from the plan drafter on purpose. +A reviewer that shares the author's blind spots reviews its own reasoning back +at itself, and the plan phase has no gate, only this review. +""" + +from __future__ import annotations + +from typing import Any, Mapping + +from agent_flow.agent_runtime import AGENTS_FIELD, AgentConfig, resolve_agent_config +from agent_flow.config import CLAUDE_CODE_DEFAULT_MODEL, CODEX_DEFAULT_MODEL, BackendKind + +#: Every role a turn can be dispatched to. ``runner`` makes no judgement: it +#: relays a command to a scheduler only an agent can reach. It is listed +#: because it is still a turn on some backend. +ROLES: tuple[str, ...] = ( + "plan_drafter", + "plan_reviewer", + "smith", + "assembler", + "reviewer", + "auditor", + "runner", +) + +#: The backend a role runs on when the ``agents`` block does not choose one. +DEFAULT_BACKENDS: Mapping[str, BackendKind] = {role: "claude-code" for role in ROLES} | { + "plan_reviewer": "codex" +} + + +def _default_model(backend: BackendKind) -> str: + return CODEX_DEFAULT_MODEL if backend == "codex" else CLAUDE_CODE_DEFAULT_MODEL + + +def resolve_role_agent(task_data: Mapping[str, Any], role: str) -> AgentConfig: + """Resolve one role's backend, model, reasoning effort and extra MCP servers.""" + backend = DEFAULT_BACKENDS[role] + return resolve_agent_config( + task_data, role, default_backend=backend, default_model=_default_model(backend) + ) + + +def validate_role_defaults(data: Mapping[str, Any]) -> list[str]: + """Reject an ``agents.defaults.model`` that would reach two backends. + + The shared resolver applies ``defaults.model`` to every role whose backend + matches ``defaults.backend``, and when ``defaults.backend`` is absent it + compares against each role's own default. The roles here do not share a + default backend, so a bare ``defaults.model`` would send a Claude model + name to the Codex plan reviewer, or the reverse. + """ + agents = data.get(AGENTS_FIELD) + if not isinstance(agents, Mapping): + return [] + defaults = agents.get("defaults") + if isinstance(defaults, Mapping) and "model" in defaults and "backend" not in defaults: + return [ + "'agents.defaults.model' needs 'agents.defaults.backend' here: the " + "staircase roles do not share a default backend (plan_reviewer runs on " + "codex, the others on claude-code), so a model name alone would be " + "sent to both. Set 'agents.defaults.backend', or set the model per role " + "under 'agents.roles'." + ] + return [] + + +def shares_the_drafters_model(task_data: Mapping[str, Any]) -> str: + """A warning when the plan reviewer resolves to the plan drafter's model, else ``""``. + + Legal, because a host may not reach a second provider. But it is a weaker + review, so the run says so instead of recording a cross-model review that + did not happen. + """ + drafter = resolve_role_agent(task_data, "plan_drafter") + reviewer = resolve_role_agent(task_data, "plan_reviewer") + if (drafter.backend, drafter.model) != (reviewer.backend, reviewer.model): + return "" + return ( + f"plan_reviewer and plan_drafter both resolve to {reviewer.backend}/{reviewer.model}. " + "The plan review is then not cross-model: the reviewer shares the drafter's blind " + "spots. Route plan_reviewer to another backend or model under 'agents.roles' " + "if this host can reach one." + ) + + +__all__ = [ + "DEFAULT_BACKENDS", + "ROLES", + "resolve_role_agent", + "shares_the_drafters_model", + "validate_role_defaults", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/runtime.py b/agent-flow/agent_flow/workflows/staircase_onboard/runtime.py index a18fb50e718b..df9097a6bebc 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/runtime.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/runtime.py @@ -53,7 +53,7 @@ KERNEL_SOURCES_PATH = Path(__file__).parent / "kernel_sources.yaml" KERNEL_SOURCE_REQUIRED: tuple[str, ...] = ("summary", "tier", "availability", "certifying") -KERNEL_SOURCE_OPTIONAL: tuple[str, ...] = ("registers", "implementation") +KERNEL_SOURCE_OPTIONAL: tuple[str, ...] = ("registers", "implementation", "claims_files_containing") #: Areas every run needs regardless of what the checkpoint declares. A target #: always sits somewhere in the parallel topology and is always captured or @@ -215,11 +215,21 @@ class KernelSource: certifying: str registers: tuple[str, ...] = () implementation: tuple[str, ...] = () + #: Markers that claim a `custom_ops/*.py` file `registers` does not name: + #: a file containing every one of them belongs to this source. For a + #: source whose files carry no naming convention, so a tree that adds one + #: does not need this registry edited to classify it. + claims_files_containing: tuple[str, ...] = () @property def pointers(self) -> tuple[str, ...]: return (*self.registers, *self.implementation) + def claims(self, text: str) -> bool: + """Whether a registration file's source text belongs to this source by content.""" + markers = self.claims_files_containing + return bool(markers) and all(marker in text for marker in markers) + def render(self) -> str: lines = [f"### {self.name} (tier {self.tier})", "", " ".join(self.summary.split()), ""] if self.registers: @@ -228,6 +238,9 @@ def render(self) -> str: lines.append( f"- **implemented in**: {', '.join(f'`{p}`' for p in self.implementation)}" ) + if self.claims_files_containing: + markers = " and ".join(f"`{m}`" for m in self.claims_files_containing) + lines.append(f"- **also any `custom_ops/*.py` containing**: {markers}") lines.append(f"- **availability**: {' '.join(self.availability.split())}") lines.append(f"- **certifying one**: {' '.join(self.certifying.split())}") return "\n".join(lines) @@ -253,6 +266,10 @@ def load_kernel_sources(path: str | Path = KERNEL_SOURCES_PATH) -> dict[str, Ker errors.append(f"{name}: missing required field '{required}'") if not entry.get("registers") and not entry.get("implementation"): errors.append(f"{name}: needs at least one of 'registers' or 'implementation'") + markers = entry.get("claims_files_containing", []) + if not isinstance(markers, list) or not all(isinstance(m, str) and m for m in markers): + errors.append(f"{name}: 'claims_files_containing' must be a list of non-empty strings") + markers = [] sources[str(name)] = KernelSource( name=str(name), summary=str(entry.get("summary", "")), @@ -261,6 +278,7 @@ def load_kernel_sources(path: str | Path = KERNEL_SOURCES_PATH) -> dict[str, Ker certifying=str(entry.get("certifying", "")), registers=tuple(str(p) for p in entry.get("registers", ())), implementation=tuple(str(p) for p in entry.get("implementation", ())), + claims_files_containing=tuple(markers), ) if errors: raise RuntimeIndexError(f"{path} is malformed:\n - " + "\n - ".join(errors)) @@ -268,7 +286,10 @@ def load_kernel_sources(path: str | Path = KERNEL_SOURCES_PATH) -> dict[str, Ker def source_for_stem( - stem: str, sources: "Mapping[str, KernelSource] | None" = None + stem: str, + sources: "Mapping[str, KernelSource] | None" = None, + *, + text: str | None = None, ) -> KernelSource | None: """Map a `custom_ops/.py` file back to the source it represents. @@ -279,13 +300,20 @@ def source_for_stem( ``cpp_thop`` is returned for the C++ plane, which the probe reports for any op no Python file claims --- a conclusion, not a fallback. + + ``text`` is the file's source. Given, a file no path names can still be + claimed by content (``claims_files_containing``); a named path always wins. """ sources = load_kernel_sources() if sources is None else sources for source in sources.values(): for path in source.registers: if Path(path).stem == stem or path.rstrip("/").endswith(stem): return source - return sources.get(stem) + if stem in sources: + return sources[stem] + if text is not None: + return next((source for source in sources.values() if source.claims(text)), None) + return None def format_kernel_sources(sources: "Mapping[str, KernelSource]") -> str: diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/SKILL.md index 3b8957965d8a..7d3b0e2a07e2 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/SKILL.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-feasibility/SKILL.md @@ -170,10 +170,10 @@ wall-clock. The tag changes only what a **failure** blocks. An untagged row reads as `core`, and a `core` row that fails stops the run before assembly. A tagged row that fails is deferred instead, and retried once before its unit's first assembler turn; the run stops only if that retry fails too, and then at -that unit rather than before core. Measured on DeepSeek-V4.1: a 33-dispatch -queue with six feature-only rows — among them the Engram host-gather ports, -the ViT attention verify, a DSpark-only routing surface and the MXFP4 drafter -MoE — any one of which could close the run before core started. +that unit rather than before core. Measured on one run: a 33-dispatch queue +with six feature-only rows — host-offload gathers, a vision-tower verify and +speculator-only surfaces among them — any one of which could close the run +before core started. - Tag the **first** unit that needs the entry, in `task.yaml` feature order. - Leave the tag off when core uses the entry at all, or when you are unsure. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-bounded-replay/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-bounded-replay/SKILL.md index a58e85aa332b..a03b237b11f2 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-bounded-replay/SKILL.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-bounded-replay/SKILL.md @@ -1,46 +1,44 @@ --- name: staircase-modeling-bounded-replay -description: Load when adding DeepSeek-V4.1's SWA Bounded Replay to a target — `decoder_replay` (the layers after the last KV-source layer run over only the last `window` prompt tokens) or `encoder_replay` (a prefix-cache hit rebuilds sliding-window KV by re-forwarding the last `window` cached tokens). Deliberately not output-identical, and gated by a work witness, because a unit that implements nothing passes accuracy too. +description: Load when adding sliding-window (SWA) bounded replay to a target — `decoder_replay` (the layers after the last KV-source layer run over only the last `window` prompt tokens) or `encoder_replay` (a prefix-cache hit rebuilds sliding-window KV by re-forwarding the last `window` cached tokens). Deliberately not output-identical, and gated by a work witness, because a unit that implements nothing passes accuracy too. --- # staircase-modeling-bounded-replay -`--features decoder_replay` and `--features encoder_replay`. Both come from -DeepSeek-V4.1-Flash's tech report (§3.2.2, "SWA Bounded Replay"), where the -model is a **Causal Encoder-Decoder**: layers 0-19 are the "encoder", layers -20-39 the "decoder", and decoder global KV is projected from the encoder's -final hidden state rather than from each decoder layer's own. "Encoder" here -is those causal layers — **never the vision tower** — and "replay" is neither -CUDA-graph replay nor recurrent-state replay. - -## The definitions, verbatim - -TR §3.2.2: - -> "SWA Bounded Replay instead replays only the most recent n_win tokens and -> truncates SWA to the replay segment, accepting approximate states: for a -> replay starting at position r, a query at position i attends to SWA keys in -> [max(r, i−n_win+1), i]." - -> "**Encoder SWA Bounded Replay** makes prefix caching depend only on global -> KV, allowing SWA KV to be removed from the persistent KV cache. When the -> encoder SWA KV is missing, we replay the last n_win tokens of the cached -> prefix and process them together with the uncached suffix. The replayed -> tokens regenerate only SWA KV, reusing the cached global KV without -> recomputation or overwriting, while the uncached suffix generates both -> global KV and SWA KV. By design, the replayed prefix state is approximate…" - -> "**Decoder SWA Bounded Replay** bounds the decoder forward pass to n_win -> tokens, nearly halving total prefill computation. Under CED, decoder global -> KV is projected from the final encoder hidden states. The only obstacle to -> ending prefill at the encoder is decoder SWA KV … at every prefill, we -> replay the last n_win tokens of the prompt, feed their encoder outputs -> through the decoder layers under the same SWA truncation, and use the -> resulting decoder SWA KV only for decoding, not for prefix caching. … not -> mathematically equivalent … we additionally simulate the same replay during -> post-training for train-aware adaptation." - -`n_win` is the checkpoint's sliding window: 128 for DeepSeek-V4.1. +`--features decoder_replay` and `--features encoder_replay`. Both belong to a +model that keeps two kinds of KV, global (often compressed) KV and +sliding-window (SWA) KV, and that was post-trained to tolerate truncated SWA +state. Such a model may be a **causal encoder-decoder**: the first block of +layers is the "encoder", the rest the "decoder", and decoder global KV is +projected from the encoder's final hidden state rather than from each decoder +layer's own. "Encoder" here is those causal layers — **never the vision +tower** — and "replay" is neither CUDA-graph replay nor recurrent-state +replay. The checkpoint's tech report says whether it was trained for either +feature; quote it in the plan. + +## The definitions + +**Bounded replay** replays only the most recent `n_win` tokens and truncates +SWA to the replay segment, accepting approximate states: for a replay starting +at position `r`, a query at position `i` attends to SWA keys in +`[max(r, i - n_win + 1), i]`. + +**Encoder bounded replay** makes prefix caching depend only on global KV, so +SWA KV can leave the persistent cache. On a prefix hit the last `n_win` cached +tokens are replayed together with the uncached suffix. The replayed tokens +regenerate only SWA KV, reusing the cached global KV without recomputing or +overwriting it; the uncached suffix generates both. The replayed prefix state +is approximate by design. + +**Decoder bounded replay** bounds the decoder's forward pass to `n_win` +tokens. Decoder global KV comes from the encoder's output, so the only thing +standing between prefill and stopping at the encoder is decoder SWA KV: at +every prefill the last `n_win` prompt tokens' encoder outputs run through the +decoder layers under the same SWA truncation, and the decoder SWA KV that +produces serves decoding only, never prefix caching. Not mathematically +equivalent to the full computation; the training simulated it. + +`n_win` is the checkpoint's sliding window, read from its config. ## The property that shapes the gates @@ -61,14 +59,13 @@ declared before assembly. See **Gates**. Target-local: no runtime component, no scheduler change. -- **Which layers are bounded.** The ones after the last KV-source layer - (`kv_source_layer_ids` in the config; its last entry is 20 on - DeepSeek-V4.1). The boundary is a decision, not a fact: the tech report's - decoder starts at layer 20, while SGLang's implementation computes - `late_layer_start = max(kv_source_layer_ids) + 1` and runs layer 20 over - every prompt token, because layer 20's compressor is what produces the - decoder's global KV. Pick one, say why, and record it — the plan's open - risks carry it. +- **Which layers are bounded.** The decoder's: for example the layers after + the last KV-source layer (the last entry of a `kv_source_layer_ids` config + key). The boundary is a decision, not a fact. A tech report may start the + decoder *at* the last KV-source layer, while an implementation may start it + one layer later and run that layer over every prompt token, because its + compressor is what produces the decoder's global KV. Pick one, say why, and + record it — the plan's open risks carry it. - **Context phase.** Every layer up to the boundary — including the last KV-source layer's compressor and index-K projection — runs over all prompt tokens. The bounded layers run over only the last `min(N, window)` prompt @@ -84,15 +81,15 @@ Target-local: no runtime component, no scheduler change. accuracy. - **What it cannot return.** The bounded layers never computed rows before the tail, so prompt logprobs, context logits and full hidden-state capture do - not exist under this feature. SGLang refuses them; the target must fail fast - rather than return rows it did not compute. + not exist under this feature. The target must fail fast rather than return + rows it did not compute. - The decoder SWA KV the replay produces serves decoding only — it is not something prefix caching may reuse. -SGLang's help text for its flag states the contract compactly: "after the last -kv_source layer, run the remaining layers over only the last window_size -tokens of a prefill. Main and indexer KV stay exact; nothing is replayed. -Deterministic for a fixed prompt and chunk size." +The contract, compactly: after the last KV-source layer, run the remaining +layers over only the last `window` tokens of a prefill. Main and indexer KV +stay exact; nothing is replayed. Deterministic for a fixed prompt and chunk +size. ## encoder_replay — runtime and target, and the runtime half is new @@ -107,8 +104,8 @@ Deterministic for a fixed prompt and chunk size." An odd boundary on a ratio-2 layer splits a compressed pair: one half is cached and the other was never computed, and the result is a quietly wrong -global entry rather than an error. SGLang refuses the combination; the witness -checks it when the spec declares `reuse_alignment: 2`. +global entry rather than an error. Refuse the combination in the runtime; the +witness checks it when the spec declares `reuse_alignment`. **TRT has none of the runtime half today.** `DeepseekV4CacheManager` in `tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py` @@ -118,16 +115,15 @@ cannot own the scheduler, so name that work in the plan as runtime work, with its owner — a target that fakes it by re-forwarding the whole prefix on a hit is exactly what the witness's accounting check catches. -SGLang's implementation (branch `origin/dsv4.1`, -`python/sglang/srt/model_executor/encoder_swa_replay.py`) is the reference for -the mechanics: it re-runs the cached tokens `[end - window, end)` through the -layers with the compressor skipped, keeps SWA in a per-request ring, and -charges at most `window` replay tokens per request to the scheduler. +Where another engine already implements it, read that for the mechanics and +name it in the plan. The shape is: re-run the cached tokens +`[end - window, end)` through the layers with the compressor skipped, keep SWA +in a per-request ring, and charge at most `window` replay tokens per request +to the scheduler. ## Interaction rules, and saying which combinations work -From SGLang's argument checks for this model -(`python/sglang/srt/arg_groups/deepseek_v4_hook.py`, branch `origin/dsv4.1`): +Combinations a published implementation of these features refuses: | feature | refused together with | |---|---| @@ -159,7 +155,7 @@ boot → accuracy → **replay witness**, all under `TRTLLM_MODELING_V2=require` **The switch.** Each feature is turned off by an environment variable the target reads at construction — for example -`TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY=0` — and the spec names it as the +`MYMODEL_DECODER_REPLAY=0` — and the spec names it as the witness's `contrast_env`. Off must mean the full computation, not an error: the witness runs both. @@ -213,14 +209,14 @@ features: [decoder_replay, encoder_replay] accuracy: witness: decoder_replay: - contrast_env: {TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY: "0"} - window: 128 + contrast_env: {MYMODEL_DECODER_REPLAY: "0"} + window: 128 # the checkpoint's sliding window prompt_tokens: 4096 # > window; at most 16384 encoder_replay: - contrast_env: {TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY: "0"} + contrast_env: {MYMODEL_ENCODER_REPLAY: "0"} window: 128 prefix_tokens: 4096 # > window; at most 16384 - reuse_alignment: 2 # the compress ratio + reuse_alignment: 2 # e.g. the compress ratio ``` `runtime.max_seq_len` (through `extra_flags`) must admit the longest witness diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-host-offload/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-host-offload/SKILL.md index df233f65a506..d19c950edfe2 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-host-offload/SKILL.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-host-offload/SKILL.md @@ -14,9 +14,8 @@ The first instance is Qwen3.8-Flash-Next's PLE n-gram embedding table: **53.7 GB of a 124 GB checkpoint**. Whether it is resident decides how much of a 284 GB device is left for KV. That is the test for whether this feature applies: *a single tensor whose size is comparable to the rest of the model*. -It is not a general memory-pressure remedy. DeepSeek-V4.1's Engram table is -the second: two modules of 384M rows x 256 FP8 plus E8M0 scales, **202.76 GB**, -larger than the rest of the checkpoint. +It is not a general memory-pressure remedy. A hashed n-gram memory table +larger than the rest of its checkpoint is the same case again. Read `staircase-modeling-core/references/adapter-rules.md` first if the offload is expressed as an adapter rather than inside the forward. @@ -118,34 +117,32 @@ one host copy that every rank on the node maps and registers. - rank 0 creates the table in one anonymous shared region (`memfd_create`), backed by huge pages (`madvise(MADV_HUGEPAGE)`); -- the other ranks on the node map the same region — SGLang reaches it through - `/proc//fd/`, which needs every rank in one PID namespace, so one +- the other ranks on the node map the same region; reaching it through + `/proc//fd/` needs every rank in one PID namespace, so one container per node; - every rank `cudaHostRegister`s the mapping and resolves its own device pointer, before capture as above; - every GPU can then gather **any** row: no sharding, no zero-fill, no allreduce or reduce-scatter, and the lookup is bitwise the resident one. -At 202.76 GB for DeepSeek-V4.1's Engram table, per-rank copies on a four-GPU -node would be 811 GB of host memory; the shared layout is 202.76 GB, once. -SGLang ships exactly this behind `SGLANG_DSV41_ENGRAM_HOST_TABLE_LAYOUT=shared` -(implementation in `python/sglang/srt/layers/engram.py`, branch -`origin/dsv4.1`); TRT does not have it yet, so it is new target code, and the -`ple.py` pieces above are the parts it can follow. +For a table of `T` bytes on a node of `G` GPUs, per-rank copies cost `G x T` +of host memory and the shared layout costs `T`, once. Where TRT has no shared +layout yet, it is new target code, and the `ple.py` pieces above are the parts +it can follow. What to check before choosing it, none of which is about correctness: | check | why | |---|---| -| huge pages | SGLang's own note: without them "expect ~10x slower lookups" — TLB misses on a 200 GB mapping | +| huge pages | without them expect lookups an order of magnitude slower: TLB misses on a mapping of hundreds of GB | | NUMA | a GB300 node puts four GPUs on two Grace sockets; one copy lives on one socket (or interleaved), and the far GPUs read it over the socket link | | host RAM | the table once, plus pinned staging for weight loading, must fit beside everything else the node holds | | lifetime | rank 0 owns the region; the others must map it before rank 0 can exit, and a restarted rank 0 invalidates every mapping | The residency witness is unchanged by the layout — it compares per-rank device memory with the offload on and off — but `expected_mib` is the OFF -leg's **per-rank** footprint: DeepSeek-V4.1's table row-sharded at TP4 is -202.76 GB / 4, **48,340 MiB** per rank. +leg's **per-rank** footprint: a table row-sharded at TP4 contributes a +quarter of its bytes, in MiB, per rank. ### The fallback: a host `index_select` diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-kv-quant/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-kv-quant/SKILL.md index 9b1f4c375459..83fadfa6c633 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-kv-quant/SKILL.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-kv-quant/SKILL.md @@ -1,17 +1,17 @@ --- name: staircase-modeling-kv-quant -description: Load when switching a target's KV pools from bf16 containers of the reference's fake-quantized values to the checkpoint's native storage formats (DeepSeek-V4.1's FP4 main KV, MXFP4 indexer K, FP8 SWA KV). Output-identical by construction when attention dequantizes exactly, so boot and accuracy cannot see it; it needs a bytes-per-token witness against a declared figure. +description: Load when switching a target's KV pools from bf16 containers of the reference's fake-quantized values to the checkpoint's native storage formats (for example FP4 compressed KV, MXFP4 indexer keys, FP8 sliding-window KV). Output-identical by construction when attention dequantizes exactly, so boot and accuracy cannot see it; it needs a bytes-per-token witness against a declared figure. --- # staircase-modeling-kv-quant `--features kv_quant`. Core stores the **same values** the reference computes — the reference fake-quantizes KV to its storage formats and attends over the -dequantized result — but in bf16 pools, because the only sm_103 sparse -attention kernels that serve DeepSeek-V4.1 are exact only over bf16 KV. This -unit switches the pools to the formats the checkpoint was trained for. It is -what a request for "FP4 KV cache", "FP8 attention" or "890 bytes of KV per -token" actually asks for. +dequantized result — but in bf16 pools, typically because the attention +kernels that serve the model are exact only over bf16 KV. This unit switches +the pools to the formats the checkpoint was trained for. It is what a request +for "FP4 KV cache", "FP8 attention" or "N bytes of KV per token" actually asks +for. ## The rule: store the reference's value set, never a different one @@ -25,51 +25,46 @@ Substitutions that look like this feature and are not: | format | why it is a different value set | |---|---| -| `kv_cache_config.dtype: fp8` as a per-tensor FP8 cache | one scalar scale, not UE8M0 per 32 channels | -| `fp8_ds_mla` — DeepSeek-V4's 584 B layout | 448 FP8 + 64 bf16 RoPE dims with block-64 scales; V4.1 quantizes all 512 dims, RoPE included, after RoPE (`tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/footer_scale_kv.py` is the V4 layout) | -| TRT's `nvfp4` compressed pool as DeepSeek-V4 uses it | NVFP4 with a second-level global scale and a 64-dim residual; V4.1 has no global scale (it is 1) and all 512 dims in FP4 | +| `kv_cache_config.dtype: fp8` as a per-tensor FP8 cache | one scalar scale, not a block scale per group of channels | +| an in-tree preset with the same bit budget | its scale type, scale granularity, quantized dims (is RoPE included, and before or after rotation?) or second-level global scale differ from the reference's; compare each, never the bit width alone | -The DeepSeek-V4 machinery in +The sparse-attention cache machinery in `tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/` — the `KVCacheDtype` presets in `compressor.py`, `DeepseekV4CacheManager` and its `get_cache_bytes_per_token` in `cache_manager.py` — is the closest thing in -tree, and it encodes the **V4** layouts. Read it to learn the interface; do -not assume its byte formulas describe V4.1. +tree to a multi-pool quantized KV cache. Read it to learn the interface; do +not assume its byte formulas describe the checkpoint in hand. -## DeepSeek-V4.1's formats +## Deriving the checkpoint's formats -From the tech report (§2.4.4) and the reference's fake-quantization: +From the checkpoint's tech report and the reference's fake-quantization, write +down one row per KV pool before touching the target: | pool | who owns one | format | bytes per entry | |---|---|---|---| -| SWA KV | every layer (40), the window only | FP8 E4M3 per 32 channels with a UE8M0 scale, over all 512 dims including RoPE | 512 + 16 = **528** | -| main (compressed global) KV | the KV-source layers only: 2, 8, 14 (compress ratio 2) and 20 (ratio 1) | FP4 E2M1 with an E4M3 scale per 16, **no global scale** | 256 + 32 = **288** | -| indexer K | the same KV-source layers | MXFP4: E2M1 with a UE8M0 scale per 32, 128 dims | 64 + 4 = **68** | - -**Only layers 2, 8, 14 and 20 own compressed and indexer pools.** Every other -layer reads a source layer's pool through `kv_source_layer_ids`. A pool -allocated per compressed layer (38 of them) comes to 10324 B/token, more than -ten times the figure below, and is the sign the sharing was never wired. - -**Bytes per token** is the persistent, growing part — global KV only: - -``` -3 x (288 + 68) / 2 + (288 + 68) = 534 + 356 = 890 B/token - ratio-2 layers ratio-1 layer -``` - -The tech report's headline figure. The same formula gives 3200 for bf16 -containers of these values (core) and 1490 with FP8 main KV. The SWA window -is per request, not per token — 40 layers x 128 x 528 B, about 2.7 MB — and so -is the compressor's carry-over state; neither belongs in the figure. +| e.g. sliding-window KV | which layers, and whether only the window is held | element type, scale type and group size, which dims | payload + scales | +| e.g. compressed global KV | which layers **own** a pool, and at what compress ratio | | | +| e.g. indexer keys | | | | + +**Find out which layers own a pool and which read another layer's.** A model +whose layers share pools through a source-layer mapping (a config key such as +`kv_source_layer_ids`) allocates far fewer pools than it has layers. A pool +allocated per layer multiplies the figure, and is the sign the sharing was +never wired. + +**Bytes per token** is the persistent, growing part, global KV only: for each +pool-owning layer, its entry size divided by its compress ratio, summed. The +same formula over bf16 containers of the same values gives core's figure, +which the contrast leg must exceed. A sliding window is per request, not per +token, and so is any compressor carry-over state; neither belongs in the +figure. ## Attention over native storage -The kernels that serve V4.1 on sm_103 — trtllm-gen's H512 DynamicTokenSparse -path through FlashInfer's `trtllm_batch_decode_sparse_mla_dsv4`, and the -catalog's `thop_attention` — are exact only over bf16 KV. The native formats -are consumed by SGLang's FlashMLA fork, which TRT does not vendor, or by a -**dequantize-and-gather to bf16** step ahead of the kernel. So: +The attention kernels that serve the model may be exact only over bf16 KV. +The native formats are then consumed by a kernel that reads them directly, +where one exists, or by a **dequantize-and-gather to bf16** step ahead of the +kernel. So: - native storage + exact dequantization + bf16 Q is **output-identical** to core. Declare `exact: true` (the default) and the witness demands token @@ -86,7 +81,7 @@ boot → accuracy → **kv_quant witness**, all under `TRTLLM_MODELING_V2=requir once per rank after its KV pools are allocated: ``` -MODELING_V2_WITNESS kv_bytes_per_token=890 +MODELING_V2_WITNESS kv_bytes_per_token=656 ``` Computed from the pools **as allocated** — the tensors' dtypes and shapes, @@ -98,8 +93,9 @@ compressor state are per request and excluded. **The verdict**, against the spec's declared figure: - native storage reports `expected_bytes_per_token` within `bytes_tolerance` - (default 2%) — enough for row padding, not for a layout change: padding the - 4-byte indexer scale to 16 moves 890 to 920; + (default 2%) — enough for row padding, not for a layout change: padding a + 4-byte scale to 16 bytes on one pool already moves the figure by a few + percent; - with `exact: true`, the leg under `contrast_env` (bf16 storage) reports MORE bytes per token — otherwise it selected nothing and the next check compares native storage with itself — and the greedy continuations are @@ -111,16 +107,16 @@ features: [kv_quant] accuracy: witness: kv_quant: - contrast_env: {TRTLLM_DSV41_KV_STORAGE: bf16} # the leg WITH bf16 storage - expected_bytes_per_token: 890 + contrast_env: {MYMODEL_KV_STORAGE: bf16} # the leg WITH bf16 storage + expected_bytes_per_token: 656 # derived above, before the run exact: true # bytes_tolerance: 0.02 max_batch_size: 1 # both legs one request at a time ``` Compare at batch 1 -- `max_batch_size: 1` in the witness block runs both legs -there: SGLang reports that this model's output is not bitwise stable across -batch composition, and the identity check must measure storage, not batching. +there: a model's output is often not bitwise stable across batch composition, +and the identity check must measure storage, not batching. ## What the gates still do not tell you @@ -131,8 +127,7 @@ the pools hold the checkpoint's formats and that the output did not move. ## Report additions - each pool: owner layers, format, bytes per entry, and the allocated dtype -- measured bytes per token against 890, and the bf16 leg's figure +- measured bytes per token against the declared figure, and the bf16 leg's figure - `exact`, and if false, which attention path is lossy and why it was chosen - how attention reads the native values (dequantize-gather, or which kernel) -- confirmation that only the KV-source layers allocate compressed and indexer - pools +- confirmation that only the pool-owning layers allocate pools diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/SKILL.md index 1d926e6006b6..da0233dc9f21 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/SKILL.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/SKILL.md @@ -10,8 +10,8 @@ Covers every `--features` value that makes the target draft its own tokens: skill because they share a binding surface and a gate; they differ in what the draft path computes, which is `references/algorithms.md`. -**A checkpoint's `mtp.*` namespace does not mean MTP.** DeepSeek-V4.1 stores -its DSpark block drafter there, with `num_nextn_predict_layers: 3` beside it. +**A checkpoint's `mtp.*` namespace does not mean MTP.** A checkpoint may store +a DSpark block drafter there, with `num_nextn_predict_layers` beside it. Read what the stages compute before choosing the feature — the DSpark section of `references/algorithms.md` says how to tell — because `mtp` over DSpark weights either fails to load or drafts garbage that verification quietly diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/algorithms.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/algorithms.md index d534727d87a2..efbbdaa907d1 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/algorithms.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-modeling-speculator/references/algorithms.md @@ -19,7 +19,7 @@ draft-target is a different shape and is not covered. | **MTP** | the checkpoint's own MTP module, replayed autoregressively `max_draft_len` times | the MTP module in the checkpoint | `references/multi-token-prediction.md` is the depth. `num_nextn_predict_layers` is a module count, **not** the draft-length knob | | **Eagle3** | a separate small drafter reading the target's hidden states | a separate drafter checkpoint | needs hidden states exported from the trunk — a forward-shape change the other modes do not require | | **DFlash** | as upstream `_torch/speculative/dflash.py` | | not yet exercised by any target here | -| **DSpark** | the checkpoint's own block drafter: one 3-stage pass drafts a whole block, then a sequential Markov refine | the drafter stages — on DeepSeek-V4.1 stored as `mtp.*` | **not MTP, whatever the namespace says.** See the DSpark section below | +| **DSpark** | the checkpoint's own block drafter: one 3-stage pass drafts a whole block, then a sequential Markov refine | the drafter stages, which a checkpoint may store as `mtp.*` | **not MTP, whatever the namespace says.** See the DSpark section below | | n-gram, suffix automaton, PARD | no weights at all — drafts from the emitted prefix | none | no weight-loading change, so the gate reduces to acceptance alone | Only MTP has been carried end to end here. Everything about the others below is @@ -59,11 +59,10 @@ reference is still required, because "drafts nothing, accepts nothing" and ## DSpark — a block drafter that lives where MTP would -DeepSeek-V4.1-Flash is the first checkpoint here. Its tech report (§2.4.3): -"We omit the MTP module during backbone pre-training and use DSpark for -speculative decoding." The drafter nevertheless ships under `mtp.0..2` with -`num_nextn_predict_layers: 3` in the config, which is exactly what makes the -wrong feature easy to pick. +A checkpoint can ship a DSpark drafter under `mtp.*` with +`num_nextn_predict_layers` in its config, which is exactly what makes the +wrong feature easy to pick. Its tech report says which drafter it trained; +quote that in the plan. ### Telling it from MTP @@ -71,7 +70,7 @@ Any one of these settles it: - the config carries `dspark_*` keys — `dspark_block_size`, `dspark_noise_token_id`, `dspark_target_layer_ids`, `dspark_markov_rank` - (DeepSeek-V4.1 nests them under `text_config`); + (a multimodal config may nest them under `text_config`); - the `mtp.*` weights include `main_proj`, `markov_head`, `confidence_head` and lack MTP's signature: no `e_proj` / `h_proj` / `enorm` / `hnorm` / `shared_head`; @@ -79,8 +78,8 @@ Any one of these settles it: module. Then the feature is `dspark` and never `mtp`. `MTPDecodingConfig` would read -`num_nextn_predict_layers: 3`, pick vanilla sequential MTP, and build three -chained one-token heads over weights trained as one 3-stage block pass: it +`num_nextn_predict_layers: N`, pick vanilla sequential MTP, and build N +chained one-token heads over weights trained as one N-stage block pass: it either fails to load the missing MTP keys or drafts garbage that verification rejects — and the second passes accuracy. If a request says "MTP" for such a checkpoint, it means "speculate with the bundled drafter", which is `dspark`; @@ -92,29 +91,28 @@ From the checkpoint's reference forward. Each item is a place a transliteration from an MTP or Eagle mental model goes wrong **and still drafts**, so the only symptom is a lower acceptance length. -1. **One pass drafts the whole block.** The input is a `block_size`-token block - `[x, NOISE, NOISE, NOISE, NOISE]` at block size 5: `x` is the token the - target just emitted, NOISE is `dspark_noise_token_id` (128799). The three - stages run **once**, in sequence, over all five positions — each a full - block (sliding-window attention plus a 128-expert top-3 MoE, with - hyper-connections) — and the head produces base logits for all five - positions in parallel. Not five passes; not three heads chained. +1. **One pass drafts the whole block.** The input is a `block_size`-token + block `[x, NOISE, ..., NOISE]`: `x` is the token the target just emitted, + NOISE is `dspark_noise_token_id`. The stages run **once**, in sequence, + over every block position — each stage a full transformer block as the + reference defines it — and the head produces base logits for all positions + in parallel. Not one pass per position; not chained heads. 2. **Attention inside the block is non-causal.** Every block position attends to every other block position plus the filled ring slots. A causal mask here is the most likely error to carry over from a one-token drafter. -3. **The ring holds TARGET state.** Each stage owns a sliding-window ring - (window 128). It is filled from the target's hidden states: the *input* of - target layers `dspark_target_layer_ids` (37, 38, 39), averaged over the - hyper-connection copies, concatenated, projected once by stage 0's - `main_proj` + `main_norm`, then per stage through `wkv` and `kv_norm` with - plain RoPE. Inputs, not outputs — vLLM's DeepSeek-V4 path captures - outputs, which is a different checkpoint's contract. The block's own K/V - are transient and never enter the ring. After a multi-token acceptance, - every newly committed position's target state must reach the ring before - the next draft, not just the last one. +3. **The ring holds TARGET state.** Each stage owns a sliding-window ring, + filled from the target's hidden states at the layers + `dspark_target_layer_ids` names, combined and projected as the reference + does it. Read from the reference whether a ring takes each layer's + **input or output**, and how parallel residual streams are reduced: that is + a per-checkpoint contract, and another engine's choice for a different + checkpoint is no evidence. The block's own K/V are transient and never + enter the ring. After a multi-token acceptance, every newly committed + position's target state must reach the ring before the next draft, not + just the last one. 4. **The Markov head is applied sequentially.** After the parallel pass, for - each position in order: add a rank-256 bigram bias computed from the - *previous* position's token, then choose this position's token. Applied in + each position in order: add a low-rank (`dspark_markov_rank`) bigram bias + computed from the *previous* position's token, then choose this position's token. Applied in parallel it would condition on NOISE tokens. 5. **The confidence head only schedules.** It predicts per-position acceptance probabilities so a scheduler can choose how many drafts each request @@ -171,7 +169,7 @@ absolute claim. - The `TorchLlmArgs` DSpark validation reads `dspark_*` at the top level of the drafter's `config.json`, or inside `dspark_config` / `dflash_config` — **not - under `text_config`**, where DeepSeek-V4.1 keeps them. So + under `text_config`**, where a multimodal checkpoint may keep them. So `target_layer_ids`, `mask_token_id` and `markov_rank` stay unset unless the spec states them, and with `target_layer_ids` unset `DSparkSpecMetadata` captures nothing: the drafter sees no target context and acceptance falls diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml b/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml index 697e69e25bb1..bfa584572433 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml +++ b/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml @@ -84,10 +84,9 @@ parallel: tp1 # 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 (DeepSeek-V4.1: `dspark_*` keys in its config) asks for -# `dspark` even when the request says "MTP"; `mtp` over those weights builds -# three chained heads the checkpoint never trained. The two are refused -# together. +# 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: [] trtllm_repo_path: /path/to/TensorRT-LLM @@ -449,7 +448,7 @@ accuracy: # # witness: # host_offload: - # contrast_env: {QWEN_PLE_OFFLOAD: "0"} # the leg WITHOUT the offload + # contrast_env: {MYMODEL_HOST_OFFLOAD: "0"} # the leg WITHOUT the offload # expected_mib: 48800 # 20 M rows x 2560 dims, fp8 # disagg: # contrast_env: {TRTLLM_DISAGG: "1"} # the leg WITH disagg @@ -462,8 +461,8 @@ accuracy: # max_batch_size: 1 # one schedule on BOTH legs # # The environment of BOTH legs (`contrast_env` is leg B's alone). # legs_env: - # TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY: "0" - # TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY: "0" + # 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 @@ -495,8 +494,10 @@ accuracy: # names a lossy attention path (an FP8-query kernel) and drops that leg. # # kv_quant: - # contrast_env: {TRTLLM_DSV41_KV_STORAGE: bf16} # the leg WITH bf16 storage - # expected_bytes_per_token: 890 # 3 x (288 + 68) / 2 + (288 + 68) + # 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 @@ -509,14 +510,14 @@ accuracy: # under 16384, since the prompt travels as one argv string. # # decoder_replay: - # contrast_env: {TRTLLM_DSV41_DECODER_SWA_BOUNDED_REPLAY: "0"} # replay OFF - # window: 128 + # contrast_env: {MYMODEL_DECODER_REPLAY: "0"} # replay OFF + # window: 128 # the checkpoint's sliding window # prompt_tokens: 4096 # encoder_replay: - # contrast_env: {TRTLLM_DSV41_ENCODER_SWA_BOUNDED_REPLAY: "0"} # replay OFF + # contrast_env: {MYMODEL_ENCODER_REPLAY: "0"} # replay OFF # window: 128 # prefix_tokens: 4096 - # reuse_alignment: 2 # the compress ratio; an odd boundary splits a pair + # reuse_alignment: 2 # e.g. a KV compression ratio the boundary may not split # The ABSOLUTE `acceptance_length` floor a speculator unit falls back to # when the `=off` leg cannot exist. The acceptance gate normally compares @@ -626,3 +627,33 @@ accuracy: # max_attempts_per_item: 1 # a second attempt at identical work is not progress # max_new_items: 16 # surfaces one round may add to the queue # max_parallel_verifiers: 4 # concurrent verify smiths + +# ------------------------------------------------------- agent backends +# Which backend, model and reasoning effort each role runs on. Omit the block +# for the defaults: plan_reviewer on codex, every other role on claude-code, +# each at its backend's default model. The roles are plan_drafter, +# plan_reviewer, smith, assembler, reviewer, auditor and runner. +# +# The plan reviewer runs on a different model from the plan drafter on +# purpose: the plan phase has no gate, only this review, and a reviewer on the +# drafter's model shares its blind spots. The CLI warns when the two resolve to +# the same backend and model. +# +# `defaults` applies to every role and `roles.` overrides it per field. +# A `defaults.model` needs a `defaults.backend` beside it, because the roles do +# not share a default backend. Under `--agent-transport external` a driver runs +# the turn; the role's configured backend and model reach it as `agent` in the +# turn's request file, and choosing a matching model is the driver's job. +# agents: +# defaults: +# backend: claude-code # claude-code | codex +# model: claude-opus-5-5 +# reasoning_effort: max +# 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} diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py b/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py index e01014e2092f..ddbc795f1819 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py @@ -34,6 +34,10 @@ import yaml +from agent_flow.agent_runtime import AGENTS_FIELD, validate_agents + +from .roles import ROLES, validate_role_defaults + REQUIRED_PATH_FIELDS: tuple[str, ...] = ( "checkpoint_path", "trtllm_repo_path", @@ -330,6 +334,8 @@ REMEDIATION_FIELD, SERVING_ENVELOPE_FIELD, FEATURE_MATRIX_FIELD, + # Per-role backend, model and reasoning effort; see `roles.py`. + AGENTS_FIELD, ) REMEDIATION_KEYS: tuple[str, ...] = ( @@ -1129,8 +1135,7 @@ def _validate_kv_quant_witness(entry: dict[str, Any], errors: list[str], where: 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 (890 for " - f"DeepSeek-V4.1 NVFP4); got {expected!r}" + f"computed from the checkpoint's storage formats before the run; got {expected!r}" ) exact = entry.setdefault("exact", True) if not isinstance(exact, bool): @@ -1567,6 +1572,8 @@ def load_and_validate_task_yaml(path: str | Path) -> dict[str, Any]: # said it was done", which is a much longer trip back to the cause. _validate_remediation(data.get(REMEDIATION_FIELD), errors) reject_unknown_keys(data, TOP_LEVEL_KEYS, errors) + errors.extend(validate_agents(data, ROLES)) + errors.extend(validate_role_defaults(data)) for name in REQUIRED_PATH_FIELDS: if name not in data: @@ -1627,8 +1634,9 @@ def load_and_validate_task_yaml(path: str | Path) -> dict[str, Any]: 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. A DeepSeek-V4.1 `mtp.*` is DSpark, " - f"so its 'MTP' request is spelled `dspark`." + 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 ENVIRONMENT_FIELD not in data: diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/vocabulary.py b/agent-flow/agent_flow/workflows/staircase_onboard/vocabulary.py index 9c77a63e1f09..f0dbc0c71daf 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/vocabulary.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/vocabulary.py @@ -140,10 +140,9 @@ class QueueItem: #: failed dispatch blocks, not when the dispatch runs: every row is #: dispatched in the smith phase, and only a failure that core needs #: stops the run there. A feature-only failure is deferred to the start - #: of its unit and retried once. Measured on the DeepSeek-V4.1 run: 33 - #: dispatches queued, six of them needed only by dspark / multimodal / - #: host_offload / kv_quant, and any one of those failing closed the run - #: `smith-blocked` before core -- which needed none of them -- had begun. + #: of its unit and retried once. Without the tag, any one row that only a + #: feature unit needs could close the run `smith-blocked` before core, + #: which does not need it, had begun. unit: str = UNIT_CORE @classmethod diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py b/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py index 4d7ed7d85310..72097830ceeb 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py @@ -40,8 +40,6 @@ from rich.markup import escape as escape_markup from agent_flow import ( - CLAUDE_CODE_DEFAULT_MODEL, - CODEX_DEFAULT_MODEL, AgentLayer, AgentLayerConfig, BackendConfig, @@ -50,6 +48,7 @@ ) from agent_flow.console import print_message, print_rule from agent_flow.logger import get_logger +from agent_flow.prompts import dump_prompt_bundle from agent_flow.runtime import CommandExecutor, TurnExecutor, TurnRequest from . import envelope as envelope_mod @@ -89,7 +88,8 @@ latest_entry, mark_decision_dispatch, ) -from .prompts import PromptBundle, build_staircase_onboard_prompts +from .prompts import PROMPTS_DIRNAME, PromptBundle, build_staircase_onboard_prompts +from .roles import resolve_role_agent from .runtime import ( areas_for_demands, format_areas, @@ -135,12 +135,6 @@ needed_by, ) -# Which backend each role runs on. The PlanReviewer is deliberately not the -# same model as the PlanDrafter — a reviewer that shares the author's blind -# spots reviews its own reasoning back at itself. Every other role needs the -# skill system, which is a Claude Code feature. -_CODEX_ROLES = frozenset({"plan_reviewer"}) - #: The role that submits a command to a scheduler the orchestrator cannot #: reach. It is the only one that returns no typed decision, because it makes #: none — see `prompts/runner.py` and `runtime/commands/delegated.py`. @@ -477,8 +471,9 @@ def _make_agent( required_tools: list[str], session_mode: str = "stateless", ) -> AgentLayer: - backend_kind = "codex" if name in _CODEX_ROLES else "claude-code" - model = CODEX_DEFAULT_MODEL if name in _CODEX_ROLES else CLAUDE_CODE_DEFAULT_MODEL + # Backend, model and reasoning effort come from `task.yaml`'s `agents` + # block, over the defaults in `roles.py`. + agent = resolve_role_agent(self.task_data, name) # Under an enterprise MCP configuration the in-process tool server # cannot be registered at all, so the role reports by writing a file # and the Stop hook requires the write rather than the tool. `Write` @@ -491,15 +486,25 @@ def _make_agent( hooks = _compose_required_tools_hooks( (["Write"] if self._decision_files else required_tools) if required_tools else [] ) + layer_required: tuple[str, ...] = () + if agent.backend == "codex": + # Codex refuses per-client Claude SDK hooks, so the requirement + # moves to the layer's backend-neutral policy. A decision that + # travels as a file has no tool to require on Codex, which writes + # files with its own tools; the file is validated when it is read. + hooks = None + layer_required = () if self._decision_files else tuple(required_tools) return AgentLayer( AgentLayerConfig( name=name, system_prompt=system_prompt, backend=BackendConfig( - kind=backend_kind, - model=model, + kind=agent.backend, + model=agent.model, + reasoning_effort=agent.reasoning_effort, tools=tools, hooks=hooks, + extra_mcp_servers=agent.extra_mcp_servers, # The workspace is the session cwd so the run's own # materialized skills are the ones the agent discovers, # rather than whatever happens to be configured in the @@ -507,6 +512,7 @@ def _make_agent( cwd=self.workspace, ), session=SessionConfig(mode=session_mode), + required_tools=layer_required, ) ) @@ -534,11 +540,14 @@ def _run_turn( rather than the executor -- ten smiths are ten agents, not ten turns of one. `bound` is a no-op for transports with no session concept. """ + agent = resolve_role_agent(self.task_data, role) request = TurnRequest( role=role, what=what, prompt=prompt, system_prompt=getattr(self.prompts, role, "") if self.prompts else "", + backend=agent.backend, + model=agent.model, decision_file=( str(decision_path(self.workspace, role, iteration)) if self._decision_files and iteration is not None @@ -784,6 +793,10 @@ def _run_preflight(self, state: WorkflowState, log) -> bool: decision_files=self._decision_files, commit_script=str(commit_script), ) + # What every role is told, verbatim and per launch: the composition + # depends on the resolved environment and the spec, so it otherwise + # exists only in this process's memory. + dump_prompt_bundle(self.prompts, self.workspace / PROMPTS_DIRNAME) # Written here, from what preflight has just resolved, so that the # first agent turn of the run already has them. Both files are # referenced by `_context_block`, which every dispatch carries. diff --git a/agent-flow/tests/runtime/test_turn_executors.py b/agent-flow/tests/runtime/test_turn_executors.py index 264d22af4d95..a7282ed4581a 100644 --- a/agent-flow/tests/runtime/test_turn_executors.py +++ b/agent-flow/tests/runtime/test_turn_executors.py @@ -351,3 +351,94 @@ def test_an_answered_or_different_request_is_never_adopted(tmp_path): TurnRequest(role="smith", what="onboard a", prompt="p", decision_file="d/1") ) assert outcome.request_id == "003-smith" + + +def _wait_for(path, timeout_s: float = 5.0) -> None: + deadline = time.monotonic() + timeout_s + while not path.exists(): + assert time.monotonic() < deadline, f"{path} never appeared" + time.sleep(0.005) + + +def test_a_concurrent_caller_asking_alike_is_never_handed_a_siblings_request(tmp_path): + """Two in-flight turns of one process are two questions, even when they read alike. + + The ordering is forced rather than left to the scheduler: the second + caller looks only after the first caller's request file is on disk. That + is the interleaving in which the twin check used to adopt the sibling's + request, so both callers waited on one answer and one prompt was never + delivered. + """ + # Bounded, and daemon threads: a regression must fail here, not hang the run. + executor = FileDispatchTurnExecutor(tmp_path, poll_s=0.005, timeout_s=10) + alike = dict(role="smith", what="onboard a", decision_file=None) + outcomes: dict[str, TurnOutcome] = {} + + def ask(key: str) -> threading.Thread: + prompt = f"{key} prompt" + thread = threading.Thread( + target=lambda: outcomes.__setitem__( + key, executor.run(TurnRequest(prompt=prompt, **alike)) + ), + daemon=True, + ) + thread.start() + return thread + + first = ask("first") + _wait_for(executor.root / "001-smith.request.json") + second = ask("second") + _wait_for(executor.root / "002-smith.request.json") + + executor.answer("001-smith", "answer one") + executor.answer("002-smith", "answer two") + first.join(timeout=5) + second.join(timeout=5) + + assert outcomes["first"].request_id == "001-smith" + assert outcomes["second"].request_id == "002-smith" + assert outcomes["first"].text == "answer one" + assert outcomes["second"].text == "answer two" + assert (executor.root / "002-smith.prompt.md").read_text(encoding="utf-8") == "second prompt" + + +def test_a_retry_in_the_same_process_waits_on_its_own_unanswered_request(tmp_path): + """Once its caller has given up, a request is an orphan like a dead process's. + + Re-asking would duplicate work a driver may still be doing, exactly as on + a resume. + """ + executor = FileDispatchTurnExecutor(tmp_path, poll_s=0.005, timeout_s=0.05) + request = TurnRequest(role="smith", what="onboard a", prompt="p", decision_file="d/1") + with pytest.raises(TurnTimeout): + executor.run(request) + + def answer_once_adopted(): + # Answering before the retry has looked would make 001 an answered + # request, which is rightly never adopted; wait until the retry holds it. + deadline = time.monotonic() + 5 + while "001-smith" not in executor._in_flight and time.monotonic() < deadline: + time.sleep(0.005) + executor.answer("001-smith", "late answer") + + executor.timeout_s = 5.0 + threading.Thread(target=answer_once_adopted, daemon=True).start() + outcome = executor.run(request) + assert outcome.request_id == "001-smith" + assert outcome.text == "late answer" + assert not (executor.root / "002-smith.request.json").exists() + + +def test_the_request_file_names_the_backend_the_role_is_configured_for(tmp_path): + """A driver can only honour a role's model if the request says which one it is.""" + import json + + executor = FileDispatchTurnExecutor(tmp_path, poll_s=0.005) + _answer_after(executor, "001-plan_reviewer", "ok") + executor.run( + TurnRequest(role="plan_reviewer", what="r1", prompt="p", backend="codex", model="m") + ) + request = json.loads( + (executor.root / "001-plan_reviewer.request.json").read_text(encoding="utf-8") + ) + assert request["agent"] == {"backend": "codex", "model": "m"} diff --git a/agent-flow/tests/workflows/staircase_onboard/test_delivery.py b/agent-flow/tests/workflows/staircase_onboard/test_delivery.py index 9be242ff4e0a..1a66b95d9464 100644 --- a/agent-flow/tests/workflows/staircase_onboard/test_delivery.py +++ b/agent-flow/tests/workflows/staircase_onboard/test_delivery.py @@ -280,19 +280,44 @@ def test_verify_delivery_includes_the_git_check(): not (_REPO_ROOT / delivery.MODELING_V2_PACKAGE).is_dir(), reason="not running inside a TensorRT-LLM checkout with modeling_v2 present", ) -def test_the_real_merged_target_passes_every_check(): +@pytest.mark.parametrize( + ("target_rel", "family", "model_key", "accuracy_test"), + [ + (TARGET_REL, "gpt_oss", "openai/gpt-oss-120b", "test_modeling_v2_gpt_oss.py"), + ( + f"{delivery.MODELING_V2_PACKAGE}/models/deepseek_v3/r1_0528_nvfp4__sm_103__dep4", + "deepseek_v3", + "deepseek-ai/DeepSeek-R1-0528", + "test_modeling_v2_deepseek_v3.py", + ), + ], + ids=["gpt_oss", "deepseek_v3"], +) +def test_the_real_merged_target_passes_every_check(target_rel, family, model_key, accuracy_test): """The path conventions must match what the repository actually uses. - `gpt_oss / sm_103 / tp1` is merged and gated, so every check here has a - known-good answer. If this fails, a convention moved and the checks are - describing a layout that no longer exists. + Both targets are merged and gated, so every check here has a known-good + answer. If this fails, a convention moved and the checks are describing a + layout that no longer exists. The DeepSeek one is multi-rank, so it also + holds the convention for the `comm` entries' op-matrix tests. """ report = delivery.verify_delivery( repo_root=_REPO_ROOT, - target_rel=TARGET_REL, - family="gpt_oss", + target_rel=target_rel, + family=family, reference_yaml=f"{delivery.ACCURACY_REFERENCES_DIR}/gsm8k.yaml", - model_key="openai/gpt-oss-120b", - accuracy_test="test_modeling_v2_gpt_oss.py", + model_key=model_key, + accuracy_test=accuracy_test, ) assert report.passed, report.render() + + +def test_a_multi_rank_entry_is_satisfied_by_its_op_matrix_test(tmp_path): + tests = tmp_path / delivery.CATALOG_TESTS_DIR / "comm" + tests.mkdir(parents=True) + (tests / "test_modeling_v2_allgather_op_matrix.py").write_text("", encoding="utf-8") + [check] = delivery.check_catalog_tests(tmp_path, [("comm", "allgather")]) + assert check.passed, check.detail + [missing] = delivery.check_catalog_tests(tmp_path, [("comm", "reducescatter")]) + assert not missing.passed + assert missing.detail.endswith("comm/test_modeling_v2_reducescatter.py") diff --git a/agent-flow/tests/workflows/staircase_onboard/test_dspark.py b/agent-flow/tests/workflows/staircase_onboard/test_dspark.py index 21067d029132..8546bb9b1ae5 100644 --- a/agent-flow/tests/workflows/staircase_onboard/test_dspark.py +++ b/agent-flow/tests/workflows/staircase_onboard/test_dspark.py @@ -3,8 +3,8 @@ """DSpark as its own speculator kind, and never a spelling of MTP. -DeepSeek-V4.1 ships its DSpark block drafter as `mtp.*` with -`num_nextn_predict_layers: 3`. Every test here guards one of the ways that +A checkpoint can ship its DSpark block drafter as `mtp.*` beside +`num_nextn_predict_layers`. Every test here guards one of the ways that namespace leads a run to gate the wrong drafter, or the right drafter at the wrong depth, or not at all. """ @@ -12,6 +12,7 @@ from __future__ import annotations import ast +import textwrap from pathlib import Path import pytest @@ -105,7 +106,7 @@ def run(command, env=None, timeout=0): if "TRTLLM_MODELING_V2=off" in command: return ShellResult( returncode=1, - stderr="RuntimeError: Unknown model architecture: DeepseekV41ForCausalLM", + stderr="RuntimeError: Unknown model architecture: NewArchForCausalLM", ) return ShellResult(returncode=0, stdout=accept_line) @@ -153,23 +154,87 @@ def test_below_the_declared_floor_fails(tmp_path): # --------------------------------------------------------------- quickstart -def _function_source(name: str) -> str: - source = QUICKSTART.read_text(encoding="utf-8") +def _function_source(name: str, source: str | None = None) -> str: + source = QUICKSTART.read_text(encoding="utf-8") if source is None else source node = next( n for n in ast.parse(source).body if isinstance(n, ast.FunctionDef) and n.name == name ) return ast.get_source_segment(source, node) or "" -def test_quickstart_builds_a_dspark_config(): - """The gates emit `--spec_decode_algo=DSPARK`; quickstart has to honour it.""" - setup = _function_source("setup_llm") - branch = setup[setup.index('spec_decode_algo == "DSPARK"') :] - branch = branch[: branch.index("elif")] +def _dspark_branch(source: str) -> str | None: + """The `DSPARK` arm of `setup_llm`'s speculator chain, or None when it has none. + + Quickstart falls through to `spec_config = None` for an algorithm it does + not know, so a checkout without this arm runs the dspark unit's boot and + acceptance commands with no speculator at all. + """ + setup = _function_source("setup_llm", source) + marker = 'spec_decode_algo == "DSPARK"' + if marker not in setup: + return None + branch = setup[setup.index(marker) :] + return branch[: branch.index("elif")] + + +#: The shape of a quickstart that honours `--spec_decode_algo=DSPARK`. A +#: fixture, so the contract is checked whether or not this checkout's +#: quickstart has the arm yet. +QUICKSTART_WITH_DSPARK = textwrap.dedent( + """\ + def setup_llm(args, **kwargs): + spec_decode_algo = args.spec_decode_algo.upper( + ) if args.spec_decode_algo is not None else None + if spec_decode_algo == 'MTP': + spec_config = MTPDecodingConfig(max_draft_len=args.spec_decode_max_draft_len) + elif spec_decode_algo == "DSPARK": + spec_config = DSparkDecodingConfig( + max_draft_len=args.spec_decode_max_draft_len, + speculative_model=args.draft_model_dir) + elif spec_decode_algo == "NGRAM": + spec_config = NGramDecodingConfig(max_draft_len=args.spec_decode_max_draft_len) + else: + spec_config = None + """ +) + + +def _assert_honours_dspark(branch: str | None) -> None: + assert branch is not None, "setup_llm has no DSPARK arm" assert "DSparkDecodingConfig(" in branch assert "max_draft_len=args.spec_decode_max_draft_len" in branch +def test_a_quickstart_with_a_dspark_arm_builds_the_config_at_the_flag_depth(): + """The gates emit `--spec_decode_algo=DSPARK`; quickstart has to honour it.""" + _assert_honours_dspark(_dspark_branch(QUICKSTART_WITH_DSPARK)) + + +def test_a_quickstart_without_a_dspark_arm_is_told_apart(): + without = QUICKSTART_WITH_DSPARK.replace('"DSPARK"', '"OTHER"') + assert _dspark_branch(without) is None + + +def test_a_dspark_arm_at_another_depth_is_rejected(): + fixed = QUICKSTART_WITH_DSPARK.replace( + "max_draft_len=args.spec_decode_max_draft_len,", "max_draft_len=1," + ) + with pytest.raises(AssertionError): + _assert_honours_dspark(_dspark_branch(fixed)) + + +@pytest.mark.skipif(not QUICKSTART.is_file(), reason="not inside a TensorRT-LLM checkout") +def test_this_checkouts_quickstart_honours_dspark_once_it_has_the_arm(): + branch = _dspark_branch(QUICKSTART.read_text(encoding="utf-8")) + if branch is None: + pytest.skip( + "this checkout's quickstart_advanced.py has no DSPARK arm yet, so a dspark " + "unit's boot and acceptance commands run without a speculator; the " + "acceptance gate then finds no accept-rate line and fails closed" + ) + _assert_honours_dspark(branch) + + def test_the_accept_rate_line_does_not_depend_on_the_algorithm(): """The acceptance gate parses `Accept rate: ... (tokens=T, iterations=I)`. diff --git a/agent-flow/tests/workflows/staircase_onboard/test_kernel_sources.py b/agent-flow/tests/workflows/staircase_onboard/test_kernel_sources.py index eae0880b5f70..6d021b2dae38 100644 --- a/agent-flow/tests/workflows/staircase_onboard/test_kernel_sources.py +++ b/agent-flow/tests/workflows/staircase_onboard/test_kernel_sources.py @@ -70,11 +70,11 @@ def test_every_registration_file_in_the_tree_is_classified(sources): A new `custom_ops/_custom_ops.py` lands and nothing tells an agent what kind of kernel it holds. This fails when that happens. """ - known = {Path(path).stem for source in sources.values() for path in source.registers} unclassified = sorted( py.stem for py in CUSTOM_OPS.glob("*.py") - if py.stem not in known and py.stem not in NOT_A_SOURCE + if py.stem not in NOT_A_SOURCE + and source_for_stem(py.stem, sources, text=py.read_text(encoding="utf-8")) is None ) assert not unclassified, ( "custom_ops files no kernel source claims --- an agent meeting an op " @@ -104,8 +104,10 @@ def test_attribution_resolves_real_catalog_ops_to_real_sources(sources): pytest.skip("catalog index not present") registrations = {} + texts = {} for py in sorted(CUSTOM_OPS.glob("*.py")): - for name in re.findall(r'"trtllm::([a-zA-Z0-9_]+)', py.read_text(encoding="utf-8")): + texts[py.stem] = py.read_text(encoding="utf-8") + for name in re.findall(r'"trtllm::([a-zA-Z0-9_]+)', texts[py.stem]): registrations.setdefault(name, py.stem) entries = yaml.safe_load(index.read_text(encoding="utf-8"))["entries"] @@ -121,13 +123,66 @@ def test_attribution_resolves_real_catalog_ops_to_real_sources(sources): stem = registrations.get(op, "cpp_thop") if stem in NOT_A_SOURCE: continue # a registration mechanism; provenance is elsewhere - if source_for_stem(stem, sources) is None: + if source_for_stem(stem, sources, text=texts.get(stem)) is None: unresolved.append(f"{op} (registered by {stem})") assert not unresolved, "catalog ops whose provenance cannot be resolved:\n " + "\n ".join( unresolved ) +_TRITON_OP_FILE = """ +import torch +import triton + + +@triton.jit +def _kernel(x_ptr): + pass + + +@torch.library.custom_op("trtllm::new_triton_op", mutates_args=()) +def new_triton_op(x: torch.Tensor) -> torch.Tensor: + return x +""" + + +def test_a_triton_custom_op_file_no_path_names_is_claimed_by_content(sources): + """A model's onboarding adds these one op at a time, under names nobody predicted. + + Naming each in the registry made every such addition fail this suite + until the registry was edited to match the model's tree. + """ + source = source_for_stem("new_triton_op", sources, text=_TRITON_OP_FILE) + assert source is not None and source.name == "triton_custom_op" + + +def test_a_file_with_only_one_marker_is_not_claimed(sources): + plain_op = _TRITON_OP_FILE.replace("@triton.jit", "") + assert source_for_stem("new_plain_op", sources, text=plain_op) is None + + +def test_a_named_path_wins_over_content(sources): + """Content is the fallback for unnamed files, never a reclassification.""" + source = source_for_stem("cute_dsl_custom_ops", sources, text=_TRITON_OP_FILE) + assert source.name == "cute_dsl" + + +def test_the_rendered_block_names_the_content_markers(sources): + rendered = sources["triton_custom_op"].render() + assert "`@triton.jit` and `torch.library.custom_op`" in rendered + + +def test_schema_rejects_content_markers_that_are_not_strings(tmp_path): + bad = tmp_path / "k.yaml" + bad.write_text( + "x:\n summary: s\n tier: 1\n availability: a b c d e\n certifying: a b c d e\n" + " registers: [a.py]\n claims_files_containing: [1]\n", + encoding="utf-8", + ) + with pytest.raises(RuntimeIndexError, match="claims_files_containing"): + load_kernel_sources(bad) + + def test_the_cpp_plane_is_the_conclusion_for_an_unclaimed_op(sources): """Not a fallback: an op no Python file registers IS a C++ op.""" assert source_for_stem("cpp_thop", sources) is not None diff --git a/agent-flow/tests/workflows/staircase_onboard/test_package_hygiene.py b/agent-flow/tests/workflows/staircase_onboard/test_package_hygiene.py index 1f2281b2a91a..d3017a202dde 100644 --- a/agent-flow/tests/workflows/staircase_onboard/test_package_hygiene.py +++ b/agent-flow/tests/workflows/staircase_onboard/test_package_hygiene.py @@ -64,3 +64,50 @@ def test_no_tracked_file_names_a_personal_path(root): if PERSONAL_PATH.search(line): offences.append(f"{path.relative_to(root)}:{number}: {line.strip()[:100]}") assert not offences, "personal paths:\n " + "\n ".join(offences) + + +def test_raw_run_reports_do_not_ship_and_their_lessons_do(): + """Run reports stay out of the package; the lessons drawn from them are linked and shipped.""" + assert not (PACKAGE / "findings").exists() + assert "(LESSONS.md)" in (PACKAGE / "README.md").read_text(encoding="utf-8") + pyproject = PACKAGE.parents[2] / "pyproject.toml" + assert '"LESSONS.md"' in pyproject.read_text(encoding="utf-8") + + +def test_the_readme_flags_table_is_the_cli(): + """Every flag the CLI accepts is in the README's table, and the table names no other.""" + import argparse + + from agent_flow.workflows.staircase_onboard import cli + + captured: list[argparse.ArgumentParser] = [] + original = argparse.ArgumentParser.parse_args + + def capture(self, args=None, namespace=None): + captured.append(self) + return original(self, ["--task", "t.yaml"], namespace) + + argparse.ArgumentParser.parse_args = capture + try: + cli._parse_args([]) + finally: + argparse.ArgumentParser.parse_args = original + accepted = { + option + for action in captured[0]._actions + for option in action.option_strings + if option.startswith("--") and option != "--help" + } + + readme = (PACKAGE / "README.md").read_text(encoding="utf-8") + table = readme.split("## Flags", 1)[1].split("\n\n", 2)[1] + documented = set(re.findall(r"`(--[a-z-]+)`", table)) + assert documented == accepted + for subcommand in cli.SUBCOMMANDS: + assert f"`staircase-onboard {subcommand} " in readme + + +def test_the_readme_has_no_placeholder_left(): + readme = (PACKAGE / "README.md").read_text(encoding="utf-8") + assert "To be written" not in readme + assert "workstream" not in readme diff --git a/agent-flow/tests/workflows/staircase_onboard/test_prompt_snapshot.py b/agent-flow/tests/workflows/staircase_onboard/test_prompt_snapshot.py new file mode 100644 index 000000000000..45ce25afe42c --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_prompt_snapshot.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The composed role prompts are snapshotted into the workspace at every launch. + +Same contract as the perf workflows (`agent_flow.dump_prompt_bundle`). What +is specific here is *when*: a staircase prompt carries the launch environment +that preflight resolves, so the snapshot is taken after preflight has composed +the bundle, and a dry run already leaves one behind. +""" + +from __future__ import annotations + +import dataclasses + +from agent_flow.workflows.staircase_onboard.prompts import DEFAULT_PROMPTS, PROMPTS_DIRNAME + +from .test_workflow import Script, _install, _workflow + + +def _snapshot(tmp_path) -> dict[str, str]: + directory = tmp_path / "ws" / PROMPTS_DIRNAME + return {path.stem: path.read_text(encoding="utf-8") for path in directory.glob("*.md")} + + +def test_a_dry_run_snapshots_every_role_as_composed(tmp_path, task_file, monkeypatch): + workflow = _workflow(tmp_path, dry_run=True) + _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task_file) + + snapshot = _snapshot(tmp_path) + roles = {field.name for field in dataclasses.fields(workflow.prompts)} + assert set(snapshot) == roles + for role in roles: + assert snapshot[role] == getattr(workflow.prompts, role), role + # Composed, not the bare defaults: the resolved launch environment is in it. + assert snapshot["assembler"] != DEFAULT_PROMPTS.assembler + assert "/opt/venv/bin/python" in snapshot["assembler"] + + +def test_an_injected_bundle_is_the_one_snapshotted(tmp_path, task_file, monkeypatch): + """The snapshot records what the agents are built from, whoever composed it.""" + bundle = DEFAULT_PROMPTS.with_extensions(reviewer="INJECTED FOR THIS RUN") + workflow = _workflow(tmp_path, dry_run=True, prompts=bundle) + _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task_file) + + assert _snapshot(tmp_path)["reviewer"].endswith("INJECTED FOR THIS RUN") diff --git a/agent-flow/tests/workflows/staircase_onboard/test_roles.py b/agent-flow/tests/workflows/staircase_onboard/test_roles.py new file mode 100644 index 000000000000..a1217f7da8ce --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_roles.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Per-role backend routing for the staircase roles. + +The routing itself is the shared `agents` block (`agent_flow.agent_runtime`). +What is specific here is the default assignment, in which the plan reviewer is +on Codex, the validation of a bare `defaults.model` that this assignment needs, +and the way a layer is built for a role that resolves to Codex. +""" + +from __future__ import annotations + +import pytest + +from agent_flow.config import CLAUDE_CODE_DEFAULT_MODEL, CODEX_DEFAULT_MODEL +from agent_flow.workflows.staircase_onboard import roles, task_schema +from agent_flow.workflows.staircase_onboard.prompts import PromptBundle + +from .test_task_schema import _valid_spec, _write +from .test_workflow import _workflow + + +def test_every_prompt_role_is_routable(): + """A role with a system prompt is a role with a backend, and the reverse.""" + assert set(roles.ROLES) == set(PromptBundle.__dataclass_fields__) + + +def test_without_an_agents_block_the_plan_reviewer_runs_on_codex(): + reviewer = roles.resolve_role_agent({}, "plan_reviewer") + assert (reviewer.backend, reviewer.model) == ("codex", CODEX_DEFAULT_MODEL) + for role in set(roles.ROLES) - {"plan_reviewer"}: + agent = roles.resolve_role_agent({}, role) + assert (agent.backend, agent.model) == ("claude-code", CLAUDE_CODE_DEFAULT_MODEL), role + assert agent.reasoning_effort is None + + +def test_a_role_override_changes_only_that_role(): + task = { + "agents": { + "roles": { + "plan_reviewer": {"backend": "claude-code", "model": "claude-other"}, + "smith": {"backend": "codex", "reasoning_effort": "high"}, + } + } + } + reviewer = roles.resolve_role_agent(task, "plan_reviewer") + smith = roles.resolve_role_agent(task, "smith") + assembler = roles.resolve_role_agent(task, "assembler") + assert (reviewer.backend, reviewer.model) == ("claude-code", "claude-other") + assert (smith.backend, smith.model, smith.reasoning_effort) == ( + "codex", + CODEX_DEFAULT_MODEL, + "high", + ) + assert (assembler.backend, assembler.model) == ("claude-code", CLAUDE_CODE_DEFAULT_MODEL) + + +def test_defaults_with_a_backend_reach_every_role_including_the_reviewer(): + task = {"agents": {"defaults": {"backend": "claude-code", "model": "claude-x"}}} + for role in roles.ROLES: + agent = roles.resolve_role_agent(task, role) + assert (agent.backend, agent.model) == ("claude-code", "claude-x"), role + + +def test_the_spec_accepts_an_agents_block(tmp_path): + spec = _valid_spec(tmp_path) + spec["agents"] = { + "defaults": {"backend": "claude-code", "model": "claude-x"}, + "roles": {"plan_reviewer": {"backend": "codex", "reasoning_effort": "high"}}, + } + data = task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert data["agents"]["roles"]["plan_reviewer"]["backend"] == "codex" + + +@pytest.mark.parametrize( + ("agents", "message"), + [ + ({"roles": {"planner": {"backend": "codex"}}}, "unknown role(s) ['planner']"), + ({"roles": {"smith": {"backend": "gemini"}}}, "must be 'claude-code' or 'codex'"), + ({"roles": {"smith": {"temperature": 0}}}, "unknown field(s) ['temperature']"), + ({"defaults": {"model": "claude-x"}}, "needs 'agents.defaults.backend'"), + ("codex", "'agents' must be a mapping"), + ], +) +def test_the_spec_rejects_a_malformed_agents_block(tmp_path, agents, message): + spec = _valid_spec(tmp_path) + spec["agents"] = agents + with pytest.raises(task_schema.TaskSchemaError) as excinfo: + task_schema.load_and_validate_task_yaml(_write(tmp_path, spec)) + assert message in str(excinfo.value) + + +def test_a_same_model_review_is_reported_not_hidden(): + same = {"agents": {"defaults": {"backend": "claude-code"}}} + assert "not cross-model" in roles.shares_the_drafters_model(same) + assert roles.shares_the_drafters_model({}) == "" + + +def test_a_claude_role_keeps_its_stop_hook(tmp_path): + workflow = _workflow(tmp_path) + # Pinned: the default depends on whether this host has an enterprise MCP config. + workflow._decision_files = False + workflow.task_data = {"agents": {"roles": {"assembler": {"reasoning_effort": "high"}}}} + layer = workflow._make_agent( + "assembler", "system", required_tools=["append_assembler_progress"] + ) + backend = layer.config.backend + assert (backend.kind, backend.model, backend.reasoning_effort) == ( + "claude-code", + CLAUDE_CODE_DEFAULT_MODEL, + "high", + ) + assert backend.hooks and backend.hooks["Stop"] + assert backend.cwd == workflow.workspace + assert layer.config.required_tools == () + + +def test_a_codex_role_requires_its_tool_without_claude_hooks(tmp_path): + """Codex refuses per-client Claude SDK hooks, so a Stop hook would fail the turn.""" + workflow = _workflow(tmp_path) + workflow._decision_files = False + workflow.task_data = { + "agents": { + "roles": { + "plan_reviewer": { + "extra_mcp_servers": {"kb": {"type": "http", "url": "https://kb.test/mcp"}} + } + } + } + } + layer = workflow._make_agent( + "plan_reviewer", + "system", + required_tools=["append_plan_reviewer_progress"], + session_mode="persistent", + ) + backend = layer.config.backend + assert (backend.kind, backend.model) == ("codex", CODEX_DEFAULT_MODEL) + assert backend.hooks is None + assert backend.extra_mcp_servers == {"kb": {"type": "http", "url": "https://kb.test/mcp"}} + assert layer.config.required_tools == ("append_plan_reviewer_progress",) + + +def test_a_codex_role_with_a_decision_file_requires_no_tool(tmp_path): + """Codex has no `Write` tool; the decision file is validated when it is read.""" + workflow = _workflow(tmp_path) + workflow._decision_files = True + layer = workflow._make_agent( + "plan_reviewer", "system", required_tools=["append_plan_reviewer_progress"] + ) + assert layer.config.backend.hooks is None + assert layer.config.required_tools == () + + +def test_every_turn_carries_the_roles_backend_whatever_the_transport(tmp_path): + """The external transport hands the turn to a driver that chooses its own model.""" + from agent_flow.runtime import TurnOutcome + + class Capture: + def __init__(self): + self.requests = [] + + def run(self, request): + self.requests.append(request) + return TurnOutcome(role=request.role, text="", transport="capture") + + workflow = _workflow(tmp_path) + workflow._turns = Capture() + workflow.task_data = {"agents": {"roles": {"smith": {"backend": "codex", "model": "m"}}}} + workflow._run_turn("plan_reviewer", "round 1", "prompt") + workflow._run_turn("smith", "onboard x", "prompt") + sent = [(r.role, r.backend, r.model) for r in workflow._turns.requests] + assert sent == [("plan_reviewer", "codex", CODEX_DEFAULT_MODEL), ("smith", "codex", "m")] diff --git a/agent-flow/tests/workflows/staircase_onboard/test_substrate_docs.py b/agent-flow/tests/workflows/staircase_onboard/test_substrate_docs.py new file mode 100644 index 000000000000..c953c4499baf --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_substrate_docs.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The operator-facing docs say what the substrate does, in the words its output uses.""" + +from __future__ import annotations + +import re +from pathlib import Path + +import agent_flow.workflows.staircase_onboard as pkg +from agent_flow.workflows.staircase_onboard.substrate import scheduler + +PACKAGE = Path(pkg.__file__).parent + + +def _doc(name: str) -> str: + return (PACKAGE / name).read_text(encoding="utf-8") + + +def test_the_safe_kill_rule_is_written_down_where_an_operator_looks(): + """An operator facing a hung job must find the rule before reaching for kill.""" + text = _doc("TROUBLESHOOTING.md") + assert "safe-kill rule" in text + assert "spoolctl.sh cancel" in text + assert "_errhandler status=-61" in text + + +def test_every_verdict_reason_the_scheduler_writes_is_documented(): + """A receipt says `reason=`; the operator greps for exactly that.""" + text = _doc("TROUBLESHOOTING.md") + reasons = set( + re.findall(r'Verdict\(\s*RC_[A-Z_]+,\s*"([a-z_]+)"', Path(scheduler.__file__).read_text()) + ) + reasons |= {"orphaned", "host_memory", "memory_pressure"} + undocumented = sorted(r for r in reasons if f"reason={r}" not in text and f"`{r}`" not in text) + assert not undocumented, f"verdict reasons missing from TROUBLESHOOTING.md: {undocumented}" + + +def test_every_return_code_the_docs_quote_is_the_schedulers(): + text = _doc("TROUBLESHOOTING.md") + quoted = {int(code) for code in re.findall(r"\(rc (\d+)\)", text)} + codes = {getattr(scheduler, name) for name in dir(scheduler) if name.startswith("RC_")} + assert quoted and quoted <= codes, quoted - codes + + +def test_the_build_page_sends_rebuilds_and_worktrees_through_the_substrate(): + text = _doc("BUILD_AND_TEST.md") + assert "incremental_build.sh" in text and "fresh inode" in text + assert "--tree " in text + assert "staircase-onboard session render" in text + + +def test_the_readme_describes_the_substrate(): + text = _doc("README.md") + section = text[text.index("## Execution substrate\n") :].split("\n## ")[0] + for fragment in ("--tree", "safe-kill rule", "staircase-onboard driver", "session render"): + assert fragment in section diff --git a/agent-flow/tests/workflows/staircase_onboard/test_system_prompt_spill.py b/agent-flow/tests/workflows/staircase_onboard/test_system_prompt_spill.py new file mode 100644 index 000000000000..691c4813b8bf --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_system_prompt_spill.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A staircase role's system prompt reaches the Claude CLI through the spill. + +The Claude backend writes a system prompt above a threshold to a file under +the session cwd and passes a pointer on argv instead, because Linux caps one +argv string at 128 KiB. Staircase prompts are composed per run from the role +doctrine, the launch environment, the embargo and the decision-file blocks, so +their size is not fixed. These tests follow a role layer's own client options +into the backend, so a change to the layer's cwd or to how it builds its client +would show up here. +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from agent_flow.backends import claude_code as cc_mod +from agent_flow.backends import create_backend + +from .test_workflow import _workflow + + +async def _launched_options(monkeypatch, layer) -> Any: + captured: dict[str, Any] = {} + + class FakeSdkClient: + def __init__(self, options): + captured["options"] = options + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + return None + + monkeypatch.setattr(cc_mod, "ClaudeSDKClient", FakeSdkClient) + monkeypatch.setattr(cc_mod, "create_sdk_mcp_server", lambda **_: object()) + backend = create_backend(layer.config.backend) + async with backend.create_client(**layer._client_options(layer.config.system_prompt)): + pass + return captured["options"] + + +@pytest.mark.parametrize("decision_files", [False, True]) +async def test_an_oversized_role_prompt_is_spilled_into_the_run_workspace( + tmp_path, monkeypatch, decision_files +): + workflow = _workflow(tmp_path) + workflow._decision_files = decision_files + prompt = workflow.prompts.assembler + "\n" + "x" * (cc_mod.SYSTEM_PROMPT_FILE_THRESHOLD + 1) + layer = workflow._make_agent("assembler", prompt, required_tools=["append_assembler_progress"]) + + options = await _launched_options(monkeypatch, layer) + + spilled = list((workflow.workspace / ".agent-flow").glob("system-prompt-*.md")) + assert len(spilled) == 1 + assert spilled[0].read_text(encoding="utf-8") == prompt + pointer = options.system_prompt["append"] + assert str(spilled[0]) in pointer + assert len(pointer.encode("utf-8")) < 4096 + + +async def test_a_role_prompt_under_the_threshold_is_passed_inline(tmp_path, monkeypatch): + workflow = _workflow(tmp_path) + workflow._decision_files = True + layer = workflow._make_agent( + "auditor", workflow.prompts.auditor, required_tools=["append_auditor_progress"] + ) + + options = await _launched_options(monkeypatch, layer) + + assert options.system_prompt["append"] == workflow.prompts.auditor + assert not (workflow.workspace / ".agent-flow").exists() diff --git a/agent-flow/tests/workflows/staircase_onboard/test_worked_examples.py b/agent-flow/tests/workflows/staircase_onboard/test_worked_examples.py new file mode 100644 index 000000000000..1ef22d044cf1 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_worked_examples.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Worked examples stay examples: valid spec fragments that no role is handed. + +`examples/` holds one checkpoint's values for the generic features. The +skills, the registries and the example spec describe the features without +naming a model, because a run that onboards that model to measure the workflow +must derive those values itself. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import task_schema + +from .test_workflow import _workflow + +PACKAGE = Path(task_schema.__file__).parent +EXAMPLES = sorted((PACKAGE / "examples").glob("*.md")) +YAML_BLOCK = re.compile(r"```yaml\n(.*?)```", re.S) + +#: What a role is handed or can be pointed at by the run itself. +ROLE_FACING = ( + *sorted((PACKAGE / "skills").rglob("*.md")), + PACKAGE / "runtime_index.yaml", + PACKAGE / "kernel_sources.yaml", + PACKAGE / "task.example.yaml", + *sorted((PACKAGE / "prompts").glob("*.py")), +) + +#: Identifiers of the checkpoints the worked examples describe. +EXAMPLE_CHECKPOINT_MARKERS = ("V4.1", "DSV41", "dsv4.1", "deepseek_v41") + + +def test_there_is_a_worked_example(): + assert EXAMPLES + + +@pytest.mark.parametrize("example", EXAMPLES, ids=lambda p: p.name) +def test_every_spec_fragment_in_an_example_uses_real_keys(example): + """A fragment an operator copies must not fail validation on a misspelled key.""" + blocks = YAML_BLOCK.findall(example.read_text(encoding="utf-8")) + assert blocks, f"{example.name} has no yaml fragment" + for block in blocks: + data = yaml.safe_load(block) + assert isinstance(data, dict), block + assert set(data) <= set(task_schema.TOP_LEVEL_KEYS), block + accuracy = data.get("accuracy") or {} + assert set(accuracy) <= set(task_schema.ACCURACY_KEYS), block + for feature in data.get("features") or (): + assert feature in task_schema.FEATURE_SKILLS, feature + for feature in accuracy.get("witness") or {}: + assert feature in task_schema.FEATURE_SKILLS, feature + + +def test_a_run_materializes_skills_and_never_examples(tmp_path): + workflow = _workflow(tmp_path) + destination = workflow._materialize_skills() + copied = {path.name for path in destination.rglob("*")} + assert copied + assert not copied & {path.name for path in EXAMPLES} + + +@pytest.mark.parametrize("path", ROLE_FACING, ids=lambda p: str(p.relative_to(PACKAGE))) +def test_role_facing_material_names_no_example_checkpoint(path): + text = path.read_text(encoding="utf-8") + found = [marker for marker in EXAMPLE_CHECKPOINT_MARKERS if marker in text] + assert not found, f"{path.relative_to(PACKAGE)} names {found}; move it to examples/" From 11f29acc6bd18f1277e73164e23ff5f657d3110d Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 12 Oct 2026 00:56:07 -0700 Subject: [PATCH 04/70] [TRTLLM-16460][feat] staircase: commit helper, pre-merge pack, source-embargo hooks, unit overlap and delivery audit commit.sh commits as the repository user with a matching DCO sign-off and refuses untracked links, sandbox placeholders and paths outside the policy; the pre-merge pack (staircase-onboard premerge lint|ci|copyright|hygiene|split|resign|pr-open|pack); source_embargo enforced by PreToolUse/PostToolUse hooks with an attempt log (prompt instruction for backends without hooks); opt-in --overlap-units that starts the next unit's assembler in its own worktree (F027); and a delivery tail that audits the run's commits and writes the pre-merge pack. Squashed from user/williamj/staircase-polish @ c8109d6f90 (8 commits): - [feat] staircase: commit helper signs as the repository user and refuses links, placeholders and stray paths - [feat] staircase: pre-merge pack as `staircase-onboard premerge ` - [feat] staircase: enforce source_embargo with PreToolUse/PostToolUse hooks and log every attempt - [feat] staircase: opt-in --overlap-units starts the next unit's assembler in its own worktree (F027) - [feat] staircase: delivery audits the run's commits and ends with the pre-merge pack - [doc] staircase: name the worked example's embargo in the package: spelling - [fix] staircase: package: embargo entries stay inside the package and read as reference material - [fix] staircase: the AST proof's label for a file that does not parse is spelled "unparsable" Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../workflows/staircase_onboard/cli.py | 101 ++- .../workflows/staircase_onboard/committer.py | 682 ++++++++++++++++-- .../workflows/staircase_onboard/delivery.py | 90 +++ .../workflows/staircase_onboard/embargo.py | 569 +++++++++++++++ .../examples/deepseek-v4.1.md | 4 +- .../workflows/staircase_onboard/overlap.py | 541 ++++++++++++++ .../staircase_onboard/premerge/__init__.py | 35 + .../staircase_onboard/premerge/_git.py | 60 ++ .../staircase_onboard/premerge/ci_case.py | 337 +++++++++ .../staircase_onboard/premerge/cli.py | 369 ++++++++++ .../staircase_onboard/premerge/copyright.py | 91 +++ .../staircase_onboard/premerge/lint.py | 411 +++++++++++ .../staircase_onboard/premerge/pack.py | 159 ++++ .../staircase_onboard/premerge/pr_open.py | 173 +++++ .../staircase_onboard/premerge/resign.py | 142 ++++ .../staircase_onboard/premerge/split.py | 308 ++++++++ .../staircase_onboard/prompts/_common.py | 23 +- .../staircase_onboard/task.example.yaml | 26 +- .../staircase_onboard/task_schema.py | 19 +- .../workflows/staircase_onboard/transport.py | 10 +- .../workflows/staircase_onboard/workflow.py | 44 +- .../staircase_onboard/test_committer.py | 475 +++++++++++- .../staircase_onboard/test_delivery_tail.py | 190 +++++ .../staircase_onboard/test_dspark.py | 5 +- .../staircase_onboard/test_embargo.py | 376 ++++++++++ .../staircase_onboard/test_overlap.py | 355 +++++++++ .../staircase_onboard/test_premerge.py | 613 ++++++++++++++++ 27 files changed, 6103 insertions(+), 105 deletions(-) create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/embargo.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/overlap.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/premerge/__init__.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/premerge/_git.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/premerge/ci_case.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/premerge/cli.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/premerge/copyright.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/premerge/lint.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/premerge/pack.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/premerge/pr_open.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/premerge/resign.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/premerge/split.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_delivery_tail.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_embargo.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_overlap.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_premerge.py diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/cli.py b/agent-flow/agent_flow/workflows/staircase_onboard/cli.py index 60acd4243299..48e3753c872d 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/cli.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/cli.py @@ -5,11 +5,19 @@ import argparse import os -import subprocess import sys from pathlib import Path from .agent_home import ensure_agent_home +from .committer import ( + ENV_EMAIL, + ENV_NAME, + CommitIdentity, + CommitPolicy, + CommitPolicyError, + resolve_commit_identity, + trial_commit, +) from .roles import shares_the_drafters_model from .state import STATE_FILENAME from .task_schema import TaskSchemaError, load_and_validate_task_yaml @@ -24,6 +32,7 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: "Bring one checkpoint to a gated modeling_v2 target: " "preflight -> plan -> smith -> assemble -> deliver." ), + epilog="Other commands: `staircase-onboard premerge --help` (the pre-merge pack).", ) parser.add_argument( "--task", @@ -143,6 +152,16 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: "machine, the queue, decision validation, all three gates — stays " "here either way.", ) + parser.add_argument( + "--overlap-units", + action="store_true", + help="Start the next unit's assembler in its own git worktree as soon as the " + "current unit is UNIT_READY, while the current unit is still being gated. The " + "early result is adopted only if it ended UNIT_READY, committed cleanly, rebases " + "onto HEAD without conflict and fast-forwards the checkout; anything else falls " + "back to a sequential assembly and keeps the worktree for inspection. Off by " + "default.", + ) parser.add_argument( "--agent-turn-timeout", type=int, @@ -153,30 +172,33 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: return parser.parse_args(argv) -# The identity commits are made under. Needed because redirecting $HOME takes -# the real ~/.gitconfig off git's search path, and this checkout has no -# identity of its own -- so without it every smith commit dies with "Please -# tell me who you are", which is a new infrastructure failure traded for the -# one just fixed. Not an AI attribution: the sign-off is git's own. -def _git_config(key: str) -> str: - """The operator's own git setting, read before $HOME is redirected.""" - try: - out = subprocess.run(["git", "config", key], capture_output=True, text=True, timeout=10) - except (OSError, subprocess.SubprocessError): - return "" - return out.stdout.strip() - +def _commit_identity(task_data: dict) -> CommitIdentity: + """The identity every commit of the run is authored and signed off under. -GIT_IDENTITY_NAME = ( - os.environ.get("STAIRCASE_GIT_NAME") - or _git_config("user.name") - or os.environ.get("USER", "staircase") -) -GIT_IDENTITY_EMAIL = ( - os.environ.get("STAIRCASE_GIT_EMAIL") - or _git_config("user.email") - or f"{GIT_IDENTITY_NAME}@nvidia.com" -) + The repository's configured user, read before $HOME is redirected + (redirecting it takes the real ~/.gitconfig off git's search path), and + proven by one trial commit through the generated helper: a commit that + lands without a sign-off matching its author fails the target + repository's DCO check, and finding that out after the run has made fifty + of them is the expensive way. Exits 2 when there is no identity, or when + this environment's git changes what the helper writes. + """ + repo = str(task_data.get("trtllm_repo_path") or "") + try: + identity = resolve_commit_identity(repo) + except CommitPolicyError as exc: + print(f"error: {exc}", file=sys.stderr) + sys.exit(2) + policy = CommitPolicy.from_task(task_data, identity=identity) + problems = trial_commit(policy) + if problems: + print( + "error: a trial commit through the commit helper does not meet the commit " + f"rules under {identity.render()}:\n - " + "\n - ".join(problems), + file=sys.stderr, + ) + sys.exit(2) + return identity #: Open file objects whose flock this process holds, kept alive deliberately. @@ -227,8 +249,25 @@ def _claim_workspace(workspace: Path): return handle +def _premerge(argv: list[str]) -> int: + from .premerge.cli import main as premerge_main + + return premerge_main(argv) + + +#: ``staircase-onboard ...`` for the tools that are not a run. +#: A run is still ``staircase-onboard --task ...``; a first argument named +#: here dispatches instead, and everything after it is the subcommand's own. +SUBCOMMANDS = { + "premerge": _premerge, +} + + def main(argv: list[str] | None = None) -> None: """Run the staircase-onboard workflow as a CLI.""" + argv = list(sys.argv[1:] if argv is None else argv) + if argv and argv[0] in SUBCOMMANDS: + sys.exit(SUBCOMMANDS[argv[0]](argv[1:])) args = _parse_args(argv) # Validate at the CLI boundary: an invalid spec must abort before any @@ -241,6 +280,10 @@ def main(argv: list[str] | None = None) -> None: sys.exit(2) if warning := shares_the_drafters_model(task_data): print(f"warning: {warning}", file=sys.stderr) + identity = _commit_identity(task_data) + # Pinned for the rest of the process, so the workflow resolves the same + # identity after $HOME has been redirected below. + os.environ[ENV_NAME], os.environ[ENV_EMAIL] = identity.name, identity.email # A sandboxed session's $HOME is read-only, and the first thing a spawned # `claude` CLI does is take a lock beside its credential store. Without @@ -252,7 +295,7 @@ def main(argv: list[str] | None = None) -> None: staged = ensure_agent_home( Path(args.workspace), - git_identity=(GIT_IDENTITY_NAME, GIT_IDENTITY_EMAIL), + git_identity=(identity.name, identity.email), ) if staged is not None: staged.apply() @@ -260,7 +303,12 @@ def main(argv: list[str] | None = None) -> None: transport = None if args.agent_transport == "external": - transport = ExternalAgentTransport(args.workspace, timeout_s=args.agent_turn_timeout) + transport = ExternalAgentTransport( + args.workspace, + timeout_s=args.agent_turn_timeout, + source_embargo=task_data.get("source_embargo") or (), + repo_root=task_data.get("trtllm_repo_path"), + ) print( f"→ role turns go through {transport.root}; a driver session must " f"answer each `*.request.json` with a `*.result.md` and a `*.done`" @@ -280,6 +328,7 @@ def main(argv: list[str] | None = None) -> None: dry_run=args.dry_run, skip_substrate_gate=args.skip_substrate_gate, agent_transport=transport, + overlap_units=args.overlap_units, ) as workflow: workflow.run(args.task) diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/committer.py b/agent-flow/agent_flow/workflows/staircase_onboard/committer.py index bed1898dc8f8..07dd19bbefbc 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/committer.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/committer.py @@ -23,12 +23,34 @@ 1. **Serialized.** An ``flock`` around the whole stage-and-commit, with a bounded wait, so concurrency becomes a short queue rather than an error. -2. **Scoped.** ``git commit -- `` commits the working-tree content of - exactly those paths on top of ``HEAD``, whatever else sits in the index. -3. **Verified.** After committing it diffs the new commit against its parent - and **fails** if the file list is not exactly what was asked for. This is - the part that matters: misattribution stops being something a reviewer - notices later and becomes something the tool refuses to do. +2. **Scoped.** Only explicit, repo-relative file paths are accepted: no + directory, no glob, no absolute path, no ``..``. ``git commit -- `` + then commits the working-tree content of exactly those paths on top of + ``HEAD``, whatever else sits in the index. +3. **Filtered.** A path is refused before anything is staged when it is + + - a symlink git does not track, or lies beneath one. A worktree links the + main checkout's native build (``tensorrt_llm/bindings`` and friends); + ``.gitignore`` names the directory form only, so the link form shows up + as untracked and one ``git add -A`` committed it; + - a sandbox placeholder: a ``.mcp.json`` or anything under a ``.claude/`` + or ``.cc-writes/`` directory, which the operator's sandbox leaves in + every directory its shell enters (one was committed); + - outside the policy's allowed path prefixes, or inside a denied one (the + run's own workspace, when it sits inside the repository). + +4. **Signed as the repository's user.** The author, the committer and the + ``Signed-off-by`` trailer are one identity: the repository's configured + ``user.name`` / ``user.email``, resolved when the script is installed and + baked into it. A commit whose sign-off email differs from its author + email fails the DCO check, and a re-sign under the operator's identity + produced exactly that for 45 commits once. +5. **Verified.** After committing it checks the new commit against its + parent and **undoes it** unless the file list is exactly what was asked + for, no symlink was added, every ``Signed-off-by`` carries the author's + email, and no ``Co-authored-by`` trailer is present (unless the policy + allows one). Misattribution stops being something a reviewer notices later + and becomes something the tool refuses to do. A private ``GIT_INDEX_FILE`` per dispatch was the other candidate and is a trap. It does stop one agent's ``add`` from seeing another's staged files, @@ -42,7 +64,11 @@ import os import shlex import stat +import subprocess +import tempfile +from dataclasses import dataclass from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence #: How long an agent waits for the commit lock before giving up. Generous #: because the section is milliseconds of git and the alternative — an agent @@ -51,47 +77,341 @@ SCRIPT_NAME = "commit.sh" +#: The task.yaml block that configures the policy. Optional; every key has a +#: default. +COMMIT_POLICY_FIELD = "commit_policy" +COMMIT_POLICY_KEYS: tuple[str, ...] = ("allowed_paths", "allow_co_authors", "base_ref") + +#: Where an onboarding writes, as repo-relative prefixes. The target, the +#: catalog and its tests live under ``tensorrt_llm/`` and ``tests/``; a new +#: kernel may need ``cpp/``; a deployment example ``examples/``. Everything +#: else in the repository is somebody else's. +DEFAULT_ALLOWED_PATHS: tuple[str, ...] = ("tensorrt_llm/", "tests/", "cpp/", "examples/") + +#: File names the operator's sandbox drops into every directory its shell +#: enters, and directory names it creates there (F002). +SANDBOX_PLACEHOLDER_NAMES: tuple[str, ...] = (".mcp.json",) +SANDBOX_PLACEHOLDER_DIRS: tuple[str, ...] = (".claude", ".cc-writes") + +#: Exit codes. 2 is a refusal before anything was staged, 1 a commit that was +#: made and undone (or failed), 75 a lock that was never granted. +EXIT_REFUSED = 2 +EXIT_UNDONE = 1 +EXIT_LOCK_TIMEOUT = 75 + +ENV_NAME = "STAIRCASE_GIT_NAME" +ENV_EMAIL = "STAIRCASE_GIT_EMAIL" + + +class CommitPolicyError(ValueError): + """The commit policy cannot be established, so nothing may commit.""" + + +@dataclass(frozen=True) +class CommitIdentity: + """Who authors, commits and signs off every commit the run makes.""" + + name: str + email: str + #: Where the identity came from, for the preflight report. + source: str = "" + + def render(self) -> str: + return f"{self.name} <{self.email}>" + + +@dataclass(frozen=True) +class CommitPolicy: + """What the commit helper accepts, and as whom it commits. + + ``identity`` is ``None`` only where none could be resolved when the + script was installed. The script then resolves it from the repository's + git configuration at commit time and **refuses** if there is none — it + never invents one from ``$USER``. + + An empty ``allowed_paths`` means "anywhere in the repository"; the + placeholder and symlink refusals apply either way. A run's policy comes + from :meth:`from_task`, which defaults to :data:`DEFAULT_ALLOWED_PATHS`. + """ + + identity: CommitIdentity | None = None + allowed_paths: tuple[str, ...] = () + denied_paths: tuple[str, ...] = () + allow_co_authors: bool = False + base_ref: str = "" + + @classmethod + def from_task( + cls, + task_data: Mapping[str, Any], + *, + workspace: str | Path | None = None, + identity: CommitIdentity | None = None, + ) -> "CommitPolicy": + """The policy a run's ``task.yaml`` asks for. + + The identity is resolved from the task's repository unless one is + given. A workspace inside the repository is denied: the run's own + records are never part of a delivery. + """ + repo = str(task_data.get("trtllm_repo_path") or "") + block = task_data.get(COMMIT_POLICY_FIELD) or {} + if identity is None and repo: + try: + identity = resolve_commit_identity(repo) + except CommitPolicyError: + identity = None + allowed = block.get("allowed_paths") + denied: list[str] = [] + if workspace is not None and repo: + rel = _relative_inside(Path(workspace), Path(repo)) + if rel: + denied.append(rel.rstrip("/") + "/") + return cls( + identity=identity, + allowed_paths=( + DEFAULT_ALLOWED_PATHS + if allowed is None + else tuple(_as_prefix(str(p)) for p in allowed) + ), + denied_paths=tuple(denied), + allow_co_authors=bool(block.get("allow_co_authors", False)), + base_ref=str(block.get("base_ref") or ""), + ) + + +def validate_commit_policy_block(block: Any, errors: list[str]) -> None: + """Validate ``task.yaml``'s optional ``commit_policy`` block. + + Problems are appended to ``errors`` so the schema can report every one + of them at once. + """ + if block is None: + return + if not isinstance(block, dict): + errors.append(f"'{COMMIT_POLICY_FIELD}' must be a mapping, got {type(block).__name__}") + return + unknown = sorted(set(block) - set(COMMIT_POLICY_KEYS)) + if unknown: + errors.append( + f"'{COMMIT_POLICY_FIELD}' has unknown key(s) {unknown}; known keys are " + f"{list(COMMIT_POLICY_KEYS)}" + ) + allowed = block.get("allowed_paths") + if allowed is not None: + if not isinstance(allowed, list) or not all(isinstance(p, str) for p in allowed): + errors.append(f"'{COMMIT_POLICY_FIELD}.allowed_paths' must be a list of strings") + else: + for entry in allowed: + problem = _path_problem(entry) + if problem: + errors.append( + f"'{COMMIT_POLICY_FIELD}.allowed_paths' entry {entry!r}: {problem}" + ) + if "allow_co_authors" in block and not isinstance(block["allow_co_authors"], bool): + errors.append(f"'{COMMIT_POLICY_FIELD}.allow_co_authors' must be true or false") + if "base_ref" in block and not isinstance(block["base_ref"], str): + errors.append(f"'{COMMIT_POLICY_FIELD}.base_ref' must be a string (a git revision)") + + +def _path_problem(entry: str) -> str: + if not entry.strip(): + return "is empty" + if entry.startswith("/"): + return "must be repo-relative" + if any(part in ("..", ".") for part in entry.strip("/").split("/")): + return "must not contain '.' or '..' components" + return "" + + +def _as_prefix(entry: str) -> str: + """A repo-relative prefix: a bare file name stays one; a directory gets its slash.""" + entry = entry.strip() + return entry if entry.endswith("/") or "." in entry.rsplit("/", 1)[-1] else entry + "/" + + +def _relative_inside(path: Path, root: Path) -> str: + """``path`` relative to ``root`` when it lies inside it, else ``""``.""" + try: + rel = path.resolve().relative_to(root.resolve()) + except ValueError: + return "" + text = rel.as_posix() + return "" if text in ("", ".") else text + + +def _git_config(repo: str | Path, key: str) -> str: + try: + out = subprocess.run( + ["git", "-C", str(repo), "config", "--get", key], + capture_output=True, + text=True, + timeout=30, + ) + except (OSError, subprocess.SubprocessError): + return "" + return out.stdout.strip() if out.returncode == 0 else "" + + +def resolve_commit_identity( + repo: str | Path, *, environ: Mapping[str, str] | None = None +) -> CommitIdentity: + """The repository's configured user, or the operator's explicit override. + + ``STAIRCASE_GIT_NAME`` / ``STAIRCASE_GIT_EMAIL`` win when **both** are + set: an operator who sets one means to set the identity, and half an + override is how author and sign-off came apart before. Otherwise the + repository's own ``git config`` answers (local, then global — git's + order). There is deliberately no ``$USER`` fallback: a login name with a + guessed domain is an identity no DCO check will match. + """ + env = os.environ if environ is None else environ + name, email = env.get(ENV_NAME, "").strip(), env.get(ENV_EMAIL, "").strip() + if name and email: + return CommitIdentity(name=name, email=email, source=f"${ENV_NAME} / ${ENV_EMAIL}") + name, email = _git_config(repo, "user.name"), _git_config(repo, "user.email") + if name and email: + return CommitIdentity(name=name, email=email, source=f"git config of {repo}") + raise CommitPolicyError( + f"no commit identity for {repo}: git config user.name={name!r}, user.email={email!r}. " + f"Every commit the run makes is authored and signed off as the repository's user, " + f"so set `git -C {repo} config user.name/user.email` (or export both {ENV_NAME} and " + f"{ENV_EMAIL}) before starting the run." + ) + + _SCRIPT = """#!/usr/bin/env bash # Generated by the staircase-onboard orchestrator. Do not edit. # # usage: {name} [path...] # -# Commits exactly the named paths, under a lock, and verifies afterwards that -# the commit contains exactly those paths and nothing else. +# Commits exactly the named files, under a lock, as the repository's user +# with a DCO sign-off, and verifies afterwards that the commit contains +# exactly those files and carries exactly the trailers the repository wants. set -uo pipefail REPO={repo} LOCK={lock} TIMEOUT={timeout} +# The identity is baked when the script is installed. Empty means it could +# not be resolved then; it is resolved from the repository's git config at +# commit time instead, and the commit is refused if there is none. +COMMIT_NAME={commit_name} +COMMIT_EMAIL={commit_email} +ALLOW_CO_AUTHORS={allow_co_authors} +# Repo-relative prefixes a commit may touch (empty: anywhere in the repo), +# and prefixes it may never touch. +ALLOWED_PATHS=({allowed_paths}) +DENIED_PATHS=({denied_paths}) + +refuse() {{ + echo "commit.sh: REFUSED: $*" >&2 + echo " Nothing was staged or committed. Fix the call and run it again." >&2 + exit {exit_refused} +}} if [ "$#" -lt 2 ]; then echo "usage: $0 [path...]" >&2 - exit 2 + exit {exit_refused} fi MESSAGE="$1"; shift -for path in "$@"; do - if [ ! -e "$REPO/$path" ] && [ ! -e "$path" ]; then - echo "commit.sh: $path does not exist; commit by path relative to the repo root" >&2 - exit 2 +if [ "$ALLOW_CO_AUTHORS" != 1 ] && printf '%s\\n' "$MESSAGE" | grep -qiE '^[[:space:]]*co-authored-by:'; then + refuse "the message carries a Co-authored-by trailer; commits in this repository carry none" +fi +if printf '%s\\n' "$MESSAGE" | grep -qiE '^[[:space:]]*signed-off-by:'; then + refuse "do not write a Signed-off-by line; this script adds it as the repository's user" +fi + +tracked() {{ + git -C "$REPO" ls-files --error-unmatch -- "$1" >/dev/null 2>&1 +}} + +tracked_symlink() {{ + local line + line=$(git -C "$REPO" ls-files -s -- "$1" | head -n 1) + [ "${{line%% *}}" = 120000 ] && [ "${{line#*$'\\t'}}" = "$1" ] +}} + +check_path() {{ + local p="$1" acc="" part base + local -a parts + case "$p" in + "") refuse "an empty path" ;; + /*) refuse "$p is absolute; name files relative to the repo root ($REPO)" ;; + *'*'*|*'?'*|*'['*) refuse "$p is a glob; name every file explicitly" ;; + esac + case "/$p/" in + */../*|*/./*|*//*) refuse "$p is not a plain repo-relative path (no '.', '..' or '//')" ;; + esac + base="${{p##*/}}" + case "$base" in + {placeholder_names}) refuse "$p is a sandbox placeholder file, never part of a commit" ;; + esac + case "/$p/" in + {placeholder_dirs}) refuse "$p lies in a sandbox placeholder directory, never part of a commit" ;; + esac + IFS='/' read -r -a parts <<< "$p" + for part in "${{parts[@]}}"; do + acc="${{acc:+$acc/}}$part" + if [ -L "$REPO/$acc" ] && ! tracked_symlink "$acc"; then + refuse "$acc is a symlink git does not track (a worktree's link to a build output," \\ + "typically); it is never committed. Name real files only, and do not stage the link" + fi + done + if [ -d "$REPO/$p" ] && [ ! -L "$REPO/$p" ]; then + refuse "$p is a directory; name each file in it explicitly" + fi + if [ ! -e "$REPO/$p" ] && [ ! -L "$REPO/$p" ] && ! tracked "$p"; then + echo "commit.sh: $p does not exist; commit by path relative to the repo root" >&2 + exit {exit_refused} fi + local d + for d in "${{DENIED_PATHS[@]}}"; do + case "$p" in + "$d"*) refuse "$p is under $d, which this run never commits (its own workspace)" ;; + esac + done + if [ "${{#ALLOWED_PATHS[@]}}" -gt 0 ]; then + local a ok=0 + for a in "${{ALLOWED_PATHS[@]}}"; do + case "$p" in + "$a"*) ok=1 ;; + esac + done + if [ "$ok" != 1 ]; then + refuse "$p is outside the paths this run may commit: ${{ALLOWED_PATHS[*]}}" + fi + fi +}} + +for path in "$@"; do + check_path "$path" done +if [ -z "$COMMIT_NAME" ]; then COMMIT_NAME=$(git -C "$REPO" config --get user.name 2>/dev/null); fi +if [ -z "$COMMIT_EMAIL" ]; then COMMIT_EMAIL=$(git -C "$REPO" config --get user.email 2>/dev/null); fi +if [ -z "$COMMIT_NAME" ] || [ -z "$COMMIT_EMAIL" ]; then + refuse "no commit identity: the repository has no user.name/user.email and none was baked in" +fi + exec {fd}>"$LOCK" if ! flock -w "$TIMEOUT" {fd}; then echo "commit.sh: waited ${{TIMEOUT}}s for the commit lock and did not get it." >&2 echo "Another agent is committing. This is a symptom to report, not to work around." >&2 - exit 75 + exit {exit_lock} fi cd "$REPO" || exit 1 -# The identity git would use by itself comes before $USER: a login name is -# not a name a DCO sign-off can carry. -COMMIT_NAME=${{STAIRCASE_GIT_NAME:-$(git -C "$REPO" config user.name 2>/dev/null)}} -COMMIT_NAME=${{COMMIT_NAME:-${{USER:-staircase}}}} -COMMIT_EMAIL=${{STAIRCASE_GIT_EMAIL:-$(git -C "$REPO" config user.email 2>/dev/null)}} -COMMIT_EMAIL=${{COMMIT_EMAIL:-${{COMMIT_NAME}}@nvidia.com}} +# One identity for author, committer and sign-off. Exported rather than set +# with `-c user.*`, because an inherited GIT_AUTHOR_* / GIT_COMMITTER_* beats +# `-c`, and a sign-off is the committer's identity: an operator environment +# that exports either would split them. Never written to config: a redirected +# $HOME takes the real ~/.gitconfig off git's search path, and a stale +# `.git/config.lock` on Lustre can refuse `git config --local` outright. +export GIT_AUTHOR_NAME="$COMMIT_NAME" GIT_AUTHOR_EMAIL="$COMMIT_EMAIL" +export GIT_COMMITTER_NAME="$COMMIT_NAME" GIT_COMMITTER_EMAIL="$COMMIT_EMAIL" # `git add` so a new file is known to the index; `git commit -- ` so # the commit carries the working-tree content of exactly these paths on top @@ -100,31 +420,50 @@ echo "commit.sh: git add failed" >&2 exit 1 fi -# Identity passed on the command line, never written to config. `-s` adds -# the DCO sign-off the target repository requires of every commit. -# -# Two independent reasons, both measured. A redirected $HOME (see -# agent_home.py) takes the real ~/.gitconfig off git's search path; and this -# checkout cannot accept `git config --local` at all, because a stale -# `.git/config.lock` on Lustre refuses both `rm` ("Device or resource busy") -# and `chmod` ("Read-only file system"). Without this every commit dies with -# -# fatal: unable to auto-detect email address (got 'user@host.(none)') -# -# which commit.sh reported as "nothing to commit is the usual cause" --- so -# each of seven concurrent smiths rediscovered it separately and worked -# around it by hand. -if ! git -c user.name="$COMMIT_NAME" -c user.email="$COMMIT_EMAIL" \ +if ! git -c user.name="$COMMIT_NAME" -c user.email="$COMMIT_EMAIL" \\ commit -s -q -m "$MESSAGE" -- "$@"; then echo "commit.sh: git commit failed (nothing to commit is the usual cause)" >&2 exit 1 fi SHA=$(git rev-parse HEAD) -ACTUAL=$(git diff-tree --no-commit-id --name-only -r "$SHA" | LC_ALL=C sort) -WANTED=$(printf '%s\\n' "$@" | LC_ALL=C sort) +ACTUAL=$(git diff-tree --root --no-commit-id --name-only -r "$SHA" | LC_ALL=C sort) +WANTED=$(printf '%s\\n' "$@" | LC_ALL=C sort -u) +PROBLEMS=() if [ "$ACTUAL" != "$WANTED" ]; then - echo "commit.sh: REFUSING TO LEAVE THIS COMMIT AS IS." >&2 + PROBLEMS+=("the commit contains [$(echo $ACTUAL)], not exactly the paths asked for [$(echo $WANTED)]") +fi +NEW_LINKS=$(git diff-tree --root --no-commit-id -r --raw "$SHA" \\ + | awk '$2 == "120000" && $1 != ":120000" {{ sub(/^[^\\t]*\\t/, ""); print }}') +if [ -n "$NEW_LINKS" ]; then + PROBLEMS+=("the commit adds a symlink: $(echo $NEW_LINKS)") +fi +AUTHOR_EMAIL=$(git log -1 --format=%ae "$SHA") +COMMITTER_EMAIL=$(git log -1 --format=%ce "$SHA") +SIGNOFFS=$(git log -1 --format='%(trailers:key=Signed-off-by,valueonly)' "$SHA" | sed '/^$/d') +if [ -z "$SIGNOFFS" ]; then + PROBLEMS+=("the commit has no Signed-off-by trailer") +fi +while IFS= read -r line; do + [ -z "$line" ] && continue + case "$line" in + *"<$AUTHOR_EMAIL>") ;; + *) PROBLEMS+=("sign-off '$line' does not carry the author email <$AUTHOR_EMAIL>") ;; + esac +done <<< "$SIGNOFFS" +if [ "$AUTHOR_EMAIL" != "$COMMIT_EMAIL" ] || [ "$COMMITTER_EMAIL" != "$COMMIT_EMAIL" ]; then + PROBLEMS+=("author <$AUTHOR_EMAIL> / committer <$COMMITTER_EMAIL> is not the repository user <$COMMIT_EMAIL>") +fi +if [ "$ALLOW_CO_AUTHORS" != 1 ]; then + CO=$(git log -1 --format='%(trailers:key=Co-authored-by,valueonly)' "$SHA" | sed '/^$/d') + if [ -n "$CO" ]; then + PROBLEMS+=("the commit carries Co-authored-by: $(echo $CO)") + fi +fi + +if [ "${{#PROBLEMS[@]}}" -gt 0 ]; then + echo "commit.sh: REFUSING TO LEAVE THIS COMMIT AS IS ($SHA)." >&2 + printf ' - %s\\n' "${{PROBLEMS[@]}}" >&2 echo " asked for:" >&2 printf ' %s\\n' $WANTED >&2 echo " commit $SHA actually contains:" >&2 @@ -134,11 +473,16 @@ # runs it — and then the explicit reset below runs it a second time, # rewinding a commit that had nothing to do with this one. echo ' Undoing it with: git reset --soft HEAD^' >&2 - echo " Your work stays in the tree. Report this: it means a path you named" >&2 - echo " matched more than you meant, or another agent's file was already" >&2 - echo " staged for one of them." >&2 - git reset --soft HEAD^ - exit 1 + echo " Your work stays in the tree. Report this: a path you named matched more" >&2 + echo " than you meant, another agent's file was already staged, or a hook or" >&2 + echo " the environment changed what git wrote." >&2 + if git rev-parse -q --verify "$SHA^" >/dev/null; then + git reset -q --soft HEAD^ + else + # A root commit has no parent to reset to; the branch goes back to unborn. + git update-ref -d HEAD + fi + exit {exit_undone} fi echo "commit.sh: $SHA" @@ -146,17 +490,43 @@ """ -def render_script(repo: str | Path, lock: str | Path, *, timeout: int = LOCK_TIMEOUT_S) -> str: - """The script body, with this run's repo and lock baked in. +def _bash_words(values: Iterable[str]) -> str: + return " ".join(shlex.quote(v) for v in values) + + +def _case_alternatives(patterns: Iterable[str]) -> str: + return "|".join(patterns) + + +def render_script( + repo: str | Path, + lock: str | Path, + *, + timeout: int = LOCK_TIMEOUT_S, + policy: CommitPolicy | None = None, +) -> str: + """The script body, with this run's repo, lock and policy baked in. Baked rather than passed so an agent has no argument it can get wrong, and so the same call works from any working directory. """ + policy = policy or CommitPolicy() + identity = policy.identity return _SCRIPT.format( name=SCRIPT_NAME, repo=shlex.quote(str(repo)), lock=shlex.quote(str(lock)), timeout=int(timeout), + commit_name=shlex.quote(identity.name if identity else ""), + commit_email=shlex.quote(identity.email if identity else ""), + allow_co_authors=1 if policy.allow_co_authors else 0, + allowed_paths=_bash_words(policy.allowed_paths), + denied_paths=_bash_words(policy.denied_paths), + placeholder_names=_case_alternatives(SANDBOX_PLACEHOLDER_NAMES), + placeholder_dirs=_case_alternatives(f"*/{d}/*" for d in SANDBOX_PLACEHOLDER_DIRS), + exit_refused=EXIT_REFUSED, + exit_undone=EXIT_UNDONE, + exit_lock=EXIT_LOCK_TIMEOUT, # An explicit high fd, so the redirect cannot collide with a # descriptor the surrounding tool already holds open. fd=9, @@ -164,7 +534,11 @@ def render_script(repo: str | Path, lock: str | Path, *, timeout: int = LOCK_TIM def install_committer( - workspace: str | Path, repo: str | Path, *, timeout: int = LOCK_TIMEOUT_S + workspace: str | Path, + repo: str | Path, + *, + timeout: int = LOCK_TIMEOUT_S, + policy: CommitPolicy | None = None, ) -> Path: """Write the commit helper into the workspace and return its path.""" workspace = Path(workspace) @@ -172,12 +546,194 @@ def install_committer( bin_dir.mkdir(parents=True, exist_ok=True) lock = workspace / ".commit.lock" path = bin_dir / SCRIPT_NAME - path.write_text(render_script(repo, lock, timeout=timeout), encoding="utf-8") + path.write_text(render_script(repo, lock, timeout=timeout, policy=policy), encoding="utf-8") path.chmod(path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) lock.touch(exist_ok=True) return path +def trial_commit(policy: CommitPolicy, *, scratch: str | Path | None = None) -> list[str]: + """Commit once through the helper in a throwaway repository; return the problems. + + The check F022 asked for: not "the template passes ``-s``" but "a commit + made through the helper, in this environment, carries the trailers the + target repository requires". The environment is the part a template + cannot vouch for — a global ``core.hooksPath`` whose ``commit-msg`` hook + appends a trailer, an exported ``GIT_COMMITTER_EMAIL``, a + ``format.signOff`` setting — and every one of those changes what lands. + + Returns an empty list when the commit is signed off by its author, the + author is the policy's identity, and no ``Co-authored-by`` slipped in. + """ + if policy.identity is None: + return ["no commit identity could be resolved; the run cannot sign its commits"] + with tempfile.TemporaryDirectory(prefix="staircase-commit-trial-", dir=scratch) as tmp: + root = Path(tmp) + repo = root / "repo" + repo.mkdir() + init = subprocess.run( + ["git", "-C", str(repo), "init", "-q"], capture_output=True, text=True + ) + if init.returncode != 0: + return [f"git init failed: {init.stderr.strip()}"] + (repo / "trial.txt").write_text("trial\n", encoding="utf-8") + trial_policy = CommitPolicy( + identity=policy.identity, + allowed_paths=(), + allow_co_authors=policy.allow_co_authors, + ) + script = install_committer(root / "ws", repo, policy=trial_policy) + result = subprocess.run( + [str(script), "[None][chore] staircase commit-helper trial", "trial.txt"], + capture_output=True, + text=True, + timeout=120, + ) + if result.returncode != 0: + return [f"the trial commit failed (rc {result.returncode}): {result.stderr.strip()}"] + findings = audit_commits( + repo, "HEAD", allow_co_authors=policy.allow_co_authors, include_root=True + ) + problems = [f.problem for f in findings] + author = subprocess.run( + ["git", "-C", str(repo), "log", "-1", "--format=%an <%ae>"], + capture_output=True, + text=True, + ).stdout.strip() + if author != policy.identity.render(): + problems.append( + f"the trial commit was authored {author!r}, not {policy.identity.render()!r}" + ) + return problems + + +@dataclass(frozen=True) +class CommitFinding: + """One way one commit falls short of the repository's commit rules.""" + + sha: str + subject: str + problem: str + + def render(self) -> str: + return f"{self.sha[:10]} {self.subject[:70]}: {self.problem}" + + +def _email_of(trailer_value: str) -> str: + value = trailer_value.strip() + if "<" in value and value.endswith(">"): + return value.rsplit("<", 1)[1][:-1].strip() + return "" + + +def audit_commits( + repo: str | Path, + rev_range: str, + *, + allow_co_authors: bool = False, + include_root: bool = False, + paths: Sequence[str] = (), +) -> list[CommitFinding]: + """Check every non-merge commit in ``rev_range`` against the commit rules. + + The same rules ``commit.sh`` enforces one commit at a time, applied to a + whole branch, because not every commit goes through ``commit.sh``: the + operator's own commits, a cherry-picked fix, a rebase that dropped a + trailer. Per commit: + + - a ``Signed-off-by`` carrying the author's email exists (what the DCO + check matches; an additional sign-off is not a violation here, though + ``commit.sh`` itself only ever writes one); + - no ``Co-authored-by``, unless allowed; + - no symlink is added, and no sandbox placeholder file is added. + + ``rev_range`` is anything ``git log`` accepts (``base..HEAD``, or a single + revision with ``include_root=True`` for "this commit and its ancestors"). + ``paths`` limits the audit to commits touching those repo-relative + prefixes — a run's own commits, on a branch that also carries other work. + """ + fmt = ( + "%x1e%H%x1f%ae%x1f%s" + "%x1f%(trailers:key=Signed-off-by,valueonly,separator=%x1d)" + "%x1f%(trailers:key=Co-authored-by,valueonly,separator=%x1d)%x1f" + ) + args = [ + "git", + "-C", + str(repo), + "log", + "--no-merges", + "--raw", + "--no-renames", + f"--format={fmt}", + ] + if not include_root and ".." not in rev_range: + args.append("-1") + args.append(rev_range) + if paths: + args += ["--", *(p.rstrip("/") for p in paths)] + out = subprocess.run(args, capture_output=True, text=True) + if out.returncode != 0: + raise CommitPolicyError(f"git log {rev_range} failed: {out.stderr.strip()}") + findings: list[CommitFinding] = [] + for record in out.stdout.split("\x1e"): + if not record.strip(): + continue + fields = record.split("\x1f") + if len(fields) < 6: + continue + sha, author_email, subject, signoffs, coauthors, raw = fields[:6] + signed = [s for s in signoffs.split("\x1d") if s.strip()] + if not signed: + findings.append(CommitFinding(sha, subject, "no Signed-off-by")) + elif not any(_email_of(v).lower() == author_email.strip().lower() for v in signed): + found = ", ".join(repr(v.strip()) for v in signed) + findings.append( + CommitFinding( + sha, + subject, + f"no sign-off carries the author email <{author_email}> (found {found})", + ) + ) + if not allow_co_authors and any(c.strip() for c in coauthors.split("\x1d")): + findings.append(CommitFinding(sha, subject, "carries Co-authored-by")) + for line in raw.splitlines(): + if not line.startswith(":") or "\t" not in line: + continue + meta, path = line.split("\t", 1) + modes = meta.split() + if len(modes) >= 2 and modes[1] == "120000" and modes[0] != ":120000": + findings.append(CommitFinding(sha, subject, f"adds a symlink: {path}")) + if is_sandbox_placeholder(path) and not (len(modes) >= 2 and modes[1] == "000000"): + findings.append(CommitFinding(sha, subject, f"adds a sandbox placeholder: {path}")) + return findings + + +def is_sandbox_placeholder(path: str) -> bool: + """Whether ``path`` (repo-relative) is a file the operator's sandbox leaves behind.""" + parts = path.strip("/").split("/") + return parts[-1] in SANDBOX_PLACEHOLDER_NAMES or any( + part in SANDBOX_PLACEHOLDER_DIRS for part in parts[:-1] + ) + + +def resolve_base(repo: str | Path, candidates: Sequence[str]) -> str: + """The merge base of ``HEAD`` with the first candidate ref that resolves, or ``""``.""" + for ref in candidates: + if not ref: + continue + out = subprocess.run( + ["git", "-C", str(repo), "merge-base", "HEAD", ref], capture_output=True, text=True + ) + if out.returncode == 0 and out.stdout.strip(): + return out.stdout.strip() + return "" + + +#: Where a branch's commits are counted from when the policy names no base. +DEFAULT_BASE_REFS: tuple[str, ...] = ("origin/main", "upstream/main", "main") + + def commit_instruction(script: str | Path) -> str: """The block every committing role's prompt carries.""" return ( @@ -198,11 +754,18 @@ def commit_instruction(script: str | Path) -> str: "file of your unit, no directories, no globs that could match a file you " "did not write.\n" "\n" + "It refuses, before staging anything: a directory, a glob, an absolute path; " + "a symlink git does not track (a worktree's link to the native build); a " + "sandbox placeholder (`.mcp.json`, anything under `.claude/`); and any path " + "outside the ones this run may commit. It commits and signs off as the " + "repository's configured user — do not write a `Signed-off-by` line, and " + "write no `Co-authored-by` line.\n" + "\n" "It takes a lock, so you may wait behind another agent; that is normal and " "is not an error. It then checks that the commit contains exactly the paths " - "you named, and **undoes the commit** if it does not, leaving your work in " - "the tree. If that happens, report it rather than retrying: it means a path " - "you named swept in more than you meant.\n" + "you named and carries the sign-off, and **undoes the commit** if it does " + "not, leaving your work in the tree. If that happens, report it rather than " + "retrying: it means a path you named swept in more than you meant.\n" "\n" "One commit per unit of work, so attribution and revert stay per-unit. No " "other git command that writes — never `push`, never a branch operation, " @@ -234,10 +797,29 @@ def lock_is_held(workspace: str | Path) -> bool: __all__ = [ + "COMMIT_POLICY_FIELD", + "COMMIT_POLICY_KEYS", + "CommitFinding", + "CommitIdentity", + "CommitPolicy", + "CommitPolicyError", + "DEFAULT_ALLOWED_PATHS", + "DEFAULT_BASE_REFS", + "EXIT_LOCK_TIMEOUT", + "EXIT_REFUSED", + "EXIT_UNDONE", "LOCK_TIMEOUT_S", + "SANDBOX_PLACEHOLDER_DIRS", + "SANDBOX_PLACEHOLDER_NAMES", "SCRIPT_NAME", + "audit_commits", "commit_instruction", "install_committer", + "is_sandbox_placeholder", "lock_is_held", "render_script", + "resolve_base", + "resolve_commit_identity", + "trial_commit", + "validate_commit_policy_block", ] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/delivery.py b/agent-flow/agent_flow/workflows/staircase_onboard/delivery.py index e92aefc40819..2201a2b36140 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/delivery.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/delivery.py @@ -39,6 +39,14 @@ import yaml +from .committer import ( + DEFAULT_BASE_REFS, + CommitPolicy, + CommitPolicyError, + audit_commits, + resolve_base, +) + MODELING_V2_PACKAGE = "tensorrt_llm/_torch/_experimental/modeling_v2" CATALOG_DIR = f"{MODELING_V2_PACKAGE}/catalog" CATALOG_TESTS_DIR = "tests/unittest/_torch/modeling_v2" @@ -391,6 +399,83 @@ def check_tree_committed(repo_root: Path, paths: Sequence[str]) -> list[Check]: ] +def check_commit_hygiene(repo_root: Path, policy: CommitPolicy) -> list[Check]: + """The run's commits meet the commit rules ``commit.sh`` enforces one at a time. + + Not every commit on a delivered branch went through ``commit.sh`` — the + operator's own fixes, a cherry-pick, a rebase that dropped a trailer — and + CI's DCO check reads them all. So the same audit runs over every commit + since the merge base that touches the paths the run may commit + (``commit_policy.allowed_paths``): a sign-off carrying the author's email, + no ``Co-authored-by`` unless allowed, no symlink, no sandbox placeholder. + + Returns no check outside a git work tree, or when no base resolves (a + fixture is not a delivery, and a branch with no base has nothing to + compare); the pre-merge pack says which. + """ + if not (repo_root / ".git").exists(): + return [] + base = resolve_base(repo_root, [policy.base_ref] if policy.base_ref else DEFAULT_BASE_REFS) + if not base: + return [] + try: + findings = audit_commits( + repo_root, + f"{base}..HEAD", + allow_co_authors=policy.allow_co_authors, + paths=policy.allowed_paths, + ) + except CommitPolicyError as exc: + return [Check(name="git:commit-hygiene", passed=False, detail=str(exc))] + return [ + Check( + name="git:commit-hygiene", + passed=not findings, + detail=( + f"every commit since {base[:10]} under {', '.join(policy.allowed_paths) or 'the repo'} " + f"is signed off by its author, with no co-author, symlink or placeholder" + if not findings + else f"{len(findings)} finding(s); `staircase-onboard premerge resign` adds missing " + f"sign-offs on a new branch:\n " + "\n ".join(f.render() for f in findings[:12]) + ), + ) + ] + + +def write_premerge_pack( + repo_root: Path | str, + reports_dir: Path | str, + *, + family: str, + policy: CommitPolicy | None = None, +) -> tuple[Path, list[str]]: + """The delivery tail: the cheap pre-merge checks and the remaining commands. + + Writes ``/premerge.md`` and returns its path with a short + summary. Never fails a delivery: what it reports is the operator's + remaining work (lint, the CI cases, the split, the PR), and the one + blocking finding — commit hygiene — is already a delivery check. + """ + from .premerge.pack import build_pack, summarize, write_pack + + policy = policy or CommitPolicy() + root = Path(repo_root) + if (root / ".git").exists(): + report = build_pack( + root, + test_pattern=f"test_modeling_v2_{family}", + base_ref=policy.base_ref, + allow_co_authors=policy.allow_co_authors, + paths=policy.allowed_paths, + ) + else: + from .premerge.pack import PackReport + + report = PackReport(repo=str(root), notes=["not a git work tree; nothing to check"]) + path = write_pack(report, reports_dir) + return path, list(summarize(report)) + + def verify_delivery( *, repo_root: Path | str, @@ -400,6 +485,7 @@ def verify_delivery( model_key: str, accuracy_test: str, extra_entries: Sequence[tuple[str, str]] = (), + commit_policy: CommitPolicy | None = None, ) -> DeliveryReport: """Run every mechanical delivery check and return the report. @@ -430,6 +516,8 @@ def verify_delivery( checks.extend(check_accuracy_reference(root, reference_yaml, model_key)) checks.extend(check_test_registration(root, accuracy_test)) checks.extend(check_tree_committed(root, (target_rel, MODELING_V2_PACKAGE, CATALOG_TESTS_DIR))) + if commit_policy is not None: + checks.extend(check_commit_hygiene(root, commit_policy)) return DeliveryReport(checks=tuple(checks)) @@ -459,6 +547,7 @@ def summarize_for_auditor(report: DeliveryReport, dump_excerpt: str = "") -> str "check_accuracy_reference", "check_catalog_entries", "check_catalog_tests", + "check_commit_hygiene", "check_routing", "check_target", "check_test_registration", @@ -467,4 +556,5 @@ def summarize_for_auditor(report: DeliveryReport, dump_excerpt: str = "") -> str "imported_catalog_entries", "summarize_for_auditor", "verify_delivery", + "write_premerge_pack", ] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/embargo.py b/agent-flow/agent_flow/workflows/staircase_onboard/embargo.py new file mode 100644 index 000000000000..8fca7cba0718 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/embargo.py @@ -0,0 +1,569 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``source_embargo``, enforced: a tool-call hook, not only a sentence in a prompt. + +``task.yaml`` names files and trees no agent may read — an implementation of +this exact checkpoint that already exists, here or in a sibling engine's +checkout. Every role prompt says so (``prompts._common.source_embargo_block``), +and a prompt is a request. This module makes it a mechanism for every role +the orchestrator spawns in process, through the agent backend's native +hooks (``BackendConfig.hooks``): + +**PreToolUse — refuse the read.** + +- ``Read`` / ``Edit`` / ``MultiEdit`` / ``Write`` / ``NotebookEdit`` on an + embargoed file, or anywhere inside an embargoed tree (symlinks resolved). +- ``Grep`` whose search root is an embargoed file or lies inside an + embargoed tree. +- ``Bash`` whose command names an embargoed **file** (absolute, + repo-relative or by its file name, or through a glob that matches it), or + names an embargoed **tree** together with a reading or searching program + (``cat``, ``grep``, ``sed``, ``python -c`` ...). Naming a tree to *run* + it — starting a sibling engine as the accuracy oracle — stays allowed: + running the other implementation is expected; reading it is not. + +**PostToolUse — redact what a broad search swept in.** A ``Grep`` or a +``grep -r`` over a parent directory is legitimate and can still return lines +from an embargoed file. Every output line (and every file-list element) that +names an embargoed path is removed before the model sees it, the output +keeps its shape, and the agent is told how many lines were withheld. + +**Every refusal and every redaction is logged**, one JSON object per line, +to ``/embargo/attempts.jsonl``: the role, the tool, the matched +entry, and what was asked. + +What it cannot see: a command that reads an embargoed file without naming +it or any parent search root that output lines would name (``rg -h``, a +script that opens a computed path). The prompt rule still covers those, and +the log records everything the hook did see. + +Where no hook can run, the embargo is an instruction instead, appended to the +role's system prompt by :func:`embargo_instruction`: a role on a backend +that cannot run this process's SDK callbacks (Codex refuses per-client +hooks; see :func:`source_embargo_enforcement`), and a turn delivered over the +external transport, which runs in a session this process does not own +(:func:`embargo_transport_note`). + +Entries are repo-relative, absolute (a tree outside the repository), or +``package:`` — relative to this workflow package wherever it is +installed, for the package's own worked examples. +""" + +from __future__ import annotations + +import fnmatch +import json +import os +import re +import shlex +import threading +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + +EMBARGO_DIRNAME = "embargo" +ATTEMPTS_FILENAME = "attempts.jsonl" + +#: Tools whose input names one file. +FILE_TOOLS: dict[str, tuple[str, ...]] = { + "Read": ("file_path",), + "Edit": ("file_path",), + "MultiEdit": ("file_path",), + "Write": ("file_path",), + "NotebookEdit": ("notebook_path",), + "NotebookRead": ("notebook_path",), +} +#: Tools whose output may carry another file's lines. +REDACTED_TOOLS: tuple[str, ...] = ("Grep", "Bash") + +#: Programs that read, search or copy what they are pointed at. Naming an +#: embargoed *tree* is only refused alongside one of these (or an inline +#: interpreter, below); naming an embargoed *file* is refused outright. +READ_PROGRAMS: frozenset[str] = frozenset( + { + "cat", "tac", "nl", "less", "more", "most", "head", "tail", "sed", "awk", "gawk", + "grep", "egrep", "fgrep", "rg", "ag", "ack", "git", "xargs", "cut", "sort", "uniq", + "diff", "cmp", "comm", "strings", "xxd", "od", "hexdump", "bat", "view", "vi", "vim", + "nvim", "nano", "emacs", "ed", "cp", "rsync", "scp", "tar", "zip", "jq", "yq", + "base64", "dd", "paste", "fold", "fmt", "pr", "ctags", "pygmentize", + } +) # fmt: skip +#: Interpreters that read a path only when handed inline code. ``python -m +#: vllm...`` from inside an embargoed checkout runs it, which is expected. +INLINE_INTERPRETERS: frozenset[str] = frozenset({"python", "python3", "perl", "ruby", "node"}) + +_GLOB_CHARS = set("*?[") + +#: An entry spelled ``package:`` is relative to this workflow package, +#: wherever it is installed — e.g. ``package:examples/.md``, a worked +#: example a run that onboards that model must not read. +PACKAGE_PREFIX = "package:" +PACKAGE_ROOT = Path(__file__).resolve().parent + +#: Backends whose sessions run this process's SDK hook callbacks. Every other +#: backend gets the embargo as an instruction in the role's system prompt. +HOOK_BACKENDS: frozenset[str] = frozenset({"claude-code"}) + + +class EmbargoEntryError(ValueError): + """An embargo entry that names nothing it is allowed to name.""" + + +def package_entry_problem(entry: str) -> str: + """Why a ``package:`` entry is unusable, or ``""``. + + It must name something strictly **inside** the workflow package: + ``package:../perf_analyze`` resolves to a sibling workflow, and an entry + for the package root would embargo the workflow's own skills. + """ + root = PACKAGE_ROOT.resolve() + target = (PACKAGE_ROOT / entry[len(PACKAGE_PREFIX) :].lstrip("/")).resolve() + if target == root: + return "names the whole workflow package; name a file or directory inside it" + if root not in target.parents: + return f"resolves to {target}, outside the workflow package {root}" + return "" + + +def resolve_entry(entry: str, repo_root: str | Path | None = None) -> Path: + """The absolute path an embargo entry names (symlinks not resolved). + + Raises :class:`EmbargoEntryError` for a ``package:`` entry that escapes + the workflow package; the task schema reports that before a run starts. + """ + if entry.startswith(PACKAGE_PREFIX): + problem = package_entry_problem(entry) + if problem: + raise EmbargoEntryError(f"{entry!r} {problem}") + return PACKAGE_ROOT / entry[len(PACKAGE_PREFIX) :].lstrip("/") + raw = Path(entry) + if raw.is_absolute() or repo_root is None: + return raw + return Path(repo_root) / raw + + +def display_entry(entry: str, repo_root: str | Path | None = None) -> str: + """How a role prompt names an entry. + + Repo and absolute entries pass through. A ``package:`` entry keeps its + prefix, so the prompt's embargo block can give it its own wording — it + is the workflow's reference material, not an implementation of the + checkpoint — followed by the path an agent would actually meet: + repo-relative when the package sits inside the repository, else absolute. + """ + if not entry.startswith(PACKAGE_PREFIX): + return entry + absolute = resolve_entry(entry, repo_root) + shown = str(absolute) + if repo_root is not None: + try: + shown = absolute.resolve().relative_to(Path(repo_root).resolve()).as_posix() + except ValueError: + pass + return f"{PACKAGE_PREFIX}{shown}" + + +def display_entries(task_data: Mapping[str, Any]) -> list[str]: + """``source_embargo`` as the role prompts render it.""" + repo = task_data.get("trtllm_repo_path") or None + return [display_entry(str(e), repo) for e in (task_data.get("source_embargo") or [])] + + +@dataclass(frozen=True) +class EmbargoEntry: + """One ``source_embargo`` entry, resolved.""" + + entry: str + path: Path + is_dir: bool + #: Spellings that, found in a command or an output line, name this entry. + forms: tuple[str, ...] + + +@dataclass(frozen=True) +class Verdict: + allowed: bool + entry: str = "" + reason: str = "" + + +def _forms( + entry: str, absolute: Path, resolved: Path, repo: Path | None, is_dir: bool +) -> tuple[str, ...]: + forms = {str(absolute), str(resolved)} + if entry.startswith(PACKAGE_PREFIX): + forms.add(entry[len(PACKAGE_PREFIX) :].strip("/")) + try: + forms.add(resolved.relative_to(PACKAGE_ROOT.parent.parent.parent).as_posix()) + except ValueError: + pass + elif not entry.startswith("/"): + forms.add(entry.rstrip("/")) + if repo is not None: + for root in {repo, repo.resolve()}: + try: + forms.add(resolved.relative_to(root).as_posix()) + except ValueError: + pass + if not is_dir: + forms.add(absolute.name) + return tuple(sorted(f for f in forms if f and f != ".")) + + +def _pattern(forms: Iterable[str]) -> re.Pattern[str]: + alternatives = "|".join(re.escape(f) for f in sorted(forms, key=len, reverse=True)) + return re.compile(rf"(? None: + self.repo_root = Path(repo_root) if repo_root else None + self.log_path = Path(log_path) if log_path else None + self._lock = threading.Lock() + resolved: list[EmbargoEntry] = [] + for entry in entries: + absolute = resolve_entry(entry, self.repo_root) + real = absolute.resolve() + is_dir = real.is_dir() or entry.endswith("/") + resolved.append( + EmbargoEntry( + entry=entry, + path=real, + is_dir=is_dir, + forms=_forms(entry, absolute, real, self.repo_root, is_dir), + ) + ) + self.entries: tuple[EmbargoEntry, ...] = tuple(resolved) + self._patterns = [(e, _pattern(e.forms)) for e in self.entries] + + @classmethod + def from_task( + cls, task_data: Mapping[str, Any], workspace: str | Path | None = None + ) -> "SourceEmbargo | None": + entries = [str(p) for p in (task_data.get("source_embargo") or [])] + if not entries: + return None + log = Path(workspace) / EMBARGO_DIRNAME / ATTEMPTS_FILENAME if workspace else None + return cls(entries, repo_root=task_data.get("trtllm_repo_path") or None, log_path=log) + + # ------------------------------------------------------------- matching + + def match_path(self, path: str, cwd: str | Path | None = None) -> EmbargoEntry | None: + """The entry ``path`` falls under, symlinks resolved, or ``None``.""" + if not path: + return None + raw = Path(os.path.expanduser(path)) + if not raw.is_absolute(): + raw = Path(cwd or self.repo_root or Path.cwd()) / raw + candidates = {Path(os.path.abspath(raw)), raw.resolve()} + for entry in self.entries: + for candidate in candidates: + if candidate == entry.path or (entry.is_dir and entry.path in candidate.parents): + return entry + return None + + def mentioned(self, text: str) -> EmbargoEntry | None: + """The first entry ``text`` names in any of its spellings.""" + for entry, pattern in self._patterns: + if pattern.search(text): + return entry + return None + + def _glob_match(self, token: str, cwd: str | Path | None) -> EmbargoEntry | None: + """The embargoed file a shell glob would expand to, if any. + + Tried against the session's cwd and the repository root both: the + command may ``cd`` first, and over-matching only refuses a read. + """ + if not _GLOB_CHARS & set(token): + return None + if token.startswith("/"): + patterns = [token] + else: + bases = {str(b) for b in (cwd, self.repo_root) if b} + patterns = [str(Path(b) / token) for b in bases] or [token] + for entry in self.entries: + if any(fnmatch.fnmatch(str(entry.path), p) for p in patterns): + return entry + return None + + # ------------------------------------------------------------- the checks + + def check_tool( + self, tool: str, tool_input: Mapping[str, Any], cwd: str | Path | None = None + ) -> Verdict: + """Whether one tool call may run.""" + if tool in FILE_TOOLS: + for key in FILE_TOOLS[tool]: + hit = self.match_path(str(tool_input.get(key) or ""), cwd) + if hit: + return self._deny(hit, f"{tool} on an embargoed path") + return Verdict(True) + if tool == "Grep": + root = str(tool_input.get("path") or "") + hit = self.match_path(root, cwd) if root else None + if hit: + return self._deny(hit, "Grep inside an embargoed path") + return Verdict(True) + if tool == "Bash": + return self._check_command(str(tool_input.get("command") or ""), cwd) + return Verdict(True) + + def _check_command(self, command: str, cwd: str | Path | None) -> Verdict: + try: + tokens = shlex.split(command, posix=True) + except ValueError: + tokens = command.split() + programs = {Path(t).name for t in tokens if t and not t.startswith("-")} + inline = bool(programs & INLINE_INTERPRETERS) and bool({"-c", "-e"} & set(tokens)) + reads = bool(programs & READ_PROGRAMS) or inline or "<" in command + hit = self.mentioned(command) + if hit is not None and (not hit.is_dir or reads): + return self._deny(hit, "a shell command that names an embargoed path") + for token in tokens: + globbed = self._glob_match(token, cwd) + if globbed is not None and reads: + return self._deny(globbed, "a shell glob that matches an embargoed path") + return Verdict(True) + + def _deny(self, entry: EmbargoEntry, what: str) -> Verdict: + return Verdict( + False, + entry.entry, + f"Refused by the run's source embargo: {what} (`{entry.entry}`). This path is off " + f"limits to read for the whole run — not to open, grep, cat or copy. Running the " + f"implementation as a black box is still expected. This attempt has been logged; " + f"continue without it.", + ) + + # ------------------------------------------------------------- redaction + + def redact(self, value: Any) -> tuple[Any, int]: + """``value`` with every line / element naming an embargoed path removed. + + The structure is preserved (same keys, same types), because a hook's + replacement output must match the tool's own output schema. + """ + if isinstance(value, str): + lines = value.splitlines(keepends=True) + kept = [line for line in lines if self.mentioned(line) is None] + return "".join(kept), len(lines) - len(kept) + if isinstance(value, list): + out, removed = [], 0 + for item in value: + if isinstance(item, str) and self.mentioned(item) is not None: + removed += 1 + continue + new, n = self.redact(item) + out.append(new) + removed += n + return out, removed + if isinstance(value, dict): + out_d, removed = {}, 0 + for key, item in value.items(): + new, n = self.redact(item) + out_d[key] = new + removed += n + return out_d, removed + return value, 0 + + # ------------------------------------------------------------- the log + + def log(self, **event: Any) -> None: + if self.log_path is None: + return + record = {"ts": datetime.now().isoformat(timespec="seconds"), **event} + line = json.dumps(record, sort_keys=True, default=str) + "\n" + with self._lock: + self.log_path.parent.mkdir(parents=True, exist_ok=True) + with self.log_path.open("a", encoding="utf-8") as handle: + handle.write(line) + + def attempts(self) -> list[dict[str, Any]]: + if self.log_path is None or not self.log_path.is_file(): + return [] + return [ + json.loads(line) + for line in self.log_path.read_text(encoding="utf-8").splitlines() + if line + ] + + +def _summary(tool_input: Mapping[str, Any]) -> str: + for key in ("command", "file_path", "notebook_path", "path", "pattern"): + if tool_input.get(key): + return str(tool_input[key])[:300] + return json.dumps(dict(tool_input), default=str)[:300] + + +def embargo_hooks(embargo: SourceEmbargo, *, role: str) -> dict[str, list[Any]]: + """``BackendConfig.hooks`` entries enforcing ``embargo`` for one role's sessions.""" + from claude_agent_sdk.types import HookMatcher + + async def pre_tool_use(input_data, _tool_use_id, _context): + tool = str(input_data.get("tool_name") or "") + tool_input = input_data.get("tool_input") or {} + verdict = embargo.check_tool(tool, tool_input, input_data.get("cwd")) + if verdict.allowed: + return {} + embargo.log( + action="deny", + role=role, + tool=tool, + entry=verdict.entry, + detail=_summary(tool_input), + session=input_data.get("session_id"), + agent_id=input_data.get("agent_id"), + ) + return { + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "deny", + "permissionDecisionReason": verdict.reason, + } + } + + async def post_tool_use(input_data, _tool_use_id, _context): + tool = str(input_data.get("tool_name") or "") + if tool not in REDACTED_TOOLS: + return {} + redacted, removed = embargo.redact(input_data.get("tool_response")) + if not removed: + return {} + tool_input = input_data.get("tool_input") or {} + embargo.log( + action="redact", + role=role, + tool=tool, + removed=removed, + detail=_summary(tool_input), + session=input_data.get("session_id"), + agent_id=input_data.get("agent_id"), + ) + return { + "hookSpecificOutput": { + "hookEventName": "PostToolUse", + "updatedToolOutput": redacted, + "additionalContext": ( + f"The run's source embargo withheld {removed} line(s) of this {tool} output " + f"that came from embargoed paths. Do not try to read them another way; the " + f"attempt is logged." + ), + } + } + + pre_matcher = "|".join([*FILE_TOOLS, "Grep", "Bash"]) + return { + "PreToolUse": [HookMatcher(matcher=pre_matcher, hooks=[pre_tool_use])], + "PostToolUse": [HookMatcher(matcher="|".join(REDACTED_TOOLS), hooks=[post_tool_use])], + } + + +def merge_hook_maps(*maps: Mapping[str, list[Any]] | None) -> dict[str, list[Any]] | None: + """Concatenate hook maps event by event; ``None`` when there is nothing.""" + merged: dict[str, list[Any]] = {} + for hook_map in maps: + for event, matchers in (hook_map or {}).items(): + merged.setdefault(event, []).extend(matchers) + return merged or None + + +def source_embargo_hooks( + task_data: Mapping[str, Any], workspace: str | Path, *, role: str +) -> dict[str, list[Any]] | None: + """The hooks for ``role``, or ``None`` when the task embargoes nothing.""" + embargo = SourceEmbargo.from_task(task_data, workspace) + return embargo_hooks(embargo, role=role) if embargo is not None else None + + +def embargo_instruction( + paths: Sequence[str], *, repo_root: str | Path | None = None, why: str +) -> str: + """The embargo, restated for a session no hook of this process reaches. + + ``why`` says which: a turn on the external transport runs in a session + the orchestrator does not own; a role on a backend outside + :data:`HOOK_BACKENDS` cannot run SDK callbacks. The role prompt still + carries the embargo; this says that nothing will refuse the read, gives + absolute paths, and asks for any exposure to be reported, not used. + """ + if not paths: + return "" + listed = "\n".join(f"- `{resolve_entry(str(p), repo_root)}`" for p in paths) + return ( + "\n\n---\n\n" + "## The source embargo is not enforced by a hook here — you enforce it\n\n" + "In process, on a backend that runs them, a hook refuses every Read, Grep and shell " + f"read of these paths and logs the attempt. {why}, so nothing will stop you:\n\n" + + listed + + "\n\nA directory embargoes the whole tree beneath it. Do not open, grep, cat, copy or " + "quote any of it, and keep broad searches out of it (exclude it, or search narrower " + "roots). If a tool output shows you content from one anyway, stop reading it, do not " + "use it, and list each exposure under `## Embargo exposures` in your result. " + "Running the implementation as a black box remains expected.\n" + ) + + +def embargo_transport_note(paths: Sequence[str], *, repo_root: str | Path | None = None) -> str: + """:func:`embargo_instruction` for a turn on the external transport.""" + return embargo_instruction( + paths, + repo_root=repo_root, + why="This turn runs on the external transport, in a session the orchestrator does not own", + ) + + +def source_embargo_enforcement( + task_data: Mapping[str, Any], workspace: str | Path, *, role: str, backend: str +) -> tuple[dict[str, list[Any]] | None, str]: + """``(hooks, system-prompt suffix)`` enforcing the task's embargo for one role. + + A backend in :data:`HOOK_BACKENDS` gets the PreToolUse/PostToolUse hooks + and no suffix; any other backend (Codex refuses per-client callbacks) gets + no hooks and the instruction. Both are empty when nothing is embargoed. + """ + embargo = SourceEmbargo.from_task(task_data, workspace) + if embargo is None: + return None, "" + if backend in HOOK_BACKENDS: + return embargo_hooks(embargo, role=role), "" + note = embargo_instruction( + [e.entry for e in embargo.entries], + repo_root=embargo.repo_root, + why=f"This role runs on the `{backend}` backend, which cannot run this process's hooks", + ) + embargo.log(action="instruct", role=role, backend=backend, detail="no hook on this backend") + return None, note + + +__all__ = [ + "ATTEMPTS_FILENAME", + "EMBARGO_DIRNAME", + "EmbargoEntry", + "EmbargoEntryError", + "SourceEmbargo", + "Verdict", + "HOOK_BACKENDS", + "PACKAGE_PREFIX", + "PACKAGE_ROOT", + "display_entries", + "display_entry", + "embargo_hooks", + "embargo_instruction", + "embargo_transport_note", + "merge_hook_maps", + "package_entry_problem", + "resolve_entry", + "source_embargo_enforcement", + "source_embargo_hooks", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/examples/deepseek-v4.1.md b/agent-flow/agent_flow/workflows/staircase_onboard/examples/deepseek-v4.1.md index 3b3fbc4e42e9..cf0aca003a1e 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/examples/deepseek-v4.1.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/examples/deepseek-v4.1.md @@ -12,8 +12,8 @@ this file shows one model's values filled in. **This file is for operators writing a `task.yaml`.** The run does not materialize `examples/` into its workspace, so no role is handed it. A run that onboards this checkpoint, or a sibling of it, to measure the workflow -should add this file to `source_embargo`, because the facts below are what -the run is meant to derive. +should embargo it (`source_embargo: [package:examples/deepseek-v4.1.md]`), +because the facts below are what the run is meant to derive. Sources: the DeepSeek-V4.1-Flash tech report, the checkpoint's own reference code, and SGLang's implementation (branch `origin/dsv4.1`). diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/overlap.py b/agent-flow/agent_flow/workflows/staircase_onboard/overlap.py new file mode 100644 index 000000000000..9eace145f517 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/overlap.py @@ -0,0 +1,541 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit overlap (``--overlap-units``): start the next unit's assembler early, in its own worktree. + +Units are assembled strictly in order, and that order is right: core is green +before the first feature. But a unit's *gates* — boot over every member, the +reviewer, accuracy over every member, the witness — take hours, and while they +run the next unit's assembler is idle, although nothing it does depends on +those gates passing. Certification-only rejects never change the target, so +the work it would do is usually valid as it stands. + +So, opt-in: + +1. As soon as unit N's assembler writes ``UNIT_READY`` (before N's gates), + unit N+1's assembler starts in a **separate git worktree** branched from + ``HEAD``, with the main checkout's untracked native build linked in, its + own commit helper, its own progress file and its own decision path. It + runs in the background; nothing it writes can be read as unit N's. +2. When N's gates have passed and the loop reaches N+1, the early result is + **adopted** only if every one of these holds: the turn ended + ``UNIT_READY``, it committed something, its worktree is clean, its + commits rebase onto the current ``HEAD`` without conflict (N's later + turns may have moved it), and the main checkout fast-forwards to them. + Adoption appends the early decision to the run's progress and goes + straight to N+1's gates, the way a resume onto an unchanged + ``UNIT_READY`` does. +3. **Anything else falls back to sequential**: the main checkout is left + exactly as it was, the early worktree is kept for inspection (with its + diff saved), and N+1's assembler is dispatched normally. + +Not covered, deliberately: a run that ends (or a process that dies) while an +early turn is still running leaves its worktree on disk; a resumed process +never adopts a worktree it did not start, and says so. Remove it with +``git worktree remove --force /overlap//tree``. +""" + +from __future__ import annotations + +import os +import threading +from dataclasses import dataclass, field, replace +from pathlib import Path +from typing import Any, Callable + +from .committer import CommitIdentity, CommitPolicy, commit_instruction, install_committer +from .premerge._git import GitError, git + +OVERLAP_DIRNAME = "overlap" +BRANCH_PREFIX = "staircase-overlap" + +#: Untracked build outputs a worktree needs to import the package, found as +#: the git-ignored entries directly under this directory of the main checkout. +NATIVE_ROOT = "tensorrt_llm" +_NOT_NATIVE = ("__pycache__",) + + +class OverlapError(RuntimeError): + """The early worktree could not be prepared.""" + + +@dataclass +class OverlapJob: + """One early assembler turn, and what came of it.""" + + unit: str + after: str + root: Path + tree: Path + branch: str + base: str + commit_script: Path + thread: threading.Thread | None = None + decision: str = "" + entry: dict[str, Any] = field(default_factory=dict) + error: str = "" + linked: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class AdoptOutcome: + adopted: bool + detail: str + + +def native_artifacts(repo: str | Path) -> list[str]: + """Git-ignored entries directly under ``tensorrt_llm/`` in ``repo``.""" + out = git( + repo, + "ls-files", + "--others", + "--ignored", + "--exclude-standard", + "--directory", + "--", + f"{NATIVE_ROOT}/", + check=False, + ) + found = [] + for line in out.splitlines(): + path = line.strip().rstrip("/") + parts = path.split("/") + if len(parts) == 2 and parts[0] == NATIVE_ROOT and parts[1] not in _NOT_NATIVE: + if not parts[1].endswith(".egg-info"): + found.append(path) + return found + + +class UnitOverlap: + """The git side of unit overlap: worktrees, linking, integration, cleanup. + + ``turn_runner(job, prompt)`` runs the early turn and returns its typed + decision entry (a mapping with ``decision``); the workflow supplies one, + a test supplies a fake. + """ + + def __init__( + self, + repo: str | Path, + workspace: str | Path, + *, + policy: CommitPolicy | None = None, + link_native: bool = True, + turn_runner: Callable[[OverlapJob, str], dict[str, Any] | None] | None = None, + ) -> None: + self.repo = Path(repo) + self.workspace = Path(workspace) + self.policy = policy or CommitPolicy() + self.link_native = link_native + self.turn_runner = turn_runner + self.jobs: dict[str, OverlapJob] = {} + + # ------------------------------------------------------------- set-up + + def _env(self) -> dict[str, str]: + env = dict(os.environ) + identity: CommitIdentity | None = self.policy.identity + if identity is not None: + env.update(GIT_COMMITTER_NAME=identity.name, GIT_COMMITTER_EMAIL=identity.email) + return env + + def open_worktree(self, unit: str, *, after: str) -> OverlapJob: + """A worktree for ``unit`` at the main checkout's ``HEAD``, ready to commit into.""" + root = self.workspace / OVERLAP_DIRNAME / unit + tree = root / "tree" + if tree.exists(): + raise OverlapError(f"{tree} already exists (a previous early start); remove it first") + base = git(self.repo, "rev-parse", "HEAD").strip() + branch = f"{BRANCH_PREFIX}/{unit}-{base[:10]}" + root.mkdir(parents=True, exist_ok=True) + try: + git( + self.repo, + "worktree", + "add", + "-q", + "-b", + branch, + str(tree), + base, + env={ + **os.environ, + "GIT_LFS_SKIP_SMUDGE": "1", + }, + ) + except GitError as exc: + raise OverlapError(str(exc)) from exc + linked = self._link_native(tree) if self.link_native else [] + script = install_committer( + root, + tree, + policy=replace(self.policy, denied_paths=()), + ) + job = OverlapJob( + unit=unit, + after=after, + root=root, + tree=tree, + branch=branch, + base=base, + commit_script=script, + linked=linked, + ) + self.jobs[unit] = job + return job + + def _link_native(self, tree: Path) -> list[str]: + linked = [] + for rel in native_artifacts(self.repo): + target = tree / rel + if target.exists() or target.is_symlink(): + continue + target.parent.mkdir(parents=True, exist_ok=True) + target.symlink_to(self.repo / rel) + linked.append(rel) + return linked + + # ------------------------------------------------------------- running + + def start(self, job: OverlapJob, prompt: str) -> None: + """Run the early turn in a daemon thread; the result lands on ``job``.""" + if self.turn_runner is None: + raise OverlapError("no turn runner configured") + + def run() -> None: + try: + entry = self.turn_runner(job, prompt) or {} + except Exception as exc: # noqa: BLE001 + # The thread boundary: whatever the turn raised has to reach + # the adopting thread as a reason to fall back, not vanish + # with a daemon thread. + job.error = f"{type(exc).__name__}: {exc}" + return + job.entry = dict(entry) + job.decision = str(entry.get("decision", "")).strip().upper() + + job.thread = threading.Thread(target=run, name=f"overlap-{job.unit}", daemon=True) + job.thread.start() + + def wait(self, unit: str, timeout: float | None = None) -> OverlapJob | None: + job = self.jobs.get(unit) + if job is not None and job.thread is not None: + job.thread.join(timeout) + return job + + # ------------------------------------------------------------- landing + + def integrate(self, job: OverlapJob) -> AdoptOutcome: + """Rebase the early commits onto ``HEAD`` and fast-forward the main checkout. + + Touches the main checkout only in the final fast-forward, and only + after every check passed; a failure at any step leaves it unchanged. + """ + if job.thread is not None and job.thread.is_alive(): + return AdoptOutcome(False, "the early turn is still running") + if job.error: + return AdoptOutcome(False, f"the early turn failed: {job.error}") + if job.decision != "UNIT_READY": + return AdoptOutcome(False, f"the early turn ended {job.decision or 'with no decision'}") + dirty = [ + line + for line in git(job.tree, "status", "--porcelain", "--untracked-files=no").splitlines() + if line.strip() + ] + if dirty: + return AdoptOutcome(False, f"uncommitted work in the early worktree: {dirty[:5]}") + ahead = git(job.tree, "rev-list", f"{job.base}..HEAD").split() + if not ahead: + return AdoptOutcome(False, "the early turn committed nothing") + head = git(self.repo, "rev-parse", "HEAD").strip() + if head != job.base: + try: + git(job.tree, "rebase", "-q", "--onto", head, job.base, env=self._env()) + except GitError as exc: + git(job.tree, "rebase", "--abort", check=False) + return AdoptOutcome( + False, f"its commits do not rebase onto HEAD {head[:10]}: {exc}" + ) + tip = git(job.tree, "rev-parse", "HEAD").strip() + try: + git(self.repo, "merge", "--ff-only", "-q", tip) + except GitError as exc: + return AdoptOutcome( + False, f"the main checkout does not fast-forward to {tip[:10]}: {exc}" + ) + return AdoptOutcome(True, f"{len(ahead)} commit(s) landed; HEAD {head[:10]} -> {tip[:10]}") + + def remove(self, job: OverlapJob) -> None: + """Remove an adopted job's worktree and branch (its commits are on HEAD now).""" + git(self.repo, "worktree", "remove", "--force", str(job.tree), check=False) + git(self.repo, "branch", "-D", job.branch, check=False) + self.jobs.pop(job.unit, None) + + def keep_for_inspection(self, job: OverlapJob, why: str) -> Path: + """Leave a job that was not adopted on disk, with its diff and the reason.""" + note = job.root / "NOT_ADOPTED.md" + diff = git(job.tree, "diff", job.base, check=False) if job.tree.exists() else "" + (job.root / "early.patch").write_text(diff, encoding="utf-8") + note.write_text( + f"# {job.unit}: early start not adopted\n\n- why: {why}\n- worktree: `{job.tree}`\n" + f"- branch: `{job.branch}` (from `{job.base[:10]}`)\n- diff: `early.patch`\n\n" + f"Remove with `git -C {self.repo} worktree remove --force {job.tree}` and " + f"`git -C {self.repo} branch -D {job.branch}`.\n", + encoding="utf-8", + ) + self.jobs.pop(job.unit, None) + return note + + def stale(self, unit: str) -> Path | None: + """A worktree for ``unit`` this process did not start, if one is on disk.""" + tree = self.workspace / OVERLAP_DIRNAME / unit / "tree" + return tree if unit not in self.jobs and tree.exists() else None + + +# ----------------------------------------------------------------- the glue +# +# Everything below needs the orchestrator; it is here rather than in +# workflow.py so that file only gains its two call sites. + + +def early_prompt(workflow: Any, job: OverlapJob) -> str: + """Unit ``job.unit``'s first assembler prompt, re-pointed at the worktree.""" + repo = str(workflow.task_data["trtllm_repo_path"]) + body = workflow._assembler_prompt(job.unit, 1).replace( + f"Repository: {repo}\n", f"Repository: {job.tree}\n", 1 + ) + preface = ( + f"## Early start: `{job.unit}` overlaps `{job.after}`'s gates\n\n" + f"`{job.after}` is UNIT_READY and is being gated now. You are assembling " + f"`{job.unit}` ahead of that, in its **own worktree**:\n\n" + f"- every path that names the repository means `{job.tree}` for this turn — edit " + f"only there, never in `{repo}`;\n" + f"- commit only through `{job.commit_script}` (the script in your system prompt " + f"commits into the other checkout);\n" + f"- the native build is linked in from the main checkout, so imports work; do not " + f"rebuild it from here.\n\n" + f"If `{job.after}`'s gates force a change, your commits are rebased onto it before " + f"they land; a conflict discards this turn and the unit is assembled again " + f"afterwards, so keep to `{job.unit}`'s own files.\n\n" + ) + context = workflow._context_block() + return ( + preface + + body + + (f"\n\n{context}\n" if context else "") + + "\n\n" + + commit_instruction(job.commit_script) + ) + + +def run_early_turn(workflow: Any, job: OverlapJob, prompt: str) -> dict[str, Any] | None: + """Run the early turn through the run's own transport; return its decision entry.""" + from agent_flow.runtime import TurnRequest + + from .progress import ( + ProgressContext, + build_progress_tools, + collect_decision_file, + decision_path, + init_progress_file, + latest_entry, + mark_decision_dispatch, + ) + from .roles import resolve_role_agent + + ctx = ProgressContext(path=job.root / "progress.yaml", current_iteration=1) + init_progress_file(ctx.path) + decision_file = None + if workflow._decision_files: + target = decision_path(job.root, "assembler", 1) + decision_file = str(target) + mark_decision_dispatch(job.root, "assembler", 1) + prompt = ( + f"{prompt}\n\n---\n\n**Write your decision file to exactly this path when you " + f"finish:**\n\n`{target}`\n" + ) + layer = None + if workflow._agent_transport is None: + layer = workflow._make_agent( + "assembler", + workflow.prompts.assembler, + required_tools=["append_assembler_progress"], + session_mode="stateless", + ) + config = getattr(layer, "config", None) + if config is not None and not workflow._decision_files: + from agent_flow import AgentLayer + + tools = build_progress_tools(ctx).get("assembler") + layer = AgentLayer(replace(config, backend=replace(config.backend, tools=tools))) + what = f"{job.unit} (early start, overlapping {job.after})" + agent = resolve_role_agent(workflow.task_data, "assembler") + request = TurnRequest( + role="assembler", + what=what, + prompt=prompt, + system_prompt=getattr(workflow.prompts, "assembler", ""), + backend=agent.backend, + model=agent.model, + decision_file=decision_file, + iteration=1, + ) + report = workflow.journal.record_dispatch("assembler", what, prompt) + try: + executor = workflow._turns if layer is None else workflow._turns.bound(layer) + text = executor.run(request).text + finally: + if layer is not None and hasattr(layer, "__exit__"): + layer.__exit__(None, None, None) + workflow.journal.record_result(report, text) + if workflow._decision_files: + collect_decision_file(ctx, job.root, "assembler", 1) + return latest_entry(ctx.path, "assembler") + + +def maybe_start(workflow: Any, state: Any, unit: str, log: Any) -> bool: + """After ``unit``'s UNIT_READY: start the next unit early, if that is safe.""" + from agent_flow.console import print_message + + from .state import ITEM_DONE + from .task_schema import assembly_units + + overlap = _overlap_of(workflow) + if overlap is None: + return False + units = assembly_units(workflow.task_data) + if unit not in units or units.index(unit) + 1 >= len(units): + return False + nxt = units[units.index(unit) + 1] + if nxt in state.units_done or nxt in overlap.jobs: + return False + owed = [row for row in state.queue if row.get("deferred_to") and row.get("status") != ITEM_DONE] + reason = "" + if owed: + reason = f"{len(owed)} deferred smith dispatch(es) must land first" + elif not workflow._target_marker(): + reason = "the target's commit could not be established" + else: + dirt = git( + overlap.repo, + "status", + "--porcelain", + "--untracked-files=no", + "--", + workflow._preflight.target_rel, + check=False, + ).strip() + if dirt: + reason = "the target has uncommitted changes, which a worktree would not carry" + if not reason: + try: + job = overlap.open_worktree(nxt, after=unit) + except (OverlapError, GitError) as exc: + reason = f"no worktree: {exc}" + if reason: + print_message( + f"[bold yellow]→ overlap: not starting {nxt} early — {reason}[/bold yellow]", log + ) + workflow.journal.event("overlap not started", f"{nxt}: {reason}") + return False + overlap.start(job, early_prompt(workflow, job)) + print_message( + f"[bold cyan]→ overlap: {unit} is UNIT_READY; {nxt}'s assembler starts now in " + f"{job.tree} while {unit} is gated[/bold cyan]", + log, + ) + workflow.journal.event( + "overlap started", + f"{nxt} in {job.tree} from {job.base[:10]}; native links: {len(job.linked)}", + ) + return True + + +def adopt(workflow: Any, state: Any, unit: str, log: Any) -> bool: + """Before ``unit``'s first turn: land its early result, or fall back to sequential. + + Returns True when adopted; the caller's ordinary resume path then sees an + unchanged ``UNIT_READY`` and goes straight to the gates. + """ + from agent_flow.console import print_message + + overlap = _overlap_of(workflow) + if overlap is None: + return False + stale = overlap.stale(unit) + if stale is not None: + print_message( + f"[bold yellow]→ overlap: {stale} was left by a previous process and is not " + f"adopted; {unit} is assembled normally[/bold yellow]", + log, + ) + workflow.journal.event("overlap stale worktree ignored", str(stale)) + return False + job = overlap.jobs.get(unit) + if job is None: + return False + if job.thread is not None and job.thread.is_alive(): + print_message( + f"[bold cyan]→ overlap: waiting for {unit}'s early turn to finish[/bold cyan]", log + ) + overlap.wait(unit) + outcome = overlap.integrate(job) + if not outcome.adopted: + note = overlap.keep_for_inspection(job, outcome.detail) + print_message( + f"[bold yellow]→ overlap: {unit}'s early start is not adopted ({outcome.detail}); " + f"falling back to a sequential assembly. Kept for inspection: {note}[/bold yellow]", + log, + ) + workflow.journal.event("overlap fell back", f"{unit}: {outcome.detail}") + return False + entry = {**job.entry, "iteration": 1, "overlap": True, "unit": unit} + entry.setdefault("agent", "assembler") + workflow._progress_ctx.append("assembler", entry) + state.unit_marker[unit] = workflow._target_marker() + workflow._resumed_into_unit = True + workflow._checkpoint(state) + overlap.remove(job) + print_message( + f"[bold green]✔ overlap: {unit}'s early UNIT_READY adopted — {outcome.detail}[/bold green]", + log, + ) + workflow.journal.event("overlap adopted", f"{unit}: {outcome.detail}") + return True + + +def _overlap_of(workflow: Any) -> UnitOverlap | None: + """The run's overlap, built on first use; ``None`` unless ``--overlap-units``.""" + if not getattr(workflow, "_overlap_units", False): + return None + if getattr(workflow, "_overlap", None) is None: + workflow._overlap = build_overlap(workflow) + return workflow._overlap + + +def build_overlap(workflow: Any) -> UnitOverlap: + """The run's :class:`UnitOverlap`, wired to the run's transport.""" + repo = workflow.task_data["trtllm_repo_path"] + overlap = UnitOverlap( + repo, + workflow.workspace, + policy=CommitPolicy.from_task(workflow.task_data, workspace=workflow.workspace), + ) + overlap.turn_runner = lambda job, prompt: run_early_turn(workflow, job, prompt) + return overlap + + +__all__ = [ + "AdoptOutcome", + "BRANCH_PREFIX", + "OVERLAP_DIRNAME", + "OverlapError", + "OverlapJob", + "UnitOverlap", + "adopt", + "build_overlap", + "early_prompt", + "maybe_start", + "native_artifacts", + "run_early_turn", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/premerge/__init__.py b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/__init__.py new file mode 100644 index 000000000000..f4776db455fe --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/__init__.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The pre-merge pack: what a delivered branch still owes before it is a PR. + +Every step here was first done by hand at the end of an onboarding, after +every unit was green, and every one of them found something: + +- :mod:`.lint` — the repository's pre-commit hooks over the changed files, + twice; the formatting is committed only when the second pass is clean, and + only with an AST proof that it changed nothing but formatting. +- :mod:`.ci_case` — a test-db line run the way CI runs it (the line itself as + the test list, CI's pytest flags, CI's working directory), recording the + ``HEAD`` it measured at start and at end. +- :mod:`.copyright` — NVIDIA headers on new source files, current year on + modified ones. +- :mod:`.split` — the branch split by path into reviewable PR stacks, with + the checks that the stacks together reproduce the branch. +- :mod:`.resign` — a non-destructive DCO re-sign that signs each commit as + its own author. +- :mod:`.pr_open` — opening a draft PR; opt-in, and dry-run unless told + otherwise. Nothing in this package pushes. + +``staircase-onboard premerge ...`` runs any of them; :mod:`.pack` +writes the cheap ones into a report at the end of delivery. + +Target-path flattening (``flatten_target.sh`` in the run that produced this +pack) is not ported: the workflow now places targets at +``models//____/`` from the start, so there +is no nested layout left to flatten. +""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/premerge/_git.py b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/_git.py new file mode 100644 index 000000000000..ecbef5a38058 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/_git.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The one way this package runs git.""" + +from __future__ import annotations + +import subprocess +from pathlib import Path +from typing import Mapping + + +class GitError(RuntimeError): + """A git command this package depends on failed.""" + + +def git( + repo: str | Path, + *args: str, + env: Mapping[str, str] | None = None, + input_: str | None = None, + check: bool = True, +) -> str: + """Run ``git -C repo `` and return its stdout. + + Raises :class:`GitError` naming the command and git's own message when + ``check`` is set and git exits non-zero. + """ + result = subprocess.run( + ["git", "-C", str(repo), *args], + capture_output=True, + text=True, + env=dict(env) if env is not None else None, + input=input_, + ) + if check and result.returncode != 0: + raise GitError(f"git {' '.join(args)} failed: {result.stderr.strip()}") + return result.stdout + + +def rev_parse(repo: str | Path, rev: str) -> str: + return git(repo, "rev-parse", "--verify", f"{rev}^{{commit}}").strip() + + +def dirty_tracked(repo: str | Path) -> list[str]: + """Tracked files with uncommitted changes (untracked files are ignored).""" + out = git(repo, "status", "--porcelain", "--untracked-files=no") + return [line for line in out.splitlines() if line.strip()] + + +def changed_files(repo: str | Path, base: str, head: str = "HEAD") -> list[str]: + """Files ``head`` changes relative to ``base``, deletions excluded. + + The same set CI's pre-merge release check hands to pre-commit. + """ + out = git(repo, "diff", "--name-only", "--diff-filter=d", base, head) + return [line for line in out.splitlines() if line.strip()] + + +__all__ = ["GitError", "changed_files", "dirty_tracked", "git", "rev_parse"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/premerge/ci_case.py b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/ci_case.py new file mode 100644 index 000000000000..3739d3e8fec3 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/ci_case.py @@ -0,0 +1,337 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Run a test-db line the way CI runs it, and record what it measured. + +The difference between "the gate passed" and "CI will pass" was, in the run +this was built from, four accuracy cases collected and then **skipped** on +every CI stage: no stage exported the switch the gate had exported. Running +the case as CI does — not as the gate does — is what found it. + +What "as CI does" means here (``jenkins/L0_Test.groovy``, +``getPytestBaseCommandLine``): + +- the test-db line itself, ``TIMEOUT (n)`` included, is the ``--test-list``; +- CI's pytest flags, from ``tests/integration/defs`` as the working + directory, with the stage name as ``--test-prefix`` and the repository's + waives file; +- nothing exported beyond what CI exports (``LLM_ROOT``, + ``LLM_BACKEND_ROOT``, ``LLM_MODELS_ROOT``, ``MODEL_CACHE_DIR``, + ``COLUMNS``); the case's own environment travels with the case. + +What differs is stated in the generated script: how ranks are launched +(``--launcher``, e.g. ``mpirun -n 4 ... trtllm-llmapi-launch`` where the +container forbids ``MPI_Comm_spawn``) and where the models root points. + +The script prints ``HEAD``, the branch and the number of modified tracked +files at start **and** at end, plus where ``tensorrt_llm`` resolved: a result +is only evidence for the tree it measured, and a branch that moved while the +case ran has no single tree. +""" + +from __future__ import annotations + +import json +import re +import shlex +import subprocess +import sys +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Sequence + +import yaml + +TEST_DB_DIR = "tests/integration/test_lists/test-db" +DEFS_DIR = "tests/integration/defs" +WAIVES = "tests/integration/test_lists/waives.txt" +STAGE_MAPPER = "scripts/test_to_stage_mapping.py" + +#: CI's own pytest flags, in CI's order, minus the reporting plugins that +#: only make sense inside Jenkins. ``{...}`` fields are filled per case. +CI_PYTEST_FLAGS: tuple[str, ...] = ( + "-vv", + "-rA", + "--timeout-method=thread", + "--apply-test-list-correction", + "--timeout={timeout}", + "--rootdir={defs}", + "--test-prefix={stage}", + "--waives-file={waives}", + "--output-dir={out}/", + "--csv={out}/report.csv", + "-o", + "junit_logging=out-err", + "-p", + "no:cacheprovider", + "--unittest-markexpr={markexpr}", +) + +#: Variables a cluster session carries and CI's job does not; an MPI launcher +#: that inherits them attaches to the allocation instead of starting its own +#: world, and an abnormal exit then takes the allocation down with it. +SCRUBBED_ENV: tuple[str, ...] = ( + "SLURM_PROCID", + "SLURM_NTASKS", + "SLURM_LOCALID", + "SLURM_NODEID", + "SLURM_STEP_ID", + "SLURM_TASK_PID", + "SLURM_STEP_NUM_TASKS", +) + + +@dataclass(frozen=True) +class DbEntry: + """One test-db line, with the block it sits in.""" + + yml: str + line: str + test_id: str + terms: dict[str, Any] = field(default_factory=dict) + system_gpu_count: tuple[int | None, int | None] = (None, None) + + @property + def timeout_minutes(self) -> int | None: + match = re.search(r"TIMEOUT\s*\((\d+)\)", self.line) + return int(match.group(1)) if match else None + + +def _blocks(data: Any) -> list[dict[str, Any]]: + """The ``{condition, tests}`` blocks of one test-db file, whatever its top level is.""" + if isinstance(data, dict): + out: list[dict[str, Any]] = [] + for key, value in data.items(): + if key != "version": + out.extend(_blocks(value)) + return out + if isinstance(data, list): + if all(isinstance(item, str) for item in data): + return [{"tests": data}] + return [item for item in data if isinstance(item, dict)] + return [] + + +def find_test_db_entries(repo: str | Path, pattern: str) -> list[DbEntry]: + """Every test-db line whose test id contains ``pattern``.""" + found: list[DbEntry] = [] + for yml in sorted((Path(repo) / TEST_DB_DIR).glob("*.yml")): + try: + data = yaml.safe_load(yml.read_text(encoding="utf-8")) + except yaml.YAMLError: + continue + for block in _blocks(data): + condition = block.get("condition") or {} + if not isinstance(condition, dict): + condition = {} + gpus = ((condition.get("ranges") or {}).get("system_gpu_count")) or {} + for line in block.get("tests") or []: + line = str(line).strip() + test_id = line.split(" TIMEOUT", 1)[0].strip() + if pattern in test_id: + found.append( + DbEntry( + yml=yml.name, + line=line, + test_id=test_id, + terms=dict(condition.get("terms") or {}), + system_gpu_count=(gpus.get("gte"), gpus.get("lte")), + ) + ) + return found + + +def resolve_stages(repo: str | Path, test_id: str, *, python: str = sys.executable) -> list[str]: + """The CI stages that run ``test_id``, from the repository's own mapper. + + Empty when the mapper is missing or cannot answer; the caller then needs + an explicit ``--stage``. + """ + mapper = Path(repo) / STAGE_MAPPER + if not mapper.is_file(): + return [] + try: + out = subprocess.run( + [python, str(mapper), "--tests", test_id], + cwd=str(repo), + capture_output=True, + text=True, + timeout=300, + ) + except (OSError, subprocess.SubprocessError): + return [] + if out.returncode != 0: + return [] + return [ + line.strip() for line in out.stdout.splitlines() if re.match(r"^[\w.-]+$", line.strip()) + ] + + +@dataclass(frozen=True) +class CiCase: + """Everything needed to run one test-db line as CI would.""" + + tree: str + entry: DbEntry + stage: str + out_dir: str + models_root: str = "" + launcher: str = "" + timeout_s: int = 3600 + collect_only: bool = False + extra_env: dict[str, str] = field(default_factory=dict) + #: Put the tree itself on ``PYTHONPATH``: needed for a worktree whose + #: package is not the installed one. Empty unsets ``PYTHONPATH``, as CI. + pythonpath: str = "" + + @property + def markexpr(self) -> str: + return "cpu_only" if self.stage.startswith("CPU-") else "not cpu_only" + + +def pytest_argv(case: CiCase, test_list: str) -> list[str]: + """The pytest command line CI would build for this case.""" + tree = case.tree.rstrip("/") + fields = { + "timeout": str(case.timeout_s), + "defs": f"{tree}/{DEFS_DIR}", + "stage": case.stage, + "waives": f"{tree}/{WAIVES}", + "out": case.out_dir.rstrip("/"), + "markexpr": case.markexpr, + } + argv = ["pytest", *(flag.format(**fields) for flag in CI_PYTEST_FLAGS)] + if case.collect_only: + argv = ["pytest", "--collect-only", "-q", *argv[1:]] + return [*argv, f"--test-list={test_list}"] + + +def render_script(case: CiCase) -> str: + """A self-contained bash script that runs ``case`` and records its tree.""" + tree = case.tree.rstrip("/") + out = case.out_dir.rstrip("/") + test_list = f"{out}/test_list.txt" + argv = pytest_argv(case, test_list) + launch = shlex.split(case.launcher) if case.launcher and not case.collect_only else [] + command = " ".join(shlex.quote(a) for a in [*launch, *argv]) + if launch: + command = "env " + " ".join(f"-u {v}" for v in SCRUBBED_ENV) + " " + command + exports = { + "LLM_ROOT": tree, + "LLM_BACKEND_ROOT": f"{tree}/triton_backend", + "COLUMNS": "300", + "PYTHONUNBUFFERED": "1", + } + if case.models_root: + exports["LLM_MODELS_ROOT"] = case.models_root + exports["MODEL_CACHE_DIR"] = case.models_root + exports.update(case.extra_env) + export_lines = "\n".join(f"export {k}={shlex.quote(v)}" for k, v in exports.items()) + pythonpath_line = ( + f"export PYTHONPATH={shlex.quote(case.pythonpath)}" + if case.pythonpath + else "unset PYTHONPATH" + ) + record = shlex.quote(f"{out}/ci_case.json") + head_probe = ( + f'"HEAD=$(git -C {shlex.quote(tree)} rev-parse --short=10 HEAD) ' + f"branch=$(git -C {shlex.quote(tree)} rev-parse --abbrev-ref HEAD) " + f'dirty=$(git -C {shlex.quote(tree)} status --porcelain --untracked-files=no | wc -l)"' + ) + return f"""#!/usr/bin/env bash +# Generated by `staircase-onboard premerge ci`. Runs one test-db line as CI does. +# test-db: {case.entry.yml} +# line: {case.entry.line} +# stage: {case.stage} +# Differs from CI in: launcher={case.launcher or "(none: pytest directly)"}; +# LLM_MODELS_ROOT={case.models_root or "(inherited)"}. +set -u +mkdir -p {shlex.quote(out)} +printf '%s\\n' {shlex.quote(case.entry.line)} > {shlex.quote(test_list)} +{export_lines} +{pythonpath_line} +cd {shlex.quote(tree + "/" + DEFS_DIR)} || exit 2 +START_TREE={head_probe} +echo "CI_CASE_START $(date -u +%FT%TZ) case={case.entry.test_id} stage={case.stage} $START_TREE host=$(hostname)" +python -c "import importlib.util as u; print('CI_CASE tensorrt_llm resolves to', u.find_spec('tensorrt_llm').origin)" +{command} +rc=$? +END_TREE={head_probe} +echo "CI_CASE_END $(date -u +%FT%TZ) rc=$rc $END_TREE" +python3 - "$rc" "$START_TREE" "$END_TREE" <<'PY' > {record} +import json, sys +rc, start, end = sys.argv[1:4] +print(json.dumps({{"rc": int(rc), "start": start, "end": end, + "same_tree": start == end, "line": {json.dumps(case.entry.line)}, + "stage": {json.dumps(case.stage)}}}, indent=2)) +PY +exit $rc +""" + + +def write_case(case: CiCase) -> Path: + """Write the script (and the case description) into ``case.out_dir``.""" + out = Path(case.out_dir) + out.mkdir(parents=True, exist_ok=True) + script = out / "run_ci_case.sh" + script.write_text(render_script(case), encoding="utf-8") + script.chmod(0o755) + (out / "case.json").write_text(json.dumps(asdict(case), indent=2) + "\n", encoding="utf-8") + return script + + +def case_dir_name(entry: DbEntry) -> str: + tag = re.sub(r"[^A-Za-z0-9_]+", "-", entry.test_id.rsplit("::", 1)[-1]).strip("-") + return f"ci-{tag}-{datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')}" + + +def build_cases( + tree: str | Path, + pattern: str, + *, + out_root: str | Path, + stage: str = "", + models_root: str = "", + launcher: str = "", + collect_only: bool = False, + pythonpath: str = "", + python: str = sys.executable, +) -> list[CiCase]: + """One :class:`CiCase` per matching test-db line.""" + cases = [] + for entry in find_test_db_entries(tree, pattern): + chosen = stage or next(iter(resolve_stages(tree, entry.test_id, python=python)), "") + if not chosen: + chosen = "UNKNOWN-STAGE" + cases.append( + CiCase( + tree=str(tree), + entry=entry, + stage=chosen, + out_dir=str(Path(out_root) / case_dir_name(entry)), + models_root=models_root, + launcher=launcher, + collect_only=collect_only, + pythonpath=pythonpath, + ) + ) + return cases + + +def describe(entries: Sequence[DbEntry]) -> str: + return "\n".join(f"{e.yml}: {e.line} terms={e.terms}" for e in entries) + + +__all__ = [ + "CI_PYTEST_FLAGS", + "CiCase", + "DbEntry", + "build_cases", + "describe", + "find_test_db_entries", + "pytest_argv", + "render_script", + "resolve_stages", + "write_case", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/premerge/cli.py b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/cli.py new file mode 100644 index 000000000000..8f538fa3d782 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/cli.py @@ -0,0 +1,369 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``staircase-onboard premerge ``: the pre-merge pack, one step at a time. + +Exit codes: 0 the step passed (or, for a dry run, would pass); 1 it ran and +found something; 2 it could not run (bad arguments, dirty tree, no git). +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from datetime import datetime +from pathlib import Path + +from ..committer import ( + DEFAULT_BASE_REFS, + CommitPolicyError, + audit_commits, + resolve_base, + resolve_commit_identity, +) +from ._git import GitError, git +from .ci_case import build_cases, describe, find_test_db_entries, write_case +from .copyright import check_copyright +from .lint import LintError, run_lint +from .pack import build_pack, write_pack +from .pr_open import PrOpenError, PrRequest, open_pr +from .resign import ResignError, resign_branch +from .split import SplitError, parse_stack, split_branch + + +def _base(repo: str, explicit: str) -> str: + base = resolve_base(repo, [explicit] if explicit else list(DEFAULT_BASE_REFS)) + if not base: + print( + f"error: no base: {explicit or ' / '.join(DEFAULT_BASE_REFS)} does not resolve in {repo}", + file=sys.stderr, + ) + sys.exit(2) + return base + + +def _stamp() -> str: + return datetime.now().strftime("%Y%m%d-%H%M%S") + + +def _default_out(repo: str) -> Path: + """Inside the repository's git directory: kept, and never committable.""" + try: + return Path(git(repo, "rev-parse", "--absolute-git-dir").strip()) / "staircase-premerge" + except GitError: + return Path(repo) / ".staircase-premerge" + + +def _cmd_lint(args) -> int: + base = _base(args.repo, args.base) + try: + identity = resolve_commit_identity(args.repo) + except CommitPolicyError as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + log_dir = Path(args.log_dir or _default_out(args.repo) / f"lint-{_stamp()}") + try: + result = run_lint( + args.repo, + base=base, + log_dir=log_dir, + pre_commit=args.pre_commit, + dry_run=args.dry_run, + message=args.message, + skip_hooks=args.skip or (), + auto_skip_type_check=not args.run_type_check, + allow_semantic_changes=args.allow_semantic_changes, + identity=identity, + ) + except (LintError, GitError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + print(result.render()) + if args.copyright: + print(check_copyright(args.repo, base, "HEAD", year=datetime.now().year).render()) + return 0 if result.ok else 1 + + +def _cmd_ci(args) -> int: + if args.list: + print( + describe(find_test_db_entries(args.repo, args.case)) + or f"no test-db line contains {args.case!r}" + ) + return 0 + out_root = Path(args.out or _default_out(args.repo)) + cases = build_cases( + args.repo, + args.case, + out_root=out_root, + stage=args.stage, + models_root=args.models_root, + launcher=args.launcher, + collect_only=args.collect_only, + pythonpath=args.repo if args.pythonpath_tree else "", + ) + if len(cases) != 1: + print( + f"error: expected exactly one test-db line containing {args.case!r}, found {len(cases)}:\n" + + describe([c.entry for c in cases]), + file=sys.stderr, + ) + return 2 + script = write_case(cases[0]) + print( + f"ci: {cases[0].entry.yml}: {cases[0].entry.line}\n stage {cases[0].stage}\n script {script}" + ) + if not args.run: + print(" not run (pass --run, or submit the script where the case's GPUs are)") + return 0 + return subprocess.run(["bash", str(script)]).returncode + + +def _cmd_copyright(args) -> int: + base = _base(args.repo, args.base) + report = check_copyright(args.repo, base, args.head, year=args.year or datetime.now().year) + print(report.render()) + return 0 if report.passed else 1 + + +def _cmd_hygiene(args) -> int: + base = _base(args.repo, args.base) + findings = audit_commits( + args.repo, f"{base}..{args.head}", allow_co_authors=args.allow_co_authors + ) + for finding in findings: + print(finding.render()) + print(f"hygiene: {len(findings)} finding(s) in {base[:10]}..{args.head}") + return 0 if not findings else 1 + + +def _cmd_split(args) -> int: + try: + stacks = [parse_stack(s) for s in args.stack] + forbid = dict(f.split("=", 1) for f in args.forbid or ()) + result = split_branch( + args.repo, + base=args.base, + src=args.src, + stacks=stacks, + drop_paths=tuple(args.drop_path) if args.drop_path else ("tensorrt_llm/bindings",), + stacked=args.stacked, + force=args.force, + forbid=forbid, + ) + except (SplitError, GitError, ValueError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + print(result.render()) + return 0 if result.ok else 1 + + +def _cmd_resign(args) -> int: + try: + result = resign_branch(args.repo, src=args.src, new_branch=args.new, base=args.base) + except (ResignError, GitError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + text = result.render_markdown(args.src) + if args.map_out: + Path(args.map_out).write_text(text, encoding="utf-8") + print(f"map: {args.map_out}") + print(text.split("\n| #", 1)[0]) + return 0 if result.ok else 1 + + +def _cmd_pr_open(args) -> int: + request = PrRequest( + upstream=args.upstream, + owner=args.owner, + branch=args.branch, + title=args.title, + body=Path(args.body_file).read_text(encoding="utf-8"), + base=args.base, + draft=not args.ready, + ) + try: + report = open_pr(request, yes=args.yes, check_remote=args.check) + except PrOpenError as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + shown = dict(report) + shown["payload"] = {**report["payload"], "body": report["payload"]["body"][:300] + " ..."} + print(json.dumps(shown, indent=2)) + if report["dry_run"]: + print("pr-open: DRY RUN. Nothing was sent. Add --yes to create the pull request.") + return 0 + + +def _cmd_pack(args) -> int: + report = build_pack(args.repo, test_pattern=args.tests, base_ref=args.base) + path = write_pack(report, args.out or _default_out(args.repo)) + print(f"pack: {path}") + return 0 if report.hygiene_ok else 1 + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="staircase-onboard premerge", + description="The pre-merge pack: lint, CI-case reproduction, copyright, commit " + "audit, PR split, DCO re-sign, and an opt-in PR opener. Nothing here pushes.", + ) + sub = parser.add_subparsers(dest="step", required=True) + + def repo_arg(p): + p.add_argument("--repo", default=".", help="The checkout (default: the current directory).") + + p = sub.add_parser( + "lint", + help="pre-commit twice over the changed files; commit formatting only when pass 2 is clean", + ) + repo_arg(p) + p.add_argument( + "--base", default="", help=f"Diff base (default: merge-base with {DEFAULT_BASE_REFS[0]})." + ) + p.add_argument( + "--dry-run", action="store_true", help="Run both passes and the AST proof; commit nothing." + ) + p.add_argument("--message", default="[None][chore] Apply pre-commit formatting") + p.add_argument("--skip", action="append", help="A hook id to skip (repeatable).") + p.add_argument( + "--run-type-check", + action="store_true", + help="Do not auto-skip type-check where mypy is not importable.", + ) + p.add_argument( + "--allow-semantic-changes", + action="store_true", + help="Commit even when the AST proof finds a non-formatting change.", + ) + p.add_argument("--pre-commit", default="pre-commit") + p.add_argument("--log-dir", default="") + p.add_argument("--copyright", action="store_true", help="Also report copyright headers.") + p.set_defaults(func=_cmd_lint) + + p = sub.add_parser("ci", help="write (and with --run, run) one test-db line as CI runs it") + repo_arg(p) + p.add_argument("--case", required=True, help="Substring of exactly one test-db test id.") + p.add_argument( + "--stage", + default="", + help="CI stage name (default: from scripts/test_to_stage_mapping.py).", + ) + p.add_argument( + "--launcher", + default="", + help="Rank launcher prefix, e.g. 'mpirun -n 4 --oversubscribe --allow-run-as-root trtllm-llmapi-launch'.", + ) + p.add_argument( + "--models-root", default="", help="LLM_MODELS_ROOT for the case (default: inherited)." + ) + p.add_argument( + "--pythonpath-tree", + action="store_true", + help="Put the tree on PYTHONPATH (a worktree whose package is not the installed one).", + ) + p.add_argument("--collect-only", action="store_true") + p.add_argument( + "--out", + default="", + help="Where case directories go (default: /staircase-premerge).", + ) + p.add_argument("--list", action="store_true", help="List matching test-db lines and stop.") + p.add_argument( + "--run", action="store_true", help="Run the script here instead of only writing it." + ) + p.set_defaults(func=_cmd_ci) + + p = sub.add_parser( + "copyright", help="NVIDIA headers on new files, current year on modified ones" + ) + repo_arg(p) + p.add_argument("--base", default="") + p.add_argument("--head", default="HEAD") + p.add_argument("--year", type=int, default=0) + p.set_defaults(func=_cmd_copyright) + + p = sub.add_parser( + "hygiene", help="sign-off by author, no co-author, no symlink, no placeholder, per commit" + ) + repo_arg(p) + p.add_argument("--base", default="") + p.add_argument("--head", default="HEAD") + p.add_argument("--allow-co-authors", action="store_true") + p.set_defaults(func=_cmd_hygiene) + + p = sub.add_parser( + "split", help="split base..src into PR stacks by path, with consistency checks" + ) + repo_arg(p) + p.add_argument("--base", required=True) + p.add_argument("--src", required=True) + p.add_argument( + "--stack", + action="append", + required=True, + help="NAME=BRANCH:prefix[,prefix...]; '*' for the rest. Repeatable; order matters with --stacked.", + ) + p.add_argument( + "--stacked", action="store_true", help="Build each stack on the previous one's tip." + ) + p.add_argument("--force", action="store_true", help="Reset existing stack branches.") + p.add_argument( + "--drop-path", + action="append", + help="Remove a path from every commit (default tensorrt_llm/bindings).", + ) + p.add_argument( + "--forbid", action="append", help="NAME=REGEX that must not match in that stack's paths." + ) + p.set_defaults(func=_cmd_split) + + p = sub.add_parser("resign", help="new branch with every commit signed off by its own author") + repo_arg(p) + p.add_argument("--src", required=True) + p.add_argument("--new", required=True) + p.add_argument("--base", required=True) + p.add_argument("--map-out", default="") + p.set_defaults(func=_cmd_resign) + + p = sub.add_parser( + "pr-open", help="OPT-IN: open a draft PR via the GitHub API (dry run unless --yes)" + ) + p.add_argument("--upstream", default="NVIDIA/TensorRT-LLM") + p.add_argument( + "--owner", required=True, help="The fork's owner (the branch must already be pushed there)." + ) + p.add_argument("--branch", required=True) + p.add_argument("--title", required=True) + p.add_argument("--body-file", required=True) + p.add_argument("--base", default="main") + p.add_argument( + "--ready", action="store_true", help="Open as ready for review instead of draft." + ) + p.add_argument( + "--check", + action="store_true", + help="Read-only GETs: does the fork branch exist, is a PR open already.", + ) + p.add_argument("--yes", action="store_true", help="Actually create the pull request.") + p.set_defaults(func=_cmd_pr_open) + + p = sub.add_parser( + "pack", help="the cheap checks plus the remaining commands, as a markdown report" + ) + repo_arg(p) + p.add_argument("--base", default="") + p.add_argument("--tests", default="", help="Substring naming the model's test-db lines.") + p.add_argument("--out", default="") + p.set_defaults(func=_cmd_pack) + return parser + + +def main(argv: list[str] | None = None) -> int: + args = _parser().parse_args(argv) + return int(args.func(args) or 0) + + +__all__ = ["main"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/premerge/copyright.py b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/copyright.py new file mode 100644 index 000000000000..32b1fc3eb35b --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/copyright.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""NVIDIA copyright headers on the files a branch changes. + +New source files need an NVIDIA copyright header; modified ones need the +current year in it (the repository's AGENTS.md and CODING_GUIDELINES.md). +"Source" is a compiled or interpreted file; documents and data are not +checked. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from pathlib import Path + +from ._git import git + +SOURCE_SUFFIXES: tuple[str, ...] = ( + ".py", + ".pyi", + ".cpp", + ".cc", + ".cu", + ".cuh", + ".h", + ".hpp", + ".sh", + ".groovy", +) + +#: How far into a file a header may start (a shebang, an encoding line). +HEADER_LINES = 25 + +_HEADER = re.compile(r"Copyright[^\n]*?(\d{4})(?:\s*-\s*(\d{4}))?[^\n]*NVIDIA", re.I) + + +@dataclass +class CopyrightReport: + """Files without an NVIDIA header, and files whose header is out of date.""" + + year: int + missing: list[tuple[str, str]] = field(default_factory=list) + stale: list[tuple[str, str, str]] = field(default_factory=list) + + @property + def passed(self) -> bool: + return not self.missing and not self.stale + + def render(self) -> str: + lines = [f"== missing NVIDIA header ({len(self.missing)})"] + lines += [f" {status} {path}" for status, path in self.missing] + lines.append(f"== header year older than {self.year} ({len(self.stale)})") + lines += [f" {status} {path}: {why}" for status, path, why in self.stale] + return "\n".join(lines) + + +def check_copyright( + repo: str | Path, + base: str, + head: str = "HEAD", + *, + year: int, + suffixes: tuple[str, ...] = SOURCE_SUFFIXES, +) -> CopyrightReport: + """Check every source file ``head`` adds or modifies relative to ``base``.""" + report = CopyrightReport(year=year) + out = git(repo, "diff", "--name-status", "--no-renames", base, head) + for line in out.splitlines(): + if "\t" not in line: + continue + status, path = line.split("\t", 1) + if status.startswith("D") or not path.endswith(suffixes): + continue + text = git(repo, "show", f"{head}:{path}", check=False) + top = "\n".join(text.splitlines()[:HEADER_LINES]) + match = _HEADER.search(top) + if not match: + if status.startswith("A") or not re.search(r"Copyright", top, re.I): + report.missing.append((status, path)) + else: + report.stale.append((status, path, "non-NVIDIA or unparsed header")) + continue + years = [int(y) for y in match.groups() if y] + if max(years) < year: + report.stale.append((status, path, match.group(0)[:90])) + return report + + +__all__ = ["CopyrightReport", "SOURCE_SUFFIXES", "check_copyright"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/premerge/lint.py b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/lint.py new file mode 100644 index 000000000000..7771ced02478 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/lint.py @@ -0,0 +1,411 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The lint pass: the repository's pre-commit hooks over what the branch changed. + +Two passes over the same file set, the set CI's release check gives +pre-commit (``git diff --name-only --diff-filter=d HEAD``): + +1. **Pass 1** lets the formatters rewrite files in place. +2. **Pass 2** runs the same hooks again. Anything still failing is a real + error no formatter can fix, and nothing is committed: the tree is + restored and the formatting diff is kept as a patch. + +When pass 2 is clean, every rewritten Python file is put through an **AST +proof** — parse it before and after and compare — so a "formatting" commit +cannot carry a semantic change unnoticed. Differences a formatter +legitimately makes are classified rather than waved through: import order, +an unused import removed, whitespace inside a docstring. Anything else is +``different`` and refuses the commit unless explicitly allowed. The proof's +counts go into the commit message. + +The commit is staged by explicit path, signed off, and authored as the +repository's user, then audited with the same rules ``commit.sh`` enforces. + +Environment it copes with, all measured: + +- ``~/.cache/pre-commit`` may be read-only (an agent sandbox). The hook + environments can stay there; only the store's database and lock need + writing, so ``PRE_COMMIT_HOME`` points at a writable copy of ``db.db``. +- The ``type-check`` hook needs a Python that can import mypy and the + repository's native modules; on a login node it cannot, so it is skipped + there (and the skip is reported) unless asked otherwise. +""" + +from __future__ import annotations + +import ast +import os +import re +import shutil +import subprocess +from collections import Counter +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Sequence + +from ..committer import CommitIdentity, audit_commits +from ._git import changed_files, dirty_tracked, git + +DEFAULT_MESSAGE = "[None][chore] Apply pre-commit formatting" + +#: Classes of change the AST proof accepts in a formatting commit. +FORMATTING_CLASSES: tuple[str, ...] = ( + "unchanged", + "identical", + "docstring-whitespace", + "import-order", + "import-removed", + "not-python", +) + + +class LintError(RuntimeError): + """The lint pass cannot start (a dirty tree, no pre-commit, ...).""" + + +# ------------------------------------------------------------------ AST proof + + +def _dump(tree: ast.AST) -> str: + return ast.dump(tree, include_attributes=False) + + +def _normalize_docstrings(tree: ast.AST) -> ast.AST: + for node in ast.walk(tree): + if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)): + body = getattr(node, "body", []) + if ( + body + and isinstance(body[0], ast.Expr) + and isinstance(body[0].value, ast.Constant) + and isinstance(body[0].value.value, str) + ): + body[0].value.value = " ".join(body[0].value.value.split()) + return tree + + +def _explode_import(node: ast.Import | ast.ImportFrom) -> list[tuple[Any, ...]]: + if isinstance(node, ast.Import): + return [("import", a.name, a.asname) for a in node.names] + return [("from", node.module, node.level, a.name, a.asname) for a in node.names] + + +def _pull_imports(tree: ast.AST) -> list[tuple[Any, ...]]: + """Remove every import statement from ``tree`` and return them, exploded per name.""" + found: list[tuple[Any, ...]] = [] + for node in ast.walk(tree): + for attr in ("body", "orelse", "finalbody"): + stmts = getattr(node, attr, None) + if not isinstance(stmts, list): + continue + kept = [] + for stmt in stmts: + if isinstance(stmt, (ast.Import, ast.ImportFrom)): + found.extend(_explode_import(stmt)) + else: + kept.append(stmt) + setattr(node, attr, kept) + return found + + +def classify_python_change(before: str, after: str) -> str: + """How ``after`` differs from ``before``, as far as Python is concerned. + + ``unchanged`` / ``identical`` (same AST) / ``docstring-whitespace`` / + ``import-order`` / ``import-removed`` are what formatters do; + ``different`` is anything else, and ``unparsable`` a file one side of + which is not Python at all. + """ + if before == after: + return "unchanged" + try: + a, b = ast.parse(before), ast.parse(after) + except SyntaxError: + return "unparsable" + if _dump(a) == _dump(b): + return "identical" + a, b = _normalize_docstrings(a), _normalize_docstrings(b) + if _dump(a) == _dump(b): + return "docstring-whitespace" + imports_a, imports_b = _pull_imports(a), _pull_imports(b) + if _dump(a) != _dump(b): + return "different" + if Counter(imports_a) == Counter(imports_b): + return "import-order" + if not Counter(imports_b) - Counter(imports_a): + return "import-removed" + return "different" + + +@dataclass +class AstProof: + """The classification of every file a formatting pass rewrote.""" + + entries: list[tuple[str, str]] = field(default_factory=list) + + @property + def counts(self) -> Counter: + return Counter(cls for _, cls in self.entries) + + @property + def unexplained(self) -> list[tuple[str, str]]: + return [(p, c) for p, c in self.entries if c not in FORMATTING_CLASSES] + + @property + def ok(self) -> bool: + return not self.unexplained + + def summary(self) -> str: + return ", ".join(f"{n} {cls}" for cls, n in sorted(self.counts.items())) or "no files" + + def render(self) -> str: + lines = [f"AST proof: {self.summary()}", ""] + lines += [f"- `{path}`: {cls}" for path, cls in self.entries] + return "\n".join(lines) + "\n" + + +def prove_formatting(repo: str | Path, paths: Sequence[str], *, rev: str = "HEAD") -> AstProof: + """Compare each path at ``rev`` with its working-tree content.""" + proof = AstProof() + for path in paths: + if not path.endswith((".py", ".pyi")): + proof.entries.append((path, "not-python")) + continue + before = git(repo, "show", f"{rev}:{path}", check=False) + after_path = Path(repo) / path + after = after_path.read_text(encoding="utf-8") if after_path.is_file() else "" + proof.entries.append((path, classify_python_change(before, after))) + return proof + + +# ------------------------------------------------------------------ the pass + + +@dataclass +class LintResult: + status: str + files: list[str] = field(default_factory=list) + log_dir: Path | None = None + rewritten: list[str] = field(default_factory=list) + pass2_failures: list[str] = field(default_factory=list) + proof: AstProof | None = None + commit: str = "" + skipped_hooks: list[str] = field(default_factory=list) + notes: list[str] = field(default_factory=list) + + @property + def ok(self) -> bool: + return self.status in ("nothing", "clean", "committed", "dry-run") + + def render(self) -> str: + lines = [f"lint: {self.status} ({len(self.files)} files, {len(self.rewritten)} rewritten)"] + if self.skipped_hooks: + lines.append(f" skipped hooks: {', '.join(self.skipped_hooks)}") + if self.pass2_failures: + lines.append(f" pass 2 failing hooks: {', '.join(self.pass2_failures)}") + if self.proof is not None: + lines.append(f" {self.proof.render().splitlines()[0]}") + for path, cls in self.proof.unexplained: + lines.append(f" NOT FORMATTING: {path} ({cls})") + if self.commit: + lines.append(f" committed {self.commit[:10]}") + lines += [f" {note}" for note in self.notes] + if self.log_dir is not None: + lines.append(f" logs: {self.log_dir}") + return "\n".join(lines) + + +_FAILED_HOOK = re.compile(r"^(?P.+?)\.{3,}.*Failed\s*$") + + +def failing_hooks(log_text: str) -> list[str]: + """Hook names pre-commit reported as ``Failed``.""" + return [m.group("name").strip() for m in map(_FAILED_HOOK.match, log_text.splitlines()) if m] + + +def writable_pre_commit_home(log_dir: Path, *, home: Path | None = None) -> Path: + """A ``PRE_COMMIT_HOME`` this process can write. + + The usual store when it is writable; otherwise a fresh directory holding a + copy of the store's database, which still names the hook environments in + the read-only store by absolute path — so nothing is re-installed. + """ + store = (home or Path.home()) / ".cache" / "pre-commit" + probe = store / ".staircase-write-probe" + try: + probe.touch() + except OSError: + pass + else: + probe.unlink(missing_ok=True) + return store + fallback = log_dir / "pre-commit-home" + fallback.mkdir(parents=True, exist_ok=True) + if (store / "db.db").is_file(): + shutil.copy2(store / "db.db", fallback / "db.db") + return fallback + + +def _mypy_importable(python: str) -> bool: + try: + result = subprocess.run( + [python, "-c", "import mypy.main"], capture_output=True, text=True, timeout=120 + ) + except (OSError, subprocess.SubprocessError): + return False + return result.returncode == 0 + + +Runner = Callable[..., "subprocess.CompletedProcess[str]"] + + +def run_lint( + repo: str | Path, + *, + base: str, + log_dir: str | Path, + pre_commit: str = "pre-commit", + dry_run: bool = False, + message: str = DEFAULT_MESSAGE, + skip_hooks: Sequence[str] = (), + auto_skip_type_check: bool = True, + python: str = "python3", + allow_semantic_changes: bool = False, + identity: CommitIdentity | None = None, + runner: Runner = subprocess.run, +) -> LintResult: + """Run both passes and, when the second is clean, commit the formatting. + + Refuses to start on a tree with modified tracked files: the formatting + commit must contain only formatting. + """ + repo = Path(repo) + log = Path(log_dir) + log.mkdir(parents=True, exist_ok=True) + dirty = dirty_tracked(repo) + if dirty: + raise LintError( + "tracked files are modified; commit or stash them first:\n " + "\n ".join(dirty) + ) + files = changed_files(repo, base) + (log / "files.txt").write_text("\n".join(files) + ("\n" if files else ""), encoding="utf-8") + result = LintResult(status="nothing", files=files, log_dir=log) + if not files: + return result + + env = dict(os.environ) + env["PRE_COMMIT_HOME"] = str(writable_pre_commit_home(log)) + skips = list(skip_hooks) + if auto_skip_type_check and "type-check" not in skips and not _mypy_importable(python): + skips.append("type-check") + result.notes.append( + "type-check skipped: mypy is not importable here; run it where the repository's " + "Python is (a compute session)" + ) + if skips: + env["SKIP"] = ",".join(filter(None, [env.get("SKIP", ""), *skips])) + result.skipped_hooks = skips + + def one_pass(n: int) -> tuple[int, str]: + try: + proc = runner( + [pre_commit, "run", "--show-diff-on-failure", "--files", *files], + cwd=str(repo), + env=env, + capture_output=True, + text=True, + ) + except FileNotFoundError as exc: + raise LintError(f"{pre_commit} is not runnable: {exc}") from exc + text = (proc.stdout or "") + (proc.stderr or "") + (log / f"pass{n}.log").write_text(text, encoding="utf-8") + return proc.returncode, text + + one_pass(1) + rewritten = [line for line in git(repo, "diff", "--name-only").splitlines() if line.strip()] + result.rewritten = rewritten + rc2, text2 = one_pass(2) + + def keep_patch_and_restore() -> None: + (log / "formatting.patch").write_text(git(repo, "diff"), encoding="utf-8") + current = [line for line in git(repo, "diff", "--name-only").splitlines() if line] + if current: + git(repo, "checkout", "--", *current) + + if rc2 != 0: + result.pass2_failures = failing_hooks(text2) or [f"pre-commit exited {rc2}"] + keep_patch_and_restore() + result.status = "pass2-failed" + return result + # Pass 2 may itself have rewritten something pass 1 left (rare, but a + # formatter that is not idempotent is exactly what this guards against). + rewritten = [line for line in git(repo, "diff", "--name-only").splitlines() if line.strip()] + result.rewritten = rewritten + if not rewritten: + result.status = "clean" + return result + + proof = prove_formatting(repo, rewritten) + result.proof = proof + (log / "ast-proof.md").write_text(proof.render(), encoding="utf-8") + if dry_run: + keep_patch_and_restore() + result.status = "dry-run" + return result + if not proof.ok and not allow_semantic_changes: + keep_patch_and_restore() + result.status = "semantic-change" + return result + + body = ( + f"Output of the repository's pre-commit hooks over the {len(files)} files this " + f"branch changes; {len(rewritten)} rewritten. No hand edits.\n\n" + f"AST proof: {proof.summary()}." + ) + commit_env = dict(os.environ) + if identity is not None: + commit_env.update( + GIT_AUTHOR_NAME=identity.name, + GIT_AUTHOR_EMAIL=identity.email, + GIT_COMMITTER_NAME=identity.name, + GIT_COMMITTER_EMAIL=identity.email, + ) + git(repo, "add", "--", *rewritten) + git( + repo, + "commit", + "-s", + "-q", + "--no-verify", + "-F", + "-", + "--", + *rewritten, + env=commit_env, + input_=f"{message}\n\n{body}\n", + ) + sha = git(repo, "rev-parse", "HEAD").strip() + findings = audit_commits(repo, sha) + if findings: + git(repo, "reset", "-q", "--soft", "HEAD^") + result.status = "hygiene-failed" + result.notes += [f.render() for f in findings] + return result + result.commit = sha + result.status = "committed" + return result + + +__all__ = [ + "AstProof", + "DEFAULT_MESSAGE", + "FORMATTING_CLASSES", + "LintError", + "LintResult", + "classify_python_change", + "failing_hooks", + "prove_formatting", + "run_lint", + "writable_pre_commit_home", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/premerge/pack.py b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/pack.py new file mode 100644 index 000000000000..a02736304409 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/pack.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The pre-merge report written at the end of delivery. + +The steps that need nothing but git run here, every delivery: the commit +audit over the branch, the copyright check, and the list of test-db lines +that name the delivered model, each with the exact command that runs it the +way CI does. The steps that need a GPU, the network or a decision — the lint +pass, the CI cases themselves, the split, opening the PR — are written into +the report as commands, not run: they belong to the operator, after the +last unit lands. +""" + +from __future__ import annotations + +import shlex +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Sequence + +from ..committer import DEFAULT_BASE_REFS, CommitFinding, audit_commits, resolve_base +from ._git import GitError, git +from .ci_case import DbEntry, find_test_db_entries +from .copyright import CopyrightReport, check_copyright + +PACK_FILENAME = "premerge.md" + + +@dataclass +class PackReport: + repo: str + base: str = "" + head: str = "" + hygiene: list[CommitFinding] = field(default_factory=list) + copyright: CopyrightReport | None = None + ci_entries: list[DbEntry] = field(default_factory=list) + commands: list[str] = field(default_factory=list) + notes: list[str] = field(default_factory=list) + + @property + def hygiene_ok(self) -> bool: + return not self.hygiene + + def render_markdown(self) -> str: + out = [ + "# Pre-merge pack", + "", + f"- repository: `{self.repo}`", + f"- HEAD: `{self.head[:10] or 'unknown'}`; base: `{self.base[:10] or 'not established'}`", + f"- written: {datetime.now().isoformat(timespec='seconds')}", + "", + f"## Commit audit ({'clean' if self.hygiene_ok else f'{len(self.hygiene)} finding(s)'})", + "", + ] + out += [f"- {f.render()}" for f in self.hygiene] or [ + "Every commit is signed off by its author; no co-author, symlink or placeholder." + ] + out += ["", "## Copyright headers", ""] + if self.copyright is None: + out.append("Not checked (no base).") + else: + out += ["```", self.copyright.render(), "```"] + out += ["", f"## CI lines naming this model ({len(self.ci_entries)})", ""] + out += [f"- `{e.yml}`: `{e.line}`" for e in self.ci_entries] or [ + "None found — CI would never run this model's accuracy test." + ] + out += ["", "## Remaining steps (run by the operator, in order)", ""] + out += [f"{n}. `{c}`" for n, c in enumerate(self.commands, 1)] + if self.notes: + out += ["", "## Notes", ""] + [f"- {n}" for n in self.notes] + return "\n".join(out) + "\n" + + +def build_pack( + repo: str | Path, + *, + test_pattern: str = "", + base_ref: str = "", + year: int | None = None, + allow_co_authors: bool = False, + paths: Sequence[str] = (), +) -> PackReport: + """Run the cheap checks and assemble the operator's remaining commands. + + ``paths`` scopes the commit audit to commits touching them (a run's own + commits); empty audits every commit since the base. + """ + repo = str(repo) + report = PackReport(repo=repo) + try: + report.head = git(repo, "rev-parse", "HEAD").strip() + except GitError as exc: + report.notes.append(f"not a git checkout: {exc}") + return report + candidates = [base_ref] if base_ref else list(DEFAULT_BASE_REFS) + report.base = resolve_base(repo, candidates) + if report.base: + report.hygiene = audit_commits( + repo, f"{report.base}..HEAD", allow_co_authors=allow_co_authors, paths=paths + ) + report.copyright = check_copyright( + repo, report.base, "HEAD", year=year or datetime.now().year + ) + else: + report.notes.append( + f"no base: none of {candidates} resolves here, so the commit audit and the " + f"copyright check did not run. Set `commit_policy.base_ref`." + ) + if test_pattern: + report.ci_entries = find_test_db_entries(repo, test_pattern) + q = shlex.quote + base = report.base or "" + report.commands = [ + f"staircase-onboard premerge lint --repo {q(repo)} --base {base} --dry-run", + f"staircase-onboard premerge lint --repo {q(repo)} --base {base}", + *( + f"staircase-onboard premerge ci --repo {q(repo)} --case {q(e.test_id)} " + f"--launcher '' --out " + for e in report.ci_entries + ), + f"staircase-onboard premerge hygiene --repo {q(repo)} --base {base}", + f"staircase-onboard premerge split --repo {q(repo)} --base {base} --src HEAD " + f"--stack 'model=:*' [--stack 'name=:,...'] [--stacked]", + "staircase-onboard premerge pr-open --upstream NVIDIA/TensorRT-LLM --owner " + "--branch --title '[JIRA][type] ...' --body-file " + "(dry run; add --yes to open)", + ] + if report.hygiene: + report.commands.insert( + 0, + f"staircase-onboard premerge resign --repo {q(repo)} --src HEAD --new " + f"--base {base}", + ) + return report + + +def write_pack(report: PackReport, directory: str | Path) -> Path: + path = Path(directory) / PACK_FILENAME + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(report.render_markdown(), encoding="utf-8") + return path + + +def summarize(report: PackReport) -> Sequence[str]: + """One line per finding class, for a console.""" + lines = [ + f"commit audit: {'clean' if report.hygiene_ok else f'{len(report.hygiene)} finding(s)'}" + ] + if report.copyright is not None: + lines.append( + f"copyright: {len(report.copyright.missing)} missing, {len(report.copyright.stale)} stale" + ) + lines.append(f"CI lines naming the model: {len(report.ci_entries)}") + return lines + + +__all__ = ["PACK_FILENAME", "PackReport", "build_pack", "summarize", "write_pack"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/premerge/pr_open.py b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/pr_open.py new file mode 100644 index 000000000000..f68e0039c365 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/pr_open.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Open one draft pull request through the GitHub REST API. Opt-in only. + +**This module never pushes and never opens anything by default.** Without +``yes=True`` it validates the request and returns what it *would* send; the +network is not touched unless ``check_remote`` (read-only ``GET`` requests) +or ``yes`` (one ``POST``) is asked for explicitly. Pushing the branch to the +fork is the operator's own step, done by hand beforehand. + +It exists because the login node this was built on has no ``gh`` binary; the +token is read from ``gh``'s own ``hosts.yml`` (``$GH_CONFIG_DIR``, default +``~/.config/gh``) or from ``$GH_TOKEN`` / ``$GITHUB_TOKEN``, and is never +printed or logged. +""" + +from __future__ import annotations + +import json +import os +import re +import urllib.parse +import urllib.request +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Mapping + +import yaml + +API = "https://api.github.com" + +#: ``[JIRA/NVBUG/None][type] description`` — the repository's PR title rule. +TITLE_PATTERN = re.compile( + r"^\[(?:[A-Z][A-Z0-9]+-\d+|https://nvbugs/\d+|nvbug/\d+|None)(?:,[^\]]+)*\]" + r"\[(?:feat|fix|doc|docs|test|chore|perf|refactor|infra|ci|build|style|revert)\] \S" +) + + +class PrOpenError(ValueError): + """The request is not valid, or no token is available.""" + + +@dataclass(frozen=True) +class PrRequest: + upstream: str + owner: str + branch: str + title: str + body: str + base: str = "main" + draft: bool = True + + def payload(self) -> dict[str, Any]: + return { + "title": self.title, + "head": f"{self.owner}:{self.branch}", + "base": self.base, + "body": self.body, + "draft": self.draft, + } + + +def validate(request: PrRequest) -> list[str]: + problems = [] + if not TITLE_PATTERN.match(request.title): + problems.append( + f"title {request.title!r} does not match `[JIRA/NVBUG/None][type] description`" + ) + if not request.body.strip(): + problems.append("the body is empty") + if "/" not in request.upstream: + problems.append(f"upstream {request.upstream!r} is not OWNER/REPO") + if not request.owner or not request.branch: + problems.append("the fork owner and the branch are both required") + if re.search(r"(?im)^co-authored-by:", request.body): + problems.append("the body carries a Co-authored-by line") + return problems + + +def read_token(environ: Mapping[str, str] | None = None) -> str: + """The GitHub token, from the environment or ``gh``'s hosts file. Never printed.""" + env = os.environ if environ is None else environ + for name in ("GH_TOKEN", "GITHUB_TOKEN"): + if env.get(name): + return env[name] + config = Path( + env.get("GH_CONFIG_DIR") or Path(env.get("HOME", "~")).expanduser() / ".config/gh" + ) + hosts = config / "hosts.yml" + if not hosts.is_file(): + raise PrOpenError(f"no token: $GH_TOKEN/$GITHUB_TOKEN unset and {hosts} does not exist") + data = yaml.safe_load(hosts.read_text(encoding="utf-8")) or {} + host = data.get("github.com") or {} + user = host.get("user") + token = ((host.get("users") or {}).get(user) or {}).get("oauth_token") or host.get( + "oauth_token" + ) + if not token: + raise PrOpenError(f"no oauth_token for github.com in {hosts}") + return str(token) + + +Opener = Callable[[urllib.request.Request], Any] + + +def _call(opener: Opener, method: str, path: str, token: str, data: dict | None = None) -> Any: + request = urllib.request.Request( + f"{API}{path}", + method=method, + data=json.dumps(data).encode("utf-8") if data is not None else None, + headers={ + "Authorization": f"Bearer {token}", + "Accept": "application/vnd.github+json", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + with opener(request) as response: + return json.loads(response.read().decode("utf-8") or "null") + + +def open_pr( + request: PrRequest, + *, + yes: bool = False, + check_remote: bool = False, + token: str | None = None, + opener: Opener = urllib.request.urlopen, +) -> dict[str, Any]: + """Validate, optionally check the remote, and only with ``yes`` create the PR. + + Returns a report: ``{"dry_run": bool, "payload": ..., "checks": [...], + "url": ...}``. Raises :class:`PrOpenError` on an invalid request. + """ + problems = validate(request) + if problems: + raise PrOpenError("; ".join(problems)) + report: dict[str, Any] = {"dry_run": not yes, "payload": request.payload(), "checks": []} + if not (yes or check_remote): + return report + token = token or read_token() + fork_repo = request.upstream.split("/", 1)[1] + branch = urllib.parse.quote(request.branch, safe="") + found = _call(opener, "GET", f"/repos/{request.owner}/{fork_repo}/branches/{branch}", token) + report["checks"].append( + f"fork branch {request.owner}:{request.branch}: " + + ( + f"at {found['commit']['sha'][:10]}" + if isinstance(found, dict) and "commit" in found + else "MISSING" + ) + ) + head = urllib.parse.quote(f"{request.owner}:{request.branch}", safe=":") + existing = _call( + opener, "GET", f"/repos/{request.upstream}/pulls?state=open&head={head}", token + ) + urls = [p.get("html_url") for p in existing or [] if isinstance(p, dict)] + report["checks"].append(f"open PRs with this head: {urls or 'none'}") + if not yes: + return report + if urls: + raise PrOpenError(f"an open PR already uses this head: {urls}") + created = _call(opener, "POST", f"/repos/{request.upstream}/pulls", token, request.payload()) + report["url"] = created.get("html_url") if isinstance(created, dict) else None + report["response"] = ( + {k: created.get(k) for k in ("number", "state", "draft")} + if isinstance(created, dict) + else created + ) + return report + + +__all__ = ["PrOpenError", "PrRequest", "TITLE_PATTERN", "open_pr", "read_token", "validate"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/premerge/resign.py b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/resign.py new file mode 100644 index 000000000000..36646517046b --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/resign.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Re-sign a branch's commits for the DCO check, on a NEW branch. + +Non-destructive: the source branch is untouched, no worktree is used and no +hook runs. Every commit in ``base..src`` (linear history only) is rebuilt +with :program:`git commit-tree` on its own tree, so content cannot change; +only the message gains a trailer, and only where one is missing: + + Signed-off-by: + +and the rebuilt commit's committer is that same author. **Signing as the +operator is the mistake this exists to avoid**: the DCO check matches the +sign-off email to the author email, and one re-sign under the operator's +configured identity left 45 commits with a sign-off that matched nobody. + +A commit is rewritten when it has no ``Signed-off-by`` carrying its author's +email; every descendant of a rewritten commit gets a new hash as well. The +result maps every old hash to its new one, so records that cite the old +hashes can be translated. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from pathlib import Path + +from ..committer import audit_commits +from ._git import git + + +class ResignError(RuntimeError): + """The re-sign cannot be done safely.""" + + +@dataclass +class ResignResult: + base: str + new_branch: str + mapping: list[tuple[str, str, bool, str]] = field(default_factory=list) + unsigned_left: int = 0 + same_tree: bool = False + + @property + def rewritten(self) -> int: + return sum(1 for old, new, _, _ in self.mapping if old != new) + + @property + def ok(self) -> bool: + return self.same_tree and self.unsigned_left == 0 + + def render_markdown(self, src: str) -> str: + lines = [ + f"# DCO re-sign: `{src}` -> `{self.new_branch}`", + "", + f"- Base: `{self.base[:10]}`. Rewritten commits: {self.rewritten} of {len(self.mapping)}.", + f"- Commits still failing the sign-off audit on `{self.new_branch}`: {self.unsigned_left}. " + f"Final tree identical to `{src}`: {'yes' if self.same_tree else 'NO'}.", + f"- `{src}` is unchanged. Use this map to translate hashes cited against it.", + "", + "| # | old | new | re-signed | subject |", + "|---|---|---|---|---|", + ] + for n, (old, new, signed, subject) in enumerate(self.mapping, 1): + lines.append( + f"| {n} | `{old[:10]}` | `{new[:10]}` | {'yes' if signed else 'no'} | " + f"{subject.replace('|', chr(92) + '|')} |" + ) + return "\n".join(lines) + "\n" + + +def _needs_signoff(repo: str, commit: str) -> bool: + return any( + "Signed-off-by" in f.problem or "sign-off" in f.problem for f in audit_commits(repo, commit) + ) + + +def resign_branch(repo: str | Path, *, src: str, new_branch: str, base: str) -> ResignResult: + """Create ``new_branch`` from ``src`` with every commit signed off by its author.""" + repo = str(repo) + if git(repo, "rev-parse", "--verify", "-q", f"refs/heads/{new_branch}", check=False).strip(): + raise ResignError(f"branch {new_branch} exists; refusing to overwrite it") + base_sha = git(repo, "merge-base", src, base).strip() + src_sha = git(repo, "rev-parse", src).strip() + merges = git(repo, "rev-list", "--merges", f"{base_sha}..{src_sha}").split() + if merges: + raise ResignError( + f"{len(merges)} merge commit(s) in {base_sha[:10]}..{src}; only linear history is re-signed" + ) + commits = git(repo, "rev-list", "--reverse", f"{base_sha}..{src_sha}").split() + result = ResignResult(base=base_sha, new_branch=new_branch) + parent = base_sha + changed = False + for commit in commits: + fields = git( + repo, "log", "-1", "--format=%an%x00%ae%x00%aI%x00%cn%x00%ce%x00%cI%x00%T%x00%s", commit + ).split("\x00") + an, ae, ad, cn, ce, cd, tree, subject = (f.strip() for f in fields) + sign = _needs_signoff(repo, commit) + if not sign and not changed: + result.mapping.append((commit, commit, False, subject)) + parent = commit + continue + message = git(repo, "log", "-1", "--format=%B", commit) + if sign: + message = git( + repo, + "interpret-trailers", + "--if-exists", + "addIfDifferent", + "--trailer", + f"Signed-off-by: {an} <{ae}>", + input_=message, + ) + env = { + **os.environ, + "GIT_AUTHOR_NAME": an, + "GIT_AUTHOR_EMAIL": ae, + "GIT_AUTHOR_DATE": ad, + # Signed as the author, so the committer is the author too. + "GIT_COMMITTER_NAME": an if sign else cn, + "GIT_COMMITTER_EMAIL": ae if sign else ce, + "GIT_COMMITTER_DATE": cd, + } + new = git(repo, "commit-tree", tree, "-p", parent, env=env, input_=message).strip() + result.mapping.append((commit, new, sign, subject)) + parent = new + changed = True + git(repo, "branch", new_branch, parent) + result.unsigned_left = len( + [f for f in audit_commits(repo, f"{base_sha}..{new_branch}") if "sign" in f.problem.lower()] + ) + result.same_tree = ( + git(repo, "rev-parse", f"{src_sha}^{{tree}}").strip() + == git(repo, "rev-parse", f"{new_branch}^{{tree}}").strip() + ) + return result + + +__all__ = ["ResignError", "ResignResult", "resign_branch"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/premerge/split.py b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/split.py new file mode 100644 index 000000000000..8ca3d2a70d53 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/premerge/split.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Split a branch into reviewable PR stacks by path, commit by commit. + +A stack is a name, a branch and a set of repo-relative path prefixes; one +stack may claim ``*`` — everything no other stack claims. Every commit in +``base..src`` (first-parent order) is replayed onto the stacks its paths +belong to. A commit touching several stacks is split into one commit per +stack, each sharing its message, author and committer, so attribution and +sign-offs survive unchanged. + +Nothing touches a worktree or a checked-out branch: each commit is rebuilt +with ``git apply --cached`` into a temporary index and ``git commit-tree``. +The branches are created (or, with ``force``, reset) to point at the +results. With ``stacked``, each stack is built on the previous stack's tip, +in the order given; otherwise every stack starts at ``base``. + +``drop_paths`` removes a path from every commit (default +``tensorrt_llm/bindings``, a symlink one worktree committed by accident), so +it appears in no PR; a commit left with no change is dropped. + +Checks, every one of which must pass (``SplitResult.ok``): + +- each stack equals ``src`` inside its paths (dropped paths excepted); +- each stack equals its parent (``base``, or the previous stack) outside them; +- no stack commit touches a dropped path; +- every ``forbid`` pattern (``stack -> regex``) finds nothing in that stack's + paths at its tip — e.g. "the model PR does not import agent_flow"; +- every stack commit passes the commit audit (sign-off by author, no + co-author, no symlink, no placeholder). +""" + +from __future__ import annotations + +import os +import tempfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import Mapping, Sequence + +from ..committer import audit_commits +from ._git import git + +REST = "*" +DEFAULT_DROP: tuple[str, ...] = ("tensorrt_llm/bindings",) + + +class SplitError(ValueError): + """The stack specification is not usable.""" + + +@dataclass(frozen=True) +class Stack: + name: str + branch: str + paths: tuple[str, ...] + + @property + def is_rest(self) -> bool: + return self.paths == (REST,) + + +def parse_stack(spec: str) -> Stack: + """``NAME=BRANCH:prefix[,prefix...]``; the prefix list ``*`` means "the rest".""" + try: + name, rest = spec.split("=", 1) + branch, paths = rest.split(":", 1) + except ValueError as exc: + raise SplitError(f"stack {spec!r} is not NAME=BRANCH:prefix[,prefix...]") from exc + prefixes = tuple(p.strip() for p in paths.split(",") if p.strip()) + if not name.strip() or not branch.strip() or not prefixes: + raise SplitError(f"stack {spec!r} is missing a name, a branch or a path") + if REST in prefixes and len(prefixes) > 1: + raise SplitError(f"stack {spec!r}: '*' cannot be combined with other prefixes") + return Stack(name=name.strip(), branch=branch.strip(), paths=prefixes) + + +def _under(path: str, prefix: str) -> bool: + prefix = prefix.rstrip("/") + return path == prefix or path.startswith(prefix + "/") + + +@dataclass +class SplitResult: + tips: dict[str, str] = field(default_factory=dict) + counts: dict[str, int] = field(default_factory=dict) + split_commits: int = 0 + dropped_commits: int = 0 + checks: list[tuple[str, bool, str]] = field(default_factory=list) + + @property + def ok(self) -> bool: + return all(passed for _, passed, _ in self.checks) + + def render(self) -> str: + lines = [ + "split: " + + ", ".join( + f"{name} {self.counts.get(name, 0)} -> {tip[:10]}" + for name, tip in self.tips.items() + ) + + f"; {self.split_commits} split, {self.dropped_commits} dropped" + ] + for label, passed, detail in self.checks: + lines.append( + f" {'ok ' if passed else 'FAIL'} {label}" + ("" if passed else f"\n{detail}") + ) + return "\n".join(lines) + + +def _replay( + repo: str, tip: str, commit: str, paths: Sequence[str], env_base: Mapping[str, str] +) -> str: + """Commit ``commit``'s changes to ``paths`` on top of ``tip``; return the new commit.""" + fd, index = tempfile.mkstemp(prefix="premerge_split_index_") + os.close(fd) + os.unlink(index) + env = {**env_base, "GIT_INDEX_FILE": index} + try: + git(repo, "read-tree", tip, env=env) + patch = git( + repo, + "diff", + "--no-renames", + "--binary", + "--full-index", + f"{commit}^", + commit, + "--", + *paths, + ) + if patch: + git(repo, "apply", "--cached", "--whitespace=nowarn", env=env, input_=patch) + tree = git(repo, "write-tree", env=env).strip() + finally: + if os.path.exists(index): + os.unlink(index) + meta = git(repo, "log", "-1", "--format=%an%x00%ae%x00%aI%x00%cn%x00%ce%x00%cI", commit).split( + "\x00" + ) + message = git(repo, "log", "-1", "--format=%B", commit) + commit_env = { + **env_base, + "GIT_AUTHOR_NAME": meta[0], + "GIT_AUTHOR_EMAIL": meta[1], + "GIT_AUTHOR_DATE": meta[2], + "GIT_COMMITTER_NAME": meta[3], + "GIT_COMMITTER_EMAIL": meta[4], + "GIT_COMMITTER_DATE": meta[5].strip(), + } + return git(repo, "commit-tree", tree, "-p", tip, env=commit_env, input_=message).strip() + + +def split_branch( + repo: str | Path, + *, + base: str, + src: str, + stacks: Sequence[Stack], + drop_paths: Sequence[str] = DEFAULT_DROP, + stacked: bool = False, + force: bool = False, + forbid: Mapping[str, str] | None = None, + allow_co_authors: bool = False, +) -> SplitResult: + """Replay ``base..src`` onto ``stacks`` and verify the result.""" + repo = str(repo) + names = [s.name for s in stacks] + if len(set(names)) != len(names): + raise SplitError(f"stack names repeat: {names}") + rests = [s for s in stacks if s.is_rest] + if len(rests) > 1: + raise SplitError("at most one stack may claim '*'") + claimed = [p for s in stacks if not s.is_rest for p in s.paths] + + def owner(path: str) -> str | None: + for stack in stacks: + if not stack.is_rest and any(_under(path, p) for p in stack.paths): + return stack.name + return rests[0].name if rests else None + + env_base = dict(os.environ) + base_sha = git(repo, "rev-parse", base).strip() + src_sha = git(repo, "rev-parse", src).strip() + commits = git(repo, "rev-list", "--reverse", "--first-parent", f"{base_sha}..{src_sha}").split() + result = SplitResult( + tips={s.name: base_sha for s in stacks}, counts={s.name: 0 for s in stacks} + ) + unowned: list[str] = [] + for commit in commits: + files = [ + f + for f in git( + repo, "show", "--no-renames", "--name-only", "--format=", commit + ).splitlines() + if f and not any(_under(f, d) for d in drop_paths) + ] + if not files: + result.dropped_commits += 1 + continue + parts: dict[str, list[str]] = {} + for f in files: + who = owner(f) + if who is None: + unowned.append(f"{commit[:10]} {f}") + continue + parts.setdefault(who, []).append(f) + if len(parts) > 1: + result.split_commits += 1 + for stack in stacks: + if stack.name in parts: + result.tips[stack.name] = _replay( + repo, result.tips[stack.name], commit, parts[stack.name], env_base + ) + result.counts[stack.name] += 1 + if unowned: + raise SplitError( + "paths no stack claims (add a '*' stack, or claim them):\n " + + "\n ".join(unowned[:20]) + ) + + parents = {s.name: base_sha for s in stacks} + if stacked: + tip = base_sha + for stack in stacks: + stack_commits = git( + repo, "rev-list", "--reverse", f"{base_sha}..{result.tips[stack.name]}" + ).split() + parents[stack.name] = tip + for commit in stack_commits: + tip = _replay(repo, tip, commit, _pathspec(stack, claimed), env_base) + result.tips[stack.name] = tip + + for stack in stacks: + git(repo, "branch", *(["-f"] if force else []), stack.branch, result.tips[stack.name]) + + drop_excludes = [f":(exclude){d}" for d in drop_paths] + for stack in stacks: + inside = _pathspec(stack, claimed) + outside = _complement(stack, claimed) + tip = result.tips[stack.name] + same_inside = git(repo, "diff", "--stat", src_sha, tip, "--", *inside, *drop_excludes) + result.checks.append( + (f"{stack.name} == src inside its paths", not same_inside.strip(), same_inside) + ) + same_outside = git( + repo, "diff", "--stat", parents[stack.name], tip, "--", *outside, *drop_excludes + ) + result.checks.append( + ( + f"{stack.name} == its parent outside its paths", + not same_outside.strip(), + same_outside, + ) + ) + leaked = [ + d + for d in drop_paths + if git(repo, "log", "--format=%h", f"{parents[stack.name]}..{tip}", "--", d).strip() + ] + result.checks.append( + ( + f"{stack.name} never touches {', '.join(drop_paths) or '(nothing dropped)'}", + not leaked, + str(leaked), + ) + ) + pattern = (forbid or {}).get(stack.name) + if pattern: + hits = git(repo, "grep", "-n", "-E", pattern, tip, "--", *inside, check=False) + result.checks.append( + (f"{stack.name} has no match for /{pattern}/", not hits.strip(), hits) + ) + findings = audit_commits( + repo, f"{parents[stack.name]}..{tip}", allow_co_authors=allow_co_authors + ) + result.checks.append( + ( + f"{stack.name} commits pass the commit audit", + not findings, + "\n".join(f.render() for f in findings[:20]), + ) + ) + return result + + +def _pathspec(stack: Stack, claimed: Sequence[str]) -> list[str]: + if stack.is_rest: + return [".", *(f":(exclude){p.rstrip('/')}" for p in claimed)] + return [p.rstrip("/") for p in stack.paths] + + +def _complement(stack: Stack, claimed: Sequence[str]) -> list[str]: + if stack.is_rest: + return [p.rstrip("/") for p in claimed] or [":(exclude)."] + return [".", *(f":(exclude){p.rstrip('/')}" for p in stack.paths)] + + +__all__ = [ + "DEFAULT_DROP", + "REST", + "SplitError", + "SplitResult", + "Stack", + "parse_stack", + "split_branch", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/_common.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/_common.py index f7c70374bd99..c071a09f6da8 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/_common.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/_common.py @@ -18,6 +18,8 @@ import re from typing import Sequence +from ..embargo import PACKAGE_PREFIX + def tidy(text: str) -> str: """Normalize a composed prompt so it reads as one document. @@ -300,9 +302,18 @@ def source_embargo_block(paths: Sequence[str]) -> str: """ if not paths: return "" - listed = "\n".join(f"- `{p}`" for p in paths) - out_of_repo = [p for p in paths if p.startswith("/")] - in_repo = [p for p in paths if not p.startswith("/")] + # `package:` (see `embargo.display_entry`) is the workflow's own + # reference material — a worked example of this model — not an + # implementation of the checkpoint, so it gets its own premise and no + # "running it is expected" paragraph. + reference = [p[len(PACKAGE_PREFIX) :] for p in paths if p.startswith(PACKAGE_PREFIX)] + sources = [p for p in paths if not p.startswith(PACKAGE_PREFIX)] + listed = "\n".join( + [f"- `{p}`" for p in sources] + + [f"- `{p}` (workflow reference material)" for p in reference] + ) + out_of_repo = [p for p in sources if p.startswith("/")] + in_repo = [p for p in sources if not p.startswith("/")] premise = [] if in_repo: @@ -312,6 +323,12 @@ def source_embargo_block(paths: Sequence[str]) -> str: "This checkpoint is implemented in another engine's checkout on this " "filesystem, outside this repository." ) + if reference: + premise.append( + "This workflow ships reference material about this model — a worked example " + "of onboarding it — which you must not read: deriving the target is the " + "measurement, and an answer key would void it." + ) running = [] if in_repo: running.append( diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml b/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml index bfa584572433..361fc7455d15 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml +++ b/agent-flow/agent_flow/workflows/staircase_onboard/task.example.yaml @@ -126,10 +126,17 @@ trtllm_repo_path: /path/to/TensorRT-LLM # # 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. +# reading is forbidden, and it is enforced: a tool hook refuses every Read, +# Grep and shell read of these paths, strips their lines from broad searches, +# and logs each attempt to `/embargo/attempts.jsonl`. +# +# 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 @@ -657,3 +664,20 @@ accuracy: # # 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 +# STAIRCASE_GIT_NAME / STAIRCASE_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 diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py b/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py index ddbc795f1819..a09b5542f0d6 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/task_schema.py @@ -36,6 +36,8 @@ from agent_flow.agent_runtime import AGENTS_FIELD, validate_agents +from .committer import COMMIT_POLICY_FIELD, validate_commit_policy_block +from .embargo import PACKAGE_PREFIX, package_entry_problem, resolve_entry from .roles import ROLES, validate_role_defaults REQUIRED_PATH_FIELDS: tuple[str, ...] = ( @@ -336,6 +338,7 @@ FEATURE_MATRIX_FIELD, # Per-role backend, model and reasoning effort; see `roles.py`. AGENTS_FIELD, + COMMIT_POLICY_FIELD, ) REMEDIATION_KEYS: tuple[str, ...] = ( @@ -761,13 +764,24 @@ def _validate_source_embargo(block: Any, errors: list[str], repo: Any) -> None: """ 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, or " - f"absolute for a tree outside the repo" + 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: @@ -1571,6 +1585,7 @@ def load_and_validate_task_yaml(path: str | Path) -> dict[str, Any]: # 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. _validate_remediation(data.get(REMEDIATION_FIELD), errors) + validate_commit_policy_block(data.get(COMMIT_POLICY_FIELD), errors) reject_unknown_keys(data, TOP_LEVEL_KEYS, errors) errors.extend(validate_agents(data, ROLES)) errors.extend(validate_role_defaults(data)) diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/transport.py b/agent-flow/agent_flow/workflows/staircase_onboard/transport.py index 605b63afd387..68493f1f96eb 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/transport.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/transport.py @@ -21,6 +21,7 @@ from __future__ import annotations from pathlib import Path +from typing import Sequence from agent_flow.runtime.turns import ( DEFAULT_TURN_TIMEOUT_S, @@ -30,6 +31,8 @@ TurnTimeout, ) +from .embargo import embargo_transport_note + #: The historical name, kept because it is what the workflow, the CLI and #: their tests raise and catch. A `TurnTimeout` from any transport is the #: same event, so this is an alias rather than a parallel hierarchy. @@ -90,12 +93,17 @@ def __init__( *, poll_s: float = 5.0, timeout_s: float = DEFAULT_TURN_TIMEOUT_S, + source_embargo: Sequence[str] = (), + repo_root: Path | str | None = None, ) -> None: + # A driver-run agent is outside this process's tool hooks, so the + # embargo they enforce in process is restated here as an instruction. + embargo_note = embargo_transport_note(tuple(source_embargo), repo_root=repo_root) super().__init__( workspace, poll_s=poll_s, timeout_s=timeout_s, - transport_note=staircase_transport_note, + transport_note=lambda role: staircase_transport_note(role) + embargo_note, ) diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py b/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py index 72097830ceeb..e52b411b2504 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py @@ -53,9 +53,11 @@ from . import envelope as envelope_mod from . import gates as gates_mod +from . import overlap as overlap_mod from . import remediation -from .committer import install_committer -from .delivery import DeliveryReport, summarize_for_auditor, verify_delivery +from .committer import CommitPolicy, install_committer +from .delivery import DeliveryReport, summarize_for_auditor, verify_delivery, write_premerge_pack +from .embargo import display_entries, merge_hook_maps, source_embargo_enforcement from .environment import ( RunEnvironment, image_import_warning, @@ -300,6 +302,7 @@ def __init__( command_runner: Any | None = None, shell_runner: Any | None = None, agent_transport: Any | None = None, + overlap_units: bool = False, ) -> None: self.workspace = Path(workspace) self.task_path = self.workspace / "task.yaml" @@ -336,6 +339,10 @@ def __init__( # queue, the decision validation and every gate; only the turn # moves. self._agent_transport = agent_transport + #: Opt-in unit overlap (overlap.py): the next unit's assembler starts + #: in its own worktree once this unit is UNIT_READY. + self._overlap_units = overlap_units + self._overlap: Any = None #: How a role turn is delivered, resolved once. Both transports #: answer the same `TurnExecutor` interface, so nothing downstream #: of `_run_turn` branches on which one is in force. @@ -494,6 +501,13 @@ def _make_agent( # files with its own tools; the file is validated when it is read. hooks = None layer_required = () if self._decision_files else tuple(required_tools) + # source_embargo: a tool hook where the backend runs SDK callbacks, an + # instruction in the system prompt where it cannot (embargo.py). + embargo_hooks_map, embargo_note = source_embargo_enforcement( + self.task_data, self.workspace, role=name, backend=agent.backend + ) + hooks = merge_hook_maps(hooks, embargo_hooks_map) + system_prompt += embargo_note return AgentLayer( AgentLayerConfig( name=name, @@ -782,13 +796,17 @@ def _run_preflight(self, state: WorkflowState, log) -> bool: # followed and still produced two cross-attributed commits — so the # rule is now a script that takes a lock and refuses a commit whose # contents are not exactly what was asked for. - commit_script = install_committer(self.workspace, self.task_data["trtllm_repo_path"]) + commit_script = install_committer( + self.workspace, + self.task_data["trtllm_repo_path"], + policy=CommitPolicy.from_task(self.task_data, workspace=self.workspace), + ) if self._prompts_override is None: self.prompts = build_staircase_onboard_prompts( environment=result.environment, reviewer_mode=self.reviewer_mode, - source_embargo=self.task_data.get("source_embargo") or (), + source_embargo=display_entries(self.task_data), scope=str(self.task_data.get("scope") or "full"), decision_files=self._decision_files, commit_script=str(commit_script), @@ -1601,6 +1619,7 @@ def _run_assemble_phase(self, state: WorkflowState, log) -> bool: # BLOCKED on a vocabulary gap the run already knew about. if not self._land_deferred(state, unit, log): return False + overlap_mod.adopt(self, state, unit, log) self._progress_ctx.current_iteration = state.unit_index + 1 if not self._assemble_one_unit(state, unit, log): @@ -1842,6 +1861,7 @@ def _assemble_one_unit(self, state: WorkflowState, unit: str, log) -> bool: self._checkpoint(state) if not self._fan_out_certification_gaps(state, "assembler", unit, log): continue + overlap_mod.maybe_start(self, state, unit, log) # The configurations this unit is gated over, resolved once. A # `unified` segment is ONE codebase over a FAMILY of parallel @@ -2174,7 +2194,20 @@ def _run_deliver_phase(self, state: WorkflowState, log) -> bool: return False boot = self.gate_log.last("boot") - return self._run_remediation_cycle(state, report, boot, log) + if not self._run_remediation_cycle(state, report, boot, log): + return False + # The delivery tail: what the branch still owes before it is a PR. + pack, summary = write_premerge_pack( + self.task_data["trtllm_repo_path"], + self.workspace / "reports", + family=self._preflight.family, + policy=CommitPolicy.from_task(self.task_data, workspace=self.workspace), + ) + self.journal.section( + "pre-merge pack", "\n".join(f"- {line}" for line in summary) + f"\n\n`{pack}`" + ) + print_message(f"→ pre-merge pack: {pack} ({'; '.join(summary)})", log) + return True # ------------------------------------------------------- remediation cycle @@ -2520,6 +2553,7 @@ def _delivery_report(self, state: WorkflowState): ), model_key=self._model_key(), accuracy_test=f"test_modeling_v2_{self._preflight.family}.py", + commit_policy=CommitPolicy.from_task(self.task_data, workspace=self.workspace), ) return DeliveryReport(checks=(*report.checks, *self._certification_checks())) diff --git a/agent-flow/tests/workflows/staircase_onboard/test_committer.py b/agent-flow/tests/workflows/staircase_onboard/test_committer.py index dd3a82731c8b..ef0ce29850ce 100644 --- a/agent-flow/tests/workflows/staircase_onboard/test_committer.py +++ b/agent-flow/tests/workflows/staircase_onboard/test_committer.py @@ -16,11 +16,19 @@ import pytest from agent_flow.workflows.staircase_onboard.committer import ( + DEFAULT_ALLOWED_PATHS, SCRIPT_NAME, + CommitIdentity, + CommitPolicy, + CommitPolicyError, + audit_commits, commit_instruction, install_committer, lock_is_held, render_script, + resolve_commit_identity, + trial_commit, + validate_commit_policy_block, ) @@ -112,11 +120,11 @@ def test_concurrent_commits_all_land_and_none_contaminate(repo, commit): assert len(files) == 1 # one file per commit, no contamination -def test_a_commit_that_swept_more_than_asked_is_undone(repo, commit): - """The verification, exercised by pointing at a directory. +def test_a_directory_is_refused_before_anything_is_staged(repo, commit): + """A directory pathspec is the exact mistake the prose always warned about. - A directory pathspec is the exact mistake the prose has always warned - about. Here it stops being advice: the script refuses the result. + It used to be caught after the commit and undone; it is now refused up + front, so nothing is staged and nothing has to be rewound. """ sub = repo / "pkg" sub.mkdir() @@ -125,12 +133,41 @@ def test_a_commit_that_swept_more_than_asked_is_undone(repo, commit): before = _git(repo, "rev-parse", "HEAD").stdout.strip() result = commit("sloppy", "pkg") + assert result.returncode == 2 + assert "is a directory" in result.stderr + assert _git(repo, "rev-parse", "HEAD").stdout.strip() == before + assert _git(repo, "diff", "--cached", "--name-only").stdout.strip() == "" + assert (sub / "mine.py").exists() + + +def _sweeping_hook(repo, tmp_path, extra): + """A pre-commit hook that stages ``extra`` into the commit being made. + + The one way left for a commit to carry more than was asked once every + argument is an explicit file: something else touches the index while + git commits. The verification after the commit is what catches it. + """ + hooks = tmp_path / "hooks" + hooks.mkdir(exist_ok=True) + hook = hooks / "pre-commit" + hook.write_text(f"#!/bin/sh\ngit add -- {extra}\n") + hook.chmod(0o755) + _git(repo, "config", "core.hooksPath", str(hooks)) + + +def test_a_commit_that_swept_more_than_asked_is_undone(tmp_path, repo, commit): + (repo / "mine.py").write_text("mine\n") + (repo / "theirs.py").write_text("theirs\n") + _sweeping_hook(repo, tmp_path, "theirs.py") + before = _git(repo, "rev-parse", "HEAD").stdout.strip() + + result = commit("mine", "mine.py") assert result.returncode == 1 assert "REFUSING" in result.stderr - assert "pkg/theirs.py" in result.stderr + assert "theirs.py" in result.stderr # HEAD is back where it was, and the work is still in the tree. assert _git(repo, "rev-parse", "HEAD").stdout.strip() == before - assert (sub / "mine.py").exists() + assert (repo / "mine.py").exists() def test_a_path_that_does_not_exist_is_rejected_before_the_lock(repo, commit): @@ -256,7 +293,7 @@ def test_the_new_recipe_explicitly_supersedes_the_old_one(tmp_path, repo): assert smith.index("Commit **by explicit file path**") < smith.index("this replaces the") -def test_the_undo_rewinds_exactly_one_commit(repo, commit): +def test_the_undo_rewinds_exactly_one_commit(tmp_path, repo, commit): """Not two. The refusal message quoted the reset command inside a double-quoted @@ -271,11 +308,10 @@ def test_the_undo_rewinds_exactly_one_commit(repo, commit): _git(repo, "commit", "-q", "-m", f"history {i}") log_before = _git(repo, "log", "--format=%H").stdout.split() - sub = repo / "pkg" - sub.mkdir() - (sub / "mine.py").write_text("mine\n") - (sub / "theirs.py").write_text("theirs\n") - assert commit("sloppy", "pkg").returncode == 1 + (repo / "mine.py").write_text("mine\n") + (repo / "theirs.py").write_text("theirs\n") + _sweeping_hook(repo, tmp_path, "theirs.py") + assert commit("sloppy", "mine.py").returncode == 1 assert _git(repo, "log", "--format=%H").stdout.split() == log_before @@ -329,3 +365,418 @@ def test_every_commit_carries_a_dco_sign_off_under_gits_own_identity(repo, commi repo, "log", "-1", "--format=%(trailers:key=Signed-off-by,valueonly)" ).stdout.strip() assert trailers == "Test " + + +# ------------------------------------------------- what is refused up front + + +@pytest.fixture +def isolated_git_env(tmp_path, monkeypatch): + """No global or system git configuration, and no identity override. + + So a test that asks "what identity does git have" gets the repository's + own answer, not the developer's ~/.gitconfig. + """ + home = tmp_path / "home" + home.mkdir() + monkeypatch.setenv("HOME", str(home)) + monkeypatch.setenv("GIT_CONFIG_NOSYSTEM", "1") + monkeypatch.setenv("GIT_CONFIG_GLOBAL", str(home / ".gitconfig")) + for name in ( + "STAIRCASE_GIT_NAME", + "STAIRCASE_GIT_EMAIL", + "GIT_AUTHOR_NAME", + "GIT_AUTHOR_EMAIL", + "GIT_COMMITTER_NAME", + "GIT_COMMITTER_EMAIL", + ): + monkeypatch.delenv(name, raising=False) + return home + + +def test_the_bindings_symlink_is_never_committed(tmp_path, repo, commit): + """The accident this refusal exists for. + + A worktree links the main checkout's native build. `.gitignore` names + `tensorrt_llm/bindings/` — the directory form — so the link form is + untracked and visible, and a `git add -A` committed it. Naming it, or + any path beneath it, is refused before anything is staged; a real file + beside it still commits. + """ + (repo / ".gitignore").write_text("tensorrt_llm/bindings/\n") + _git(repo, "add", ".gitignore") + _git(repo, "commit", "-q", "-m", "ignore the build output") + build = tmp_path / "main-checkout" / "tensorrt_llm" / "bindings" + build.mkdir(parents=True) + (build / "__init__.pyi").write_text("# stub\n") + pkg = repo / "tensorrt_llm" + pkg.mkdir() + (pkg / "bindings").symlink_to(build) + (pkg / "modeling.py").write_text("x = 1\n") + # The premise: git sees the link as an ordinary untracked path. + assert "tensorrt_llm/bindings" in _git(repo, "status", "--porcelain", "-uall").stdout + + for path in ("tensorrt_llm/bindings", "tensorrt_llm/bindings/__init__.pyi"): + result = commit("add the link", path) + assert result.returncode == 2, result.stderr + assert "symlink git does not track" in result.stderr + assert _git(repo, "diff", "--cached", "--name-only").stdout.strip() == "" + + result = commit("the real file", "tensorrt_llm/modeling.py") + assert result.returncode == 0, result.stderr + assert _files_in(repo) == ["tensorrt_llm/modeling.py"] + assert "120000" not in _git(repo, "ls-tree", "-r", "HEAD").stdout + + +def test_a_tracked_symlink_may_still_be_updated(repo, commit): + """Only links git does not already track are refused.""" + (repo / "target_a.txt").write_text("a\n") + (repo / "link").symlink_to("target_a.txt") + _git(repo, "add", "target_a.txt", "link") + _git(repo, "commit", "-q", "-m", "a tracked link") + (repo / "link").unlink() + (repo / "link").symlink_to("seed.txt") + result = commit("retarget", "link") + assert result.returncode == 0, result.stderr + + +@pytest.mark.parametrize( + "placeholder", + [".mcp.json", "pkg/.mcp.json", "pkg/.claude/settings.json", "pkg/.claude/.cc-writes/x"], +) +def test_sandbox_placeholders_are_refused(repo, commit, placeholder): + """F002: one was committed; the sandbox leaves them wherever its shell goes.""" + path = repo / placeholder + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("") + result = commit("placeholder", placeholder) + assert result.returncode == 2 + assert "placeholder" in result.stderr + + +@pytest.mark.parametrize( + "path,why", + [ + ("*.py", "glob"), + ("./a.py", "plain repo-relative"), + ("pkg/../a.py", "plain repo-relative"), + ("/etc/passwd", "absolute"), + ], +) +def test_only_plain_explicit_paths_are_accepted(repo, commit, path, why): + (repo / "a.py").write_text("a\n") + (repo / "pkg").mkdir() + result = commit("bad path", path) + assert result.returncode == 2 + assert why in result.stderr + + +def test_paths_outside_the_allowed_prefixes_are_refused(tmp_path, repo): + policy = CommitPolicy(allowed_paths=("tensorrt_llm/", "tests/")) + script = install_committer(tmp_path / "ws", repo, policy=policy) + (repo / "README.md").write_text("readme\n") + (repo / "tests").mkdir() + (repo / "tests" / "test_x.py").write_text("x\n") + + refused = subprocess.run([str(script), "m", "README.md"], capture_output=True, text=True) + assert refused.returncode == 2 + assert "outside the paths this run may commit" in refused.stderr + + # One bad path refuses the whole call: nothing partial lands. + mixed = subprocess.run( + [str(script), "m", "tests/test_x.py", "README.md"], capture_output=True, text=True + ) + assert mixed.returncode == 2 + assert _git(repo, "log", "--format=%s", "-1").stdout.strip() == "seed" + + ok = subprocess.run([str(script), "m", "tests/test_x.py"], capture_output=True, text=True) + assert ok.returncode == 0, ok.stderr + + +def test_the_runs_own_workspace_is_denied_when_it_sits_inside_the_repo(tmp_path, repo): + workspace = repo / "onboard" / "workspace" + workspace.mkdir(parents=True) + policy = CommitPolicy.from_task({"trtllm_repo_path": str(repo)}, workspace=workspace) + assert policy.denied_paths == ("onboard/workspace/",) + assert policy.allowed_paths == DEFAULT_ALLOWED_PATHS + script = install_committer( + workspace, repo, policy=CommitPolicy(denied_paths=policy.denied_paths) + ) + (workspace / "journal.md").write_text("notes\n") + result = subprocess.run( + [str(script), "m", "onboard/workspace/journal.md"], capture_output=True, text=True + ) + assert result.returncode == 2 + assert "its own workspace" in result.stderr + + +def test_a_deleted_tracked_file_can_be_committed_by_path(repo, commit): + (repo / "seed.txt").unlink() + result = commit("drop seed", "seed.txt") + assert result.returncode == 0, result.stderr + assert _git(repo, "ls-tree", "-r", "--name-only", "HEAD").stdout.strip() == "" + + +# ------------------------------------------------- identity and trailers + + +def test_the_message_may_not_carry_co_authors_or_its_own_sign_off(repo, commit): + (repo / "a.py").write_text("a\n") + co = commit("[None][feat] x\n\nCo-authored-by: Someone ", "a.py") + assert co.returncode == 2 + assert "Co-authored-by" in co.stderr + so = commit("[None][feat] x\n\nSigned-off-by: Other ", "a.py") + assert so.returncode == 2 + assert "Signed-off-by" in so.stderr + + +def test_co_authors_pass_when_the_policy_allows_them(tmp_path, repo): + script = install_committer(tmp_path / "ws", repo, policy=CommitPolicy(allow_co_authors=True)) + (repo / "a.py").write_text("a\n") + result = subprocess.run( + [str(script), "[None][feat] x\n\nCo-authored-by: Pair ", "a.py"], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + + +def test_author_committer_and_sign_off_are_one_baked_identity(tmp_path, repo, monkeypatch): + """F043: the sign-off email must equal the author email, whatever the shell exports. + + An inherited GIT_AUTHOR_* / GIT_COMMITTER_* beats `-c user.*`, and a + sign-off is the committer's identity, so an operator environment that + exports either used to split them. + """ + identity = CommitIdentity(name="Repo User", email="repo.user@example.invalid") + script = install_committer(tmp_path / "ws", repo, policy=CommitPolicy(identity=identity)) + monkeypatch.setenv("GIT_AUTHOR_NAME", "login") + monkeypatch.setenv("GIT_AUTHOR_EMAIL", "login@host.invalid") + monkeypatch.setenv("GIT_COMMITTER_EMAIL", "operator@host.invalid") + monkeypatch.setenv("STAIRCASE_GIT_EMAIL", "runtime-override@host.invalid") + (repo / "a.py").write_text("a\n") + result = subprocess.run([str(script), "m", "a.py"], capture_output=True, text=True) + assert result.returncode == 0, result.stderr + fmt = "%an <%ae>|%cn <%ce>|%(trailers:key=Signed-off-by,valueonly)" + author, committer, signoff = ( + _git(repo, "log", "-1", f"--format={fmt}").stdout.strip().split("|") + ) + assert author == committer == signoff.strip() == "Repo User " + + +def test_with_no_identity_anywhere_the_commit_is_refused(tmp_path, isolated_git_env): + """No `$USER@nvidia.com` invented: that is an identity no DCO check matches.""" + repo = tmp_path / "bare-identity" + repo.mkdir() + _git(repo, "init", "-q") + (repo / "a.py").write_text("a\n") + script = install_committer(tmp_path / "ws2", repo) + result = subprocess.run([str(script), "m", "a.py"], capture_output=True, text=True) + assert result.returncode == 2 + assert "no commit identity" in result.stderr + + +def test_a_hook_that_appends_a_co_author_gets_the_commit_undone(tmp_path, repo, commit): + hooks = tmp_path / "msg-hooks" + hooks.mkdir() + hook = hooks / "commit-msg" + hook.write_text('#!/bin/sh\nprintf "\\nCo-authored-by: Bot \\n" >> "$1"\n') + hook.chmod(0o755) + _git(repo, "config", "core.hooksPath", str(hooks)) + before = _git(repo, "rev-parse", "HEAD").stdout.strip() + (repo / "a.py").write_text("a\n") + result = commit("m", "a.py") + assert result.returncode == 1 + assert "Co-authored-by" in result.stderr + assert _git(repo, "rev-parse", "HEAD").stdout.strip() == before + + +def test_identity_resolution_prefers_a_complete_override_then_the_repo(tmp_path, repo): + assert resolve_commit_identity(repo, environ={}).render() == "Test " + half = {"STAIRCASE_GIT_NAME": "Only Name"} + assert resolve_commit_identity(repo, environ=half).render() == "Test " + both = {"STAIRCASE_GIT_NAME": "Op", "STAIRCASE_GIT_EMAIL": "op@example.invalid"} + assert resolve_commit_identity(repo, environ=both).render() == "Op " + + +def test_identity_resolution_fails_loudly_without_one(tmp_path, isolated_git_env): + repo = tmp_path / "noid" + repo.mkdir() + _git(repo, "init", "-q") + with pytest.raises(CommitPolicyError, match="no commit identity"): + resolve_commit_identity(repo, environ={}) + + +# ------------------------------------------------- preflight trial + + +def test_the_trial_commit_is_clean_in_a_clean_environment(isolated_git_env): + identity = CommitIdentity(name="Repo User", email="repo.user@example.invalid") + assert trial_commit(CommitPolicy(identity=identity)) == [] + + +def test_the_trial_commit_catches_a_global_hook_that_adds_a_trailer(tmp_path, isolated_git_env): + """What a template cannot vouch for: the environment's own git configuration.""" + hooks = tmp_path / "global-hooks" + hooks.mkdir() + hook = hooks / "commit-msg" + hook.write_text('#!/bin/sh\nprintf "\\nCo-authored-by: Bot \\n" >> "$1"\n') + hook.chmod(0o755) + (isolated_git_env / ".gitconfig").write_text(f"[core]\n\thooksPath = {hooks}\n") + identity = CommitIdentity(name="Repo User", email="repo.user@example.invalid") + problems = trial_commit(CommitPolicy(identity=identity)) + assert problems and "Co-authored-by" in " ".join(problems) + + +def test_the_trial_without_an_identity_is_a_problem(): + assert trial_commit(CommitPolicy()) != [] + + +# ------------------------------------------------- a whole branch + + +def _commit_raw(repo, message, path, content="x\n", env=None): + (repo / path).parent.mkdir(parents=True, exist_ok=True) + (repo / path).write_text(content) + _git(repo, "add", path) + subprocess.run( + ["git", "-C", str(repo), "commit", "-q", "-m", message], + check=True, + capture_output=True, + env=env, + ) + return _git(repo, "rev-parse", "HEAD").stdout.strip() + + +def test_audit_commits_finds_every_kind_of_violation(tmp_path, repo): + base = _git(repo, "rev-parse", "HEAD").stdout.strip() + good = _commit_raw(repo, "good\n\nSigned-off-by: Test ", "good.py") + unsigned = _commit_raw(repo, "unsigned", "unsigned.py") + mismatched = _commit_raw( + repo, "mismatch\n\nSigned-off-by: Operator ", "mismatch.py" + ) + coauthored = _commit_raw( + repo, + "co\n\nCo-authored-by: Bot \nSigned-off-by: Test ", + "co.py", + ) + (repo / "link").symlink_to("good.py") + _git(repo, "add", "link") + _git(repo, "commit", "-q", "-s", "-m", "link") + link = _git(repo, "rev-parse", "HEAD").stdout.strip() + placeholder = _commit_raw( + repo, "ph\n\nSigned-off-by: Test ", "pkg/.mcp.json", content="" + ) + + findings = audit_commits(repo, f"{base}..HEAD") + by_sha: dict[str, list[str]] = {} + for finding in findings: + by_sha.setdefault(finding.sha, []).append(finding.problem) + assert good not in by_sha + assert by_sha[unsigned] == ["no Signed-off-by"] + assert "no sign-off carries the author email" in by_sha[mismatched][0] + assert by_sha[coauthored] == ["carries Co-authored-by"] + assert by_sha[link] == ["adds a symlink: link"] + assert by_sha[placeholder] == ["adds a sandbox placeholder: pkg/.mcp.json"] + assert audit_commits(repo, f"{base}..HEAD", allow_co_authors=True) != findings + scoped = audit_commits(repo, f"{base}..HEAD", paths=("pkg/",)) + assert {f.sha for f in scoped} == {placeholder} + + +def test_commit_policy_block_validation_reports_every_problem(): + errors: list[str] = [] + validate_commit_policy_block( + {"allowed_paths": ["/abs", "a/../b", "tests/"], "allow_co_authors": "yes", "nope": 1}, + errors, + ) + joined = "\n".join(errors) + assert "unknown key" in joined + assert "must be repo-relative" in joined + assert "'..'" in joined + assert "allow_co_authors" in joined + errors = [] + validate_commit_policy_block( + {"allowed_paths": ["tensorrt_llm/"], "base_ref": "origin/main"}, errors + ) + assert errors == [] + + +def test_the_policy_reads_the_task_block(): + policy = CommitPolicy.from_task( + { + "trtllm_repo_path": "", + "commit_policy": {"allowed_paths": ["tensorrt_llm", "setup.py"], "base_ref": "main"}, + }, + identity=CommitIdentity("A", "a@example.invalid"), + ) + assert policy.allowed_paths == ("tensorrt_llm/", "setup.py") + assert policy.base_ref == "main" + assert policy.identity.email == "a@example.invalid" + + +# ------------------------------------------------- wired into the run + + +def test_the_cli_refuses_to_start_without_an_identity(tmp_path, isolated_git_env): + from agent_flow.workflows.staircase_onboard.cli import _commit_identity + + repo = tmp_path / "noid" + repo.mkdir() + _git(repo, "init", "-q") + with pytest.raises(SystemExit) as exc: + _commit_identity({"trtllm_repo_path": str(repo)}) + assert exc.value.code == 2 + + +def test_the_cli_takes_the_repositorys_identity(repo, isolated_git_env): + from agent_flow.workflows.staircase_onboard.cli import _commit_identity + + assert _commit_identity({"trtllm_repo_path": str(repo)}).render() == "Test " + + +def test_the_run_installs_its_helper_under_the_task_policy(tmp_path, monkeypatch): + """Preflight writes `bin/commit.sh` with the task's paths and the repo's identity.""" + import json + + import yaml + + from .test_workflow import Script, _build_repo, _install, _workflow + + repo = _build_repo(tmp_path) + _git(repo, "init", "-q") + _git(repo, "config", "user.name", "Repo User") + _git(repo, "config", "user.email", "repo.user@example.invalid") + checkpoint = tmp_path / "ckpt" + checkpoint.mkdir() + (checkpoint / "config.json").write_text(json.dumps({"architectures": ["GptOssForCausalLM"]})) + 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/t", "CLAUDE_CODE_TMPDIR": "/scratch/t"}, + }, + "accuracy": { + "benchmark": "gsm8k", + "reference_yaml": "tests/integration/defs/accuracy/references/gsm8k.yaml", + "model_key": "openai/gpt-oss-120b", + "tolerance": 5.0, + }, + "commit_policy": {"allowed_paths": ["tensorrt_llm/", "tests/"]}, + } + task_file = tmp_path / "task.yaml" + task_file.write_text(yaml.safe_dump(spec), encoding="utf-8") + monkeypatch.delenv("STAIRCASE_GIT_NAME", raising=False) + monkeypatch.delenv("STAIRCASE_GIT_EMAIL", raising=False) + + workflow = _workflow(tmp_path, dry_run=True) + _install(monkeypatch, workflow, Script(None)) + with workflow: + workflow.run(task_file) + body = (tmp_path / "ws" / "bin" / "commit.sh").read_text() + assert "ALLOWED_PATHS=(tensorrt_llm/ tests/)" in body + assert "COMMIT_EMAIL=repo.user@example.invalid" in body diff --git a/agent-flow/tests/workflows/staircase_onboard/test_delivery_tail.py b/agent-flow/tests/workflows/staircase_onboard/test_delivery_tail.py new file mode 100644 index 000000000000..893de29f06ec --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_delivery_tail.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Delivery's tail: the commit audit as a delivery check, and the pre-merge pack.""" + +from __future__ import annotations + +import json +import os +import subprocess + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import delivery +from agent_flow.workflows.staircase_onboard.committer import CommitPolicy +from agent_flow.workflows.staircase_onboard.state import load_state + +from .test_delivery import TARGET_REL, _build_repo +from .test_workflow import Script, _install, _workflow, _write_vocabulary + + +def _git(repo, *args, env=None): + return subprocess.run( + ["git", "-C", str(repo), *args], capture_output=True, text=True, check=True, env=env + ).stdout + + +def _signed_repo(tmp_path): + """`_build_repo`'s trees, committed on `main`, with a `work` branch checked out.""" + repo = _build_repo(tmp_path) + _git(repo, "init", "-q", "-b", "main") + _git(repo, "config", "user.name", "Repo User") + _git(repo, "config", "user.email", "repo.user@example.invalid") + _git(repo, "config", "commit.gpgsign", "false") + _git(repo, "add", "-A") + _git(repo, "commit", "-q", "-s", "-m", "base") + _git(repo, "checkout", "-q", "-b", "work") + return repo + + +def _commit(repo, rel, text, *, signoff=True, env=None): + path = repo / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text) + _git(repo, "add", rel) + _git(repo, "commit", "-q", *(["-s"] if signoff else []), "-m", f"touch {rel}", env=env) + + +# ------------------------------------------------------------------ the check + + +def test_an_unsigned_commit_under_the_runs_paths_fails_delivery(tmp_path): + repo = _signed_repo(tmp_path) + _commit(repo, f"{TARGET_REL}/extra.py", "x\n", signoff=False) + checks = delivery.check_commit_hygiene(repo, CommitPolicy(allowed_paths=("tensorrt_llm/",))) + assert [c.name for c in checks] == ["git:commit-hygiene"] + assert not checks[0].passed + assert "no Signed-off-by" in checks[0].detail and "premerge resign" in checks[0].detail + + +def test_a_sign_off_that_is_not_the_authors_fails_delivery(tmp_path): + repo = _signed_repo(tmp_path) + env = { + **os.environ, + "GIT_AUTHOR_NAME": "login", + "GIT_AUTHOR_EMAIL": "login@host.invalid", + } + _commit( + repo, f"{TARGET_REL}/extra.py", "x\n", env=env + ) # signed as the committer, not the author + (check,) = delivery.check_commit_hygiene(repo, CommitPolicy(allowed_paths=("tensorrt_llm/",))) + assert not check.passed and "author email " in check.detail + + +def test_commits_outside_the_runs_paths_are_not_the_runs_to_answer_for(tmp_path): + repo = _signed_repo(tmp_path) + _commit(repo, "agent-flow/x.py", "x\n", signoff=False) + _commit(repo, f"{TARGET_REL}/extra.py", "x\n") + (check,) = delivery.check_commit_hygiene(repo, CommitPolicy(allowed_paths=("tensorrt_llm/",))) + assert check.passed, check.detail + # Audited over the whole repository, the same branch fails. + (whole,) = delivery.check_commit_hygiene(repo, CommitPolicy()) + assert not whole.passed + + +def test_no_git_or_no_base_makes_no_claim(tmp_path): + assert delivery.check_commit_hygiene(tmp_path, CommitPolicy()) == [] + repo = _signed_repo(tmp_path / "x") + assert delivery.check_commit_hygiene(repo, CommitPolicy(base_ref="no-such-ref")) == [] + + +def test_verify_delivery_runs_the_audit_only_when_given_a_policy(tmp_path): + repo = _signed_repo(tmp_path) + kwargs = dict( + repo_root=repo, + target_rel=TARGET_REL, + family="gpt_oss", + reference_yaml=f"{delivery.ACCURACY_REFERENCES_DIR}/gsm8k.yaml", + model_key="openai/gpt-oss-120b", + accuracy_test="test_modeling_v2_gpt_oss.py", + ) + names = [c.name for c in delivery.verify_delivery(**kwargs).checks] + assert "git:commit-hygiene" not in names + with_policy = delivery.verify_delivery(**kwargs, commit_policy=CommitPolicy()) + assert "git:commit-hygiene" in [c.name for c in with_policy.checks] + + +# ------------------------------------------------------------------ the pack + + +def test_the_pack_is_written_into_the_reports(tmp_path): + repo = _signed_repo(tmp_path) + _commit(repo, f"{TARGET_REL}/extra.py", "x = 1\n") + path, summary = delivery.write_premerge_pack(repo, tmp_path / "reports", family="gpt_oss") + assert path == tmp_path / "reports" / "premerge.md" + text = path.read_text() + assert "## Commit audit (clean)" in text + assert "premerge lint" in text and "premerge split" in text + assert any(line.startswith("copyright:") for line in summary) + # A fixture that is not a repository gets a report that says so. + plain, _ = delivery.write_premerge_pack(tmp_path / "nogit", tmp_path / "r2", family="gpt_oss") + assert "not a git work tree" in plain.read_text() + + +# ------------------------------------------------------------------ the run + + +@pytest.fixture +def git_task(tmp_path): + repo = _signed_repo(tmp_path) + checkpoint = tmp_path / "gpt-oss-120b" + checkpoint.mkdir() + (checkpoint / "config.json").write_text(json.dumps({"architectures": ["GptOssForCausalLM"]})) + 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/t", "CLAUDE_CODE_TMPDIR": "/scratch/t"}, + }, + "accuracy": { + "benchmark": "gsm8k", + "reference_yaml": "tests/integration/defs/accuracy/references/gsm8k.yaml", + "model_key": "openai/gpt-oss-120b", + "tolerance": 5.0, + }, + } + task_file = tmp_path / "task.yaml" + task_file.write_text(yaml.safe_dump(spec), encoding="utf-8") + return repo, task_file + + +def _run(tmp_path, monkeypatch, task_file): + workflow = _workflow(tmp_path) + calls = _install(monkeypatch, workflow, Script(None)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + return load_state(tmp_path / "ws" / ".staircase_onboard_state.json"), calls + + +def test_a_delivered_run_ends_with_the_pre_merge_pack(tmp_path, monkeypatch, git_task): + _, task_file = git_task + state, _ = _run(tmp_path, monkeypatch, task_file) + assert state.verdict == "delivered" + pack = tmp_path / "ws" / "reports" / "premerge.md" + assert pack.is_file() and "## Remaining steps" in pack.read_text() + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "pre-merge pack" in journal and "git:commit-hygiene" in journal + + +def test_an_unsigned_run_commit_fails_delivery(tmp_path, monkeypatch, git_task): + repo, task_file = git_task + _commit(repo, f"{TARGET_REL}/unsigned.py", "x\n", signoff=False) + state, calls = _run(tmp_path, monkeypatch, task_file) + assert state.verdict == "delivery-failed" + assert "auditor" not in [n for n, _ in calls] + assert not (tmp_path / "ws" / "reports" / "premerge.md").exists() diff --git a/agent-flow/tests/workflows/staircase_onboard/test_dspark.py b/agent-flow/tests/workflows/staircase_onboard/test_dspark.py index 8546bb9b1ae5..63500fb329a7 100644 --- a/agent-flow/tests/workflows/staircase_onboard/test_dspark.py +++ b/agent-flow/tests/workflows/staircase_onboard/test_dspark.py @@ -228,9 +228,8 @@ def test_this_checkouts_quickstart_honours_dspark_once_it_has_the_arm(): branch = _dspark_branch(QUICKSTART.read_text(encoding="utf-8")) if branch is None: pytest.skip( - "this checkout's quickstart_advanced.py has no DSPARK arm yet, so a dspark " - "unit's boot and acceptance commands run without a speculator; the " - "acceptance gate then finds no accept-rate line and fails closed" + "this checkout's quickstart_advanced.py has no DSPARK arm yet; preflight's " + "check_speculation_support refuses a dspark unit until it does" ) _assert_honours_dspark(branch) diff --git a/agent-flow/tests/workflows/staircase_onboard/test_embargo.py b/agent-flow/tests/workflows/staircase_onboard/test_embargo.py new file mode 100644 index 000000000000..34ea71020625 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_embargo.py @@ -0,0 +1,376 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``source_embargo`` as a tool-call hook: refusals, redactions, the log, the wiring.""" + +from __future__ import annotations + +import asyncio +import json + +import pytest + +from agent_flow.workflows.staircase_onboard.embargo import ( + SourceEmbargo, + embargo_hooks, + embargo_transport_note, + merge_hook_maps, + source_embargo_hooks, +) + +REL_FILE = "tensorrt_llm/_torch/models/modeling_famv4.py" + + +@pytest.fixture +def tree(tmp_path): + repo = tmp_path / "repo" + models = repo / "tensorrt_llm" / "_torch" / "models" + models.mkdir(parents=True) + (models / "modeling_famv4.py").write_text("SECRET = 'the existing implementation'\n") + (models / "modeling_famv41.py").write_text("ALSO_FINE = 1\n") + (models / "modeling_other.py").write_text("FINE = 1\n") + sibling = tmp_path / "vllm" + (sibling / "vllm" / "model_executor").mkdir(parents=True) + (sibling / "vllm" / "model_executor" / "fam.py").write_text("SECRET = 2\n") + workspace = repo / "onboard" / "ws" + workspace.mkdir(parents=True) + embargo = SourceEmbargo( + [REL_FILE, str(sibling)], + repo_root=repo, + log_path=workspace / "embargo" / "attempts.jsonl", + ) + return repo, sibling, workspace, embargo + + +# ------------------------------------------------------------------ refusals + + +def test_file_tools_are_refused_on_an_embargoed_file_or_tree(tree, tmp_path): + repo, sibling, _, embargo = tree + for tool, key in (("Read", "file_path"), ("Edit", "file_path"), ("Write", "file_path")): + assert not embargo.check_tool(tool, {key: str(repo / REL_FILE)}).allowed + assert not embargo.check_tool("NotebookEdit", {"notebook_path": str(repo / REL_FILE)}).allowed + inside = sibling / "vllm" / "model_executor" / "fam.py" + verdict = embargo.check_tool("Read", {"file_path": str(inside)}) + assert not verdict.allowed and verdict.entry == str(sibling) + assert "logged" in verdict.reason + # Relative to the session's cwd, and through a symlink. + assert not embargo.check_tool("Read", {"file_path": REL_FILE}, cwd=repo).allowed + link = tmp_path / "innocent.py" + link.symlink_to(repo / REL_FILE) + assert not embargo.check_tool("Read", {"file_path": str(link)}).allowed + + +def test_neighbours_are_not_embargoed(tree): + repo, _, _, embargo = tree + models = repo / "tensorrt_llm" / "_torch" / "models" + for name in ("modeling_famv41.py", "modeling_other.py"): + assert embargo.check_tool("Read", {"file_path": str(models / name)}).allowed + assert embargo.check_tool("Glob", {"pattern": "**/modeling_*.py"}).allowed + + +def test_grep_inside_an_embargoed_path_is_refused_and_a_broad_grep_is_not(tree): + repo, sibling, _, embargo = tree + assert not embargo.check_tool("Grep", {"pattern": "x", "path": str(sibling / "vllm")}).allowed + assert not embargo.check_tool("Grep", {"pattern": "x", "path": str(repo / REL_FILE)}).allowed + assert embargo.check_tool("Grep", {"pattern": "x", "path": str(repo)}).allowed + assert embargo.check_tool("Grep", {"pattern": "x"}).allowed + + +@pytest.mark.parametrize( + "command", + [ + "cat {repo}/" + REL_FILE, + "sed -n 1,40p " + REL_FILE, + "grep -n SECRET modeling_famv4.py", + "head {repo}/tensorrt_llm/_torch/models/modeling_fam*.py", + "python3 -c \"print(open('{repo}/" + REL_FILE + "').read())\"", + "grep -rn SECRET {sibling}", + "cd {sibling} && cat vllm/model_executor/fam.py", + "git -C {repo} show HEAD:" + REL_FILE, + "python -c 'import pathlib; print(pathlib.Path(\"{sibling}/vllm\").read_text())'", + "less < {repo}/" + REL_FILE, + ], +) +def test_shell_reads_of_an_embargoed_path_are_refused(tree, command): + repo, sibling, _, embargo = tree + verdict = embargo.check_tool( + "Bash", {"command": command.format(repo=repo, sibling=sibling)}, cwd=repo + ) + assert not verdict.allowed, command + + +@pytest.mark.parametrize( + "command", + [ + "TRTLLM_MODELING_V2=off python examples/llm-api/quickstart_advanced.py", + "cd {sibling} && python -m vllm.entrypoints.openai.api_server --port 8000", + "PYTHONPATH={sibling} python3 -m vllm.entrypoints.openai.api_server", + "cat {repo}/tensorrt_llm/_torch/models/modeling_famv41.py", + "pytest tests/unittest/test_modeling_famv4.py", + "grep -rn forward {repo}/tensorrt_llm", + ], +) +def test_running_the_implementation_and_unrelated_reads_stay_allowed(tree, command): + repo, sibling, _, embargo = tree + verdict = embargo.check_tool( + "Bash", {"command": command.format(repo=repo, sibling=sibling)}, cwd=repo + ) + assert verdict.allowed, command + + +# ------------------------------------------------------------------ redaction + + +def test_a_broad_search_has_the_embargoed_lines_removed_and_keeps_its_shape(tree): + repo, sibling, _, embargo = tree + grep_out = { + "mode": "content", + "numFiles": 2, + "filenames": [str(repo / REL_FILE), str(repo / "tensorrt_llm/x.py")], + "content": ( + f"{repo}/{REL_FILE}:1:SECRET = 'the existing implementation'\n" + f"{repo}/tensorrt_llm/x.py:3:SECRET_NOT = 0\n" + f"{sibling}/vllm/model_executor/fam.py:1:SECRET = 2\n" + ), + } + redacted, removed = embargo.redact(grep_out) + assert removed == 3 + assert set(redacted) == set(grep_out) + assert redacted["filenames"] == [str(repo / "tensorrt_llm/x.py")] + assert redacted["content"] == f"{repo}/tensorrt_llm/x.py:3:SECRET_NOT = 0\n" + assert redacted["numFiles"] == 2 # numbers are left alone; only lines and names go + + bash_out = { + "stdout": f"{REL_FILE}:1:SECRET\nother.py:2:x\n", + "stderr": "", + "interrupted": False, + } + redacted, removed = embargo.redact(bash_out) + assert ( + removed == 1 and redacted["stdout"] == "other.py:2:x\n" and redacted["interrupted"] is False + ) + assert embargo.redact({"stdout": "nothing here\n"}) == ({"stdout": "nothing here\n"}, 0) + + +# ------------------------------------------------------------------ the hooks + + +def test_the_hooks_refuse_redact_and_log_every_attempt(tree): + repo, _, workspace, embargo = tree + hooks = embargo_hooks(embargo, role="smith") + pre = hooks["PreToolUse"][0] + post = hooks["PostToolUse"][0] + assert "Read" in pre.matcher and "Bash" in pre.matcher and "Grep" in pre.matcher + + denied = asyncio.run( + pre.hooks[0]( + { + "tool_name": "Read", + "tool_input": {"file_path": str(repo / REL_FILE)}, + "cwd": str(workspace), + }, + "t1", + None, + ) + ) + assert denied["hookSpecificOutput"]["permissionDecision"] == "deny" + allowed = asyncio.run( + pre.hooks[0]( + {"tool_name": "Read", "tool_input": {"file_path": str(repo / "x.py")}}, "t2", None + ) + ) + assert allowed == {} + + response = {"stdout": f"{repo}/{REL_FILE}:1:SECRET\nkept\n", "stderr": "", "interrupted": False} + redacted = asyncio.run( + post.hooks[0]( + { + "tool_name": "Bash", + "tool_input": {"command": "grep -rn SECRET ."}, + "tool_response": response, + }, + "t3", + None, + ) + ) + out = redacted["hookSpecificOutput"] + assert out["updatedToolOutput"]["stdout"] == "kept\n" + assert "withheld 1 line" in out["additionalContext"] + + attempts = [ + json.loads(line) + for line in (workspace / "embargo" / "attempts.jsonl").read_text().splitlines() + ] + assert [a["action"] for a in attempts] == ["deny", "redact"] + assert attempts[0]["role"] == "smith" and attempts[0]["tool"] == "Read" + assert attempts[0]["entry"] == REL_FILE + assert attempts[1]["removed"] == 1 + + +def test_no_embargo_means_no_hooks_and_maps_merge_per_event(tmp_path): + assert source_embargo_hooks({"source_embargo": []}, tmp_path, role="smith") is None + assert merge_hook_maps(None, None) is None + merged = merge_hook_maps({"Stop": [1]}, {"Stop": [2], "PreToolUse": [3]}, None) + assert merged == {"Stop": [1, 2], "PreToolUse": [3]} + + +# ------------------------------------------------------------------ wiring + + +def test_every_in_process_role_gets_the_embargo_its_backend_can_enforce(tree, tmp_path): + """A hook on claude-code; on codex, which refuses SDK callbacks, the instruction.""" + from agent_flow.workflows.staircase_onboard.workflow import StaircaseOnboardWorkflow + + repo, _, _, _ = tree + workflow = StaircaseOnboardWorkflow(workspace=tmp_path / "ws") + workflow.task_data = {"trtllm_repo_path": str(repo), "source_embargo": [REL_FILE]} + smith = workflow._make_agent("smith", "system", required_tools=["append_smith_progress"]) + assert smith.config.backend.kind == "claude-code" + assert set(smith.config.backend.hooks) == {"Stop", "PreToolUse", "PostToolUse"} + assert "not enforced by a hook" not in smith.config.system_prompt + + reviewer = workflow._make_agent("plan_reviewer", "system", required_tools=[]) + assert reviewer.config.backend.kind == "codex" + assert not reviewer.config.backend.hooks + assert "not enforced by a hook here" in reviewer.config.system_prompt + assert "`codex` backend" in reviewer.config.system_prompt + assert f"`{repo / REL_FILE}`" in reviewer.config.system_prompt + attempts = SourceEmbargo.from_task(workflow.task_data, tmp_path / "ws").attempts() + assert [(a["action"], a["role"]) for a in attempts] == [("instruct", "plan_reviewer")] + + workflow.task_data = {"trtllm_repo_path": str(repo), "source_embargo": []} + plain = workflow._make_agent("smith", "system", required_tools=["append_smith_progress"]) + assert set(plain.config.backend.hooks) == {"Stop"} + assert ( + workflow._make_agent("plan_reviewer", "system", required_tools=[]).config.system_prompt + == "system" + ) + + +def test_a_codex_role_routed_by_the_task_is_instructed_too(tree, tmp_path): + from agent_flow.workflows.staircase_onboard.embargo import source_embargo_enforcement + + repo, _, workspace, _ = tree + task = {"trtllm_repo_path": str(repo), "source_embargo": [REL_FILE]} + hooks, note = source_embargo_enforcement(task, workspace, role="smith", backend="codex") + assert hooks is None and "`codex` backend" in note + hooks, note = source_embargo_enforcement(task, workspace, role="smith", backend="claude-code") + assert set(hooks) == {"PreToolUse", "PostToolUse"} and note == "" + + +def test_package_relative_entries_name_the_workflows_own_files(tree, tmp_path): + """`package:` is relative to the workflow package, wherever it is installed.""" + from agent_flow.workflows.staircase_onboard.embargo import ( + PACKAGE_ROOT, + display_entry, + resolve_entry, + ) + + repo, _, workspace, _ = tree + entry = "package:task.example.yaml" + assert resolve_entry(entry) == PACKAGE_ROOT / "task.example.yaml" + embargo = SourceEmbargo([entry], repo_root=repo, log_path=workspace / "e.jsonl") + target = PACKAGE_ROOT / "task.example.yaml" + assert not embargo.check_tool("Read", {"file_path": str(target)}).allowed + assert not embargo.check_tool("Bash", {"command": f"cat {target}"}).allowed + assert embargo.check_tool("Read", {"file_path": str(PACKAGE_ROOT / "README.md")}).allowed + # Outside the repository it is shown by its absolute path; inside, repo-relative. + # The prefix is kept for the prompt's own wording; outside the repository + # the path is absolute, inside it repo-relative. + assert display_entry(entry, repo) == f"package:{target}" + assert ( + display_entry(entry, PACKAGE_ROOT.parent) == "package:staircase_onboard/task.example.yaml" + ) + assert display_entry(REL_FILE, repo) == REL_FILE + + +def test_the_schema_accepts_package_entries_that_exist_and_only_those(tree): + from agent_flow.workflows.staircase_onboard.task_schema import _validate_source_embargo + + repo, _, _, _ = tree + errors: list[str] = [] + _validate_source_embargo(["package:task.example.yaml", REL_FILE], errors, str(repo)) + assert errors == [] + _validate_source_embargo(["package:examples/no-such-model.md"], errors, str(repo)) + assert len(errors) == 1 and "workflow package" in errors[0] + + +def test_the_external_transport_restates_the_embargo_as_an_instruction(tree, tmp_path): + from agent_flow.workflows.staircase_onboard.transport import ExternalAgentTransport + + repo, sibling, _, _ = tree + transport = ExternalAgentTransport( + tmp_path / "ws", poll_s=0.01, source_embargo=[REL_FILE, str(sibling)], repo_root=repo + ) + note = transport._transport_note("assembler") + assert "append_assembler_progress" in note # the existing note is kept + assert "not enforced by a hook here" in note and "external transport" in note + assert f"`{repo}/{REL_FILE}`" in note and f"`{sibling}`" in note + assert "Embargo exposures" in note + bare = ExternalAgentTransport(tmp_path / "ws2", poll_s=0.01) + assert "embargo" not in bare._transport_note("assembler").lower() + assert embargo_transport_note([]) == "" + + +def test_the_embargo_is_built_from_the_task(tree): + repo, sibling, workspace, _ = tree + embargo = SourceEmbargo.from_task( + {"trtllm_repo_path": str(repo), "source_embargo": [REL_FILE, str(sibling)]}, workspace + ) + assert embargo is not None + assert embargo.log_path == workspace / "embargo" / "attempts.jsonl" + assert [e.is_dir for e in embargo.entries] == [False, True] + assert SourceEmbargo.from_task({"source_embargo": []}, workspace) is None + + +@pytest.mark.parametrize( + "entry,why", + [ + ("package:../perf_analyze", "outside the workflow package"), + ("package:examples/../../perf_analyze", "outside the workflow package"), + ("package:../../workflows", "outside the workflow package"), + ("package:", "whole workflow package"), + ("package:.", "whole workflow package"), + ], +) +def test_a_package_entry_may_not_escape_the_package(tree, entry, why): + """`package:../perf_analyze` resolved to the sibling workflow and passed validation.""" + from agent_flow.workflows.staircase_onboard.embargo import EmbargoEntryError, resolve_entry + from agent_flow.workflows.staircase_onboard.task_schema import _validate_source_embargo + + repo, _, _, _ = tree + errors: list[str] = [] + _validate_source_embargo([entry], errors, str(repo)) + assert len(errors) == 1 and why in errors[0], errors + with pytest.raises(EmbargoEntryError, match=why): + resolve_entry(entry) + with pytest.raises(EmbargoEntryError): + SourceEmbargo([entry], repo_root=repo) + + +def test_the_role_prompt_calls_a_package_entry_reference_material(tree): + """Not "this checkpoint already has an implementation in this repository".""" + from agent_flow.workflows.staircase_onboard.embargo import PACKAGE_ROOT, display_entries + from agent_flow.workflows.staircase_onboard.prompts import source_embargo_block + + repo, _, _, _ = tree + inside = { + "trtllm_repo_path": str(PACKAGE_ROOT.parent), + "source_embargo": ["package:task.example.yaml"], + } + block = source_embargo_block(display_entries(inside)) + assert "- `staircase_onboard/task.example.yaml` (workflow reference material)" in block + assert "reference material about this model" in block + assert "already has an implementation in this repository" not in block + assert "Running the in-repository implementation" not in block + + mixed = { + "trtllm_repo_path": str(repo), + "source_embargo": [REL_FILE, "package:task.example.yaml"], + } + block = source_embargo_block(display_entries(mixed)) + assert "already has an implementation in this repository" in block + assert "reference material about this model" in block + assert f"- `{REL_FILE}`\n" in block + assert f"- `{PACKAGE_ROOT / 'task.example.yaml'}` (workflow reference material)" in block diff --git a/agent-flow/tests/workflows/staircase_onboard/test_overlap.py b/agent-flow/tests/workflows/staircase_onboard/test_overlap.py new file mode 100644 index 000000000000..5ccd7af81ddb --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_overlap.py @@ -0,0 +1,355 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``--overlap-units``: the git mechanics against real repositories, then the run. + +Every fallback is asserted to leave the main checkout exactly where it was: +the whole point of the opt-in is that a bad early start costs nothing but +the turn it spent. +""" + +from __future__ import annotations + +import json +import subprocess + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import overlap as overlap_mod +from agent_flow.workflows.staircase_onboard.committer import CommitIdentity, CommitPolicy +from agent_flow.workflows.staircase_onboard.overlap import UnitOverlap, native_artifacts +from agent_flow.workflows.staircase_onboard.state import load_state + +from .test_delivery import TARGET_REL, _build_repo +from .test_workflow import Script, _install, _workflow, _write_vocabulary + +IDENTITY = CommitIdentity(name="Repo User", email="repo.user@example.invalid") + + +def _git(repo, *args, check=True): + return subprocess.run( + ["git", "-C", str(repo), *args], capture_output=True, text=True, check=check + ).stdout + + +def _init(repo): + _git(repo, "init", "-q", "-b", "main") + _git(repo, "config", "user.name", IDENTITY.name) + _git(repo, "config", "user.email", IDENTITY.email) + _git(repo, "config", "commit.gpgsign", "false") + + +def _commit_all(repo, message): + _git(repo, "add", "-A") + _git(repo, "commit", "-q", "-s", "-m", message) + return _git(repo, "rev-parse", "HEAD").strip() + + +@pytest.fixture +def main_repo(tmp_path): + repo = tmp_path / "main" + (repo / "tensorrt_llm").mkdir(parents=True) + _init(repo) + (repo / ".gitignore").write_text("tensorrt_llm/bindings/\ntensorrt_llm/*.so\n") + (repo / "tensorrt_llm" / "__init__.py").write_text("") + (repo / "tensorrt_llm" / "shared.py").write_text("A = 1\n") + _commit_all(repo, "seed") + # The native build: ignored in the main checkout. + (repo / "tensorrt_llm" / "bindings").mkdir() + (repo / "tensorrt_llm" / "bindings" / "__init__.pyi").write_text("") + (repo / "tensorrt_llm" / "libnative.so").write_text("elf") + return repo + + +@pytest.fixture +def overlap(tmp_path, main_repo): + return UnitOverlap(main_repo, tmp_path / "ws", policy=CommitPolicy(identity=IDENTITY)) + + +def _commit_in(job, rel, text): + path = job.tree / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text) + result = subprocess.run( + [str(job.commit_script), f"[None][feat] {rel}", rel], capture_output=True, text=True + ) + assert result.returncode == 0, result.stderr + + +# ------------------------------------------------------------------ mechanics + + +def test_the_native_build_is_found_and_linked_into_the_worktree(overlap, main_repo): + assert native_artifacts(main_repo) == ["tensorrt_llm/bindings", "tensorrt_llm/libnative.so"] + job = overlap.open_worktree("mtp", after="core") + assert job.base == _git(main_repo, "rev-parse", "HEAD").strip() + assert job.linked == ["tensorrt_llm/bindings", "tensorrt_llm/libnative.so"] + link = job.tree / "tensorrt_llm" / "bindings" + assert ( + link.is_symlink() and link.resolve() == (main_repo / "tensorrt_llm" / "bindings").resolve() + ) + # The link is exactly the thing the worktree's commit helper refuses. + refused = subprocess.run( + [str(job.commit_script), "m", "tensorrt_llm/bindings"], capture_output=True, text=True + ) + assert refused.returncode == 2 and "symlink git does not track" in refused.stderr + with pytest.raises(overlap_mod.OverlapError, match="already exists"): + overlap.open_worktree("mtp", after="core") + + +def test_an_early_result_on_an_unmoved_head_fast_forwards(overlap, main_repo): + job = overlap.open_worktree("mtp", after="core") + _commit_in(job, "tensorrt_llm/mtp.py", "MTP = 1\n") + job.decision = "UNIT_READY" + outcome = overlap.integrate(job) + assert outcome.adopted, outcome.detail + assert (main_repo / "tensorrt_llm" / "mtp.py").read_text() == "MTP = 1\n" + assert _git( + main_repo, "log", "-1", "--format=%an <%ae>|%(trailers:key=Signed-off-by,valueonly)" + ).strip() == (f"{IDENTITY.render()}|{IDENTITY.render()}") + overlap.remove(job) + assert not job.tree.exists() + assert job.branch not in _git(main_repo, "branch", "--list") + + +def test_an_early_result_rebases_onto_a_head_that_moved(overlap, main_repo): + job = overlap.open_worktree("mtp", after="core") + _commit_in(job, "tensorrt_llm/mtp.py", "MTP = 1\n") + # core's later turn lands on the main checkout meanwhile. + (main_repo / "tensorrt_llm" / "core_fix.py").write_text("FIX = 1\n") + moved = _commit_all(main_repo, "core fix") + job.decision = "UNIT_READY" + outcome = overlap.integrate(job) + assert outcome.adopted, outcome.detail + assert _git(main_repo, "rev-parse", "HEAD~1").strip() == moved + assert (main_repo / "tensorrt_llm" / "core_fix.py").exists() + assert (main_repo / "tensorrt_llm" / "mtp.py").exists() + + +def test_a_conflict_falls_back_and_leaves_the_main_checkout_alone(overlap, main_repo): + job = overlap.open_worktree("mtp", after="core") + _commit_in(job, "tensorrt_llm/shared.py", "A = 'mtp'\n") + early_tip = _git(job.tree, "rev-parse", "HEAD").strip() + (main_repo / "tensorrt_llm" / "shared.py").write_text("A = 'core'\n") + head = _commit_all(main_repo, "core changes the same line") + job.decision = "UNIT_READY" + outcome = overlap.integrate(job) + assert not outcome.adopted and "do not rebase" in outcome.detail + assert _git(main_repo, "rev-parse", "HEAD").strip() == head + assert (main_repo / "tensorrt_llm" / "shared.py").read_text() == "A = 'core'\n" + assert _git(job.tree, "rev-parse", "HEAD").strip() == early_tip # rebase aborted + note = overlap.keep_for_inspection(job, outcome.detail) + assert "not adopted" in note.read_text() + assert "A = 'mtp'" in (job.root / "early.patch").read_text() + assert job.tree.exists() # kept, not deleted + + +@pytest.mark.parametrize( + "setup,expected", + [ + ("blocked", "ended BLOCKED"), + ("nothing", "committed nothing"), + ("dirty", "uncommitted work"), + ("error", "the early turn failed"), + ], +) +def test_every_other_early_outcome_falls_back(overlap, main_repo, setup, expected): + job = overlap.open_worktree("mtp", after="core") + head = _git(main_repo, "rev-parse", "HEAD").strip() + job.decision = "UNIT_READY" + if setup == "blocked": + _commit_in(job, "tensorrt_llm/mtp.py", "x\n") + job.decision = "BLOCKED" + elif setup == "dirty": + _commit_in(job, "tensorrt_llm/mtp.py", "x\n") + (job.tree / "tensorrt_llm" / "shared.py").write_text("A = 'half-written'\n") + elif setup == "error": + job.error = "TurnTimeout: nobody answered" + outcome = overlap.integrate(job) + assert not outcome.adopted and expected in outcome.detail + assert _git(main_repo, "rev-parse", "HEAD").strip() == head + + +def test_the_early_turn_runs_in_a_thread_and_its_failure_is_reported(overlap): + def boom(job, prompt): + raise RuntimeError("the driver went away") + + overlap.turn_runner = boom + job = overlap.open_worktree("mtp", after="core") + overlap.start(job, "prompt") + overlap.wait("mtp", timeout=30) + assert job.error == "RuntimeError: the driver went away" + assert not overlap.integrate(job).adopted + + +# ------------------------------------------------------------------ the run + + +@pytest.fixture +def run_setup(tmp_path): + repo = _build_repo(tmp_path) + _init(repo) + _commit_all(repo, "the delivered trees") + checkpoint = tmp_path / "gpt-oss-120b" + checkpoint.mkdir() + (checkpoint / "config.json").write_text(json.dumps({"architectures": ["GptOssForCausalLM"]})) + spec = { + "checkpoint_path": str(checkpoint), + "arch": "sm_103", + "parallel": "tp1", + "trtllm_repo_path": str(repo), + "target": {"family": "gpt_oss", "checkpoint_segment": "gpt_oss_120b"}, + "features": ["mtp"], + "environment": { + "kind": "local", + "python": "/opt/venv/bin/python", + "env": {"TMPDIR": "/scratch/t", "CLAUDE_CODE_TMPDIR": "/scratch/t"}, + }, + "accuracy": { + "benchmark": "gsm8k", + "reference_yaml": "tests/integration/defs/accuracy/references/gsm8k.yaml", + "model_key": "openai/gpt-oss-120b", + "tolerance": 5.0, + }, + } + task_file = tmp_path / "task.yaml" + task_file.write_text(yaml.safe_dump(spec), encoding="utf-8") + return repo, task_file + + +def _run(tmp_path, monkeypatch, task_file, early): + workflow = _workflow(tmp_path, overlap_units=True) + calls = _install(monkeypatch, workflow, Script(None)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + monkeypatch.setattr(overlap_mod, "run_early_turn", lambda wf, job, prompt: early(job, prompt)) + with workflow: + workflow.run(task_file) + return workflow, calls + + +def test_the_next_unit_is_assembled_early_and_adopted(tmp_path, monkeypatch, run_setup): + repo, task_file = run_setup + seen = {} + + def early(job, prompt): + seen["prompt"] = prompt + _commit_in(job, f"{TARGET_REL}/mtp_head.py", "MTP = 1\n") + return {"decision": "UNIT_READY", "summary": "mtp head, early"} + + workflow, calls = _run(tmp_path, monkeypatch, task_file, early) + + assemblers = [p for n, p in calls if n == "assembler"] + assert len(assemblers) == 1 and "Unit: **core**" in assemblers[0] + assert "Early start: `mtp` overlaps `core`" in seen["prompt"] + assert "Unit: **mtp**" in seen["prompt"] + state = load_state(tmp_path / "ws" / ".staircase_onboard_state.json") + assert state.units_done == ["core", "mtp"] + assert (repo / TARGET_REL / "mtp_head.py").read_text() == "MTP = 1\n" + assert not (tmp_path / "ws" / "overlap" / "mtp" / "tree").exists() + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "overlap started" in journal and "overlap adopted" in journal + progress = yaml.safe_load((tmp_path / "ws" / "progress.yaml").read_text()) + adopted = [ + e + for entries in progress.values() + if isinstance(entries, list) + for e in entries + if e.get("overlap") + ] + assert adopted and adopted[0]["unit"] == "mtp" + + +def test_an_early_start_that_is_not_ready_falls_back_to_sequential( + tmp_path, monkeypatch, run_setup +): + repo, task_file = run_setup + head_before = _git(repo, "rev-parse", "HEAD").strip() + + def early(job, prompt): + return {"decision": "BLOCKED", "summary": "needs core's final shape"} + + workflow, calls = _run(tmp_path, monkeypatch, task_file, early) + + assemblers = [p for n, p in calls if n == "assembler"] + assert len(assemblers) == 2 and "Unit: **mtp**" in assemblers[1] + state = load_state(tmp_path / "ws" / ".staircase_onboard_state.json") + assert state.units_done == ["core", "mtp"] + assert _git(repo, "rev-parse", "HEAD").strip() == head_before + assert (tmp_path / "ws" / "overlap" / "mtp" / "NOT_ADOPTED.md").is_file() + assert "overlap fell back" in (tmp_path / "ws" / "journal.md").read_text() + + +def test_without_the_flag_nothing_overlaps(tmp_path, monkeypatch, run_setup): + _, task_file = run_setup + workflow = _workflow(tmp_path) + calls = _install(monkeypatch, workflow, Script(None)) + original = workflow._run_plan_phase + + def plan_then_write(state, log): + (tmp_path / "ws").mkdir(parents=True, exist_ok=True) + _write_vocabulary(tmp_path / "ws", []) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_plan_phase", plan_then_write) + with workflow: + workflow.run(task_file) + assert len([n for n, _ in calls if n == "assembler"]) == 2 + assert not (tmp_path / "ws" / "overlap").exists() + + +def test_the_early_turn_goes_through_the_runs_transport_as_the_assembler( + tmp_path, overlap, main_repo +): + """The real turn runner: own decision path, own progress, the role's backend and model.""" + from agent_flow.runtime import TurnOutcome + + seen = {} + + class Transport: + def run(self, request): + seen["request"] = request + from pathlib import Path + + Path(request.decision_file).write_text("decision: UNIT_READY\nsummary: early\n") + return TurnOutcome(role=request.role, text="done") + + def bound(self, layer): + return self + + class Journal: + def record_dispatch(self, role, what, prompt): + seen["what"] = what + return "report" + + def record_result(self, report, text): + seen["result"] = text + + class Prompts: + assembler = "assembler system prompt" + + class Workflow: + task_data = {"trtllm_repo_path": str(main_repo)} + prompts = Prompts() + journal = Journal() + _decision_files = True + _agent_transport = object() + _turns = Transport() + + job = overlap.open_worktree("mtp", after="core") + entry = overlap_mod.run_early_turn(Workflow(), job, "the prompt") + request = seen["request"] + assert entry["decision"] == "UNIT_READY" + assert request.role == "assembler" and request.backend == "claude-code" + assert request.decision_file.startswith(str(job.root)) + assert str(job.root) in request.prompt and request.system_prompt == "assembler system prompt" + assert seen["what"] == "mtp (early start, overlapping core)" + # Written to the job's own progress file, never the run's. + assert (job.root / "progress.yaml").is_file() diff --git a/agent-flow/tests/workflows/staircase_onboard/test_premerge.py b/agent-flow/tests/workflows/staircase_onboard/test_premerge.py new file mode 100644 index 000000000000..109cbb5caad7 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_premerge.py @@ -0,0 +1,613 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The pre-merge pack, against real git repositories in temporary directories. + +pre-commit, pytest and the GitHub API are stood in for by small fakes; git is +real, because every step here is a statement about what git will record. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import textwrap +from pathlib import Path + +import pytest + +from agent_flow.workflows.staircase_onboard.committer import CommitIdentity +from agent_flow.workflows.staircase_onboard.premerge import ci_case, pack +from agent_flow.workflows.staircase_onboard.premerge.copyright import check_copyright +from agent_flow.workflows.staircase_onboard.premerge.lint import ( + LintError, + classify_python_change, + run_lint, + writable_pre_commit_home, +) +from agent_flow.workflows.staircase_onboard.premerge.pr_open import ( + PrOpenError, + PrRequest, + open_pr, + read_token, + validate, +) +from agent_flow.workflows.staircase_onboard.premerge.resign import ResignError, resign_branch +from agent_flow.workflows.staircase_onboard.premerge.split import ( + SplitError, + parse_stack, + split_branch, +) + +AUTHOR = "Repo User " +IDENTITY = CommitIdentity(name="Repo User", email="repo.user@example.invalid") + + +def _git(repo, *args, env=None, check=True): + return subprocess.run( + ["git", "-C", str(repo), *args], capture_output=True, text=True, check=check, env=env + ).stdout + + +def _env(name="Repo User", email="repo.user@example.invalid", committer=None): + env = dict(os.environ) + cn, ce = committer or (name, email) + env.update( + GIT_AUTHOR_NAME=name, + GIT_AUTHOR_EMAIL=email, + GIT_COMMITTER_NAME=cn, + GIT_COMMITTER_EMAIL=ce, + ) + return env + + +def _commit(repo, message, files, *, signoff=True, env=None): + for path, content in files.items(): + target = repo / path + target.parent.mkdir(parents=True, exist_ok=True) + if content is None: + _git(repo, "rm", "-q", path) + continue + target.write_text(content) + _git(repo, "add", path) + args = ["commit", "-q", "-m", message] + (["-s"] if signoff else []) + _git(repo, *args, env=env or _env()) + return _git(repo, "rev-parse", "HEAD").strip() + + +@pytest.fixture +def repo(tmp_path): + path = tmp_path / "repo" + path.mkdir() + _git(path, "init", "-q", "-b", "main") + _git(path, "config", "user.name", "Repo User") + _git(path, "config", "user.email", "repo.user@example.invalid") + _git(path, "config", "commit.gpgsign", "false") + _commit(path, "seed", {"README.md": "seed\n"}) + _git(path, "checkout", "-q", "-b", "work") + return path + + +# ------------------------------------------------------------------ AST proof + + +@pytest.mark.parametrize( + "before,after,expected", + [ + ("x = 1\n", "x = 1\n", "unchanged"), + ("x=1\n", "x = 1\n", "identical"), + ("x = 'a'\n", 'x = "a"\n', "identical"), + ('"""Doc string\n here."""\n', '"""Doc string here."""\n', "docstring-whitespace"), + ("import os\nimport sys\n", "import sys\nimport os\n", "import-order"), + ("from a import x, y\n", "from a import y\nfrom a import x\n", "import-order"), + ("import os\nimport sys\nx = 1\n", "import sys\nx = 1\n", "import-removed"), + ("x = 1\n", "x = 2\n", "different"), + ("import os\n", "import os\nimport sys\n", "different"), + ("x = 1\n", "def (:\n", "unparsable"), + ], +) +def test_the_ast_proof_classifies_what_formatters_do(before, after, expected): + assert classify_python_change(before, after) == expected + + +# ------------------------------------------------------------------ lint + + +def _fake_pre_commit(tmp_path, body): + """A `pre-commit` that runs ``body`` over every file and records its env.""" + script = tmp_path / "fake-pre-commit" + script.write_text( + textwrap.dedent( + f"""\ + #!/usr/bin/env bash + echo "PRE_COMMIT_HOME=$PRE_COMMIT_HOME SKIP=${{SKIP:-}}" >> {tmp_path}/pre-commit-env.log + shift; shift; shift # run --show-diff-on-failure --files + rc=0 + for f in "$@"; do + before=$(cat "$f") + {body} + [ "$before" != "$(cat "$f")" ] && rc=1 + done + exit $rc + """ + ) + ) + script.chmod(0o755) + return str(script) + + +def test_lint_commits_formatting_only_when_pass_two_is_clean(tmp_path, repo): + _commit(repo, "feature", {"tensorrt_llm/a.py": "x = 1 \n", "notes.txt": "n\n"}) + base = _git(repo, "rev-parse", "main").strip() + fake = _fake_pre_commit(tmp_path, "sed -i 's/[[:space:]]*$//' \"$f\"") + result = run_lint( + repo, + base=base, + log_dir=tmp_path / "log", + pre_commit=fake, + python="/nonexistent/python", + identity=IDENTITY, + ) + assert result.status == "committed", result.render() + assert result.rewritten == ["tensorrt_llm/a.py"] + assert result.proof.summary() == "1 identical" + assert "type-check" in result.skipped_hooks + message = _git(repo, "log", "-1", "--format=%B") + assert "AST proof: 1 identical" in message + assert _git(repo, "log", "-1", "--format=%an <%ae>").strip() == AUTHOR + assert ( + _git(repo, "log", "-1", "--format=%(trailers:key=Signed-off-by,valueonly)").strip() + == AUTHOR + ) + assert _git(repo, "show", "--name-only", "--format=", "HEAD").split() == ["tensorrt_llm/a.py"] + assert "SKIP=type-check" in (tmp_path / "pre-commit-env.log").read_text() + + +def test_lint_restores_the_tree_and_commits_nothing_when_pass_two_fails(tmp_path, repo): + _commit(repo, "feature", {"a.py": "x = 1 \n"}) + head = _git(repo, "rev-parse", "HEAD").strip() + fake = _fake_pre_commit( + tmp_path, "sed -i 's/[[:space:]]*$//' \"$f\"; grep -q 'x' \"$f\" && rc=1" + ) + result = run_lint(repo, base="main", log_dir=tmp_path / "log", pre_commit=fake, python="/x") + assert result.status == "pass2-failed" + assert not result.ok + assert _git(repo, "rev-parse", "HEAD").strip() == head + assert _git(repo, "status", "--porcelain", "--untracked-files=no").strip() == "" + assert (tmp_path / "log" / "formatting.patch").read_text().strip() + + +def test_lint_refuses_a_formatting_commit_that_changes_semantics(tmp_path, repo): + _commit(repo, "feature", {"a.py": "x = 1\n"}) + head = _git(repo, "rev-parse", "HEAD").strip() + fake = _fake_pre_commit(tmp_path, "sed -i 's/x = 1/x = 2/' \"$f\"") + result = run_lint(repo, base="main", log_dir=tmp_path / "log", pre_commit=fake, python="/x") + assert result.status == "semantic-change" + assert result.proof.unexplained == [("a.py", "different")] + assert _git(repo, "rev-parse", "HEAD").strip() == head + assert (repo / "a.py").read_text() == "x = 1\n" + + +def test_lint_dry_run_proves_and_restores(tmp_path, repo): + _commit(repo, "feature", {"a.py": "x=1\n"}) + fake = _fake_pre_commit(tmp_path, "sed -i 's/x=1/x = 1/' \"$f\"") + result = run_lint( + repo, base="main", log_dir=tmp_path / "log", pre_commit=fake, python="/x", dry_run=True + ) + assert result.status == "dry-run" + assert (repo / "a.py").read_text() == "x=1\n" + assert "identical" in (tmp_path / "log" / "ast-proof.md").read_text() + + +def test_lint_refuses_a_dirty_tree(tmp_path, repo): + _commit(repo, "feature", {"a.py": "x = 1\n"}) + (repo / "a.py").write_text("x = 3\n") + with pytest.raises(LintError, match="tracked files are modified"): + run_lint(repo, base="main", log_dir=tmp_path / "log", pre_commit="true") + + +def test_a_read_only_pre_commit_store_gets_a_writable_copy_of_its_database(tmp_path): + store = tmp_path / "home" / ".cache" / "pre-commit" + store.mkdir(parents=True) + (store / "db.db").write_text("db") + store.chmod(0o555) + try: + if os.access(store, os.W_OK): + pytest.skip("running as a user that can write a 0555 directory") + home = writable_pre_commit_home(tmp_path / "log", home=tmp_path / "home") + assert home == tmp_path / "log" / "pre-commit-home" + assert (home / "db.db").read_text() == "db" + finally: + store.chmod(0o755) + assert writable_pre_commit_home(tmp_path / "log2", home=tmp_path / "home") == store + + +# ------------------------------------------------------------------ CI cases + + +_TEST_DB = """\ +version: 0.0.1 +l0_fake_multi_gpus: +- condition: + ranges: + system_gpu_count: + gte: 4 + lte: 4 + terms: + stage: post_merge + backend: pytorch + tests: + - unittest/_torch/attention TIMEOUT (120) + - accuracy/test_modeling_v2_fam.py::TestFam::test_gsm8k[dep4] TIMEOUT (90) + - accuracy/test_modeling_v2_fam.py::TestFam::test_gsm8k[tp4] +""" + + +@pytest.fixture +def ci_repo(repo): + db = repo / ci_case.TEST_DB_DIR + db.mkdir(parents=True) + (db / "l0_fake_multi_gpus.yml").write_text(_TEST_DB) + (repo / "tests/integration/defs").mkdir(parents=True) + _commit(repo, "test-db", {f"{ci_case.TEST_DB_DIR}/l0_fake_multi_gpus.yml": _TEST_DB}) + return repo + + +def test_test_db_lines_are_found_with_their_block(ci_repo): + entries = ci_case.find_test_db_entries(ci_repo, "test_modeling_v2_fam") + assert [e.test_id.rsplit("::", 1)[-1] for e in entries] == [ + "test_gsm8k[dep4]", + "test_gsm8k[tp4]", + ] + assert entries[0].timeout_minutes == 90 + assert entries[0].terms == {"stage": "post_merge", "backend": "pytorch"} + assert entries[0].system_gpu_count == (4, 4) + + +def test_the_stage_comes_from_the_repositorys_own_mapper(ci_repo): + mapper = ci_repo / ci_case.STAGE_MAPPER + mapper.parent.mkdir(parents=True, exist_ok=True) + mapper.write_text("print('FAKE-4_GPUs-PyTorch-Post-Merge-1')\n") + stages = ci_case.resolve_stages(ci_repo, "accuracy/x.py::T::t") + assert stages == ["FAKE-4_GPUs-PyTorch-Post-Merge-1"] + + +def test_the_ci_script_runs_the_line_as_ci_does_and_records_its_tree(tmp_path, ci_repo): + cases = ci_case.build_cases( + ci_repo, + "test_gsm8k[dep4]", + out_root=tmp_path / "out", + stage="FAKE-Stage-1", + launcher="mpirun -n 4 trtllm-llmapi-launch", + models_root="/models", + ) + assert len(cases) == 1 + script = ci_case.write_case(cases[0]) + body = script.read_text() + assert subprocess.run(["bash", "-n", str(script)]).returncode == 0 + assert "--test-prefix=FAKE-Stage-1" in body + assert "--apply-test-list-correction" in body + assert "TIMEOUT (90)" in body # the line itself is the test list + assert "env -u SLURM_PROCID" in body + assert "export LLM_MODELS_ROOT=/models" in body + assert f"cd {ci_repo}/tests/integration/defs" in body + assert body.count("rev-parse --short=10 HEAD") == 2 # at start and at end + + # Run it, with stand-ins for the launcher and pytest that record their argv. + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + for name in ("mpirun", "pytest"): + stub = bin_dir / name + stub.write_text( + f'#!/usr/bin/env bash\necho "{name} $*" >> {tmp_path}/argv.log\n' + + ('shift 3; exec "$@"\n' if name == "mpirun" else "exit 0\n") + ) + stub.chmod(0o755) + (bin_dir / "trtllm-llmapi-launch").write_text('#!/usr/bin/env bash\nexec "$@"\n') + (bin_dir / "trtllm-llmapi-launch").chmod(0o755) + env = {**os.environ, "PATH": f"{bin_dir}:{os.environ['PATH']}"} + run = subprocess.run(["bash", str(script)], capture_output=True, text=True, env=env) + assert run.returncode == 0, run.stdout + run.stderr + assert "CI_CASE_START" in run.stdout and "CI_CASE_END" in run.stdout + record = json.loads((Path(cases[0].out_dir) / "ci_case.json").read_text()) + assert record["rc"] == 0 and record["same_tree"] is True + argv = (tmp_path / "argv.log").read_text() + assert "--test-list=" in argv and "--unittest-markexpr=not cpu_only" in argv + listed = (Path(cases[0].out_dir) / "test_list.txt").read_text().strip() + assert listed == "accuracy/test_modeling_v2_fam.py::TestFam::test_gsm8k[dep4] TIMEOUT (90)" + + +def test_premerge_ci_insists_on_exactly_one_line(tmp_path, ci_repo): + from agent_flow.workflows.staircase_onboard.premerge.cli import main + + assert main(["ci", "--repo", str(ci_repo), "--case", "test_gsm8k", "--stage", "S"]) == 2 + assert ( + main( + [ + "ci", + "--repo", + str(ci_repo), + "--case", + "test_gsm8k[tp4]", + "--stage", + "S", + "--out", + str(tmp_path), + ] + ) + == 0 + ) + + +# ------------------------------------------------------------------ copyright + + +def test_copyright_finds_missing_and_stale_headers(repo): + header = "# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.\n" + _commit(repo, "old file", {"old.py": header + "x = 1\n"}) + _git(repo, "branch", "-f", "main", "HEAD") + fresh = "# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.\n" + _commit( + repo, + "changes", + { + "old.py": header + "x = 2\n", + "new_no_header.py": "y = 1\n", + "new_ok.py": fresh + "z = 1\n", + "doc.md": "not source\n", + }, + ) + report = check_copyright(repo, "main", "HEAD", year=2026) + assert report.missing == [("A", "new_no_header.py")] + assert [p for _, p, _ in report.stale] == ["old.py"] + assert not report.passed + + +# ------------------------------------------------------------------ split + + +def test_split_by_path_reproduces_the_branch_and_drops_the_bindings_link(tmp_path, repo): + base = _git(repo, "rev-parse", "HEAD").strip() + _commit(repo, "workflow", {"agent-flow/w.py": "w = 1\n"}) + _commit(repo, "model", {"tensorrt_llm/m.py": "m = 1\n"}) + _commit(repo, "both", {"agent-flow/w.py": "w = 2\n", "tensorrt_llm/m.py": "m = 2\n"}) + build = tmp_path / "build" + build.mkdir() + (repo / "tensorrt_llm" / "bindings").symlink_to(build) + _git(repo, "add", "tensorrt_llm/bindings") + _git(repo, "commit", "-q", "-s", "-m", "accident", env=_env()) + _git(repo, "rm", "-q", "--cached", "tensorrt_llm/bindings") + _git(repo, "commit", "-q", "-s", "-m", "untrack the link", env=_env()) + + stacks = [parse_stack("model=pr-a:*"), parse_stack("flow=pr-b:agent-flow/,AGENTS.md")] + result = split_branch( + repo, base=base, src="work", stacks=stacks, forbid={"model": "import agent_flow"} + ) + assert result.ok, result.render() + assert result.counts == {"model": 2, "flow": 2} + assert result.split_commits == 1 + assert result.dropped_commits == 2 + assert _git(repo, "ls-tree", "-r", "--name-only", "pr-a").split() == [ + "README.md", + "tensorrt_llm/m.py", + ] + assert _git(repo, "ls-tree", "-r", "--name-only", "pr-b").split() == [ + "README.md", + "agent-flow/w.py", + ] + # Attribution and sign-off survive the replay. + assert ( + _git( + repo, "log", "-1", "--format=%an <%ae>|%(trailers:key=Signed-off-by,valueonly)", "pr-b" + ).strip() + == f"{AUTHOR}|{AUTHOR}" + ) + + stacked = split_branch( + repo, + base=base, + src="work", + stacks=[parse_stack("model=pr-a2:*"), parse_stack("flow=pr-b2:agent-flow/")], + stacked=True, + ) + assert stacked.ok, stacked.render() + ancestry = subprocess.run( + ["git", "-C", str(repo), "merge-base", "--is-ancestor", "pr-a2", "pr-b2"] + ) + assert ancestry.returncode == 0 + assert _git(repo, "rev-parse", "pr-b2^{tree}") == _git(repo, "rev-parse", "work^{tree}") + + +def test_a_forbidden_import_fails_the_split(repo): + base = _git(repo, "rev-parse", "HEAD").strip() + _commit(repo, "leak", {"tensorrt_llm/m.py": "import agent_flow\n", "agent-flow/w.py": "w\n"}) + result = split_branch( + repo, + base=base, + src="work", + stacks=[parse_stack("model=x-a:*"), parse_stack("flow=x-b:agent-flow/")], + forbid={"model": r"(^|\s)import agent_flow"}, + ) + assert not result.ok + assert any("no match" in label and not ok for label, ok, _ in result.checks) + + +def test_a_path_no_stack_claims_is_an_error(repo): + base = _git(repo, "rev-parse", "HEAD").strip() + _commit(repo, "x", {"other/z.py": "z\n"}) + with pytest.raises(SplitError, match="no stack claims"): + split_branch(repo, base=base, src="work", stacks=[parse_stack("flow=y-b:agent-flow/")]) + + +def test_stack_specs_are_validated(): + with pytest.raises(SplitError): + parse_stack("no-branch") + with pytest.raises(SplitError): + parse_stack("a=b:*,x/") + + +# ------------------------------------------------------------------ resign + + +def test_resign_signs_each_commit_as_its_own_author_on_a_new_branch(repo): + base = _git(repo, "rev-parse", "HEAD").strip() + signed = _commit(repo, "signed", {"a.py": "a\n"}) + unsigned = _commit( + repo, + "unsigned by the orchestrator", + {"b.py": "b\n"}, + signoff=False, + env=_env("login", "login@host.invalid", committer=("Operator", "op@example.invalid")), + ) + after = _commit(repo, "after", {"c.py": "c\n"}) + result = resign_branch(repo, src="work", new_branch="work-signed", base="main") + assert result.ok + assert result.unsigned_left == 0 + assert [(o, n == o, s) for o, n, s, _ in result.mapping][:1] == [(signed, True, False)] + assert result.mapping[1][0] == unsigned and result.mapping[1][2] is True + assert result.mapping[2][0] == after and result.mapping[2][1] != after # descendant rewritten + new_unsigned = result.mapping[1][1] + fmt = "%an <%ae>|%cn <%ce>|%(trailers:key=Signed-off-by,valueonly)" + an, cn, so = _git(repo, "log", "-1", f"--format={fmt}", new_unsigned).strip().split("|") + assert an == cn == so.strip() == "login " + assert _git(repo, "rev-parse", "work").strip() == after # source untouched + with pytest.raises(ResignError, match="exists"): + resign_branch(repo, src="work", new_branch="work-signed", base="main") + assert result.base == base + assert "| 2 |" in result.render_markdown("work") + + +# ------------------------------------------------------------------ PR opener + + +def _request(**kw): + defaults = dict( + upstream="NVIDIA/TensorRT-LLM", + owner="someone", + branch="user/someone/x", + title="[TRTLLM-1][feat] staircase: x", + body="body\n", + ) + return PrRequest(**{**defaults, **kw}) + + +def test_the_pr_opener_validates_the_title_and_body(): + assert validate(_request()) == [] + assert validate(_request(title="feat: x")) + assert validate(_request(body="x\n\nCo-authored-by: Bot \n")) + with pytest.raises(PrOpenError): + open_pr(_request(title="bad title")) + + +def test_the_pr_opener_is_a_dry_run_that_touches_no_network_by_default(): + def opener(_request): + raise AssertionError("the network was touched") + + report = open_pr(_request(), opener=opener) + assert report["dry_run"] is True + assert report["payload"]["draft"] is True + assert report["payload"]["head"] == "someone:user/someone/x" + + +class _Response: + def __init__(self, payload): + self.payload = payload + + def read(self): + return json.dumps(self.payload).encode() + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + +def test_the_pr_opener_posts_only_with_yes_and_never_reports_the_token(): + calls = [] + + def opener(request): + calls.append((request.get_method(), request.full_url, request.data)) + if request.get_method() == "POST": + return _Response( + { + "html_url": "https://github.com/x/pull/1", + "number": 1, + "state": "open", + "draft": True, + } + ) + if "/branches/" in request.full_url: + return _Response({"commit": {"sha": "abcdef1234567890"}}) + return _Response([]) + + checked = open_pr(_request(), check_remote=True, token="SECRET-TOKEN", opener=opener) + assert [m for m, _, _ in calls] == ["GET", "GET"] + assert checked["dry_run"] is True + created = open_pr(_request(), yes=True, token="SECRET-TOKEN", opener=opener) + assert created["url"] == "https://github.com/x/pull/1" + assert [m for m, _, _ in calls].count("POST") == 1 + assert json.loads(calls[-1][2])["draft"] is True + assert "SECRET-TOKEN" not in json.dumps(created) + json.dumps(checked) + + +def test_the_pr_opener_refuses_when_a_pr_is_already_open(): + def opener(request): + if "/branches/" in request.full_url: + return _Response({"commit": {"sha": "abc"}}) + return _Response([{"html_url": "https://github.com/x/pull/9"}]) + + with pytest.raises(PrOpenError, match="already"): + open_pr(_request(), yes=True, token="t", opener=opener) + + +def test_the_token_comes_from_the_environment_or_ghs_hosts_file(tmp_path): + assert read_token({"GH_TOKEN": "a"}) == "a" + config = tmp_path / "gh" + config.mkdir() + (config / "hosts.yml").write_text( + "github.com:\n user: me\n users:\n me:\n oauth_token: from-hosts\n" + ) + assert read_token({"GH_CONFIG_DIR": str(config)}) == "from-hosts" + with pytest.raises(PrOpenError): + read_token({"GH_CONFIG_DIR": str(tmp_path / "none")}) + + +# ------------------------------------------------------------------ pack and CLI + + +def test_the_pack_reports_the_audit_the_headers_and_the_ci_lines(tmp_path, ci_repo): + _commit(ci_repo, "unsigned", {"tensorrt_llm/m.py": "m\n"}, signoff=False) + report = pack.build_pack( + ci_repo, test_pattern="test_modeling_v2_fam", base_ref="main", year=2026 + ) + assert len(report.hygiene) == 1 and "no Signed-off-by" in report.hygiene[0].problem + assert report.copyright is not None and report.copyright.missing == [("A", "tensorrt_llm/m.py")] + assert len(report.ci_entries) == 2 + path = pack.write_pack(report, tmp_path / "reports") + text = path.read_text() + assert "premerge resign" in text # the remedy is the first remaining step + assert "test_gsm8k[dep4]" in text + + +def test_the_pack_without_a_base_says_so(tmp_path): + lone = tmp_path / "lone" + lone.mkdir() + _git(lone, "init", "-q", "-b", "trunk") + _commit(lone, "seed", {"x": "x\n"}) + report = pack.build_pack(lone, base_ref="nope") + assert report.base == "" + assert "no base" in report.render_markdown() + + +def test_the_cli_dispatches_premerge(repo): + from agent_flow.workflows.staircase_onboard import cli + + _commit(repo, "unsigned", {"z.py": "z\n"}, signoff=False) + with pytest.raises(SystemExit) as exc: + cli.main(["premerge", "hygiene", "--repo", str(repo), "--base", "main"]) + assert exc.value.code == 1 + with pytest.raises(SystemExit) as exc: + cli.main(["premerge", "hygiene", "--repo", str(repo), "--base", "HEAD"]) + assert exc.value.code == 0 From 1cdd01798a775386b4dffae8758d91db93f8ac47 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 12 Oct 2026 00:56:12 -0700 Subject: [PATCH 05/70] [TRTLLM-16460][feat] staircase: evidence ledger, stale-verdict invalidation and reports generated from it An append-only ledger of gate verdicts keyed by tree, member, config and covered code paths; verdicts invalidated by later commits are queued for re-run and the final regression is bound to the delivered tree. progress.json/html and the final report's gate tables are generated from the ledger. Decision files are --.yaml and never overwritten (F033), journal ordinals are atomic, and catalog/index.yaml is generated from per-entry fragments (F013/F014). Squashed from user/williamj/staircase-polish @ c8109d6f90 (10 commits): - [feat] staircase: append-only evidence ledger for gate verdicts - [feat] staircase: stale-verdict invalidation, re-run queue and final-regression binding - [feat] staircase: generate progress.json/html and the final report's gate tables from the ledger - [fix] staircase: decision files are --.yaml and never overwritten (F033) - [fix] staircase: Journal.next_ordinal is atomic across threads and processes - [feat] staircase: generate catalog/index.yaml from per-entry fragments (F013/F014) - [feat] staircase: the ledger records envelope probes, early checks and delivery checks - [feat] staircase: a verdict's config fingerprint includes each role's backend and model - [feat] staircase: the pre-merge pack states the final-regression binding and how to re-check it - [fix] staircase: catalog index sync is serialized and never undoes a commit that landed Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../staircase_onboard/catalog_index.py | 706 +++++++++ .../workflows/staircase_onboard/evidence.py | 558 +++++++ .../workflows/staircase_onboard/freshness.py | 542 +++++++ .../workflows/staircase_onboard/gates.py | 6 + .../workflows/staircase_onboard/journal.py | 56 +- .../workflows/staircase_onboard/ledger.py | 1176 +++++++++++++++ .../workflows/staircase_onboard/progress.py | 192 ++- .../staircase_onboard/prompts/_common.py | 6 +- .../staircase_onboard/prompts/smith.py | 9 +- .../workflows/staircase_onboard/report.py | 1337 +++++++++++++++++ .../skills/staircase-kernel-onboard/SKILL.md | 8 +- .../references/decomposition.md | 2 +- .../references/entry-spec.md | 17 +- .../skills/staircase-kernel-verify/SKILL.md | 7 +- .../workflows/staircase_onboard/workflow.py | 83 +- .../staircase_onboard/test_catalog_index.py | 418 ++++++ .../staircase_onboard/test_decision_files.py | 122 +- .../staircase_onboard/test_evidence.py | 201 +++ .../staircase_onboard/test_freshness.py | 472 ++++++ .../staircase_onboard/test_journal.py | 117 ++ .../staircase_onboard/test_ledger.py | 531 +++++++ .../staircase_onboard/test_report.py | 367 +++++ 22 files changed, 6815 insertions(+), 118 deletions(-) create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/catalog_index.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/evidence.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/freshness.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/ledger.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/report.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_catalog_index.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_evidence.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_freshness.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_journal.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_ledger.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_report.py diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/catalog_index.py b/agent-flow/agent_flow/workflows/staircase_onboard/catalog_index.py new file mode 100644 index 000000000000..1897722bee5a --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/catalog_index.py @@ -0,0 +1,706 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``catalog/index.yaml``, generated from one fragment per entry. + +The index is the one file every smith touched. Each appended its row and +committed the whole file, so the first smith to commit swept in its +siblings' uncommitted rows: ``HEAD`` then named entries whose files were not +in ``HEAD``, and a later smith's message named rows it did not write (F013). +Every smith of the DeepSeek-V4.1 run spent effort documenting that. + +Now each entry's row is its own file:: + + catalog/index.d/_header.yaml # everything above `entries:`, and the category order + catalog/index.d//.yaml # {path, impl, summary} of one entry + +A smith writes and commits its own fragment with its entry, so no two +smiths ever commit the same file. The orchestrator regenerates +``index.yaml`` from the fragments **committed in HEAD** -- never from the +working tree, where siblings' fragments may be half-written -- and commits +that one file by itself (:func:`sync`). Generation also leaves out a +fragment whose entry wrapper is not in ``HEAD``, so the index can never name +an entry that a clone would not have. + +:func:`migrate` splits an existing hand-maintained index into the header and +fragments (preserving the row order and the category dividers). :func:`sync` +does the same by itself, in a commit of its own, the first time HEAD holds a +fragment and no header, so a run that adds no catalog entry never touches +the catalog. :func:`problems` is the check: the index matches its fragments, +each fragment sits where its ``path`` says, and every fragment's entry exists. +""" + +from __future__ import annotations + +import argparse +import fcntl +import json +import os +import re +import subprocess +import sys +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Mapping, Sequence + +import yaml + +from .delivery import CATALOG_DIR + +INDEX_FILENAME = "index.yaml" +FRAGMENTS_DIRNAME = "index.d" +HEADER_FILENAME = "_header.yaml" +#: In the workspace: serializes every sync of one run, across threads and processes. +LOCK_FILENAME = ".catalog-index.lock" +#: The comment the generated index carries just above its rows. +GENERATED_NOTICE = ( + "# GENERATED from index.d/ (python -m agent_flow.workflows.staircase_onboard.catalog_index).\n" + "# Each entry's row is its own file, index.d//.yaml: edit that, not this list.\n" +) +#: Fragment keys rendered first, in this order; any others follow, sorted. +ROW_KEYS = ("path", "impl", "summary") +#: Kept in a fragment to preserve a migrated index's row order; never rendered. +ORDER_KEY = "order" +_DIVIDER_WIDTH = 67 +_GIT_TIMEOUT_S = 60 + +CommitFn = Callable[[str, Sequence[str]], "tuple[bool, str]"] + + +# ------------------------------------------------------------------ model + + +@dataclass(frozen=True) +class Header: + """Everything of ``index.yaml`` that is not an entry row.""" + + preamble: str = "" + category_order: tuple[str, ...] = () + category_dividers: Mapping[str, str] = field(default_factory=dict) + + def to_yaml(self) -> str: + data = { + "preamble": self.preamble, + "category_order": list(self.category_order), + "category_dividers": dict(self.category_dividers), + } + return ( + "# The header of catalog/index.yaml: everything above `entries:`, and the category\n" + "# order. Each entry's row is its own file in this directory.\n" + + yaml.dump(data, Dumper=_Dumper, sort_keys=False, allow_unicode=True, width=1000) + ) + + @classmethod + def from_yaml(cls, text: str) -> "Header": + data = yaml.safe_load(text) or {} + if not isinstance(data, Mapping): + raise ValueError(f"{HEADER_FILENAME} must be a mapping") + return cls( + preamble=str(data.get("preamble") or ""), + category_order=tuple(str(c) for c in data.get("category_order") or ()), + category_dividers={ + str(k): str(v) for k, v in (data.get("category_dividers") or {}).items() + }, + ) + + +class _Dumper(yaml.SafeDumper): + """Writes multi-line strings as literal blocks, so the preamble stays readable.""" + + +def _str_representer(dumper: yaml.SafeDumper, value: str) -> yaml.Node: + style = "|" if "\n" in value else None + return dumper.represent_scalar("tag:yaml.org,2002:str", value, style=style) + + +_Dumper.add_representer(str, _str_representer) + + +@dataclass(frozen=True) +class Snapshot: + """The catalog as one tree holds it: header, fragments, files, and the index text.""" + + header: Header | None + #: Fragment location (relative to the catalog) -> its row. + fragments: Mapping[str, Mapping[str, Any]] + #: Every file under the catalog, relative to it. + files: frozenset[str] + index_text: str | None + + +def fragment_rel(entry_path: str) -> str: + """Where the fragment of the entry at ``entry_path`` (``norm/x.py``) lives.""" + stem = entry_path[:-3] if entry_path.endswith(".py") else entry_path + return f"{FRAGMENTS_DIRNAME}/{stem}.yaml" + + +def _header_rel() -> str: + return f"{FRAGMENTS_DIRNAME}/{HEADER_FILENAME}" + + +def _load_fragment(text: str) -> dict[str, Any] | None: + try: + data = yaml.safe_load(text) + except yaml.YAMLError: + return None + return dict(data) if isinstance(data, Mapping) else None + + +# ---------------------------------------------------------------- snapshots + + +def snapshot_worktree(catalog_dir: str | Path) -> Snapshot: + """The catalog as the working tree holds it.""" + root = Path(catalog_dir) + files = frozenset( + p.relative_to(root).as_posix() + for p in root.rglob("*") + if p.is_file() and "__pycache__" not in p.parts + ) + header = None + header_path = root / _header_rel() + if header_path.is_file(): + header = Header.from_yaml(header_path.read_text(encoding="utf-8")) + fragments: dict[str, Mapping[str, Any]] = {} + for rel in sorted(files): + if rel.startswith(f"{FRAGMENTS_DIRNAME}/") and rel.endswith(".yaml"): + if rel == _header_rel(): + continue + row = _load_fragment((root / rel).read_text(encoding="utf-8")) + fragments[rel] = row if row is not None else {} + index = root / INDEX_FILENAME + return Snapshot( + header=header, + fragments=fragments, + files=files, + index_text=index.read_text(encoding="utf-8") if index.is_file() else None, + ) + + +def _git(repo: str | Path, *args: str, stdin: str | None = None) -> str | None: + try: + completed = subprocess.run( + ["git", "-C", str(repo), *args], + capture_output=True, + text=True, + timeout=_GIT_TIMEOUT_S, + check=False, + input=stdin, + ) + except (OSError, subprocess.SubprocessError): + return None + return completed.stdout if completed.returncode == 0 else None + + +def _read_blobs(repo: str | Path, blobs: Sequence[str]) -> dict[str, str]: + """Contents of ``blobs`` by id, through one ``git cat-file --batch``.""" + if not blobs: + return {} + try: + completed = subprocess.run( + ["git", "-C", str(repo), "cat-file", "--batch"], + input=("\n".join(blobs) + "\n").encode(), + capture_output=True, + timeout=_GIT_TIMEOUT_S, + check=False, + ) + except (OSError, subprocess.SubprocessError): + return {} + out = completed.stdout + contents: dict[str, str] = {} + pos = 0 + while pos < len(out): + newline = out.index(b"\n", pos) + header = out[pos:newline].decode().split() + pos = newline + 1 + if len(header) < 3 or header[1] == "missing": + continue + size = int(header[2]) + contents[header[0]] = out[pos : pos + size].decode("utf-8", errors="replace") + pos += size + 1 + return contents + + +def snapshot_at( + repo: str | Path, rev: str = "HEAD", *, catalog_rel: str = CATALOG_DIR +) -> Snapshot | None: + """The catalog as commit ``rev`` holds it, or None when git cannot answer.""" + listing = _git(repo, "ls-tree", "-r", "-z", rev, "--", catalog_rel) + if listing is None: + return None + prefix = catalog_rel.rstrip("/") + "/" + blobs: dict[str, str] = {} + for record in listing.split("\0"): + meta, _, path = record.partition("\t") + parts = meta.split() + if len(parts) == 3 and parts[1] == "blob" and path.startswith(prefix): + blobs[path[len(prefix) :]] = parts[2] + wanted = [ + rel + for rel in blobs + if rel == INDEX_FILENAME + or (rel.startswith(f"{FRAGMENTS_DIRNAME}/") and rel.endswith(".yaml")) + ] + contents = _read_blobs(repo, [blobs[rel] for rel in wanted]) + text = {rel: contents.get(blobs[rel], "") for rel in wanted} + header = Header.from_yaml(text[_header_rel()]) if _header_rel() in text else None + fragments: dict[str, Mapping[str, Any]] = {} + for rel in sorted(text): + if rel in (INDEX_FILENAME, _header_rel()): + continue + row = _load_fragment(text[rel]) + fragments[rel] = row if row is not None else {} + return Snapshot( + header=header, + fragments=fragments, + files=frozenset(blobs), + index_text=text.get(INDEX_FILENAME), + ) + + +# --------------------------------------------------------------- rendering + + +def _scalar(value: Any) -> str: + """One YAML scalar on one line: plain when that reads back unchanged, else quoted.""" + if isinstance(value, str): + plain = ( + value + and value == value.strip() + and "\n" not in value + and ": " not in value + and " #" not in value + and value[0] not in "!&*-?{}[],#|>@`\"'%" + ) + if plain: + try: + if yaml.safe_load(value) == value: + return value + except yaml.YAMLError: + pass + return json.dumps(value, ensure_ascii=False) + return json.dumps(value, ensure_ascii=False) + + +def _divider(category: str) -> str: + text = f" # ─── {category} " + return text + "─" * max(3, _DIVIDER_WIDTH - len(text)) + + +def _render_row(row: Mapping[str, Any]) -> list[str]: + keys = [k for k in ROW_KEYS if k in row] + sorted( + k for k in row if k not in ROW_KEYS and k != ORDER_KEY + ) + lines = [] + for index, key in enumerate(keys): + lead = " - " if index == 0 else " " + # Summaries are always double-quoted, as the hand-written index had them. + value = json.dumps(row[key], ensure_ascii=False) if key == "summary" else _scalar(row[key]) + lines.append(f"{lead}{key}: {value}") + return lines + + +def renderable_rows(snapshot: Snapshot) -> list[Mapping[str, Any]]: + """Fragments that sit where their ``path`` says and whose entry exists.""" + rows = [] + for rel, row in snapshot.fragments.items(): + path = str(row.get("path") or "") + if path and fragment_rel(path) == rel and path in snapshot.files: + rows.append(row) + return rows + + +def render(snapshot: Snapshot) -> str: + """The ``index.yaml`` text the snapshot's header and fragments define.""" + header = snapshot.header or Header() + by_category: dict[str, list[Mapping[str, Any]]] = {} + for row in renderable_rows(snapshot): + by_category.setdefault(str(row["path"]).split("/", 1)[0], []).append(row) + order = [c for c in header.category_order if c in by_category] + sorted( + c for c in by_category if c not in header.category_order + ) + out = [header.preamble.rstrip("\n") + "\n\n" if header.preamble else "", GENERATED_NOTICE] + out.append("entries:\n") + blocks = [] + for category in order: + rows = sorted( + by_category[category], + key=lambda r: ( + 0 if isinstance(r.get(ORDER_KEY), int) else 1, + r.get(ORDER_KEY) if isinstance(r.get(ORDER_KEY), int) else 0, + str(r["path"]), + ), + ) + lines = [header.category_dividers.get(category) or _divider(category)] + for index, row in enumerate(rows): + if index: + lines.append("") + lines.extend(_render_row(row)) + blocks.append("\n".join(lines)) + out.append("\n\n".join(blocks) + ("\n" if blocks else "")) + return "".join(out) + + +# --------------------------------------------------------------- migration + + +_DIVIDER_LINE = re.compile(r"^\s*#\s*─+\s*(?P.*?)\s*─+\s*$") + + +def split_index(text: str) -> tuple[Header, list[dict[str, Any]]]: + """Split a hand-maintained ``index.yaml`` into its header and ordered rows.""" + lines = text.splitlines(keepends=True) + start = next((i for i, line in enumerate(lines) if line.startswith("entries:")), None) + if start is None: + raise ValueError(f"{INDEX_FILENAME} has no top-level `entries:`") + notice = set(GENERATED_NOTICE.splitlines(keepends=True)) + preamble = "".join(line for line in lines[:start] if line not in notice) + data = yaml.safe_load(text) or {} + rows = [dict(r) for r in data.get("entries") or () if isinstance(r, Mapping) and r.get("path")] + order: list[str] = [] + dividers: dict[str, str] = {} + pending_divider: str | None = None + for line in lines[start + 1 :]: + if _DIVIDER_LINE.match(line): + pending_divider = line.rstrip("\n") + continue + match = re.match(r"^\s*-\s*path:\s*(\S+)", line) + if match: + category = match.group(1).strip("'\"").split("/", 1)[0] + if category not in order: + order.append(category) + if pending_divider is not None: + dividers[category] = pending_divider + pending_divider = None + for row in rows: + category = str(row["path"]).split("/", 1)[0] + if category not in order: + order.append(category) + return Header(preamble=preamble, category_order=tuple(order), category_dividers=dividers), rows + + +def _dump_fragment(row: Mapping[str, Any]) -> str: + return yaml.safe_dump(dict(row), sort_keys=False, allow_unicode=True, width=1000) + + +def migration_files(snapshot: Snapshot) -> dict[str, str]: + """The files that turn ``snapshot`` into a fragment-generated catalog. + + The header, when there is none yet, and a fragment for every index row + that has none: a row a smith appended by hand is adopted rather than + dropped by the next regeneration. Row order is kept through ``order``. + """ + if snapshot.index_text is None: + return {} + header, rows = split_index(snapshot.index_text) + files: dict[str, str] = {} + if snapshot.header is None: + files[_header_rel()] = header.to_yaml() + for position, row in enumerate(rows, start=1): + rel = fragment_rel(str(row["path"])) + if rel in snapshot.fragments or rel in files: + continue + fragment = {k: row[k] for k in row} + if snapshot.header is None: + fragment[ORDER_KEY] = position + files[rel] = _dump_fragment(fragment) + return files + + +def migrate(catalog_dir: str | Path) -> list[Path]: + """Split the working tree's ``index.yaml`` into fragments, then regenerate it.""" + root = Path(catalog_dir) + written = [] + for rel, text in migration_files(snapshot_worktree(root)).items(): + path = root / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text, encoding="utf-8") + written.append(path) + index = root / INDEX_FILENAME + index.write_text(render(snapshot_worktree(root)), encoding="utf-8") + written.append(index) + return written + + +# ------------------------------------------------------------------- check + + +def problems(snapshot: Snapshot) -> list[str]: + """Everything that makes ``snapshot``'s index not a pure function of its fragments.""" + if snapshot.header is None: + return [f"{FRAGMENTS_DIRNAME}/{HEADER_FILENAME} is missing: the index is not migrated"] + found: list[str] = [] + seen: dict[str, str] = {} + for rel, row in snapshot.fragments.items(): + path = str(row.get("path") or "") + if not path: + found.append(f"{rel}: no `path` (or not a YAML mapping)") + continue + if fragment_rel(path) != rel: + found.append(f"{rel}: names {path}, whose fragment belongs at {fragment_rel(path)}") + if path not in snapshot.files: + found.append(f"{rel}: entry {path} does not exist") + if path in seen: + found.append(f"{rel}: {path} is also claimed by {seen[path]}") + seen[path] = rel + if snapshot.index_text != render(snapshot): + found.append( + f"{INDEX_FILENAME} is not what its fragments generate; regenerate it " + f"(python -m agent_flow.workflows.staircase_onboard.catalog_index generate)" + ) + return found + + +# -------------------------------------------------------------------- sync + + +@dataclass(frozen=True) +class SyncResult: + action: str # skipped | current | committed | deferred | failed + detail: str = "" + paths: tuple[str, ...] = () + + +def commit_with_script(script: str | Path, repo: str | Path) -> CommitFn: + """A commit function that goes through the workspace's ``commit.sh``.""" + + def commit(message: str, paths: Sequence[str]) -> tuple[bool, str]: + try: + completed = subprocess.run( + [str(script), message, *paths], + cwd=str(repo), + capture_output=True, + text=True, + timeout=900, + check=False, + ) + except (OSError, subprocess.SubprocessError) as exc: + return False, f"{type(exc).__name__}: {exc}" + return completed.returncode == 0, (completed.stdout + completed.stderr).strip() + + return commit + + +def sync( + repo: str | Path, + *, + commit: CommitFn, + catalog_rel: str = CATALOG_DIR, + lock_path: str | Path | None = None, +) -> SyncResult: + """Make ``index.yaml`` in HEAD what HEAD's fragments generate, in one commit of its own. + + Reads HEAD, not the working tree: a sibling's fragment that is written + but not committed is not part of HEAD, so the index never names it. Once + HEAD has a fragment but no header, migrates first (header plus a fragment + per row); a committed index row with no fragment is adopted. Does nothing + while HEAD has no fragment at all. Leaves an ``index.yaml`` with + uncommitted edits alone -- someone else is writing it. + + Idempotent: a second call on the same HEAD is ``current``. With + ``lock_path``, calls from several threads or processes are serialized + under an ``flock`` on that file, so two of them can never both migrate; + the one that waited reads the HEAD the other committed. A commit that + reports failure is checked against HEAD before anything is undone, and + the undo touches only the files this call created and that HEAD does not + hold. + """ + if lock_path is None: + return _sync(repo, commit=commit, catalog_rel=catalog_rel) + lock = Path(lock_path) + lock.parent.mkdir(parents=True, exist_ok=True) + fd = os.open(lock, os.O_RDWR | os.O_CREAT, 0o644) + try: + fcntl.flock(fd, fcntl.LOCK_EX) + return _sync(repo, commit=commit, catalog_rel=catalog_rel) + finally: + fcntl.flock(fd, fcntl.LOCK_UN) + os.close(fd) + + +def _sync(repo: str | Path, *, commit: CommitFn, catalog_rel: str) -> SyncResult: + head = snapshot_at(repo, "HEAD", catalog_rel=catalog_rel) + if head is None: + return SyncResult("skipped", "not a git work tree") + if head.header is None and not head.fragments: + # Migrated lazily, by the first fragment a smith commits: a run that + # adds no catalog entry never touches the catalog. + return SyncResult("skipped", "no fragment committed yet; nothing to generate") + root = Path(repo) / catalog_rel + index = root / INDEX_FILENAME + current = index.read_text(encoding="utf-8") if index.is_file() else None + if current != head.index_text: + return SyncResult("deferred", f"{INDEX_FILENAME} has uncommitted edits; left alone") + + new_files = migration_files(head) + fragments = dict(head.fragments) + header = head.header + for rel, text in new_files.items(): + if rel == _header_rel(): + header = Header.from_yaml(text) + else: + fragments[rel] = _load_fragment(text) or {} + target = Snapshot(header=header, fragments=fragments, files=head.files, index_text=None) + rendered = render(target) + if rendered == current and not new_files: + return SyncResult("current", "index.yaml matches its fragments") + + for rel in new_files: + if (root / rel).exists(): + return SyncResult("deferred", f"{rel} exists uncommitted; left alone") + for rel, text in new_files.items(): + path = root / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(text, encoding="utf-8") + index.write_text(rendered, encoding="utf-8") + paths = [f"{catalog_rel}/{rel}" for rel in new_files] + [f"{catalog_rel}/{INDEX_FILENAME}"] + message = ( + "chore(catalog): split index.yaml into per-entry fragments under index.d/" + if head.header is None + else "chore(catalog): regenerate index.yaml from index.d/ fragments" + ) + ok, output = commit(message, paths) + if ok: + return SyncResult("committed", output.splitlines()[0] if output else "", tuple(paths)) + after = snapshot_at(repo, "HEAD", catalog_rel=catalog_rel) + if ( + after is not None + and after.index_text == rendered + and all(rel in after.files for rel in new_files) + ): + # The commit landed and only its report failed: nothing to undo. + return SyncResult("committed", f"landed; the helper said: {output[-500:]}", tuple(paths)) + committed = after.files if after is not None else frozenset() + undo = [rel for rel in new_files if rel not in committed] + # Unstage first (a helper that resets its own commit leaves the paths + # staged), then remove what this call created, then put the index back. + _git( + repo, + "reset", + "-q", + "--", + *(f"{catalog_rel}/{rel}" for rel in undo), + f"{catalog_rel}/{INDEX_FILENAME}", + ) + for rel in undo: + path = root / rel + path.unlink(missing_ok=True) + for parent in path.parents: + if parent == root or not parent.is_dir() or any(parent.iterdir()): + break + parent.rmdir() + if current is not None and (after is None or after.index_text == head.index_text): + index.write_text(current, encoding="utf-8") + return SyncResult("failed", output[-2000:], tuple(paths)) + + +# --------------------------------------------------------------- workflow + + +class CatalogIndexHook: + """The orchestrator's side: sync before a turn reads the index, check at delivery.""" + + def __init__(self, workflow: Any) -> None: + self._workflow = workflow + + def _repo(self) -> str: + return str((getattr(self._workflow, "task_data", None) or {}).get("trtllm_repo_path") or "") + + def sync(self, log: Any = None) -> SyncResult: + """Regenerate and commit the index if HEAD's fragments moved. Never raises. + + Commits through the workspace's ``commit.sh``, so it takes the same + lock as every agent's commit. Without that script (before preflight + has installed it) nothing is written. + """ + from .committer import SCRIPT_NAME + + repo = self._repo() + script = Path(self._workflow.workspace) / "bin" / SCRIPT_NAME + if not repo or not script.is_file(): + return SyncResult("skipped", "no repository or no commit helper yet") + result = sync( + repo, + commit=commit_with_script(script, repo), + lock_path=Path(self._workflow.workspace) / LOCK_FILENAME, + ) + if result.action in ("committed", "failed", "deferred"): + self._workflow.journal.event(f"catalog index {result.action}", result.detail[:500]) + return result + + def checks(self) -> list[Any]: + """Delivery checks: the delivered index is exactly what HEAD's fragments generate.""" + from .delivery import Check + + repo = self._repo() + head = snapshot_at(repo) if repo else None + if head is None or (head.header is None and not head.fragments): + return [] + found = problems(head) + return [ + Check( + name="catalog:index-generated", + passed=not found, + detail="; ".join(found[:10]) + or f"{INDEX_FILENAME} is generated from {len(head.fragments)} fragment(s)", + ) + ] + + def extend(self, report: Any) -> Any: + """``report`` with :meth:`checks` appended.""" + from .delivery import DeliveryReport + + return DeliveryReport(checks=(*report.checks, *self.checks())) + + +# ----------------------------------------------------------------------- CLI + + +def main(argv: Sequence[str] | None = None) -> int: + """``migrate``, ``generate`` and ``check`` a working tree's catalog index.""" + parser = argparse.ArgumentParser(prog="staircase-catalog-index", description=main.__doc__) + parser.add_argument("command", choices=("migrate", "generate", "check")) + parser.add_argument("--repo", required=True, help="the repository root") + parser.add_argument("--catalog-rel", default=CATALOG_DIR) + args = parser.parse_args(argv) + root = Path(args.repo) / args.catalog_rel + if args.command == "migrate": + for path in migrate(root): + sys.stdout.write(f"wrote {path}\n") + return 0 + if args.command == "generate": + (root / INDEX_FILENAME).write_text(render(snapshot_worktree(root)), encoding="utf-8") + sys.stdout.write(f"wrote {root / INDEX_FILENAME}\n") + return 0 + found = problems(snapshot_worktree(root)) + for problem in found: + sys.stdout.write(f"PROBLEM {problem}\n") + sys.stdout.write(f"{len(found)} problem(s)\n") + return 1 if found else 0 + + +__all__ = [ + "CatalogIndexHook", + "FRAGMENTS_DIRNAME", + "GENERATED_NOTICE", + "HEADER_FILENAME", + "Header", + "INDEX_FILENAME", + "Snapshot", + "SyncResult", + "commit_with_script", + "fragment_rel", + "migrate", + "migration_files", + "problems", + "render", + "renderable_rows", + "snapshot_at", + "snapshot_worktree", + "split_index", + "sync", +] + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/evidence.py b/agent-flow/agent_flow/workflows/staircase_onboard/evidence.py new file mode 100644 index 000000000000..72f20dfa2c54 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/evidence.py @@ -0,0 +1,558 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The orchestrator's side of the evidence ledger. + +A verdict needs context no gate has: the repository and target it measured, +the unit being gated, the environment every gate command runs under, and the +spec. The orchestrator has all of it. :class:`RunEvidence` binds the two: +``gates.GateLog`` calls :meth:`RunEvidence.record` with every result, and the +result goes into the workspace ledger with that context attached. + +At delivery, :meth:`RunEvidence.settle_delivery` makes the evidence describe +the tree being delivered. It re-runs every gate whose latest verdict is +stale, failed, or was measured on another tree or native build, records the +full regression once every required verdict passes on this tree, and refuses +delivery when either is not achieved. + +The workflow holds one instance and calls into it from a few hook points, so +``workflow.py`` itself carries no ledger logic. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, replace +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Callable, Iterable + +import yaml +from rich.markup import escape as escape_markup + +from agent_flow.console import print_message + +from . import envelope as envelope_mod +from . import gates as gates_mod +from .freshness import ( + RERUN_QUEUE_FILENAME, + Assessment, + RerunEntry, + assess, + complete_regression, + drain, + last_regression, + record_regression, + regression_binding, + rerun_queue, + write_rerun_queue, +) +from .ledger import ( + LEDGER_FILENAME, + Ledger, + LedgerPolicy, + NativeBuild, + TreeIdentity, + Verdict, + VerdictContext, + hash_paths, + measure, + native_build, + record_gate_result, + tree_identity, +) +from .report import ( + FINAL_REPORT, + PROGRESS_HTML, + PROGRESS_JSON, + check_consistency, + collect_facts, + write_reports, +) +from .roles import ROLES, resolve_role_agent +from .task_schema import assembly_units, members_for + +#: The verdict a run closes with when its evidence does not describe the tree +#: it would deliver. Remediable: the repair (a fix, then a re-run of the +#: failing gate) happens outside the run, and re-entering delivery re-checks. +VERDICT_EVIDENCE = "delivery-failed:evidence" + +_QUEUE_LINES_SHOWN = 20 + +#: Ledger keys for verdicts that are not engine gates. They carry no unit, so +#: they are reported and kept as history but never re-run or required by the +#: final regression: an early check re-runs after every unit by itself, and +#: the delivery checks run again on every delivery. +EARLY_CHECK_PREFIX = "early:" +DELIVERY_CHECK_PREFIX = "delivery:" + +#: ``at B=8 T=2048 ...: <boot detail>`` -- how an envelope probe names its point. +_PROBE_POINT = re.compile(r"^at (?P<point>.+?): ") + + +@dataclass(frozen=True) +class _CheckResult: + """A check shaped like a gate result, so it records through the one entry point.""" + + name: str + passed: bool + detail: str + finding: bool = False + + +def role_agents(task_data: Any) -> dict[str, dict[str, Any]]: + """The backend and model every role is routed to, as the ``agents`` block resolves them.""" + agents: dict[str, dict[str, Any]] = {} + for role in ROLES: + agent = resolve_role_agent(task_data or {}, role) + agents[role] = {"backend": agent.backend, "model": agent.model} + if agent.reasoning_effort: + agents[role]["reasoning_effort"] = agent.reasoning_effort + return agents + + +def gate_key(result: Any) -> str: + """The ledger key for ``result``: its name, qualified by the point for an envelope probe. + + Every probe of a unit at one member is named ``envelope``; keyed by name + alone, the last probe would supersede the others and a failing point + could hide behind a passing one. + """ + name = str(getattr(result, "name", "") or "gate") + if name == envelope_mod.PROBE_GATE: + match = _PROBE_POINT.match(str(getattr(result, "detail", "") or "")) + if match: + return f"{name}[{match.group('point')}]" + return name + + +#: Where ``--clean`` moves the previous run's evidence. Moved, not deleted: +#: the ledger is append-only, and a fresh run does not change what an old +#: one measured. +ARCHIVE_DIRNAME = "archive" +_ARCHIVED = (LEDGER_FILENAME, RERUN_QUEUE_FILENAME, FINAL_REPORT, PROGRESS_JSON, PROGRESS_HTML) + + +class RunEvidence: + """One run's ledger, and the hooks the orchestrator calls.""" + + def __init__( + self, workflow: Any, *, policy: LedgerPolicy | None = None, clean: bool = False + ) -> None: + self._workflow = workflow + self.ledger = Ledger.in_workspace(workflow.workspace) + self.policy = policy or LedgerPolicy() + #: Verdict ids already reported as stale in the journal. + self._warned: set[str] = set() + if clean: + self.archive() + + def archive(self) -> list[Path]: + """Move the previous run's ledger and reports under ``archive/``, stamped.""" + workspace = Path(self._workflow.workspace) + stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + moved = [] + for name in _ARCHIVED: + path = workspace / name + if path.is_file(): + target = workspace / ARCHIVE_DIRNAME / f"{stamp}-{name}" + target.parent.mkdir(parents=True, exist_ok=True) + path.rename(target) + moved.append(target) + return moved + + # -------------------------------------------------------------- reports + + def refresh(self, state: Any = None) -> None: + """Regenerate the reports from the ledger, and journal verdicts that just went stale. + + Called after every checkpoint, so the page never trails the run by + more than one step. The git-history switch check is left to delivery: + it costs a ``git diff`` per verdict commit, and this runs often. + A report that cannot be written is journaled rather than raised; it + is regenerated at the next checkpoint and must not stop a run. + """ + wf = self._workflow + try: + facts = collect_facts( + wf.workspace, + repo_root=self.repo_root() or None, + task_data=getattr(wf, "task_data", None) or None, + scope=self.reportable(), + policy=self.policy, + ) + write_reports(wf.workspace, facts=facts, final_report=bool(getattr(state, "done", 0))) + except (OSError, ValueError, yaml.YAMLError) as exc: + wf.journal.event("reports not regenerated", f"{type(exc).__name__}: {exc}") + return + newly = [ + item + for item in facts.assessments.values() + if item.stale and item.verdict.id not in self._warned + ] + if newly: + self._warned.update(item.verdict.id for item in newly) + wf.journal.section( + f"evidence: {len(newly)} verdict(s) went stale", + "\n".join(f"- {item.verdict.label()}: {'; '.join(item.reasons)}" for item in newly) + + "\n\nDelivery re-runs them on the delivered tree.", + ) + + # ------------------------------------------------------------ recording + + def unit(self) -> str: + """The unit whose gates are running now, as the orchestrator last set it.""" + ctx = getattr(self._workflow, "_progress_ctx", None) + return str(getattr(ctx, "current_unit", "") or "") + + def repo_root(self) -> str: + return str((getattr(self._workflow, "task_data", None) or {}).get("trtllm_repo_path") or "") + + def context(self, unit: str | None = None) -> VerdictContext: + """Everything a verdict records that the gate result does not carry.""" + wf = self._workflow + task = dict(getattr(wf, "task_data", None) or {}) + unit = self.unit() if unit is None else unit + preflight = getattr(wf, "_preflight", None) + env = dict(((task.get("environment") or {}).get("env")) or {}) + if getattr(wf, "environment", None) is not None: + env = dict(wf._gate_env().env) + features: list[str] = [] + if unit and task: + features = list(wf._features_through(unit)) + return VerdictContext( + repo_root=self.repo_root(), + target_rel=str(getattr(preflight, "target_rel", "") or ""), + unit=unit, + env=env, + features=features, + task_data=task or None, + policy=self.policy, + agents=role_agents(task), + ) + + def record(self, result: Any) -> Verdict: + """The ``GateLog`` sink: put one gate result into the ledger.""" + return record_gate_result(self.ledger, result, self.context(), gate=gate_key(result)) + + def record_checks(self, prefix: str, checks: Iterable[_CheckResult]) -> list[Verdict]: + """Record checks that are not engine gates, under ``prefix``, against one reading. + + They carry no unit and no target dependencies: they are about the + repository as a whole, and their freshness is that they run again. + """ + context = replace(self.context(unit=""), target_rel="") + measurement = measure(context) + return [ + record_gate_result( + self.ledger, check, context, gate=f"{prefix}{check.name}", measurement=measurement + ) + for check in checks + ] + + def record_early_checks(self, report: Any) -> Any: + """Record an early-checks report: a ``warn`` is a FINDING, ``ok``/``info`` a PASS.""" + self.record_checks( + EARLY_CHECK_PREFIX, + ( + _CheckResult( + name=str(f.name), + passed=True, + detail=str(f.detail), + finding=str(f.level) == "warn", + ) + for f in getattr(report, "findings", ()) + ), + ) + return report + + def record_delivery_checks(self, report: Any) -> Any: + """Record every delivery check (the capacity audit and matrix evidence among them).""" + self.record_checks( + DELIVERY_CHECK_PREFIX, + ( + _CheckResult(name=str(c.name), passed=bool(c.passed), detail=str(c.detail)) + for c in getattr(report, "checks", ()) + ), + ) + return report + + # ------------------------------------------------------------ freshness + + def in_scope(self) -> Callable[[Verdict], bool]: + """Whether the current spec still requires a verdict: its unit, and its member.""" + wf = self._workflow + task = getattr(wf, "task_data", None) or {} + units = assembly_units(task) + world = wf._gate_env().world_size if getattr(wf, "environment", None) is not None else 1 + + plan = getattr(getattr(wf, "_preflight", None), "certification", None) + + def keep(verdict: Verdict) -> bool: + if verdict.unit not in units: + return False + if verdict.gate.startswith(f"{envelope_mod.PROBE_GATE}[") and plan is not None: + return self.probe_for(verdict.gate, verdict.unit, verdict.member) is not None + return not verdict.member or verdict.member in members_for(task, verdict.unit, world) + + return keep + + def reportable(self) -> Callable[[Verdict], bool]: + """What the reports show: every required verdict, and every unit-less check.""" + keep = self.in_scope() + return lambda verdict: not verdict.unit or keep(verdict) + + def probe_for(self, gate: str, unit: str, member: str) -> Any: + """The certification plan's probe a ledger key names, or None.""" + plan = getattr(getattr(self._workflow, "_preflight", None), "certification", None) + for probe in plan.probes_for(unit) if plan is not None else (): + if probe.member == member and gate == f"{envelope_mod.PROBE_GATE}[{probe.point()}]": + return probe + return None + + def current(self) -> tuple[TreeIdentity, NativeBuild]: + """The tree and native build the ledger's verdicts would be delivered on.""" + repo = self.repo_root() + deps: set[str] = set() + for verdict in self.ledger.latest().values(): + deps.update(verdict.dependencies) + tree = tree_identity(repo, dependencies=hash_paths(repo, sorted(deps))) + return tree, native_build(repo, globs=self.policy.native_globs) + + def assess(self, *, check_history: bool = True) -> list[Assessment]: + """The latest in-scope verdicts, judged against the tree as it is now.""" + keep = self.in_scope() + return assess( + [v for v in self.ledger.latest().values() if keep(v)], + self.repo_root(), + task_data=getattr(self._workflow, "task_data", None), + policy=self.policy, + check_history=check_history, + ) + + def queue(self) -> tuple[list[RerunEntry], TreeIdentity, NativeBuild]: + """What delivery must run again before the evidence describes this tree.""" + tree, native = self.current() + return rerun_queue(self.assess(), tree=tree, native=native), tree, native + + # ------------------------------------------------------------- re-runs + + def rerunner(self, entry: RerunEntry, log: Any = None) -> Callable[[], Any] | None: + """A callable that re-measures ``entry`` and records it, or None if unknown. + + Only gates the orchestrator itself fires are known here, invoked the + way the assembly loop invokes them. A gate that cannot be re-run from + here stays queued, and delivery refuses on it rather than guessing. + """ + from .workflow import WITNESS_GATED_UNITS + + wf = self._workflow + unit, member = entry.unit, entry.member + if entry.gate.startswith(f"{envelope_mod.PROBE_GATE}["): + probe = self.probe_for(entry.gate, unit, member) + if probe is None: + return None + return lambda: wf.gate_log.record( + envelope_mod.run_envelope_probe( + wf._gate_env(), + task_data=wf.task_data, + probe=probe, + spec=gates_mod.BootSpec.load(wf.boot_spec_path), + runner=wf._shell_runner, + features=wf._features_through(unit), + ) + ) + if entry.gate.startswith("boot"): + return lambda: wf.gate_log.record( + gates_mod.run_boot_gate( + wf._gate_env(), + task_data=wf.task_data, + spec=gates_mod.BootSpec.load(wf.boot_spec_path), + runner=wf._shell_runner, + features=wf._features_through(unit), + unit=unit, + member=member, + ) + ) + if entry.gate.startswith("accuracy"): + return lambda: wf.gate_log.record( + gates_mod.run_accuracy_gate( + wf._gate_env(), + task_data=wf.task_data, + reference=wf._reference_score(unit), + runner=wf._shell_runner, + unit=unit, + member=member, + ) + ) + if entry.gate == "acceptance": + return lambda: wf.gate_log.record( + gates_mod.run_acceptance_gate( + wf._gate_env(), + task_data=wf.task_data, + spec=gates_mod.BootSpec.load(wf.boot_spec_path), + runner=wf._shell_runner, + unit=unit, + features=wf._features_through(unit), + member=member, + ) + ) + if unit in WITNESS_GATED_UNITS: + # Records itself; returns None on a misconfigured block. + return lambda: wf._run_witness(unit, member, log) + return None + + def rerun(self, entry: RerunEntry, log: Any = None) -> bool | None: + """Re-measure one queued verdict. See :func:`freshness.drain` for the answer.""" + run = self.rerunner(entry, log) + if run is None: + return None + wf = self._workflow + wf._progress_ctx.current_unit = entry.unit + wf.journal.event("gate re-run for delivery", f"{entry.gate}@{entry.member} [{entry.unit}]") + result = run() + if result is None: + return False + wf.journal.event("gate", result.summary()) + print_message(f" re-run: {escape_markup(result.summary())}", log) + return True + + # ------------------------------------------------------------- delivery + + def settle_delivery(self, state: Any, log: Any = None) -> bool: + """Bring the evidence onto the delivery tree, or close the run saying why not. + + 1. Queue every in-scope latest verdict that is stale, did not pass, + or was measured on another tree or native build. + 2. Re-run the queue. + 3. Record a full regression if every required verdict now passed on + this tree and build. + 4. Refuse unless nothing is still queued and the tree under merge is + the tree of the last full regression. + + Outside a git work tree there is no tree to bind, so only freshness is + checked; that is the same "absence of git is absence of a claim" rule + the delivery checks follow. + """ + wf = self._workflow + if not self.repo_root(): + return True + queue, tree, native = self.queue() + write_rerun_queue(wf.workspace, queue) + if queue: + wf.journal.section( + f"evidence: {len(queue)} gate(s) to re-run before delivery", _render_queue(queue) + ) + print_message( + f"[bold cyan]→ evidence: re-running {len(queue)} gate(s) so every verdict " + f"describes tree {tree.short()}[/bold cyan]", + log, + ) + drain(queue, lambda entry: self.rerun(entry, log)) + queue, tree, native = self.queue() + write_rerun_queue(wf.workspace, queue) + + problems: list[str] = [] + if queue: + problems.append( + f"{len(queue)} verdict(s) still do not hold on {tree.short()}:\n" + + _render_queue(queue) + ) + else: + self._record_regression(tree, native) + binding = regression_binding(self.ledger, self.repo_root(), policy=self.policy) + if binding.claim and not binding.ok: + problems.append(f"final regression binding: {binding.detail}") + self.refresh(state) + contradictions = check_consistency(wf.workspace) + if contradictions: + problems.append( + f"{len(contradictions)} report statement(s) contradict the ledger:\n" + + "\n".join(f"- {c.render()}" for c in contradictions) + ) + + body = "\n\n".join(problems) if problems else f"ok: {binding.detail}" + wf.journal.section("evidence at delivery", body) + if problems: + print_message( + "[bold red]✘ delivery refused: the evidence does not describe the tree being " + "delivered.[/bold red]\n" + escape_markup(body), + log, + ) + wf._close_run(state, log, verdict=VERDICT_EVIDENCE) + return False + print_message(f"[bold green]✔ evidence: {escape_markup(binding.detail)}[/bold green]", log) + return True + + def annotate_premerge(self, pack: str | Path) -> str: + """Append the final-regression binding to the pre-merge pack; return a one-line summary. + + The pack's remaining steps include a lint pass and a split into PR + branches, and both can produce a tree no regression ran on. So the + pack says which tree is bound, and gives the command that re-checks + whatever is finally merged. + """ + repo = self.repo_root() + binding = regression_binding(self.ledger, repo, policy=self.policy) + workspace = Path(self._workflow.workspace) + command = ( + f"python -m agent_flow.workflows.staircase_onboard.freshness check-merge " + f"{workspace} --repo {repo} --rev <the branch or commit being merged>" + ) + state = "bound" if binding.ok else "NOT bound" + section = [ + "", + f"## Evidence binding ({state})", + "", + binding.detail + ".", + "", + "A lint fix, a split or any later commit changes the tree; re-check the exact rev " + "that merges, and re-enter delivery (which re-runs the gates on it) if it is refused:", + "", + f"```\n{command}\n```", + "", + ] + path = Path(pack) + if path.is_file(): + with path.open("a", encoding="utf-8") as handle: + handle.write("\n".join(section)) + return f"evidence {state}: {binding.detail}" + + def _record_regression(self, tree: TreeIdentity, native: NativeBuild) -> None: + latest = self.ledger.latest() + regression = complete_regression(latest, tree, native, scope=self.in_scope()) + if regression is None: + return + last = last_regression(self.ledger) + if ( + last is not None + and last.tree.key == regression.tree.key + and last.native.fingerprint == regression.native.fingerprint + and set(last.verdict_ids) == set(regression.verdict_ids) + ): + return + record_regression(self.ledger, regression) + self._workflow.journal.event( + "full regression recorded", + f"{len(regression.keys)} verdicts on {tree.short()} (native {native.short()})", + ) + + +def _render_queue(queue: list[RerunEntry]) -> str: + lines = [ + f"- {e.gate}@{e.member or '-'} [{e.unit}]: {'; '.join(e.reasons)}" + for e in queue[:_QUEUE_LINES_SHOWN] + ] + if len(queue) > _QUEUE_LINES_SHOWN: + lines.append(f"- (+{len(queue) - _QUEUE_LINES_SHOWN} more in rerun_queue.json)") + return "\n".join(lines) + + +__all__ = [ + "DELIVERY_CHECK_PREFIX", + "EARLY_CHECK_PREFIX", + "RunEvidence", + "VERDICT_EVIDENCE", + "gate_key", + "role_agents", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/freshness.py b/agent-flow/agent_flow/workflows/staircase_onboard/freshness.py new file mode 100644 index 000000000000..d1ef61626589 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/freshness.py @@ -0,0 +1,542 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Evidence freshness: which recorded verdicts still hold on the tree as it is now. + +The ledger (:mod:`.ledger`) stores each verdict with the code it depends on. +This module turns that into three answers: + +1. **Which verdicts are stale.** :func:`assess` compares each verdict against + the current tree: its dependency files, its dependency set, the native + build, the unit's spec, and the defaults of the switches it did not set + itself. A later commit that touches the target makes every earlier unit's + verdict on that target stale. A flipped default makes stale every verdict + that ran under the old one, including those of earlier units. +2. **What to run again.** :func:`rerun_queue` lists the latest verdict of + each ``(gate, unit, member)`` that is stale, did not pass, or -- at + delivery -- was measured on a different tree than the one being + delivered. :func:`drain` runs that queue through a caller-supplied + re-runner and returns what is still owed. +3. **Whether the tree under merge was regressed.** A *full regression* is a + set of passing verdicts covering every required key, all measured on one + tree and one native build. :func:`complete_regression` recognises one and + the ledger records it; :func:`regression_binding` refuses a merge whose + tree or native build differs from the last one. This is the rule the + DeepSeek-V4.1 run lacked: its final head carried a functional change and a + lint series that landed after the last regression, so no gate ever ran on + the code that was delivered. + +Reading the tree costs a git call or two and one hash per dependency file. +The switch-history check adds one ``git diff -G`` per distinct verdict commit. +""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import sys +import threading +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any, Callable, Iterable, Mapping, Sequence + +from .ledger import ( + DEFAULT_SWITCH_PREFIXES, + KIND_REGRESSION, + SCHEMA_VERSION, + Ledger, + LedgerPolicy, + NativeBuild, + TreeIdentity, + Verdict, + dependency_paths, + harvest_paths, + harvest_switch_defaults, + hash_paths, + is_git_repo, + load_task, + native_build, + run_git, + spec_fingerprint, + tree_identity, + utc_now, +) + +RERUN_QUEUE_FILENAME = "rerun_queue.json" +_REASON_PATH_CAP = 5 + + +# -------------------------------------------------------------- regressions + + +@dataclass(frozen=True) +class Regression: + """A complete set of passing verdicts, every one measured on one tree and build.""" + + tree: TreeIdentity + native: NativeBuild + verdict_ids: tuple[str, ...] + keys: tuple[tuple[str, str, str], ...] + timestamp: str = "" + + def to_record(self) -> dict[str, Any]: + return { + "kind": KIND_REGRESSION, + "schema": SCHEMA_VERSION, + "timestamp": self.timestamp or utc_now(), + "tree": self.tree.to_json(), + "native_build": self.native.to_json(), + "verdict_ids": list(self.verdict_ids), + "keys": [list(k) for k in self.keys], + } + + @classmethod + def from_record(cls, raw: Mapping[str, Any]) -> "Regression": + return cls( + tree=TreeIdentity.from_json(raw.get("tree")), + native=NativeBuild.from_json(raw.get("native_build")), + verdict_ids=tuple(str(v) for v in raw.get("verdict_ids") or ()), + keys=tuple( + tuple(str(p) for p in k) # type: ignore[misc] + for k in raw.get("keys") or () + if isinstance(k, (list, tuple)) and len(k) == 3 + ), + timestamp=str(raw.get("timestamp") or ""), + ) + + +def regressions(ledger: Ledger) -> list[Regression]: + return [Regression.from_record(r) for r in ledger.records() if r.get("kind") == KIND_REGRESSION] + + +def last_regression(ledger: Ledger) -> Regression | None: + found = regressions(ledger) + return found[-1] if found else None + + +def record_regression(ledger: Ledger, regression: Regression) -> Regression: + record = regression.to_record() + ledger.append(record) + return replace(regression, timestamp=record["timestamp"]) + + +# -------------------------------------------------------------- invalidation + + +@dataclass(frozen=True) +class Assessment: + """A verdict and the reasons it no longer holds (none means fresh).""" + + verdict: Verdict + reasons: tuple[str, ...] = () + + @property + def stale(self) -> bool: + return bool(self.reasons) + + +def _capped(prefix: str, paths: Sequence[str]) -> str: + shown = ", ".join(paths[:_REASON_PATH_CAP]) + more = f" (+{len(paths) - _REASON_PATH_CAP} more)" if len(paths) > _REASON_PATH_CAP else "" + return f"{prefix}: {shown}{more}" + + +def switch_history_flips( + repo_root: str | Path, + commit: str, + *, + prefixes: Sequence[str] = DEFAULT_SWITCH_PREFIXES, +) -> dict[str, tuple[str | None, str | None]]: + """Switch defaults that differ between ``commit`` and the working tree. + + Only files whose diff touches a line naming a switch are read, which + ``git diff -G`` selects without reading anything else. Returns + ``{name: (old_default, new_default)}``; ``{}`` when git cannot answer. + """ + if not commit or not prefixes: + return {} + pattern = "(" + "|".join(re.escape(p) for p in prefixes) + ")" + out = run_git(repo_root, "diff", "--name-only", "-G", pattern, commit, "--", "*.py") + if out is None: + return {} + root = Path(repo_root) + old_all: dict[str, set[str | None]] = {} + new_all: dict[str, set[str | None]] = {} + for rel in (line.strip() for line in out.splitlines() if line.strip()): + old_text = run_git(repo_root, "show", f"{commit}:{rel}") or "" + try: + new_text = (root / rel).read_text(encoding="utf-8") + except (OSError, UnicodeDecodeError): + new_text = "" + for name, default in harvest_switch_defaults(old_text, prefixes=prefixes).items(): + old_all.setdefault(name, set()).add(default) + for name, default in harvest_switch_defaults(new_text, prefixes=prefixes).items(): + new_all.setdefault(name, set()).add(default) + + def merged(values: set[str | None] | None) -> str | None: + if not values: + return None + if len(values) == 1: + return next(iter(values)) + return " | ".join(sorted("None" if v is None else v for v in values)) + + flips: dict[str, tuple[str | None, str | None]] = {} + for name in sorted(set(old_all) & set(new_all)): + old, new = merged(old_all[name]), merged(new_all[name]) + if old != new: + flips[name] = (old, new) + return flips + + +def assess( + verdicts: Iterable[Verdict], + repo_root: str | Path, + *, + task_data: Mapping[str, Any] | None = None, + policy: LedgerPolicy = LedgerPolicy(), + native: NativeBuild | None = None, + check_history: bool = True, +) -> list[Assessment]: + """Decide, against the tree at ``repo_root`` now, which verdicts still hold. + + A verdict is stale when: + + - a dependency file's content differs from what was measured; + - a file joined the target's dependency set (a new module, a newly + imported catalog entry); + - the native build's fingerprint differs (both known); + - ``task_data`` is given and the unit's spec projection differs; + - a switch the verdict did not set explicitly has a different default + now. Checked two ways: against the defaults recorded with the verdict, + and, with ``check_history``, against git history from the verdict's + commit, which also sees switches in files outside the recorded scope. + """ + root = Path(repo_root) + native = native if native is not None else native_build(root, globs=policy.native_globs) + current_deps: dict[str, list[str]] = {} + history: dict[str, dict[str, tuple[str | None, str | None]]] = {} + out: list[Assessment] = [] + for verdict in verdicts: + reasons: list[str] = [] + recorded = dict(verdict.dependencies) + now = hash_paths(root, recorded) + changed = sorted(p for p, h in recorded.items() if now.get(p) != h) + if changed: + reasons.append(_capped("code changed", changed)) + if verdict.target_rel: + if verdict.target_rel not in current_deps: + current_deps[verdict.target_rel] = dependency_paths( + root, verdict.target_rel, extra=policy.extra_dependencies + ) + added = sorted(set(current_deps[verdict.target_rel]) - set(recorded)) + if added: + reasons.append(_capped("new dependency", added)) + if ( + verdict.native.known + and native.known + and verdict.native.fingerprint != native.fingerprint + ): + reasons.append(f"native build changed: {verdict.native.short()} -> {native.short()}") + if task_data is not None and verdict.spec_fingerprint: + if spec_fingerprint(task_data, verdict.unit) != verdict.spec_fingerprint: + reasons.append(f"task spec for unit {verdict.unit or '-'} changed") + flips: dict[str, tuple[str | None, str | None]] = {} + if verdict.switch_defaults: + scope = [*recorded, *policy.switch_sources] + current = harvest_paths(root, scope, prefixes=policy.switch_prefixes) + for name, old in verdict.switch_defaults.items(): + if name in current and current[name] != old: + flips[name] = (old, current[name]) + if check_history and verdict.tree.commit: + if verdict.tree.commit not in history: + history[verdict.tree.commit] = switch_history_flips( + root, verdict.tree.commit, prefixes=policy.switch_prefixes + ) + flips.update(history[verdict.tree.commit]) + for name, (old, new) in sorted(flips.items()): + if name in verdict.env: + continue + reasons.append(f"default switch flipped: {name} {old!r} -> {new!r}") + out.append(Assessment(verdict=verdict, reasons=tuple(reasons))) + return out + + +# ------------------------------------------------------------- re-run queue + + +@dataclass(frozen=True) +class RerunEntry: + gate: str + unit: str + member: str + reasons: tuple[str, ...] + verdict_id: str = "" + + @property + def key(self) -> tuple[str, str, str]: + return (self.gate, self.unit, self.member) + + def to_json(self) -> dict[str, Any]: + return { + "gate": self.gate, + "unit": self.unit, + "member": self.member, + "reasons": list(self.reasons), + "verdict_id": self.verdict_id, + } + + @classmethod + def from_json(cls, raw: Mapping[str, Any]) -> "RerunEntry": + return cls( + gate=str(raw.get("gate") or ""), + unit=str(raw.get("unit") or ""), + member=str(raw.get("member") or ""), + reasons=tuple(str(r) for r in raw.get("reasons") or ()), + verdict_id=str(raw.get("verdict_id") or ""), + ) + + +def rerun_queue( + assessments: Iterable[Assessment], + *, + tree: TreeIdentity | None = None, + native: NativeBuild | None = None, + scope: Callable[[Verdict], bool] | None = None, +) -> list[RerunEntry]: + """What must be measured again, in the order the verdicts were recorded. + + A latest verdict is queued when it is stale, when it did not pass, or -- + given ``tree`` -- when it was measured on a different tree or native + build than ``tree``/``native``. The last is what binds the final + regression to the tree under merge: a verdict carried over from an + earlier commit is evidence about that commit. ``scope`` drops verdicts + the run no longer requires (a unit or member the spec no longer lists). + """ + queue: list[RerunEntry] = [] + for item in assessments: + verdict = item.verdict + if scope is not None and not scope(verdict): + continue + reasons = list(item.reasons) + if not verdict.passed: + reasons.append(f"last verdict {verdict.outcome}") + if tree is not None and tree.known: + if verdict.tree.key != tree.key: + reasons.append(f"measured on tree {verdict.tree.short()}, not on {tree.short()}") + elif native is not None and verdict.native.fingerprint != native.fingerprint: + reasons.append( + f"measured on native build {verdict.native.short()}, not on {native.short()}" + ) + if reasons: + queue.append( + RerunEntry( + gate=verdict.gate, + unit=verdict.unit, + member=verdict.member, + reasons=tuple(reasons), + verdict_id=verdict.id, + ) + ) + return queue + + +def write_rerun_queue(workspace: str | Path, entries: Sequence[RerunEntry]) -> Path: + """Persist the queue beside the ledger, atomically.""" + path = Path(workspace) / RERUN_QUEUE_FILENAME + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(f".{path.name}.{os.getpid()}.{threading.get_ident()}.tmp") + tmp.write_text( + json.dumps( + {"generated_at": utc_now(), "entries": [e.to_json() for e in entries]}, + indent=2, + ensure_ascii=False, + ) + + "\n", + encoding="utf-8", + ) + os.replace(tmp, path) + return path + + +def read_rerun_queue(workspace: str | Path) -> list[RerunEntry]: + path = Path(workspace) / RERUN_QUEUE_FILENAME + try: + raw = json.loads(path.read_text(encoding="utf-8")) + except (FileNotFoundError, json.JSONDecodeError): + return [] + return [RerunEntry.from_json(e) for e in raw.get("entries") or () if isinstance(e, Mapping)] + + +def drain( + queue: Sequence[RerunEntry], run: Callable[[RerunEntry], bool | None] +) -> list[RerunEntry]: + """Re-run each entry with ``run`` and return the ones still owed. + + ``run`` returns True when it re-measured the entry (the new verdict is in + the ledger whatever it says), False when the re-run failed, and None when + it does not know how to re-run that gate. Both of the latter stay queued. + """ + remaining: list[RerunEntry] = [] + for entry in queue: + outcome = run(entry) + if outcome is not True: + note = "re-run failed" if outcome is False else "no automatic re-run for this gate" + remaining.append(replace(entry, reasons=(*entry.reasons, note))) + return remaining + + +# ---------------------------------------------------------------- regression + + +def complete_regression( + latest: Mapping[tuple[str, str, str], Verdict], + tree: TreeIdentity, + native: NativeBuild, + *, + scope: Callable[[Verdict], bool] | None = None, +) -> Regression | None: + """A regression record if every in-scope latest verdict passed on ``tree``/``native``.""" + chosen = [v for v in latest.values() if scope is None or scope(v)] + if not chosen or not tree.known: + return None + for verdict in chosen: + if not verdict.passed or verdict.tree.key != tree.key: + return None + if verdict.native.fingerprint != native.fingerprint: + return None + return Regression( + tree=tree, + native=native, + verdict_ids=tuple(v.id for v in chosen), + keys=tuple(v.key for v in chosen), + timestamp=utc_now(), + ) + + +@dataclass(frozen=True) +class Binding: + """Whether the tree under merge is the tree of the last full regression.""" + + ok: bool + detail: str + #: False outside a git work tree, where no claim about a tree can be made. + claim: bool = True + + +def regression_binding( + ledger: Ledger, + repo_root: str | Path, + *, + rev: str = "HEAD", + native_root: str | Path | None = None, + policy: LedgerPolicy = LedgerPolicy(), +) -> Binding: + """Compare the tree under merge, and its native build, with the last regression. + + The tree under merge is the commit ``rev`` -- uncommitted changes are not + merged -- and its native build is the one under ``native_root`` (the repo + itself by default). + """ + if not is_git_repo(repo_root): + return Binding(ok=True, claim=False, detail="not a git work tree; no tree to bind") + merge_tree = tree_identity(repo_root, rev=rev) + merge_native = native_build(native_root or repo_root, globs=policy.native_globs) + last = last_regression(ledger) + if last is None: + return Binding( + ok=False, + detail=( + f"no full regression is recorded; the tree under merge is " + f"{merge_tree.short()} (native {merge_native.short()})" + ), + ) + problems: list[str] = [] + if last.tree.key != merge_tree.key: + problems.append( + f"the last full regression ran on {last.tree.short()} " + f"({last.tree.tree[:12] or '?'}), the tree under merge is {merge_tree.short()} " + f"({merge_tree.tree[:12]})" + ) + if last.native.fingerprint != merge_native.fingerprint: + problems.append( + f"the last full regression ran on native build {last.native.short()}, " + f"the tree under merge has {merge_native.short()}" + ) + if problems: + return Binding(ok=False, detail="; ".join(problems)) + return Binding( + ok=True, + detail=( + f"the last full regression ({len(last.keys)} verdicts, {last.timestamp}) ran on " + f"{merge_tree.short()} with native {merge_native.short()}" + ), + ) + + +# ----------------------------------------------------------------------- CLI + + +def main(argv: Sequence[str] | None = None) -> int: + """``status`` and ``check-merge`` over a workspace's ledger.""" + parser = argparse.ArgumentParser(prog="staircase-freshness", description=main.__doc__) + sub = parser.add_subparsers(dest="command", required=True) + + status = sub.add_parser("status", help="list stale or failing verdicts and the re-run queue") + status.add_argument("workspace") + status.add_argument("--repo", required=True) + status.add_argument("--task", help="task.yaml, to detect spec changes") + status.add_argument("--write-queue", action="store_true") + + merge = sub.add_parser( + "check-merge", help="refuse unless the tree under merge was fully regressed" + ) + merge.add_argument("workspace") + merge.add_argument("--repo", required=True) + merge.add_argument("--rev", default="HEAD") + merge.add_argument("--native-root", help="where the native build lives, if not in --repo") + + args = parser.parse_args(argv) + ledger = Ledger.in_workspace(args.workspace) + if args.command == "status": + assessed = assess(ledger.latest().values(), args.repo, task_data=load_task(args.task)) + queue = rerun_queue(assessed) + for item in assessed: + state = "STALE" if item.stale else "fresh" + sys.stdout.write( + f"{state:5} {item.verdict.outcome:7} {item.verdict.label()} " + f"tree {item.verdict.tree.short()}\n" + ) + for reason in item.reasons: + sys.stdout.write(f" - {reason}\n") + if args.write_queue: + sys.stdout.write(f"queue: {write_rerun_queue(args.workspace, queue)}\n") + return 1 if queue else 0 + binding = regression_binding(ledger, args.repo, rev=args.rev, native_root=args.native_root) + sys.stdout.write(f"{'ok' if binding.ok else 'REFUSED'}: {binding.detail}\n") + return 0 if binding.ok else 1 + + +__all__ = [ + "Assessment", + "Binding", + "RERUN_QUEUE_FILENAME", + "Regression", + "RerunEntry", + "assess", + "complete_regression", + "drain", + "last_regression", + "read_rerun_queue", + "record_regression", + "regression_binding", + "regressions", + "rerun_queue", + "switch_history_flips", + "write_rerun_queue", +] + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/gates.py b/agent-flow/agent_flow/workflows/staircase_onboard/gates.py index 3c0a7c93a587..02b0ef8335b2 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/gates.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/gates.py @@ -1792,9 +1792,15 @@ class GateLog: """Every gate a run has fired, in order. The journal's raw material.""" results: list[GateResult] = field(default_factory=list) + #: Called with every result as it is recorded. The workflow points it at + #: the evidence ledger (``ledger.record_gate_result``), so every verdict a + #: run produces enters the ledger here and nowhere else. + sink: Callable[[GateResult], Any] | None = None def record(self, result: GateResult) -> GateResult: self.results.append(result) + if self.sink is not None: + self.sink(result) return result def last(self, name_prefix: str) -> GateResult | None: diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/journal.py b/agent-flow/agent_flow/workflows/staircase_onboard/journal.py index 6ecdbbe0a033..d7a72ced9557 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/journal.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/journal.py @@ -21,13 +21,18 @@ from __future__ import annotations +import os import re +import shutil import threading from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path _REPORT_NAME = re.compile(r"^(\d+)-") +#: Beside ``reports/``, not in it: readers of the reports directory expect +#: nothing but reports there. +_CLAIMS_DIRNAME = ".dispatch-ordinals" def _utc_now() -> str: @@ -63,6 +68,8 @@ def open(self, *, slug: str, header_lines: list[str] | None = None) -> None: if self.path.exists() and self.path.read_text(encoding="utf-8").strip(): self.event("run resumed", f"slug={slug}") return + # A new journal is a new run: its dispatch numbers start again at 01. + shutil.rmtree(self.claims_dir, ignore_errors=True) lines = [f"# Onboarding journal — {slug}", ""] lines.extend(header_lines or []) lines.append("") @@ -86,23 +93,52 @@ def section(self, title: str, body: str) -> None: with self.path.open("a", encoding="utf-8") as handle: handle.write(f"\n## {_utc_now()} — {title}\n\n{body.rstrip()}\n") + @property + def claims_dir(self) -> Path: + """Where each dispatch number is claimed, one empty file per number.""" + return self.path.parent / _CLAIMS_DIRNAME + + def _highest_ordinal(self) -> int: + seen = [0] + if self.reports_dir.is_dir(): + seen.extend( + int(match.group(1)) + for entry in self.reports_dir.iterdir() + if (match := _REPORT_NAME.match(entry.name)) + ) + if self.claims_dir.is_dir(): + seen.extend( + int(entry.name) for entry in self.claims_dir.iterdir() if entry.name.isdigit() + ) + return max(seen) + def next_ordinal(self) -> int: - """The next dispatch number, counted across the whole run. + """Claim and return the next dispatch number, counted across the whole run. Sequential across every agent, so ``NN`` names exactly one dispatch and the next ordinal is visible without counting. + + **Claiming is the allocation**, and it is atomic across threads and + processes: the number is taken by creating ``.dispatch-ordinals/<NN>`` + with ``O_CREAT | O_EXCL``, which exactly one caller can win. Reading + the highest number and then writing the report, as this used to, let + two concurrent smiths of one batch both read the same maximum and + both record dispatch ``NN`` (duplicate numbers in ``journal.md``, + measured on the DeepSeek-V4.1 run). A loser moves to the next number. """ - if not self.reports_dir.is_dir(): - return 1 - seen = [ - int(match.group(1)) - for entry in self.reports_dir.iterdir() - if (match := _REPORT_NAME.match(entry.name)) - ] - return max(seen, default=0) + 1 + self.claims_dir.mkdir(parents=True, exist_ok=True) + candidate = self._highest_ordinal() + 1 + while True: + try: + fd = os.open(self.claims_dir / str(candidate), os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except FileExistsError: + candidate += 1 + continue + os.close(fd) + return candidate def record_dispatch(self, agent: str, what: str, prompt: str) -> Path: - """Write a dispatch prompt down **before it is sent**, and reserve its ordinal. + """Write a dispatch prompt down **before it is sent**, under a freshly claimed ordinal. An agent's context does not survive a session cut, so a child's prompt only survives if its parent wrote it first. This is also why diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/ledger.py b/agent-flow/agent_flow/workflows/staircase_onboard/ledger.py new file mode 100644 index 000000000000..0792f1c523ec --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/ledger.py @@ -0,0 +1,1176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The results ledger: every gate verdict, tied to the code it was measured on. + +A gate verdict is a claim about a specific piece of code under a specific +configuration. Once the code moves, the claim describes code that no longer +exists. A record that keeps only "accuracy PASS" therefore goes stale without +anyone noticing, and a report written from such records can say "16/16 PASS" +about a tree no gate ever ran on. + +This module stores each verdict together with what it depends on: + +- the gate, unit and member, the measured value and the threshold it was + judged against, and when it was measured; +- the **tree**: the commit, its tree object, and any uncommitted changes + among the files the verdict depends on; +- the **native build**: a content fingerprint of the compiled libraries the + engine loaded; +- the **dependency paths**: the target's own files, the family-level shared + modules beside it, and every catalog entry the target imports, each with + its git blob hash at measurement time; +- the **configuration**: the gate command, the explicit environment switches + (both fingerprinted), the switch defaults the code declared at the time, + and a fingerprint of the part of ``task.yaml`` that applies to the unit. + +From that, :mod:`.freshness` decides later, against any tree, whether a +verdict still holds, and which gates delivery has to run again. + +The ledger is one JSON Lines file in the workspace, and it is append-only. +Nothing in it is ever rewritten. A newer verdict for the same +``(gate, unit, member)`` key supersedes an older one by position in the file. + +Gates do not import this module. ``gates.GateLog`` carries a sink, and the +workflow points that sink at :func:`record_gate_result`. That is the single +place verdicts enter the ledger. A harness outside the workflow uses the same +function or the ``record`` subcommand of this module's CLI. +""" + +from __future__ import annotations + +import argparse +import ast +import fcntl +import fnmatch +import hashlib +import json +import os +import re +import shlex +import subprocess +import sys +import threading +import time +from dataclasses import dataclass, field, replace +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + +from .delivery import CATALOG_DIR, MODELING_V2_PACKAGE, imported_catalog_entries +from .task_schema import assembly_units + +LEDGER_FILENAME = "ledger.jsonl" +SCHEMA_VERSION = 1 + +KIND_VERDICT = "verdict" +KIND_REGRESSION = "regression" + +OUTCOME_PASS = "PASS" +OUTCOME_FAIL = "FAIL" +#: Passed its threshold but flagged for a human. A finding does not block +#: delivery, but it must never disappear from a summary. +OUTCOME_FINDING = "FINDING" +OUTCOMES: tuple[str, ...] = (OUTCOME_PASS, OUTCOME_FAIL, OUTCOME_FINDING) + +#: Environment variables whose name starts with one of these are switches: +#: their explicit values are recorded, and their defaults in code are tracked. +DEFAULT_SWITCH_PREFIXES: tuple[str, ...] = ("TRTLLM_",) + +#: Where the compiled libraries the engine loads live, relative to the repo. +DEFAULT_NATIVE_GLOBS: tuple[str, ...] = ("tensorrt_llm/libs/*.so", "tensorrt_llm/*.so") + +#: ``task.yaml`` blocks that describe the substrate or the run's own +#: bookkeeping rather than what a gate measures. A change to them does not +#: invalidate a verdict. +#: ``agents`` is here too: which model writes the next change does not change +#: what an existing verdict measured. It is recorded in the verdict's config +#: (:attr:`VerdictContext.agents`) instead, as provenance. So is +#: ``trtllm_repo_path``: the code is compared by content, so the same code in +#: another checkout is the same evidence, and different code is caught anyway. +SPEC_EXCLUDED_KEYS: frozenset[str] = frozenset( + {"remediation", "source_embargo", "agents", "trtllm_repo_path"} +) +#: The ``environment`` keys that change what a gate measures. The rest +#: (partition, spool, mounts, time limits) only change where it runs. +SPEC_ENVIRONMENT_KEYS: tuple[str, ...] = ("python", "env", "container_image") + +_GIT_TIMEOUT_S = 30 +#: Native libraries at most this large are hashed whole; larger ones are +#: sampled. ``libtensorrt_llm.so`` is over 1 GB, and hashing it whole on a +#: networked filesystem would put tens of seconds in front of every verdict. +_FULL_HASH_LIMIT = 8 * 1024 * 1024 +_SAMPLE_EDGE = 1024 * 1024 +_SAMPLE_CHUNK = 64 * 1024 +_SAMPLE_COUNT = 16 + + +def utc_now() -> str: + return datetime.now(timezone.utc).isoformat(timespec="seconds") + + +def _canonical(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str) + + +def fingerprint(value: Any) -> str: + """A stable short digest of any JSON-able value.""" + return "sha256:" + hashlib.sha256(_canonical(value).encode("utf-8")).hexdigest()[:24] + + +def run_git(repo_root: str | Path, *args: str, stdin: str | None = None) -> str | None: + """Ask git one question about ``repo_root``. ``None`` when it cannot answer. + + Every failure collapses to ``None``: no repository, no ``git`` binary, a + hung filesystem. A ledger that raised here would turn a tooling problem + into a failed gate. + """ + try: + completed = subprocess.run( + ["git", "-C", str(repo_root), *args], + capture_output=True, + text=True, + timeout=_GIT_TIMEOUT_S, + check=False, + input=stdin, + ) + except (OSError, subprocess.SubprocessError): + return None + if completed.returncode != 0: + return None + return completed.stdout + + +def is_git_repo(repo_root: str | Path) -> bool: + out = run_git(repo_root, "rev-parse", "--is-inside-work-tree") + return bool(out) and out.strip() == "true" + + +# ----------------------------------------------------------------- hashing + +#: A file modified this recently is not cached: on a filesystem with coarse +#: timestamps a second write in the same tick, of the same size, would +#: otherwise read back as the first (git's "racy" entries, the same answer). +_RACY_WINDOW_NS = 2_000_000_000 + + +def _stat_key(path: Path, stat: os.stat_result, *extra: Any) -> tuple[Any, ...] | None: + """A cache key for ``path`` as ``stat`` describes it, or None if too fresh to trust.""" + if time.time_ns() - stat.st_mtime_ns < _RACY_WINDOW_NS: + return None + return (str(path), stat.st_ino, stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns, *extra) + + +_blob_cache: dict[tuple[Any, ...], str] = {} +_blob_lock = threading.Lock() + + +def git_blob_hash(data: bytes) -> str: + """The id git gives ``data`` as a blob, without asking git.""" + return hashlib.sha1(f"blob {len(data)}\0".encode() + data).hexdigest() + + +def blob_hash(path: str | Path) -> str | None: + """The git blob id of the file at ``path``, or ``None`` if there is no file. + + The same id ``git ls-tree`` reports for a committed file, so a working-tree + file can be compared against ``HEAD`` without a ``git status`` over the + whole worktree (a minute on the networked filesystem these runs use). + Cached by ``(path, mtime, size)``. + """ + path = Path(path) + try: + stat = path.stat() + except OSError: + return None + if not path.is_file(): + return None + key = _stat_key(path, stat) + if key is not None: + with _blob_lock: + cached = _blob_cache.get(key) + if cached is not None: + return cached + try: + digest = git_blob_hash(path.read_bytes()) + except OSError: + return None + if key is not None: + with _blob_lock: + _blob_cache[key] = digest + return digest + + +# ------------------------------------------------------------ dependencies + + +def _is_dependency_file(path: Path) -> bool: + return ( + path.is_file() + and "__pycache__" not in path.parts + and not path.name.startswith(".") + and path.suffix not in (".pyc", ".pyo") + ) + + +def _expand(root: Path, rel: str) -> list[str]: + """Repo-relative files at ``rel``: the file itself, or every file under it.""" + path = root / rel + if path.is_file(): + return [Path(rel).as_posix()] + if not path.is_dir(): + return [] + return [p.relative_to(root).as_posix() for p in path.rglob("*") if _is_dependency_file(p)] + + +def dependency_paths( + repo_root: str | Path, target_rel: str, *, extra: Sequence[str] = () +) -> list[str]: + """The repo-relative files a verdict on ``target_rel`` depends on. + + - every file under the target directory; + - the family-level files directly beside it (shared modules, routing); + - the wrapper of every catalog entry any target module imports; + - ``extra``: further files or directories the caller names. + + The catalog contracts (``.md``) are not included: a contract records a + certification, and editing it does not change what the engine computes. + """ + root = Path(repo_root) + found: set[str] = set() + if target_rel: + target = root / target_rel + found.update(_expand(root, target_rel)) + family = target.parent + if family.is_dir(): + found.update( + p.relative_to(root).as_posix() for p in family.iterdir() if _is_dependency_file(p) + ) + for rel in sorted(found): + if not rel.endswith(".py") or not rel.startswith(target_rel): + continue + try: + entries = imported_catalog_entries(root / rel) + except (OSError, SyntaxError, UnicodeDecodeError, ValueError): + continue + for category, name in entries: + wrapper = f"{CATALOG_DIR}/{category}/{name}.py" + if (root / wrapper).is_file(): + found.add(wrapper) + for rel in extra: + found.update(_expand(root, rel)) + return sorted(found) + + +def hash_paths(repo_root: str | Path, paths: Iterable[str]) -> dict[str, str | None]: + root = Path(repo_root) + return {rel: blob_hash(root / rel) for rel in paths} + + +# ------------------------------------------------------------- tree identity + + +@dataclass(frozen=True) +class TreeIdentity: + """Which code a verdict was measured on. + + ``tree`` is the commit's tree object rather than the commit. Two commits + with the same content (a re-signed or rebased commit, say) have the same + tree, and evidence measured on one is evidence about the other. + + ``dirty`` lists the files among the verdict's dependencies whose + working-tree content differs from the commit. ``dirty_digest`` pins that + content, so two measurements over the same uncommitted edit compare equal + and two over different edits do not. + """ + + commit: str = "" + tree: str = "" + describe: str = "" + dirty: tuple[str, ...] = () + dirty_digest: str = "" + + @property + def known(self) -> bool: + return bool(self.tree) + + @property + def key(self) -> str: + """One token that compares equal exactly when the measured code does.""" + if not self.tree: + return "" + return f"{self.tree}+{self.dirty_digest}" if self.dirty else self.tree + + def short(self) -> str: + if not self.known: + return "unknown" + label = self.describe or self.commit[:10] + return f"{label}+dirty" if self.dirty else label + + def to_json(self) -> dict[str, Any]: + return { + "commit": self.commit, + "tree": self.tree, + "describe": self.describe, + "dirty": list(self.dirty), + "dirty_digest": self.dirty_digest, + } + + @classmethod + def from_json(cls, raw: Any) -> "TreeIdentity": + raw = raw if isinstance(raw, Mapping) else {} + return cls( + commit=str(raw.get("commit") or ""), + tree=str(raw.get("tree") or ""), + describe=str(raw.get("describe") or ""), + dirty=tuple(str(p) for p in raw.get("dirty") or ()), + dirty_digest=str(raw.get("dirty_digest") or ""), + ) + + +def committed_blobs(repo_root: str | Path, rev: str, paths: Sequence[str]) -> dict[str, str]: + """``path -> blob id`` at ``rev`` for those of ``paths`` that exist there.""" + if not paths: + return {} + out = run_git(repo_root, "ls-tree", "-r", "-z", rev, "--", *paths) + blobs: dict[str, str] = {} + for record in (out or "").split("\0"): + meta, _, path = record.partition("\t") + parts = meta.split() + if len(parts) == 3 and parts[1] == "blob": + blobs[path] = parts[2] + return blobs + + +def tree_identity( + repo_root: str | Path, + *, + rev: str = "HEAD", + dependencies: Mapping[str, str | None] | None = None, +) -> TreeIdentity: + """The identity of ``rev`` plus any uncommitted change among ``dependencies``. + + ``dependencies`` maps repo-relative paths to their working-tree blob ids, + as :func:`hash_paths` returns them. Pass none to describe the commit alone, + which is what the tree under merge is. + """ + out = run_git(repo_root, "rev-parse", rev, f"{rev}^{{tree}}") + lines = (out or "").split() + if len(lines) < 2: + return TreeIdentity() + described = (run_git(repo_root, "describe", "--always", rev) or "").strip() + dirty: list[str] = [] + if dependencies: + committed = committed_blobs(repo_root, rev, sorted(dependencies)) + dirty = sorted( + path + for path, blob in dependencies.items() + if blob != committed.get(path) and not (blob is None and path not in committed) + ) + digest = "" + if dirty: + digest = fingerprint({path: dependencies[path] for path in dirty}) # type: ignore[index] + return TreeIdentity( + commit=lines[0], tree=lines[1], describe=described, dirty=tuple(dirty), dirty_digest=digest + ) + + +# ------------------------------------------------------------- native build + + +@dataclass(frozen=True) +class NativeBuild: + """A content fingerprint of the compiled libraries the engine loads. + + ``fingerprint`` is empty when no library was found, which is the normal + state of a worktree that borrows another checkout's build. + """ + + fingerprint: str = "" + libraries: tuple[str, ...] = () + + @property + def known(self) -> bool: + return bool(self.fingerprint) + + def short(self) -> str: + return self.fingerprint.split(":", 1)[-1][:12] if self.fingerprint else "none" + + def to_json(self) -> dict[str, Any]: + return {"fingerprint": self.fingerprint, "libraries": list(self.libraries)} + + @classmethod + def from_json(cls, raw: Any) -> "NativeBuild": + raw = raw if isinstance(raw, Mapping) else {} + return cls( + fingerprint=str(raw.get("fingerprint") or ""), + libraries=tuple(str(x) for x in raw.get("libraries") or ()), + ) + + +_native_cache: dict[tuple[Any, ...], str] = {} +_native_lock = threading.Lock() + + +def _sampled_digest(path: Path, size: int) -> str: + """Hash a library whole when small, at fixed offsets when large. + + The samples are the two ends and evenly spaced chunks between them, plus + the size. A rebuild that changes code changes the size or the symbol and + section tables at the ends; an identical reinstall changes nothing here, + which is the point of hashing content rather than mtimes. + """ + sha = hashlib.sha256(str(size).encode()) + with path.open("rb") as handle: + if size <= _FULL_HASH_LIMIT: + sha.update(handle.read()) + return sha.hexdigest() + offsets = [0, max(0, size - _SAMPLE_EDGE)] + step = size // (_SAMPLE_COUNT + 1) + offsets[1:1] = [step * (i + 1) for i in range(_SAMPLE_COUNT)] + for index, offset in enumerate(offsets): + handle.seek(offset) + length = _SAMPLE_EDGE if index in (0, len(offsets) - 1) else _SAMPLE_CHUNK + sha.update(handle.read(length)) + return sha.hexdigest() + + +def native_build( + repo_root: str | Path, *, globs: Sequence[str] = DEFAULT_NATIVE_GLOBS +) -> NativeBuild: + """Fingerprint the native libraries under ``repo_root``.""" + root = Path(repo_root) + rows: list[str] = [] + for pattern in globs: + for path in sorted(root.glob(pattern)): + try: + stat = path.stat() + except OSError: + continue + if not path.is_file(): + continue + key = _stat_key(path.resolve(), stat) + digest = None + if key is not None: + with _native_lock: + digest = _native_cache.get(key) + if digest is None: + try: + digest = _sampled_digest(path, stat.st_size) + except OSError: + continue + if key is not None: + with _native_lock: + _native_cache[key] = digest + rows.append(f"{path.relative_to(root).as_posix()} {stat.st_size} {digest[:16]}") + rows = sorted(set(rows)) + if not rows: + return NativeBuild() + return NativeBuild(fingerprint=fingerprint(rows), libraries=tuple(rows)) + + +# ---------------------------------------------------------------- switches + +_ENV_FUNCS = {"get", "getenv", "setdefault"} + + +def _dotted(node: ast.AST) -> str: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + head = _dotted(node.value) + return f"{head}.{node.attr}" if head else node.attr + return "" + + +def _default_repr(node: ast.AST | None) -> str | None: + if node is None: + return None + if isinstance(node, ast.Constant): + return None if node.value is None else str(node.value) + return f"<expr {ast.unparse(node)}>" + + +def harvest_switch_defaults( + source: str, *, prefixes: Sequence[str] = DEFAULT_SWITCH_PREFIXES +) -> dict[str, str | None]: + """Every switch ``source`` reads, with the default it falls back to. + + Recognises ``os.environ.get(NAME, default)``, ``os.getenv(NAME, default)``, + ``os.environ.setdefault(NAME, default)`` and ``os.environ[NAME]``, for + names starting with one of ``prefixes``. A read with no default records + ``None``; a read that has none at all (``os.environ[NAME]``) records + ``"<required>"``. A default that is not a literal is recorded as its + source text, so a change to the expression still reads as a change. A + name read with two different defaults records both, joined with ``|``. + """ + try: + tree = ast.parse(source) + except (SyntaxError, ValueError): + return {} + found: dict[str, set[str | None]] = {} + + def add(name: Any, default: str | None) -> None: + if isinstance(name, str) and name.startswith(tuple(prefixes)): + found.setdefault(name, set()).add(default) + + for node in ast.walk(tree): + if isinstance(node, ast.Call) and isinstance(node.func, (ast.Attribute, ast.Name)): + dotted = _dotted(node.func) + last = dotted.rsplit(".", 1)[-1] + owner = dotted.rsplit(".", 1)[0] if "." in dotted else "" + is_env = last in _ENV_FUNCS and (owner.endswith("environ") or last == "getenv") + if not is_env or not node.args or not isinstance(node.args[0], ast.Constant): + continue + default_node = node.args[1] if len(node.args) > 1 else None + for keyword in node.keywords: + if keyword.arg in ("default", "value"): + default_node = keyword.value + add(node.args[0].value, _default_repr(default_node)) + elif isinstance(node, ast.Subscript) and _dotted(node.value).endswith("environ"): + key = node.slice + if isinstance(key, ast.Constant) and isinstance(node.ctx, ast.Load): + add(key.value, "<required>") + return { + name: ( + next(iter(values)) + if len(values) == 1 + else " | ".join(sorted("None" if v is None else v for v in values)) + ) + for name, values in sorted(found.items()) + } + + +_harvest_cache: dict[tuple[Any, ...], dict[str, str | None]] = {} +_harvest_lock = threading.Lock() + + +def harvest_paths( + repo_root: str | Path, + paths: Iterable[str], + *, + prefixes: Sequence[str] = DEFAULT_SWITCH_PREFIXES, +) -> dict[str, str | None]: + """Merge :func:`harvest_switch_defaults` over the Python files at ``paths``. + + ``paths`` may name directories, which are walked. + """ + root = Path(repo_root) + merged: dict[str, set[str | None]] = {} + files: set[str] = set() + for rel in paths: + files.update(p for p in _expand(root, rel) if p.endswith(".py")) + for rel in sorted(files): + path = root / rel + try: + stat = path.stat() + except OSError: + continue + key = _stat_key(path, stat, tuple(prefixes)) + harvested = None + if key is not None: + with _harvest_lock: + harvested = _harvest_cache.get(key) + if harvested is None: + try: + text = path.read_text(encoding="utf-8") + except (OSError, UnicodeDecodeError): + continue + harvested = harvest_switch_defaults(text, prefixes=prefixes) + if key is not None: + with _harvest_lock: + _harvest_cache[key] = harvested + for name, default in harvested.items(): + merged.setdefault(name, set()).add(default) + return { + name: ( + next(iter(values)) + if len(values) == 1 + else " | ".join(sorted("None" if v is None else v for v in values)) + ) + for name, values in sorted(merged.items()) + } + + +def command_env(command: str) -> dict[str, str]: + """``NAME=value`` assignments anywhere in a shell command line.""" + try: + tokens = shlex.split(command) + except ValueError: + tokens = command.split() + env: dict[str, str] = {} + for token in tokens: + match = re.fullmatch(r"([A-Za-z_][A-Za-z0-9_]*)=(.*)", token, flags=re.S) + if match: + env[match.group(1)] = match.group(2) + return env + + +def switch_env(env: Mapping[str, str], prefixes: Sequence[str]) -> dict[str, str]: + return {k: str(v) for k, v in sorted(env.items()) if str(k).startswith(tuple(prefixes))} + + +# -------------------------------------------------------------------- spec + + +def _project(value: Any, unit: str, units: frozenset[str]) -> Any: + """Keep only what applies to ``unit`` in a per-unit mapping, recursively. + + A mapping whose keys are all unit names (or ``default``) is the per-unit + form several ``task.yaml`` blocks accept. For one unit, only its own entry + and ``default`` matter, so another unit's settings never invalidate it. + """ + if isinstance(value, Mapping): + keys = {str(k) for k in value} + if keys and keys <= units | {"default"}: + value = {k: v for k, v in value.items() if str(k) in (unit, "default")} + projected = {str(k): _project(v, unit, units) for k, v in value.items()} + return {k: v for k, v in projected.items() if v not in ({}, None)} + if isinstance(value, (list, tuple)): + return [_project(v, unit, units) for v in value] + return value + + +def spec_projection(task_data: Mapping[str, Any], unit: str) -> dict[str, Any]: + """The part of ``task.yaml`` that can change what ``unit``'s gates measure.""" + units = assembly_units(task_data) + known = frozenset(units) + projected: dict[str, Any] = {} + for key, value in task_data.items(): + if key in SPEC_EXCLUDED_KEYS: + continue + if key == "features": + # A unit's gates run with the features assembled so far; adding a + # later feature does not change what an earlier unit measures. + if unit in units: + value = units[1 : units.index(unit) + 1] + elif key == "environment" and isinstance(value, Mapping): + value = {k: value[k] for k in SPEC_ENVIRONMENT_KEYS if k in value} + projected[str(key)] = _project(value, unit, known) + return projected + + +def spec_fingerprint(task_data: Mapping[str, Any] | None, unit: str) -> str: + if not task_data: + return "" + return fingerprint(spec_projection(task_data, unit)) + + +# ------------------------------------------------------------------ verdicts + + +@dataclass(frozen=True) +class LedgerPolicy: + """What a run treats as a dependency and as a switch.""" + + switch_prefixes: tuple[str, ...] = DEFAULT_SWITCH_PREFIXES + #: Files or directories, besides each verdict's own dependencies, whose + #: switch defaults are recorded with every verdict. + switch_sources: tuple[str, ...] = (MODELING_V2_PACKAGE,) + #: Files or directories every verdict depends on, on top of the target. + extra_dependencies: tuple[str, ...] = () + native_globs: tuple[str, ...] = DEFAULT_NATIVE_GLOBS + + +@dataclass(frozen=True) +class Verdict: + """One gate verdict, with everything needed to decide later whether it still holds.""" + + gate: str + unit: str + member: str + outcome: str + value: float | None = None + threshold: float | None = None + reference: float | None = None + tolerance: float | None = None + detail: str = "" + timestamp: str = "" + duration_s: float = 0.0 + target_rel: str = "" + tree: TreeIdentity = field(default_factory=TreeIdentity) + native: NativeBuild = field(default_factory=NativeBuild) + dependencies: Mapping[str, str | None] = field(default_factory=dict) + #: Explicit switch values the gate ran under (task env plus command). + env: Mapping[str, str] = field(default_factory=dict) + #: The switch defaults the code declared at measurement time. + switch_defaults: Mapping[str, str | None] = field(default_factory=dict) + config: Mapping[str, Any] = field(default_factory=dict) + config_fingerprint: str = "" + spec_fingerprint: str = "" + #: The gate's own one-token tree description, kept for humans. + tree_state: str = "" + dirty_override: bool = False + log: str = "" + id: str = "" + + @property + def key(self) -> tuple[str, str, str]: + return (self.gate, self.unit, self.member) + + @property + def passed(self) -> bool: + """Whether the verdict clears its threshold. A finding does.""" + return self.outcome in (OUTCOME_PASS, OUTCOME_FINDING) + + def label(self) -> str: + member = f"@{self.member}" if self.member else "" + return f"{self.gate}{member} [{self.unit or '-'}]" + + def value_text(self) -> str: + if self.value is None: + return self.detail[:120] + if self.threshold is None: + return f"{self.value:.4f}" + return f"{self.value:.4f} vs {self.threshold:.4f}" + + def to_record(self) -> dict[str, Any]: + record: dict[str, Any] = { + "kind": KIND_VERDICT, + "schema": SCHEMA_VERSION, + "gate": self.gate, + "unit": self.unit, + "member": self.member, + "outcome": self.outcome, + "value": self.value, + "threshold": self.threshold, + "reference": self.reference, + "tolerance": self.tolerance, + "detail": self.detail, + "timestamp": self.timestamp, + "duration_s": self.duration_s, + "target_rel": self.target_rel, + "tree": self.tree.to_json(), + "native_build": self.native.to_json(), + "dependencies": dict(sorted(self.dependencies.items())), + "env": dict(sorted(self.env.items())), + "switch_defaults": dict(sorted(self.switch_defaults.items())), + "config": dict(self.config), + "config_fingerprint": self.config_fingerprint, + "spec_fingerprint": self.spec_fingerprint, + "tree_state": self.tree_state, + "dirty_override": self.dirty_override, + "log": self.log, + } + record["id"] = self.id or fingerprint(record).split(":", 1)[1][:16] + return record + + @classmethod + def from_record(cls, raw: Mapping[str, Any]) -> "Verdict": + def number(key: str) -> float | None: + value = raw.get(key) + try: + return None if value is None else float(value) + except (TypeError, ValueError): + return None + + outcome = str(raw.get("outcome") or OUTCOME_FAIL).upper() + return cls( + gate=str(raw.get("gate") or ""), + unit=str(raw.get("unit") or ""), + member=str(raw.get("member") or ""), + outcome=outcome if outcome in OUTCOMES else OUTCOME_FAIL, + value=number("value"), + threshold=number("threshold"), + reference=number("reference"), + tolerance=number("tolerance"), + detail=str(raw.get("detail") or ""), + timestamp=str(raw.get("timestamp") or ""), + duration_s=float(raw.get("duration_s") or 0.0), + target_rel=str(raw.get("target_rel") or ""), + tree=TreeIdentity.from_json(raw.get("tree")), + native=NativeBuild.from_json(raw.get("native_build")), + dependencies={ + str(k): (None if v is None else str(v)) + for k, v in (raw.get("dependencies") or {}).items() + }, + env={str(k): str(v) for k, v in (raw.get("env") or {}).items()}, + switch_defaults={ + str(k): (None if v is None else str(v)) + for k, v in (raw.get("switch_defaults") or {}).items() + }, + config=dict(raw.get("config") or {}), + config_fingerprint=str(raw.get("config_fingerprint") or ""), + spec_fingerprint=str(raw.get("spec_fingerprint") or ""), + tree_state=str(raw.get("tree_state") or ""), + dirty_override=bool(raw.get("dirty_override", False)), + log=str(raw.get("log") or ""), + id=str(raw.get("id") or ""), + ) + + +# --------------------------------------------------------------------- store + + +class Ledger: + """The append-only JSON Lines store. + + An append is one ``write`` of one line to a file opened ``O_APPEND``, + under an exclusive ``flock`` and a thread lock, so concurrent recorders + in one process or in several never interleave. A torn final line (a + process killed mid-write) is skipped on read and counted in + :attr:`skipped`. + """ + + def __init__(self, path: str | Path) -> None: + self.path = Path(path) + self.skipped = 0 + self._lock = threading.Lock() + + @classmethod + def in_workspace(cls, workspace: str | Path) -> "Ledger": + return cls(Path(workspace) / LEDGER_FILENAME) + + def append(self, record: Mapping[str, Any]) -> None: + data = (json.dumps(record, sort_keys=True, ensure_ascii=False) + "\n").encode("utf-8") + with self._lock: + self.path.parent.mkdir(parents=True, exist_ok=True) + fd = os.open(self.path, os.O_WRONLY | os.O_APPEND | os.O_CREAT, 0o644) + try: + fcntl.flock(fd, fcntl.LOCK_EX) + view = memoryview(data) + while view: + written = os.write(fd, view) + view = view[written:] + os.fsync(fd) + finally: + fcntl.flock(fd, fcntl.LOCK_UN) + os.close(fd) + + def records(self) -> list[dict[str, Any]]: + self.skipped = 0 + try: + text = self.path.read_text(encoding="utf-8") + except FileNotFoundError: + return [] + out: list[dict[str, Any]] = [] + for line in text.splitlines(): + if not line.strip(): + continue + try: + record = json.loads(line) + except json.JSONDecodeError: + self.skipped += 1 + continue + if isinstance(record, dict): + out.append(record) + else: + self.skipped += 1 + return out + + def verdicts(self) -> list[Verdict]: + return [Verdict.from_record(r) for r in self.records() if r.get("kind") == KIND_VERDICT] + + def latest(self) -> dict[tuple[str, str, str], Verdict]: + """The newest verdict per ``(gate, unit, member)``, in first-seen key order.""" + latest: dict[tuple[str, str, str], Verdict] = {} + for verdict in self.verdicts(): + latest[verdict.key] = verdict + return latest + + def digest(self) -> str: + """A digest of the ledger's bytes, for reports to say what they were made from.""" + try: + return "sha256:" + hashlib.sha256(self.path.read_bytes()).hexdigest()[:24] + except FileNotFoundError: + return "" + + def record_verdict(self, verdict: Verdict) -> Verdict: + record = verdict.to_record() + self.append(record) + return replace(verdict, id=record["id"]) + + +# ----------------------------------------------------------------- recording + + +@dataclass(frozen=True) +class VerdictContext: + """What the recorder knows that a gate result does not.""" + + repo_root: str = "" + target_rel: str = "" + unit: str = "" + #: The explicit environment every gate command runs under (``task.yaml``'s + #: ``environment.env``). Assignments inside the command itself are added. + env: Mapping[str, str] = field(default_factory=dict) + features: Sequence[str] = () + task_data: Mapping[str, Any] | None = None + policy: LedgerPolicy = field(default_factory=LedgerPolicy) + log: str = "" + #: The backend and model each role is routed to (``role -> {backend, + #: model, ...}``): who wrote the code a verdict measured. Part of the + #: config fingerprint, not of freshness. + agents: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + + +def _number(value: Any) -> float | None: + try: + return None if value is None else float(value) + except (TypeError, ValueError): + return None + + +def outcome_of(result: Any) -> str: + """PASS, FAIL or FINDING for a gate result of any vintage. + + A result may carry an explicit ``outcome``; one that says ``finding`` is a + finding; otherwise ``passed`` decides. + """ + explicit = str(getattr(result, "outcome", "") or "").upper() + if explicit in OUTCOMES: + return explicit + if getattr(result, "finding", False): + return OUTCOME_FINDING + return OUTCOME_PASS if getattr(result, "passed", False) else OUTCOME_FAIL + + +@dataclass(frozen=True) +class Measurement: + """What a verdict is measured against, read once from the tree. + + Separate from the result so a batch of verdicts taken at one moment (the + checks of one delivery report) reads the tree once. + """ + + tree: TreeIdentity + native: NativeBuild + dependencies: Mapping[str, str | None] + switch_defaults: Mapping[str, str | None] + + +def measure(context: VerdictContext) -> Measurement: + """Fingerprint the tree, the native build, the dependencies and the switch defaults now.""" + policy = context.policy + root = context.repo_root + if not root: + return Measurement(TreeIdentity(), NativeBuild(), {}, {}) + paths = dependency_paths(root, context.target_rel, extra=policy.extra_dependencies) + dependencies = hash_paths(root, paths) + return Measurement( + tree=tree_identity(root, dependencies=dependencies), + native=native_build(root, globs=policy.native_globs), + dependencies=dependencies, + switch_defaults=harvest_paths( + root, [*paths, *policy.switch_sources], prefixes=policy.switch_prefixes + ), + ) + + +def verdict_from_result( + result: Any, + context: VerdictContext, + *, + gate: str | None = None, + measurement: Measurement | None = None, +) -> Verdict: + """Build a :class:`Verdict` from a gate result, fingerprinting the tree now. + + ``result`` is duck-typed against ``gates.GateResult`` so a gate added + later records without this module changing: ``name``, ``passed`` and + ``detail`` are required; ``member``, ``measured``, ``reference``, + ``tolerance``, ``threshold``, ``command``, ``duration_s``, + ``tree_state``, ``dirty_override``, ``outcome`` and ``finding`` are read + when present. ``gate`` overrides the result's name as the ledger key, + for a gate that runs at several points under one name. + """ + policy = context.policy + measured = measurement or measure(context) + command = str(getattr(result, "command", "") or "") + explicit = {str(k): str(v) for k, v in dict(context.env or {}).items()} + explicit.update(command_env(command)) + switches = switch_env(explicit, policy.switch_prefixes) + reference = _number(getattr(result, "reference", None)) + tolerance = _number(getattr(result, "tolerance", None)) + threshold = _number(getattr(result, "threshold", None)) + if threshold is None and reference is not None: + threshold = reference - (tolerance or 0.0) + gate = gate or str(getattr(result, "name", "") or "gate") + member = str(getattr(result, "member", "") or "") + config = { + "gate": gate, + "unit": context.unit, + "member": member, + "features": [str(f) for f in context.features], + "env": switches, + } + if context.agents: + config["agents"] = {str(role): dict(agent) for role, agent in context.agents.items()} + return Verdict( + gate=gate, + unit=context.unit, + member=member, + outcome=outcome_of(result), + value=_number(getattr(result, "measured", None)), + threshold=threshold, + reference=reference, + tolerance=tolerance, + detail=str(getattr(result, "detail", "") or "")[:2000], + timestamp=utc_now(), + duration_s=float(getattr(result, "duration_s", 0.0) or 0.0), + target_rel=context.target_rel, + tree=measured.tree, + native=measured.native, + dependencies=measured.dependencies, + env=switches, + switch_defaults=measured.switch_defaults, + config=config, + config_fingerprint=fingerprint({**config, "command": command, "env": explicit}), + spec_fingerprint=spec_fingerprint(context.task_data, context.unit), + tree_state=str(getattr(result, "tree_state", "") or ""), + dirty_override=bool(getattr(result, "dirty_override", False)), + log=context.log, + ) + + +def record_gate_result( + ledger: Ledger, + result: Any, + context: VerdictContext, + *, + gate: str | None = None, + measurement: Measurement | None = None, +) -> Verdict: + """Record one gate result. **The** entry point for gates and harnesses. + + Stable on purpose: ``gates.GateLog`` calls a sink with each result, and + the workflow's sink is this function with the run's context bound. + """ + return ledger.record_verdict( + verdict_from_result(result, context, gate=gate, measurement=measurement) + ) + + +def select(verdicts: Iterable[Verdict], *, gate="*", unit="*", member="*") -> list[Verdict]: + """Verdicts whose gate, unit and member match the ``fnmatch`` patterns.""" + return [ + v + for v in verdicts + if fnmatch.fnmatchcase(v.gate, gate) + and fnmatch.fnmatchcase(v.unit, unit) + and fnmatch.fnmatchcase(v.member, member) + ] + + +# ----------------------------------------------------------------------- CLI + + +@dataclass(frozen=True) +class _RecordedResult: + """A gate result handed in on the command line, shaped like ``gates.GateResult``.""" + + name: str + passed: bool + detail: str + member: str = "" + measured: float | None = None + threshold: float | None = None + command: str = "" + outcome: str = "" + duration_s: float = 0.0 + + +def load_task(path: str | Path | None) -> dict[str, Any] | None: + """Read a ``task.yaml`` for its spec fingerprint, or ``None``.""" + if not path: + return None + import yaml + + data = yaml.safe_load(Path(path).read_text(encoding="utf-8")) + return data if isinstance(data, dict) else None + + +def main(argv: Sequence[str] | None = None) -> int: + """``record``: put a verdict measured outside the workflow into a workspace's ledger. + + For a harness that runs a gate by hand -- a regression sweep, an A/B -- + so its verdict carries the same fingerprints as the workflow's own and is + judged by the same freshness rules. + """ + parser = argparse.ArgumentParser(prog="staircase-ledger", description=main.__doc__) + sub = parser.add_subparsers(dest="command", required=True) + record = sub.add_parser("record", help="record a verdict measured outside the workflow") + record.add_argument("workspace") + record.add_argument("--repo", required=True) + record.add_argument("--target-rel", default="") + record.add_argument("--gate", required=True) + record.add_argument("--unit", required=True) + record.add_argument("--member", default="") + record.add_argument("--outcome", choices=OUTCOMES, required=True) + record.add_argument("--value", type=float) + record.add_argument("--threshold", type=float) + record.add_argument("--detail", default="") + record.add_argument("--command-line", default="", help="the command the gate ran") + record.add_argument("--env", action="append", default=[], help="NAME=value, repeatable") + record.add_argument("--task", help="task.yaml, for the spec fingerprint") + record.add_argument("--log", default="", help="where the gate's log is") + args = parser.parse_args(argv) + + env = dict(item.split("=", 1) for item in args.env if "=" in item) + verdict = record_gate_result( + Ledger.in_workspace(args.workspace), + _RecordedResult( + name=args.gate, + passed=args.outcome != OUTCOME_FAIL, + detail=args.detail, + member=args.member, + measured=args.value, + threshold=args.threshold, + command=args.command_line, + outcome=args.outcome, + ), + VerdictContext( + repo_root=args.repo, + target_rel=args.target_rel, + unit=args.unit, + env=env, + task_data=load_task(args.task), + log=args.log, + ), + ) + sys.stdout.write(f"recorded {verdict.id}: {verdict.label()} {verdict.outcome}\n") + return 0 + + +__all__ = [ + "DEFAULT_NATIVE_GLOBS", + "DEFAULT_SWITCH_PREFIXES", + "KIND_REGRESSION", + "KIND_VERDICT", + "LEDGER_FILENAME", + "Ledger", + "LedgerPolicy", + "Measurement", + "NativeBuild", + "OUTCOMES", + "OUTCOME_FAIL", + "OUTCOME_FINDING", + "OUTCOME_PASS", + "TreeIdentity", + "Verdict", + "VerdictContext", + "blob_hash", + "command_env", + "dependency_paths", + "fingerprint", + "harvest_paths", + "harvest_switch_defaults", + "hash_paths", + "is_git_repo", + "load_task", + "measure", + "native_build", + "outcome_of", + "record_gate_result", + "select", + "spec_fingerprint", + "spec_projection", + "run_git", + "tree_identity", + "utc_now", + "verdict_from_result", +] + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/progress.py b/agent-flow/agent_flow/workflows/staircase_onboard/progress.py index 4fef63083d86..a067bc53a51d 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/progress.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/progress.py @@ -33,6 +33,8 @@ from __future__ import annotations +import os +import re import threading from dataclasses import dataclass, field from datetime import datetime @@ -205,10 +207,14 @@ class ProgressContext: ``current_iteration`` is set before each dispatch so an entry is stamped without the agent having to pass (or guess) it. ``lock`` serializes the read-modify-write, which a concurrent smith batch would otherwise race. + + ``current_unit`` is the assembly unit in progress, or empty outside + assembly. Gate verdicts are attributed to it in the evidence ledger. """ path: Path current_iteration: int = 0 + current_unit: str = "" lock: threading.Lock = field(default_factory=threading.Lock, repr=False, compare=False) def append(self, agent: str, entry: dict[str, Any]) -> None: @@ -424,64 +430,145 @@ def parse_uncovered_records(raw: Any) -> list[dict[str, Any]]: return rows -def decision_path(workspace: Path, agent: str, iteration: int) -> Path: - """Where a role writes its decision when MCP tools are unavailable.""" - return Path(workspace) / DECISIONS_DIRNAME / f"{agent}-{iteration}.yaml" +#: The unit a plan-stage or deliver-stage decision is filed under. Those +#: stages are not per unit, and naming them keeps their files apart from every +#: assembly unit's. +PLAN_DECISION_UNIT = "plan" +DELIVER_DECISION_UNIT = "deliver" +#: The unit a build-stage decision is filed under when none is in progress. +DEFAULT_DECISION_UNIT = "core" + + +def decision_unit(agent: str, unit: str = "") -> str: + """The unit ``agent``'s decision file is named for. + Plan and deliver roles file under their stage; build roles (smith, + assembler, reviewer) under the unit they work for. + """ + stage = _STAGE_BY_AGENT.get(agent) + if stage == PLAN_STAGE: + return PLAN_DECISION_UNIT + if stage == DELIVER_STAGE: + return DELIVER_DECISION_UNIT + return re.sub(r"[^A-Za-z0-9_.]+", "_", unit or DEFAULT_DECISION_UNIT) + + +def decision_path( + workspace: Path, agent: str, iteration: int, *, unit: str = "", attempt: int = 1 +) -> Path: + """Where a role writes its decision when MCP tools are unavailable. + + ``decisions/<unit>-<role>-<turn>.yaml``. The unit is part of the name + because turn numbers restart with every unit: named by role and turn + alone, the multimodal assembler's turn 1 overwrote core's, and reviewers + overwrote earlier units' decisions three times in one day (F033). + + A turn number can still come round again for the same unit and role -- a + re-entered unit restarts its turn count, a smith row is retried under its + own number, a ``--clean`` run starts over in the same workspace. So a + second dispatch of the same ``(unit, role, turn)`` gets ``attempt`` 2 and + the turn token ``<turn>.2``, and so on: :func:`reserve_decision` never + hands out a name that was used before, so no decision file is ever + overwritten. + """ + token = f"{iteration}.{attempt}" if attempt > 1 else str(iteration) + name = f"{decision_unit(agent, unit)}-{agent}-{token}" + return Path(workspace) / DECISIONS_DIRNAME / f"{name}.yaml" -def _decision_marker_path(workspace: Path, agent: str, iteration: int) -> Path: + +def _decision_marker_path(decision: Path) -> Path: """The dotfile stamped beside a decision when its turn is dispatched. A dotfile, and beside the decision rather than under a directory of its own, because the two must share a filesystem — see - ``mark_decision_dispatch`` for why that is the whole mechanism. + ``reserve_decision`` for why that is the whole mechanism. + """ + return decision.with_name(f".{decision.stem}.dispatched") + + +def reserve_decision(workspace: Path, agent: str, iteration: int, *, unit: str = "") -> Path: + """Claim a decision path no earlier dispatch used, and stamp its dispatch marker. + + The first attempt whose decision file and marker both do not exist is + claimed by creating the marker with ``O_CREAT | O_EXCL``, so two + concurrent dispatches (or two processes) can never be handed the same + path. Returns the path the turn must write. + + The marker also orders the turn against its answer. It is compared by + **mtime against the decision file on the same filesystem**, never against + a clock: an agent may answer from a different node whose clock is skewed + by minutes, and both files live under the workspace. + """ + attempt = 1 + while True: + path = decision_path(workspace, agent, iteration, unit=unit, attempt=attempt) + marker = _decision_marker_path(path) + if not path.exists(): + marker.parent.mkdir(parents=True, exist_ok=True) + try: + os.close(os.open(marker, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o644)) + except FileExistsError: + pass + else: + return path + attempt += 1 + + +def mark_decision_dispatch(workspace: Path, agent: str, iteration: int, *, unit: str = "") -> Path: + """Reserve this turn's decision path before it is dispatched. See :func:`reserve_decision`. + + Before reservation existed, ``(agent, iteration)`` was the whole name and + a turn that ended *without* writing a decision left the **previous** + turn's verdict at its path, looking fresh. That happened five times in one + onboarding, the worst being three audit verdicts read back from two + passes earlier, which made the audit loop unable to converge. A reserved + path is new, so it holds nothing until this turn writes it. """ - return Path(workspace) / DECISIONS_DIRNAME / f".{agent}-{iteration}.dispatched" - - -def mark_decision_dispatch(workspace: Path, agent: str, iteration: int) -> None: - """Stamp the moment this turn was dispatched, so a stale answer is detectable. - - ``(agent, iteration)`` is **not unique across a run**: ``iteration`` is 1 - for every turn of the deliver stage, and it restarts per unit in assembly. - So two different turns routinely write the same path — see - ``decision_path`` — and a turn that ends *without* writing one leaves the - **previous** turn's verdict sitting there looking fresh, which - ``collect_decision_file`` then reports as this turn's answer. - - That is not hypothetical. It happened five times in one onboarding: ten - smith decisions collected from the wrong place, an assembler told to write - the previous turn's path, decision paths colliding across units, and — - worst — three audit verdicts read back as a verdict from two passes - earlier, which made the audit loop structurally unable to converge, - because a ``CLEAN`` could not end the run. Every one of those was silent. - - This run saw its cousin: a smith committed 380 lines with a passing - receipt and then died without writing its decision file, blocking a batch - for 58 minutes. Staleness detection is what turns that class of failure - from silent into visible. - - The marker is compared by **mtime against the decision file on the same - filesystem**, never against a clock. An agent may answer from a different - node whose clock is skewed by minutes; both files live under the - workspace, so their mtimes are ordered by the one filesystem that wrote - them. + return reserve_decision(workspace, agent, iteration, unit=unit) + + +def decision_file_instruction(path: Path) -> str: + """The block a dispatch prompt ends with: the one exact path this turn writes. + + Stated in the dispatch rather than left to the system prompt's pattern, so + the path is exact and un-guessable: the role is told the absolute file to + write, not a pattern to instantiate. """ - marker = _decision_marker_path(workspace, agent, iteration) - marker.parent.mkdir(parents=True, exist_ok=True) - marker.write_text("", encoding="utf-8") + return ( + "---\n\n" + "**Write your decision file to exactly this path when you finish:**\n\n" + f"`{path}`\n\n" + "Nothing downstream reads your prose. This file is how this turn " + "reaches the orchestrator.\n" + ) -def decision_is_stale(workspace: Path, agent: str, iteration: int) -> bool: - """Whether the decision file at this path predates the turn that wants it. +def current_decision(workspace: Path, agent: str, iteration: int, *, unit: str = "") -> Path: + """The path the latest dispatch of ``(unit, agent, iteration)`` was told to write. + + The highest reserved attempt. With nothing reserved -- a run that never + stamped one -- it is the first attempt's path. + """ + latest = decision_path(workspace, agent, iteration, unit=unit) + attempt = 1 + while True: + attempt += 1 + candidate = decision_path(workspace, agent, iteration, unit=unit, attempt=attempt) + if not (_decision_marker_path(candidate).exists() or candidate.exists()): + return latest + latest = candidate + + +def decision_is_stale(workspace: Path, agent: str, iteration: int, *, unit: str = "") -> bool: + """Whether the current decision file predates the dispatch that reserved it. False when either file is absent: a missing decision is the caller's "no decision" case, and a missing marker means this run never stamped one — in which case the old unconditional behaviour is the honest fallback, rather than silently refusing every decision a run ever writes. """ - decision = decision_path(workspace, agent, iteration) - marker = _decision_marker_path(workspace, agent, iteration) + decision = current_decision(workspace, agent, iteration, unit=unit) + marker = _decision_marker_path(decision) if not decision.is_file() or not marker.is_file(): return False return decision.stat().st_mtime < marker.stat().st_mtime @@ -492,6 +579,8 @@ def collect_decision_file( workspace: Path, agent: str, iteration: int, + *, + unit: str = "", ) -> str | None: """Read, validate and record the decision file ``agent`` wrote. @@ -513,14 +602,13 @@ def collect_decision_file( silence to interpret. """ allowed = DECISIONS_BY_AGENT[agent] - path = decision_path(workspace, agent, iteration) + path = current_decision(workspace, agent, iteration, unit=unit) if not path.is_file(): return None # A file older than this turn's dispatch marker belongs to a PREVIOUS turn - # that happened to share this path. Reading it would report that turn's - # verdict as this one's — the failure `mark_decision_dispatch` exists to - # make visible, and one that is silent every other way. - if decision_is_stale(workspace, agent, iteration): + # that shared this path. Reservation makes that impossible for paths it + # hands out; the check stays for files written before it existed. + if decision_is_stale(workspace, agent, iteration, unit=unit): return None try: @@ -540,6 +628,7 @@ def collect_decision_file( "timestamp": _now_iso(), "summary": str(raw.get("summary", "")).strip(), "decision": decision, + "decision_file": path.name, } for key in EXTRA_KEYS_BY_AGENT.get(agent, ()): value = raw.get(key) @@ -826,13 +915,22 @@ def build_progress_tools(ctx: ProgressContext) -> dict[str, list[Any]]: "SMITH_DECISIONS", "UNCOVERED_OPTIONAL", "UNCOVERED_REQUIRED", + "DEFAULT_DECISION_UNIT", + "DELIVER_DECISION_UNIT", + "PLAN_DECISION_UNIT", "build_progress_tools", + "collect_decision_file", + "current_decision", + "decision_file_instruction", "decision_is_stale", + "decision_path", + "decision_unit", "find_entries", "init_progress_file", "latest_decision", "latest_entry", "mark_decision_dispatch", "read_progress", + "reserve_decision", "write_progress", ] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/_common.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/_common.py index c071a09f6da8..c9652e953a97 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/_common.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/_common.py @@ -248,10 +248,12 @@ def decision_file_block(agent: str, decisions: Sequence[str], extra: Sequence[st "`Write` tool to:\n" "\n" "```\n" - f"<workspace>/decisions/{agent}-<iteration>.yaml\n" + f"<workspace>/decisions/<unit>-{agent}-<turn>.yaml\n" "```\n" "\n" - "where `<workspace>` and `<iteration>` are given in your dispatch prompt. " + "The exact path is given at the end of your dispatch prompt. Write that " + "file and no other: every turn gets a path no earlier turn used, and a " + "decision written anywhere else is never read. " "The file is YAML with exactly these keys:\n" "\n" "```yaml\n" diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/smith.py b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/smith.py index 03fac5a2a883..5ad3105d3861 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/prompts/smith.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/prompts/smith.py @@ -39,8 +39,8 @@ to answer specific questions; you read it to build the thing they will later read *instead of* the source. -Also yours: the catalog's contracts and `index.yaml`, the generated op -inventory, and the values your dispatch prompt hands you. +Also yours: the catalog's contracts and your entry's index fragment, the +generated op inventory, and the values your dispatch prompt hands you. Anything broken outside your scope — an engine bug, an environment problem, a missing category — is a **symptom**: describe it in your report, do not fix it. @@ -48,7 +48,10 @@ ## What you may write - `tensorrt_llm/_torch/_experimental/modeling_v2/catalog/<category>/<name>.md` - and `.py`, plus `catalog/index.yaml` + and `.py`, plus the entry's index fragment + `catalog/index.d/<category>/<name>.yaml`. Never `catalog/index.yaml` itself: + it is generated from the fragments, and a smith that commits it sweeps in + its siblings' rows - that entry's GPU test at `tests/unittest/_torch/modeling_v2/<category>/test_modeling_v2_<name>.py` - on a macro verdict only, the decomposition report your skill specifies diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/report.py b/agent-flow/agent_flow/workflows/staircase_onboard/report.py new file mode 100644 index 000000000000..cc152fcbf759 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/report.py @@ -0,0 +1,1337 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reports generated from the ledger, and the check that they agree with it. + +The DeepSeek-V4.1 run kept ``progress.json``, the gap-closure report and the +final report by hand. They drifted apart in ten recorded ways. One summary +said "16/16 PASS, no finding" over a regression table that held a FINDING. +The progress page still listed an A/B as pending three hours after an +addendum closed it, and named a branch and a target commit that were no +longer current. Every one of those was a sentence written about evidence +instead of from it. + +So the facts in a report are rendered, never typed: + +- ``progress.json`` and ``progress.html`` (the latter offline and + self-contained, ported from that run's ``scripts/render_progress.py``) are + generated from the ledger, the checkpoint and the journal; +- the final report's gate tables are generated into marked regions of + ``final-report.md``, and everything outside those regions is the + operator's narrative. + +Narrative still exists, because judgement does: feature notes, workstream +items, prose conclusions. It lives in ``progress.overlay.yaml`` and outside +the generated regions, and :func:`check_consistency` reads it against the +ledger. The check fails on "no finding" while a FINDING verdict exists, an +"N/N PASS" or "all gates pass" while a verdict fails or is stale, a pending +item whose evidence is in fact fresh and passing, a closed item whose +evidence is not, a "final tree" that is not the tree of the last +regression, and a ``progress.json`` that was generated from an older ledger +or was edited after it was generated. +""" + +from __future__ import annotations + +import argparse +import html +import json +import os +import re +import sys +import threading +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Iterable, Mapping, Sequence + +import yaml + +from .freshness import ( + Assessment, + Binding, + Regression, + RerunEntry, + assess, + last_regression, + regression_binding, + rerun_queue, +) +from .ledger import ( + OUTCOME_FAIL, + OUTCOME_FINDING, + OUTCOME_PASS, + Ledger, + LedgerPolicy, + Verdict, + load_task, + select, + utc_now, +) +from .state import ( + STAGE_ASSEMBLE, + STAGE_DELIVER, + STAGE_PLAN, + STAGE_PREFLIGHT, + STAGE_SMITH, + STATE_FILENAME, + load_state, +) +from .task_schema import assembly_units + +PROGRESS_JSON = "progress.json" +PROGRESS_HTML = "progress.html" +OVERLAY_FILENAME = "progress.overlay.yaml" +FINAL_REPORT = "final-report.md" +SCHEMA_VERSION = 1 + +STATUS_DELIVERED = "DELIVERED" +STATUS_STALE = "STALE" +STATUS_BLOCKED = "BLOCKED" +STATUS_IN_PROGRESS = "IN PROGRESS" +STATUS_NOT_STARTED = "NOT STARTED" +STATUS_DONE = "DONE" +STATUS_FRESH = "FRESH" + +#: Item statuses that claim the work is finished, and those that claim it is not. +CLOSED_STATUSES = frozenset( + {"DONE", "CLOSED", "PASS", "GATED-PASS", "DELIVERED", "CONFIRMED", "RESOLVED"} +) +OPEN_STATUSES = frozenset( + {"PENDING", "IN PROGRESS", "OPEN", "NOT STARTED", "RUNNING", "QUEUED", "PARTIAL", "BLOCKED"} +) + +_STAGES = (STAGE_PREFLIGHT, STAGE_PLAN, STAGE_SMITH, STAGE_ASSEMBLE, STAGE_DELIVER) +_GENERATED = re.compile( + r"<!-- BEGIN GENERATED: (?P<name>[a-z-]+) -->\n.*?<!-- END GENERATED: (?P=name) -->", + re.S, +) +_JOURNAL_EVENT = re.compile(r"^- `(?P<time>[^`]+)` \*\*(?P<what>.+?)\*\*(?: — (?P<detail>.*))?$") +_TIMELINE_LIMIT = 200 + + +# ------------------------------------------------------------------- facts + + +@dataclass(frozen=True) +class Facts: + """Everything a report says, read from the ledger, the checkpoint and the journal.""" + + latest: tuple[Verdict, ...] = () + assessments: Mapping[tuple[str, str, str], Assessment] = field(default_factory=dict) + queue: tuple[RerunEntry, ...] = () + runs: Mapping[tuple[str, str, str], int] = field(default_factory=dict) + regression: Regression | None = None + binding: Binding | None = None + ledger_digest: str = "" + verdict_count: int = 0 + state: Mapping[str, Any] = field(default_factory=dict) + units: tuple[str, ...] = () + repo: str = "" + events: tuple[Mapping[str, str], ...] = () + + @property + def findings(self) -> list[Verdict]: + return [v for v in self.latest if v.outcome == OUTCOME_FINDING] + + @property + def failing(self) -> list[Verdict]: + return [v for v in self.latest if v.outcome == OUTCOME_FAIL] + + @property + def stale(self) -> list[Verdict]: + return [v for v in self.latest if self.is_stale(v)] + + def is_stale(self, verdict: Verdict) -> bool: + item = self.assessments.get(verdict.key) + return bool(item and item.stale) + + def holds(self, verdict: Verdict) -> bool: + """Passed, and still describes the tree.""" + return verdict.passed and not self.is_stale(verdict) + + def counts(self) -> dict[str, int]: + return { + "total": len(self.latest), + "pass": sum(v.outcome == OUTCOME_PASS for v in self.latest), + "finding": len(self.findings), + "fail": len(self.failing), + "stale": len(self.stale), + "queued": len(self.queue), + } + + def status(self) -> str: + """The run's one-word status, derived rather than declared.""" + state = self.state + if not state: + return STATUS_NOT_STARTED + if state.get("done"): + if state.get("verdict") != "delivered": + return STATUS_BLOCKED + if self.stale or self.failing: + return STATUS_STALE + if self.binding is not None and self.binding.claim and not self.binding.ok: + return STATUS_STALE + return STATUS_DELIVERED + return STATUS_IN_PROGRESS + + +def _read_journal_events(path: Path) -> list[dict[str, str]]: + try: + lines = path.read_text(encoding="utf-8").splitlines() + except (FileNotFoundError, UnicodeDecodeError): + return [] + events = [] + for line in lines: + match = _JOURNAL_EVENT.match(line) + if match: + events.append( + { + "time": match.group("time"), + "event": match.group("what") + + (f" — {match.group('detail')}" if match.group("detail") else ""), + } + ) + return events[-_TIMELINE_LIMIT:] + + +def collect_facts( + workspace: str | Path, + *, + repo_root: str | Path | None = None, + task_data: Mapping[str, Any] | None = None, + scope: Callable[[Verdict], bool] | None = None, + policy: LedgerPolicy | None = None, + check_history: bool = False, +) -> Facts: + """Read the workspace's ledger, checkpoint and journal into :class:`Facts`. + + ``repo_root`` and ``task_data`` default to the workspace's own + ``task.yaml``. ``check_history`` adds the git-history switch check, which + costs a ``git diff`` per verdict commit; delivery wants it, a page + refresh after every checkpoint does not. + """ + workspace = Path(workspace) + if task_data is None: + try: + task_data = load_task(workspace / "task.yaml") + except (OSError, yaml.YAMLError): + task_data = None + task_data = task_data or {} + repo = str(repo_root or task_data.get("trtllm_repo_path") or "") + policy = policy or LedgerPolicy() + + ledger = Ledger.in_workspace(workspace) + verdicts = ledger.verdicts() + runs: dict[tuple[str, str, str], int] = {} + for verdict in verdicts: + runs[verdict.key] = runs.get(verdict.key, 0) + 1 + latest = [v for v in ledger.latest().values() if scope is None or scope(v)] + assessed: list[Assessment] = [] + binding = None + if repo: + assessed = assess( + latest, repo, task_data=task_data or None, policy=policy, check_history=check_history + ) + binding = regression_binding(ledger, repo, policy=policy) + state: dict[str, Any] = {} + state_path = workspace / STATE_FILENAME + if state_path.is_file(): + try: + state = dict(load_state(state_path).__dict__) + except (OSError, ValueError, KeyError): + state = {} + return Facts( + latest=tuple(latest), + assessments={a.verdict.key: a for a in assessed}, + queue=tuple(rerun_queue(assessed)) if assessed else (), + runs=runs, + regression=last_regression(ledger), + binding=binding, + ledger_digest=ledger.digest(), + verdict_count=len(verdicts), + state=state, + units=tuple(assembly_units(task_data)) if task_data else (), + repo=repo, + events=tuple(_read_journal_events(workspace / "journal.md")), + ) + + +# -------------------------------------------------------------- progress.json + + +def headline(facts: Facts) -> str: + """The one-sentence summary, rendered from counts so it cannot overstate them.""" + c = facts.counts() + if not c["total"]: + return "No verdict recorded yet." + parts = [ + f"{c['total']} verdict(s): {c['pass']} PASS, {c['finding']} FINDING, {c['fail']} FAIL", + f"{c['stale']} stale" if c["stale"] else "none stale", + ] + if facts.regression is not None: + bound = "" + if facts.binding is not None and facts.binding.claim: + bound = ", bound to HEAD" if facts.binding.ok else ", NOT the current HEAD" + parts.append( + f"last full regression {len(facts.regression.keys)} verdict(s) on " + f"{facts.regression.tree.short()}{bound}" + ) + else: + parts.append("no full regression recorded") + return "; ".join(parts) + "." + + +def _stage_rows(state: Mapping[str, Any]) -> list[dict[str, str]]: + current = str(state.get("stage") or STAGE_PREFLIGHT) + done = bool(state.get("done")) + verdict = str(state.get("verdict") or "") + position = _STAGES.index(current) if current in _STAGES else 0 + rows = [] + for index, stage in enumerate(_STAGES): + if index < position: + status = STATUS_DONE + elif index > position: + status = STATUS_NOT_STARTED + elif done and verdict == "delivered": + status = STATUS_DELIVERED + elif done: + status = STATUS_BLOCKED + else: + status = STATUS_IN_PROGRESS + rows.append({"stage": stage, "status": status}) + return rows + + +def _unit_rows(facts: Facts) -> list[dict[str, Any]]: + state = facts.state + done = set(state.get("units_done") or ()) + in_assembly = state.get("stage") == STAGE_ASSEMBLE and not state.get("done") + index = int(state.get("unit_index") or 0) + rows = [] + for position, unit in enumerate(facts.units): + verdicts = [v for v in facts.latest if v.unit == unit] + holding = sum(facts.holds(v) for v in verdicts) + if unit in done: + status = STATUS_DONE if holding == len(verdicts) else STATUS_STALE + elif in_assembly and position == index: + status = STATUS_IN_PROGRESS + else: + status = STATUS_NOT_STARTED + rows.append( + { + "unit": unit, + "status": status, + "gates": f"{holding}/{len(verdicts)} hold" if verdicts else "no verdict", + } + ) + return rows + + +def gate_row(facts: Facts, verdict: Verdict) -> dict[str, Any]: + item = facts.assessments.get(verdict.key) + return { + "unit": verdict.unit, + "member": verdict.member, + "gate": verdict.gate, + "result": verdict.outcome, + "fresh": STATUS_STALE if item and item.stale else STATUS_FRESH, + "value_vs_anchor": verdict.value_text(), + "tree": verdict.tree.short(), + "native": verdict.native.short(), + "timestamp": verdict.timestamp, + "runs": facts.runs.get(verdict.key, 1), + "reasons": list(item.reasons) if item else [], + "id": verdict.id, + } + + +def _item_evidence(facts: Facts, item: Mapping[str, Any]) -> tuple[list[Verdict], list[str]]: + """The verdicts an overlay item cites, and what keeps them from closing it.""" + selectors = item.get("evidence") or [] + if isinstance(selectors, Mapping): + selectors = [selectors] + matched: list[Verdict] = [] + for selector in selectors: + if not isinstance(selector, Mapping): + continue + for verdict in select( + facts.latest, + gate=str(selector.get("gate", "*")), + unit=str(selector.get("unit", "*")), + member=str(selector.get("member", "*")), + ): + if verdict not in matched: + matched.append(verdict) + gaps = [ + f"{v.label()} is {'stale' if facts.is_stale(v) else v.outcome}" + for v in matched + if not facts.holds(v) + ] + queued = {e.key for e in facts.queue} + gaps.extend(f"{v.label()} is queued for a re-run" for v in matched if v.key in queued) + return matched, gaps + + +def _items(facts: Facts, overlay: Mapping[str, Any]) -> list[dict[str, Any]]: + rows = [] + for item in overlay.get("items") or (): + if not isinstance(item, Mapping): + continue + row = {str(k): v for k, v in item.items()} + if item.get("evidence"): + matched, gaps = _item_evidence(facts, item) + if not matched: + row["evidence_status"] = "cites no verdict" + elif gaps: + row["evidence_status"] = "open: " + "; ".join(gaps) + else: + row["evidence_status"] = f"closed by {len(matched)} fresh passing verdict(s)" + rows.append(row) + return rows + + +def build_progress( + facts: Facts, overlay: Mapping[str, Any] | None = None, *, workspace: str | Path = "" +) -> dict[str, Any]: + """The ``progress.json`` document. Every fact in it comes from ``facts``.""" + overlay = dict(overlay or {}) + state = facts.state + meta = dict(overlay.get("meta") or {}) + meta.update( + { + "title": meta.get("title") or f"Onboarding progress — {state.get('slug') or 'run'}", + "slug": state.get("slug", ""), + "branch": state.get("branch", ""), + "repo": facts.repo, + "stage": state.get("stage", ""), + "verdict": state.get("verdict", ""), + "status": facts.status(), + "last_updated": utc_now(), + "final_report": FINAL_REPORT + if workspace and (Path(workspace) / FINAL_REPORT).is_file() + else meta.get("final_report", ""), + } + ) + regression = facts.regression + doc: dict[str, Any] = { + "schema": SCHEMA_VERSION, + "generated_at": meta["last_updated"], + "generated_from": { + "ledger": "ledger.jsonl", + "ledger_digest": facts.ledger_digest, + "verdicts": facts.verdict_count, + "state": STATE_FILENAME, + }, + "meta": meta, + "summary": {"headline": headline(facts), "gates": facts.counts()}, + "pipeline_stages": _stage_rows(state), + "units": _unit_rows(facts), + "gates": { + "note": "Latest verdict per (gate, unit, member), generated from the ledger.", + "rows": [gate_row(facts, v) for v in facts.latest], + }, + "findings": [ + { + "gate": v.gate, + "unit": v.unit, + "member": v.member, + "value_vs_anchor": v.value_text(), + "detail": v.detail[:300], + } + for v in facts.findings + ], + "rerun_queue": [e.to_json() for e in facts.queue], + "regression": { + "last": None + if regression is None + else { + "tree": regression.tree.short(), + "tree_object": regression.tree.tree, + "commit": regression.tree.commit, + "native": regression.native.short(), + "verdicts": len(regression.keys), + "timestamp": regression.timestamp, + }, + "binding": None + if facts.binding is None + else { + "ok": facts.binding.ok, + "claim": facts.binding.claim, + "detail": facts.binding.detail, + }, + }, + "items": _items(facts, overlay), + "timeline": [dict(e) for e in facts.events] + list(overlay.get("timeline") or ()), + } + for key in ("features", "streams", "workflow_feedback", "notes"): + if overlay.get(key): + doc[key] = overlay[key] + return doc + + +# ------------------------------------------------------------- consistency + + +@dataclass(frozen=True) +class Contradiction: + """A statement in a report that the ledger refutes.""" + + where: str + claim: str + fact: str + + def render(self) -> str: + return f"{self.where}: says {self.claim!r}, but {self.fact}" + + +_NO_FINDING = re.compile( + r"\bno\s+(?:new\s+|open\s+)?findings?\b|\bfindings?\s*[:=]\s*(?:none|0)\b", re.I +) +_ALL_PASS = re.compile(r"\ball\s+(?:\d+\s+)?(?:gates?|verdicts?)\s+(?:pass(?:ed)?|green)\b", re.I) +_N_OF_N = re.compile(r"\b(\d+)\s*/\s*(\d+)\s+(?:gates?\s+)?(?:PASS(?:ED)?|green)\b", re.I) +_TREE_CLAIM = re.compile( + r"\b(?:final|delivered|regression|merge)\s+(?:tree|head|commit)\b[^0-9a-f`]{0,12}`?" + r"(?P<sha>[0-9a-f]{7,40})\b", + re.I, +) + + +def check_claims(text: str, facts: Facts, *, where: str) -> list[Contradiction]: + """Prose claims in ``text`` that the ledger refutes.""" + problems: list[Contradiction] = [] + counts = facts.counts() + for match in _NO_FINDING.finditer(text): + if facts.findings: + names = ", ".join(v.label() for v in facts.findings[:5]) + problems.append( + Contradiction(where, match.group(0), f"the ledger holds a FINDING: {names}") + ) + not_holding = [v for v in facts.latest if not facts.holds(v)] + for match in _ALL_PASS.finditer(text): + if not_holding: + names = ", ".join(v.label() for v in not_holding[:5]) + problems.append(Contradiction(where, match.group(0), f"these do not hold now: {names}")) + for match in _N_OF_N.finditer(text): + claimed, total = int(match.group(1)), int(match.group(2)) + if claimed != total: + continue + sizes = {counts["total"]} + if facts.regression is not None: + sizes.add(len(facts.regression.keys)) + if total not in sizes: + problems.append( + Contradiction( + where, + match.group(0), + f"the ledger's latest verdicts number {counts['total']} and the last " + f"regression {len(facts.regression.keys) if facts.regression else 0}", + ) + ) + elif not_holding: + names = ", ".join(v.label() for v in not_holding[:5]) + problems.append(Contradiction(where, match.group(0), f"these do not hold now: {names}")) + regression = facts.regression + for match in _TREE_CLAIM.finditer(text): + sha = match.group("sha").lower() + if regression is None: + problems.append(Contradiction(where, match.group(0), "no full regression is recorded")) + elif not ( + regression.tree.commit.startswith(sha) + or regression.tree.tree.startswith(sha) + or sha in regression.tree.describe + ): + problems.append( + Contradiction( + where, + match.group(0), + f"the last full regression ran on {regression.tree.short()} " + f"({regression.tree.commit[:12]})", + ) + ) + return problems + + +def check_items(items: Iterable[Mapping[str, Any]], facts: Facts) -> list[Contradiction]: + """Overlay items whose stated status their cited evidence refutes.""" + problems: list[Contradiction] = [] + for item in items: + if not isinstance(item, Mapping) or not item.get("evidence"): + continue + name = f"item {item.get('id') or item.get('title') or '?'}" + status = str(item.get("status") or "").upper() + matched, gaps = _item_evidence(facts, item) + if status in CLOSED_STATUSES: + if not matched: + problems.append(Contradiction(name, status, "its evidence matches no verdict")) + elif gaps: + problems.append(Contradiction(name, status, "; ".join(gaps))) + elif status in OPEN_STATUSES and matched and not gaps: + problems.append( + Contradiction( + name, + status, + f"it is in fact closed: {len(matched)} fresh passing verdict(s) " + f"({', '.join(v.label() for v in matched[:3])})", + ) + ) + return problems + + +def _strings(value: Any, path: str) -> Iterable[tuple[str, str]]: + if isinstance(value, str): + yield path, value + elif isinstance(value, Mapping): + for key, inner in value.items(): + yield from _strings(inner, f"{path}.{key}") + elif isinstance(value, (list, tuple)): + for index, inner in enumerate(value): + yield from _strings(inner, f"{path}[{index}]") + + +def check_overlay(overlay: Mapping[str, Any], facts: Facts) -> list[Contradiction]: + """The operator's narrative and items, read against the ledger.""" + problems = check_items(overlay.get("items") or (), facts) + for key, value in overlay.items(): + if key == "timeline": + continue + for where, text in _strings(value, f"overlay.{key}"): + problems.extend(check_claims(text, facts, where=where)) + return problems + + +def check_progress(doc: Mapping[str, Any], facts: Facts) -> list[Contradiction]: + """A ``progress.json`` on disk, read against the ledger it claims to come from.""" + problems: list[Contradiction] = [] + digest = str((doc.get("generated_from") or {}).get("ledger_digest") or "") + if digest != facts.ledger_digest: + problems.append( + Contradiction( + PROGRESS_JSON, + f"generated from ledger {digest or '(none)'}", + f"the ledger is now {facts.ledger_digest or '(empty)'}; the page is stale", + ) + ) + status = str((doc.get("meta") or {}).get("status") or "") + if status and status != facts.status(): + problems.append( + Contradiction(f"{PROGRESS_JSON} meta.status", status, f"the facts say {facts.status()}") + ) + latest = {v.key: v for v in facts.latest} + for row in (doc.get("gates") or {}).get("rows") or (): + key = (str(row.get("gate")), str(row.get("unit")), str(row.get("member"))) + verdict = latest.get(key) + where = f"{PROGRESS_JSON} gate {key[0]}@{key[2]} [{key[1]}]" + if verdict is None: + problems.append(Contradiction(where, str(row.get("result")), "no such verdict")) + continue + if str(row.get("result")) != verdict.outcome: + problems.append( + Contradiction(where, str(row.get("result")), f"the ledger says {verdict.outcome}") + ) + fresh = STATUS_STALE if facts.is_stale(verdict) else STATUS_FRESH + if row.get("fresh") and row.get("fresh") != fresh: + problems.append(Contradiction(where, str(row.get("fresh")), f"it is {fresh}")) + queued = {e.key for e in facts.queue} + for entry in doc.get("rerun_queue") or (): + key = (str(entry.get("gate")), str(entry.get("unit")), str(entry.get("member"))) + if key not in queued: + problems.append( + Contradiction( + f"{PROGRESS_JSON} rerun_queue", + f"{key[0]}@{key[2]} [{key[1]}] pending", + "it is in fact closed: its latest verdict holds", + ) + ) + summary = doc.get("summary") or {} + if isinstance(summary, Mapping): + for where, text in _strings(summary, f"{PROGRESS_JSON} summary"): + problems.extend(check_claims(text, facts, where=where)) + return problems + + +def narrative(text: str) -> str: + """``text`` with every generated region removed: what a human wrote.""" + return _GENERATED.sub("", text) + + +def check_consistency( + workspace: str | Path, + *, + facts: Facts | None = None, + documents: Sequence[str | Path] = (), +) -> list[Contradiction]: + """Every report in the workspace, plus ``documents``, read against the ledger.""" + workspace = Path(workspace) + facts = facts or collect_facts(workspace) + problems: list[Contradiction] = [] + overlay = read_overlay(workspace) + problems.extend(check_overlay(overlay, facts)) + progress = workspace / PROGRESS_JSON + if progress.is_file(): + try: + problems.extend(check_progress(json.loads(progress.read_text("utf-8")), facts)) + except json.JSONDecodeError as exc: + problems.append(Contradiction(PROGRESS_JSON, "a JSON document", f"it is not: {exc}")) + for path in [workspace / FINAL_REPORT, *map(Path, documents)]: + if path.is_file(): + text = narrative(path.read_text(encoding="utf-8")) + problems.extend(check_claims(text, facts, where=path.name)) + unique: dict[str, Contradiction] = {} + for problem in problems: + unique.setdefault(problem.render(), problem) + return list(unique.values()) + + +# ------------------------------------------------------------- final report + + +def gate_table_markdown(facts: Facts) -> str: + if not facts.latest: + return "_No verdict recorded._" + lines = [ + "| unit | member | gate | result | fresh | value vs threshold | tree | measured | runs |", + "|---|---|---|---|---|---|---|---|---|", + ] + for verdict in facts.latest: + row = gate_row(facts, verdict) + cells = [ + row["unit"] or "-", + row["member"] or "-", + row["gate"], + f"**{row['result']}**" if row["result"] != OUTCOME_PASS else row["result"], + row["fresh"], + row["value_vs_anchor"].replace("|", "\\|"), + f"`{row['tree']}`", + row["timestamp"], + str(row["runs"]), + ] + lines.append("| " + " | ".join(cells) + " |") + return "\n".join(lines) + + +def findings_markdown(facts: Facts) -> str: + if not facts.findings: + return "FINDING verdicts: 0." + return "\n".join( + f"- **{v.label()}**: {v.value_text()} — {v.detail[:300]}" for v in facts.findings + ) + + +def freshness_markdown(facts: Facts) -> str: + stale = [facts.assessments[v.key] for v in facts.stale] + lines = [f"Stale verdicts: {len(stale)}. Queued for a re-run: {len(facts.queue)}."] + for item in stale: + lines.append(f"- {item.verdict.label()}: {'; '.join(item.reasons)}") + return "\n".join(lines) + + +def regression_markdown(facts: Facts) -> str: + regression = facts.regression + if regression is None: + lines = ["No full regression is recorded."] + else: + lines = [ + f"Last full regression: {len(regression.keys)} verdict(s) on tree " + f"`{regression.tree.short()}` (commit `{regression.tree.commit[:12]}`, tree object " + f"`{regression.tree.tree[:12]}`), native build `{regression.native.short()}`, " + f"recorded {regression.timestamp}." + ] + if facts.binding is not None: + verdict = "holds" if facts.binding.ok else "DOES NOT HOLD" + lines.append(f"Binding to the tree under merge {verdict}: {facts.binding.detail}.") + return "\n\n".join(lines) + + +_SECTIONS: tuple[tuple[str, str, Callable[[Facts], str]], ...] = ( + ("summary", "", headline), + ("gate-table", "## Gates", gate_table_markdown), + ("findings", "## Findings", findings_markdown), + ("freshness", "## Evidence freshness", freshness_markdown), + ("regression", "## Final regression", regression_markdown), +) + + +def _region(name: str, body: str) -> str: + return f"<!-- BEGIN GENERATED: {name} -->\n{body.rstrip()}\n<!-- END GENERATED: {name} -->" + + +def render_final_report(facts: Facts, existing: str | None = None) -> str: + """Refresh the generated regions of a final report, or write a new one. + + Text outside the ``BEGIN/END GENERATED`` markers is kept as written; a + region the existing report lacks is appended. + """ + if existing is None: + title = f"# Final report — {facts.state.get('slug') or 'onboarding run'}" + parts = [title] + for name, heading, render in _SECTIONS: + if heading: + parts.append(heading) + parts.append(_region(name, render(facts))) + parts.append( + "## Narrative\n\nWritten by the operator. Everything outside the generated " + "regions is read against the ledger by `report.check_consistency`." + ) + return "\n\n".join(parts) + "\n" + text = existing + for name, heading, render in _SECTIONS: + region = _region(name, render(facts)) + pattern = re.compile( + rf"<!-- BEGIN GENERATED: {name} -->\n.*?<!-- END GENERATED: {name} -->", re.S + ) + if pattern.search(text): + text = pattern.sub(lambda _m, r=region: r, text) + else: + text = text.rstrip() + "\n\n" + (f"{heading}\n\n" if heading else "") + region + "\n" + return text + + +# -------------------------------------------------------------------- HTML + +# Status -> (css class, icon). Color is never the only signal: every badge +# also carries the icon and the literal status text. +STATUS_MAP = { + "NOT STARTED": ("status-neutral", "○"), + "RESEARCH": ("status-info", "◐"), + "IN PROGRESS": ("status-warning", "◑"), + "GATED-PASS": ("status-good", "◆"), + "DONE": ("status-good", "●"), + "PASS": ("status-good", "●"), + "FRESH": ("status-good", "✓"), + "BLOCKED": ("status-critical", "✕"), + "FAIL": ("status-critical", "✕"), + "FINDING": ("status-serious", "▲"), + "STALE": ("status-serious", "↻"), + "QUEUED": ("status-warning", "◇"), + "DESCOPED": ("status-neutral", "—"), + "OPEN": ("status-warning", "◇"), + "PENDING": ("status-warning", "◇"), + "DELIVERED": ("status-good", "★"), + "CONFIRMED": ("status-good", "✓"), + "PARTIAL": ("status-warning", "◑"), +} + + +def esc(value: Any) -> str: + return html.escape(str(value if value is not None else ""), quote=True) + + +def badge(status: str) -> str: + status = str(status or "NOT STARTED") + css_class, icon = STATUS_MAP.get(status.upper(), ("status-neutral", "○")) + return ( + f'<span class="badge {css_class}"><span class="badge-icon" aria-hidden="true">{icon}' + f'</span><span class="badge-label">{esc(status)}</span></span>' + ) + + +def _link(path: Any) -> str: + if not path: + return '<span class="muted">—</span>' + return f'<a href="{esc(path)}">{esc(path)}</a>' + + +def _table(headers: Sequence[str], rows: Iterable[Sequence[str]], empty: str) -> str: + body = "".join("<tr>" + "".join(f"<td>{c}</td>" for c in row) + "</tr>" for row in rows) + if not body: + return f'<p class="muted">{esc(empty)}</p>' + head = "".join(f"<th>{esc(h)}</th>" for h in headers) + return ( + f'<div class="table-scroll"><table class="data-table"><thead><tr>{head}</tr></thead>' + f"<tbody>{body}</tbody></table></div>" + ) + + +def _card(title: str, body: str) -> str: + return f'<section class="card"><h2>{esc(title)}</h2>{body}</section>' + + +def render_html(doc: Mapping[str, Any]) -> str: + """Render a ``progress.json`` document into one offline page.""" + meta = doc.get("meta") or {} + summary = doc.get("summary") or {} + counts = summary.get("gates") or {} + regression = doc.get("regression") or {} + binding = regression.get("binding") or {} + consistency = doc.get("consistency") or {} + + meta_rows = "".join( + f'<div class="meta-row"><div class="meta-key">{esc(k)}</div>' + f'<div class="meta-val">{esc(v)}</div></div>' + for k, v in meta.items() + if k not in ("title", "status", "last_updated") and v not in ("", None) + ) + tiles = "".join( + f'<div class="tile"><div class="tile-label">{esc(label)}</div>' + f'<div class="tile-value">{esc(counts.get(key, 0))}</div></div>' + for key, label in ( + ("total", "verdicts"), + ("pass", "pass"), + ("finding", "finding"), + ("fail", "fail"), + ("stale", "stale"), + ("queued", "re-run queue"), + ) + ) + parts = [ + _card( + "Overview", + f'<p class="headline">{esc(summary.get("headline"))}</p>' + f'<div class="tile-row">{tiles}</div>' + f'<div class="meta-grid" style="margin-top:12px">{meta_rows}</div>', + ) + ] + problems = consistency.get("problems") or [] + if problems: + parts.append( + _card( + "Consistency check: contradictions", + '<ul class="feedback-list">' + + "".join(f"<li>{badge('BLOCKED')} {esc(p)}</li>" for p in problems) + + "</ul>", + ) + ) + parts.append( + _card( + "Pipeline stages", + _table( + ("Stage", "Status"), + ((esc(r["stage"]), badge(r["status"])) for r in doc.get("pipeline_stages") or ()), + "No stages recorded.", + ), + ) + ) + parts.append( + _card( + "Units", + _table( + ("Unit", "Status", "Gates"), + ( + (esc(r["unit"]), badge(r["status"]), esc(r["gates"])) + for r in doc.get("units") or () + ), + "No units.", + ), + ) + ) + gate_rows = (doc.get("gates") or {}).get("rows") or () + parts.append( + _card( + "Gates", + _table( + ( + "Unit", + "Member", + "Gate", + "Result", + "Fresh", + "Value vs threshold", + "Tree", + "Measured", + "Runs", + "Why stale", + ), + ( + ( + esc(r.get("unit")), + esc(r.get("member")), + esc(r.get("gate")), + badge(r.get("result")), + badge(r.get("fresh")), + esc(r.get("value_vs_anchor")), + f"<code>{esc(r.get('tree'))}</code>", + esc(r.get("timestamp")), + esc(r.get("runs")), + esc("; ".join(r.get("reasons") or ())), + ) + for r in gate_rows + ), + "No verdict recorded yet.", + ), + ) + ) + findings = doc.get("findings") or () + parts.append( + _card( + "Findings", + _table( + ("Gate", "Unit", "Member", "Value", "Detail"), + ( + ( + esc(f["gate"]), + esc(f["unit"]), + esc(f["member"]), + esc(f["value_vs_anchor"]), + esc(f["detail"]), + ) + for f in findings + ), + "FINDING verdicts: 0.", + ), + ) + ) + parts.append( + _card( + "Re-run queue", + _table( + ("Gate", "Unit", "Member", "Why"), + ( + (esc(e["gate"]), esc(e["unit"]), esc(e["member"]), esc("; ".join(e["reasons"]))) + for e in doc.get("rerun_queue") or () + ), + "Nothing queued.", + ), + ) + ) + last = regression.get("last") or {} + regression_body = ( + f"<p>{esc(last.get('verdicts'))} verdict(s) on <code>{esc(last.get('tree'))}" + f"</code>, native <code>{esc(last.get('native'))}</code>, {esc(last.get('timestamp'))}</p>" + if last + else '<p class="muted">No full regression recorded.</p>' + ) + if binding: + regression_body += f"<p>{badge('DONE' if binding.get('ok') else 'BLOCKED')} {esc(binding.get('detail'))}</p>" + parts.append(_card("Final regression", regression_body)) + items = doc.get("items") or () + if items: + parts.append( + _card( + "Items", + _table( + ("Item", "Title", "Status", "Evidence"), + ( + ( + esc(i.get("id")), + esc(i.get("title")), + badge(i.get("status")), + esc(i.get("evidence_status", "narrative only")), + ) + for i in items + ), + "", + ), + ) + ) + if doc.get("features"): + parts.append( + _card( + "Feature scorecard", + _table( + ("Feature", "Meaning", "Mapping", "Status", "Evidence", "Notes"), + ( + ( + esc(r.get("feature")), + esc(r.get("meaning")), + esc(r.get("mapping")), + badge(r.get("status")), + _link(r.get("evidence")), + esc(r.get("notes")), + ) + for r in doc["features"] + ), + "", + ), + ) + ) + if doc.get("streams"): + parts.append( + _card( + "Parallel streams", + _table( + ("Stream", "Detail", "Status", "Evidence"), + ( + ( + esc(r.get("stream")), + esc(r.get("detail")), + badge(r.get("status")), + _link(r.get("evidence")), + ) + for r in doc["streams"] + ), + "", + ), + ) + ) + if doc.get("workflow_feedback"): + parts.append( + _card( + "Workflow feedback", + '<ul class="feedback-list">' + + "".join( + f"<li><code>{esc(i.get('id'))}</code> <strong>{esc(i.get('title'))}</strong> " + f"{badge(i.get('status'))}" + + (f" — {_link(i.get('link'))}" if i.get("link") else "") + + "</li>" + for i in doc["workflow_feedback"] + ) + + "</ul>", + ) + ) + timeline = doc.get("timeline") or () + parts.append( + _card( + "Timeline", + '<ul class="timeline-list">' + + "".join( + f'<li><span class="timeline-time">{esc(e.get("time"))}</span>' + f'<span class="timeline-event">{esc(e.get("event"))}</span></li>' + for e in timeline + ) + + "</ul>" + if timeline + else '<p class="muted">No events recorded.</p>', + ) + ) + final = meta.get("final_report") + return _PAGE.format( + title=esc(meta.get("title", "Onboarding progress")), + status_badge=badge(meta.get("status")) if meta.get("status") else "", + last_updated=esc(meta.get("last_updated", "")), + final_report_link=f' — <a href="{esc(final)}">final report</a>' if final else "", + body="\n".join(parts), + ) + + +_PAGE = """<!DOCTYPE html> +<html lang="en"> +<head> +<meta charset="utf-8"> +<meta name="viewport" content="width=device-width, initial-scale=1"> +<title>{title} + + + +
+
+
+

{title} {status_badge}

+
Last updated: {last_updated}{final_report_link}
+
+ +
+{body} +
Generated from ledger.jsonl by staircase_onboard.report — offline, no external resources.
+
+ + + +""" + + +# ------------------------------------------------------------------- writing + +_write_lock = threading.Lock() + + +def read_overlay(workspace: str | Path) -> dict[str, Any]: + """The operator's narrative for the progress page, or ``{}``.""" + path = Path(workspace) / OVERLAY_FILENAME + try: + data = yaml.safe_load(path.read_text(encoding="utf-8")) + except FileNotFoundError: + return {} + except yaml.YAMLError as exc: + return {"notes": [f"{OVERLAY_FILENAME} does not parse: {exc}"]} + return dict(data) if isinstance(data, Mapping) else {} + + +def _atomic_write(path: Path, text: str) -> None: + tmp = path.with_name(f".{path.name}.{os.getpid()}.{threading.get_ident()}.tmp") + tmp.write_text(text, encoding="utf-8") + os.replace(tmp, path) + + +def write_reports( + workspace: str | Path, + *, + facts: Facts | None = None, + final_report: bool = False, +) -> dict[str, Any]: + """Regenerate ``progress.json``/``progress.html``, and the final report's tables. + + The final report is written when ``final_report`` is set or the file + already exists; its narrative is kept. The consistency check's result is + embedded in ``progress.json`` so the page shows its own contradictions. + """ + workspace = Path(workspace) + facts = facts or collect_facts(workspace) + overlay = read_overlay(workspace) + with _write_lock: + report = workspace / FINAL_REPORT + if final_report or report.is_file(): + existing = report.read_text(encoding="utf-8") if report.is_file() else None + _atomic_write(report, render_final_report(facts, existing)) + doc = build_progress(facts, overlay, workspace=workspace) + problems = check_overlay(overlay, facts) + if report.is_file(): + problems.extend( + check_claims( + narrative(report.read_text(encoding="utf-8")), facts, where=FINAL_REPORT + ) + ) + doc["consistency"] = {"ok": not problems, "problems": [p.render() for p in problems]} + _atomic_write( + workspace / PROGRESS_JSON, json.dumps(doc, indent=2, ensure_ascii=False) + "\n" + ) + _atomic_write(workspace / PROGRESS_HTML, render_html(doc)) + return doc + + +# ----------------------------------------------------------------------- CLI + + +def main(argv: Sequence[str] | None = None) -> int: + """Regenerate a workspace's reports, or check them (and other documents) for contradictions.""" + parser = argparse.ArgumentParser(prog="staircase-report", description=main.__doc__) + parser.add_argument("workspace") + parser.add_argument("--repo", help="the repository, if not the workspace task.yaml's") + parser.add_argument("--final-report", action="store_true", help="also write final-report.md") + parser.add_argument( + "--check", + nargs="*", + metavar="DOC", + help="do not write; check the workspace's reports and these documents, exit 1 on any " + "contradiction", + ) + args = parser.parse_args(argv) + facts = collect_facts(args.workspace, repo_root=args.repo, check_history=True) + if args.check is not None: + problems = check_consistency(args.workspace, facts=facts, documents=args.check) + for problem in problems: + sys.stdout.write(f"CONTRADICTION {problem.render()}\n") + sys.stdout.write(f"{len(problems)} contradiction(s)\n") + return 1 if problems else 0 + doc = write_reports(args.workspace, facts=facts, final_report=args.final_report) + sys.stdout.write(f"{doc['summary']['headline']}\n") + return 0 + + +__all__ = [ + "CLOSED_STATUSES", + "Contradiction", + "FINAL_REPORT", + "Facts", + "OPEN_STATUSES", + "OVERLAY_FILENAME", + "PROGRESS_HTML", + "PROGRESS_JSON", + "build_progress", + "check_claims", + "check_consistency", + "check_items", + "check_overlay", + "check_progress", + "collect_facts", + "gate_table_markdown", + "headline", + "narrative", + "read_overlay", + "render_final_report", + "render_html", + "write_reports", +] + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/SKILL.md index cd418665a311..dfdf2c3e537f 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/SKILL.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/SKILL.md @@ -54,8 +54,10 @@ Every rule below that says "every file of the entry" means both trees. it passes, or until you have established a genuine op failure. A fake pass is not a product. 5. **Record the receipt** from that observed run only. -6. **Sync - `tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml`.** +6. **Write the entry's index fragment** + `tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.d//.yaml` + (shape in `references/entry-spec.md`). Never edit `catalog/index.yaml`: the + orchestrator regenerates it from the committed fragments. 7. **Do not format or lint during iteration.** The repo's `pre-commit` runs once, at delivery, and covers formatting, import order and the general lint rules in one pass. What the retired per-entry recipe taught is why that pass @@ -77,7 +79,7 @@ Every rule below that says "every file of the entry" means both trees. git add tensorrt_llm/_torch/_experimental/modeling_v2/catalog//.md \ tensorrt_llm/_torch/_experimental/modeling_v2/catalog//.py \ tests/unittest/_torch/modeling_v2//test_modeling_v2_.py \ - tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml + tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.d//.yaml git commit -m "feat(catalog): onboard — " ``` diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decomposition.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decomposition.md index fde59f29aabe..43fe9ccdb21a 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decomposition.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/decomposition.md @@ -6,7 +6,7 @@ lookup delegating to a module `forward`, module calls, tensor glue beyond the single call, a phase dispatch, or nested calls into other registered ops. A macro verdict ends the onboarding path for this run. Onboard nothing: no -entry files, no `index.yaml` change. Produce the report instead, land it, and +entry files, no index fragment. Produce the report instead, land it, and dispatch one child per constituent. `decompositions/mla-custom-op.md` is both a worked example and the shape to diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/entry-spec.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/entry-spec.md index 1b36a006a780..a30aea153b03 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/entry-spec.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-onboard/references/entry-spec.md @@ -112,13 +112,18 @@ Receipts live in the contract frontmatter and are test outcomes, not claims: Derive the arch key from the device the test actually ran on (`torch.cuda.get_device_capability()`, e.g. `(10, 0)` → `sm_100`). **A receipt is valid only against the files of its own entry**: it must post-date the last write to every one of them, and an entry spans two trees, so the check looks in both. A write to any of the three makes the receipt stale — never rewrite a receipt without rerunning the test. -## index.yaml +## Index fragment -Append to `tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml`, -matching the existing shape: +Write the entry's row as its own file, +`tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.d//.yaml`, +and commit it with the entry: ```yaml - - path: norm/rms_norm.py - impl: torch.ops.trtllm. - summary: "One-line compute-role summary" +path: norm/rms_norm.py +impl: torch.ops.trtllm. +summary: "One-line compute-role summary" ``` + +`catalog/index.yaml` is generated from these fragments by the orchestrator and +must not be edited by hand. When several smiths run at once, each commits only +its own fragment, so no commit carries another entry's row. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/SKILL.md b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/SKILL.md index 82db861b6d27..508a3a9ab9a8 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/SKILL.md +++ b/agent-flow/agent_flow/workflows/staircase_onboard/skills/staircase-kernel-verify/SKILL.md @@ -60,11 +60,14 @@ target from a wrong answer into a loud failure. ```bash git add tensorrt_llm/_torch/_experimental/modeling_v2/catalog//.md \ - tests/unittest/_torch/modeling_v2//test_modeling_v2_.py \ - tensorrt_llm/_torch/_experimental/modeling_v2/catalog/index.yaml + tests/unittest/_torch/modeling_v2//test_modeling_v2_.py git commit -m "feat(catalog): certify at " ``` +A verify does not touch `catalog/index.yaml`, which is generated. If the +entry's one-line summary has to change, edit and commit its fragment, +`catalog/index.d//.yaml`. + The wrapper usually stays untouched. If it does not — if covering this surface needs a signature change — that is worth saying loudly in the report, because every existing caller is affected. diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py b/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py index e52b411b2504..3dbfd890d937 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/workflow.py @@ -55,6 +55,7 @@ from . import gates as gates_mod from . import overlap as overlap_mod from . import remediation +from .catalog_index import CatalogIndexHook from .committer import CommitPolicy, install_committer from .delivery import DeliveryReport, summarize_for_auditor, verify_delivery, write_premerge_pack from .embargo import display_entries, merge_hook_maps, source_embargo_enforcement @@ -64,6 +65,7 @@ resolve_environment, shares_allocation_with_spool, ) +from .evidence import RunEvidence from .execution import select_command_executor, select_turn_executor, shell_runner_for from .journal import Journal from .preflight import ( @@ -84,7 +86,8 @@ ProgressContext, build_progress_tools, collect_decision_file, - decision_path, + current_decision, + decision_file_instruction, init_progress_file, latest_decision, latest_entry, @@ -427,7 +430,11 @@ def __init__( init_progress_file(self.progress_path) self.journal = Journal(path=self.journal_path) - self.gate_log = gates_mod.GateLog() + #: The evidence ledger: every gate result enters it through the log's sink. + self.evidence = RunEvidence(self, clean=clean) + self.gate_log = gates_mod.GateLog(sink=self.evidence.record) + #: catalog/index.yaml, regenerated from the per-entry fragments in HEAD. + self._catalog_index = CatalogIndexHook(self) self._progress_ctx = ProgressContext(path=self.progress_path) self._tools = build_progress_tools(self._progress_ctx) @@ -538,6 +545,7 @@ def _run_turn( *, iteration: int | None = None, layer: Any | None = None, + unit: str = "", ) -> str: """Run one role turn, wherever this run's transport delivers it. @@ -563,7 +571,7 @@ def _run_turn( backend=agent.backend, model=agent.model, decision_file=( - str(decision_path(self.workspace, role, iteration)) + str(current_decision(self.workspace, role, iteration, unit=unit)) if self._decision_files and iteration is not None else None ), @@ -1093,6 +1101,7 @@ def _run_early_checks(self, stage: str, log) -> None: ) path = self.workspace / "early_checks.md" path.write_text(report.render(), encoding="utf-8") + self.evidence.record_early_checks(report) warnings = report.warnings self.journal.section( f"early delivery checks ({stage})", @@ -1737,6 +1746,7 @@ def _features_through(self, unit: str) -> list[str]: return [str(u) for u in units[1 : units.index(unit) + 1]] def _assemble_one_unit(self, state: WorkflowState, unit: str, log) -> bool: + self._progress_ctx.current_unit = unit # `unit_turns` bounds a self-heal loop, so it has to count WORK. It # used to count entries into this loop, and those are not the same # number: `task.yaml` is loaded once in `_init_state`, so every spec @@ -2184,7 +2194,10 @@ def _run_deliver_phase(self, state: WorkflowState, log) -> bool: print_rule("[bold magenta]deliver[/bold magenta]", log) self._progress_ctx.current_iteration = 1 - report = self._delivery_report(state) + self._catalog_index.sync(log) + report = self.evidence.record_delivery_checks( + self._catalog_index.extend(self._delivery_report(state)) + ) self.journal.section("delivery checks", f"```\n{report.render()}\n```") print_message(report.render(), log) if not report.passed: @@ -2196,13 +2209,20 @@ def _run_deliver_phase(self, state: WorkflowState, log) -> bool: boot = self.gate_log.last("boot") if not self._run_remediation_cycle(state, report, boot, log): return False - # The delivery tail: what the branch still owes before it is a PR. + # The delivery tail. The evidence first: the remediation cycle may + # still have committed, so only now is the delivered tree known, and + # every verdict has to describe it -- a stale one is re-run, and a tree + # no full regression ran on refuses delivery. + if not self.evidence.settle_delivery(state, log): + return False + # Then what the branch still owes before it is a PR. pack, summary = write_premerge_pack( self.task_data["trtllm_repo_path"], self.workspace / "reports", family=self._preflight.family, policy=CommitPolicy.from_task(self.task_data, workspace=self.workspace), ) + summary.append(self.evidence.annotate_premerge(pack)) self.journal.section( "pre-merge pack", "\n".join(f"- {line}" for line in summary) + f"\n\n`{pack}`" ) @@ -2763,6 +2783,7 @@ def _refresh_workspace_spec(self, source: Path, log) -> None: def _checkpoint(self, state: WorkflowState) -> None: save_state(self.state_path, state) + self.evidence.refresh(state) @staticmethod def _materialize(source: str | Path, destination: Path) -> None: @@ -2898,7 +2919,9 @@ def _run_one_smith(self, item: QueueItem, note: str = "") -> str: decision on disk is the failure being retried. """ if self._decision_files and not note and item.n not in self._smiths_dispatched: - adopted = collect_decision_file(self._progress_ctx, self.workspace, "smith", item.n) + adopted = collect_decision_file( + self._progress_ctx, self.workspace, "smith", item.n, unit=item.unit + ) if adopted is not None: print_message( f"[bold cyan]→ smith {item.n} ({item.target}): adopting {adopted} written " @@ -2919,10 +2942,10 @@ def _run_one_smith(self, item: QueueItem, note: str = "") -> str: try: prompt = self._smith_prompt(item, note) if self._decision_files: - # Smiths number by queue item rather than by turn, and a batch - # re-run reuses those numbers, so the same staleness hazard - # applies here. - mark_decision_dispatch(self.workspace, "smith", item.n) + # Smiths number by queue item rather than by turn, and a retry + # reuses that number, so the retry is reserved a path of its own. + target = mark_decision_dispatch(self.workspace, "smith", item.n, unit=item.unit) + prompt = f"{prompt}\n\n{decision_file_instruction(target)}" report = self.journal.record_dispatch("smith", f"{item.mode} {item.target}", prompt) result = self._run_turn( "smith", @@ -2930,17 +2953,19 @@ def _run_one_smith(self, item: QueueItem, note: str = "") -> str: prompt, iteration=item.n, layer=layer, + unit=item.unit, ) self.journal.record_result(report, result) if self._decision_files: decision = collect_decision_file( - self._progress_ctx, self.workspace, "smith", item.n + self._progress_ctx, self.workspace, "smith", item.n, unit=item.unit ) if decision is None: print_message( f"[bold yellow]! smith {item.n} ({item.target}) ended its turn " f"without a valid decision file at " - f"{decision_path(self.workspace, 'smith', item.n)}[/bold yellow]", + f"{current_decision(self.workspace, 'smith', item.n, unit=item.unit)}" + f"[/bold yellow]", None, ) return result @@ -2968,6 +2993,9 @@ def _read_instruction(self) -> str: def _dispatch(self, role: str, *, what: str, prompt: str) -> str: """Run one agent turn, recording the prompt **before** it is sent.""" iteration = self._progress_ctx.current_iteration + unit = self._progress_ctx.current_unit + # Every role that reads the catalog index reads it current. + self._catalog_index.sync() # Every role, one place. The smith path does not come through here -- # it owns its own layer per dispatch -- so `_smith_prompt` carries # the same block explicitly. @@ -2975,39 +3003,26 @@ def _dispatch(self, role: str, *, what: str, prompt: str) -> str: if context: prompt = f"{prompt}\n\n{context}\n" if self._decision_files: - # Stated here rather than only in the system prompt so the path - # is exact and un-guessable: the role is told the absolute file - # to write, not a pattern to instantiate. - target = decision_path(self.workspace, role, iteration) - prompt = ( - f"{prompt}\n\n" - f"---\n\n" - f"**Write your decision file to exactly this path when you finish:**\n\n" - f"`{target}`\n\n" - f"Nothing downstream reads your prose. This file is how this turn " - f"reaches the orchestrator.\n" - ) - # Stamp BEFORE the turn runs. `(role, iteration)` is not unique - # across a run -- `iteration` is 1 for every deliver turn and - # restarts per unit in assembly -- so a turn that ends without - # writing a decision leaves the PREVIOUS turn's verdict sitting at - # this path looking fresh. The marker is what lets - # `collect_decision_file` tell those apart. - mark_decision_dispatch(self.workspace, role, iteration) + # Reserved BEFORE the turn runs, and never a path an earlier turn + # used -- see `progress.reserve_decision`. + target = mark_decision_dispatch(self.workspace, role, iteration, unit=unit) + prompt = f"{prompt}\n\n{decision_file_instruction(target)}" report = self.journal.record_dispatch(role, what, prompt) - result = self._run_turn(role, what, prompt, iteration=iteration) + result = self._run_turn(role, what, prompt, iteration=iteration, unit=unit) self.journal.record_result(report, result) if self._decision_files: - decision = collect_decision_file(self._progress_ctx, self.workspace, role, iteration) + decision = collect_decision_file( + self._progress_ctx, self.workspace, role, iteration, unit=unit + ) if decision is None: # Deliberately not silent and deliberately not fatal here: # the caller reads the decision back out of progress.yaml and # already knows what to do with its absence. print_message( f"[bold yellow]! {role} ended its turn without a valid decision " - f"file at {decision_path(self.workspace, role, iteration)}" + f"file at {current_decision(self.workspace, role, iteration, unit=unit)}" f"[/bold yellow]", None, ) diff --git a/agent-flow/tests/workflows/staircase_onboard/test_catalog_index.py b/agent-flow/tests/workflows/staircase_onboard/test_catalog_index.py new file mode 100644 index 000000000000..044b875d9259 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_catalog_index.py @@ -0,0 +1,418 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``catalog/index.yaml`` is generated from per-entry fragments (F013/F014). + +Concurrent smiths each appended their row to the one shared index and +committed the whole file, so the first to commit swept in its siblings' +uncommitted rows: HEAD named entries whose files were not in HEAD, and a +later smith's commit message named rows it did not write. +""" + +from __future__ import annotations + +import shutil +import subprocess +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import catalog_index as ci +from agent_flow.workflows.staircase_onboard.committer import install_committer +from agent_flow.workflows.staircase_onboard.delivery import CATALOG_DIR + +from .test_delivery import _build_repo + +REPO_ROOT = Path(__file__).resolve().parents[4] + + +def _git(root: Path, *args: str) -> str: + return subprocess.run( + ["git", "-C", str(root), *args], check=True, capture_output=True, text=True + ).stdout + + +def _commit_all(root: Path, message: str) -> None: + _git(root, "add", "-A") + _git(root, "commit", "-qm", message) + + +@pytest.fixture +def repo(tmp_path): + root = _build_repo(tmp_path) + _git(root, "init", "-q") + _git(root, "config", "user.email", "t@e.st") + _git(root, "config", "user.name", "t") + _commit_all(root, "seed") + return root + + +@pytest.fixture +def commit(tmp_path, repo): + script = install_committer(tmp_path / "ws", repo) + return ci.commit_with_script(script, repo) + + +def _catalog(repo: Path) -> Path: + return repo / CATALOG_DIR + + +def _entry(repo: Path, category: str, name: str, summary: str) -> list[str]: + """What a smith writes: the wrapper, the contract, and its own fragment.""" + catalog = _catalog(repo) + (catalog / category).mkdir(parents=True, exist_ok=True) + (catalog / category / f"{name}.py").write_text("def f(): ...\n", encoding="utf-8") + (catalog / category / f"{name}.md").write_text("# contract\n", encoding="utf-8") + fragment = catalog / ci.fragment_rel(f"{category}/{name}.py") + fragment.parent.mkdir(parents=True, exist_ok=True) + fragment.write_text( + yaml.safe_dump( + { + "path": f"{category}/{name}.py", + "impl": f"torch.ops.trtllm.{name}", + "summary": summary, + } + ), + encoding="utf-8", + ) + return [ + f"{CATALOG_DIR}/{category}/{name}.py", + f"{CATALOG_DIR}/{category}/{name}.md", + f"{CATALOG_DIR}/{ci.fragment_rel(f'{category}/{name}.py')}", + ] + + +def _index_paths(text: str) -> list[str]: + return [row["path"] for row in yaml.safe_load(text)["entries"]] + + +def _head_index(repo: Path) -> str: + return _git(repo, "show", f"HEAD:{CATALOG_DIR}/index.yaml") + + +def _last_commit_files(repo: Path) -> list[str]: + return sorted(_git(repo, "diff-tree", "--no-commit-id", "--name-only", "-r", "HEAD").split()) + + +# --------------------------------------------------------------- migration + + +def test_the_real_catalog_migrates_without_losing_a_row_or_a_comment(tmp_path): + """Run against this checkout's catalog: the conventions have to be the ones in use.""" + source = REPO_ROOT / CATALOG_DIR + if not (source / "index.yaml").is_file(): + pytest.skip("no catalog in this checkout") + copy = tmp_path / "catalog" + shutil.copytree(source, copy) + original = (copy / "index.yaml").read_text(encoding="utf-8") + if (copy / ci.FRAGMENTS_DIRNAME).exists(): + pytest.skip("this checkout's catalog is already migrated") + + ci.migrate(copy) + regenerated = (copy / "index.yaml").read_text(encoding="utf-8") + + assert yaml.safe_load(regenerated)["entries"] == yaml.safe_load(original)["entries"] + # Byte-identical apart from the notice that says the file is generated. + assert regenerated.replace(ci.GENERATED_NOTICE, "") == original + assert ci.problems(ci.snapshot_worktree(copy)) == [] + + +def test_migration_is_idempotent(tmp_path): + catalog = _catalog(_build_repo(tmp_path)) + ci.migrate(catalog) + before = {p: p.read_text() for p in catalog.rglob("*.yaml")} + ci.migrate(catalog) + assert {p: p.read_text() for p in catalog.rglob("*.yaml")} == before + + +def test_a_new_entry_sorts_after_the_migrated_ones_in_its_category(tmp_path): + root = _build_repo(tmp_path) + ci.migrate(_catalog(root)) + _entry(root, "norm", "aaa_first_by_name", "a later entry") + text = ci.render(ci.snapshot_worktree(_catalog(root))) + paths = _index_paths(text) + assert paths.index("norm/aaa_first_by_name.py") == paths.index("norm/flashinfer_rmsnorm.py") + 1 + + +# ------------------------------------------------------------------- check + + +def test_a_hand_edit_to_the_index_is_a_problem(tmp_path): + catalog = _catalog(_build_repo(tmp_path)) + ci.migrate(catalog) + index = catalog / "index.yaml" + index.write_text(index.read_text() + "\n - path: norm/ghost.py\n", encoding="utf-8") + (problem,) = ci.problems(ci.snapshot_worktree(catalog)) + assert "not what its fragments generate" in problem + + +def test_misplaced_duplicate_and_orphaned_fragments_are_problems(tmp_path): + catalog = _catalog(_build_repo(tmp_path)) + ci.migrate(catalog) + (catalog / "index.d" / "norm" / "elsewhere.yaml").write_text( + "path: norm/flashinfer_rmsnorm.py\n", encoding="utf-8" + ) + (catalog / "index.d" / "gemm" / "ghost.yaml").write_text( + "path: gemm/ghost.py\n", encoding="utf-8" + ) + found = "\n".join(ci.problems(ci.snapshot_worktree(catalog))) + assert "whose fragment belongs at index.d/norm/flashinfer_rmsnorm.yaml" in found + assert "also claimed by" in found + assert "entry gemm/ghost.py does not exist" in found + + +def test_an_unmigrated_catalog_says_so(): + snapshot = ci.Snapshot(header=None, fragments={}, files=frozenset(), index_text="entries: []\n") + assert "not migrated" in ci.problems(snapshot)[0] + + +def test_the_check_cli_exit_codes(tmp_path, capsys): + root = _build_repo(tmp_path) + assert ci.main(["check", "--repo", str(root)]) == 1 + assert ci.main(["migrate", "--repo", str(root)]) == 0 + assert ci.main(["check", "--repo", str(root)]) == 0 + assert "0 problem(s)" in capsys.readouterr().out + + +# -------------------------------------------------------------------- sync + + +def test_without_a_fragment_the_catalog_is_not_touched(repo, commit): + """A run that adds no catalog entry must not commit anything to the catalog.""" + head = _git(repo, "rev-parse", "HEAD") + assert ci.sync(repo, commit=commit).action == "skipped" + assert _git(repo, "rev-parse", "HEAD") == head + assert not (_catalog(repo) / "index.d").exists() + + +def _first_smith(repo, commit): + paths = _entry(repo, "norm", "first_fragment", "the first smith of the run") + ok, _ = commit("feat(catalog): onboard first_fragment", paths) + assert ok + + +def test_the_first_fragment_migrates_the_catalog_in_one_commit_of_its_own(repo, commit): + _first_smith(repo, commit) + result = ci.sync(repo, commit=commit) + assert result.action == "committed" + files = _last_commit_files(repo) + assert f"{CATALOG_DIR}/index.yaml" in files + assert f"{CATALOG_DIR}/index.d/_header.yaml" in files + assert all(f.startswith(f"{CATALOG_DIR}/index") for f in files) + paths = _index_paths(_head_index(repo)) + assert "norm/first_fragment.py" in paths and "gemm/cublas_mm.py" in paths + assert ci.problems(ci.snapshot_at(repo)) == [] + assert ci.sync(repo, commit=commit).action == "current" + + +def test_concurrent_smiths_never_cross_attribute(repo, commit): + """F013 as measured, replayed: smith A commits while smith B is still writing.""" + _first_smith(repo, commit) + ci.sync(repo, commit=commit) + a_paths = _entry(repo, "norm", "rmsnorm_a", "smith A's entry") + _entry(repo, "gemm", "gemm_b", "smith B's entry, not committed yet") + ok, _ = commit("feat(catalog): onboard rmsnorm_a", a_paths) + assert ok + assert _last_commit_files(repo) == sorted(a_paths) + + result = ci.sync(repo, commit=commit) + assert result.action == "committed" + assert _last_commit_files(repo) == [f"{CATALOG_DIR}/index.yaml"] + head_index = _head_index(repo) + assert "norm/rmsnorm_a.py" in _index_paths(head_index) + # B's fragment is on disk but not in HEAD, so HEAD's index does not name it. + assert "gemm/gemm_b.py" not in _index_paths(head_index) + + +def test_a_fragment_whose_entry_is_not_committed_is_left_out(repo, commit): + _first_smith(repo, commit) + ci.sync(repo, commit=commit) + paths = _entry(repo, "norm", "half_landed", "fragment committed, wrapper not") + ok, _ = commit("feat(catalog): fragment only", [paths[2]]) + assert ok + ci.sync(repo, commit=commit) + assert "norm/half_landed.py" not in _index_paths(_head_index(repo)) + assert any("does not exist" in p for p in ci.problems(ci.snapshot_at(repo))) + + +def test_a_row_committed_the_old_way_is_adopted_not_dropped(repo, commit): + _first_smith(repo, commit) + ci.sync(repo, commit=commit) + catalog = _catalog(repo) + (catalog / "norm" / "legacy.py").write_text("def f(): ...\n", encoding="utf-8") + index = catalog / "index.yaml" + index.write_text( + index.read_text() + + '\n - path: norm/legacy.py\n impl: torch.ops.trtllm.legacy\n summary: "old way"\n', + encoding="utf-8", + ) + _commit_all(repo, "smith that still edits index.yaml") + + assert ci.sync(repo, commit=commit).action == "committed" + assert f"{CATALOG_DIR}/index.d/norm/legacy.yaml" in _last_commit_files(repo) + assert "norm/legacy.py" in _index_paths(_head_index(repo)) + + +def test_an_index_with_uncommitted_edits_is_left_alone(repo, commit): + _first_smith(repo, commit) + index = _catalog(repo) / "index.yaml" + index.write_text(index.read_text() + "# someone is editing\n", encoding="utf-8") + assert ci.sync(repo, commit=commit).action == "deferred" + assert index.read_text().endswith("# someone is editing\n") + + +def test_a_failed_commit_restores_the_tree(repo, commit): + _first_smith(repo, commit) + before = (_catalog(repo) / "index.yaml").read_text() + result = ci.sync(repo, commit=lambda message, paths: (False, "lock timeout")) + assert result.action == "failed" + assert (_catalog(repo) / "index.yaml").read_text() == before + assert not (_catalog(repo) / "index.d" / ci.HEADER_FILENAME).exists() + assert _git(repo, "status", "--porcelain", "--", CATALOG_DIR) == "" + + +def test_outside_git_nothing_is_written(tmp_path): + root = _build_repo(tmp_path) + result = ci.sync(root, commit=lambda message, paths: (True, "")) + assert result.action == "skipped" + assert not (_catalog(root) / "index.d").exists() + + +# ------------------------------------------------------------ in the workflow + + +def test_a_run_generates_the_index_and_delivery_checks_it(tmp_path, task_file, monkeypatch): + from agent_flow.workflows.staircase_onboard.state import load_state + + from .test_workflow import Script, _install, _plan_then_vocabulary, _workflow + + root = Path(yaml.safe_load(task_file.read_text())["trtllm_repo_path"]) + # An earlier smith's entry, committed with its fragment the new way. + paths = _entry(root, "norm", "earlier_smith", "committed before this run") + _git(root, "add", *paths) + _git(root, "commit", "-qm", "feat(catalog): onboard earlier_smith") + workflow = _workflow(tmp_path) + _install(monkeypatch, workflow, Script(None)) + _plan_then_vocabulary(workflow, tmp_path, monkeypatch, []) + with workflow: + workflow.run(task_file) + + assert load_state(tmp_path / "ws" / ".staircase_onboard_state.json").verdict == "delivered" + assert "index.d/_header.yaml" in _git(root, "ls-files", CATALOG_DIR) + assert "norm/earlier_smith.py" in _index_paths(_head_index(root)) + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "catalog index committed" in journal + assert "[ok] catalog:index-generated" in journal + + +# ------------------------------------------- idempotent, and safe under concurrency + + +def _migration_commits(repo: Path) -> list[str]: + return [ + line + for line in _git(repo, "log", "--format=%s").splitlines() + if line.startswith("chore(catalog): split index.yaml") + ] + + +def test_a_second_sync_on_the_same_head_changes_nothing(tmp_path, repo, commit): + _first_smith(repo, commit) + lock = tmp_path / "ws" / ci.LOCK_FILENAME + assert ci.sync(repo, commit=commit, lock_path=lock).action == "committed" + head = _git(repo, "rev-parse", "HEAD") + assert ci.sync(repo, commit=commit, lock_path=lock).action == "current" + assert _git(repo, "rev-parse", "HEAD") == head + + +def test_racing_syncs_migrate_exactly_once(tmp_path, repo, commit): + """Two orchestrators (or a resume racing a dying run) on one repo.""" + import threading + + _first_smith(repo, commit) + lock = tmp_path / "ws" / ci.LOCK_FILENAME + results: list[str] = [] + barrier = threading.Barrier(6) + + def run(): + barrier.wait() + results.append(ci.sync(repo, commit=commit, lock_path=lock).action) + + threads = [threading.Thread(target=run) for _ in range(6)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert sorted(results) == ["committed"] + ["current"] * 5 + assert len(_migration_commits(repo)) == 1 + assert _git(repo, "status", "--porcelain", "--", CATALOG_DIR) == "" + + +def test_smiths_landing_concurrently_with_syncs_never_cross_attribute(tmp_path, repo, commit): + """F013 under load: six smiths commit while syncs regenerate the index.""" + import threading + + lock = tmp_path / "ws" / ci.LOCK_FILENAME + errors: list[str] = [] + landed: dict[str, list[str]] = {} + + def smith(index: int) -> None: + paths = _entry(repo, "norm", f"concurrent_{index}", f"smith {index}") + ok, output = commit(f"feat(catalog): onboard concurrent_{index}", paths) + if not ok: + errors.append(output) + landed[f"concurrent_{index}"] = sorted(paths) + + def syncer() -> None: + for _ in range(3): + if ci.sync(repo, commit=commit, lock_path=lock).action == "failed": + errors.append("sync failed") + + threads = [threading.Thread(target=smith, args=(i,)) for i in range(6)] + threads += [threading.Thread(target=syncer) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert errors == [] + ci.sync(repo, commit=commit, lock_path=lock) + assert ci.sync(repo, commit=commit, lock_path=lock).action == "current" + + paths = _index_paths(_head_index(repo)) + assert all(f"norm/{name}.py" in paths for name in landed) + assert len(_migration_commits(repo)) == 1 + for name, files in landed.items(): + sha = _git(repo, "log", "--format=%H", "--grep", f"onboard {name}$").split()[0] + changed = sorted( + _git(repo, "diff-tree", "--no-commit-id", "--name-only", "-r", sha).split() + ) + assert changed == files, name + assert ci.problems(ci.snapshot_at(repo)) == [] + assert _git(repo, "status", "--porcelain", "--", CATALOG_DIR) == "" + + +def test_a_commit_that_landed_but_reported_failure_is_not_undone(repo, commit): + _first_smith(repo, commit) + + def lands_then_complains(message, paths): + commit(message, paths) + return False, "post-commit check complained" + + result = ci.sync(repo, commit=lands_then_complains) + assert result.action == "committed" + assert (_catalog(repo) / "index.d" / ci.HEADER_FILENAME).is_file() + assert _git(repo, "status", "--porcelain", "--", CATALOG_DIR) == "" + + +def test_a_helper_that_stages_then_fails_leaves_nothing_staged(repo, commit): + """A helper that resets its own commit leaves the paths staged; the undo unstages them.""" + _first_smith(repo, commit) + + def stages_then_fails(message, paths): + _git(repo, "add", "--", *paths) + return False, "refused after staging" + + assert ci.sync(repo, commit=stages_then_fails).action == "failed" + assert _git(repo, "status", "--porcelain", "--", CATALOG_DIR) == "" + assert ci.sync(repo, commit=commit).action == "committed" diff --git a/agent-flow/tests/workflows/staircase_onboard/test_decision_files.py b/agent-flow/tests/workflows/staircase_onboard/test_decision_files.py index aeaaabe6a72d..ba22aef24f84 100644 --- a/agent-flow/tests/workflows/staircase_onboard/test_decision_files.py +++ b/agent-flow/tests/workflows/staircase_onboard/test_decision_files.py @@ -30,12 +30,14 @@ EXTRA_KEYS_BY_AGENT, ProgressContext, collect_decision_file, + current_decision, decision_is_stale, decision_path, init_progress_file, latest_decision, latest_entry, mark_decision_dispatch, + reserve_decision, ) from agent_flow.workflows.staircase_onboard.prompts import build_staircase_onboard_prompts from agent_flow.workflows.staircase_onboard.prompts._common import decision_file_block @@ -162,18 +164,17 @@ def test_a_decision_written_after_the_marker_is_collected(tmp_path): def test_a_decision_older_than_the_marker_is_refused(tmp_path): - """The failure this exists for, and it happened five times in one run. + """The failure the marker exists for, on a path written before reservation. - `(agent, iteration)` is not unique across a run — `iteration` is 1 for - every deliver-stage turn and restarts per unit in assembly. A turn that - ends without writing leaves the PREVIOUS turn's verdict at the same path, - looking fresh. Three audit verdicts were read back that way, including a - `CLEAN` that could therefore never end the run. + A legacy workspace can hold a marker and a decision at the same path, in + the wrong order. The decision predates the dispatch that wants it, so it + belongs to a previous turn and is refused. """ ctx = _context(tmp_path) stale = _write_decision(tmp_path, "auditor", 1, "UNCOVERED") os.utime(stale, (1_000_000, 1_000_000)) # written long before this dispatch - mark_decision_dispatch(tmp_path, "auditor", 1) + marker = stale.with_name(f".{stale.stem}.dispatched") + marker.write_text("", encoding="utf-8") assert decision_is_stale(tmp_path, "auditor", 1) assert collect_decision_file(ctx, tmp_path, "auditor", 1) is None @@ -181,6 +182,24 @@ def test_a_decision_older_than_the_marker_is_refused(tmp_path): assert latest_decision(ctx.path, "auditor", DECISIONS_BY_AGENT["auditor"]) is None +def test_a_previous_turns_decision_is_never_read_as_this_ones(tmp_path): + """The five silent misreads of one onboarding, now structurally impossible. + + Three audit verdicts were read back from two passes earlier because every + deliver turn shared one path. A new dispatch is reserved a new path, so + the old verdict is neither read as this turn's nor overwritten by it. + """ + ctx = _context(tmp_path) + old = _write_decision(tmp_path, "auditor", 1, "UNCOVERED") + target = mark_decision_dispatch(tmp_path, "auditor", 1) + + assert target != old + assert collect_decision_file(ctx, tmp_path, "auditor", 1) is None + target.write_text("summary: s\ndecision: CLEAN\n", encoding="utf-8") + assert collect_decision_file(ctx, tmp_path, "auditor", 1) == "CLEAN" + assert old.read_text() == "summary: s\ndecision: UNCOVERED\n" + + def test_without_a_marker_the_old_behaviour_stands(tmp_path): """A run that never stamped one must not have its decisions refused. @@ -203,7 +222,7 @@ def test_the_marker_does_not_collide_with_the_decision_it_guards(tmp_path): """A dotfile beside it, so a marker is never mistaken for an answer.""" mark_decision_dispatch(tmp_path, "smith", 4) written = sorted(p.name for p in decision_path(tmp_path, "smith", 4).parent.iterdir()) - assert written == [".smith-4.dispatched"] + assert written == [".core-smith-4.dispatched"] def test_staleness_is_judged_by_mtime_not_by_the_clock(): @@ -217,6 +236,89 @@ def test_staleness_is_judged_by_mtime_not_by_the_clock(): assert "time.time" not in src and "datetime.now" not in src +# ------------------------------------- F033: named by unit, never overwritten + + +def test_decision_files_are_named_by_unit_role_and_turn(tmp_path): + assert decision_path(tmp_path, "assembler", 2, unit="multimodal").name == ( + "multimodal-assembler-2.yaml" + ) + assert decision_path(tmp_path, "reviewer", 1, unit="core").name == "core-reviewer-1.yaml" + assert decision_path(tmp_path, "smith", 7, unit="dspark").name == "dspark-smith-7.yaml" + # Plan and deliver are not per unit; they file under their stage. + assert decision_path(tmp_path, "plan_drafter", 1, unit="mtp").name == ( + "plan-plan_drafter-1.yaml" + ) + assert decision_path(tmp_path, "auditor", 3).name == "deliver-auditor-3.yaml" + + +def test_two_units_turn_one_do_not_share_a_file(tmp_path): + """F033 as measured: the multimodal assembler's turn 1 overwrote core's.""" + core = mark_decision_dispatch(tmp_path, "assembler", 1, unit="core") + core.write_text("summary: core\ndecision: UNIT_READY\n", encoding="utf-8") + multimodal = mark_decision_dispatch(tmp_path, "assembler", 1, unit="multimodal") + multimodal.write_text("summary: mm\ndecision: BLOCKED\n", encoding="utf-8") + assert core != multimodal + assert "core" in core.read_text() + + +def test_a_reused_turn_number_gets_a_new_attempt_not_the_old_file(tmp_path): + """A re-entered unit restarts its turn count; a retried smith keeps its number.""" + first = mark_decision_dispatch(tmp_path, "smith", 7) + first.write_text("summary: s\ndecision: FAILED_RECEIPT\n", encoding="utf-8") + second = mark_decision_dispatch(tmp_path, "smith", 7) + third = mark_decision_dispatch(tmp_path, "smith", 7) + assert [first.name, second.name, third.name] == [ + "core-smith-7.yaml", + "core-smith-7.2.yaml", + "core-smith-7.3.yaml", + ] + assert current_decision(tmp_path, "smith", 7) == third + + +def test_concurrent_reservations_never_share_a_path(tmp_path): + import threading + + got: list[Path] = [] + lock = threading.Lock() + + def reserve(): + path = reserve_decision(tmp_path, "reviewer", 1, unit="core") + with lock: + got.append(path) + + threads = [threading.Thread(target=reserve) for _ in range(16)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert len(set(got)) == 16 + + +def test_the_collected_entry_names_its_file(tmp_path): + ctx = _context(tmp_path) + path = mark_decision_dispatch(tmp_path, "reviewer", 2, unit="mtp") + path.write_text("summary: s\ndecision: APPROVE\n", encoding="utf-8") + assert collect_decision_file(ctx, tmp_path, "reviewer", 2, unit="mtp") == "APPROVE" + assert latest_entry(ctx.path, "reviewer")["decision_file"] == "mtp-reviewer-2.yaml" + + +def test_a_dispatch_prompt_names_the_reserved_path(tmp_path, task_file, monkeypatch): + """The role is told the one exact file, including the unit and any attempt suffix.""" + from .test_workflow import Script, _install, _workflow + + workflow = _workflow(tmp_path) + calls = _install(monkeypatch, workflow, Script(None)) + workflow._decision_files = True + workflow._progress_ctx.current_iteration = 1 + workflow._progress_ctx.current_unit = "mtp" + workflow._dispatch("reviewer", what="mtp", prompt="Review it.") + workflow._dispatch("reviewer", what="mtp", prompt="Review it again.") + prompts = [prompt for role, prompt in calls if role == "reviewer"] + assert "mtp-reviewer-1.yaml`" in prompts[0] + assert "mtp-reviewer-1.2.yaml`" in prompts[1] + + # ------------------------------------------------------------------ prompts @@ -233,7 +335,7 @@ def test_every_role_gets_a_decision_block_when_the_mode_is_on(): for role, decisions in DECISIONS_BY_AGENT.items(): prompt = getattr(bundle, role) assert "Ending your turn" in prompt, role - assert f"decisions/{role}-.yaml" in prompt, role + assert f"decisions/-{role}-.yaml" in prompt, role for value in decisions: assert f"`{value}`" in prompt, (role, value) @@ -248,7 +350,7 @@ def test_roles_with_extra_keys_are_told_about_them(): def test_no_decision_block_when_the_tools_are_available(): """The tools are the preferred mechanism; this mode is the fallback.""" bundle = build_staircase_onboard_prompts(decision_files=False) - assert "decisions/plan_drafter-.yaml" not in bundle.plan_drafter + assert "decisions/-plan_drafter-.yaml" not in bundle.plan_drafter # ------------------------------------------------- the workspace spec copy diff --git a/agent-flow/tests/workflows/staircase_onboard/test_evidence.py b/agent-flow/tests/workflows/staircase_onboard/test_evidence.py new file mode 100644 index 000000000000..b2c97288f8ab --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_evidence.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The orchestrator's side of the ledger: which verdicts it records, and under what key. + +Besides the engine gates, a run produces three more kinds of verdict: the +envelope probes (one boot per declared envelope point), the early delivery +checks (warnings after the plan and after every unit), and the delivery +checks, which include the capacity audit. All of them are recorded, and none +may hide another. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import pytest + +from agent_flow.workflows.staircase_onboard import envelope as envelope_mod +from agent_flow.workflows.staircase_onboard.delivery import Check, DeliveryReport +from agent_flow.workflows.staircase_onboard.early_checks import EarlyFinding, EarlyReport +from agent_flow.workflows.staircase_onboard.evidence import ( + DELIVERY_CHECK_PREFIX, + EARLY_CHECK_PREFIX, + RunEvidence, + gate_key, + role_agents, +) +from agent_flow.workflows.staircase_onboard.freshness import RerunEntry +from agent_flow.workflows.staircase_onboard.gates import GateLog, GateResult +from agent_flow.workflows.staircase_onboard.journal import Journal +from agent_flow.workflows.staircase_onboard.ledger import OUTCOME_FINDING, OUTCOME_PASS +from agent_flow.workflows.staircase_onboard.progress import ProgressContext + +from .test_delivery import TARGET_REL, _build_repo + + +@dataclass +class _Probe: + unit: str + member: str + label: str + + def point(self) -> str: + return self.label + + +@dataclass +class _Plan: + probes: list[_Probe] = field(default_factory=list) + + def probes_for(self, unit: str) -> list[_Probe]: + return [p for p in self.probes if p.unit == unit] + + +@dataclass +class _Preflight: + target_rel: str = TARGET_REL + certification: Any = None + + +class _Workflow: + """Just the attributes RunEvidence reads.""" + + def __init__(self, tmp_path: Path, *, plan: _Plan | None = None, agents: Any = None) -> None: + self.workspace = tmp_path / "ws" + self.workspace.mkdir(parents=True, exist_ok=True) + repo = _build_repo(tmp_path) + self.task_data: dict[str, Any] = { + "trtllm_repo_path": str(repo), + "features": ["mtp"], + "environment": {"kind": "local", "python": "/py", "env": {"TRTLLM_X": "1"}}, + } + if agents is not None: + self.task_data["agents"] = agents + self.environment = None + self._preflight = _Preflight(certification=plan) + self._progress_ctx = ProgressContext(path=self.workspace / "progress.yaml") + self.journal = Journal(path=self.workspace / "journal.md") + self.boot_spec_path = self.workspace / "boot.yaml" + self._shell_runner = None + self.evidence = RunEvidence(self) + self.gate_log = GateLog(sink=self.evidence.record) + + def _features_through(self, unit: str) -> list[str]: + return [] if unit == "core" else ["mtp"] + + +def _probe_result(point: str, member: str = "tep4", passed: bool = True) -> GateResult: + return GateResult( + name=envelope_mod.PROBE_GATE, passed=passed, member=member, detail=f"at {point}: booted" + ) + + +# --------------------------------------------------------------- envelope + + +def test_each_envelope_point_is_its_own_ledger_key(): + """Keyed by name alone, a failing point would hide behind the next passing one.""" + assert gate_key(_probe_result("B=8 T=2048")) == "envelope[B=8 T=2048]" + assert gate_key(GateResult(name="envelope", passed=True, detail="no point")) == "envelope" + assert gate_key(GateResult(name="boot", passed=True, detail="at x: y")) == "boot" + + +def test_two_probes_at_one_member_are_both_latest(tmp_path): + wf = _Workflow(tmp_path) + wf._progress_ctx.current_unit = "core" + wf.gate_log.record(_probe_result("B=8 T=2048", passed=False)) + wf.gate_log.record(_probe_result("B=256 T=16384")) + latest = wf.evidence.ledger.latest() + assert latest[("envelope[B=8 T=2048]", "core", "tep4")].outcome == "FAIL" + assert latest[("envelope[B=256 T=16384]", "core", "tep4")].outcome == OUTCOME_PASS + + +def test_a_stale_probe_is_re_run_at_its_own_point(tmp_path, monkeypatch): + probe = _Probe("core", "tep4", "B=8 T=2048") + wf = _Workflow(tmp_path, plan=_Plan([probe])) + calls: list[Any] = [] + + def run_probe(environment, *, task_data, probe, spec, runner, features): + calls.append((probe, features)) + return _probe_result(probe.point()) + + monkeypatch.setattr(envelope_mod, "run_envelope_probe", run_probe) + monkeypatch.setattr(type(wf), "_gate_env", lambda self: None, raising=False) + entry = RerunEntry("envelope[B=8 T=2048]", "core", "tep4", ("code changed: x",)) + assert wf.evidence.rerun(entry) is True + assert calls == [(probe, [])] + assert ("envelope[B=8 T=2048]", "core", "tep4") in wf.evidence.ledger.latest() + + +def test_a_probe_the_plan_no_longer_has_is_neither_required_nor_re_run(tmp_path): + wf = _Workflow(tmp_path, plan=_Plan([_Probe("core", "tep4", "B=256 T=16384")])) + wf._progress_ctx.current_unit = "core" + wf.gate_log.record(_probe_result("B=8 T=2048")) + (verdict,) = wf.evidence.ledger.verdicts() + assert not wf.evidence.in_scope()(verdict) + entry = RerunEntry(verdict.gate, "core", "tep4", ()) + assert wf.evidence.rerunner(entry) is None + + +# ------------------------------------------------------------------ checks + + +def test_early_check_warnings_are_findings_and_carry_no_unit(tmp_path): + wf = _Workflow(tmp_path) + report = EarlyReport( + stage="after the core unit", + findings=( + EarlyFinding("tests:accuracy-file", "warn", "no accuracy test yet"), + EarlyFinding("capacity:max_num_tokens", "ok", "bound to the engine argument"), + ), + ) + assert wf.evidence.record_early_checks(report) is report + latest = wf.evidence.ledger.latest() + warn = latest[(f"{EARLY_CHECK_PREFIX}tests:accuracy-file", "", "")] + ok = latest[(f"{EARLY_CHECK_PREFIX}capacity:max_num_tokens", "", "")] + assert warn.outcome == OUTCOME_FINDING and ok.outcome == OUTCOME_PASS + assert warn.dependencies == {} # about the repository, not about the target's code + assert not wf.evidence.in_scope()(warn) # never re-run, never required by the regression + assert wf.evidence.reportable()(warn) + + +def test_delivery_checks_including_the_capacity_audit_are_recorded(tmp_path): + wf = _Workflow(tmp_path) + report = DeliveryReport( + checks=( + Check("capacity:KV_POOL_ROWS", False, "constant 4096 < envelope 8192"), + Check("matrix:dspark-x-disagg", True, "pass on abc1234"), + ) + ) + assert wf.evidence.record_delivery_checks(report) is report + latest = wf.evidence.ledger.latest() + assert latest[(f"{DELIVERY_CHECK_PREFIX}capacity:KV_POOL_ROWS", "", "")].outcome == "FAIL" + assert latest[(f"{DELIVERY_CHECK_PREFIX}matrix:dspark-x-disagg", "", "")].outcome == "PASS" + + +# ------------------------------------------------------------------ agents + + +def test_each_roles_backend_and_model_are_in_the_config_fingerprint(tmp_path): + default = _Workflow(tmp_path / "a") + default._progress_ctx.current_unit = "core" + first = default.evidence.record(GateResult(name="boot", passed=True, detail="ok")) + assert first.config["agents"]["plan_reviewer"]["backend"] == "codex" + assert first.config["agents"]["assembler"]["backend"] == "claude-code" + + routed = _Workflow(tmp_path / "b", agents={"roles": {"assembler": {"model": "other-model"}}}) + routed._progress_ctx.current_unit = "core" + second = routed.evidence.record(GateResult(name="boot", passed=True, detail="ok")) + assert second.config["agents"]["assembler"]["model"] == "other-model" + assert second.config_fingerprint != first.config_fingerprint + # Which model writes the next change does not invalidate what was measured. + assert second.spec_fingerprint == first.spec_fingerprint + + +@pytest.mark.parametrize("role", ["plan_drafter", "smith", "assembler", "reviewer", "auditor"]) +def test_role_agents_covers_every_role(role): + assert role in role_agents({}) diff --git a/agent-flow/tests/workflows/staircase_onboard/test_freshness.py b/agent-flow/tests/workflows/staircase_onboard/test_freshness.py new file mode 100644 index 000000000000..698417860df0 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_freshness.py @@ -0,0 +1,472 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for evidence freshness, the re-run queue and the final-regression binding. + +Each failure mode here happened on the DeepSeek-V4.1 run: + +- a later unit edited the target, so the earlier units' gate verdicts were + about code that no longer existed, and the report still quoted them; +- a default switch flipped (decoder replay became the multimodal default) + and nothing re-gated the units measured under the old default; +- a witness's settings changed after its verdict of record, and delivery + accepted the old verdict; +- the delivered head carried a functional change and a lint series that + landed after the last full regression, so no gate ran on what shipped. +""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import freshness +from agent_flow.workflows.staircase_onboard.freshness import ( + RerunEntry, + assess, + complete_regression, + drain, + last_regression, + read_rerun_queue, + record_regression, + regression_binding, + rerun_queue, + switch_history_flips, + write_rerun_queue, +) +from agent_flow.workflows.staircase_onboard.gates import GateResult +from agent_flow.workflows.staircase_onboard.ledger import ( + Ledger, + VerdictContext, + native_build, + record_gate_result, + tree_identity, +) + +from .test_delivery import TARGET_REL, _build_repo + +_SWITCH_LINE = '\nimport os\nREPLAY = os.environ.get("TRTLLM_REPLAY", "{}")\n' + + +def _git(root: Path, *args: str) -> str: + return subprocess.run( + ["git", "-C", str(root), *args], check=True, capture_output=True, text=True + ).stdout + + +def _commit_all(root: Path, message: str) -> None: + _git(root, "add", "-A") + _git(root, "commit", "-qm", message) + + +@pytest.fixture +def repo(tmp_path): + root = _build_repo(tmp_path) + modeling = root / TARGET_REL / "modeling.py" + modeling.write_text(modeling.read_text() + _SWITCH_LINE.format("0"), encoding="utf-8") + _git(root, "init", "-q") + _git(root, "config", "user.email", "t@e.st") + _git(root, "config", "user.name", "t") + _commit_all(root, "seed") + return root + + +def _spec(**overrides): + spec = { + "features": ["mtp"], + "accuracy": {"benchmark": "gsm8k", "witness": {"mtp": {"contrast_env": {"X": "0"}}}}, + "environment": {"kind": "local", "python": "/py", "env": {}}, + } + spec.update(overrides) + return spec + + +def _record(ledger, root, *, unit="core", gate="accuracy[gsm8k]", member="tp4", **kw): + passed = kw.pop("passed", True) + env = kw.pop("env", {}) + task = kw.pop("task", _spec()) + return record_gate_result( + ledger, + GateResult(name=gate, passed=passed, detail="d", member=member, **kw), + VerdictContext( + repo_root=str(root), target_rel=TARGET_REL, unit=unit, env=env, task_data=task + ), + ) + + +def _reasons(ledger, root, **kw): + (item,) = assess(ledger.latest().values(), root, **kw) + return item.reasons + + +# ---------------------------------------------------------------- staleness + + +def test_an_untouched_tree_leaves_a_verdict_fresh(tmp_path, repo): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + assert _reasons(ledger, repo, task_data=_spec()) == () + + +def test_a_commit_that_touches_a_dependency_makes_the_verdict_stale(tmp_path, repo): + """The core verdict, after a later unit rewrote the target it measured.""" + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + modeling = repo / TARGET_REL / "modeling.py" + modeling.write_text(modeling.read_text() + "\n# mtp increment\n", encoding="utf-8") + _commit_all(repo, "mtp") + + reasons = _reasons(ledger, repo) + assert any(r.startswith("code changed:") and "modeling.py" in r for r in reasons) + + +def test_a_change_to_an_imported_catalog_entry_makes_the_verdict_stale(tmp_path, repo): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + wrapper = next(repo.rglob("cublas_mm.py")) + wrapper.write_text("def f(): return 2\n", encoding="utf-8") + assert any("cublas_mm.py" in r for r in _reasons(ledger, repo)) + + +def test_a_change_elsewhere_in_the_repo_leaves_the_verdict_fresh(tmp_path, repo): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + (repo / "README.md").write_text("docs\n", encoding="utf-8") + _commit_all(repo, "docs") + assert _reasons(ledger, repo) == () + + +def test_a_newly_imported_entry_is_a_new_dependency(tmp_path, repo): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + extra = repo / TARGET_REL / "vision.py" + extra.write_text("X = 1\n", encoding="utf-8") + reasons = _reasons(ledger, repo) + assert any(r.startswith("new dependency:") and "vision.py" in r for r in reasons) + + +def test_a_rebuilt_native_library_makes_the_verdict_stale(tmp_path, repo): + lib = repo / "tensorrt_llm" / "libs" / "libtensorrt_llm.so" + lib.parent.mkdir(parents=True) + lib.write_bytes(b"\x7fELF-one") + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + lib.write_bytes(b"\x7fELF-two") + assert any(r.startswith("native build changed") for r in _reasons(ledger, repo)) + + +def test_a_changed_spec_for_the_unit_makes_its_verdict_stale(tmp_path, repo): + """F039's follow-up: a witness whose settings changed after its verdict of record.""" + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo, unit="mtp", gate="equivalence") + changed = _spec( + accuracy={"benchmark": "gsm8k", "witness": {"mtp": {"contrast_env": {"X": "1"}}}} + ) + assert any( + "task spec for unit mtp changed" in r for r in _reasons(ledger, repo, task_data=changed) + ) + + +def test_another_units_spec_change_leaves_a_verdict_fresh(tmp_path, repo): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo, unit="core") + changed = _spec( + accuracy={"benchmark": "gsm8k", "witness": {"mtp": {"contrast_env": {"X": "1"}}}} + ) + assert _reasons(ledger, repo, task_data=changed) == () + + +def test_a_flipped_default_in_a_dependency_is_named(tmp_path, repo): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + modeling = repo / TARGET_REL / "modeling.py" + modeling.write_text( + modeling.read_text().replace(_SWITCH_LINE.format("0"), _SWITCH_LINE.format("1")), + encoding="utf-8", + ) + reasons = _reasons(ledger, repo) + assert "default switch flipped: TRTLLM_REPLAY '0' -> '1'" in reasons + + +def test_a_flipped_default_outside_the_target_is_found_through_history(tmp_path, repo): + """The default lives in a runtime module no verdict lists as a dependency. + + This is gap-analysis D5: flipping a default must re-gate the units + measured under the old one, even though their own files did not change. + """ + runtime = repo / "tensorrt_llm" / "_torch" / "pyexecutor" / "replay.py" + runtime.parent.mkdir(parents=True) + runtime.write_text('import os\nON = os.environ.get("TRTLLM_DECODER_REPLAY", "0")\n') + _commit_all(repo, "runtime switch") + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + + runtime.write_text('import os\nON = os.environ.get("TRTLLM_DECODER_REPLAY", "1")\n') + _commit_all(repo, "replay on by default") + + commit = ledger.verdicts()[0].tree.commit + assert switch_history_flips(repo, commit) == {"TRTLLM_DECODER_REPLAY": ("0", "1")} + assert "default switch flipped: TRTLLM_DECODER_REPLAY '0' -> '1'" in _reasons(ledger, repo) + assert _reasons(ledger, repo, check_history=False) == () + + +def test_a_switch_the_gate_set_explicitly_does_not_flip_with_its_default(tmp_path, repo): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo, env={"TRTLLM_REPLAY": "0"}) + modeling = repo / TARGET_REL / "modeling.py" + modeling.write_text( + modeling.read_text().replace(_SWITCH_LINE.format("0"), _SWITCH_LINE.format("1")), + encoding="utf-8", + ) + _commit_all(repo, "flip") + reasons = _reasons(ledger, repo) + assert not [r for r in reasons if "TRTLLM_REPLAY" in r] + # it is still stale, for the honest reason: the file changed + assert any(r.startswith("code changed:") for r in reasons) + + +# -------------------------------------------------------------- re-run queue + + +def test_the_queue_holds_stale_and_failing_verdicts_only(tmp_path, repo): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo, gate="boot", member="tp4") + _record(ledger, repo, gate="accuracy[gsm8k]", member="tp4", passed=False) + _record(ledger, repo, gate="boot", member="ep4") + queue = rerun_queue(assess(ledger.latest().values(), repo)) + assert [(e.gate, e.member) for e in queue] == [("accuracy[gsm8k]", "tp4")] + assert queue[0].reasons == ("last verdict FAIL",) + + +def test_at_delivery_a_verdict_from_another_tree_is_queued(tmp_path, repo): + """Evidence carried over from an earlier commit is evidence about that commit.""" + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + (repo / "runtime.py").write_text("RETRY = 3\n", encoding="utf-8") + _commit_all(repo, "a functional change outside the target") + + assessed = assess(ledger.latest().values(), repo) + assert rerun_queue(assessed) == [] + tree = tree_identity(repo) + (entry,) = rerun_queue(assessed, tree=tree, native=native_build(repo)) + assert "measured on tree" in entry.reasons[0] + + +def test_scope_drops_verdicts_the_spec_no_longer_requires(tmp_path, repo): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo, member="tp8", passed=False) + queue = rerun_queue(assess(ledger.latest().values(), repo), scope=lambda v: v.member != "tp8") + assert queue == [] + + +def test_drain_keeps_what_could_not_be_re_run_and_says_why(): + queue = [RerunEntry("boot", "core", "tp4", ("stale",)), RerunEntry("perf", "core", "tp4", ())] + answers = {"boot": False, "perf": None} + remaining = drain(queue, lambda e: answers[e.gate]) + assert [e.reasons[-1] for e in remaining] == [ + "re-run failed", + "no automatic re-run for this gate", + ] + assert drain(queue, lambda e: True) == [] + + +def test_the_queue_round_trips_through_the_workspace(tmp_path): + entries = [RerunEntry("boot", "core", "tp4", ("code changed: x",), "abc")] + path = write_rerun_queue(tmp_path, entries) + assert json.loads(path.read_text())["entries"][0]["gate"] == "boot" + assert read_rerun_queue(tmp_path) == entries + + +# ----------------------------------------------------------------- binding + + +def _regress(ledger, repo): + tree, native = tree_identity(repo), native_build(repo) + regression = complete_regression(ledger.latest(), tree, native) + assert regression is not None + return record_regression(ledger, regression) + + +def test_a_regression_needs_every_verdict_passing_on_this_tree(tmp_path, repo): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo, gate="boot") + _record(ledger, repo, gate="accuracy[gsm8k]", passed=False) + tree, native = tree_identity(repo), native_build(repo) + assert complete_regression(ledger.latest(), tree, native) is None + _record(ledger, repo, gate="accuracy[gsm8k]") + assert complete_regression(ledger.latest(), tree, native) is not None + + +def test_the_tree_of_the_last_regression_binds_the_merge(tmp_path, repo): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + _regress(ledger, repo) + assert regression_binding(ledger, repo).ok + + # The DeepSeek-V4.1 shape: one more functional commit after the regression. + (repo / "runtime.py").write_text("RETRY = 3\n", encoding="utf-8") + _commit_all(repo, "VMM create retries") + binding = regression_binding(ledger, repo) + assert not binding.ok + assert "the last full regression ran on" in binding.detail + + +def test_a_re_signed_head_is_still_bound(tmp_path, repo): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + _regress(ledger, repo) + _git(repo, "commit", "-q", "--amend", "-s", "-m", "seed, signed") + assert regression_binding(ledger, repo).ok + + +def test_a_rebuilt_native_build_unbinds_the_merge(tmp_path, repo): + lib = repo / "tensorrt_llm" / "libs" / "libth_common.so" + lib.parent.mkdir(parents=True) + lib.write_bytes(b"\x7fELF-1") + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + _regress(ledger, repo) + lib.write_bytes(b"\x7fELF-2") + binding = regression_binding(ledger, repo) + assert not binding.ok and "native build" in binding.detail + + +def test_no_regression_at_all_refuses(tmp_path, repo): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + binding = regression_binding(ledger, repo) + assert not binding.ok and "no full regression" in binding.detail + + +def test_outside_git_the_binding_makes_no_claim(tmp_path): + binding = regression_binding(Ledger.in_workspace(tmp_path / "ws"), tmp_path) + assert binding.ok and not binding.claim + + +def test_check_merge_cli_exit_codes(tmp_path, repo, capsys): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + args = ["check-merge", str(tmp_path / "ws"), "--repo", str(repo)] + assert freshness.main(args) == 1 + _regress(ledger, repo) + assert freshness.main(args) == 0 + assert last_regression(ledger) is not None + assert "ok:" in capsys.readouterr().out + + +def test_status_cli_lists_stale_verdicts(tmp_path, repo, capsys): + ledger = Ledger.in_workspace(tmp_path / "ws") + _record(ledger, repo) + args = ["status", str(tmp_path / "ws"), "--repo", str(repo), "--write-queue"] + assert freshness.main(args) == 0 + (repo / TARGET_REL / "modeling.py").write_text("# rewritten\n", encoding="utf-8") + assert freshness.main(args) == 1 + out = capsys.readouterr().out + assert "STALE" in out and "code changed" in out + assert read_rerun_queue(tmp_path / "ws")[0].gate == "accuracy[gsm8k]" + + +# ------------------------------------------------------- delivery, end to end + + +def _delivered_run(tmp_path, task_file, monkeypatch, *, before_deliver=None, shell=None): + from agent_flow.workflows.staircase_onboard.state import load_state + + from .test_workflow import Script, _install, _plan_then_vocabulary, _workflow, passing_shell + + # `task_file` is built on this module's `repo` fixture, already a git repo. + repo = Path(yaml.safe_load(task_file.read_text())["trtllm_repo_path"]) + + calls: list[str] = [] + base = shell or passing_shell() + + def counting(command, env=None, timeout=0): + if "trtllm-eval" in command or "quickstart" in command: + calls.append(command) + return base(command, env, timeout) + + workflow = _workflow(tmp_path, shell_runner=counting) + _install(monkeypatch, workflow, Script(None)) + _plan_then_vocabulary(workflow, tmp_path, monkeypatch, []) + if before_deliver is not None: + original = workflow._run_deliver_phase + + def deliver(state, log): + before_deliver(repo) + return original(state, log) + + monkeypatch.setattr(workflow, "_run_deliver_phase", deliver) + with workflow: + workflow.run(task_file) + state = load_state(tmp_path / "ws" / ".staircase_onboard_state.json") + return state, calls, repo, Ledger.in_workspace(tmp_path / "ws") + + +def test_a_clean_delivery_records_the_regression_and_re_runs_nothing( + tmp_path, task_file, monkeypatch +): + state, calls, repo, ledger = _delivered_run(tmp_path, task_file, monkeypatch) + assert state.verdict == "delivered" + assert len(calls) == 2 # boot + accuracy for core, once each + regression = last_regression(ledger) + assert regression is not None + assert regression.tree.key == tree_identity(repo).key + assert regression_binding(ledger, repo).ok + + +def test_the_pre_merge_pack_states_the_binding_and_how_to_re_check_it( + tmp_path, task_file, monkeypatch +): + """The pack's remaining steps (lint, split) can change the tree after delivery.""" + state, _, repo, _ = _delivered_run(tmp_path, task_file, monkeypatch) + assert state.verdict == "delivered" + pack = (tmp_path / "ws" / "reports" / "premerge.md").read_text() + assert "## Evidence binding (bound)" in pack + assert "freshness check-merge" in pack and f"--repo {repo}" in pack + assert "evidence bound:" in (tmp_path / "ws" / "journal.md").read_text() + + +def test_a_change_before_delivery_re_runs_the_stale_gates(tmp_path, task_file, monkeypatch): + def edit_target(repo): + modeling = repo / TARGET_REL / "modeling.py" + modeling.write_text(modeling.read_text() + "\n# late fix\n", encoding="utf-8") + _commit_all(repo, "late fix") + + state, calls, repo, ledger = _delivered_run( + tmp_path, task_file, monkeypatch, before_deliver=edit_target + ) + assert state.verdict == "delivered" + assert len(calls) == 4 # both gates again, on the delivered tree + gates = [v for v in ledger.latest().values() if v.unit] + assert all(v.tree.key == tree_identity(repo).key for v in gates) + assert regression_binding(ledger, repo).ok + journal = (tmp_path / "ws" / "journal.md").read_text() + assert "gate(s) to re-run before delivery" in journal + + +def test_a_re_run_that_fails_refuses_delivery(tmp_path, task_file, monkeypatch): + from .test_workflow import passing_shell + + scores = iter([90.0, 10.0]) + good = passing_shell() + + def shell(command, env=None, timeout=0): + if "trtllm-eval" in command: + return passing_shell(score=next(scores))(command, env, timeout) + return good(command, env, timeout) + + def edit_target(repo): + modeling = repo / TARGET_REL / "modeling.py" + modeling.write_text(modeling.read_text() + "\n# regressed\n", encoding="utf-8") + _commit_all(repo, "regression") + + state, _, _, ledger = _delivered_run( + tmp_path, task_file, monkeypatch, before_deliver=edit_target, shell=shell + ) + assert state.verdict == "delivery-failed:evidence" + assert read_rerun_queue(tmp_path / "ws")[0].gate.startswith("accuracy") + assert last_regression(ledger) is None diff --git a/agent-flow/tests/workflows/staircase_onboard/test_journal.py b/agent-flow/tests/workflows/staircase_onboard/test_journal.py new file mode 100644 index 000000000000..2745b379a97e --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_journal.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Dispatch numbers are unique across threads and processes. + +A smith batch records its dispatches concurrently. Reading the highest +number and then writing the report let two smiths take the same ``NN``, +which the DeepSeek-V4.1 run found as duplicate dispatch numbers in its +journal (final report, section 6.3). +""" + +from __future__ import annotations + +import multiprocessing +import re +import threading +from pathlib import Path + +from agent_flow.workflows.staircase_onboard.journal import Journal + + +def _journal(tmp_path: Path) -> Journal: + journal = Journal(path=tmp_path / "journal.md") + journal.open(slug="t") + return journal + + +def test_ordinals_are_sequential_and_follow_existing_reports(tmp_path): + journal = _journal(tmp_path) + (journal.reports_dir / "07-smith-old.md").write_text("x", encoding="utf-8") + assert journal.next_ordinal() == 8 + assert journal.next_ordinal() == 9 + path = journal.record_dispatch("assembler", "core", "prompt") + assert path.name.startswith("10-assembler-") + + +def test_a_claimed_number_is_never_handed_out_again(tmp_path): + """Even when its report was never written: the claim, not the report, owns it.""" + journal = _journal(tmp_path) + claimed = journal.next_ordinal() + assert journal.next_ordinal() == claimed + 1 + assert Journal(path=tmp_path / "journal.md").next_ordinal() == claimed + 2 + + +def test_concurrent_dispatches_in_threads_get_distinct_numbers(tmp_path): + journal = _journal(tmp_path) + barrier = threading.Barrier(24) + + def dispatch(index: int) -> None: + barrier.wait() + journal.record_dispatch("smith", f"entry {index}", "prompt") + + threads = [threading.Thread(target=dispatch, args=(i,)) for i in range(24)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + numbers = sorted(int(p.name.split("-")[0]) for p in journal.reports_dir.glob("*.md")) + assert numbers == list(range(1, 25)) + journaled = re.findall(r"dispatch (\d+) · smith", journal.path.read_text(encoding="utf-8")) + assert len(set(journaled)) == 24 + + +def _dispatch_from_process(path: str, index: int) -> None: + journal = Journal(path=Path(path)) + for n in range(6): + journal.record_dispatch("smith", f"p{index} entry {n}", "prompt") + + +def test_concurrent_dispatches_in_processes_get_distinct_numbers(tmp_path): + """Two orchestrator processes on one workspace, or a resume racing a dying run.""" + journal = _journal(tmp_path) + context = multiprocessing.get_context("fork") + procs = [ + context.Process(target=_dispatch_from_process, args=(str(journal.path), i)) + for i in range(4) + ] + for proc in procs: + proc.start() + for proc in procs: + proc.join() + numbers = sorted(int(p.name.split("-")[0]) for p in journal.reports_dir.glob("*.md")) + assert numbers == list(range(1, 25)) + + +def test_latest_report_ignores_the_claims(tmp_path): + journal = _journal(tmp_path) + first = journal.record_dispatch("smith", "norm x", "p") + second = journal.record_dispatch("smith", "norm x", "p") + assert journal.latest_report("smith", "norm x") == second != first + + +def test_the_reports_directory_holds_only_reports(tmp_path): + """Readers list ``reports/`` and open every entry; the claims live beside it.""" + journal = _journal(tmp_path) + journal.record_dispatch("smith", "norm x", "p") + assert [p.is_file() for p in journal.reports_dir.iterdir()] == [True] + + +def test_a_new_journal_numbers_from_one_again(tmp_path): + """``--clean`` removes the journal and the reports; the claims go with them.""" + journal = _journal(tmp_path) + journal.record_dispatch("smith", "a", "p") + journal.record_dispatch("smith", "b", "p") + journal.path.unlink() + for report in journal.reports_dir.iterdir(): + report.unlink() + fresh = _journal(tmp_path) + assert fresh.next_ordinal() == 1 + + +def test_a_resumed_journal_keeps_counting(tmp_path): + journal = _journal(tmp_path) + journal.record_dispatch("smith", "a", "p") + resumed = Journal(path=journal.path) + resumed.open(slug="t") + assert resumed.next_ordinal() == 2 diff --git a/agent-flow/tests/workflows/staircase_onboard/test_ledger.py b/agent-flow/tests/workflows/staircase_onboard/test_ledger.py new file mode 100644 index 000000000000..f550ef36b8c7 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_ledger.py @@ -0,0 +1,531 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the results ledger: what a verdict records, and how it is stored. + +The ledger exists because a verdict without the code it was measured on can +be attributed to any code. The DeepSeek-V4.1 run's records cited evidence +from more than ten trees and none of them said which verdict belonged to +which, so a final report could call a head "16/16 PASS" that no gate had run +on. Every test here pins one thing a verdict must carry for that to be +decidable later. +""" + +from __future__ import annotations + +import json +import multiprocessing +import subprocess +import threading +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import ledger as ledger_mod +from agent_flow.workflows.staircase_onboard.delivery import CATALOG_DIR +from agent_flow.workflows.staircase_onboard.gates import GateLog, GateResult +from agent_flow.workflows.staircase_onboard.ledger import ( + OUTCOME_FAIL, + OUTCOME_FINDING, + OUTCOME_PASS, + Ledger, + VerdictContext, + blob_hash, + command_env, + dependency_paths, + harvest_switch_defaults, + native_build, + record_gate_result, + spec_fingerprint, + tree_identity, + verdict_from_result, +) + +from .test_delivery import TARGET_REL, _build_repo + + +def _git(root: Path, *args: str) -> str: + return subprocess.run( + ["git", "-C", str(root), *args], check=True, capture_output=True, text=True + ).stdout + + +def _git_repo(root: Path) -> Path: + _git(root, "init", "-q") + _git(root, "config", "user.email", "t@e.st") + _git(root, "config", "user.name", "t") + _git(root, "add", "-A") + _git(root, "commit", "-qm", "seed") + return root + + +def _result(**fields) -> GateResult: + fields.setdefault("name", "accuracy[gsm8k]") + fields.setdefault("passed", True) + fields.setdefault("detail", "ok") + return GateResult(**fields) + + +# ------------------------------------------------------------------- storage + + +def test_a_verdict_round_trips_and_the_file_is_jsonl(tmp_path): + ledger = Ledger.in_workspace(tmp_path) + verdict = record_gate_result( + ledger, + _result(member="tep4", measured=92.5, reference=92.3, tolerance=5.0), + VerdictContext(unit="core"), + ) + lines = (tmp_path / "ledger.jsonl").read_text().splitlines() + assert len(lines) == 1 + assert json.loads(lines[0])["id"] == verdict.id + (back,) = ledger.verdicts() + assert back.key == ("accuracy[gsm8k]", "core", "tep4") + assert back.value == 92.5 + assert back.threshold == pytest.approx(87.3) + assert back.outcome == OUTCOME_PASS + + +def test_the_ledger_is_append_only_and_the_newest_verdict_per_key_wins(tmp_path): + """Nothing is rewritten: a re-run adds a line, and the old one stays as history.""" + ledger = Ledger.in_workspace(tmp_path) + context = VerdictContext(unit="core") + record_gate_result(ledger, _result(member="tp4", passed=False, detail="bad"), context) + record_gate_result(ledger, _result(member="tp4", passed=True), context) + record_gate_result(ledger, _result(name="boot", member="tp4"), context) + + assert len(ledger.verdicts()) == 3 + latest = ledger.latest() + assert [k for k in latest] == [("accuracy[gsm8k]", "core", "tp4"), ("boot", "core", "tp4")] + assert latest[("accuracy[gsm8k]", "core", "tp4")].outcome == OUTCOME_PASS + + +def test_a_torn_last_line_is_skipped_and_counted(tmp_path): + """A process killed mid-append must not make the whole ledger unreadable.""" + ledger = Ledger.in_workspace(tmp_path) + record_gate_result(ledger, _result(), VerdictContext(unit="core")) + with ledger.path.open("a", encoding="utf-8") as handle: + handle.write('{"kind": "verdict", "gate": "bo') + assert len(ledger.verdicts()) == 1 + assert ledger.skipped == 1 + + +def test_concurrent_appends_from_threads_never_interleave(tmp_path): + ledger = Ledger.in_workspace(tmp_path) + + def write(index: int) -> None: + for n in range(20): + record_gate_result( + ledger, _result(member=f"m{index}", detail="x" * 5000), VerdictContext(unit=str(n)) + ) + + threads = [threading.Thread(target=write, args=(i,)) for i in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert len(ledger.verdicts()) == 160 + assert ledger.skipped == 0 + + +def _append_from_process(path: str, index: int) -> None: + ledger = Ledger(path) + for n in range(20): + record_gate_result( + ledger, _result(member=f"p{index}", detail="y" * 5000), VerdictContext(unit=str(n)) + ) + + +def test_concurrent_appends_from_processes_never_interleave(tmp_path): + """Two orchestrators, or a harness beside the run, share one ledger file.""" + path = str(tmp_path / "ledger.jsonl") + context = multiprocessing.get_context("fork") + procs = [context.Process(target=_append_from_process, args=(path, i)) for i in range(4)] + for proc in procs: + proc.start() + for proc in procs: + proc.join() + ledger = Ledger(path) + assert len(ledger.verdicts()) == 80 + assert ledger.skipped == 0 + + +# --------------------------------------------------------- outcome and values + + +def test_a_finding_is_its_own_outcome_and_still_passes(tmp_path): + """The run that motivated this reported "no finding" over a FINDING verdict.""" + + class Flagged: + name = "acceptance" + passed = True + detail = "below the paired reference" + finding = True + + verdict = verdict_from_result(Flagged(), VerdictContext(unit="dspark")) + assert verdict.outcome == OUTCOME_FINDING + assert verdict.passed + + +def test_an_explicit_outcome_wins_and_a_failed_gate_is_fail(): + class Explicit: + name = "perf" + passed = True + detail = "" + outcome = "finding" + + assert verdict_from_result(Explicit(), VerdictContext()).outcome == OUTCOME_FINDING + assert verdict_from_result(_result(passed=False), VerdictContext()).outcome == OUTCOME_FAIL + + +def test_an_explicit_threshold_is_kept_as_measured(): + verdict = verdict_from_result( + _result(name="perf", measured=1.03, reference=1.0, tolerance=0.1), + VerdictContext(), + ) + assert verdict.threshold == pytest.approx(0.9) + + class WithThreshold: + name = "perf" + passed = True + detail = "" + measured = 1.03 + threshold = 1.05 + + assert verdict_from_result(WithThreshold(), VerdictContext()).threshold == 1.05 + + +# ------------------------------------------------------------- dependencies + + +def test_dependencies_are_the_target_its_family_and_its_imported_catalog_entries(tmp_path): + root = _build_repo(tmp_path) + family = (root / TARGET_REL).parent + (family / "shared_rope.py").write_text("X = 1\n", encoding="utf-8") + (root / TARGET_REL / "__pycache__").mkdir() + (root / TARGET_REL / "__pycache__" / "modeling.cpython-312.pyc").write_bytes(b"\0") + + paths = dependency_paths(root, TARGET_REL) + + assert f"{TARGET_REL}/modeling.py" in paths + assert f"{TARGET_REL}/weights.py" in paths + assert f"{family.relative_to(root).as_posix()}/shared_rope.py" in paths + assert f"{CATALOG_DIR}/norm/flashinfer_rmsnorm.py" in paths + assert f"{CATALOG_DIR}/gemm/cublas_mm.py" in paths + assert f"{CATALOG_DIR}/torch/embedding.py" in paths + # a contract is a certification, not code the engine runs + assert not [p for p in paths if p.endswith(".md")] + assert not [p for p in paths if "__pycache__" in p] + + +def test_a_target_that_does_not_parse_still_records_its_files(tmp_path): + """A half-written target is exactly what a failing gate measures.""" + root = _build_repo(tmp_path, modeling="def broken(:\n") + paths = dependency_paths(root, TARGET_REL) + assert f"{TARGET_REL}/modeling.py" in paths + assert not [p for p in paths if p.startswith(CATALOG_DIR)] + + +def test_blob_hash_is_the_id_git_gives_the_file(tmp_path): + path = tmp_path / "f.py" + path.write_text("print('hi')\n", encoding="utf-8") + expected = subprocess.run( + ["git", "hash-object", str(path)], check=True, capture_output=True, text=True + ).stdout.strip() + assert blob_hash(path) == expected + assert blob_hash(tmp_path / "missing.py") is None + + +# ------------------------------------------------------------- tree identity + + +def test_tree_identity_names_the_commit_and_its_tree(tmp_path): + root = _git_repo(_build_repo(tmp_path)) + identity = tree_identity(root) + assert identity.commit == _git(root, "rev-parse", "HEAD").strip() + assert identity.tree == _git(root, "rev-parse", "HEAD^{tree}").strip() + assert identity.key == identity.tree + assert not identity.dirty + + +def test_a_re_signed_commit_with_the_same_content_is_the_same_tree(tmp_path): + """A DCO re-sign rewrote 225 commits; evidence about their content still holds.""" + root = _git_repo(_build_repo(tmp_path)) + before = tree_identity(root) + _git(root, "commit", "-q", "--amend", "-s", "-m", "seed, signed") + after = tree_identity(root) + assert before.commit != after.commit + assert before.key == after.key + + +def test_an_uncommitted_dependency_is_part_of_the_identity(tmp_path): + root = _git_repo(_build_repo(tmp_path)) + paths = dependency_paths(root, TARGET_REL) + clean = tree_identity(root, dependencies=ledger_mod.hash_paths(root, paths)) + + (root / TARGET_REL / "modeling.py").write_text("# edited\n", encoding="utf-8") + first = tree_identity(root, dependencies=ledger_mod.hash_paths(root, paths)) + (root / TARGET_REL / "modeling.py").write_text("# edited again\n", encoding="utf-8") + second = tree_identity(root, dependencies=ledger_mod.hash_paths(root, paths)) + + assert first.dirty == (f"{TARGET_REL}/modeling.py",) + assert len({clean.key, first.key, second.key}) == 3 + + +def test_outside_git_the_identity_is_unknown_rather_than_an_error(tmp_path): + identity = tree_identity(tmp_path) + assert not identity.known + assert identity.short() == "unknown" + + +# --------------------------------------------------------------- native build + + +def _lib(root: Path, name: str, payload: bytes) -> Path: + path = root / "tensorrt_llm" / "libs" / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(payload) + return path + + +def test_the_native_fingerprint_follows_content_not_mtime(tmp_path): + lib = _lib(tmp_path, "libtensorrt_llm.so", b"\x7fELF" + b"a" * 1000) + first = native_build(tmp_path) + assert first.known and len(first.libraries) == 1 + + lib.touch() # a reinstall of the same bytes + assert native_build(tmp_path).fingerprint == first.fingerprint + + lib.write_bytes(b"\x7fELF" + b"b" * 1000) + assert native_build(tmp_path).fingerprint != first.fingerprint + + +def test_a_large_library_is_sampled_and_a_change_at_either_end_is_seen(tmp_path, monkeypatch): + monkeypatch.setattr(ledger_mod, "_FULL_HASH_LIMIT", 1024) + monkeypatch.setattr(ledger_mod, "_SAMPLE_EDGE", 64) + monkeypatch.setattr(ledger_mod, "_SAMPLE_CHUNK", 16) + payload = bytearray(b"z" * 100_000) + lib = _lib(tmp_path, "libth_common.so", bytes(payload)) + base = native_build(tmp_path).fingerprint + payload[-1:] = b"q" + lib.write_bytes(bytes(payload)) + assert native_build(tmp_path).fingerprint != base + + +def test_no_native_libraries_is_an_empty_fingerprint(tmp_path): + assert not native_build(tmp_path).known + + +# ------------------------------------------------------------------ switches + + +def test_switch_defaults_are_read_from_the_code(): + source = """ +import os +A = os.environ.get("TRTLLM_A", "1") +B = os.getenv("TRTLLM_B") +C = os.environ["TRTLLM_C"] +D = os.environ.get("HOME", "/root") +E = int(os.environ.get("TRTLLM_E", 0)) +F = os.environ.get("TRTLLM_F", DEFAULT_F) +G = os.environ.setdefault("TRTLLM_G", "on") +""" + assert harvest_switch_defaults(source) == { + "TRTLLM_A": "1", + "TRTLLM_B": None, + "TRTLLM_C": "", + "TRTLLM_E": "0", + "TRTLLM_F": "", + "TRTLLM_G": "on", + } + + +def test_a_switch_read_with_two_defaults_records_both(): + source = 'import os\nx = os.environ.get("TRTLLM_X", "1")\ny = os.environ.get("TRTLLM_X", "0")\n' + assert harvest_switch_defaults(source) == {"TRTLLM_X": "0 | 1"} + + +def test_a_file_that_does_not_parse_has_no_switches(): + assert harvest_switch_defaults("def broken(:\n") == {} + + +def test_command_env_reads_every_assignment(): + command = "TRTLLM_MODELING_V2=require FOO='a b' python run.py --flag=1" + assert command_env(command) == {"TRTLLM_MODELING_V2": "require", "FOO": "a b"} + + +# ---------------------------------------------------------------------- spec + + +def _spec(**overrides): + spec = { + "checkpoint_path": "/ckpt", + "features": ["mtp", "dspark"], + "accuracy": { + "benchmark": {"default": "gsm8k", "dspark": "gsm8k_long"}, + "witness": {"dspark": {"contrast_env": {"X": "0"}}}, + "tolerance": 5.0, + }, + "environment": {"kind": "local", "python": "/py", "env": {"TRTLLM_A": "1"}}, + "remediation": {"max_rounds": 2}, + } + spec.update(overrides) + return spec + + +def test_another_units_settings_do_not_change_a_units_spec(): + base = _spec() + changed = _spec( + accuracy={ + "benchmark": {"default": "gsm8k", "dspark": "other"}, + "witness": {"dspark": {"contrast_env": {"X": "1"}}}, + "tolerance": 5.0, + } + ) + assert spec_fingerprint(base, "core") == spec_fingerprint(changed, "core") + assert spec_fingerprint(base, "mtp") == spec_fingerprint(changed, "mtp") + assert spec_fingerprint(base, "dspark") != spec_fingerprint(changed, "dspark") + + +def test_adding_a_later_feature_does_not_change_an_earlier_units_spec(): + assert spec_fingerprint(_spec(), "mtp") == spec_fingerprint( + _spec(features=["mtp", "dspark", "kv_quant"]), "mtp" + ) + + +def test_substrate_and_bookkeeping_do_not_change_the_spec_but_the_env_does(): + base = _spec() + moved = _spec( + environment={"kind": "slurm", "python": "/py", "env": {"TRTLLM_A": "1"}}, + remediation={"max_rounds": 9}, + ) + assert spec_fingerprint(base, "core") == spec_fingerprint(moved, "core") + switched = _spec(environment={"kind": "local", "python": "/py", "env": {"TRTLLM_A": "0"}}) + assert spec_fingerprint(base, "core") != spec_fingerprint(switched, "core") + + +# -------------------------------------------------------- a complete verdict + + +def test_a_verdict_carries_everything_freshness_needs(tmp_path): + root = _git_repo(_build_repo(tmp_path)) + (root / TARGET_REL / "modeling.py").write_text( + (root / TARGET_REL / "modeling.py").read_text() + + '\nimport os\nREPLAY = os.environ.get("TRTLLM_REPLAY", "1")\n', + encoding="utf-8", + ) + _git(root, "commit", "-qam", "replay switch") + _lib(root, "libtensorrt_llm.so", b"\x7fELF") + + verdict = verdict_from_result( + _result( + member="tep4", + measured=92.5, + reference=92.3, + tolerance=5.0, + command="TRTLLM_MODELING_V2=require trtllm-eval --model /ckpt gsm8k", + tree_state="abc1234", + ), + VerdictContext( + repo_root=str(root), + target_rel=TARGET_REL, + unit="core", + env={"TMPDIR": "/scratch/t", "TRTLLM_A": "1"}, + task_data=_spec(), + ), + ) + + assert verdict.tree.commit == _git(root, "rev-parse", "HEAD").strip() + assert verdict.native.known + assert verdict.dependencies[f"{TARGET_REL}/modeling.py"] == blob_hash( + root / TARGET_REL / "modeling.py" + ) + assert verdict.env == {"TRTLLM_A": "1", "TRTLLM_MODELING_V2": "require"} + assert verdict.switch_defaults["TRTLLM_REPLAY"] == "1" + assert verdict.config_fingerprint.startswith("sha256:") + assert verdict.spec_fingerprint == spec_fingerprint(_spec(), "core") + assert verdict.tree_state == "abc1234" + # Only switch values are stored in the clear; the rest are fingerprinted. + assert "TMPDIR" not in json.dumps(verdict.to_record()) + + +# -------------------------------------------------------------- the one hook + + +def test_the_gate_log_sink_sees_every_result_once(): + seen: list[str] = [] + log = GateLog(sink=lambda result: seen.append(result.name)) + log.record(_result(name="boot")) + log.record(_result(name="accuracy[gsm8k]")) + assert seen == ["boot", "accuracy[gsm8k]"] + assert [r.name for r in log.results] == seen + + +def test_a_gate_log_without_a_sink_behaves_as_before(): + log = GateLog() + assert log.record(_result()).name == "accuracy[gsm8k]" + + +def test_a_workflow_run_puts_every_gate_into_the_ledger_with_its_unit( + tmp_path, task_file, monkeypatch +): + from .test_workflow import Script, _install, _plan_then_vocabulary, _workflow + + spec = yaml.safe_load(task_file.read_text()) + spec["features"] = ["mtp"] + task_file.write_text(yaml.safe_dump(spec), encoding="utf-8") + workflow = _workflow(tmp_path) + _install(monkeypatch, workflow, Script(None)) + _plan_then_vocabulary(workflow, tmp_path, monkeypatch, []) + with workflow: + workflow.run(task_file) + + # Engine gates carry their unit; the early and delivery checks carry none. + verdicts = [v for v in Ledger.in_workspace(tmp_path / "ws").verdicts() if v.unit] + keys = [(v.gate.split("[")[0], v.unit) for v in verdicts] + assert ("boot", "core") in keys and ("accuracy", "core") in keys + assert ("boot", "mtp") in keys and ("acceptance", "mtp") in keys + assert len(verdicts) == len(workflow.gate_log.results) + core_boot = next(v for v in verdicts if v.key[:2] == ("boot", "core")) + assert f"{TARGET_REL}/modeling.py" in core_boot.dependencies + assert core_boot.config["features"] == [] + + +# ----------------------------------------------------------------------- CLI + + +def test_a_harness_records_a_verdict_through_the_cli(tmp_path, capsys): + root = _build_repo(tmp_path) + workspace = tmp_path / "ws" + code = ledger_mod.main( + [ + "record", + str(workspace), + "--repo", + str(root), + "--target-rel", + TARGET_REL, + "--gate", + "accuracy[gsm8k]", + "--unit", + "core", + "--member", + "tp4", + "--outcome", + "PASS", + "--value", + "92.1", + "--threshold", + "87.3", + "--env", + "TRTLLM_X=1", + ] + ) + assert code == 0 + (verdict,) = Ledger.in_workspace(workspace).verdicts() + assert verdict.key == ("accuracy[gsm8k]", "core", "tp4") + assert verdict.threshold == 87.3 + assert verdict.env == {"TRTLLM_X": "1"} + assert f"{TARGET_REL}/modeling.py" in verdict.dependencies + assert "recorded" in capsys.readouterr().out diff --git a/agent-flow/tests/workflows/staircase_onboard/test_report.py b/agent-flow/tests/workflows/staircase_onboard/test_report.py new file mode 100644 index 000000000000..2d7149121719 --- /dev/null +++ b/agent-flow/tests/workflows/staircase_onboard/test_report.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for reports generated from the ledger, and for the consistency check. + +The reflection on the DeepSeek-V4.1 run lists ten contradictions between its +records (its section 6). The first two are reproduced here as fixtures: + +1. a summary saying "16/16 PASS, no finding" while the final regression + table held a FINDING on ``accept:tp4``; +2. a progress page listing an A/B as pending after it had closed, under a + stale branch, target commit and status. + +Generated reports cannot make either mistake, and the check catches both +when a human writes them. +""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import pytest +import yaml + +from agent_flow.workflows.staircase_onboard import report +from agent_flow.workflows.staircase_onboard.freshness import complete_regression, record_regression +from agent_flow.workflows.staircase_onboard.gates import GateResult +from agent_flow.workflows.staircase_onboard.ledger import ( + Ledger, + VerdictContext, + native_build, + record_gate_result, + tree_identity, +) +from agent_flow.workflows.staircase_onboard.report import ( + FINAL_REPORT, + OVERLAY_FILENAME, + PROGRESS_HTML, + PROGRESS_JSON, + build_progress, + check_claims, + check_consistency, + check_items, + check_progress, + collect_facts, + headline, + narrative, + render_final_report, + write_reports, +) + +from .test_delivery import TARGET_REL, _build_repo + + +def _git(root: Path, *args: str) -> str: + return subprocess.run( + ["git", "-C", str(root), *args], check=True, capture_output=True, text=True + ).stdout + + +@pytest.fixture +def repo(tmp_path): + root = _build_repo(tmp_path) + _git(root, "init", "-q") + _git(root, "config", "user.email", "t@e.st") + _git(root, "config", "user.name", "t") + _git(root, "add", "-A") + _git(root, "commit", "-qm", "seed") + return root + + +@pytest.fixture +def workspace(tmp_path, repo): + ws = tmp_path / "ws" + ws.mkdir() + spec = { + "trtllm_repo_path": str(repo), + "features": ["dspark"], + "environment": {"kind": "local", "python": "/py", "env": {}}, + } + (ws / "task.yaml").write_text(yaml.safe_dump(spec), encoding="utf-8") + return ws + + +class _Finding: + """A gate result flagged for a human, the way an acceptance gate would flag one.""" + + name = "acceptance" + passed = True + detail = "-0.1216 against the dspark gate" + member = "tp4" + measured = 2.71 + finding = True + + +def _record(ws, repo, result, unit="core"): + return record_gate_result( + Ledger.in_workspace(ws), + result, + VerdictContext(repo_root=str(repo), target_rel=TARGET_REL, unit=unit), + ) + + +def _gate(name, member="tp4", passed=True, **kw): + kw.setdefault("detail", "ok") + return GateResult(name=name, passed=passed, member=member, **kw) + + +def _sixteen_with_a_finding(ws, repo): + """The shape of the DeepSeek-V4.1 final regression: all green but one FINDING.""" + for member in ("tep4", "dep4", "tp4"): + _record(ws, repo, _gate("boot", member)) + _record(ws, repo, _gate("accuracy[gsm8k]", member, measured=92.5, reference=92.3)) + _record(ws, repo, _Finding(), unit="dspark") + ledger = Ledger.in_workspace(ws) + regression = complete_regression(ledger.latest(), tree_identity(repo), native_build(repo)) + record_regression(ledger, regression) + + +# --------------------------------------------------------------- generation + + +def test_the_headline_is_counted_not_written(workspace, repo): + _sixteen_with_a_finding(workspace, repo) + facts = collect_facts(workspace) + text = headline(facts) + assert "7 verdict(s): 6 PASS, 1 FINDING, 0 FAIL" in text + assert "bound to HEAD" in text + assert check_claims(text, facts, where="headline") == [] + + +def test_progress_json_carries_every_fact_from_the_ledger(workspace, repo): + _sixteen_with_a_finding(workspace, repo) + facts = collect_facts(workspace) + doc = build_progress(facts, workspace=workspace) + assert doc["generated_from"]["ledger_digest"] == Ledger.in_workspace(workspace).digest() + assert len(doc["gates"]["rows"]) == 7 + assert [f["gate"] for f in doc["findings"]] == ["acceptance"] + assert doc["regression"]["binding"]["ok"] + assert check_progress(doc, facts) == [] + + +def test_a_stale_verdict_shows_as_stale_and_queued(workspace, repo): + _record(workspace, repo, _gate("boot")) + (repo / TARGET_REL / "modeling.py").write_text("# rewritten\n", encoding="utf-8") + doc = build_progress(collect_facts(workspace)) + (row,) = doc["gates"]["rows"] + assert row["fresh"] == "STALE" + assert any("code changed" in r for r in row["reasons"]) + assert doc["rerun_queue"][0]["gate"] == "boot" + assert doc["summary"]["gates"]["stale"] == 1 + + +def test_the_page_is_offline_and_escapes_what_it_shows(workspace, repo): + _record(workspace, repo, _gate("boot", detail="")) + doc = write_reports(workspace) + page = (workspace / PROGRESS_HTML).read_text() + assert "" not in page + assert "http://" not in page and "https://" not in page + assert "boot" in page and doc["meta"]["title"] in page + assert json.loads((workspace / PROGRESS_JSON).read_text())["consistency"]["ok"] + + +def test_overlay_sections_are_rendered(workspace, repo): + (workspace / OVERLAY_FILENAME).write_text( + yaml.safe_dump( + { + "meta": {"model": "Some-Model", "hardware": "4x GB300"}, + "features": [{"feature": "TP4", "status": "DONE", "notes": "core"}], + "workflow_feedback": [{"id": "F041", "title": "late checks", "status": "OPEN"}], + } + ), + encoding="utf-8", + ) + doc = write_reports(workspace) + page = (workspace / PROGRESS_HTML).read_text() + assert doc["meta"]["model"] == "Some-Model" + assert "Feature scorecard" in page and "F041" in page + + +# ----------------------------------------------- the contradictions that happened + + +def test_no_finding_over_a_finding_is_a_contradiction(workspace, repo): + """Reflection section 6, row 1.""" + _sixteen_with_a_finding(workspace, repo) + facts = collect_facts(workspace) + problems = check_claims("Final regression: 7/7 PASS, no finding.", facts, where="report") + assert [p.claim for p in problems] == ["no finding"] + assert "acceptance@tp4 [dspark]" in problems[0].fact + + +def test_n_of_n_pass_over_a_stale_verdict_is_a_contradiction(workspace, repo): + _record(workspace, repo, _gate("boot")) + _record(workspace, repo, _gate("accuracy[gsm8k]")) + (repo / TARGET_REL / "weights.py").write_text("MANIFEST = (1,)\n", encoding="utf-8") + facts = collect_facts(workspace) + problems = check_claims("All gates pass; 2/2 PASS on the final tree.", facts, where="r") + assert {p.claim for p in problems} == {"All gates pass", "2/2 PASS"} + + +def test_an_n_of_n_that_counts_neither_set_is_a_contradiction(workspace, repo): + _record(workspace, repo, _gate("boot")) + facts = collect_facts(workspace) + (problem,) = check_claims("16/16 PASS", facts, where="r") + assert "number 1" in problem.fact + + +def test_a_partial_count_is_not_a_claim_of_completeness(workspace, repo): + _record(workspace, repo, _gate("boot", passed=False)) + assert check_claims("1/2 PASS so far", collect_facts(workspace), where="r") == [] + + +def test_a_final_tree_that_is_not_the_regression_tree_is_a_contradiction(workspace, repo): + """Reflection rows 5 and 8: hashes in the report that no regression ran on.""" + _sixteen_with_a_finding(workspace, repo) + facts = collect_facts(workspace) + head = _git(repo, "rev-parse", "HEAD").strip() + assert check_claims(f"the final tree is `{head[:10]}`", facts, where="r") == [] + (problem,) = check_claims("the final tree is `e8c78be0a1`", facts, where="r") + assert "last full regression ran on" in problem.fact + + +def test_a_pending_item_that_is_in_fact_closed(workspace, repo): + """Reflection row 2: the A/B still listed as pending after it closed.""" + _sixteen_with_a_finding(workspace, repo) + _record(workspace, repo, _gate("acceptance", "tp4"), unit="dspark") + facts = collect_facts(workspace) + items = [ + { + "id": "G4", + "title": "acceptance A/B", + "status": "IN PROGRESS", + "evidence": [{"gate": "acceptance", "unit": "dspark"}], + } + ] + (problem,) = check_items(items, facts) + assert "in fact closed" in problem.fact + + +def test_a_closed_item_needs_fresh_passing_evidence(workspace, repo): + _record(workspace, repo, _gate("boot", "tp4")) + _record(workspace, repo, _gate("accuracy[gsm8k]", "tp4", passed=False)) + facts = collect_facts(workspace) + items = [ + {"id": "V1", "status": "DONE", "evidence": {"unit": "core"}}, + {"id": "V2", "status": "DONE", "evidence": {"gate": "witness*"}}, + {"id": "V3", "status": "DONE"}, # narrative only: not checked + ] + problems = check_items(items, facts) + assert [p.where for p in problems] == ["item V1", "item V2"] + assert "FAIL" in problems[0].fact + assert "matches no verdict" in problems[1].fact + + +def test_a_hand_edited_or_outdated_page_is_a_contradiction(workspace, repo): + _record(workspace, repo, _gate("boot", passed=False)) + write_reports(workspace) + doc = json.loads((workspace / PROGRESS_JSON).read_text()) + doc["gates"]["rows"][0]["result"] = "PASS" + doc["rerun_queue"].append({"gate": "accuracy[gsm8k]", "unit": "core", "member": "tp4"}) + (workspace / PROGRESS_JSON).write_text(json.dumps(doc)) + _record(workspace, repo, _gate("accuracy[gsm8k]")) # the ledger moves on + + problems = check_consistency(workspace) + facts = [p.fact for p in problems] + assert any("the page is stale" in f for f in facts) + assert any("the ledger says FAIL" in f for f in facts) + assert any("in fact closed" in f for f in facts) + + +# ------------------------------------------------------------- final report + + +def test_the_final_report_regenerates_its_tables_and_keeps_the_narrative(workspace, repo): + _record(workspace, repo, _gate("boot")) + first = render_final_report(collect_facts(workspace)) + assert "" in first + written = first + "\nThe operator's conclusion: core is ready.\n" + + _record(workspace, repo, _gate("accuracy[gsm8k]", measured=91.0, reference=92.0)) + second = render_final_report(collect_facts(workspace), written) + assert "accuracy[gsm8k]" in second + assert "The operator's conclusion: core is ready." in second + assert "accuracy[gsm8k]" not in narrative(second) + + +def test_a_region_missing_from_an_existing_report_is_appended(): + text = render_final_report(report.Facts(), "# My report\n\nprose\n") + assert text.startswith("# My report\n\nprose\n") + for name in ("summary", "gate-table", "findings", "freshness", "regression"): + assert f"" in text + + +def test_the_final_reports_narrative_is_checked(workspace, repo): + _sixteen_with_a_finding(workspace, repo) + write_reports(workspace, final_report=True) + path = workspace / FINAL_REPORT + path.write_text(path.read_text() + "\nConclusion: no finding.\n", encoding="utf-8") + (problem,) = check_consistency(workspace) + assert problem.where == FINAL_REPORT + + +def test_report_check_cli_exit_codes(workspace, repo, capsys): + _sixteen_with_a_finding(workspace, repo) + assert report.main([str(workspace)]) == 0 + assert report.main([str(workspace), "--check"]) == 0 + note = workspace / "summary.md" + note.write_text("16/16 PASS, no finding\n", encoding="utf-8") + assert report.main([str(workspace), "--check", str(note)]) == 1 + assert "CONTRADICTION" in capsys.readouterr().out + + +# ------------------------------------------------------------ in the workflow + + +def _run(tmp_path, task_file, monkeypatch, **kwargs): + from agent_flow.workflows.staircase_onboard.state import load_state + + from .test_workflow import Script, _install, _plan_then_vocabulary, _workflow + + workflow = _workflow(tmp_path, **kwargs) + _install(monkeypatch, workflow, Script(None)) + _plan_then_vocabulary(workflow, tmp_path, monkeypatch, []) + with workflow: + workflow.run(task_file) + return load_state(tmp_path / "ws" / ".staircase_onboard_state.json") + + +def test_a_run_writes_its_reports_from_the_ledger(tmp_path, task_file, monkeypatch): + state = _run(tmp_path, task_file, monkeypatch) + ws = tmp_path / "ws" + assert state.verdict == "delivered" + doc = json.loads((ws / PROGRESS_JSON).read_text()) + assert doc["meta"]["status"] == "DELIVERED" + gates = {r["gate"].split("[")[0] for r in doc["gates"]["rows"] if r["unit"]} + assert gates == {"boot", "accuracy"} + assert any(r["gate"].startswith("delivery:") for r in doc["gates"]["rows"]) + assert doc["pipeline_stages"][-1]["status"] == "DELIVERED" + assert (ws / PROGRESS_HTML).is_file() + assert "" in (ws / FINAL_REPORT).read_text() + assert check_consistency(ws) == [] + + +def test_delivery_refuses_a_report_that_contradicts_the_ledger(tmp_path, task_file, monkeypatch): + ws = tmp_path / "ws" + ws.mkdir() + (ws / OVERLAY_FILENAME).write_text( + yaml.safe_dump( + {"items": [{"id": "W1", "status": "DONE", "evidence": {"gate": "witness*"}}]} + ), + encoding="utf-8", + ) + state = _run(tmp_path, task_file, monkeypatch) + assert state.verdict == "delivery-failed:evidence" + assert "contradict the ledger" in (ws / "journal.md").read_text() + + +def test_clean_archives_the_previous_runs_evidence(tmp_path, task_file, monkeypatch): + _run(tmp_path, task_file, monkeypatch) + ws = tmp_path / "ws" + before = (ws / "ledger.jsonl").read_text() + _run(tmp_path, task_file, monkeypatch, clean=True) + archived = list((ws / "archive").glob("*-ledger.jsonl")) + assert len(archived) == 1 and archived[0].read_text() == before From b822ba630b313f98eefe8df78d456835e555a71e Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Mon, 12 Oct 2026 00:56:17 -0700 Subject: [PATCH 06/70] [TRTLLM-16460][feat] staircase: gates that catch the escaped defects: statistical accuracy, per-prompt acceptance, step witness Accuracy floors from the anchor's standard error with a cross-member finding; a per-prompt acceptance gate on a natural-text workload; a scheduler-equivalence witness with an optional step-trace hook; boot workloads declared in task.yaml and matched on generated text (a real image for multimodal); launch-path accuracy smoke through each unit's own deployment entry; token witnesses with declared workloads at every member; and `staircase-onboard gate ` to run one gate or witness standalone. Squashed from user/williamj/staircase-polish @ c8109d6f90 (10 commits): - [feat] staircase: statistical accuracy floors and a cross-member finding - [feat] staircase: per-prompt acceptance gate on a natural-text workload - [feat] staircase: scheduler-equivalence witness with an optional step-trace hook - [fix] staircase: the envelope's acceptance evidence follows the speculation predicate - [feat] staircase: boot workload declared in task.yaml, matched on generated text, with a real image for multimodal - [feat] staircase: launch-path accuracy smoke through each unit's own deployment entry - [feat] staircase: token witnesses declare their workload and refuse continuations too short to compare - [feat] staircase: `staircase-onboard gate ` runs one gate or witness standalone - [feat] staircase: feature witnesses run at every member, cheapest first - [feat] staircase: delivery re-runs the new gate kinds, and the ledger records the statistical floor Formatting and lint fixes folded in for this commit's files: - [chore] staircase: apply pre-commit formatting (4 files) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- .../workflows/staircase_onboard/cli.py | 7 + .../workflows/staircase_onboard/evidence.py | 37 +- .../workflows/staircase_onboard/gate_cli.py | 289 ++++ .../staircase_onboard/gate_records.py | 179 +++ .../workflows/staircase_onboard/gates.py | 1416 +++++++++++++---- .../staircase_onboard/harness/__init__.py | 22 + .../harness/engine_driver.py | 356 +++++ .../staircase_onboard/harness/serve_smoke.py | 366 +++++ .../staircase_onboard/harness/workload_io.py | 345 ++++ .../staircase_onboard/launch_gate.py | 416 +++++ .../workflows/staircase_onboard/stats.py | 473 ++++++ .../workflows/staircase_onboard/step_check.py | 701 ++++++++ .../staircase_onboard/task.example.yaml | 180 ++- .../staircase_onboard/task_schema.py | 602 ++++++- .../workflows/staircase_onboard/witness.py | 282 +++- .../workflows/staircase_onboard/workflow.py | 260 +-- .../staircase_onboard/test_acceptance_gate.py | 557 +++++++ .../staircase_onboard/test_boot_workload.py | 298 ++++ .../staircase_onboard/test_dspark.py | 14 +- .../staircase_onboard/test_freshness.py | 4 +- .../staircase_onboard/test_gate_cli.py | 186 +++ .../staircase_onboard/test_gate_reruns.py | 74 + .../staircase_onboard/test_gate_wiring.py | 13 +- .../workflows/staircase_onboard/test_gates.py | 165 +- .../staircase_onboard/test_launch_gate.py | 381 +++++ .../staircase_onboard/test_report.py | 2 +- .../test_statistical_floor.py | 354 +++++ .../staircase_onboard/test_step_check.py | 336 ++++ .../staircase_onboard/test_substrate_gate.py | 26 +- .../staircase_onboard/test_task_schema.py | 3 + .../test_witness_workload.py | 227 +++ .../staircase_onboard/test_workflow.py | 71 +- 32 files changed, 8074 insertions(+), 568 deletions(-) create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/gate_cli.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/gate_records.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/harness/__init__.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/harness/engine_driver.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/harness/serve_smoke.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/harness/workload_io.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/launch_gate.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/stats.py create mode 100644 agent-flow/agent_flow/workflows/staircase_onboard/step_check.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_acceptance_gate.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_boot_workload.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_gate_cli.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_gate_reruns.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_launch_gate.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_statistical_floor.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_step_check.py create mode 100644 agent-flow/tests/workflows/staircase_onboard/test_witness_workload.py diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/cli.py b/agent-flow/agent_flow/workflows/staircase_onboard/cli.py index 48e3753c872d..829cf470e4ea 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/cli.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/cli.py @@ -255,11 +255,18 @@ def _premerge(argv: list[str]) -> int: return premerge_main(argv) +def _gate(argv: list[str]) -> int: + from .gate_cli import main as gate_main + + return gate_main(argv) + + #: ``staircase-onboard ...`` for the tools that are not a run. #: A run is still ``staircase-onboard --task ...``; a first argument named #: here dispatches instead, and everything after it is the subcommand's own. SUBCOMMANDS = { "premerge": _premerge, + "gate": _gate, } diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/evidence.py b/agent-flow/agent_flow/workflows/staircase_onboard/evidence.py index 72f20dfa2c54..89b41bc9627d 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/evidence.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/evidence.py @@ -367,11 +367,14 @@ def rerunner(self, entry: RerunEntry, log: Any = None) -> Callable[[], Any] | No gates_mod.run_boot_gate( wf._gate_env(), task_data=wf.task_data, - spec=gates_mod.BootSpec.load(wf.boot_spec_path), + spec=gates_mod.resolve_boot_spec( + wf.task_data, unit, wf._features_through(unit), wf.boot_spec_path + ), runner=wf._shell_runner, features=wf._features_through(unit), unit=unit, member=member, + workdir=wf.workspace / "gates" / unit, ) ) if entry.gate.startswith("accuracy"): @@ -390,11 +393,41 @@ def rerunner(self, entry: RerunEntry, log: Any = None) -> Callable[[], Any] | No gates_mod.run_acceptance_gate( wf._gate_env(), task_data=wf.task_data, - spec=gates_mod.BootSpec.load(wf.boot_spec_path), runner=wf._shell_runner, unit=unit, features=wf._features_through(unit), member=member, + workdir=wf.workspace / "gates" / unit, + records=gates_mod.records_for(wf.workspace), + ) + ) + if entry.gate == "scheduler_equivalence": + from . import step_check + + return lambda: wf.gate_log.record( + step_check.run_scheduler_equivalence( + wf._gate_env(), + task_data=wf.task_data, + runner=wf._shell_runner, + unit=unit, + member=member, + features=wf._features_through(unit), + workdir=wf.workspace / "gates" / unit / "scheduler_equivalence", + ) + ) + if entry.gate == "launch": + from . import launch_gate + + return lambda: wf.gate_log.record( + launch_gate.run_launch_gate( + wf._gate_env(), + task_data=wf.task_data, + reference=launch_gate.smoke_reference(wf.task_data, unit, wf._model_key), + runner=wf._shell_runner, + unit=unit, + member=member, + features=wf._features_through(unit), + workdir=wf.workspace / "gates" / unit / "launch", ) ) if unit in WITNESS_GATED_UNITS: diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/gate_cli.py b/agent-flow/agent_flow/workflows/staircase_onboard/gate_cli.py new file mode 100644 index 000000000000..5afd5cd54561 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/gate_cli.py @@ -0,0 +1,289 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +r"""``staircase-onboard gate ``: run one gate or witness exactly as the orchestrator does. + +An assembler that "runs the witness exactly as the workflow runs it" used to +reconstruct the command, and passed where the orchestrator's own run of the +same witness failed: the reconstruction ran from another directory. This +calls the workflow's own gate functions, with the environment, runner and +flags resolved from the same task.yaml, so there is nothing to reconstruct. + +It is also how a gate is replayed on an arbitrary TensorRT-LLM tree (a +historical commit, a worktree): ``--tree`` measures that tree instead of +``trtllm_repo_path`` -- every command runs from inside it with it first on +``PYTHONPATH`` -- and ``--results`` keeps the per-tree measurements that the +acceptance and performance gates compare a later tree against:: + + staircase-onboard gate scheduler-equivalence --task task.yaml \ + --unit dspark --member tep4 --tree .dev-worktrees/pre-fix + +Exit status: 0 when the gate passed, 1 when it failed, 2 on a usage or spec +error. ``--json`` prints the verdict as JSON instead of its summary line. +""" + +from __future__ import annotations + +import argparse +import json +import shlex +import sys +from dataclasses import asdict, replace +from pathlib import Path +from typing import Any, Callable, Mapping + +from . import gates as gates_mod +from .environment import RunEnvironment, resolve_environment +from .execution import select_command_executor, shell_runner_for +from .task_schema import ( + TaskSchemaError, + assembly_units, + load_and_validate_task_yaml, + members_for, + parse_parallel_segment, +) + +#: The gates and witnesses this command runs, by the name it takes. +GATE_NAMES: tuple[str, ...] = ( + "boot", + "accuracy", + "acceptance", + "scheduler-equivalence", + "launch", + "witness", +) + + +def _parse(argv: list[str]) -> argparse.Namespace: + parser = argparse.ArgumentParser( + prog="staircase-onboard gate", + description="Run one staircase gate or witness standalone, as the orchestrator runs it.", + ) + parser.add_argument("name", choices=GATE_NAMES, help="the gate or witness to run") + parser.add_argument("--task", type=Path, required=True, help="the run's task.yaml") + parser.add_argument( + "--unit", + default="core", + help="the assembly unit whose configuration to gate (default core); a witness " + "needs the unit whose feature it witnesses", + ) + parser.add_argument( + "--member", + default=None, + help="the parallel configuration to serve at (default: the unit's cheapest member)", + ) + parser.add_argument( + "--tree", + type=Path, + default=None, + help="measure this TensorRT-LLM tree instead of task.yaml's trtllm_repo_path: commands " + "run from inside it, with it first on PYTHONPATH", + ) + parser.add_argument( + "--python", default=None, help="the interpreter to run under (default: task.yaml's)" + ) + parser.add_argument( + "--workdir", + type=Path, + default=Path("staircase-gates"), + help="where the gate writes the files its job reads; on the shared filesystem when the " + "job runs elsewhere (a spool session). Default ./staircase-gates", + ) + parser.add_argument( + "--results", + type=Path, + default=None, + help="the results file the acceptance and performance gates compare against and append " + "to (default /gate_results.jsonl); share one across trees to compare them", + ) + parser.add_argument( + "--boot-yaml", + type=Path, + default=None, + help="a workspace boot.yaml, used when task.yaml declares no accuracy.boot", + ) + parser.add_argument( + "--features", + default=None, + help="comma-separated features the target implements (default: the unit's and every " + "earlier unit's)", + ) + parser.add_argument( + "--devices", default=None, help="comma-separated device ids (default 0..world_size-1)" + ) + parser.add_argument("--world-size", type=int, default=None) + parser.add_argument("--timeout", type=int, default=None, help="seconds, per leg where legged") + parser.add_argument( + "--mode", + choices=(gates_mod.MODE_REQUIRE, gates_mod.MODE_OFF), + default=gates_mod.MODE_REQUIRE, + help="boot only: off runs the same command on the stock path, as a control", + ) + parser.add_argument("--allow-dirty", action="store_true", help="measure a dirty target path") + parser.add_argument("--json", action="store_true", help="print the verdict as JSON") + return parser.parse_args(argv) + + +def features_for(task_data: Mapping[str, Any], unit: str) -> list[str]: + """The features a target implements once ``unit`` is assembled.""" + units = assembly_units(task_data) + if unit not in units: + return [str(f) for f in task_data.get("features") or ()] + return [str(u) for u in units[1 : units.index(unit) + 1]] + + +def _with_tree(task_data: dict[str, Any], environment: RunEnvironment, tree: Path | None): + """Point the task and the environment at ``tree``.""" + if tree is None: + return task_data, environment + root = str(tree.resolve()) + task_data = {**task_data, "trtllm_repo_path": root} + inherited = environment.env.get("PYTHONPATH") + path = f"{root}:{inherited}" if inherited else root + return task_data, replace(environment, env={**environment.env, "PYTHONPATH": path}) + + +def _in_tree(runner: Callable[..., Any], tree: Path | None) -> Callable[..., Any]: + """Run every command from inside ``tree``. + + An MPI worker resolves ``tensorrt_llm`` from its working directory ahead + of ``PYTHONPATH``, so a job started elsewhere measures another checkout on + every rank but the first. + """ + if tree is None: + return runner + prefix = f"cd {shlex.quote(str(tree.resolve()))} && " + + def run(command: str, env: Mapping[str, str] | None = None, timeout: int = 1800): + return runner(prefix + command, env, timeout) + + return run + + +def run_gate(args: argparse.Namespace) -> gates_mod.GateResult: + """Resolve everything from the task and run the named gate.""" + from . import launch_gate, step_check + from . import witness as witness_mod + from .preflight import resolve_model_key + + task_data = load_and_validate_task_yaml(args.task) + unit = args.unit + world = args.world_size + member = args.member + if member is None: + members = members_for(task_data, unit, world or 1) + member = gates_mod.order_members_by_cost(members)[0] + if world is None: + world = parse_parallel_segment(member)[1] or 1 + devices = ( + tuple(int(d) for d in args.devices.split(",")) if args.devices else tuple(range(world)) + ) + environment = resolve_environment(task_data, devices=devices, world_size=world, workload="gate") + if args.python: + environment = replace(environment, python=args.python) + task_data, environment = _with_tree(task_data, environment, args.tree) + executor = select_command_executor(environment, name_prefix=f"gate-{args.name}") + runner = _in_tree(shell_runner_for(executor), args.tree) + features = ( + [f for f in args.features.split(",") if f] + if args.features is not None + else features_for(task_data, unit) + ) + workdir = args.workdir.resolve() + records = gates_mod.GateRecords(args.results or workdir / gates_mod.RESULTS_FILENAME) + timeout = {"timeout": args.timeout} if args.timeout else {} + common = dict(task_data=task_data, runner=runner, unit=unit, member=member) + + if args.name == "boot": + return gates_mod.run_boot_gate( + environment, + spec=gates_mod.resolve_boot_spec(task_data, unit, features, args.boot_yaml), + features=features, + workdir=workdir / unit, + mode=args.mode, + allow_dirty=args.allow_dirty or None, + **timeout, + **common, + ) + if args.name == "accuracy": + reference = gates_mod.resolve_reference( + task_data, model_key=resolve_model_key(task_data, unit), unit=unit + ) + return gates_mod.run_accuracy_gate( + environment, + reference=reference, + allow_dirty=args.allow_dirty or None, + **timeout, + **common, + ) + if args.name == "acceptance": + return gates_mod.run_acceptance_gate( + environment, + features=features, + workdir=workdir / unit, + records=records, + allow_dirty=args.allow_dirty or None, + **timeout, + **common, + ) + if args.name == "scheduler-equivalence": + return step_check.run_scheduler_equivalence( + environment, + features=features, + workdir=workdir / unit / "scheduler_equivalence", + allow_dirty=args.allow_dirty or None, + **timeout, + **common, + ) + if args.name == "launch": + return launch_gate.run_launch_gate( + environment, + reference=launch_gate.smoke_reference( + task_data, unit, lambda scope: resolve_model_key(task_data, scope) + ), + features=features, + workdir=workdir / unit / "launch", + allow_dirty=args.allow_dirty or None, + **timeout, + **common, + ) + if unit not in witness_mod.WITNESS_KINDS: + raise TaskSchemaError( + f"unit {unit!r} has no witness; the witnessed features are " + f"{sorted(witness_mod.WITNESS_KINDS)}" + ) + return witness_mod.run_declared_witness( + environment, + features=features, + workdir=workdir / "witness" / unit, + boot=gates_mod.resolve_boot_spec(task_data, unit, None, args.boot_yaml), + **common, + ) + + +def invocation(task: str | Path, workdir: str | Path, unit: str) -> str: + """The command an agent uses to rerun one of the orchestrator's gates by hand.""" + return ( + f"staircase-onboard gate <{'|'.join(GATE_NAMES)}> --task {task} --unit {unit} " + f"--member --workdir {workdir}" + ) + + +def main(argv: list[str] | None = None) -> int: + """Run one gate; print its verdict; return 0 on a pass, 1 on a failure, 2 on an error.""" + args = _parse(list(sys.argv[1:] if argv is None else argv)) + try: + result = run_gate(args) + except (TaskSchemaError, KeyError, FileNotFoundError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + if args.json: + print(json.dumps(asdict(result), indent=1, default=str)) + else: + print(result.summary()) + if not result.passed and result.stdout_tail.strip(): + print(result.stdout_tail) + return 0 if result.passed else 1 + + +__all__ = ["GATE_NAMES", "features_for", "invocation", "main", "run_gate"] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/gate_records.py b/agent-flow/agent_flow/workflows/staircase_onboard/gate_records.py new file mode 100644 index 000000000000..3015827c4fa8 --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/gate_records.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""What a gate measured, on which tree, kept for the next tree to be compared against. + +Three gates judge a tree against an EARLIER one rather than against a fixed +number: the acceptance gate (per-prompt acceptance, paired), the performance +gate (one prefill chunk and one decode step) and the accuracy gate's +cross-member check, whose finding has to outlive the unit that raised it so +delivery can refuse it. They need somewhere to put a measurement, and this is +the smallest thing that works: an append-only JSON-lines file, one record per +line, never rewritten. + +The interface is deliberately three calls -- :meth:`GateRecords.append`, +:meth:`GateRecords.latest` and :func:`open_findings` -- so a fuller results +ledger can replace the file behind them without touching a gate. + +``key`` is what makes two records comparable: the unit, the member and a +fingerprint of the configuration and workload. Two records with the same key +measured the same thing on (possibly) different trees. +""" + +from __future__ import annotations + +import datetime +import json +import os +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any, Mapping + +#: The file's name in a run's workspace. +RESULTS_FILENAME = "gate_results.jsonl" + +KIND_ACCEPTANCE = "acceptance" +KIND_PERF = "perf" +KIND_FINDING = "finding" + + +@dataclass(frozen=True) +class GateRecord: + """One measurement, or one finding.""" + + kind: str + key: Mapping[str, Any] + #: ``git describe --always`` of the tree measured. + tree: str + passed: bool + values: Mapping[str, Any] = field(default_factory=dict) + #: The full commit, when git could answer; a commit range starts here. + commit: str | None = None + recorded_at: str = "" + + def to_json(self) -> str: + """One line, keys sorted, so two runs' files compare line by line.""" + return json.dumps(asdict(self), sort_keys=True, ensure_ascii=True) + + @classmethod + def from_mapping(cls, data: Mapping[str, Any]) -> "GateRecord": + """The inverse of :meth:`to_json`; unknown keys are ignored.""" + return cls( + kind=str(data.get("kind", "")), + key=dict(data.get("key") or {}), + tree=str(data.get("tree", "")), + passed=bool(data.get("passed", False)), + values=dict(data.get("values") or {}), + commit=data.get("commit"), + recorded_at=str(data.get("recorded_at", "")), + ) + + +def _now() -> str: + return datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds") + + +def _same_key(a: Mapping[str, Any], b: Mapping[str, Any]) -> bool: + return json.dumps(dict(a), sort_keys=True) == json.dumps(dict(b), sort_keys=True) + + +class GateRecords: + """An append-only JSON-lines file of :class:`GateRecord`.""" + + def __init__(self, path: str | Path) -> None: + self.path = Path(path) + + def append(self, record: GateRecord) -> GateRecord: + """Write ``record`` as one line, stamping the time when it has none. + + One ``write`` of one line on a file opened for append, so two writers + never interleave inside a record. + """ + if not record.recorded_at: + record = GateRecord(**{**asdict(record), "recorded_at": _now()}) + self.path.parent.mkdir(parents=True, exist_ok=True) + line = record.to_json() + "\n" + fd = os.open(self.path, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o644) + try: + os.write(fd, line.encode("utf-8")) + finally: + os.close(fd) + return record + + def all(self, kind: str | None = None) -> list[GateRecord]: + """Every readable record, oldest first. A torn or foreign line is skipped.""" + if not self.path.is_file(): + return [] + records: list[GateRecord] = [] + for line in self.path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line: + continue + try: + data = json.loads(line) + except ValueError: + continue + if not isinstance(data, dict): + continue + record = GateRecord.from_mapping(data) + if kind is None or record.kind == kind: + records.append(record) + return records + + def latest( + self, + kind: str, + key: Mapping[str, Any], + *, + passed_only: bool = True, + exclude_tree: str | None = None, + ) -> GateRecord | None: + """The newest record of ``kind`` with exactly ``key``. + + ``passed_only`` (the default) skips failed measurements: a regression + that failed its gate must not become the baseline the next tree is + judged against, or a second tree would pass by matching it. + ``exclude_tree`` skips records of the tree being measured now, so a + re-run of one tree is not compared with itself. + """ + for record in reversed(self.all(kind)): + if passed_only and not record.passed: + continue + if exclude_tree is not None and record.tree == exclude_tree: + continue + if _same_key(record.key, key): + return record + return None + + +def open_findings( + records: GateRecords, acknowledged: Mapping[str, str] | None = None +) -> list[GateRecord]: + """Findings whose newest record is still unresolved and not acknowledged. + + A finding is keyed (``{"finding": ..., "unit": ...}``); a later record + with the same key and ``passed`` true resolves it -- the unit was gated + again and the shift was gone. ``acknowledged`` is the CURRENT task's + ``accuracy.floor.acknowledge`` mapping, read at delivery, so adding a + reason to task.yaml and relaunching is what clears one. + """ + newest: dict[str, GateRecord] = {} + for record in records.all(KIND_FINDING): + newest[json.dumps(dict(record.key), sort_keys=True)] = record + reasons = acknowledged or {} + return [ + record + for record in newest.values() + if not record.passed and str(record.key.get("unit")) not in reasons + ] + + +__all__ = [ + "GateRecord", + "GateRecords", + "KIND_ACCEPTANCE", + "KIND_FINDING", + "KIND_PERF", + "RESULTS_FILENAME", + "open_findings", +] diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/gates.py b/agent-flow/agent_flow/workflows/staircase_onboard/gates.py index 02b0ef8335b2..01d22d5ba21a 100644 --- a/agent-flow/agent_flow/workflows/staircase_onboard/gates.py +++ b/agent-flow/agent_flow/workflows/staircase_onboard/gates.py @@ -17,7 +17,9 @@ catastrophes, not accuracy. Runs after every assembler turn. 2. **accuracy** — the release criterion. ``trtllm-eval`` against the anchor in TensorRT-LLM's own accuracy references. One-sided: - ``measured >= reference - tolerance``. Runs at a unit boundary. + ``measured >= reference - k * SE``, the SE binomial over the benchmark's + sample count (``accuracy.floor``; a fixed ``tolerance`` is still + available as an explicit rule). Runs at a unit boundary. 3. **acceptance** — required only when a feature is *distribution-preserving by construction*, speculative decoding above all. Rejection sampling holds the emitted distribution to the target model's, so a miscomputed @@ -36,6 +38,7 @@ from __future__ import annotations +import hashlib import json import os import re @@ -43,7 +46,7 @@ import subprocess import tempfile import time -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path from typing import Any, Callable, Mapping, Sequence @@ -56,9 +59,21 @@ SpoolCommandExecutor, ) +from . import stats from .delivery import MODELING_V2_PACKAGE from .environment import RunEnvironment +from .gate_records import ( + KIND_ACCEPTANCE, + KIND_FINDING, + RESULTS_FILENAME, + GateRecord, + GateRecords, + open_findings, +) +from .harness import workload_io from .task_schema import ( + BOOT_KEYS, + BOOT_RENDERS, PARALLEL_UNIFIED, TaskSchemaError, derive_family, @@ -105,6 +120,13 @@ # ``MMLU weighted average accuracy: 74.98 (14042)`` — the MMLU path has its # own reporting because it does not go through lm-eval. _MMLU_SCORE = re.compile(r"MMLU weighted average accuracy:\s*([0-9]+(?:\.[0-9]+)?)") +# 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+))?$") # ``Spec-dec stats: AL (decoded tokens/step, iteration-weighted) 3.351 (...)`` _ACCEPTANCE_LENGTH = re.compile(r"Spec-dec stats:\s*AL\b[^)]*\)\s*([0-9]+(?:\.[0-9]+)?)") # trtllm-bench reports the same quantity under its own name. @@ -383,6 +405,26 @@ class GateResult: #: cannot say that one prompt ended in EOS after two tokens while another #: accepted at the ceiling; this rides in the journal line so it can. breakdown: str = "" + #: The accuracy floor the measurement was held to, and the binomial + #: standard error and sample count it was derived from. Every accuracy + #: verdict carries the SE when the count is known, whichever rule set + #: the floor, so a reader can see how many SE a margin is. + floor: float | None = None + se: float | None = None + samples: int | None = None + #: How ``floor`` was derived, verbatim (see :meth:`.stats.Floor.describe`). + floor_rule: str = "" + #: Things the verdict wants a reader to see that do not change it -- a + #: floor not backed by a measurement, a comparison that had no baseline. + warnings: tuple[str, ...] = () + #: Whether the summary line carries ``detail`` even on a pass: true for a + #: verdict that rests on more than its headline numbers. + explained: bool = False + + @property + def threshold(self) -> float | None: + """The number ``measured`` was held to, where the gate states one (the ledger reads it).""" + return self.floor def summary(self) -> str: """One line for the journal, carrying the numbers verbatim. @@ -403,12 +445,25 @@ def summary(self) -> str: if self.dirty_override: tree = f" [tree {self.tree_state}; DIRTY TARGET PATH measured by override]" breakdown = f" [{self.breakdown}]" if self.breakdown else "" + if self.warnings: + breakdown += "".join(f" [WARNING: {w}]" for w in self.warnings) + # The numbers lead; the reasoning follows whenever the verdict is a + # failure or rests on more than the numbers (a paired comparison, an + # HTTP status), so the journal line says why without the log. + why = f" -- {self.detail}" if self.detail and (self.explained or not self.passed) else "" + if self.measured is not None and self.floor is not None: + return ( + f"[{verdict}] {label}: measured {self.measured:.4f} vs floor " + f"{self.floor:.4f} [{self.floor_rule}] (margin " + f"{self.measured - self.floor:+.4f}){why}{breakdown} in " + f"{self.duration_s:.0f}s{tree}" + ) if self.measured is not None and self.reference is not None: margin = self.measured - (self.reference - (self.tolerance or 0.0)) return ( f"[{verdict}] {label}: measured {self.measured:.4f} vs " f"reference {self.reference:.4f} - tol {self.tolerance or 0.0:.2f} " - f"(margin {margin:+.4f}){breakdown} in {self.duration_s:.0f}s{tree}" + f"(margin {margin:+.4f}){why}{breakdown} in {self.duration_s:.0f}s{tree}" ) return f"[{verdict}] {label}: {self.detail}{breakdown} ({self.duration_s:.0f}s){tree}" @@ -429,6 +484,56 @@ def verdict(**fields: Any) -> GateResult: return verdict +#: The image a multimodal unit's boot sends when the task declares none: a +#: photograph shipped in the repository's own test assets (a git-lfs file), +#: resolved against the tree under test. Never downloaded. +DEFAULT_BOOT_IMAGE = "tests/integration/test_input_files/merlion.png" +#: What the default image prompt asks, and the words any correct description +#: of that photograph (the Merlion fountain in Singapore) contains one of. +DEFAULT_BOOT_IMAGE_PROMPT = "Describe this image in one sentence." +DEFAULT_BOOT_IMAGE_KEYWORDS: tuple[str, ...] = ( + "merlion", + "lion", + "statue", + "fountain", + "singapore", + "sculpture", + "water", +) + + +@dataclass(frozen=True) +class BootPrompt: + """One boot prompt with checks of its own, images included.""" + + prompt: str + #: Every one must appear in THIS prompt's generated text (case-insensitive). + keywords: tuple[str, ...] = () + #: At least one must appear in it (case-insensitive). + keywords_any: tuple[str, ...] = () + images: tuple[str, ...] = () + render: str = "loader" + + @classmethod + def from_mapping(cls, data: Mapping[str, Any]) -> "BootPrompt": + """Read one entry of a boot ``prompts`` list written as a mapping.""" + return cls( + prompt=str(data.get("prompt", "")), + keywords=tuple(str(k) for k in data.get("keywords") or ()), + keywords_any=tuple(str(k) for k in data.get("keywords_any") or ()), + images=tuple(str(i) for i in data.get("images") or ()), + render=str(data.get("render") or "loader"), + ) + + def missing(self, text: str) -> list[str]: + """What this prompt's checks did not find in ``text``.""" + lowered = (text or "").lower() + absent = [k for k in self.keywords if k.lower() not in lowered] + if self.keywords_any and not any(k.lower() in lowered for k in self.keywords_any): + absent.append(f"any of {list(self.keywords_any)}") + return absent + + @dataclass(frozen=True) class BootSpec: """What the boot gate sends and what it insists on seeing back. @@ -439,6 +544,12 @@ class BootSpec: not yet seen this model generate anything has no business asserting what it will say. Until then the gate asserts the weaker, still useful thing: the engine cold-starts and produces non-empty output. + + ``prompts`` and ``keywords`` are the plain form: each keyword must appear + in SOME prompt's generated text. ``items`` are prompts with checks of + their own, and the only ones that carry images. Keywords are matched + against the generated texts only, never against the log: ``72`` matches + ``CPUs 72-143`` in an engine log line. """ prompts: tuple[str, ...] = ( @@ -447,21 +558,120 @@ class BootSpec: ) keywords: tuple[str, ...] = () max_tokens: int = 32 + items: tuple[BootPrompt, ...] = () + #: Where the spec came from, named in a verdict. + source: str = "the default boot prompts" + + @classmethod + def from_mapping(cls, data: Mapping[str, Any], source: str) -> "BootSpec": + """A spec from a boot block: ``prompts`` (strings or mappings), ``keywords``, ``max_tokens``.""" + plain: list[str] = [] + items: list[BootPrompt] = [] + for entry in data.get("prompts") or (): + if isinstance(entry, Mapping): + items.append(BootPrompt.from_mapping(entry)) + else: + plain.append(str(entry)) + if data.get("prompts_file"): + plain.extend(load_prompts_file(data["prompts_file"])) + return cls( + prompts=tuple(plain) if (plain or items) else cls.prompts, + keywords=tuple(str(k) for k in data.get("keywords") or ()), + max_tokens=int(data.get("max_tokens") or cls.max_tokens), + items=tuple(items), + source=source, + ) @classmethod def load(cls, path: str | Path) -> "BootSpec": + """A workspace ``boot.yaml``; the default spec when there is none.""" source = Path(path) if not source.is_file(): return cls() data = yaml.safe_load(source.read_text(encoding="utf-8")) or {} if not isinstance(data, dict): raise ValueError(f"{source} must hold a YAML mapping, got {type(data).__name__}") - prompts = tuple(str(p) for p in data.get("prompts") or cls.prompts) - return cls( - prompts=prompts or cls.prompts, - keywords=tuple(str(k) for k in data.get("keywords") or ()), - max_tokens=int(data.get("max_tokens", cls.max_tokens)), + return cls.from_mapping(data, str(source)) + + @property + def has_images(self) -> bool: + """Whether any prompt sends an image.""" + return any(item.images for item in self.items) + + def requests(self, repo_root: str | Path | None) -> list[dict[str, Any]]: + """Engine-driver requests: the plain prompts, then the items. + + A relative image path is resolved against ``repo_root``, the tree + under test, so a repository test asset names the same file on every + checkout. + """ + + def image_path(name: str) -> str: + path = Path(name) + if path.is_absolute() or repo_root is None: + return str(path) + return str(Path(repo_root) / path) + + requests = [ + {"id": f"b{n}", "prompt": prompt, "max_tokens": self.max_tokens} + for n, prompt in enumerate(self.prompts) + ] + for n, item in enumerate(self.items): + request: dict[str, Any] = { + "id": f"i{n}", + "prompt": item.prompt, + "max_tokens": self.max_tokens, + } + if item.images: + request["images"] = [image_path(i) for i in item.images] + request["render"] = item.render + requests.append(request) + return requests + + +def resolve_boot_spec( + task_data: Mapping[str, Any], + unit: str | None, + features: Sequence[str] | None = None, + fallback: str | Path | None = None, +) -> BootSpec: + """The boot workload ``unit`` runs, declared in task.yaml. + + ``accuracy.boot`` (run-wide keys, with a per-unit entry layered over + them) when task.yaml declares it, else the workspace ``boot.yaml`` at + ``fallback``, else the default. A unit whose target implements + ``multimodal`` (and whose scope is not ``text_only``) and whose spec sends + no image gets the default image prompt appended: a multimodal boot that + sends only text never exercises the vision path, and the first gate that + does is an hour-long benchmark. + """ + accuracy = task_data.get("accuracy") or {} + block = accuracy.get("boot") if isinstance(accuracy, Mapping) else None + if isinstance(block, Mapping): + merged = {k: v for k, v in block.items() if k in BOOT_KEYS} + chosen = block.get(unit) if unit is not None else None + where = "accuracy.boot" + if isinstance(chosen, Mapping): + merged.update(chosen) + where = f"accuracy.boot.{unit}" + spec = BootSpec.from_mapping(merged, where) + elif fallback is not None: + spec = BootSpec.load(fallback) + else: + spec = BootSpec() + implemented = {str(f) for f in (features or ())} + if ( + "multimodal" in implemented + and task_data.get("scope", "full") != "text_only" + and not spec.has_images + ): + image = BootPrompt( + prompt=DEFAULT_BOOT_IMAGE_PROMPT, + keywords_any=DEFAULT_BOOT_IMAGE_KEYWORDS, + images=(DEFAULT_BOOT_IMAGE,), ) + spec = replace(spec, items=(*spec.items, image)) + return spec def _default_shell_runner( @@ -1026,6 +1236,183 @@ def _with_mode(environment: RunEnvironment, mode: str) -> RunEnvironment: ) +# ------------------------------------------------------------ the engine driver +# +# The gates that need per-request structure -- token ids, whole texts, +# iterations, images -- run `harness/engine_driver.py` under the tree's own +# interpreter. It builds the engine with the tree's own quickstart from the +# same argv the quickstart gates use, so the configuration measured is +# identical, and prints one `STAIRCASE_RECORD` JSON line per request. A gate +# reads those lines and nothing else of the log. + +#: The driver script, run as a file under the measured tree's interpreter. +ENGINE_DRIVER = Path(__file__).resolve().parent / "harness" / "engine_driver.py" + +#: The quickstart flag a declared batch limit replaces. The runtime block +#: usually sets it through ``extra_flags``, so an override has to remove that +#: spelling, not just append another one. +BATCH_SIZE_FLAG = "--max_batch_size" + + +def with_max_batch_size(flags: Sequence[str], max_batch_size: int | None) -> list[str]: + """``flags`` with every ``--max_batch_size`` spelling replaced by one ``=N``. + + A token-comparing gate declares the batch limit its legs run at, and the + runtime block has usually set its own through ``extra_flags`` -- + ``--max_batch_size=8``, ``--max_batch_size 8`` as one string, or the flag + and its value as two. All three are removed, so the argv states ONE limit. + ``None`` returns ``flags`` unchanged. + """ + if max_batch_size is None: + return list(flags) + kept: list[str] = [] + skip_value = False + for flag in flags: + if skip_value: + skip_value = False + continue + word = str(flag) + if word == BATCH_SIZE_FLAG: + skip_value = True + continue + if word.startswith((f"{BATCH_SIZE_FLAG}=", f"{BATCH_SIZE_FLAG} ")): + continue + kept.append(word) + return [*kept, f"{BATCH_SIZE_FLAG}={int(max_batch_size)}"] + + +def quickstart_script(task_data: Mapping[str, Any]) -> Path: + """The measured tree's ``examples/llm-api/quickstart_advanced.py``.""" + return ( + Path(str(task_data["trtllm_repo_path"])) / "examples" / "llm-api" / "quickstart_advanced.py" + ) + + +def quickstart_words( + task_data: Mapping[str, Any], + environment: RunEnvironment, + *, + unit: str | None, + features: Any | None, + max_tokens: int, + member: str | None, + max_batch_size: int | None = None, + extra: Sequence[str] = (), + runtime_overrides: Mapping[str, Any] | None = None, +) -> list[str]: + """The quickstart argv as a list of plain words, for the engine driver's spec. + + The same flags the quickstart gates emit -- topology, the unit's runtime + block, the speculator at the unit's draft depth -- with each runtime + ``extra_flags`` entry split as the shell would split it. + ``runtime_overrides`` replaces switches of the unit's runtime block (a + leg that turns the overlap scheduler off, for one). + """ + runtime = {**resolve_runtime(task_data, unit), **dict(runtime_overrides or {})} + flags = with_max_batch_size( + [ + *runtime_flags(runtime), + *speculation_flags( + task_data.get("features") if features is None else features, + max_draft_len=resolve_max_draft_len(task_data.get("accuracy") or {}, unit), + ), + *extra, + ], + max_batch_size, + ) + words = [ + f"--model_dir={task_data['checkpoint_path']}", + f"--max_tokens={int(max_tokens)}", + *topology_flags(str(task_data["parallel"]), environment.world_size, member), + ] + for flag in flags: + words.extend(shlex.split(str(flag))) + return words + + +def gate_workdir( + workdir: str | Path | None, name: str, task_data: Mapping[str, Any] | None = None +) -> Path: + """Where a gate writes the files its job reads. + + ``workdir`` when given: the orchestrator passes a directory on the shared + filesystem, because the job may run on another host (a spool session). + Otherwise a directory under the temp dir named for the gate and the task + (its tree and checkpoint), so two runs of one gate on one task write the + same files and their commands compare equal; a fresh one without a task. + """ + if workdir is not None: + directory = Path(workdir) + elif task_data is not None: + key = f"{task_data.get('trtllm_repo_path')}|{task_data.get('checkpoint_path')}" + digest = hashlib.sha256(key.encode("utf-8")).hexdigest()[:12] + directory = Path(tempfile.gettempdir()) / f"staircase-{name}-{digest}" + else: + return Path(tempfile.mkdtemp(prefix=f"staircase-{name}-")) + directory.mkdir(parents=True, exist_ok=True) + return directory + + +def write_driver_spec( + path: str | Path, + *, + task_data: Mapping[str, Any], + argv: Sequence[str], + requests: Sequence[Mapping[str, Any]] = (), + dataset: Mapping[str, Any] | None = None, + max_tokens: int = 64, + concurrency: int = 1, + ignore_eos: bool = False, + perf: Mapping[str, Any] | None = None, +) -> Path: + """Write one engine-driver spec (see ``harness/engine_driver.py``).""" + spec = { + "quickstart": str(quickstart_script(task_data)), + "quickstart_argv": [str(a) for a in argv], + "requests": [dict(r) for r in requests], + "dataset": dict(dataset) if dataset else None, + "max_tokens": int(max_tokens), + "concurrency": int(concurrency), + "ignore_eos": bool(ignore_eos), + "perf": dict(perf) if perf else None, + } + target = Path(path) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(json.dumps(spec, indent=1, ensure_ascii=False), encoding="utf-8") + return target + + +def engine_driver_command( + environment: RunEnvironment, spec_path: str | Path, quickstart: str | Path | None = None +) -> str: + """The launch of the engine driver over ``spec_path``, as engine work. + + The quickstart and its argv are spelled on the command line as well as in + the spec (the driver takes the command line's), so a command read from a + log says which tree it built from and with which flags. The prompts stay + in the spec: a long one would not fit an argv string. + """ + argv: list[str] = [] + try: + spec = json.loads(Path(spec_path).read_text(encoding="utf-8")) + except (OSError, ValueError): + spec = {} + if quickstart is None: + quickstart = spec.get("quickstart") + argv = [str(a) for a in spec.get("quickstart_argv") or ()] + named = f" --quickstart {shlex.quote(str(quickstart))}" if quickstart else "" + flags = f" -- {' '.join(shlex.quote(a) for a in argv)}" if argv else "" + return environment.python_command( + f"{shlex.quote(str(ENGINE_DRIVER))} --spec {shlex.quote(str(spec_path))}{named}{flags}" + ) + + +def parse_driver_records(output: str) -> list[dict[str, Any]]: + """Every ``STAIRCASE_RECORD`` the engine driver printed, in request order.""" + records = workload_io.parse_marked(output, workload_io.RECORD_MARKER) + return sorted(records, key=lambda r: int(r.get("index", 0))) + + def parse_eval_score(stdout: str) -> float | None: """Extract the final accuracy from a ``trtllm-eval`` run's output. @@ -1052,6 +1439,17 @@ class AcceptanceReading: #: output carried only a summary statistic, which says nothing about how #: much work it was estimated from. requests: tuple[tuple[int, int], ...] = () + #: A digest of each request's prompt, parallel to ``requests``, where the + #: run printed one (the engine driver does). What pairs a prompt's rate + #: with the same prompt's rate in another leg or on another tree. + prompts: tuple[str, ...] = () + + def rates(self) -> dict[str, float]: + """Accepted tokens per iteration, per prompt, keyed by prompt digest.""" + names = self.prompts if len(self.prompts) == len(self.requests) else () + return { + (names[n] if names else str(n)): t / i for n, (t, i) in enumerate(self.requests) if i + } @property def draft_iterations(self) -> int | None: @@ -1080,7 +1478,10 @@ def breakdown(self) -> str: def parse_acceptance(stdout: str) -> AcceptanceReading | None: - """Extract the acceptance measurement from an eval, bench or quickstart run. + """Extract the acceptance measurement from a driver, eval, bench or quickstart run. + + The engine driver's records come first: one per prompt, with the prompt's + digest, so a rate can be paired with the same prompt's rate elsewhere. The quickstart form is **iteration-weighted across requests**: the quantity is tokens-per-iteration, so several requests combine as @@ -1091,6 +1492,16 @@ def parse_acceptance(stdout: str) -> AcceptanceReading | None: :meth:`AcceptanceReading.breakdown`, because it is the number that shows a prompt that never drafted. """ + records = parse_driver_records(stdout) + if records: + parts = tuple((int(r.get("tokens", 0)), int(r.get("iterations", 0))) for r in records) + iterations = sum(i for _, i in parts) + if iterations: + return AcceptanceReading( + value=sum(t for t, _ in parts) / iterations, + requests=parts, + prompts=tuple(str(r.get("prompt_sha", n)) for n, r in enumerate(records)), + ) for pattern in (_ACCEPTANCE_LENGTH, _ACCEPTANCE_LENGTH_ALT): matches = pattern.findall(stdout) if matches: @@ -1178,8 +1589,9 @@ def run_boot_gate( unit: str | None = None, member: str | None = None, allow_dirty: bool | None = None, + workdir: str | Path | None = None, ) -> GateResult: - """Cold-start the engine and check it generates. + """Cold-start the engine and check it generates, through the engine driver. ``member`` is the parallel configuration to serve at, and it is the only thing that makes a ``unified`` target's gate mean anything — see @@ -1193,31 +1605,24 @@ def run_boot_gate( ``features`` is **the features the target under test actually implements**, which is not the same as the features the *run* will - finish with. Defaulting it to ``task_data["features"]`` — as this - function originally did with no way to override — means the ``core`` - unit is gated with the speculator's flags before a speculator exists. - - Measured: the core target refused `--spec_decode_algo=MTP` through its - own fail-fast assert, correctly, and the assembler escalated rather than - deleting the assert to make the gate pass. Deleting it would have been - the silent wrong — one-engine MTP is drafted by the model, so a base - model under that config emits no drafts and the run would have gated a - speculator that was not there. - - Staircase's rule is that core is green *before* the first feature + finish with. Defaulting it to ``task_data["features"]`` means the + ``core`` unit is gated with the speculator's flags before a speculator + exists. Staircase's rule is that core is green *before* the first feature increment; the gate has to agree with it. + ``spec`` is the workload (see :func:`resolve_boot_spec`): text prompts, + and for a multimodal unit at least one prompt that sends an image. The + engine driver builds the engine exactly as ``quickstart_advanced.py`` + does from the same flags and prints one record per prompt; keywords are + matched against those generated texts only, never against the log, so + ``72`` cannot pass on ``CPUs 72-143``. + ``mode`` defaults to ``require`` and should stay there for any result that will be attributed to modeling_v2. It is a parameter for one reason: when a boot gate fails, the next question is always *"is this the target or is this the host?"*, and the cheapest answer is the same command under ``off`` — the stock implementation, same prompts, same - runtime flags, one variable. Without this the control has to be - reconstructed by hand, which is how a run ends up debugging a target - that was never at fault. - - ``run_accuracy_gate`` already took a ``mode`` for the same reason; the - boot gate not having one was an inconsistency, not a decision. + runtime flags, one variable. ``allow_dirty`` is the escape hatch of :func:`inspect_tree`; see there for why a dirty target path stops this gate by default. @@ -1233,85 +1638,148 @@ def run_boot_gate( _verdict = _verdict_factory(state) if state.blocked: return _verdict(name=gate_name, member=member, passed=False, detail=state.block_detail()) - repo = Path(str(task_data["trtllm_repo_path"])) - script = repo / "examples" / "llm-api" / "quickstart_advanced.py" - prompts = " ".join(f'"{p}"' for p in spec.prompts) - argv = " ".join( - [ - str(script), - f'--model_dir="{task_data["checkpoint_path"]}"', - f"--max_tokens={spec.max_tokens}", - *topology_flags(str(task_data["parallel"]), environment.world_size, member), - *runtime_flags(resolve_runtime(task_data, unit)), - *speculation_flags( - task_data.get("features") if features is None else features, - max_draft_len=resolve_max_draft_len(task_data.get("accuracy") or {}, unit), - ), - f"--prompt {prompts}", - ] + requests = spec.requests(task_data.get("trtllm_repo_path")) + argv = quickstart_words( + task_data, + environment, + unit=unit, + features=task_data.get("features") if features is None else features, + max_tokens=spec.max_tokens, + member=member, + ) + spec_path = write_driver_spec( + gate_workdir(workdir, "boot", task_data) + / f"boot.{unit or 'unit'}.{member or 'member'}.json", + task_data=task_data, + argv=argv, + requests=requests, + max_tokens=spec.max_tokens, ) - command = env.python_command(argv) + command = engine_driver_command(env, spec_path) started = time.monotonic() result = runner(command, None, timeout) elapsed = time.monotonic() - started - tail = "\n".join((result.stdout + result.stderr).splitlines()[-40:]) + combined = result.stdout + result.stderr + tail = "\n".join(combined.splitlines()[-40:]) + ran = dict(name=gate_name, member=member, command=command, duration_s=elapsed, stdout_tail=tail) if result.timed_out: return _verdict( - name=gate_name, - member=member, passed=False, detail=f"timed out after {timeout}s — a boot gate that takes this long is " f"not catching catastrophes any more", - command=command, - duration_s=elapsed, - stdout_tail=tail, + **ran, ) if not result.ok: + return _verdict(passed=False, detail=f"exited {result.returncode}", **ran) + + records = {int(r.get("index", -1)): r for r in parse_driver_records(combined)} + texts = [ + str(records[n].get("text", "")) if n in records else None for n in range(len(requests)) + ] + absent = [requests[n]["id"] for n, text in enumerate(texts) if text is None] + if absent: return _verdict( - name=gate_name, - member=member, passed=False, - detail=f"exited {result.returncode}", - command=command, - duration_s=elapsed, - stdout_tail=tail, + detail=f"the engine exited cleanly but returned no record for {len(absent)} of " + f"{len(requests)} prompt(s) ({', '.join(absent)})", + **ran, ) - - missing = [k for k in spec.keywords if k not in result.stdout] - if missing: + empty = [requests[n]["id"] for n, text in enumerate(texts) if not (text or "").strip()] + if empty: return _verdict( - name=gate_name, - member=member, passed=False, - detail=f"generated output is missing {len(missing)} expected keyword(s): {missing}", - command=command, - duration_s=elapsed, - stdout_tail=tail, + detail=f"the engine exited cleanly but generated nothing for {', '.join(empty)}", + **ran, ) - if not result.stdout.strip(): + generated = "\n".join(text or "" for text in texts).lower() + missing = [k for k in spec.keywords if k.lower() not in generated] + item_texts = texts[len(spec.prompts) :] + for item, text in zip(spec.items, item_texts): + for absent_keyword in item.missing(text or ""): + label = "the image prompt" if item.images else "a prompt" + missing.append(f"{absent_keyword} (in {label} {item.prompt[:40]!r})") + if missing: return _verdict( - name=gate_name, - member=member, passed=False, - detail="the engine exited cleanly but produced no output at all", - command=command, - duration_s=elapsed, - stdout_tail=tail, + detail=f"the generated texts are missing {len(missing)} expected keyword(s): " + f"{missing}. Keywords are matched against generated text only", + **ran, ) - checked = f"{len(spec.keywords)} keyword(s)" if spec.keywords else "non-empty generation" + images = sum(len(r.get("images") or ()) for r in requests) + checked = ( + f"{len(spec.keywords) + sum(1 for i in spec.items if i.keywords or i.keywords_any)} " + f"keyword check(s) in the generated texts" + if spec.keywords or any(i.keywords or i.keywords_any for i in spec.items) + else "non-empty generation" + ) + sent = f", {images} image(s) sent" if images else "" return _verdict( - name=gate_name, - member=member, passed=True, - detail=f"engine cold-started and produced {checked}", - command=command, - duration_s=elapsed, - stdout_tail=tail, + detail=f"engine cold-started and produced {checked} for {len(requests)} prompt(s)" + f"{sent}; workload: {spec.source}", + **ran, ) +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 resolve_floor( + accuracy: Mapping[str, Any], + unit: str | None, + benchmark: str, + reference: float, + output: str = "", +) -> tuple[stats.Floor | None, str]: + """The floor ``reference`` gets under ``accuracy.floor``, and where its count came from.""" + policy = stats.resolve_floor_policy(accuracy) + samples, source = resolve_samples(accuracy, unit, benchmark, output) + return stats.floor_for(policy, reference, samples), source + + def run_accuracy_gate( environment: RunEnvironment, *, @@ -1330,6 +1798,13 @@ def run_accuracy_gate( ``unified`` target is the whole question: one codebase serving three configurations is three anchored numbers, not one repeated three times. + **The floor is statistical by default**: ``reference - k * SE``, where SE + is the binomial standard error of a score at the reference measured on + the benchmark's sample count (see :func:`resolve_samples`) and ``k`` is + ``accuracy.floor.k`` (default 2). A fixed ``reference - tolerance`` is + the explicit alternative, ``accuracy.floor.rule: tolerance``. Every + verdict records the SE, the count and the rule. + This is the gate the tree-state guard was written for. It is the expensive one — a quarter of an hour per member — and it is the one whose output becomes the target's headline number, so it is also the one where @@ -1342,14 +1817,23 @@ def run_accuracy_gate( env = _with_mode(environment, mode) state = inspect_tree(task_data, allow_dirty=allow_dirty) _verdict = _verdict_factory(state) + planned, _ = resolve_floor(accuracy, unit, benchmark, reference) + + def floor_fields(floor: stats.Floor | None) -> dict[str, Any]: + if floor is None: + return {} + return { + "floor": floor.value, + "se": floor.se, + "samples": floor.samples, + "floor_rule": floor.describe(), + } + + name = f"accuracy[{benchmark}]" + common = dict(name=name, member=member, reference=reference, tolerance=tolerance) if state.blocked: return _verdict( - name=f"accuracy[{benchmark}]", - member=member, - passed=False, - detail=state.block_detail(), - reference=reference, - tolerance=tolerance, + passed=False, detail=state.block_detail(), **common, **floor_fields(planned) ) parts = eval_command(task_data, environment, benchmark, unit=unit, member=member) @@ -1360,71 +1844,171 @@ def run_accuracy_gate( elapsed = time.monotonic() - started combined = result.stdout + result.stderr tail = "\n".join(combined.splitlines()[-40:]) + ran = dict(command=command, duration_s=elapsed, stdout_tail=tail, **common) if result.timed_out: return _verdict( - name=f"accuracy[{benchmark}]", - member=member, - passed=False, - detail=f"timed out after {timeout}s", - command=command, - reference=reference, - tolerance=tolerance, - duration_s=elapsed, - stdout_tail=tail, + passed=False, detail=f"timed out after {timeout}s", **ran, **floor_fields(planned) ) if not result.ok: return _verdict( - name=f"accuracy[{benchmark}]", - member=member, passed=False, detail=f"trtllm-eval exited {result.returncode}", - command=command, - reference=reference, - tolerance=tolerance, - duration_s=elapsed, - stdout_tail=tail, + **ran, + **floor_fields(planned), ) measured = parse_eval_score(combined) if measured is None: return _verdict( - name=f"accuracy[{benchmark}]", - member=member, passed=False, detail="trtllm-eval exited 0 but printed no score this module recognizes; " "a gate with no number is not a pass", - command=command, - reference=reference, - tolerance=tolerance, - duration_s=elapsed, - stdout_tail=tail, + **ran, + **floor_fields(planned), ) - floor = reference - tolerance - passed = measured >= floor + floor, source = resolve_floor(accuracy, unit, benchmark, reference, combined) + if floor is None: + return _verdict( + passed=False, + detail=f"measured {measured:.4f}, but the statistical floor needs the sample " + f"count and there is {source}. A floor that cannot be computed is not a pass", + measured=measured, + **ran, + ) + passed = measured >= floor.value return _verdict( - name=f"accuracy[{benchmark}]", - member=member, passed=passed, detail=( - f"measured {measured:.4f} {'>=' if passed else '<'} {floor:.4f} " - f"(reference {reference:.4f} - tol {tolerance:.2f})" + f"measured {measured:.4f} {'>=' if passed else '<'} {floor.value:.4f} " + f"({floor.describe()}; count from {source})" ), - command=command, measured=measured, - reference=reference, - tolerance=tolerance, - duration_s=elapsed, - stdout_tail=tail, + **ran, + **floor_fields(floor), ) -def resolve_acceptance_floor(accuracy: Mapping[str, Any], unit: str | None) -> float | None: - """The absolute `acceptance_length` floor this unit declares, if any. +def check_systematic_shift( + results: Sequence[GateResult], + accuracy: Mapping[str, Any], + unit: str | None, +) -> stats.ShiftFinding | None: + """The cross-member check over one unit's accuracy verdicts. + + Every member passing its own floor is not the same as every member being + on the anchor: a defect in code the members share moves all of them the + same way, and each still clears ``reference - 2 SE``. When **every** + member sits more than one SE under the reference, that is a finding, and + delivery refuses it unless ``accuracy.floor.acknowledge.`` gives a + reason. Verdicts without a measurement or an SE are left out. + """ + policy = stats.resolve_floor_policy(accuracy) + readings = { + str(r.member): (float(r.measured), float(r.reference), float(r.se)) + for r in results + if r.measured is not None and r.reference is not None and r.se is not None + } + return stats.systematic_shift( + unit, + readings, + min_members=policy.shift_min_members, + acknowledged=policy.acknowledged, + ) - Only consulted when the `=off` leg **cannot exist**. See - :func:`run_acceptance_gate`. + +#: The key a systematic-shift finding is recorded under, per unit. +SHIFT_FINDING = "systematic_shift" + + +def records_for(workspace: str | Path) -> GateRecords: + """The results file of a run's workspace.""" + return GateRecords(Path(workspace) / RESULTS_FILENAME) + + +def note_systematic_shift( + records: GateRecords, + results: Sequence[GateResult], + accuracy: Mapping[str, Any], + unit: str | None, +) -> stats.ShiftFinding | None: + """Run :func:`check_systematic_shift` and record its outcome either way. + + A unit with no shift records a resolved entry, so a finding raised by an + earlier turn of the same unit stops blocking once a later turn is clean. + """ + finding = check_systematic_shift(results, accuracy, unit) + tree = results[-1].tree_state if results else UNKNOWN_TREE_STATE + records.append( + GateRecord( + kind=KIND_FINDING, + key={"finding": SHIFT_FINDING, "unit": str(unit)}, + tree=tree, + passed=finding is None or bool(finding.acknowledged), + values={ + "detail": finding.describe() if finding else "no systematic shift", + "members": { + str(r.member): {"measured": r.measured, "reference": r.reference, "se": r.se} + for r in results + }, + }, + ) + ) + return finding + + +def blocking_findings(records: GateRecords, accuracy: Mapping[str, Any]) -> list[GateRecord]: + """Findings that refuse delivery under the CURRENT ``accuracy.floor.acknowledge``.""" + return open_findings(records, stats.resolve_floor_policy(accuracy).acknowledged) + + +#: Prompts an acceptance workload must hold before its rates are compared. On +#: four prompts a change that only moved the generated text moved the mean by +#: +0.24 / -0.18 / -0.02 across three members; on 32 GSM8K prompts the same +#: change measured -0.06 / +0.05 at paired t -1.4 / +0.8 -- no change. A task +#: may lower it through ``accuracy.acceptance.min_prompts``, and the verdict +#: then rests on fewer pairs, visibly. +DEFAULT_ACCEPTANCE_MIN_PROMPTS = 32 +#: Prompts taken from the accuracy dataset when the task declares none. +DEFAULT_ACCEPTANCE_NUM_PROMPTS = 32 +#: The minimum drop, in accepted tokens per iteration, a paired comparison +#: must find before it calls a significant change a regression. +DEFAULT_ACCEPTANCE_MIN_EFFECT = 0.1 + + +@dataclass(frozen=True) +class AcceptanceFloor: + """An absolute acceptance floor, and the measurement it came from, if any.""" + + value: float + #: Where the reference drafter this floor derives from was measured; + #: ``None`` when the task declared a number with no measurement behind it. + source: str | None = None + #: The measured reference value the floor was derived from, when it was. + measured: float | None = None + + def warning(self) -> str | None: + """What a verdict says when the floor is not tied to a measurement.""" + if self.source or self.measured is not None: + return None + return ( + f"the acceptance floor {self.value:.2f} is not tied to a measured reference: " + f"declare `floor_source` naming the reference-drafter measurement it was " + f"derived from, or `measured_reference` with the value itself. A floor chosen " + f"from the ceiling passed every wrong drafter of the run that had one" + ) + + +def resolve_acceptance_floor_spec( + accuracy: Mapping[str, Any], unit: str | None, *, min_ratio: float = 0.95 +) -> AcceptanceFloor | None: + """The absolute floor ``unit`` declares in ``accuracy.acceptance``, if any. + + ``min_absolute`` is the floor itself. ``measured_reference`` is what a + reference drafter measured on this workload; with no ``min_absolute`` + the floor is that times ``min_ratio``, the same ratio the stock + comparison uses. ``floor_source`` names where the measurement was taken. """ block = accuracy.get("acceptance") if not isinstance(block, dict): @@ -1433,43 +2017,125 @@ def resolve_acceptance_floor(accuracy: Mapping[str, Any], unit: str | None) -> f if chosen is None: chosen = block.get("default") if isinstance(chosen, dict): - chosen = chosen.get("min_absolute") + source = chosen.get("floor_source") + measured = chosen.get("measured_reference") + measured = float(measured) if isinstance(measured, (int, float)) else None + if chosen.get("min_absolute") is not None: + try: + return AcceptanceFloor(float(chosen["min_absolute"]), source, measured) + except (TypeError, ValueError): + return None + if measured is not None: + return AcceptanceFloor(measured * min_ratio, source, measured) + return None try: - return float(chosen) + return AcceptanceFloor(float(chosen)) except (TypeError, ValueError): return None +def resolve_acceptance_floor(accuracy: Mapping[str, Any], unit: str | None) -> float | None: + """The absolute `acceptance_length` floor this unit declares, if any. + + Only consulted when the `=off` leg **cannot exist**. See + :func:`run_acceptance_gate` and :func:`resolve_acceptance_floor_spec`. + """ + floor = resolve_acceptance_floor_spec(accuracy, unit) + return floor.value if floor is not None else None + + @dataclass(frozen=True) class AcceptanceWorkload: - """What the acceptance gate sends, and how much work it must get back.""" + """What the acceptance gate sends, and how much of it it needs back.""" + #: Declared prompts; empty when ``dataset`` supplies them. prompts: tuple[str, ...] max_tokens: int min_draft_iterations: int #: Where the prompts came from, named in a verdict that blames them. source: str + #: The engine driver's ``dataset`` block, when the prompts come from one. + dataset: Mapping[str, Any] | None = None + min_prompts: int = DEFAULT_ACCEPTANCE_MIN_PROMPTS + + @property + def fingerprint(self) -> str: + """Names this workload, so two runs of it can be paired.""" + payload = json.dumps( + { + "prompts": list(self.prompts), + "dataset": dict(self.dataset or {}), + "max_tokens": self.max_tokens, + }, + sort_keys=True, + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:16] + + def requests(self) -> list[dict[str, Any]]: + """The declared prompts as engine-driver requests.""" + return [ + {"id": f"p{n}", "prompt": prompt, "max_tokens": self.max_tokens} + for n, prompt in enumerate(self.prompts) + ] + + +def acceptance_dataset( + accuracy: Mapping[str, Any], unit: str | None, setting: Callable[[str], Any] +) -> tuple[dict[str, Any] | None, str]: + """The dataset an undeclared acceptance workload is drawn from, and its name. + + ``accuracy.acceptance.dataset`` when declared; otherwise the accuracy + benchmark's own dataset -- ``accuracy.dataset_path`` when the run keeps + it locally, else the benchmark's Hugging Face id, read from the local + cache only (:data:`.harness.workload_io.KNOWN_DATASETS`). A multimodal + unit's benchmark has no text dataset, so the run-wide benchmark's is used. + """ + known: dict[str, Any] = {} + for scope in (unit, "default"): + benchmark = resolve_per_unit(accuracy, "benchmark", scope) + if isinstance(benchmark, str) and benchmark in workload_io.KNOWN_DATASETS: + known = dict(workload_io.KNOWN_DATASETS[benchmark]) + break + declared = setting("dataset") + if declared: + source, where = str(declared), "accuracy.acceptance.dataset" + elif accuracy.get("dataset_path") and known: + source, where = str(accuracy["dataset_path"]), "accuracy.dataset_path" + elif known: + source, where = str(known["source"]), "the accuracy benchmark's dataset" + else: + return None, "" + dataset = { + "source": source, + "num": int(setting("num_prompts") or DEFAULT_ACCEPTANCE_NUM_PROMPTS), + "text_field": setting("text_field") or known.get("text_field"), + "answer_field": known.get("answer_field"), + "split": setting("split") or known.get("split", "test"), + "config": setting("dataset_config") or (None if declared else known.get("config")), + "template": setting("prompt_template") or "{text}", + } + return dataset, f"{dataset['num']} prompts from {where} ({source})" def resolve_acceptance_workload( accuracy: Mapping[str, Any], unit: str | None, *, - fallback: BootSpec | None = None, max_tokens: int | None = None, ) -> AcceptanceWorkload: - """The acceptance workload ``unit`` declares, or the fallback when none does. + """The acceptance workload ``unit`` declares, or the accuracy dataset's prompts. Read from ``accuracy.acceptance``, most specific level first: the unit's own entry, then its ``default`` entry, then the block's top level (see ``task_schema.ACCEPTANCE_WORKLOAD_KEYS``). ``prompts`` and ``prompts_file`` are one setting, so the level that states either wins - for both. ``fallback`` (the boot spec) and ``max_tokens`` are what a task - that declares nothing measures, which is what every task measured before - the workload could be declared. + for both. With neither, the workload is the first ``num_prompts`` + (default 32) natural-text questions of the accuracy dataset; see + :func:`acceptance_dataset`. Raises :class:`~.task_schema.TaskSchemaError` when a declared prompts - file cannot be read. + file cannot be read, or when nothing declares a workload and the + benchmark has no dataset to draw one from. """ block = accuracy.get("acceptance") block = block if isinstance(block, dict) else {} @@ -1487,8 +2153,9 @@ def setting(key: str) -> Any: return entry[key] return None - prompts: tuple[str, ...] | None = None + prompts: tuple[str, ...] = () source = "" + dataset: dict[str, Any] | None = None for where, entry in levels: if entry.get("prompts") is not None: prompts = tuple(str(p) for p in entry["prompts"]) @@ -1498,23 +2165,89 @@ def setting(key: str) -> Any: prompts = load_prompts_file(entry["prompts_file"]) source = f"{where}.prompts_file ({entry['prompts_file']})" break - if prompts is None: - prompts = (fallback or BootSpec()).prompts - source = "the boot prompts (no `accuracy.acceptance.prompts` declared)" + if not prompts: + dataset, source = acceptance_dataset(accuracy, unit, setting) + if dataset is None: + raise TaskSchemaError( + "no acceptance workload: declare `accuracy.acceptance.prompts`, " + "`prompts_file` or `dataset` (a local path or a cached dataset id); the " + "benchmark has no dataset this gate knows how to draw prompts from" + ) declared_tokens = setting("max_tokens") - declared_minimum = setting("min_draft_iterations") if declared_tokens is None: declared_tokens = max_tokens if max_tokens is not None else DEFAULT_ACCEPTANCE_MAX_TOKENS - if declared_minimum is None: - declared_minimum = DEFAULT_MIN_DRAFT_ITERATIONS + declared_minimum = setting("min_draft_iterations") + min_prompts = setting("min_prompts") return AcceptanceWorkload( prompts=prompts, max_tokens=int(declared_tokens), - min_draft_iterations=int(declared_minimum), + min_draft_iterations=int( + DEFAULT_MIN_DRAFT_ITERATIONS if declared_minimum is None else declared_minimum + ), source=source, + dataset=dataset, + min_prompts=int(DEFAULT_ACCEPTANCE_MIN_PROMPTS if min_prompts is None else min_prompts), ) +@dataclass(frozen=True) +class RegressionPolicy: + """How a paired acceptance comparison is judged (``accuracy.acceptance.regression``).""" + + test: str = stats.TEST_T + alpha: float = stats.DEFAULT_ALPHA + min_effect: float = DEFAULT_ACCEPTANCE_MIN_EFFECT + + def compare( + self, baseline: Sequence[float], candidate: Sequence[float] + ) -> stats.PairedComparison: + """:func:`.stats.paired_test` under this policy.""" + return stats.paired_test( + baseline, candidate, test=self.test, alpha=self.alpha, min_effect=self.min_effect + ) + + +def resolve_regression_policy(accuracy: Mapping[str, Any]) -> RegressionPolicy: + """Read ``accuracy.acceptance.regression``; every key optional.""" + block = accuracy.get("acceptance") + block = (block if isinstance(block, dict) else {}).get("regression") + block = block if isinstance(block, dict) else {} + return RegressionPolicy( + test=str(block.get("test") or stats.TEST_T), + alpha=float(block.get("alpha", stats.DEFAULT_ALPHA)), + min_effect=float(block.get("min_effect", DEFAULT_ACCEPTANCE_MIN_EFFECT)), + ) + + +def tree_commit(repo_root: str | Path | None) -> str | None: + """The full commit ``HEAD`` names in ``repo_root``, or ``None`` if git cannot say.""" + if not repo_root: + return None + answer = _git(repo_root, "rev-parse", "HEAD") + return answer.strip() if answer and answer.strip() else None + + +def acceptance_key( + workload: AcceptanceWorkload, + *, + member: str | None, + speculation: Sequence[str], +) -> dict[str, Any]: + """What makes two acceptance verdicts comparable across trees. + + The member, the workload and the speculator flags -- not the unit: a + later unit that keeps speculation on is compared against the earlier + unit's verdict, which is how a regression a later feature introduces in + the draft path shows up. + """ + return { + "gate": "acceptance", + "member": str(member), + "workload": workload.fingerprint, + "speculation": " ".join(speculation), + } + + def run_acceptance_gate( environment: RunEnvironment, *, @@ -1523,51 +2256,49 @@ def run_acceptance_gate( timeout: int = BOOT_TIMEOUT_S, min_ratio: float = 0.95, unit: str | None = None, - spec: BootSpec | None = None, features: Any | None = None, max_tokens: int | None = None, member: str | None = None, allow_dirty: bool | None = None, + workdir: str | Path | None = None, + records: GateRecords | None = None, ) -> GateResult: - """Compare ``acceptance_length`` against stock in-tree modeling. - - The two runs differ in exactly one variable — ``TRTLLM_MODELING_V2`` - ``off`` versus ``require`` — and must happen **in the same session**. - A cross-session delta carries session variance into the judgement, and - staircase measured a bit-identical forward at 94.7688 and 95.0720 a day - apart. - - ``member`` fixes the topology **both** legs run at, so the one variable - stays one: a draft path whose acceptance depends on the expert-parallel - slice is a thing this gate can only see when it is asked at that member. - - The tree-state guard matters here for a reason the other two gates do not - have: **both legs import the same working tree**, so an edit landing - between them moves the ratio without moving either leg's correctness, and - the gate reports it as a draft-path defect. See :func:`inspect_tree`. - - **The workload is the task's to declare** (``accuracy.acceptance.prompts`` - or ``prompts_file``, ``max_tokens``; see - :func:`resolve_acceptance_workload`). ``spec`` and ``max_tokens`` here are - only the fallback for a task that declares none. The boot prompts are - the wrong workload for a model whose answers to them are short: - DeepSeek-V4.1 answers them in 2-3 tokens and EOS, so every request is a - context step plus one generation step, the rate reads exactly 1.00 for - ANY drafter, and a declared floor above 1 can never pass while a broken - drafter cannot be told from a working one. On long-answer prompts the - same command measured 1.91 for a drafter with a broken ring index and - 3.24 after the fix. - - **A workload too short to measure is a failure of its own**, reported as - one: fewer than ``min_draft_iterations`` generation iterations in total - (from quickstart's ``Accept rate: ... (tokens=T, iterations=I)`` lines) - fails with "workload too short to measure acceptance" instead of a low - number that reads as a drafter defect. A leg that prints only a summary - statistic carries no iteration count and is not held to it. - - Every verdict that has a reading carries the per-prompt rates and their - plain mean beside the iteration-weighted value it is judged on, in - :attr:`GateResult.breakdown`. + """Measure per-prompt acceptance, and judge it against stock, a floor and the last tree. + + **The workload** is the task's to declare (``accuracy.acceptance.prompts``, + ``prompts_file`` or ``dataset``; see :func:`resolve_acceptance_workload`) + and is otherwise the first 32 questions of the accuracy dataset: natural + text whose answers run long. It goes to the engine driver (see + :data:`ENGINE_DRIVER`), which runs each prompt alone, streamed, and + prints its tokens and iterations; the rate of a prompt is its tokens per + iteration, and the gate's number is iteration-weighted over all of them. + + **Size.** Fewer than ``min_prompts`` prompts (default 32) or fewer than + ``min_draft_iterations`` generation iterations fails as a workload too + small or too short to measure, not as a low number. + + **Three comparisons**, each only where it can be made: + + * against stock: the same command under ``TRTLLM_MODELING_V2=off``, in + the same session. It fails when the weighted ratio is under + ``min_ratio`` AND the per-prompt paired test finds the drop + significant. A new architecture has no stock path; then the declared + floor (see :func:`resolve_acceptance_floor_spec`) is used instead, and + the verdict warns when that floor is not tied to a measured reference. + * against 1.0: a target at or under it rejects every draft. + * against the previous verdict for the same member, workload and + speculator in ``records``: a paired test over the prompts both + measured (``accuracy.acceptance.regression``: t or Wilcoxon, alpha, + minimum effect). A significant drop at least the minimum effect large + fails. This replaces the fixed 0.1 delta that read noise on four + prompts as a regression. + + Every verdict with a reading carries the per-prompt rates. A passing + reading is appended to ``records`` and becomes the next tree's baseline. + + The two legs differ in exactly one variable and must run in the same + session; both import the same working tree, so the tree-state guard of + :func:`inspect_tree` stops a gate whose target path is being edited. """ state = inspect_tree(task_data, allow_dirty=allow_dirty) _verdict = _verdict_factory(state) @@ -1575,9 +2306,7 @@ def run_acceptance_gate( return _verdict(name="acceptance", member=member, passed=False, detail=state.block_detail()) accuracy = task_data.get("accuracy") or {} try: - workload = resolve_acceptance_workload( - accuracy, unit, fallback=spec or BootSpec(), max_tokens=max_tokens - ) + workload = resolve_acceptance_workload(accuracy, unit, max_tokens=max_tokens) except TaskSchemaError as exc: return _verdict( name="acceptance", @@ -1585,71 +2314,52 @@ def run_acceptance_gate( passed=False, detail=f"the declared acceptance workload cannot be read: {exc}", ) - # The QUICKSTART command, not the eval command, and not as an - # optimization. `trtllm-eval` emits no speculative statistic in this - # tree, so an acceptance gate built on it reads None from both legs and - # fails identically to a broken draft path. `quickstart_advanced.py` - # prints `Accept rate: N (tokens=T, iterations=I)` per request, which is - # the quantity this gate is about. - # - # It is also two orders of magnitude cheaper -- minutes against the nine - # a full gsm8k run costs, twice -- which matters because this gate runs - # both legs. - repo = Path(str(task_data["trtllm_repo_path"])) - script = repo / "examples" / "llm-api" / "quickstart_advanced.py" - # Shell-quoted: a declared workload is free text, and a code-listing - # prompt carries the quotes, `$` and backslashes a double-quoted word - # would expand or end on. - prompts = " ".join(shlex.quote(p) for p in workload.prompts) - argv = " ".join( - [ - str(script), - f'--model_dir="{task_data["checkpoint_path"]}"', - # Longer than the boot gate's, deliberately. Acceptance length is - # a rate, and a 32-token continuation gives it a handful of - # iterations to be estimated from -- enough to see a draft path - # collapse to 1.0, not enough to resolve a 5% shortfall. - f"--max_tokens={workload.max_tokens}", - *topology_flags(str(task_data["parallel"]), environment.world_size, member), - *runtime_flags(resolve_runtime(task_data, unit)), - *speculation_flags( - task_data.get("features") if features is None else features, - max_draft_len=resolve_max_draft_len(accuracy, unit), - ), - f"--prompt {prompts}", - ] + features = task_data.get("features") if features is None else features + draft_len = resolve_max_draft_len(accuracy, unit) + speculation = speculation_flags(features, max_draft_len=draft_len) + argv = quickstart_words( + task_data, + environment, + unit=unit, + features=features, + max_tokens=workload.max_tokens, + member=member, + ) + directory = gate_workdir(workdir, "acceptance", task_data) + spec_path = write_driver_spec( + directory / f"acceptance.{unit or 'unit'}.{member or 'member'}.json", + task_data=task_data, + argv=argv, + requests=workload.requests(), + dataset=workload.dataset, + max_tokens=workload.max_tokens, ) + policy = resolve_regression_policy(accuracy) + floor_spec = resolve_acceptance_floor_spec(accuracy, unit, min_ratio=min_ratio) started = time.monotonic() readings: dict[str, AcceptanceReading | None] = {} commands: dict[str, str] = {} stock_unavailable = False for mode in (MODE_REQUIRE, MODE_OFF): - command = _with_mode(environment, mode).python_command(argv) + command = engine_driver_command(_with_mode(environment, mode), spec_path) commands[mode] = command result = runner(command, None, timeout) combined = result.stdout + result.stderr if not result.ok: - # **The stock leg can be impossible rather than broken.** This gate - # compares the target against the same checkpoint under built-in - # modeling -- and a genuinely new architecture HAS no built-in - # modeling, so `=off` cannot load it at all. That is the same shape - # as the substrate gate's problem (see - # `preflight.UNSUPPORTED_ARCHITECTURE_MARKERS`, reused here): the - # check is unavailable exactly when onboarding is doing its job. - # - # Failing the unit for it would be wrong, and silently skipping the - # gate would be worse -- it is the ONLY gate that can see a - # miscomputed draft path, because rejection sampling makes a wrong - # draft produce correct text more slowly. So the spec may declare an - # ABSOLUTE floor instead, and the weaker claim is recorded verbatim. + # **The stock leg can be impossible rather than broken.** A + # genuinely new architecture HAS no built-in modeling, so `=off` + # cannot load it at all -- the check is unavailable exactly when + # onboarding is doing its job (see + # `preflight.UNSUPPORTED_ARCHITECTURE_MARKERS`, reused here). The + # spec may then declare an absolute floor instead, and the weaker + # claim is recorded verbatim. from .preflight import reports_unsupported_architecture - floor = resolve_acceptance_floor(accuracy, unit) unavailable = mode == MODE_OFF and reports_unsupported_architecture( f"exited {result.returncode}", "\n".join(combined.splitlines()[-200:]) ) - if unavailable and floor is not None: + if unavailable and floor_spec is not None: readings[mode] = None stock_unavailable = True break @@ -1663,7 +2373,8 @@ def run_acceptance_gate( + ( " -- this checkpoint has no built-in implementation, so the " "stock leg cannot run. Declare `accuracy.acceptance.." - "min_absolute` to gate on an absolute floor instead." + "min_absolute` (or `measured_reference`) to gate on an absolute " + "floor instead." if unavailable else "" ) @@ -1680,110 +2391,171 @@ def run_acceptance_gate( target = target_reading.value if target_reading is not None else None stock = stock_reading.value if stock_reading is not None else None breakdown = target_reading.breakdown() if target_reading is not None else "" + base = dict(name="acceptance", member=member, duration_s=elapsed) - def too_short(mode: str) -> GateResult | None: - """The verdict for a leg whose workload never exercised the drafter.""" + def too_small(mode: str) -> GateResult | None: + """The verdict for a leg whose workload cannot carry a comparison.""" reading = readings.get(mode) - counted = reading.draft_iterations if reading is not None else None - if reading is None or counted is None or counted >= workload.min_draft_iterations: + if reading is None or not reading.requests: return None - return _verdict( - name="acceptance", - member=member, - passed=False, - detail=( + counted = reading.draft_iterations or 0 + if len(reading.requests) < workload.min_prompts: + why = ( + f"workload too small to compare acceptance: the {MODELING_V2_ENV}={mode} leg " + f"measured {len(reading.requests)} prompt(s), fewer than the " + f"{workload.min_prompts} a paired comparison needs -- on four prompts a " + f"change that only moved the generated text read as a regression" + ) + elif counted < workload.min_draft_iterations: + why = ( f"workload too short to measure acceptance: the {MODELING_V2_ENV}={mode} " f"leg ran {counted} generation iteration(s) over {len(reading.requests)} " f"prompt(s), fewer than the {workload.min_draft_iterations} this gate " f"needs, so its rate ({reading.value:.2f}) says nothing about the draft " - f"path -- the answers end before drafts are verified. Workload: " - f"{workload.source}, max_tokens {workload.max_tokens}. Declare " - f"long-answer prompts in `accuracy.acceptance.prompts` (or " - f"`prompts_file`); this is a property of the workload, not of the target." - ), - command=commands[mode], - duration_s=elapsed, - breakdown=reading.breakdown(), - ) - - if stock_unavailable: - floor = resolve_acceptance_floor(accuracy, unit) - if target is None: - return _verdict( - name="acceptance", - member=member, - passed=False, - detail="could not read acceptance_length from the require leg", - command=commands[MODE_REQUIRE], - duration_s=elapsed, + f"path -- the answers end before drafts are verified" ) - short = too_short(MODE_REQUIRE) - if short is not None: - return short + else: + return None return _verdict( - name="acceptance", - member=member, - passed=target >= float(floor), + passed=False, detail=( - f"acceptance_length {target:.4f} vs declared floor {float(floor):.4f}. " - f"The {MODELING_V2_ENV}=off leg could not run: this checkpoint has no " - f"built-in implementation, so there is no stock acceptance to compare " - f"against. THIS IS THE WEAKER CLAIM -- it shows the draft path accepts " - f"at a real rate, not that it matches stock." + f"{why}. Workload: {workload.source}, max_tokens {workload.max_tokens}. " + f"This is a property of the workload, not of the target." ), - command=commands[MODE_REQUIRE], - measured=target, - reference=float(floor), - duration_s=elapsed, - breakdown=breakdown, + command=commands[mode], + breakdown=reading.breakdown(), + **base, ) - if target is None or stock is None: + + if target is None: return _verdict( - name="acceptance", - member=member, passed=False, - detail=f"could not read acceptance_length from both legs " - f"(require={target}, off={stock})", + detail=f"could not read acceptance from the require leg (off={stock})", command=commands[MODE_REQUIRE], - duration_s=elapsed, + **base, ) for mode in (MODE_REQUIRE, MODE_OFF): - short = too_short(mode) + short = too_small(mode) if short is not None: return short - stock_breakdown = stock_reading.breakdown() if stock_reading is not None else "" - if breakdown and stock_breakdown: - breakdown = f"require: {breakdown}; off: {stock_breakdown}" + + warnings: list[str] = [] + notes: list[str] = [] if target <= 1.0: return _verdict( - name="acceptance", - member=member, passed=False, - detail=f"acceptance_length {target:.4f} at or below 1.0 means every draft " - f"was rejected — the draft path is computing the wrong thing", + detail=f"acceptance_length {target:.4f} at or below 1.0 means every draft was " + f"rejected — the draft path is computing the wrong thing", command=commands[MODE_REQUIRE], measured=target, - reference=stock, - duration_s=elapsed, + reference=stock if stock is not None else (floor_spec.value if floor_spec else None), breakdown=breakdown, + **base, + ) + + reference: float | None + if stock_unavailable: + assert floor_spec is not None + reference = floor_spec.value + passed = target >= floor_spec.value + if floor_spec.warning(): + warnings.append(floor_spec.warning()) + derived = ( + f" (= measured reference {floor_spec.measured:.2f} x {min_ratio:.2f})" + if floor_spec.measured is not None and floor_spec.source is None + else "" + ) + notes.append( + f"acceptance_length {target:.4f} vs declared floor {floor_spec.value:.4f}{derived}" + + (f" (from {floor_spec.source})" if floor_spec.source else "") + + f". The {MODELING_V2_ENV}=off leg could not run: this checkpoint has no " + f"built-in implementation. THIS IS THE WEAKER CLAIM -- it shows the draft path " + f"accepts at a real rate, not that it matches stock" + ) + else: + if stock is None: + return _verdict( + passed=False, + detail=f"could not read acceptance from both legs (require={target}, off={stock})", + command=commands[MODE_REQUIRE], + **base, + ) + stock_breakdown = stock_reading.breakdown() if stock_reading is not None else "" + if breakdown and stock_breakdown: + breakdown = f"require: {breakdown}; off: {stock_breakdown}" + reference = stock + ratio = target / stock if stock else 0.0 + paired = None + if target_reading.requests and stock_reading.requests: + own, theirs = target_reading.rates(), stock_reading.rates() + shared = [p for p in own if p in theirs] + if len(shared) >= 2: + paired = policy.compare([theirs[p] for p in shared], [own[p] for p in shared]) + passed = ratio >= min_ratio or (paired is not None and not paired.significant) + notes.append( + f"acceptance_length {target:.4f} vs stock {stock:.4f} (ratio {ratio:.3f}, floor " + f"{min_ratio:.2f})" + + (f"; {paired.describe()}" if paired is not None else "") + + ( + ". Clearly above 1.0 but below the reference means subtly wrong, not merely slower" + if not passed + else "" + ) + ) + + key = acceptance_key(workload, member=member, speculation=speculation) + if records is not None and target_reading.requests: + baseline = records.latest(KIND_ACCEPTANCE, key, exclude_tree=state.description) + if baseline is None: + notes.append( + "no earlier verdict for this member and workload; this one is the baseline" + ) + else: + before = { + str(k): float(v) for k, v in (baseline.values.get("per_prompt") or {}).items() + } + now = target_reading.rates() + shared = [p for p in now if p in before] + if len(shared) < 2: + warnings.append( + f"the earlier verdict on tree {baseline.tree} shares {len(shared)} prompt(s) " + f"with this one, so the two were not compared" + ) + else: + compared = policy.compare([before[p] for p in shared], [now[p] for p in shared]) + regressed = compared.regressed + notes.append( + f"against tree {baseline.tree}: " + + ("REGRESSED, " if regressed else "") + + compared.describe() + ) + passed = passed and not regressed + records.append( + GateRecord( + kind=KIND_ACCEPTANCE, + key=key, + tree=state.description, + commit=tree_commit(task_data.get("trtllm_repo_path")), + passed=passed, + values={ + "unit": unit, + "weighted": target, + "per_prompt": target_reading.rates(), + "prompts": len(target_reading.requests), + }, + ) ) - ratio = target / stock if stock else 0.0 - passed = ratio >= min_ratio return _verdict( - name="acceptance", - member=member, passed=passed, - detail=( - f"acceptance_length {target:.4f} vs stock {stock:.4f} " - f"(ratio {ratio:.3f}, floor {min_ratio:.2f}). Clearly above 1.0 but below " - f"the reference means subtly wrong, not merely slower." - ), + detail="; ".join(notes), command=commands[MODE_REQUIRE], measured=target, - reference=stock, - duration_s=elapsed, + reference=reference, breakdown=breakdown, + warnings=tuple(warnings), + explained=True, + **base, ) @@ -1842,7 +2614,39 @@ def resolve_reference( "AcceptanceReading", "AcceptanceWorkload", "BOOT_TIMEOUT_S", + "AcceptanceFloor", + "BATCH_SIZE_FLAG", + "BOOT_RENDERS", + "BootPrompt", "BootSpec", + "DEFAULT_BOOT_IMAGE", + "DEFAULT_BOOT_IMAGE_KEYWORDS", + "DEFAULT_BOOT_IMAGE_PROMPT", + "resolve_boot_spec", + "DEFAULT_ACCEPTANCE_MIN_EFFECT", + "DEFAULT_ACCEPTANCE_MIN_PROMPTS", + "DEFAULT_ACCEPTANCE_NUM_PROMPTS", + "ENGINE_DRIVER", + "RegressionPolicy", + "SHIFT_FINDING", + "acceptance_dataset", + "acceptance_key", + "engine_driver_command", + "gate_workdir", + "parse_driver_records", + "quickstart_script", + "quickstart_words", + "resolve_acceptance_floor_spec", + "resolve_regression_policy", + "tree_commit", + "with_max_batch_size", + "write_driver_spec", + "blocking_findings", + "check_systematic_shift", + "note_systematic_shift", + "records_for", + "resolve_floor", + "resolve_samples", "DEFAULT_ACCEPTANCE_MAX_TOKENS", "DEFAULT_MIN_DRAFT_ITERATIONS", "GateLog", diff --git a/agent-flow/agent_flow/workflows/staircase_onboard/harness/__init__.py b/agent-flow/agent_flow/workflows/staircase_onboard/harness/__init__.py new file mode 100644 index 000000000000..473fb6adf0fa --- /dev/null +++ b/agent-flow/agent_flow/workflows/staircase_onboard/harness/__init__.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Programs the gates run on the compute side, under the interpreter of the tree being measured. + +Nothing in this package imports ``agent_flow``: each script runs as +`` /