diff --git a/.github/workflows/ci-l0-checks.yml b/.github/workflows/ci-l0-checks.yml index ce1e313c4..bc2969bba 100644 --- a/.github/workflows/ci-l0-checks.yml +++ b/.github/workflows/ci-l0-checks.yml @@ -55,10 +55,13 @@ jobs: - uses: actions/setup-python@v5 with: python-version: "3.11" - # pyyaml is required, not optional: run_e2e.py parses the recipe envs: block - # with it, and the recipe-env tests assert the string-preserving BaseLoader - # path (typed scalars, folded multi-line EXTRA__ARGS). Without it those - # tests would silently exercise the degraded line-scanner fallback instead. + # pyyaml is required, not optional, and two suites need it. run_e2e.py parses + # the recipe envs: block with it, and the recipe-env tests assert the + # string-preserving BaseLoader path (typed scalars, folded multi-line + # EXTRA__ARGS); without it those tests would silently exercise the degraded + # line-scanner fallback instead. The experience-store tests read/write the + # on-disk meta.yaml format (that module degrades to JSON without pyyaml, but + # then the test would be asserting a format the real store never uses). - run: pip install pytest pytest-cov coverage pyyaml # Stdlib-only tests (import run_e2e via importlib; no claude_agent_sdk / GPU). # Files are listed explicitly rather than discovered, so that adding a test @@ -95,6 +98,12 @@ jobs: e2e_workflow/scripts/tests/test_leg_runner.py \ e2e_workflow/scripts/tests/test_server_teardown.py \ e2e_workflow/scripts/tests/test_bench_e2e_teardown_lookup.py \ + e2e_workflow/scripts/tests/test_e2e_store.py \ + e2e_workflow/scripts/tests/test_kb_retract.py \ + kernel_workflow/scripts/tests/test_experience_store.py \ + kb/tests/test_store_local.py \ + kb/tests/test_attest.py \ + kernel_workflow/scripts/tests/test_kb_loop_offline.py \ ci/node/test_run_model_material_audit.py \ geak/test_bootstrap.py diff --git a/.gitignore b/.gitignore index ab52feb0c..56e1b1917 100644 --- a/.gitignore +++ b/.gitignore @@ -224,6 +224,11 @@ exp/ .torch_ext/ # bench_e2e.sh default output dir (OUT_DIR=$(pwd)/e2e_bench_out) — runtime measurements e2e_bench_out/ +# Machine-produced experience KB (warm-start store, kernel_workflow Part 4/1.7): runtime-accumulated +# best patches + meta; volume is unbounded, so default-ignore and commit selectively via a whitelist. +kb_artifacts/ +# Default root of the on-disk KB Store plane (kb/store_local.py): runtime records + artifacts. +kb_store_local/ # rocprofv3 scratch: written next to the workflow scripts when a profiling run is # launched from this directory. Build output, never source. diff --git a/README.md b/README.md index dea52229c..fab65c6f1 100644 --- a/README.md +++ b/README.md @@ -183,6 +183,7 @@ GEAK/ │ ├── roles/ knowledge/ scripts/ # gpu_lock.sh, profile_kernel.sh │ └── README.md ├── perf_knowledge/ # AMD operator × backend SOTA knowledge base (REFERENCE ONLY) +├── kb_artifacts/ # machine-produced best patches (code-carrying, gitignored, warm-start source) ├── examples/ # Example kernel tasks, benchmark comparisons, real e2e runs └── exp/ # Experiment outputs (timestamped per run) ``` diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index b16ac683b..cc9cc948b 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -101,6 +101,8 @@ Related timing/tuning args: `fast_head_deadline_ms`, `fast_head_workflow_ms`, `d | `ab_finish_retries` | `3` | A/B leg completion retries. | | `use_expert_skills` | `false` | Consult `perf_knowledge/expert_skills` (advisory priors). OFF = byte-identical. | | `perf_knowledge_dir` | sibling `perf_knowledge/` | Authoring knowledge base. | +| `kb_artifacts_dir` | sibling `kb_artifacts/` | Warm-start patch store, forwarded to every recursive kernel lane so one run shares one store. | +| `warm_start`, `warm_start_match`, `warm_start_min_speedup`, `kb_mode`, `kb_store_dir`, `kb_framework_version` | see the kernel-layer table | Forwarded verbatim to the kernel lanes. Corrective re-authors are forced to `reference` (they must keep the isolated win). | | `time_budget_s`, `initial_extra_server_args`, `initial_extra_env`, `tracelens`, `agent_timeout_ms` | — | Forwarded from the external orchestrator. | ### Example @@ -142,6 +144,13 @@ budget-controlled, each patch independently verified. | `target_language` | `triton` | Author-mode language: `triton` \| `flydsl` \| `hip` \| `ck`. | | `op_spec` | `{}` | Op specification (author mode). | | `perf_knowledge_dir` | sibling `perf_knowledge/` | Knowledge base. | +| `warm_start` | `on` | Local experience reuse from `kb_artifacts/`. `on` = read the top same-arch patches (one per optimization `direction`) through the verify gate; `reference` = prose only, no patch apply; `return_after_read` = apply+validate a candidate then return; `off` = cold start (byte-equivalent to pre-warm-start). | +| `kb_artifacts_dir` | sibling `kb_artifacts/` | Lossless best-patch sink (machine-produced, code-carrying; gitignored, runtime-accumulated). | +| `warm_start_match` | `fuzzy` | How a kernel name reaches a stored page: `exact` (canonicalized name only) \| `normalized` \| `fuzzy` (closest containing page; refuses when two pages fit equally well). Names are canonicalized first, so `fused_moe_kernel_task`, `triton_fused_moe_kernel.py` and a full task path all reach the same page. | +| `warm_start_min_speedup` | `1.05` | Floor on a stored entry's recorded speedup before it is worth a verify slot. | +| `kb_mode` | `local` | Which plane warm start reads and writes. `local` = the curated `kb_artifacts/` tree, addressed by slug. `store` = a KB Store on disk in the shape the KernelForge service uses, addressed by canonical id (`kernel:geak::rocm:::mi355x`); the write records BOTH planes, so they cannot drift. Phases, schemas and the verify gate are identical either way. | +| `kb_store_dir` | sibling `kb_store_local/` | Root of the `store`-mode plane. Ignored when `kb_mode=local`. | +| `kb_framework_version` | measured stack | ROCm `.` for the key's `framework_version` segment. Set it on a box with no `/opt/rocm`, where the measured stack is empty and every record would otherwise file under `unspecified` — splitting one kernel's history across two keys. Affects the key only, never the recorded stack. | | `update_experience` | `on` | Curate ONE distilled card into `kernel_workflow/knowledge/learned/` after a measured win (every lane run; once centrally per bake-off). `off` skips it. The sink is always this workflow's own `learned/` — an e2e-opened lane never writes into `e2e_workflow/knowledge/learned/`. Cards carry a skill-style discovery header (`name`/`description`/`keywords`/`kernels`/`platforms`/`kernel_class`/`regime`) and `learned/INDEX.md` is regenerated from it by `kernel_workflow/scripts/kb.py ... index` (`--check` = fail-if-stale). Numbers on a card are relative only (ratios, % of achievable peak) — never wall-clock. | | `workload_spec_path` | — | Workload-alignment spec; makes the primary metric the time-weighted ratio-of-sums. | | `agent_timeout_ms` | `3600000` | Per-agent timeout (1h). | diff --git a/e2e_workflow/e2e_workflow.js b/e2e_workflow/e2e_workflow.js index 3dcf124ec..33551c110 100644 --- a/e2e_workflow/e2e_workflow.js +++ b/e2e_workflow/e2e_workflow.js @@ -337,6 +337,107 @@ const ACCURACY_INPUTS = (ACCURACY_GATE !== 'none') // index/capability_index.yaml; status/perf in cards are dated evidence, not routing inputs. const KERNEL_KNOWLEDGE_DIR = String(A.perf_knowledge_dir || (WORKFLOW_DIR.replace(/\/[^/]*$/, '') + '/perf_knowledge')).replace(/\/+$/, ''); +// Warm start = the kernel layer's LOCAL experience reuse: before round 1 a lane resolves this kernel's +// own history in kb_artifacts/ and re-validates the top stored patches through the same verify gate as +// a fresh candidate. The knobs live in the kernel layer; e2e only forwards them, so that (a) one store +// is shared by every recursive lane in a run and (b) an orchestrator-supplied store reaches them at all +// (a lane's own default resolves next to its workflow dir and would ignore the e2e arg). The forwarded +// values equal the lane defaults when nothing is passed, so an unparameterized run is unchanged. +const KB_ARTIFACTS_DIR = String(A.kb_artifacts_dir || + (WORKFLOW_DIR.replace(/\/[^/]*$/, '') + '/kb_artifacts')).replace(/\/+$/, ''); +const KB_ARGS = { + kb_artifacts_dir: KB_ARTIFACTS_DIR, + warm_start: String(A.warm_start != null ? A.warm_start : 'on'), + ...(A.warm_start_match != null ? { warm_start_match: String(A.warm_start_match) } : {}), + ...(A.warm_start_min_speedup != null ? { warm_start_min_speedup: A.warm_start_min_speedup } : {}), + // Which plane the lanes read and write. Forwarded like the rest so one run uses one plane; omitted + // when unset, which leaves each lane on its own `local` default. + ...(A.kb_mode != null ? { kb_mode: String(A.kb_mode) } : {}), + ...(A.kb_store_dir != null ? { kb_store_dir: String(A.kb_store_dir) } : {}), + ...(A.kb_framework_version != null ? { kb_framework_version: String(A.kb_framework_version) } : {}), +}; +// Same off-spelling set the kernel layer accepts; an off run must stay off everywhere. +const WARM_START_OFF = ['off', 'false', 'none'].includes(KB_ARGS.warm_start.trim().toLowerCase()); + +// --------------------------------------------------------------------------- +// E2E-LEVEL WARM START — the DEPLOYMENT knowledge base. +// +// KB_ARGS above forwards the KERNEL layer's warm start (per-kernel patches, resolved by the nested +// lane). This block is the layer above it: "has anyone already tuned THIS deployment — this (model, +// gfx, framework, framework_version, precision, tp, isl/osl/conc)?" That question was rediscovered +// from scratch on every run, at a cost of hours of server launches, because nothing here ever asked +// it. Two new modules answer it: Module A reads the store and VALIDATES what it gets on this box +// before believing any of it, and Module B records this run at the end. Everything between them — +// the entire original optimization flow — is untouched. +// +// One knob governs both layers (A.warm_start, already defaulted to 'on' by KB_ARGS), with the same +// vocabulary the kernel lane accepts, so an `off` run is off everywhere: +// on (default) | read the ladder, validate the top candidate on this box, ADOPT if it wins. +// reference | read only; never bench, never adopt — the offer is prose for later agents. +// return_after_read | resolve + validate, then return before the optimization flow runs. +// off/false/none | nothing: no phase, no agent, no log, and every role prompt byte-identical. +const E2E_WARM_START = KB_ARGS.warm_start.trim().toLowerCase() || 'on'; +const E2E_WARM_START_ON = !WARM_START_OFF; +// Fast mode's premise is that all optimization comes from the head track within a wall-clock budget +// (FAST_SKIP already drops 'config'), so a warm-start config bench — a 20-40min server launch that +// this mode has explicitly declined to spend anywhere else — is forced down to reference-only. +const E2E_WARM_START_REF_ONLY = E2E_WARM_START === 'reference' || FAST_MODE; +const E2E_WARM_START_RETURN_AFTER = E2E_WARM_START === 'return_after_read'; +// Which plane the DEPLOYMENT store is read from and written to. Translated from the kernel layer's +// kb_mode vocabulary ONCE, here, so the reader and the writer can never disagree about it. +// local on-disk store only (offline rehearsal; no network, nothing permanent). +// remote the KB Store service only. +// both local first, remote as well — a network failure is reported, not fatal. +// Default `both`: the local plane costs nothing and gives the run a durable copy even when the +// service is unreachable, which on this box it intermittently is. +const E2E_KB_PLANE = ['local', 'remote', 'both'].includes(String(A.e2e_kb_plane || '').trim()) + ? String(A.e2e_kb_plane).trim() + : (String(A.kb_mode || '').trim().toLowerCase() === 'local' ? 'local' : 'both'); +const E2E_KB_STORE_DIR = String(A.kb_store_dir || + KB_ARTIFACTS_DIR.replace(/\/[^/]*$/, '') + '/kb_store_local').replace(/\/+$/, ''); +const E2E_STORE_SCRIPT = `${WORKFLOW_DIR}/scripts/e2e_store.py`; +// Every candidate costs a full server launch to reject, so a recorded near-tie is not worth benching. +const E2E_WARM_START_MIN_SPEEDUP = Number.isFinite(parseFloat(A.warm_start_min_speedup)) + ? parseFloat(A.warm_start_min_speedup) : 1.05; +// Read broadly, bench narrowly. Reading is free and breadth is exactly what makes the demoted +// reference path useful to the Architect; benching is the expensive half. So the default is +// "offer 3, measure 1". +const E2E_WARM_START_TOP_N = parseInt(A.e2e_warm_start_top_n != null ? A.e2e_warm_start_top_n : 3, 10); +const E2E_WARM_START_VALIDATE_N = parseInt( + A.e2e_warm_start_validate_n != null ? A.e2e_warm_start_validate_n : 1, 10); +// How many stored KERNELS get replayed through a fresh integrate A/B. Each one is another two server +// launches, so this is a budget, not a completeness target — the rest stay references. +const E2E_WARM_START_KERNELS_N = parseInt( + A.e2e_warm_start_kernels_n != null ? A.e2e_warm_start_kernels_n : 2, 10); +// Which recorded win-kinds can be REPLAYED from a record alone. `patch` re-applies a diff through the +// overlay; `env`/`flag` re-route the op to an implementation that already exists in this install. +// `authored` cannot: its diff is against the authoring workspace, and rebinding it needs a live seam +// the record has no way to prove still exists here. An unreplayable entry is demoted to a reference, +// never reported as a failure — it is still a real datum about what worked on this deployment. +const E2E_REPLAYABLE_KINDS = new Set(['patch', 'env', 'flag']); +// Which roles are TOLD about the warm start. Deliberately excludes e2e_integrator and +// director:validate — those two produce the run's authoritative numbers, and a stored prior in their +// context is contamination with no upside. +const WARM_START_ROLES = new Set(['system_architect', 'config_tuner']); +// Credentials for every emitted KB command. The service token is NOT present in a non-interactive +// shell (it lives in ~/.bashrc, which such a shell never sources), so each command exports it itself +// from the 0600 file. It is never passed in argv: /proc is world-readable on this box, and the +// service has no revocation story for a leaked key. +// The gateway's internal AMD CA is not in a stock container trust store, so a KB command run inside +// one fails TLS (the old workaround was `curl -k`). DETECT then heal: only when the caller has NOT +// already established trust (SSL_CERT_FILE unset) do we point urllib/requests/curl/node at the first +// readable AMD-root bundle we find. Most callers here run OUTSIDE the warm-start launcher (which sets +// these at `docker run`), so this self-heal is what keeps their KB reachable. Path-only, no CA +// content; overridable with KB_CA_BUNDLE; a no-op when SSL_CERT_FILE is already set or no bundle is +// readable (so CI and already-trusting images are byte-identical). DNS (the host has none +// in-container) is a launch concern, handled with `docker run --add-host`. +const KB_ENV_PRELUDE = + 'export KB_STORE_URL="${KB_STORE_URL:-https://global.primus-safe.amd.com/knowledge-base}"; ' + + 'export KB_STORE_TOKEN="${KB_STORE_TOKEN:-$(cat ~/.geak_kb_token 2>/dev/null)}"; ' + + 'if [ -z "${SSL_CERT_FILE:-}" ]; then for _ca in "${KB_CA_BUNDLE:-}" ' + + '/shared_nfs/hyperloom/ca/amd-ca-combined.pem "$HOME/amd-extra-ca-bundle.pem"; do ' + + '[ -n "$_ca" ] && [ -r "$_ca" ] && { export SSL_CERT_FILE="$_ca" REQUESTS_CA_BUNDLE="$_ca" ' + + 'CURL_CA_BUNDLE="$_ca" NODE_EXTRA_CA_CERTS="$_ca"; break; }; done; fi; '; // Expert skills = human-authored, validated optimization recipes (perf_knowledge/expert_skills/). They // are ADVISORY priors: a matched `validated` skill is a HIGH-PRIOR candidate that routing/integration // roles reproduce, then gate by the usual on-box A/B — it NEVER overrides measurement and NEVER reduces @@ -444,6 +545,14 @@ const SETUP_SCHEMA = obj({ server_flags: { type: 'object', additionalProperties: true }, server_env: { type: 'string' }, tp: { type: 'number' }, workload: { type: 'object', additionalProperties: true }, bench_script: { type: 'string' }, notes: { type: 'string' }, + // The four deployment dimensions the KB addresses a page by, beyond model/framework/workload which + // this script already knows. The Director establishes all four during its existing preflight (it + // has to, to launch the server at all) and previously just discarded them. They are OPTIONAL here + // and `required` is unchanged, so a director that returns none of them is still valid — the run + // simply files itself under a coarse `unknown` page, which is recoverable. A GUESS is not: the + // service has no DELETE, so a wrong-but-authoritative-looking page is permanent. + gfx: { type: 'string' }, precision: { type: 'string' }, + framework_version: { type: 'string' }, rocm_version: { type: 'string' }, }, ['eval_dir', 'baseline_throughput_tok_s']); const PROFILE_SCHEMA = obj({ @@ -472,6 +581,22 @@ const SWEEP_SCHEMA = obj({ summary: { type: 'string' }, }, ['accepted_flags', 'best_throughput_tok_s']); +// `e2e_store.py resolve` output, passed through VERBATIM. Nothing is `required` and no field is +// typed as a number: a page that was never written answers with empty candidates and a +// `read_reason`, and a recorded run legitimately has `speedup: null` when it only ever measured +// absolute throughput. Typing either would turn a truthful answer into a schema failure and a +// cold start. +const KB_RESOLVE_SCHEMA = obj({ + tried: arrStr, canonical_id: { type: 'string' }, match_tier: { type: 'string' }, + ranked_by: { type: 'string' }, candidates: arrObj, read_reason: { type: 'string' }, + plane: { type: 'string' }, + // How the offer was ORDERED (throughput, high to low, on every rung) versus which metric that + // rung crowns its champion on. Two different things that used to share one word: a coarse rung + // is ranked by absolute throughput here but still promotes on speedup, and a reader told only + // 'speedup' would mis-explain the order it is looking at. + sorted_by: { type: 'string' }, champion_metric: { type: 'string' }, +}, []); + const PLAN_SCHEMA = obj({ stop: { type: 'boolean' }, reasoning: { type: 'string' }, config_directions: arrObj, head_candidates: arrObj, kernel_candidates: arrObj, @@ -587,6 +712,80 @@ function expertSkillsBlock(role) { `only, never overriding your on-box A/B, never reducing a result below the measured baseline.`; } +// --------------------------------------------------------------------------- +// e2e warm start: identity formatting + the prompt hook. +// +// These three `let`s are the ONLY state the feature adds above Module A, and all three stay at their +// off-value (null / '') for the whole run unless Module A actually executes. Every consumer below +// keys off them, which is what makes `warm_start=off` produce a byte-identical run rather than one +// that merely behaves the same. +// --------------------------------------------------------------------------- +let KB_DIMS = null; // the deployment dimensions the Director established during preflight +let KB_REF_DIR = ''; // where Module A left its references; '' when it never ran +let KB_REF_VERDICT = ''; // what we MEASURED about the offer, threaded into later roles' Inputs +let KB_READ_PLANE = ''; // which plane ANSWERED the read, which `both` alone does not tell you + +const shq = (s) => "'" + String(s == null ? '' : s).replace(/'/g, "'\\''") + "'"; + +// The ONE place e2e KB identity argv is formatted — called by both the reader (Module A) and the +// writer (Module B). kb/identity.py's own header names the failure this prevents: a reader and a +// writer that disagree by a single segment do not raise, they address two different pages, and the +// only symptom is that history quietly stops existing. Two call sites formatting the same flags +// independently is exactly how that drift starts. +function kbIdentityFlags() { + if (!KB_DIMS) return ''; + return [`--model ${shq(KB_DIMS.model)}`, `--gfx ${shq(KB_DIMS.gfx)}`, + `--framework ${shq(BACKEND)}`, `--framework-version ${shq(KB_DIMS.framework_version)}`, + `--precision ${shq(KB_DIMS.precision)}`, `--rocm-version ${shq(KB_DIMS.rocm_version)}`, + `--tp ${SERVING_TP}`, `--isl ${ISL}`, `--osl ${OSL}`, `--conc ${CONC}`].join(' '); +} + +// --store is what the local plane writes into and is meaningless to a remote-only run; open_plane() +// hard-errors without it on plane local|both, so it is not optional there. +function kbPlaneFlags(plane) { + plane = plane || E2E_KB_PLANE; + return `--plane ${plane}` + (plane === 'remote' ? '' : ` --store ${shq(E2E_KB_STORE_DIR)}`); +} + +// A READ takes exactly one plane — `open_plane()` returns (local, remote) for `both` and cmd_resolve +// uses only the first, deliberately: merging two rankings needs a cross-plane comparability rule +// that nothing here has, and silently preferring one would let a stale local mirror shadow the +// service without saying so. So `--plane both` on a read means LOCAL, which is the opposite of what +// this workflow wants. The choice is made here instead, in the open: try the service, and fall back +// to the local store only when it has no answer. `read_reason` in the returned JSON says which one +// spoke. The credentials cannot be tested from this process (non-interactive shells never source the +// profile that sets them), so the branch lives in the emitted bash, after the prelude has exported +// whatever the box actually has. +function kbResolveScript(args) { + const invoke = (plane) => + `python3 ${shq(E2E_STORE_SCRIPT)} resolve ${kbIdentityFlags()} \\\n` + + ` ${kbPlaneFlags(plane)} ${args}`; + if (E2E_KB_PLANE !== 'both') return KB_ENV_PRELUDE + '\\\n' + invoke(E2E_KB_PLANE); + return KB_ENV_PRELUDE + ` +REMOTE_OUT='' +if [ -n "$KB_STORE_TOKEN" ]; then + REMOTE_OUT=$(${invoke('remote')} 2>/dev/null || true) + if printf '%s' "$REMOTE_OUT" | python3 -c 'import json,sys; sys.exit(0 if (json.load(sys.stdin).get("candidates") or []) else 1)' 2>/dev/null; then + printf '%s\\n' "$REMOTE_OUT"; exit 0 + fi +fi +${invoke('local')}`; +} + +// Warm-start prompt injection, mirroring expertSkillsBlock exactly: returns '' whenever the feature +// is off or the role is not a consumer, so roleAgent's output is byte-identical to the pre-feature +// build in those cases. KB_REF_DIR is assigned by Module A and nowhere else, so an off run — or one +// whose read found nothing — never reaches the template at all. +function warmStartBlock(role) { + if (!KB_REF_DIR || !WARM_START_ROLES.has(role)) return ''; + return `\n\n## Warm start (a PRIOR run's record — already measured on this box)\n` + + `Also Read ${WORKFLOW_DIR}/roles/_fragments/warm_start.md and follow it. The knowledge base ` + + `offered prior configurations for this exact deployment; they are in ${KB_REF_DIR}/, and ` + + `${KB_REF_DIR}/measured_on_this_box.md records what happened when THIS run benched them — that ` + + `file OVERRIDES the stored claims in its siblings wherever the two disagree. A stored number is ` + + `a hypothesis; only the measured column is evidence.`; +} + function roleAgent(role, phase, intro, inputs) { // BACKEND is injected for every role: any role that calls bench_e2e.sh must forward it // (BACKEND=) so the right serving adapter (scripts/adapters/.sh) is used. @@ -615,7 +814,7 @@ optimization-pool id for a serving launch — keep the two separate. ${cfg(inall)} Return ONLY the structured JSON the role file specifies (a StructuredOutput tool is forced).`; - return base + expertSkillsBlock(role); + return base + expertSkillsBlock(role) + warmStartBlock(role); } // Resilient agent wrapper: a single agent failure (transient API 502 / didn't emit StructuredOutput) @@ -777,6 +976,37 @@ async function extractWithBaseline(role, phase, intro, inputs, opts) { // verdict. gate:'incomplete' or ab_complete:false means a leg is still missing. const abDone = (integ) => !!(integ && integ.gate !== 'incomplete' && integ.ab_complete !== false); +// The KERNEL_RESULT that a given integrate A/B actually gated. Both patch spellings are read at the +// call sites because the two tracks fill different ones (head=code_patch, milestone=final_patch) and +// tryCorrectiveReauthor deliberately sets BOTH to the corrected patch. +const krOf = (inputs) => (inputs && inputs.KERNEL_RESULT) || {}; + +// Bank an accepted win, carrying the fields a KB reader needs to REPLAY it on another box. +// e2e_store.py:_accepted_kernels starts from dict(item) so anything here survives verbatim into the +// record — but it LOOKS UP `name`, `language`, `isolated_speedup` and `patch`, while this workflow +// only ever pushed `short_name`, `backend` and `isolated`. Every accepted kernel was therefore +// matched by nothing and silently dropped on the way into the store, which is why both e2e records +// already in the remote KB were written with accepted_kernels:[] and no patch at all. The old +// spellings stay (the report, the resume state and run_e2e.py all read them); the aliases are added +// alongside. `e.patch` wins over the KERNEL_RESULT's because a corrective re-author supplies the +// FIXED patch and the gated one is the broken kernel. +const bankAccepted = (list, e, kr) => { + const k = kr || {}; + list.push({ + ...e, + name: e.short_name || '', + language: e.backend || '', + isolated_speedup: Number(e.isolated) || 0, + winner_kind: e.kind || k.winner_kind || '', + patch: e.patch || k.code_patch || k.final_patch || '', + target_callable: k.target_callable || '', + source_path_in_sglang: k.source_path_in_sglang || '', + apply_env: e.apply_env || k.apply_env || '', + apply_flags: e.apply_flags || k.apply_flags || '', + pct_gpu_time: e.pct_gpu_time != null ? e.pct_gpu_time : (k.pct_gpu_time || 0), + }); +}; + // Run ONE integrate A/B and GUARANTEE both legs complete. The first call does a // normal apply+gate; if the integrator returns incomplete (ran only ref, hung, // or degraded mid-A/B) we RE-INVOKE it in resume mode to run the MISSING leg, @@ -903,6 +1133,9 @@ async function tryCorrectiveReauthor(spec) { op_spec: { op_kind: spec.op_kind, shapes: spec.shapes || {}, dtype: spec.dtype || 'bf16', regime: spec.regime || '', cuda_graph_safe: true, ...(spec.workload_path ? { workload_path: spec.workload_path } : {}) }, perf_knowledge_dir: KERNEL_KNOWLEDGE_DIR, use_expert_skills: USE_EXPERT_SKILLS ? 'true' : 'false', expert_skills_dir: EXPERT_SKILLS_DIR, + // A corrective must KEEP the isolated win; adopting a stored patch over it would be exactly the + // re-discovery the task forbids. History is still readable, just never auto-applied. + ...KB_ARGS, warm_start: WARM_START_OFF ? KB_ARGS.warm_start : 'reference', budget: KERNEL_BUDGET, gpu_ids: spec.gpu_id, exp_root: `${EVAL_DIR}/kernels/_exp`, task: `CORRECTIVE FIX — do NOT re-discover the algorithm; KEEP the ${(spec.isolated || 0).toFixed(2)}x isolated win. ` + `This kernel PASSED the isolated oracle and ENGAGED on all live workers but was REJECTED at the e2e serving gate ` + @@ -943,7 +1176,11 @@ async function tryCorrectiveReauthor(spec) { const implausible2 = ab2 && (integ2.gate === 'accepted' || integ2.gate === 'stack') && isImplausibleSpeedup(pctForGuard, fix.final_geomean, integ2); if (ab2 && (integ2.gate === 'accepted' || integ2.gate === 'stack') && integ2.e2e_throughput_tok_s > curTput && !implausible2) { - return { banked: true, integ: integ2, isolated: fix.final_geomean }; + // Carry the CORRECTED patch out to the caller's bankAccepted. The gated KERNEL_RESULT at the + // call site still holds the broken kernel, and recording that into the KB would publish a + // patch that was rejected on this very box as if it were the win. + return { banked: true, integ: integ2, isolated: fix.final_geomean, patch: fix.final_patch, + kernel_eval_dir: fix.eval_dir || spec.kernel_eval_dir || '' }; } reason = implausible2 ? `implausible_speedup (+${(integ2.e2e_delta_pct || 0).toFixed(1)}% >> Amdahl ceiling +${amdahlCeilingPct(pctForGuard, fix.final_geomean).toFixed(1)}% — corruption)` @@ -1064,7 +1301,11 @@ if (!MODEL_PATH && KERNEL_PATH) { try { const r = await workflow({ scriptPath: KERNEL_WF_SCRIPT }, laneArgs({ kernel_path: KERNEL_PATH, workflow_dir: KERNEL_WF_DIR, + // The lane defaults to triton; a pass-through caller naming another backend must reach that + // backend's own warm-start page (the slug is per-language), not triton's. + ...(A.target_language ? { target_language: String(A.target_language) } : {}), use_expert_skills: USE_EXPERT_SKILLS ? 'true' : 'false', expert_skills_dir: EXPERT_SKILLS_DIR, + ...KB_ARGS, budget: KERNEL_BUDGET, gpu_ids: GPU_IDS, task: TASK, exp_root: EXP_ROOT, apply_to_original: APPLY_TO_ORIGINAL, })); @@ -1082,6 +1323,19 @@ if (!MODEL_PATH && KERNEL_PATH) { // =========================================================================== // PHASE: Setup + Baseline profile + Strategize (gated; else load carried state) // =========================================================================== +// Module A's outputs. They are folded into the ORDINARY state variables at those variables' real +// declaration sites further down — `curTput` (~1204), `curOverlay` / `acceptedKernels` (~1248) — and +// not assigned here, because all three are in the temporal dead zone at Module A's insertion point +// and Module A must run BEFORE Profile so the Top-N is taken on the adopted config. Their off-values +// are 0 / '' / [], which is exactly what makes each fold below a no-op when the feature is off. +let kbSeedTput = 0; +let kbSeedOverlay = ''; +const kbSeedKernels = []; +// Spread into the Strategize and ConfigSweep inputs. Stays `{}` unless Module A measured something: +// `{...{}}` contributes no own-properties and preserves key order, so those two prompts are +// byte-identical to the pre-feature build on an off run. +let KB_REF_INPUTS = {}; + let EVAL_DIR, MODEL_NAME, BASELINE_TPUT, NOISE_BAND, curFlags, curEnv, profile, strategy, kernelQueue, headQueue; if (want('setup')) { phase('Setup'); @@ -1102,11 +1356,489 @@ if (want('setup')) { curEnv = INIT_ENV || (setup.server_env || ''); log(`Setup done. EVAL_DIR=${EVAL_DIR}, baseline ${BASELINE_TPUT} tok/s (noise band ${NOISE_BAND}%)`); + // ========================================================================= + // MODULE A — read the deployment KB, then VALIDATE what it says on this box. + // + // Placed HERE, and not earlier, for three reasons that are each load-bearing: + // * EVAL_DIR, bench_e2e.sh, BASELINE_TPUT and NOISE_BAND do not exist until Setup returns, and a + // candidate cannot be judged without a baseline to judge it against. + // * it is before Profile, so the Top-N is captured on the ADOPTED config. This is the same + // discipline the flow already asserts by re-profiling after ConfigSweep: a config change moves + // which kernels dominate, and routing off a stale profile optimizes the wrong ops. + // * it is INSIDE want('setup'), so a phase-partial resume cold-starts automatically rather than + // re-reading and re-benching a store on every phase invocation. + // + // Nothing below is trusted on the store's word. A stored speedup was measured on another box, on + // another day, against another baseline; the only thing that makes it actionable here is that this + // run re-measured it through the same gate a fresh idea would face. + // ========================================================================= + if (E2E_WARM_START_ON) { + phase('WarmStart'); + KB_DIMS = { + model: MODEL_NAME, + gfx: String(setup.gfx || '').trim(), + framework_version: String(setup.framework_version || A.kb_framework_version || '').trim(), + precision: String(setup.precision || '').trim(), + rocm_version: String(setup.rocm_version || '').trim(), + }; + if (!KB_DIMS.gfx) { + // Refuse to read rather than read the wrong page. An arch-less identity resolves to the + // `gfx=unknown` rung, whose records were measured on hardware we cannot establish — and the + // cost of believing one is a 20-40min server launch, not a wasted lookup. + log('[kb] warm start skipped: read_reason=missing_arch (the Director reported no gfx; a ' + + 'cross-arch config is not a candidate, it is a guess).'); + } else { + const refsDir = `${EVAL_DIR}/kb_references`; + const cacheDir = `${EVAL_DIR}/kb_cache`; + const resolved = await safeAgent( + `You are the e2e warm-start resolver. Run EXACTLY this command and return its JSON stdout ` + + `verbatim as StructuredOutput — do not add, drop, reorder, or reinterpret any field. The ` + + `command pretty-prints its JSON over several lines; return the whole object, not the first ` + + `line. A non-zero exit or an empty candidate list is a VALID answer (the page has never been ` + + `written): return what it printed, do not retry with different flags, and do not invent ` + + `candidates. It may consult the shared KB Store service first and fall back to the on-disk ` + + `store by itself; run it as one script and do not split it into separate commands.\n` + + '```bash\n' + + kbResolveScript( + `--top-n ${E2E_WARM_START_TOP_N} \\\n` + + ` --min-speedup ${E2E_WARM_START_MIN_SPEEDUP} \\\n` + + ` --refs-dir ${shq(refsDir)} --cache-dir ${shq(cacheDir)}`) + '\n' + + '```', + { phase: 'WarmStart', label: 'warm_start:resolve', schema: KB_RESOLVE_SCHEMA }) || {}; + const cands = Array.isArray(resolved.candidates) ? resolved.candidates : []; + // Log the ladder VERBATIM. On a scheme with no search, "never recorded" and "recorded under an + // address one segment different" are the same 404, and this line is the only record of which + // question was actually asked — without it a silent identity drift looks like an empty store. + KB_READ_PLANE = String(resolved.plane || ''); + log(`[kb] e2e read: plane=${resolved.plane || '?'} tried=[${(resolved.tried || []).join(' | ')}] ` + + `answered=${resolved.canonical_id || '?'} tier=${resolved.match_tier || '-'} ` + + `sorted_by=${resolved.sorted_by || resolved.ranked_by || '-'} ` + + `champion_metric=${resolved.champion_metric || '-'} reason=${resolved.read_reason || '?'} ` + + `candidates=${cands.length}`); + if (cands.length) KB_REF_DIR = refsDir; // arms warmStartBlock() for the consumer roles + + // How many of the offers are worth a server launch. On a coarser rung the stored numbers were + // measured on a DIFFERENT workload point, so they are not comparable to this run's baseline at + // all — the configs are still ideas worth handing to the Architect, but benching one on the + // strength of a non-comparable number is spending 30min to test a coin flip. Zero there + // unless the caller explicitly asks otherwise. + const exactTier = (resolved.match_tier || '') === 'exact'; + const benchN = E2E_WARM_START_REF_ONLY ? 0 + : (exactTier || A.e2e_warm_start_validate_n != null ? E2E_WARM_START_VALIDATE_N : 0); + if (cands.length && !benchN) { + log(`[kb] not benching: ${E2E_WARM_START_REF_ONLY + ? (FAST_MODE ? 'fast mode — all optimization comes from the head track' : 'warm_start=reference') + : `match tier '${resolved.match_tier}' is not exact, so the stored numbers are not ` + + 'comparable to this baseline'}. The offers stay as references.`); + } + + const verdicts = []; + for (const c of cands.slice(0, benchN)) { + // VALIDATE THROUGH THE ORIGINAL GATE. This is config_tuner:sweep — the same role, the same + // schema, the same bench_e2e.sh at the same TP/GPU, the same delta-vs-median arithmetic, the + // same parity check and the same swap-took-effect log grep the flow already trusts for a + // fresh idea. Reusing it rather than writing a warm-start harness is the whole reason a + // stored config and a proposed config are judged by identical evidence. + const stored = (c.accepted_config && typeof c.accepted_config === 'object') ? c.accepted_config : {}; + const storedFlags = String(stored.flags || ''); + const storedEnv = String(stored.env || ''); + if (!storedFlags && !storedEnv) { + verdicts.push({ ...c, measured_tok_s: null, delta_pct: null, parity: 'n/a', + outcome: 'skipped', why: 'the record carries no config to apply' }); + continue; + } + const sweep = await safeAgent( + roleAgent('config_tuner', 'sweep', + 'Validate ONE historical configuration recovered from the knowledge base. Treat it exactly ' + + 'as you would a fresh direction: same A/B, same repeats, same parity check, same ' + + 'swap-took-effect verification. TWO deviations from your role file, both deliberate:\n' + + '(1) Do NOT decompose this direction into one-axis-at-a-time trials. A stored config is an ' + + 'ALREADY-COMPOUNDED whole that was accepted together on another box; benching its knobs ' + + 'separately measures a configuration nobody has ever run, and the parts can each be ' + + 'neutral while the whole is a win (or the reverse). Apply all of it, once, as a single ' + + 'trial.\n' + + '(2) Verify the swap TOOK EFFECT before you believe a null result. This config came from a ' + + 'different framework_version: a flag that was renamed or removed upstream is accepted ' + + 'silently on the command line and then ignored, which is indistinguishable from "the ' + + 'config made no difference". Grep the server log for each flag/env actually being ' + + 'honoured, and if one is not, say so in the trial notes rather than reporting a clean no-op.', + { + EVAL_DIR, MODEL_PATH, GPU_ID: GPU_LIST[0], WORKLOAD, BASELINE_THROUGHPUT: BASELINE_TPUT, + NOISE_BAND_PCT: NOISE_BAND, E2E_REPEATS, + CONFIG_DIRECTIONS: [{ + rank: 1, + direction: 'kb_warm_start:' + (c.direction || 'unlabeled'), + axis: 'compound (recovered configuration — do not split)', + flags: storedFlags, env: storedEnv, + rationale: `Recorded under ${c.canonical_id || resolved.canonical_id} (session ` + + `${c.session_id || '?'}), where it measured ${c.throughput_tok_s != null ? c.throughput_tok_s : '?'}` + + ` tok/s (${c.speedup != null ? c.speedup + 'x' : 'speedup not recorded'}) against a ` + + `baseline of ${c.baseline_throughput_tok_s != null ? c.baseline_throughput_tok_s : '?'} tok/s. ` + + 'That number is a HYPOTHESIS about this box, not a measurement of it.', + }], + CURRENT_FLAGS: curFlags, CURRENT_ENV: curEnv, SKILL_DIR: WORKFLOW_DIR, + }), + { phase: 'WarmStart', label: `warm_start:validate:${c.session_id || 'cand'}`, schema: SWEEP_SCHEMA }); + const trial = (sweep && (sweep.trials || [])[0]) || {}; + const measured = (sweep && sweep.best_throughput_tok_s) || 0; + const parity = String(trial.parity || trial.output_parity || ''); + const deltaPct = BASELINE_TPUT ? ((measured - BASELINE_TPUT) / BASELINE_TPUT) * 100 : 0; + // Both the tuner's own judgement AND our arithmetic. A warm-start candidate is precisely + // where an agent is most tempted to ratify a stored claim it did not establish, so the + // orchestrator re-derives the ratio from the raw number rather than accepting `kept: true`. + const accept = trial.kept === true && measured > BASELINE_TPUT && + deltaPct > NOISE_BAND && parity !== 'fail'; + if (accept) { + curFlags = sweep.accepted_flags || storedFlags || curFlags; + curEnv = sweep.accepted_env || storedEnv || curEnv; + kbSeedTput = measured; + log(`[kb] ADOPTED ${c.session_id || '?'} (${c.direction || 'unlabeled'}): ` + + `${measured} tok/s, +${deltaPct.toFixed(2)}% vs baseline ${BASELINE_TPUT} (noise band ${NOISE_BAND}%).`); + } + // THREE outcomes, not two. "ran and lost" and "could not be made to run" were both spelled + // `rejected`, and they mean opposite things to the record: a loss is one box's verdict on a + // real configuration, while a config that never took effect says the record is missing + // something — a flag renamed upstream, an env the build does not honour. Only the second is + // evidence for retiring it, and the store now counts them separately (kb/attest.py). + const notes = String(trial.notes || (sweep && sweep.summary) || ''); + const inert = /\b(?:not (?:honou?red|recognized|recognised|applied|supported)|unrecognized|unrecognised|ignored|no such option|unknown (?:option|argument|flag)|renamed|removed upstream|did not take effect)\b/i.test(notes); + const outcome = accept ? 'adopted' : (!measured || inert) ? 'not_reproduced' : 'rejected'; + if (!accept) { + log(`[kb] ${outcome === 'not_reproduced' ? 'NOT REPRODUCED' : 'rejected'} ` + + `${c.session_id || '?'} (${c.direction || 'unlabeled'}): ` + + `measured ${measured || 'n/a'} tok/s vs baseline ${BASELINE_TPUT}` + + `${measured ? ` (${deltaPct >= 0 ? '+' : ''}${deltaPct.toFixed(2)}%)` : ''}` + + `${parity ? `, parity=${parity}` : ''}${inert ? ', the config never took effect' : ''}` + + ` — kept as a reference, not applied.`); + } + verdicts.push({ ...c, measured_tok_s: measured || null, delta_pct: measured ? deltaPct : null, + parity: parity || 'unknown', outcome, why: notes }); + if (accept) break; // adopt the first that passes; the rest stay references + } + // --------------------------------------------------------------------- + // Replay the stored KERNELS through the ordinary integrate gate. + // + // Same machinery as the Milestone track, verbatim: runIntegrateBothLegs, the same + // INTEGRATE_SCHEMA, the same two-launch A/B with the parity probe, the same integAccepted + // predicate. The only thing that differs is where the patch came from, and that is exactly the + // thing the A/B is there to make irrelevant. + // + // Reverting is structural rather than an operation: a candidate lives in its own overlay + // directory and is activated only by being on PYTHONPATH, so rejecting one means not adopting + // it. Nothing is ever mutated in the install. + // --------------------------------------------------------------------- + const kbKernels = []; + const seenKernel = new Set(); + for (const c of cands) { + const bundlePath = (c.bundle && c.bundle.path) || ''; + for (const k of (Array.isArray(c.accepted_kernels) ? c.accepted_kernels : [])) { + const name = String((k && (k.name || k.short_name)) || '').trim(); + if (!name || seenKernel.has(name)) continue; // one op recorded by several runs is one op + seenKernel.add(name); + kbKernels.push({ ...k, name, bundle: bundlePath, kb_session_id: c.session_id || '' }); + } + } + let replayed = 0; + const kernelVerdicts = []; + for (const k of kbKernels) { + const kind = String(k.winner_kind || k.kind || '').trim().toLowerCase(); + // `patch` is stored inside the record's own bundle as `kernels/.patch`, so the path only + // resolves if the artifacts actually came down. A record whose manifest was never committed + // downloads a knowledge document and no files, and handing the integrator a path that opens + // nothing would burn two server launches to discover it. + const patchPath = (kind === 'patch' && k.patch && k.bundle) + ? `${k.bundle}/files/${k.patch}` : ''; + // An env/flag win is only replayable if the record actually says WHICH env or flag routed the + // op. Several records in the store name the kind but carry empty apply_env/apply_flags — they + // were written before those fields were banked. Sending one to the integrator produces an A/B + // between two identical configurations: two full server launches to measure nothing, and a + // 'rejected' verdict that then libels a win which may well have been real. + const hasRouting = !!(String(k.apply_env || '').trim() || String(k.apply_flags || '').trim()); + const unreplayable = + E2E_WARM_START_REF_ONLY ? 'read-only mode' + : !E2E_REPLAYABLE_KINDS.has(kind) ? `winner_kind '${kind || 'unrecorded'}' cannot be replayed from a record` + : (kind === 'patch' && !patchPath) ? 'the record names a patch but its bundle holds no file' + : (kind !== 'patch' && !hasRouting) ? `recorded as a '${kind}' win but carries no apply_env/apply_flags, so there is nothing to re-apply` + : replayed >= E2E_WARM_START_KERNELS_N ? `replay budget of ${E2E_WARM_START_KERNELS_N} already spent` + : ''; + if (unreplayable) { + kernelVerdicts.push({ ...k, kind, outcome: 'reference', measured_delta_pct: null, + why: unreplayable }); + continue; + } + replayed++; + const isolated = Number(k.isolated_speedup) || 0; + const kbIntegrateInputs = { + EVAL_DIR, MODEL_PATH, GPU_ID: GPU_LIST[0], WORKLOAD, NOISE_BAND_PCT: NOISE_BAND, E2E_REPEATS, + KERNEL_RESULT: { + short_name: k.name, winner_kind: kind, + // Both spellings, because the two tracks read different ones and this synthetic result + // has to satisfy whichever the integrator reaches for. + code_patch: patchPath, final_patch: patchPath, + apply_env: String(k.apply_env || ''), apply_flags: String(k.apply_flags || ''), + target_callable: String(k.target_callable || ''), + source_path_in_sglang: String(k.source_path_in_sglang || ''), + verified_isolated_speedup: isolated, pct_gpu_time: Number(k.pct_gpu_time) || 0, + // task_dir is deliberately EMPTY and the provenance is declared foreign — see the intro. + task_dir: '', provenance: 'knowledge_base_replay', + }, + CURRENT_OVERLAY: kbSeedOverlay, CURRENT_FLAGS: curFlags, CURRENT_ENV: curEnv, + CURRENT_THROUGHPUT: kbSeedTput || BASELINE_TPUT, SKILL_DIR: WORKFLOW_DIR, + }; + const integ = await runIntegrateBothLegs( + 'Overlay a kernel RECOVERED FROM THE KNOWLEDGE BASE and gate it on e2e throughput. Run your ' + + 'normal A/B — same two launches, same repeats, same parity probe. Two things are different ' + + 'and you must honour both:\n' + + '(1) There is NO task_dir and therefore NO immutable oracle for this kernel. Your step-1 ' + + 'provenance re-check cannot run, because the workspace that produced this patch does not ' + + 'exist on this box. Do NOT claim provenance was verified. The FRESH parity probe against the ' + + 'reference leg is the only correctness evidence available here, so run it and report its ' + + 'result plainly — if parity fails or cannot be established, REJECT.\n' + + '(2) verified_isolated_speedup is a FOREIGN number, measured on another box against another ' + + 'baseline. It is context for what to expect, never evidence. Gate on what YOU measure.\n' + + 'As always: never mutate the installed package. The overlay on PYTHONPATH is the only ' + + 'mechanism, so a rejection costs nothing to undo — you simply do not adopt the directory.', + kbIntegrateInputs, `warm_start integrate ${k.name}`, 'WarmStart'); + const base = kbSeedTput || BASELINE_TPUT; + if (abDone(integ) && integAccepted(integ, Number(k.pct_gpu_time) || 0, isolated) && + integ.e2e_throughput_tok_s > base) { + kbSeedOverlay = integ.accepted_overlay || kbSeedOverlay; + kbSeedTput = integ.e2e_throughput_tok_s; + bankAccepted(kbSeedKernels, { + short_name: k.name, backend: String(k.language || k.backend || ''), kind, + e2e_delta_pct: integ.e2e_delta_pct, isolated, patch: patchPath, + pct_gpu_time: Number(k.pct_gpu_time) || 0, + apply_env: String(k.apply_env || ''), apply_flags: String(k.apply_flags || ''), + // Provenance stays ON the banked entry: this kernel is a recovered win re-verified here, + // not something this run discovered, and Module B must not later claim otherwise. + from_knowledge_base: true, kb_session_id: k.kb_session_id, + }, krOf(kbIntegrateInputs)); + kernelVerdicts.push({ ...k, kind, outcome: 'adopted', + measured_delta_pct: integ.e2e_delta_pct, why: integ.reason || '' }); + log(`[kb] ADOPTED kernel ${k.name} (${kind}): e2e now ${kbSeedTput} tok/s (+${integ.e2e_delta_pct}%).`); + } else { + const reason = gateRejectReason(integ, Number(k.pct_gpu_time) || 0, isolated); + kernelVerdicts.push({ ...k, kind, outcome: abDone(integ) ? 'rejected' : 'incomplete', + measured_delta_pct: integ ? integ.e2e_delta_pct : null, + why: reason || 'A/B did not complete' }); + // No corrective re-author here. That path re-optimizes a kernel this run authored and owns; + // a stored patch that fails its fresh gate is simply not this box's win, and the honest + // outcome is to hand it to the Architect as a lead rather than to repair someone else's diff. + log(`[kb] kernel ${k.name} (${kind}) not adopted: ${reason || 'A/B did not complete'} — kept as a reference.`); + } + } + // --------------------------------------------------------------------- + // Demote what did not survive into a REFERENCE for the roles that come next. + // + // The resolver's own `e2e_reference_.md` is deliberately left alone. It is the pristine + // record of what the STORE claimed, and the disagreement between that and what this box + // measured is the interesting datum — editing the claim to match the measurement destroys the + // only evidence that the two differ. So the measurements go in a sibling file that says, in as + // many words, that it overrides its neighbours. + // + // Three channels, because the first two can each be missed. The reference DIRECTORY is only + // read if an agent chooses to; the prompt BLOCK only reaches roles in WARM_START_ROLES; the + // INPUTS entry lands in `## Inputs` unconditionally and is the one that always fires. + // --------------------------------------------------------------------- + // --------------------------------------------------------------------- + // Tell the STORE what happened. Until now this loop's verdicts died with the run: the next + // box to resolve this identity saw the same optimistic record, benched it, and failed the + // same way, forever. `e2e_store.py attest` counts the attempt onto the record itself, at + // every rung, so `validations / recalls` becomes something a later curation pass can retire + // on. It moves no score and no champion — one box's failure is evidence, not a verdict. + // + // Only candidates that were actually PUT ON THIS BOX are counted. A record listed in the + // offer and never benched (`skipped`, or below benchN) has learned nothing about itself, and + // counting it would decay the very ratio the retire pass reads. + // + // Config verdicts only. A replayed kernel that failed is evidence against the KERNEL lane's + // record, not against the e2e run that once used it, and the two have separate ledgers — + // `experience_store.py attest` is where that verdict belongs. + const attestable = verdicts.filter(v => v.session_id && + ['adopted', 'rejected', 'not_reproduced'].includes(v.outcome)); + if (attestable.length) { + const cmds = attestable.map(v => + `python3 ${shq(E2E_STORE_SCRIPT)} attest ${kbIdentityFlags()} ${kbPlaneFlags()} ` + + `--session-id ${shq(v.session_id)} ` + + `--outcome ${v.outcome === 'adopted' ? 'validated' : v.outcome} ` + + (v.measured_tok_s ? `--measured-tok-s ${v.measured_tok_s} ` : '') + + `--baseline-tok-s ${BASELINE_TPUT} --parity ${shq(v.parity || 'n/a')} ` + + `--note ${shq(String(v.why || '').slice(0, 200))} ` + + `--measured-by ${shq('e2e_workflow:' + BACKEND)} --apply || true`); + try { + await safeAgent( + `You are the e2e knowledge-base attestor. Run EXACTLY these commands in order and ` + + `return {"ran": , "note": ""}. Each records ` + + `what THIS box saw when it benched a stored record. Do NOT edit them, do NOT add or ` + + `drop any, and do NOT retry a failure — a repeat would double-count the attempt, and ` + + `an over-counted failure retires a record that may still be right elsewhere.\n` + + '```bash\n' + KB_ENV_PRELUDE + cmds.join('\n') + '\n```', + { phase: 'WarmStart', label: 'kb:attest', + schema: obj({ ran: { type: 'number' }, note: { type: 'string' } }, []) }, + 1); + log(`[kb] attested ${attestable.length} benched record(s): ` + + attestable.map(v => `${v.session_id.slice(-12)}=${v.outcome}`).join(' ')); + } catch (e) { + // Non-fatal, like every other KB write. The measurements are already on disk and already + // in the reference file; losing the counter costs a future reader context, not this run. + log(`[kb] attest failed (NON-FATAL): ${String(e).slice(0, 200)}`); + } + } + + const allVerdicts = verdicts.concat(kernelVerdicts); + if (allVerdicts.length) { + const adoptedCfg = verdicts.filter(v => v.outcome === 'adopted'); + const adoptedKer = kernelVerdicts.filter(v => v.outcome === 'adopted'); + const rejected = allVerdicts.filter(v => v.outcome !== 'adopted'); + // Split out of `rejected` for the reference section below, but deliberately still counted + // inside it: for the roles that come next, "lost the A/B" and "never ran" are both "do not + // re-propose this verbatim". The distinction matters to the STORE, not to the Architect. + const notReproduced = verdicts.filter(v => v.outcome === 'not_reproduced'); + const md = [ + '# Warm start — MEASURED ON THIS BOX', + '', + 'This file OVERRIDES the stored claims in its sibling `e2e_reference_*.md` wherever the two', + 'disagree. Those files record what another box reported; this one records what happened when', + 'this run applied the same thing here, through the same gate a fresh idea faces.', + '', + `- baseline: **${BASELINE_TPUT} tok/s** (noise band ${NOISE_BAND}%)`, + `- serving: BACKEND=${BACKEND} TP=${SERVING_TP} GPU=${SERVING_GPU}, workload isl=${ISL} osl=${OSL} conc=${CONC}`, + `- identity read: \`${resolved.canonical_id || '?'}\` (match tier \`${resolved.match_tier || '-'}\`)`, + '', + '## Configurations', + '', + '| stored direction | stored claim | measured here | delta vs baseline | parity | outcome |', + '|---|---|---|---|---|---|', + ...(verdicts.length ? verdicts.map(v => + `| ${v.direction || 'unlabeled'} | ${v.throughput_tok_s != null ? v.throughput_tok_s + ' tok/s' : '?'}` + + `${v.speedup != null ? ` (${v.speedup}x)` : ''} | ` + + `${v.measured_tok_s != null ? v.measured_tok_s + ' tok/s' : 'not benched'} | ` + + `${v.delta_pct != null ? (v.delta_pct >= 0 ? '+' : '') + v.delta_pct.toFixed(2) + '%' : '—'} | ` + + `${v.parity || '—'} | **${v.outcome}** |`) + : ['| _(none offered)_ | | | | | |']), + '', + '## Kernels', + '', + '| kernel | kind | stored isolated | measured e2e delta | outcome | why |', + '|---|---|---|---|---|---|', + ...(kernelVerdicts.length ? kernelVerdicts.map(v => + `| ${v.name} | ${v.kind || '?'} | ${v.isolated_speedup ? v.isolated_speedup + 'x' : '—'} | ` + + `${v.measured_delta_pct != null ? (v.measured_delta_pct >= 0 ? '+' : '') + v.measured_delta_pct + '%' : '—'} | ` + + `**${v.outcome}** | ${String(v.why || '').replace(/\|/g, '\\|').slice(0, 160)} |`) + : ['| _(none recorded)_ | | | | | |']), + '', + // The plan's fourth requirement, made concrete: a top-N entry that could not be + // reproduced is not discarded, it is handed forward WITH its reproduction material. A + // lead the next role cannot open is a lead it will not use, so the launch script and the + // patch paths are spelled out here rather than left to be rediscovered in the bundle. + ...(notReproduced.length ? [ + '## REFERENCE ONLY — recalled but NOT reproduced here', + '', + 'These were offered by the store and benched on this box, and either produced no number', + 'at all or never took effect (a flag renamed upstream is accepted silently and then', + 'ignored). They are NOT results. They are the closest thing this deployment has to a', + 'record of what someone else got working, and their material is below so you can read', + 'what they actually did rather than guess from a direction label.', + '', + ...notReproduced.flatMap(v => { + const repro = (v.repro && typeof v.repro === 'object') ? v.repro : {}; + const root = (v.bundle && v.bundle.path) ? `${v.bundle.path}/files` : ''; + const patches = (Array.isArray(repro.kernels) ? repro.kernels : []) + .filter(k => k && k.patch).map(k => root ? `${root}/${k.patch}` : k.patch); + return [ + `### ${v.direction || 'unlabeled'} (session \`${v.session_id || '?'}\`)`, + `- why it did not reproduce: ${String(v.why || 'no measurement came back').slice(0, 300)}`, + `- stored claim: ${v.throughput_tok_s != null ? v.throughput_tok_s + ' tok/s' : '?'}` + + `${v.speedup != null ? ` (${v.speedup}x)` : ''}, recorded elsewhere`, + `- launch script: ${repro.launch + ? `\`${root ? `${root}/${repro.launch}` : repro.launch}\`` + + (repro.launch_origin === 'synthesized' + ? ' — SYNTHESIZED from the stored config, never executed as-is' + : ' — captured from the original run') + : 'none recorded (this record predates `repro`)'}`, + `- kernel patches: ${patches.length ? patches.map(p => `\`${p}\``).join(', ') : 'none carried'}` + + `${repro.kernels_without_patch ? ` (${repro.kernels_without_patch} accepted kernel(s) carry no patch)` : ''}`, + `- flags: \`${String((v.accepted_config || {}).flags || '(none)')}\``, + `- env: \`${String((v.accepted_config || {}).env || '(none)')}\``, + '', + ]; + }), + ] : []), + '## How to use this', + '', + adoptedCfg.length + ? `The adopted configuration is ALREADY in CURRENT_FLAGS / CURRENT_ENV. Do not re-propose it — ` + + `propose only things that COMPOUND on top of it.` + : `Nothing from the store was adopted as configuration, so CURRENT_FLAGS / CURRENT_ENV are ` + + `unchanged from the baseline.`, + '', + adoptedKer.length + ? `${adoptedKer.length} recovered kernel(s) are already in the active overlay and already ` + + `reflected in the profile you are routing from. Their ops are DONE; look elsewhere.` + : `No recovered kernel was adopted, so the overlay is empty and every op in the profile is ` + + `still open.`, + '', + rejected.length + ? `The ${rejected.length} rejected/unreplayed entries above are LEADS, not dead ends. A ` + + `rejection here means the whole compounded thing did not beat this baseline on this box — ` + + `it does NOT mean each knob inside it is worthless, and it does not mean the DIRECTION is ` + + `wrong. Do not re-propose any of them verbatim; do feel free to propose an individual axis ` + + `from one, or the same idea approached differently.` + : `Nothing was rejected.`, + '', + ].join('\n'); + await safeAgent( + `You are a file writer. Use the Write tool to create the file ` + + `"${refsDir}/measured_on_this_box.md" with EXACTLY the content below, verbatim. ` + + `Do NOT reformat, summarize, re-order rows, or change any number:\n\n` + + '````markdown\n' + md + '\n````\n\n' + + `Then return {"written": true, "path": "${refsDir}/measured_on_this_box.md"}.`, + { phase: 'WarmStart', label: 'warm_start:record-measurements', + schema: obj({ written: { type: 'boolean' }, path: { type: 'string' } }, []) }, + 2); + KB_REF_DIR = refsDir; // the measurements exist even if the offer list was thin + + // The channel that always fires. Terse on purpose — it goes into `## Inputs` of two roles, + // and a verdict that has to be waded through is a verdict that gets skimmed. + KB_REF_INPUTS = { + KB_REFERENCE_DIR: refsDir, + KB_REFERENCE_VERDICT: + `The knowledge base offered ${allVerdicts.length} prior result(s) for this deployment ` + + `(${resolved.canonical_id || '?'}, tier ${resolved.match_tier || '-'}). This run benched them ` + + `and recorded what actually happened in ${refsDir}/measured_on_this_box.md — read that file, ` + + `and treat it as overriding the stored claims in its siblings. ` + + (adoptedCfg.length + ? `ADOPTED config: ${adoptedCfg.map(v => v.direction || 'unlabeled').join(', ')} — it is ALREADY ` + + `in CURRENT_FLAGS/CURRENT_ENV, so propose only things that COMPOUND on top of it, never it again. ` + : `No stored config was adopted. `) + + (adoptedKer.length + ? `ADOPTED kernels: ${adoptedKer.map(v => v.name).join(', ')} — already in the overlay and already ` + + `reflected in the profile, so those ops are done. ` + : `No stored kernel was adopted. `) + + (rejected.length + ? `REJECTED/unreplayed: ${rejected.map(v => v.direction || v.name || '?').join(', ')}. Do not ` + + `re-propose any of them verbatim — the compounded whole lost on this box. Their individual ` + + `knobs may each still be a valid axis, and their declared directions are still legitimate ` + + `ideas to reach a different way.` + : '') + + (notReproduced.length + ? ` ${notReproduced.length} of those could not be reproduced AT ALL here (no number, or the ` + + `config never took effect) — they are in the REFERENCE ONLY section of that file with their ` + + `launch script and kernel patches attached. Read them as evidence of what worked somewhere ` + + `else, not as something to re-run.` + : ''), + }; + } + } + } + phase('Profile'); profile = await safeAgent( roleAgent('profiler', 'baseline', 'Capture a warm trace and emit the standardized Top-N.', { EVAL_DIR, MODEL_PATH, GPU_ID: GPU_LIST[0], WORKLOAD, ROUND: 0, - OVERLAY_PYTHONPATH: '', EXTRA_SERVER_ARGS: curFlags, EXTRA_ENV: curEnv, SKILL_DIR: WORKFLOW_DIR, + // '' unless Module A adopted a stored kernel — the baseline Top-N must be captured on the + // configuration this run will actually optimize from, overlay included. + OVERLAY_PYTHONPATH: kbSeedOverlay, EXTRA_SERVER_ARGS: curFlags, EXTRA_ENV: curEnv, SKILL_DIR: WORKFLOW_DIR, ...TRACELENS_INPUTS, ...ANALYSIS_SKILL_INPUTS, }), { phase: 'Profile', label: 'profiler:baseline', schema: PROFILE_SCHEMA }); @@ -1117,7 +1849,7 @@ if (want('setup')) { roleAgent('system_architect', 'strategize', 'Route the Top-N into config/kernel/host tracks by Amdahl.', { EVAL_DIR, PROFILE_TOPN: profile ? profile.profile_topN_json : '', BASELINE_THROUGHPUT: BASELINE_TPUT, WORKLOAD, BUDGET, HEAD_THRESHOLD_PCT, CONFIG_TUNE_ENABLED, SKILL_DIR: WORKFLOW_DIR, - ...TRACELENS_INPUTS, ...ANALYSIS_SKILL_INPUTS, + ...TRACELENS_INPUTS, ...ANALYSIS_SKILL_INPUTS, ...KB_REF_INPUTS, }), { phase: 'Strategize', label: 'architect:strategize', schema: STRATEGY_SCHEMA }); kernelQueue = (strategy && strategy.kernel_candidates) ? strategy.kernel_candidates.slice() : []; @@ -1162,14 +1894,16 @@ if (want('setup')) { // =========================================================================== // PHASE: Config sweep (Config Tuner) — FIRST, reshapes the profile // =========================================================================== -let curTput = ST.throughput || BASELINE_TPUT; +// A carried phase-partial state still wins: it is a measurement this run already made. kbSeedTput is +// 0 when Module A did not run or adopted nothing, so `0 || BASELINE_TPUT === BASELINE_TPUT`. +let curTput = ST.throughput || kbSeedTput || BASELINE_TPUT; if (want('config') && CONFIG_TUNE_ENABLED && strategy && (strategy.config_directions || []).length) { phase('ConfigSweep'); const sweep = await safeAgent( roleAgent('config_tuner', 'sweep', 'Sweep the ranked config axes one at a time; keep wins.', { EVAL_DIR, MODEL_PATH, GPU_ID: GPU_LIST[0], WORKLOAD, BASELINE_THROUGHPUT: BASELINE_TPUT, NOISE_BAND_PCT: NOISE_BAND, E2E_REPEATS, CONFIG_DIRECTIONS: strategy.config_directions, - CURRENT_FLAGS: curFlags, CURRENT_ENV: curEnv, SKILL_DIR: WORKFLOW_DIR, + CURRENT_FLAGS: curFlags, CURRENT_ENV: curEnv, SKILL_DIR: WORKFLOW_DIR, ...KB_REF_INPUTS, }), { phase: 'ConfigSweep', label: 'config_tuner:sweep', schema: SWEEP_SCHEMA }); if (sweep && sweep.best_throughput_tok_s > curTput) { @@ -1206,11 +1940,13 @@ if (want('config') && CONFIG_TUNE_ENABLED && strategy && (strategy.config_direct // Shared state carried across the head + kernel tracks (and across phase invocations via args.state). // MUST be declared BEFORE the HeadKernel block that uses them (else temporal-dead-zone ReferenceError). // --------------------------------------------------------------------------- -let curOverlay = ST.overlay || ''; // the accepted overlay carried forward +let curOverlay = ST.overlay || kbSeedOverlay || ''; // the accepted overlay carried forward let dispatched = 0; // counts ONLY kernel-optimization tasks (the budget) let milestone = 0; let noImprove = 0; -const acceptedKernels = (ST.accepted_kernels || []).slice(); +// kbSeedKernels is empty unless Module A REPLAYED a stored kernel and its fresh A/B accepted it — +// these are this run's own measurements of a recovered patch, not the store's claims about it. +const acceptedKernels = (ST.accepted_kernels || []).slice().concat(kbSeedKernels); const acceptedHeads = (ST.accepted_heads || []).slice(); // Verified-isolated wins whose e2e A/B did NOT complete (integrate agent timed // out / hung / degraded to null mid-gate). These are NOT rejections — keep them @@ -1434,29 +2170,30 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { log(`[deep] E2E GATE #${e2eGateCount} on serving {${SERVING_GPU}} TP=${SERVING_TP}: [${cands.map(c => c.uid + ' ' + c.best.toFixed(3) + 'x').join(', ')}] (overlapping co-opt on dedicated cards).`); for (const c of cands) { if (opts.final && bankedHeads.has(c.head.short_name)) { log(` [deep] FINALIZE: skip ${c.uid} -- head ${c.head.short_name} already banked (same module, cannot stack).`); continue; } + const deepInputs = { + EVAL_DIR, MODEL_PATH, GPU_ID: SERVING_GPU, WORKLOAD, NOISE_BAND_PCT: NOISE_BAND, E2E_REPEATS, + KERNEL_RESULT: { + short_name: c.head.short_name, task_dir: c.ext.task_dir, op_kind: c.ext.op_kind, lane: c.key, + winner_kind: 'patch', winner_backend: c.lang, + target_callable: c.ext.target_callable || c.head.target_callable || '', + authored_language: c.lang, authored_kernel_eval_dir: c.lastEval, + apply_env: '', apply_flags: '', code_patch: c.patch || (c.lastEval ? `${c.lastEval}/final_patch.diff` : ''), tuning_artifact: '', + verified_isolated_speedup: c.best, pct_gpu_time: c.head.pct_gpu_time, parity_note: 'expected_close', + }, + CURRENT_OVERLAY: curOverlay, CURRENT_FLAGS: curFlags, CURRENT_ENV: curEnv, + CURRENT_THROUGHPUT: curTput, SKILL_DIR: WORKFLOW_DIR, DEEP_FEEDBACK: true, + ...ACCURACY_INPUTS, + ...(opts.final && ACCURACY_GATE !== 'none' ? { ACCURACY_LIMIT: DEEP_FINAL_ACCURACY_LIMIT } : {}), // de-noise the finalize accuracy decision + }; const integ = await safeAgent( - roleAgent('e2e_integrator', 'integrate', 'Apply a deep head candidate; gate on e2e throughput; report engagement/cudagraph/mem/decode for feedback.', { - EVAL_DIR, MODEL_PATH, GPU_ID: SERVING_GPU, WORKLOAD, NOISE_BAND_PCT: NOISE_BAND, E2E_REPEATS, - KERNEL_RESULT: { - short_name: c.head.short_name, task_dir: c.ext.task_dir, op_kind: c.ext.op_kind, lane: c.key, - winner_kind: 'patch', winner_backend: c.lang, - target_callable: c.ext.target_callable || c.head.target_callable || '', - authored_language: c.lang, authored_kernel_eval_dir: c.lastEval, - apply_env: '', apply_flags: '', code_patch: c.patch || (c.lastEval ? `${c.lastEval}/final_patch.diff` : ''), tuning_artifact: '', - verified_isolated_speedup: c.best, pct_gpu_time: c.head.pct_gpu_time, parity_note: 'expected_close', - }, - CURRENT_OVERLAY: curOverlay, CURRENT_FLAGS: curFlags, CURRENT_ENV: curEnv, - CURRENT_THROUGHPUT: curTput, SKILL_DIR: WORKFLOW_DIR, DEEP_FEEDBACK: true, - ...ACCURACY_INPUTS, - ...(opts.final && ACCURACY_GATE !== 'none' ? { ACCURACY_LIMIT: DEEP_FINAL_ACCURACY_LIMIT } : {}), // de-noise the finalize accuracy decision - }), + roleAgent('e2e_integrator', 'integrate', 'Apply a deep head candidate; gate on e2e throughput; report engagement/cudagraph/mem/decode for feedback.', deepInputs), { phase: 'HeadKernel', label: `integrate ${c.uid} g${e2eGateCount}`, schema: INTEGRATE_SCHEMA }); if (integ && integ.output_parity === 'fail') { log(` [deep] ${c.uid}: REJECTED — output_parity=fail vs true baseline.`); history.ledger.push({ direction: c.uid, isolated_speedup: c.best, e2e_delta_pct: integ.e2e_delta_pct, verdict: 'dead_end', lesson: 'parity fail vs true baseline' }); } else if (integAccepted(integ, c.head.pct_gpu_time, c.best) && integ.e2e_throughput_tok_s > curTput) { curOverlay = integ.accepted_overlay || curOverlay; curTput = integ.e2e_throughput_tok_s; bankedHeads.add(c.head.short_name); - acceptedHeads.push({ short_name: c.head.short_name, op_kind: c.ext.op_kind, backend: c.lang, lane: c.key, kind: 'patch', e2e_delta_pct: integ.e2e_delta_pct, isolated: c.best }); + bankAccepted(acceptedHeads, { short_name: c.head.short_name, op_kind: c.ext.op_kind, backend: c.lang, lane: c.key, kind: 'patch', e2e_delta_pct: integ.e2e_delta_pct, isolated: c.best }, krOf(deepInputs)); log(` [deep] ${c.uid}: ACCEPTED. e2e now ${curTput} tok/s (+${integ.e2e_delta_pct}%); target ${Math.round(BASELINE_TPUT * DEEP_E2E_TARGET)} tok/s.`); history.ledger.push({ direction: c.uid, isolated_speedup: c.best, e2e_delta_pct: integ.e2e_delta_pct, verdict: 'confirmed', lesson: integ.reason || '' }); } else { @@ -1482,7 +2219,7 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { }); if (dcorr.banked) { curOverlay = dcorr.integ.accepted_overlay || curOverlay; curTput = dcorr.integ.e2e_throughput_tok_s; bankedHeads.add(c.head.short_name); - acceptedHeads.push({ short_name: c.head.short_name, op_kind: c.ext.op_kind, backend: c.lang, lane: c.key, kind: 'patch', e2e_delta_pct: dcorr.integ.e2e_delta_pct, isolated: dcorr.isolated, corrective: true }); + bankAccepted(acceptedHeads, { short_name: c.head.short_name, op_kind: c.ext.op_kind, backend: c.lang, lane: c.key, kind: 'patch', e2e_delta_pct: dcorr.integ.e2e_delta_pct, isolated: dcorr.isolated, corrective: true, patch: dcorr.patch || '' }, krOf(deepInputs)); log(` [deep] ${c.uid}: ACCEPTED after corrective re-author. e2e now ${curTput} tok/s (+${dcorr.integ.e2e_delta_pct}%).`); history.ledger.push({ direction: c.uid, isolated_speedup: dcorr.isolated, e2e_delta_pct: dcorr.integ.e2e_delta_pct, verdict: 'confirmed_corrective', lesson: `fixed: ${dreason}` }); } else { @@ -1578,6 +2315,7 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { await deepBoundedWorkflow({ scriptPath: KERNEL_WF_SCRIPT }, { kernel_path: l.ext.task_dir, workflow_dir: KERNEL_WF_DIR, mode: l.mode, target_language: l.lang, op_spec: l.opSpec, perf_knowledge_dir: KERNEL_KNOWLEDGE_DIR, use_expert_skills: USE_EXPERT_SKILLS ? 'true' : 'false', expert_skills_dir: EXPERT_SKILLS_DIR, + ...KB_ARGS, budget: DEEP_WAVE_BUDGET, max_no_improve: DEEP_WAVE_BUDGET, gpu_ids: g[0], state_dir: l.state_dir, shared_kb: l.sharedKb, global_kb: GLOBAL_KB, incremental_analyze: l.ran > 1 ? 'true' : 'false', // P2: 2nd+ burst of a lane = continuation -> skip cold re-analysis @@ -1708,6 +2446,7 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { op_spec: { op_kind: j.ext.op_kind, shapes: j.ext.shapes || {}, dtype: j.ext.dtype || 'bf16', regime: j.h.regime || '', cuda_graph_safe: true, ...(j.ext.workload_path ? { workload_path: j.ext.workload_path } : {}) }, perf_knowledge_dir: KERNEL_KNOWLEDGE_DIR, use_expert_skills: USE_EXPERT_SKILLS ? 'true' : 'false', expert_skills_dir: EXPERT_SKILLS_DIR, + ...KB_ARGS, budget: KERNEL_BUDGET, gpu_ids: g[0], exp_root: `${EVAL_DIR}/kernels/_exp`, task: `Author+optimize a ${lang} implementation of this op vs the immutable oracle (beat ${j.best_known_ms || '?'} ms). ` + `This kernel will be overlaid onto the LIVE decode path (CUDA-graph captured): its STEADY-STATE hot path MUST be ` + @@ -1756,32 +2495,33 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { st.cands.sort((a, b) => (b.isolated || 0) - (a.isolated || 0)); const cand = st.cands[0]; log(` ${h.short_name}: best candidate=${cand.source} (${(cand.isolated || 0).toFixed(2)}x, ${cand.kind}). Integrating to e2e (serial, slot {${SERVING_GPU}}).`); + const headWinnerInputs = { + EVAL_DIR, MODEL_PATH, GPU_ID: SERVING_GPU, WORKLOAD, NOISE_BAND_PCT: NOISE_BAND, E2E_REPEATS, + KERNEL_RESULT: { short_name: h.short_name, task_dir: st.ext.task_dir, op_kind: st.ext.op_kind, + winner_kind: cand.winner_kind, winner_backend: cand.source, + target_callable: st.ext.target_callable || h.target_callable || '', + authored_language: cand.language || '', authored_kernel_eval_dir: cand.kernel_eval_dir || '', + apply_env: cand.apply_env || '', apply_flags: cand.apply_flags || '', + code_patch: cand.code_patch || cand.final_patch || '', tuning_artifact: cand.tuning_artifact || '', + verified_isolated_speedup: cand.isolated || 0, pct_gpu_time: h.pct_gpu_time, + // Pass the Architect's live seam + a concrete engagement assertion so the Integrator can + // VERIFY the overlay actually binds on the live path BEFORE spending a full e2e A/B — an + // unreachable lever is then rejected in minutes (no_engagement), not hours. + live_call_seam: h.live_call_seam || '', engagement_check: h.engagement_check || '', + parity_note: cand.parity_note || 'expected_close' }, + CURRENT_OVERLAY: curOverlay, CURRENT_FLAGS: curFlags, CURRENT_ENV: curEnv, + CURRENT_THROUGHPUT: curTput, SKILL_DIR: WORKFLOW_DIR, + ENGAGEMENT_CHECK: h.engagement_check || '', + }; const integ = await runIntegrateBothLegs( - 'Apply the head-op winner; gate on e2e throughput.', { - EVAL_DIR, MODEL_PATH, GPU_ID: SERVING_GPU, WORKLOAD, NOISE_BAND_PCT: NOISE_BAND, E2E_REPEATS, - KERNEL_RESULT: { short_name: h.short_name, task_dir: st.ext.task_dir, op_kind: st.ext.op_kind, - winner_kind: cand.winner_kind, winner_backend: cand.source, - target_callable: st.ext.target_callable || h.target_callable || '', - authored_language: cand.language || '', authored_kernel_eval_dir: cand.kernel_eval_dir || '', - apply_env: cand.apply_env || '', apply_flags: cand.apply_flags || '', - code_patch: cand.code_patch || cand.final_patch || '', tuning_artifact: cand.tuning_artifact || '', - verified_isolated_speedup: cand.isolated || 0, pct_gpu_time: h.pct_gpu_time, - // Pass the Architect's live seam + a concrete engagement assertion so the Integrator can - // VERIFY the overlay actually binds on the live path BEFORE spending a full e2e A/B — an - // unreachable lever is then rejected in minutes (no_engagement), not hours. - live_call_seam: h.live_call_seam || '', engagement_check: h.engagement_check || '', - parity_note: cand.parity_note || 'expected_close' }, - CURRENT_OVERLAY: curOverlay, CURRENT_FLAGS: curFlags, CURRENT_ENV: curEnv, - CURRENT_THROUGHPUT: curTput, SKILL_DIR: WORKFLOW_DIR, - ENGAGEMENT_CHECK: h.engagement_check || '', - }, + 'Apply the head-op winner; gate on e2e throughput.', headWinnerInputs, `integrate ${h.short_name}`, 'HeadKernel'); if (integAccepted(integ, h.pct_gpu_time, cand.isolated) && integ.e2e_throughput_tok_s > curTput) { curOverlay = integ.accepted_overlay || curOverlay; if (cand.winner_kind === 'env' && cand.apply_env) curEnv = (curEnv ? curEnv + ' ' : '') + cand.apply_env; if (cand.winner_kind === 'flag' && cand.apply_flags) curFlags = (curFlags ? curFlags + ' ' : '') + cand.apply_flags; curTput = integ.e2e_throughput_tok_s; - acceptedHeads.push({ short_name: h.short_name, op_kind: st.ext.op_kind, backend: cand.source, kind: cand.winner_kind, e2e_delta_pct: integ.e2e_delta_pct, isolated: cand.isolated }); + bankAccepted(acceptedHeads, { short_name: h.short_name, op_kind: st.ext.op_kind, backend: cand.source, kind: cand.winner_kind, e2e_delta_pct: integ.e2e_delta_pct, isolated: cand.isolated }, krOf(headWinnerInputs)); log(` ${h.short_name}: ACCEPTED. e2e now ${curTput} tok/s (+${integ.e2e_delta_pct}%).`); history.ledger.push({ direction: h.short_name, isolated_speedup: cand.isolated, e2e_delta_pct: integ.e2e_delta_pct, verdict: 'confirmed', lesson: integ.reason || '' }); } else { @@ -1809,7 +2549,7 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { : { banked: false }; if (corr.banked) { curOverlay = corr.integ.accepted_overlay || curOverlay; curTput = corr.integ.e2e_throughput_tok_s; - acceptedHeads.push({ short_name: h.short_name, op_kind: st.ext.op_kind, backend: cand.source, kind: 'authored', e2e_delta_pct: corr.integ.e2e_delta_pct, isolated: corr.isolated, corrective: true }); + bankAccepted(acceptedHeads, { short_name: h.short_name, op_kind: st.ext.op_kind, backend: cand.source, kind: 'authored', e2e_delta_pct: corr.integ.e2e_delta_pct, isolated: corr.isolated, corrective: true, patch: corr.patch || '' }, krOf(headWinnerInputs)); log(` ${h.short_name}: ACCEPTED after corrective re-author (${reason}). e2e now ${curTput} tok/s (+${corr.integ.e2e_delta_pct}%).`); history.ledger.push({ direction: h.short_name, isolated_speedup: corr.isolated, e2e_delta_pct: corr.integ.e2e_delta_pct, verdict: 'confirmed_corrective', lesson: `fixed: ${reason}` }); } else { @@ -1914,6 +2654,7 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { op_spec: { op_kind: ext.op_kind, shapes: ext.shapes || {}, dtype: ext.dtype || 'bf16', regime: h.regime || '', cuda_graph_safe: true, ...(ext.workload_path ? { workload_path: ext.workload_path } : {}) }, perf_knowledge_dir: KERNEL_KNOWLEDGE_DIR, use_expert_skills: USE_EXPERT_SKILLS ? 'true' : 'false', expert_skills_dir: EXPERT_SKILLS_DIR, + ...KB_ARGS, budget: KERNEL_BUDGET, gpu_ids: h.gpu_id, exp_root: `${EVAL_DIR}/kernels/_exp`, task: `Author+optimize a ${lang} implementation of this op vs the immutable oracle (beat ${bake.best_known_ms || '?'} ms). ` + `This kernel will be overlaid onto the LIVE sglang decode path, which is CUDA-graph captured: its STEADY-STATE hot ` + @@ -2008,7 +2749,7 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { verdict: passed ? 'candidate_passed' : (abc ? 'candidate_rejected' : 'candidate_incomplete'), lesson: `${cand.winner_kind} ${cand.source}: ${integ && integ.e2e_throughput_tok_s ? `${integ.e2e_throughput_tok_s.toFixed(0)} tok/s (${(integ.e2e_delta_pct || 0).toFixed(2)}%)` : (integ ? integ.reason || integ.gate : 'null/timeout')}` }); if (passed) { - if (!bestPick || integ.e2e_throughput_tok_s > bestPick.integ.e2e_throughput_tok_s) bestPick = { cand, integ }; + if (!bestPick || integ.e2e_throughput_tok_s > bestPick.integ.e2e_throughput_tok_s) bestPick = { cand, integ, inputs }; log(` ${h.short_name}: candidate ${cand.source} PASSED e2e gate (${integ.e2e_throughput_tok_s} tok/s, +${integ.e2e_delta_pct}%).`); } else { log(` ${h.short_name}: candidate ${cand.source} ${abc ? `rejected (${gateRejectReason(integ, h.pct_gpu_time, cand.isolated)})` : `A/B incomplete (${integ ? integ.reason || integ.gate : 'null/timeout'})`}.`); @@ -2023,7 +2764,7 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { if (cand.winner_kind === 'env' && cand.apply_env) curEnv = (curEnv ? curEnv + ' ' : '') + cand.apply_env; if (cand.winner_kind === 'flag' && cand.apply_flags) curFlags = (curFlags ? curFlags + ' ' : '') + cand.apply_flags; curTput = integ.e2e_throughput_tok_s; - acceptedHeads.push({ short_name: h.short_name, op_kind: ext.op_kind, backend: cand.source, kind: cand.winner_kind, e2e_delta_pct: integ.e2e_delta_pct, isolated: cand.isolated }); + bankAccepted(acceptedHeads, { short_name: h.short_name, op_kind: ext.op_kind, backend: cand.source, kind: cand.winner_kind, e2e_delta_pct: integ.e2e_delta_pct, isolated: cand.isolated }, krOf(bestPick.inputs)); log(` ${h.short_name}: ACCEPTED best candidate=${cand.source} (${(cand.isolated || 0).toFixed(2)}x iso). e2e now ${curTput} tok/s (+${integ.e2e_delta_pct}%).`); history.ledger.push({ direction: h.short_name, isolated_speedup: cand.isolated, e2e_delta_pct: integ.e2e_delta_pct, verdict: 'confirmed', lesson: integ.reason || '' }); } else { @@ -2046,7 +2787,7 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { : { banked: false }; if (corr.banked) { curOverlay = corr.integ.accepted_overlay || curOverlay; curTput = corr.integ.e2e_throughput_tok_s; - acceptedHeads.push({ short_name: h.short_name, op_kind: ext.op_kind, backend: cand.source, kind: 'authored', e2e_delta_pct: corr.integ.e2e_delta_pct, isolated: corr.isolated, corrective: true }); + bankAccepted(acceptedHeads, { short_name: h.short_name, op_kind: ext.op_kind, backend: cand.source, kind: 'authored', e2e_delta_pct: corr.integ.e2e_delta_pct, isolated: corr.isolated, corrective: true, patch: corr.patch || '' }, krOf(headIntegrateInputs)); log(` ${h.short_name}: ACCEPTED after corrective re-author (was crash/incomplete: ${reason}). e2e now ${curTput} tok/s (+${corr.integ.e2e_delta_pct}%).`); history.ledger.push({ direction: h.short_name, isolated_speedup: corr.isolated, e2e_delta_pct: corr.integ.e2e_delta_pct, verdict: 'confirmed_corrective', lesson: `fixed crash: ${reason}` }); } else { @@ -2073,7 +2814,7 @@ if (want('head') && headQueue.length && HEAD_BUDGET > 0) { : { banked: false }; if (corr.banked) { curOverlay = corr.integ.accepted_overlay || curOverlay; curTput = corr.integ.e2e_throughput_tok_s; - acceptedHeads.push({ short_name: h.short_name, op_kind: ext.op_kind, backend: cand.source, kind: 'authored', e2e_delta_pct: corr.integ.e2e_delta_pct, isolated: corr.isolated, corrective: true }); + bankAccepted(acceptedHeads, { short_name: h.short_name, op_kind: ext.op_kind, backend: cand.source, kind: 'authored', e2e_delta_pct: corr.integ.e2e_delta_pct, isolated: corr.isolated, corrective: true, patch: corr.patch || '' }, krOf(headIntegrateInputs)); log(` ${h.short_name}: ACCEPTED after corrective re-author. e2e now ${curTput} tok/s (+${corr.integ.e2e_delta_pct}%).`); history.ledger.push({ direction: h.short_name, isolated_speedup: corr.isolated, e2e_delta_pct: corr.integ.e2e_delta_pct, verdict: 'confirmed_corrective', lesson: `fixed: ${reason}` }); } else { @@ -2172,6 +2913,7 @@ while (want('kernel') && !TIME_DEADLINE_HIT && dispatched < BUDGET && (dispatche const r = await workflow({ scriptPath: KERNEL_WF_SCRIPT }, laneArgs({ kernel_path: ext.task_dir, workflow_dir: KERNEL_WF_DIR, use_expert_skills: USE_EXPERT_SKILLS ? 'true' : 'false', expert_skills_dir: EXPERT_SKILLS_DIR, + ...KB_ARGS, budget: KERNEL_BUDGET, gpu_ids: c.gpu_id, exp_root: `${EVAL_DIR}/kernels/_exp`, task: 'Compare candidate backends ' + JSON.stringify(c.candidate_backends || []) + ' for this kernel; pick the fastest that passes the immutable unittest. ' + GRAPH_REQ + (TASK || ''), @@ -2223,7 +2965,7 @@ while (want('kernel') && !TIME_DEADLINE_HIT && dispatched < BUDGET && (dispatche if (abDone && integAccepted(integ, c.pct_gpu_time, kl.final_geomean) && integ.e2e_throughput_tok_s > curTput) { curOverlay = integ.accepted_overlay || curOverlay; curTput = integ.e2e_throughput_tok_s; - acceptedKernels.push({ short_name: c.short_name, backend: kl.note || '', e2e_delta_pct: integ.e2e_delta_pct, isolated: kl.final_geomean }); + bankAccepted(acceptedKernels, { short_name: c.short_name, backend: kl.note || '', kind: 'patch', e2e_delta_pct: integ.e2e_delta_pct, isolated: kl.final_geomean }, krOf(mileIntegrateInputs)); milestoneImproved = true; log(` ${c.short_name}: ACCEPTED. e2e now ${curTput} tok/s (+${integ.e2e_delta_pct}%).`); history.ledger.push({ direction: c.short_name, isolated_speedup: kl.final_geomean, e2e_delta_pct: integ.e2e_delta_pct, verdict: 'confirmed', lesson: integ.reason || '' }); @@ -2243,7 +2985,7 @@ while (want('kernel') && !TIME_DEADLINE_HIT && dispatched < BUDGET && (dispatche : { banked: false }; if (corr.banked) { curOverlay = corr.integ.accepted_overlay || curOverlay; curTput = corr.integ.e2e_throughput_tok_s; - acceptedKernels.push({ short_name: c.short_name, backend: kl.note || '', e2e_delta_pct: corr.integ.e2e_delta_pct, isolated: corr.isolated, corrective: true }); + bankAccepted(acceptedKernels, { short_name: c.short_name, backend: kl.note || '', kind: 'patch', e2e_delta_pct: corr.integ.e2e_delta_pct, isolated: corr.isolated, corrective: true, patch: corr.patch || '' }, krOf(mileIntegrateInputs)); milestoneImproved = true; log(` ${c.short_name}: ACCEPTED after corrective re-author (${reason}). e2e now ${curTput} tok/s (+${corr.integ.e2e_delta_pct}%).`); history.ledger.push({ direction: c.short_name, isolated_speedup: corr.isolated, e2e_delta_pct: corr.integ.e2e_delta_pct, verdict: 'confirmed_corrective', lesson: `fixed: ${reason}` }); @@ -2385,9 +3127,9 @@ if (want('final')) { if (p.track === 'head') { if (p.winner_kind === 'env' && p.apply_env) curEnv = (curEnv ? curEnv + ' ' : '') + p.apply_env; if (p.winner_kind === 'flag' && p.apply_flags) curFlags = (curFlags ? curFlags + ' ' : '') + p.apply_flags; - acceptedHeads.push({ short_name: p.short_name, op_kind: p.op_kind, backend: p.backend, kind: p.winner_kind, e2e_delta_pct: integ.e2e_delta_pct, isolated: p.isolated }); + bankAccepted(acceptedHeads, { short_name: p.short_name, op_kind: p.op_kind, backend: p.backend, kind: p.winner_kind, e2e_delta_pct: integ.e2e_delta_pct, isolated: p.isolated, pct_gpu_time: p.pct_gpu_time }, krOf(p.inputs)); } else { - acceptedKernels.push({ short_name: p.short_name, backend: p.backend || '', e2e_delta_pct: integ.e2e_delta_pct, isolated: p.isolated }); + bankAccepted(acceptedKernels, { short_name: p.short_name, backend: p.backend || '', kind: 'patch', e2e_delta_pct: integ.e2e_delta_pct, isolated: p.isolated, pct_gpu_time: p.pct_gpu_time }, krOf(p.inputs)); } curTput = integ.e2e_throughput_tok_s; finalTput = curTput; finalSpeedup = BASELINE_TPUT ? curTput / BASELINE_TPUT : 1.0; @@ -2505,6 +3247,20 @@ const carryState = { history, }; +// What this run got from the KB, and — separately — what it added on its own. The split is the +// point: a warm-started run's headline speedup is partly inherited, and reporting the whole of it as +// this run's work is how a knowledge base starts flattering itself. BASELINE_TPUT stays the true +// cold baseline throughout, so `throughput_speedup` above is still measured against the real floor. +const kbWarmStart = (E2E_WARM_START_ON && KB_DIMS) ? { + identity: KB_DIMS, plane: E2E_KB_PLANE, read_plane: KB_READ_PLANE, mode: E2E_WARM_START, + reference_dir: KB_REF_DIR, + adopted_throughput_tok_s: kbSeedTput || null, + adopted_kernels: kbSeedKernels.map(k => k.name).filter(Boolean), + // The gain that is genuinely THIS run's: everything after the warm start's own contribution. Null + // when nothing was adopted, because then the ordinary speedup already answers the question. + incremental_speedup: kbSeedTput ? Number((finalTput / kbSeedTput).toFixed(6)) : null, +} : null; + const wfReturn = { // schema_version pins the CONTRACT shape run_e2e.py reads. Bump only on a // breaking change to the keys below; run_e2e.py keys its canonical-artifact @@ -2546,8 +3302,17 @@ const wfReturn = { budget_used: dispatched, budget_total: BUDGET, final_overlay: (validation && validation.final_overlay) || (finalize && finalize.final_overlay) || curOverlay, + // FINALIZE_SCHEMA has always defined final_patch and the Finalizer has always returned it, but it + // was never surfaced here — so the single diff carrying the whole run's win reached neither the + // report nor the KB (e2e_store.py's _ARTIFACT_KEYS looks up exactly result["final_patch"]). + // final_overlay cannot stand in for it: that is a DIRECTORY, and the store's os.path.isfile filter + // drops it, so a run with no final_patch uploads no reproducible code at all. + final_patch: (finalize && finalize.final_patch) || '', final_launch_script: (validation && validation.final_launch_script) || (finalize && finalize.final_launch_script) || '', report_path: report ? report.report_path : `${EVAL_DIR}/architect_report.md`, + // Conditional spread, never an unconditional key: with the feature off this contributes no own + // property and both the returned object and the persisted file are byte-identical to before. + ...(kbWarmStart ? { kb_warm_start: kbWarmStart } : {}), state: carryState, }; @@ -2588,4 +3353,68 @@ if (EVAL_DIR) { } } +// =========================================================================== +// MODULE B — record this run in the deployment knowledge base. +// +// STRICTLY AFTER the workflow_return.json persist above, and that ordering is not stylistic. +// workflow_return.json is the pinned contract run_e2e.py depends on, documented as the workflow's +// final act precisely because this process gets SIGKILLed. Module B is the one operation in the +// whole workflow that waits on a NETWORK rather than a GPU, and putting it in front of the persist +// would trade the run's canonical artifact for a knowledge-base entry. Feeding the KB that same file +// as --result has a second benefit: the record cannot drift from the report, because they are the +// same bytes. +// +// EVERY --apply IS PERMANENT. The service exposes no DELETE for /v1/kb/*, so a wrong record cannot +// be cleaned up, only outranked. Hence the gates below are conservative by design. +// =========================================================================== +if (E2E_WARM_START_ON && KB_DIMS && KB_DIMS.gfx && want('final') && EVAL_DIR && + wfReturn.throughput_speedup > 1.0 && wfReturn.final_throughput_tok_s > 0) { + // Computed HERE, deterministically, from facts this script already holds — never asked of an + // agent. `direction` is inside _content_digest, so a label that varies between two runs of the + // same configuration mints a second session instead of replacing the first, and the page fills + // with near-identical entries that outrank each other by noise. '' (unlabeled) is the honest + // answer when nothing distinct happened; the store collapses unlabeled records with nothing. + const kbDirection = [ + (curFlags !== INIT_FLAGS || curEnv !== INIT_ENV) ? 'config' : '', + (acceptedKernels.length || acceptedHeads.length) ? 'kernels' : '', + FAST_MODE ? 'fast' : (DEEP_MODE ? 'deep' : ''), + ].filter(Boolean).join('+'); + const writeCmd = + KB_ENV_PRELUDE + + `python3 ${shq(E2E_STORE_SCRIPT)} write ${kbIdentityFlags()} ` + + `${kbPlaneFlags()} --result ${shq(EVAL_DIR + '/workflow_return.json')} ` + + `--direction ${shq(kbDirection)} --measured-by ${shq('e2e_workflow:' + BACKEND)} --apply`; + try { + const written = await safeAgent( + `You are the e2e knowledge-base writer. Run EXACTLY this command and return its JSON stdout ` + + `verbatim as StructuredOutput. It pretty-prints over several lines; return the whole object. ` + + `Do NOT edit the command, do NOT re-run it with different flags, and do NOT retry it on ` + + `failure — every write it performs is PERMANENT (the store has no delete), so a second ` + + `attempt with adjusted arguments creates a second permanent record rather than fixing the ` + + `first. If it fails, return what it printed along with the error.\n` + + '```bash\n' + writeCmd + ` | tee ${shq(EVAL_DIR + '/kb_write.json')}\n` + '```', + { phase: 'Validate', label: 'kb:write', + schema: obj({ ok: { type: 'boolean' }, applied: { type: 'boolean' }, + session_id: { type: 'string' }, files: arrStr, rungs: arrObj }, []) }, + 1); + wfReturn.kb_written = written || { ok: false, error: 'writer agent returned nothing' }; + const rungs = (written && written.rungs) || []; + log(`[kb] wrote this run: session=${(written && written.session_id) || '?'} ` + + `direction='${kbDirection}' rungs=${rungs.filter(r => r.written).length}/${rungs.length} ` + + `promoted=${rungs.filter(r => r.promoted).length}` + + `${rungs.filter(r => r.error).map(r => ` [${r.tier}: ${r.error}]`).join('')}`); + } catch (e) { + // Non-fatal, exactly like the persist above. The run's result is already on disk and already + // returned; failing to record it is a lost opportunity, not a lost measurement. + wfReturn.kb_written = { ok: false, error: String(e).slice(0, 200) }; + log(`[kb] write failed (NON-FATAL — the run's own artifacts are unaffected): ${String(e)}`); + } +} else if (E2E_WARM_START_ON && KB_DIMS) { + const why = !KB_DIMS.gfx ? 'no gfx (an arch-less record is permanent and unattributable)' + : !want('final') ? 'this is a phase-partial run, so the number is not final' + : !(wfReturn.throughput_speedup > 1.0) ? `no win to record (${wfReturn.throughput_speedup}x)` + : 'no final throughput measured'; + log(`[kb] not recording this run: ${why}.`); +} + return wfReturn; diff --git a/e2e_workflow/knowledge/backend_playbook.md b/e2e_workflow/knowledge/backend_playbook.md index c40968b45..96b7aca96 100644 --- a/e2e_workflow/knowledge/backend_playbook.md +++ b/e2e_workflow/knowledge/backend_playbook.md @@ -33,6 +33,16 @@ experience; treat the seed as priors, not gospel — the unittest is the judge. | activation (silu/gelu + mul) | fused act_and_mul (aiter/triton) | collapse into the producing GEMM epilogue if possible | | elementwise/fill/cast/copy | fuse away (host_runtime) / cuda-graph | usually shouldn't be its own kernel | +## Roofline prior calibration (predicted vs measured — one line per direction) +- 2026-08-19 · gfx950 vLLM mxfp4 grouped fused-MoE (`_matmul_ogs...swiglu`, gpt-oss-120b, decode): roofline + predicted `attainable_speedup=1.0`, `expected_e2e_gain_pct=0.0` (memory-bound, `roofline_pct` 0.95–1.0, + headroom `saturated`, confidence **low**). MEASURED e2e **+26.9%** via a whole-file Triton rewrite. The + device-time byte/FLOP roofline model was WRONG for this seam — it cannot see the win, which is + host launch-overhead / decode-seam collapse, not a byte reduction. Correct behavior held (confidence was + low → not ranked on, head not dropped). Lesson: for a fused MoE dispatcher at decode, do NOT trust a + `saturated`/`attainable=1.0` roofline verdict to size the opportunity; the launch-overhead win is invisible + to it. This is the exact "measured EXCEEDS predicted attainable" failure mode — flag loudly. + ## How to use this in a run 1. Architect reads the Profiler Top-N classification + shapes. 2. For `library_*` kernels → hand to Config Tuner with the ranked swaps above (no source edit). diff --git a/e2e_workflow/knowledge/learned/INDEX.md b/e2e_workflow/knowledge/learned/INDEX.md index 64fe751e1..7d9d59f64 100644 --- a/e2e_workflow/knowledge/learned/INDEX.md +++ b/e2e_workflow/knowledge/learned/INDEX.md @@ -15,7 +15,7 @@ Confidence (a hint strength, not authority): ★ noise/unverified · ★★ sing - [gfx950 · vLLM MXFP8 E8M0 decode-bound] dense-linear split-K/fused decode-tile Triton rewrite ★★★ **+21.8% e2e (verified, gsm8k-clean); decode-driven (converts only at high conc); grouped-MoE GEMM resists (~1.1× ceiling)** — (mxfp8-linear-decode-rewrite-gfx950.md) - [gfx942 · sglang bf16] aiter per-shape DB tune ★★★ **+2.23% e2e (verified)** — (aiter-bf16-tuned-gemm-gfx942.md) - [gfx950 · sglang bf16 dense] backend swap = NO win (hipBLASLt already fastest), but the live seam IS `aiter.tuned_gemm:gemm_a16w16` (unlike vLLM) → aiter per-shape DB tune, `AITER_CONFIG_GEMM_BF16` colon-MERGE, ZERO HBM ★★ **engages; coverage-gated (gate_up/down 0 shipped = headroom, qkv/o already shipped = noise); e2e pending, gfx942 analog +2.23%** — (dense-gemm-bf16-gfx950-sglang.md) -- [gfx950 · vLLM fp8 a8w8 blockscale] per-shape CK tune DB `AITER_CONFIG_GEMM_A8W8_BLOCKSCALE`, plus the `VLLM_ROCM_USE_AITER[_LINEAR]=1` Triton→CK swap whenever AITER is off ★★★ **+16.1% / +65.69% / +18.86% e2e (Director-validated); iso 1.86–3.17× serving-wtd — but gain is SHIPPED-coverage-gated (~1.03× when already covered): probe `use default` vs `is tuned` first** — (fp8-a8w8-blockscale-ck-tune-gfx950-vllm.md) +- [gfx950 · vLLM fp8 a8w8 blockscale] per-shape CK tune DB `AITER_CONFIG_GEMM_A8W8_BLOCKSCALE`, plus the `VLLM_ROCM_USE_AITER[_LINEAR]=1` Triton→CK swap whenever AITER is off ★★★ **+16.1% / +65.69% / +18.86% e2e (Director-validated); iso 1.86–3.17× serving-wtd — but gain is SHIPPED-coverage-gated (~1.03× when already covered): probe `use default` vs `is tuned` first; swap-only UNTUNED CK confirmed 1.513× serving-wtd on Qwen3-14B-FP8 but REGRESSES prefill 0.66× → needs the CK tune (offline tuner/ckProfiler may be ABSENT from the aiter wheel — provision to recover prefill)** — (fp8-a8w8-blockscale-ck-tune-gfx950-vllm.md) - [gfx950 · sglang fp8 a8w8 blockscale, ROCm>=7.2] live kernel is CK **bpreshuffle**, not plain-CK/Triton → tune THAT DB (`--preshuffle`, env `AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE`), no overlay/swap needed ★★★ **+20.91% e2e (Qwen3-14B-FP8 TP1, Director-validated, byte-identical) at a 53.51% head; iso 1.46× geomean, ZERO shipped coverage, ZERO HBM** — (fp8-a8w8-blockscale-bpreshuffle-ck-tune-gfx950-sglang.md) - [gfx942 · sglang fp8 a8w8 blockscale] **MANDATED LEVER = the CK skill** `gemm_tuning/fp8_gemm_tuning_sglang_aiter.md` (capture live (M,N,K) → aiter CK tuner → fp8_utils Triton→CK switch overlay + `AITER_CONFIG_GEMM_A8W8_BLOCKSCALE`); baseline = the UNTUNED Triton default, so CK-tuned is the real win. The old per-(N,K) Triton config-JSON overlay is **DEPRECATED for this op (do NOT use it — it keeps the slow Triton seam live and bypasses the skill)** — (fp8-a8w8-blockscale-overlay-gfx942.md) - [gfx950 · vLLM MXFP8 E8M0] dense `tl.dot_scaled` STATIC tiles (decode BK256/prefill BM128) ★★★ part of +12.1% e2e — (mxfp8-microscale-gemm-gfx950.md) @@ -24,12 +24,13 @@ Confidence (a hint strength, not authority): ★ noise/unverified · ★★ sing ## MoE grouped GEMM - [gfx950 · vLLM+sglang MXFP8 1×32 E8M0 grouped fused-MoE] live = editable Triton `tl.dot_scaled` (hardcoded tiles, no autotune, no split-K) → NO env/flag win; author is the lever (flydsl FIRST, then triton rewrite) ★★★ **iso 1.0× baseline, head 14.8–32.6%; two dead-ends pre-recorded (aiter bf16 DB, op_bench blockscale misroute)** — (mxfp8-e8m0-grouped-moe-gfx950.md) -- [gfx950 · vLLM mxfp4 grouped fused-MoE, decode-dominated] seam is the FUSED `triton_kernel_moe_forward`, not a standalone GEMM → author a whole-file Triton replacement ★★★ **+92.5% e2e @ 8.32% head, +1.83% @ ~15.6% (both Director byte-exact); exceeds naive Amdahl (launch-overhead-bound); isolated GEMM only 1.06–1.57× — don't size from it. CAUTION: OCP-MX emulation sub-variant needs a NATIVE-fp4 swap instead** — (moe-mxfp4-grouped-authored-gfx950-vllm.md) +- [gfx950 · vLLM mxfp4 grouped fused-MoE, decode-dominated] seam is the FUSED `triton_kernel_moe_forward`, not a standalone GEMM → author a whole-file Triton replacement ★★★ **+92.5% e2e @ 8.32% head, +26.9% @ 32.43% head, +1.83% @ ~15.6% (all byte-exact); exceeds naive Amdahl (launch-overhead-bound); isolated GEMM only 1.06–1.57× — don't size from it. CAUTION: a split_k/tile rewrite can break byte-parity → corrective re-author preserving accum order recovers the win; routing-metadata chain is not standalone-extractable (fold into rewrite); OCP-MX emulation sub-variant needs a NATIVE-fp4 swap instead** — (moe-mxfp4-grouped-authored-gfx950-vllm.md) - [gfx950 · vLLM MXFP8 E8M0] grouped `dot_scaled` STATIC tiles (GEMM1-decode BN64+BK256) ★★★ part of +12.1% e2e — (mxfp8-microscale-gemm-gfx950.md) - [gfx950 · vLLM MXFP8 E8M0 grouped MoE] author/rewrite `_mxfp8_grouped_gemm_kernel`, tune BLOCK_K 128→256 ★★ iso ~1.08x geomean (all 6 buckets); e2e unverified — (mxfp8-grouped-moe-gemm-gfx950.md) - [gfx950 · vLLM MXFP8 E8M0 grouped MoE] ROUTING: editable in-tree Triton `dot_scaled`; aiter bf16 GEMM DB is the WRONG lever ★★ iso baseline only — (mxfp8-e8m0-grouped-moe-gfx950.md) - [gfx942 · vLLM int4 W4A16] per-shape fused-MoE Triton config tune via `VLLM_TUNED_CONFIG_FOLDER` (env, ZERO HBM; N=moe_int//TP) ★★★ +11-18% e2e (10 confirms, TP8 & TP4) — (moe-int4-w4a16-tune-gfx942.md) - [gfx942+gfx950 · vLLM bf16 MoE] same per-shape config-tune lever as int4, for DENSE bf16 (dtype=None filename); nothing ships for unseen `(E,N,device)` → default fallback. Sweep per M-bucket, deploy `VLLM_TUNED_CONFIG_FOLDER` + `--max-num-batched-tokens ≈2·ISL`, ZERO HBM ★★★ **iso 1.01–1.66×/bucket, serving-wtd 1.10–1.40×, +7.01% e2e (Mixtral-8x7B gfx950 TP8, Director byte-exact); 3 confirms** — (moe-bf16-tune-gfx942.md) +- [gfx950 · vLLM fp8_w8a8 block-scale MoE] same `VLLM_TUNED_CONFIG_FOLDER` lever, dtype=fp8_w8a8+block_shape filename; nothing ships for unseen `(E,N,MI355)` → default fallback. fp8-specific: pin BLOCK_SIZE_K=128, BLOCK_SIZE_N∈{128,256}; build qc via `fp8_w8a8_moe_quant_config(...,block_shape=[128,128])`. Authored Triton rewrite of `fused_experts_impl` beat env-tune (iso 1.034×). ZERO HBM ★★ **iso 1.02–1.16×/bucket, serving-wtd ~1.03× (Qwen3.5-122B-A10B-FP8 TP2, 21.4% head); e2e = validated_no_win: milestone +1.36% non-overlapping byte-exact COLLAPSED to Director same-session +0.16% (1.0016×, overlapping); trust the Director A/B, ~20% head × iso 1.03× is inside noise; aiter fmoe_fp8_blockscale_g1u1 swap is a separate Tier-A candidate** — (moe-fp8-blockscale-tune-gfx950.md) ## attention - [gfx942 · sglang hybrid prefill] `--attention-backend triton` cheap flag win ★★★ +~5% e2e — (attention-backend-triton-gfx942.md) diff --git a/e2e_workflow/knowledge/learned/fp8-a8w8-blockscale-ck-tune-gfx950-vllm.md b/e2e_workflow/knowledge/learned/fp8-a8w8-blockscale-ck-tune-gfx950-vllm.md index c618c0fd9..b1a7cff9b 100644 --- a/e2e_workflow/knowledge/learned/fp8-a8w8-blockscale-ck-tune-gfx950-vllm.md +++ b/e2e_workflow/knowledge/learned/fp8-a8w8-blockscale-ck-tune-gfx950-vllm.md @@ -3,7 +3,7 @@ key: fp8 a8w8 blockscale GEMM · gfx950 · vLLM prefill+decode type: lever confidence: ★★★ effect: e2e +16.1% / +65.69% / +18.86% on three models; iso 1.86–3.17× serving-weighted. Gain is gated by SHIPPED aiter table coverage — ~0 when the head shapes already bind tuned (measured 1.03× on a saturated model). Probe coverage before budgeting. -last_seen: 2026-08-16 +last_seen: 2026-08-19 --- # gfx950 vLLM fp8 a8w8 blockscale — per-shape CK tune DB (no overlay needed) @@ -31,6 +31,15 @@ last_seen: 2026-08-16 the same lever landed 1.067×. The +16.1% (Qwen3.5-27B) is the weakest of the three: parity n/a under the fp8 accuracy gate, tuned-shape binds unreproduced, and a server-flag bundle confounds it. If ckProfiler is absent the CK *author* lane is unavailable, but the tune DB still applies. +- confirm (2026-08-19, Qwen3-14B-FP8 TP1, gfx950/MI355, head 67.3% GPU): swap-only aiter-linear + Triton->CK (VLLM_ROCM_USE_AITER[_LINEAR]=1), UNTUNED CK, measured on the immutable unittest + (`aiter.gemm_a8w8_blockscale`, non-transposed scale, parity rel~7e-3 << TOL 0.05): serving-weighted + 1.513x, geomean 1.28x. Regime split as the card warns — decode M1/M64 1.8-2.3x, prefill M571 REGRESSES + 0.66-0.84x (untuned CK). The per-shape CK tune (recovers prefill -> card's 1.86-3.17x) was NOT runnable + here: the offline tuner `csrc/ck_gemm_a8w8_blockscale/gemm_a8w8_blockscale_tune.py` + ckProfiler are + ABSENT from this image's aiter wheel (no csrc dir); foreign-checkout tuners exist on NFS but version- + mismatch the installed aiter -> unsafe. So shipped swap-only; e2e prefill-regression risk left to the + Integrator gate + operator-provisioned tuner. - source: exp/e2e_*Qwen3.5-27B-FP8*/ 2026-08-12; exp/e2e_*Qwen3-14B-FP8*/ 2026-08-13 (Director-validated_win TP1, head 67.96% GPU); exp/e2e_*Qwen3.5-122B-A10B-FP8*/ 2026-08-13 (Director-validated_win TP2, head 19.78% GPU); exp/e2e_*DeepSeek-V4-Flash-0731*/ 2026-08-13 diff --git a/e2e_workflow/knowledge/learned/moe-fp8-blockscale-tune-gfx950.md b/e2e_workflow/knowledge/learned/moe-fp8-blockscale-tune-gfx950.md new file mode 100644 index 000000000..1f2a29e50 --- /dev/null +++ b/e2e_workflow/knowledge/learned/moe-fp8-blockscale-tune-gfx950.md @@ -0,0 +1,53 @@ +--- +key: fp8_w8a8 block-scale fused-MoE grouped GEMM · gfx950 · vLLM +type: lever +confidence: ★★ +confirms: 2 +effect: per-shape Triton config tune (winner_kind=env, ZERO HBM) → iso 1.02–1.16× per M-bucket, serving-weighted ~1.03× (decode M64 1.026×, prefill M8192 1.041×). Same VLLM_TUNED_CONFIG_FOLDER mechanism as the int4/bf16 MoE cards, dtype segment = fp8_w8a8 + block_shape. An authored Triton rewrite of `fused_experts_impl` (Tier-C) beat the env-tune bake-off (iso 1.034×). e2e transfer did NOT clear the noise band at the FINAL gate: Director same-session A/B = +0.16% (1.0016×), ranges OVERLAP → validated_no_win, byte-exact parity 12/12. The lever ENGAGES (decode-bucket rebind fired on both TP workers) but at a 21.4% head with iso ~1.03× the Amdahl ceiling (~0.6%) is inside serving noise. +last_seen: 2026-08-20 +--- +# fp8_w8a8 block-scale fused-MoE → the memory-free vLLM config-tune lever (fp8 analog of int4/bf16 cards) + +- path: same as `moe-bf16-tune` / `moe-int4-w4a16-tune` but for fp8 block-quant. (1) check whether a + tuned config ships for `(E,N,device,fp8_w8a8,block_shape)` — vLLM ships NONE for unseen fp8-blockscale + MoE shapes on gfx950/MI355 (verified 0 configs match `*MI355*fp8_w8a8*block_shape*`), so the expert + grouped-GEMM falls back to the slow default tile (`Using default MoE config`). (2) Sweep per M-bucket + against `fused_experts`+`override_config` with fp8 weights + block scales, parity rel<1e-2. (3) Deploy + `VLLM_TUNED_CONFIG_FOLDER`, pair with `--max-num-batched-tokens ≈2·ISL` (clamp 8192..32768). +- lookup filename: `get_config_file_name(E, N, "fp8_w8a8", [128,128])` → + `E=,N=,device_name=,dtype=fp8_w8a8,block_shape=[128,128].json`. N = moe_intermediate//TP. +- fp8-SPECIFIC constraint (differs from bf16/int4 sweep): the block-scale kernel requires + `BLOCK_SIZE_K % block_k == 0` and `BLOCK_SIZE_N % block_n == 0` → pin BLOCK_SIZE_K=128, BLOCK_SIZE_N∈{128,256}. + Build the quant_config via `fp8_w8a8_moe_quant_config(w1_scale,w2_scale,block_shape=[128,128])`; weights + float8_e4m3fn, scales float32 shaped [E,ceil(2N/128),ceil(K/128)] / [E,ceil(K/128),ceil(N/128)]. +- expected gain: iso 1.02–1.16× per bucket (mid buckets M128/256 biggest at ~1.15×; decode M64 1.026×, + prefill M8192 1.041×), serving-weighted ~1.03×. ZERO extra HBM → sails the mem_footprint gate. + Naive Amdahl ceiling ~0.6% at a 21.4% head, BUT decode is graph-hidden/under-counted (profiled decode + share 0.00 → floored 0.30) so measured e2e may exceed it (cf. bf16 card +7.01% > +3.37% ceiling). +- also: editable in-tree Triton MoE present (`kernel_src/fused_moe/fused_moe.py`, hot `fused_moe_kernel` + @triton.jit + `fused_experts_impl`) → Tier-C `route=rewrite`; flydsl grouped-MoE primitives import on + this gfx950 image (`aiter.ops.flydsl.flydsl_moe_stage1/stage2`, is_flydsl_available=True) → `route=author`. + aiter ALSO ships a native fp8 block-scale fused MoE (`aiter.fmoe_fp8_blockscale_g1u1`) reachable via the + vLLM `VLLM_ROCM_USE_AITER[_MOE]` backend swap — a separate Tier-A candidate the Integrator can A/B + (mutually exclusive with the Triton config env, since aiter MoE bypasses the Triton seam). +- caution (also verify flydsl viability before routing it): for fp8-[128,128]-block MoE the aiter FlyDSL + path did NOT compile on this gfx950 image — the high-level `flydsl_moe_stage1/2` wrapper accepts only + b_dtype∈{fp4,fp8 MXFP8 per-32 e8m0, bf16xint4} and raises ValueError on bf16xbf16, and the + precision-preserving fallback (dequant fp8-block→bf16/fp16 then `compile_moe_gemm1/2`) hits an internal + DSL compiler bug (`UnboundLocalError 'a0'` in the non-int4 prefetch pipeline, `moe_gemm_2stage.py`, + reproduced across tiles). So `is_flydsl_available=True` (imports OK) is NOT sufficient — verify the + actual dtype path compiles; for [128,128]-block fp8 MoE prefer the Triton rewrite. flydsl would need a + fixed non-int4 `moe_gemm_2stage` build or an MXFP8 requant path validated to the tol. +- caution (also verify): a milestone interleaved A/B here showed +1.36% NON-overlapping byte-exact for + the authored Triton rewrite, but the Director SAME-SESSION A/B collapsed it to +0.16% (overlapping + ranges) = validated_no_win. On a decode-bound serving run always trust the Director same-session A/B + over the milestone A/B: a milestone win at a ~20% head with iso ~1.03× can be entirely serving noise + once re-measured against a fresh same-session baseline (base median rose from 2405.9 warm-start to + 2451.5 same-session — most of the apparent gain was baseline drift). Byte-exact parity is NOT evidence + of a throughput win. +- source: exp/e2e_*Qwen3.5-122B-A10B-FP8*/ 2026-08-19..08-20 (E=256, N=512, K=3072, topk=8, silu, fp8 + block[128,128], vLLM 0.26.0, TP=2, gfx950/MI355 OAM, MoE 21.4% GPU; no shipped config → default + fallback; iso per-bucket 1.015–1.158×, authored iso 1.034×; Director same-session +0.16% (1.0016×), + validated_no_win, byte-exact 12/12). + driver: `config/tune_moe_fp8_blockscale.py`; tuned artifact under + `config/moe_tuned_fp8/E=...,N=...,dtype=fp8_w8a8,block_shape=[128,128].json`. diff --git a/e2e_workflow/knowledge/learned/moe-mxfp4-grouped-authored-gfx950-vllm.md b/e2e_workflow/knowledge/learned/moe-mxfp4-grouped-authored-gfx950-vllm.md index 6a92cc361..5e69f93bf 100644 --- a/e2e_workflow/knowledge/learned/moe-mxfp4-grouped-authored-gfx950-vllm.md +++ b/e2e_workflow/knowledge/learned/moe-mxfp4-grouped-authored-gfx950-vllm.md @@ -2,9 +2,9 @@ key: moe-grouped-gemm-mxfp4 · gfx950 · decode-dominated (prefill present) type: routing confidence: ★★★ -confirms: 3 -effect: authored whole-file Triton rewrite of the fused seam = +92.5% e2e (gpt-oss-120b TP2, head 8.32% GPU) and +1.83% (DeepSeek-V4-Flash TP4, head ~15.6%) — both Director-verified byte-exact. Isolated grouped-GEMM only shows 1.06–1.57×; it structurally undercounts the live decode win. -last_seen: 2026-08-17 +confirms: 4 +effect: authored whole-file Triton rewrite of the fused seam = +92.5% e2e (gpt-oss-120b TP2, head 8.32% GPU), +26.9% (gpt-oss-120b TP2, head 32.43%, byte-exact after a corrective re-author; finalize/Director-validated_win, full-run 5799→7042 tok/s = 1.285×) and +1.83% (DeepSeek-V4-Flash TP4, head ~15.6%) — all Director/Integrator-verified byte-exact. Isolated grouped-GEMM only shows 1.06–1.57×; it structurally undercounts the live decode win. Higher head share does NOT mean bigger e2e: at 32% the MoE is memory-bound at the HBM wall so the win is the launch-overhead/decode-seam share, not micro-tuning. +last_seen: 2026-08-19 --- # MXFP4 grouped fused-MoE (gpt-oss style) — author a whole-file Triton replacement, not a GEMM swap @@ -23,7 +23,17 @@ last_seen: 2026-08-17 baseline, ≥2-shape CUDA-graph replay), then a same-session e2e A/B with BYTE-EXACT greedy parity (temp=0/seed=0/ignore_eos) against a FRESH no-overlay baseline. - caution: trust the byte-exact e2e gate over the isolated ×, and don't drop the head on a modest - isolated number. Do NOT route a fused seam to standalone-gemm-swap or dense-linear-env-overlay — + isolated number. A non-quant tile-shape rewrite (block_m 32→16 + split_k + routing-metadata reuse) + can SILENTLY BREAK byte-exact greedy parity vs a deterministic baseline (7/12 temp=0 prompts + diverged, one at the first token) even while e2e is +29% and engagement is proven on every TP + worker — a faster-but-wrong server is rejected. When it happens, route to a CORRECTIVE re-author + that preserves accumulation order (avoid split_k / order-changing reduction); it recovered byte + parity AND kept +26.9% e2e. Also verify: the routing-metadata chain + (`_topk_forward`/`pack_bitmatrix`/`_bitmatrix_*`/`_sum_bitmatrix_rows`/`_stage2_pow2`) is NOT + cleanly standalone-extractable on the v3.6.0 SparseMatrix path (split across `topk_fn` + + `make_routing_data`, returns structured RoutingData/Gather/Scatter, host-launch-bound) — fold it + INTO the whole-file rewrite rather than scheduling it as its own unittest. +- caution: Do NOT route a fused seam to standalone-gemm-swap or dense-linear-env-overlay — there is no call site to bind. flydsl would have to invert two proprietary swizzles (triton_kernels `Tensor.storage` + vLLM CDNA4 mxfp4 scale) plus the w13 shuffle: high correctness risk, and `aiter.ops.flydsl.moe_kernels` fails to import. @@ -40,4 +50,13 @@ last_seen: 2026-08-17 bakeoff only — the earlier run's per-head attribution was reversed to dead_end/implausible_speedup at review, so it backs the SEAM identity but not a speedup); exp/e2e_*Qwen3.5-397B-A17B-MXFP4*/ 2026-08-16 (the emulation sub-variant; fast path never engaged, - iso 0.972× no-op — no verified win). + iso 0.972× no-op — no verified win); + exp/e2e_*gpt-oss-120b*/ 2026-08-19 (gpt-oss-120b TP2 head 32.43%: seam = + `triton_kernel_moe_forward`/`gpt_oss_triton_kernels_moe.py` editable Triton; op_bench found NO env/flag + lever (delegated to server-flag path, oracle-only). First author (block_m 32→16 + split_k + + routing-metadata reuse) engaged 2/2 workers, e2e +29.15% (5267→6802) but REJECTED on byte-parity fail; + corrective re-author preserving accumulation order = ACCEPTED, byte-parity pass, +26.9% head e2e + (Integrator A/B 5612→7122, non-overlapping) — CONFIRMED by the finalize/Director gate: full-run + 5799.19→7042.04 tok/s = 1.2848× (+28.48%), validated_win, output parity pass. Post-win the MoE + grouped GEMM is still #1 ~17% and + memory-bound at the HBM wall (hbm_util ~0.97–1.01) — only byte-reduction headroom remains). diff --git a/e2e_workflow/roles/_fragments/warm_start.md b/e2e_workflow/roles/_fragments/warm_start.md new file mode 100644 index 000000000..890ce4183 --- /dev/null +++ b/e2e_workflow/roles/_fragments/warm_start.md @@ -0,0 +1,54 @@ +# Warm start — how to use a prior run's record + +The knowledge base holds results from earlier runs on this same deployment (model, gfx, serving +framework and version, precision, TP, and workload point). Before your phase started, the +orchestrator read that page and **benched what it found on this box**. You are being handed both +halves: what the store claimed, and what this run measured. + +## The one rule + +**A stored number is a hypothesis. Only the measured column is evidence.** + +Every file in `KB_REFERENCE_DIR/` named `e2e_reference_*.md` records what another box reported — +possibly on a different day, a different ROCm build, and against a different baseline. The file +`measured_on_this_box.md` records what happened when this run applied the same thing here, through +the same gate a fresh idea faces. **Where the two disagree, the measured file wins.** It is not a +correction of the stored record; the disagreement is itself the useful datum, which is why both +files are kept. + +Never quote a stored throughput as this deployment's number. The only throughput this run may claim +is one it measured. + +## What each outcome means for you + +**`adopted` (a configuration).** It is already in your `CURRENT_FLAGS` / `CURRENT_ENV`. It is part of +the starting point, not a proposal. Re-proposing it measures the current state against itself and +burns a server launch to learn nothing. Propose only things that **compound on top of it**. + +**`adopted` (a kernel).** It is already in the active overlay, and the profile you are routing from +was captured *with* it applied. That op is done. Its share of GPU time in your Top-N already +reflects the improvement — do not read the reduced percentage as a fresh opportunity. + +**`rejected`.** The compounded whole did not beat this baseline on this box. Three things this does +**not** mean: + +- it does not mean each knob inside it is worthless — a stored config is applied as one unit, so a + single regressing axis can sink an otherwise-good set; +- it does not mean the *direction* is wrong — the idea may simply need to be reached differently + here; +- it does not mean the record was false — it may well have been true on the box that wrote it. + +So: **do not re-propose a rejected entry verbatim.** Do feel free to propose one axis out of it, or +the same underlying idea approached another way, as an ordinary candidate that stands on its own +rationale. + +**`reference` / `incomplete`.** It was never measured here — because it could not be replayed from a +record, or because the run's replay budget was spent, or because the read was in reference-only +mode. Treat it exactly as a lead from a colleague: worth a look, worth nothing until measured. + +## A trap specific to recovered configurations + +A flag that a newer framework version renamed or removed is accepted silently on the command line +and then ignored. On the measurement that looks identical to "this config made no difference." If +you propose anything recovered from the store, **verify from the server log that the flag was +actually honoured** before you believe a null result. diff --git a/e2e_workflow/roles/director.md b/e2e_workflow/roles/director.md index 128f8720e..caffbf965 100644 --- a/e2e_workflow/roles/director.md +++ b/e2e_workflow/roles/director.md @@ -116,10 +116,21 @@ Return JSON: "tp": 1, "workload": {"isl": 1024, "osl": 1024, "conc": 64}, "bench_script": "/bench_e2e.sh", + "gfx": "", + "precision": "", + "framework_version": "", + "rocm_version": "., e.g. 7.2 — or \"\" if not established>", "notes": "sglang version, anything unusual" } ``` +The last four are the dimensions the deployment knowledge base addresses a record by, and you have +already established every one of them in step 4 to launch the server at all. **Never guess one.** +An empty string files this run under a deliberately coarse `unknown` page, which is honest and +recoverable; a plausible-looking wrong value files it under an authoritative page, and the store has +no delete — a bad record there can only be outranked, never removed. If a value is genuinely +unknown, `""` is the correct answer. + --- ## PHASE=validate diff --git a/e2e_workflow/scripts/e2e_store.py b/e2e_workflow/scripts/e2e_store.py new file mode 100644 index 000000000..2fa22d0ca --- /dev/null +++ b/e2e_workflow/scripts/e2e_store.py @@ -0,0 +1,1238 @@ +#!/usr/bin/env python3 +"""Warm start and write-back for the e2e serving lane, over the same two planes as the kernel lane. + + e2e_store.py identity --model Qwen3-397B --gfx gfx950 --framework vllm \ + --framework-version 0.26.0 --precision mxfp8 --tp 8 --isl 1024 --osl 1024 --conc 64 + e2e_store.py resolve ... --plane remote --refs-dir REFS --cache-dir CACHE + e2e_store.py write ... --plane both --store DIR --result run.json --apply + +`resolve` answers the question the e2e Director asks at Setup — "has anyone already tuned this +deployment, and what did they land on" — and `write` is what makes the next run's answer non-empty. +Addresses come from `kb.identity.e2e_canonical_ids`, never from string formatting here, because a +reader and a writer that disagree by one segment do not raise: the run just cold starts. + +WHAT AN E2E RECORD IS FOR, and why it is not a kernel record with different fields. A kernel entry +offers a patch to apply. An e2e entry mostly offers a CONFIG — env flags, server args, a launch +script — plus a list of kernels that turned out to be worth overlaying. The patch is optional and +often absent (a config-only win is a real win), so `resolve` never treats a missing artifact as a +broken record the way the kernel lane does. + +THE THREE RUNGS ANSWER THREE DIFFERENT QUESTIONS. They are not a fallback chain that happens to have +three links; each one is the correct page for a different ask, which is why all three are always +written: + + ...:mxfp8:tp_8:isl_1024:osl_1024:conc_64 "tune THIS benchmark point" + ...:mxfp8:tp_8 "given TP=8, how do I configure the server" + ...:mxfp8 "how many ways should I shard this model at all" + +Only the last one can compare TP4 against TP8, because that comparison needs them filed together. + +READS AND WRITES RANK ON DIFFERENT METRICS, deliberately, and the split is the one thing to keep +straight in this file: + + * READING (`resolve`) orders every rung by absolute `throughput_tok_s`, high to low. That is what + a Director asking "what should I run" wants on every page, including the coarse ones: it is + choosing a config to spend a server launch on, and the fastest observed deployment is the + honest first offer. The coarse rungs' numbers were measured at other workload points and are + NOT comparable to this run's baseline — `_render_reference` says so in as many words — but a + speedup ordering there is not more comparable, it just hides the incomparability behind a + ratio. `--sort-by speedup` restores the old ordering for a caller that wants it. + * WRITING (`write`) keeps the per-rung champion metric from `rung_metric()`: throughput on the + exact rung, speedup on the coarser ones. The champion pointer is a promotion gate, not a + display order, and a coarse page promoting on absolute tokens/sec would crown whichever + workload point happens to be cheapest rather than whichever run improved anything. + +Both scalars are written flat at the top of every document (the service's `sessions/top?metric=` +reads a top-level scalar and rejects a nested path), which is what lets a read rank on one while +the champion is kept on the other without a second write. + +No delete exists on the service. Every `--apply` is permanent. +""" + +import argparse +import hashlib +import json +import os +import shutil +import sys +import tempfile +import time + +# The shared KB plane lives at the repo root as the `kb` package, not beside this file. Executed as +# a CLI from an arbitrary cwd, so the root is derived from __file__ and never from the environment. +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +from kb import identity as kbid # noqa: E402 +from kb.attest import (OUTCOMES, attest_session, attestation_ok, # noqa: E402 + attestations_of, carry_attestations, retire_hint) +from kb.curate import collapse_by_direction # noqa: E402 +from kb.ladder import publish # noqa: E402 +from kb.plane import open_plane # noqa: E402 +from kb.retract import is_retired, retract_session, retraction_ok # noqa: E402 +from kb.store_local import KBStoreError, finite_speedup # noqa: E402 + +SCHEMA = "geak.e2e.v1" +THROUGHPUT_METRIC = "throughput_tok_s" # ranks the exact-workload rung, and every read +SPEEDUP_METRIC = "speedup" # champion metric on every coarser rung +DEFAULT_TOP_N = 3 +DEFAULT_SCAN = 25 +# What `resolve` orders a page by, by name on the CLI. See the module docstring for why a read and +# a write do not agree on this. +SORT_METRICS = {"throughput": THROUGHPUT_METRIC, "speedup": SPEEDUP_METRIC} +DEFAULT_SORT_BY = "throughput" + +# Which metric ranks which rung, indexed the same way e2e_canonical_ids() returns them. A ladder +# shorter than three (the run recorded no tp, or no workload shape) drops rungs from the FRONT, so +# the table is applied by counting back from the end rather than by position. +_COARSE = (SPEEDUP_METRIC, 1.0) +_EXACT = (THROUGHPUT_METRIC, 0.0) + + +def rung_metric(index: int, total: int): + """(metric, promote_floor) for rung `index` of `total`, most specific first.""" + is_exact_workload = (index == 0 and total == 3) + return _EXACT if is_exact_workload else _COARSE + + +# -- identity ------------------------------------------------------------------------------------ + + +def identity_of(a) -> dict: + return kbid.e2e_identity(a.model, a.gfx, a.framework, a.framework_version, a.precision, + tp=a.tp, isl=a.isl, osl=a.osl, conc=a.conc) + + +def ladder_of(a): + """[(canonical_id, tier, metric, floor)], most specific first. + + Tiers are named after what was DROPPED, not after how good the match is, so a caller logging + `workload_any` can see at a glance that the numbers it is looking at were measured on some other + benchmark point and must not be quoted as this deployment's throughput. + """ + cids = kbid.e2e_canonical_ids(identity_of(a)) + tiers = {3: ("exact", "workload_any", "tp_any"), 2: ("exact", "tp_any"), 1: ("exact",)}[len(cids)] + return [(cid, tier) + rung_metric(i, len(cids)) + for i, (cid, tier) in enumerate(zip(cids, tiers))] + + +# -- planes -------------------------------------------------------------------------------------- + + +# -- read ---------------------------------------------------------------------------------------- + + +def _view(candidate, cid: str, tier: str, metric: str, champion_metric: str = "") -> dict: + """One offered record, flattened to what a Director prompt actually needs.""" + knowledge = candidate.knowledge if isinstance(candidate.knowledge, dict) else {} + value = knowledge.get("value") if isinstance(knowledge.get("value"), dict) else {} + workload = value.get("workload") if isinstance(value.get("workload"), dict) else {} + ledger = attestations_of(value) + return { + "session_id": candidate.session_id, + "canonical_id": cid, + "match_tier": tier, + "ranked_by": metric, + # Which metric the CHAMPION on this page was promoted under, which on a coarse rung is not + # the one the offer is ordered by. Spelled out so `is_champion` cannot be misread as "the + # top of this list" — on a coarse rung the champion is the best speedup and the first row + # is the best throughput, and those are routinely different records. + "champion_metric": champion_metric or metric, + "score": candidate.speedup, + "throughput_tok_s": finite_speedup(knowledge.get(THROUGHPUT_METRIC)), + "speedup": finite_speedup(knowledge.get(SPEEDUP_METRIC)), + "baseline_throughput_tok_s": finite_speedup(value.get("baseline_throughput_tok_s")), + "direction": str(value.get("direction") or ""), + "workload": workload, + "accepted_config": value.get("accepted_config") if isinstance( + value.get("accepted_config"), dict) else {}, + "accepted_kernels": [k for k in (value.get("accepted_kernels") or []) + if isinstance(k, dict)][:32], + "validation_status": str(value.get("validation_status") or ""), + # How much to believe the number, as the writer judged it. Surfaced rather than used as a + # filter: an unvalidated record is still a lead worth benching, and dropping it would make + # the warm start narrower exactly where it is cheapest to be broad. Only a RETRACTED record + # is filtered, because that one has been positively declared false. + "validated": bool(value.get("validated")), + "validation_basis": str(value.get("validation_basis") or "unverified"), + "parity": str(value.get("parity") or ""), + "lifecycle": str(value.get("lifecycle") or ""), + "upstream": value.get("upstream") if isinstance(value.get("upstream"), dict) else {}, + # What this record has DONE SINCE it was written, as opposed to what its writer claimed + # about it. `validated` above is one box's judgement of its own run; these are everyone + # else's. A reader deciding whether to spend a server launch wants both, and they + # disagree often enough that collapsing them would be a lie in one direction or the other. + "validations": ledger["validations"], + "recalls": ledger["recalls"], + "not_reproduced": ledger["not_reproduced"], + "last_outcome": ledger["last_outcome"], + "retire_hint": retire_hint(value), + # How to actually run this again. Empty for records written before the field existed — + # those are the ones a reader has to reconstruct from `accepted_config` by hand. + "repro": value.get("repro") if isinstance(value.get("repro"), dict) else {}, + "is_champion": bool(candidate.is_champion), + } + + +def read_metric(a) -> str: + """The metric a READ orders every rung by. `throughput_tok_s` unless the caller says otherwise. + + Applied to the store itself and not only to the list it hands back, which matters on the remote + plane: `candidates()` pages `sessions/top?metric=` and hydrates the first `--scan` rows, so + fetching a page ordered by speedup and then re-sorting it locally by throughput would rank a + biased sample and quietly drop the fast-but-modest-ratio records that never made the cut. + """ + return SORT_METRICS.get(str(getattr(a, "sort_by", "") or DEFAULT_SORT_BY).strip().lower(), + THROUGHPUT_METRIC) + + +def _sort_key(metric: str): + """Rank order for a view list: the chosen metric first, the other as tie-break, id last. + + A record missing the metric sorts last rather than raising — a write cannot produce one (final + throughput is required), but a document written by hand or by an older lane can. + """ + other = SPEEDUP_METRIC if metric == THROUGHPUT_METRIC else THROUGHPUT_METRIC + low = float("-inf") + return lambda v: (-(v.get(metric) if v.get(metric) is not None else low), + -(v.get(other) if v.get(other) is not None else low), + v["session_id"]) + + +def cmd_resolve(a) -> dict: + ladder = ladder_of(a) + metric = read_metric(a) + # Echo the plane back. A caller that tries the service and falls back to disk otherwise cannot + # tell from the output which one answered — the ladder, the ranking and the shapes are identical + # — and "where did this candidate come from" is the first question asked when one turns out to + # be wrong. `dict(out, ...)` carries it onto every return path below. + out = {"tried": [c for c, _t, _m, _f in ladder], "canonical_id": ladder[0][0], + "match_tier": "", "ranked_by": "", "sorted_by": metric, "champion_metric": "", + "candidates": [], "read_reason": "", + "plane": str(getattr(a, "plane", "local") or "local"), "curation": {}} + last_why = "" + for cid, tier, champion_metric, floor in ladder: + # The rung's own metric opens nothing here: a read ranks every rung the same way (see + # read_metric), and the floor only ever gates a promotion, which a read never performs. + store, _mirror, why = open_plane(a, metric, floor) + if store is None: + last_why = why + continue + try: + found = store.candidates(cid, limit=max(1, int(a.scan))) + except Exception as e: + last_why = "read_failed: %s: %s" % (type(e).__name__, str(e)[:120]) + continue + # Retracted records are dropped BEFORE anything else looks at them, and before the + # direction collapse in particular: a retracted entry that happens to rank first for its + # direction would otherwise evict the surviving alternatives for that same direction, so a + # single false record could hide every good one behind it. Done client-side because it has + # to be — retraction zeroes the ranking scalar and re-points the champion, but the service + # still serves the session, and nothing in the scheme lets us ask it not to. + kept = [c for c in found if not is_retired(c.value)] + curation = {"scanned": len(found), "retired": len(found) - len(kept), + "sorted_by": metric} + # Re-sorted here even though the store already ordered by this metric, because the two + # planes order by slightly different things: the local one ranks on the document scalar, + # the remote one on the score the service computed and falls back to the document only + # when that is absent. Sorting the hydrated views is the one place both planes are + # guaranteed to agree, and collapse_by_direction's contract is that its input is already + # in rank order — it keeps the FIRST entry per direction, so a wrong order here silently + # offers the wrong member of every group. + ordered = sorted([_view(c, cid, tier, metric, champion_metric) for c in kept], + key=_sort_key(metric)) + # top_n=len(kept): e2e filters min_speedup AFTER collapse and slices to top_n below, so + # collapse must not pre-slice. It consumes only the per-idea best, not the alternates. + views, _alternates, collapsed = collapse_by_direction( + ordered, lambda v: v["direction"], lambda v: v["session_id"], len(kept)) + curation["same_direction_collapsed"] = collapsed + if a.min_speedup: + # Applied to `speedup` on every rung, including the throughput-ranked one: the floor + # asks "did this run actually improve anything", which is a question about the ratio no + # matter what the page is sorted by. A record with no speedup recorded is kept — it may + # still be a usable config — rather than silently failing an unanswerable test. + before = len(views) + views = [v for v in views + if v["speedup"] is None or v["speedup"] >= float(a.min_speedup)] + curation["below_min_speedup"] = before - len(views) + curation["min_speedup"] = float(a.min_speedup or 0.0) + if not views: + # A rung whose every candidate was curated away is NOT an empty page, and the next rung + # down is about to be tried as if it were. Carry the counts forward so the caller can + # tell "nobody has recorded this" from "everything recorded here was retracted" — + # identical read_reasons otherwise, opposite meanings. + out["curation"] = dict(curation, canonical_id=cid, tier=tier) + continue + views = views[: max(1, int(a.top_n))] + if a.cache_dir: + for view in views: + view["bundle"] = _materialize(store, cid, found, view, a.cache_dir) + if a.refs_dir: + _render_reference(a.refs_dir, cid, tier, views) + return dict(out, canonical_id=cid, match_tier=tier, ranked_by=metric, + champion_metric=champion_metric, candidates=views, read_reason="read", + curation=dict(curation, canonical_id=cid, tier=tier)) + return dict(out, read_reason=last_why or "e2e_page_not_found") + + +def _materialize(store, cid: str, found, view: dict, cache_dir: str) -> dict: + """Pull one record's artifacts down, reporting a failure instead of raising. + + An e2e record is usable without its files — the config lives in the knowledge document — so a + download problem degrades the offer rather than dropping it. + """ + candidate = next((c for c in found if c.session_id == view["session_id"]), None) + if candidate is None: + return {"error": "candidate vanished between ranking and download"} + try: + path = store.materialize(cid, candidate, cache_dir) + except (KBStoreError, OSError) as e: + return {"error": str(e)[:200]} + # Walked, not listdir'd: artifacts nest (`kernels/.patch`), and a flat listing reports the + # directory itself as if it were the file, which is exactly the name a caller would then fail to + # open. + files_root = os.path.join(path, "files") + names = sorted(os.path.relpath(os.path.join(root, f), files_root).replace(os.sep, "/") + for root, _dirs, found in os.walk(files_root) for f in found) + return {"path": path, "files": names} + + +def _kernel_line(kernels) -> str: + """`moe_stage1 (ck, 1.84x, kernels/moe_stage1.patch)` — one readable line per kernel. + + Spells out the patch path because the Director reads this prose and then has to go open the + file; a name alone sends it back to the store to ask a question this page already answered. + """ + parts = [] + for k in kernels: + bits = [b for b in (k.get("language"), + "%sx" % k["isolated_speedup"] if k.get("isolated_speedup") else "", + k.get("patch") or k.get("kernel_canonical_id") or "") if b] + parts.append("%s (%s)" % (k.get("name") or "?", ", ".join(bits)) if bits + else str(k.get("name") or "?")) + return "; ".join(parts) + + +def _track_record_line(view: dict) -> str: + """`benched 3x since: 1 reproduced, 1 no-win, 1 could not run (last: failed)`, or a plain miss.""" + if not view.get("recalls"): + return "never benched by anyone since it was recorded" + parts = ["%d reproduced a win" % view["validations"] if view["validations"] else "", + "%d could not be run at all" % view["not_reproduced"] + if view.get("not_reproduced") else ""] + detail = ", ".join(p for p in parts if p) or "none reproduced a win" + hint = view.get("retire_hint") or "" + return "benched %dx since it was recorded — %s%s" % ( + view["recalls"], detail, " (**%s**)" % hint if hint else "") + + +def _repro_line(view: dict) -> str: + """Where the launch script and the kernel patches for this record actually are. + + Spelled as paths rather than as "see the bundle" because the reader is an agent that is about + to run something, and the next thing it does after this line is open a file. A record written + before `repro` existed says so plainly instead of pointing at a path that is not there. + """ + repro = view.get("repro") or {} + root = ((view.get("bundle") or {}).get("path") or "").rstrip("/") + where = (lambda name: "%s/files/%s" % (root, name) if root else name) + launch = str(repro.get("launch") or "") + if not launch: + return ("no launch script recorded (pre-`repro` record) — rebuild it from the config below") + bits = ["`%s`%s" % (where(launch), + " (SYNTHESIZED from the config, never executed as-is)" + if repro.get("launch_origin") == "synthesized" else " (captured from the run)")] + patches = [k.get("patch") for k in (repro.get("kernels") or []) if isinstance(k, dict) + and k.get("patch")] + if patches: + bits.append("kernel patches: " + ", ".join("`%s`" % where(p) for p in patches)) + missing = int(repro.get("kernels_without_patch") or 0) + if missing: + bits.append("%d accepted kernel(s) carry NO patch here — that part of the win cannot be " + "reproduced from this record" % missing) + return "; ".join(bits) + + +def _render_reference(refs_dir: str, cid: str, tier: str, views) -> str: + """Mirror the offer into prose the Director can read, or return "" and let the read stand.""" + try: + os.makedirs(refs_dir, exist_ok=True) + key = hashlib.sha1(("|".join(v["session_id"] for v in views)).encode()).hexdigest()[:7] + path = os.path.join(refs_dir, "e2e_reference_%s.md" % key) + lines = ["# e2e warm start — `%s`" % cid, "", + "Match tier `%s`, ordered by `%s` (highest first)." + % (tier, views[0]["ranked_by"]), ""] + if tier != "exact": + lines += ["> These were measured on a DIFFERENT workload point than the one requested. " + "Treat the configs as candidates and the numbers as non-comparable — do not " + "quote them as this deployment's throughput. The ordering above is by " + "absolute throughput, so it says which deployment was fastest AT ITS OWN " + "workload point, not which one would be fastest here.", ""] + for rank, v in enumerate(views, start=1): + lines += [ + "## %d. %s%s" % (rank, v["direction"] or "unlabeled", + " (champion)" if v["is_champion"] else ""), + "- throughput: %s tok/s (baseline %s), speedup %s" % ( + v["throughput_tok_s"], v["baseline_throughput_tok_s"], v["speedup"]), + "- workload: %s" % (json.dumps(v["workload"], sort_keys=True) or "{}"), + # Spelled out rather than reduced to a word: the Director's next decision is + # whether to spend a server launch on this, and "unverified, parity n/a" is a very + # different prompt from "validated, hot A/B, parity pass" even at the same speedup. + "- validation: %s (%s, basis %s, parity %s)" % ( + "VALIDATED" if v["validated"] else "unvalidated", + v["validation_status"] or "unrecorded", + v["validation_basis"] or "unverified", v["parity"] or "unrecorded"), + # The claim above is the writer's own; this line is everybody else's experience of + # it since. A record benched three times and never reproduced is a very different + # bet from an untried one at the same speedup, and only this line says which it is. + "- track record: %s" % _track_record_line(v), + "- accepted kernels: %s" % (_kernel_line(v["accepted_kernels"]) or "none"), + "- reproduce: %s" % _repro_line(v), + "- config:", "```json", + json.dumps(v["accepted_config"], indent=2, sort_keys=True), "```", "", + ] + with open(path, "w") as handle: + handle.write("\n".join(lines)) + return path + except OSError: + return "" + + +# -- write --------------------------------------------------------------------------------------- + + +def _record_state(a, result: dict) -> dict: + """How much this record's number should be believed, recorded AT WRITE TIME. + + A reader cannot recover this later. `validation_status` alone does not answer it — a record can + say `validated_win` and still be a number nobody checked for output parity, and the two are + indistinguishable once the run's logs are gone. So the judgement is made here, once, by the + process that still has the evidence, and stored as a first-class field: + + * `validated` — did this number clear a real gate on the box that produced it. False is a + perfectly good answer and does NOT mean the record is worthless; it means a reader should + re-measure before quoting it. A parity of `n/a` is false, not true: "we did not check" is + not "we checked and it was fine". + * `validation_basis` — WHICH gate. `hot_ab` is the Director's same-session interleaved A/B + (the strong one). `cold_gate` is a fresh-server before/after (weaker: it carries the drift + between two launches). `unverified` is a number that was reported but never independently + reproduced. These are not comparable, so a reader that mixes them must at least know it is. + * `parity` — pass | fail | n/a, verbatim. A faster server that answers differently is a + regression, and this is the only field that says whether anyone looked. + * `lifecycle` — `active` (believed) or `candidate` (recorded, unproven). The third value, + `retracted`, is written only by kb/retract.py and never by a fresh write. + * `retained` — the curation flag both lanes' readers filter on. True at birth; retraction + flips it. Written explicitly rather than left absent so `retained is False` stays a + three-state test (true / false / never stated) instead of degrading to a falsy check. + + Every field is overridable from the CLI because the caller sometimes knows better than the + result JSON — a backfill from an old run has evidence this function cannot see, and forcing it + to fake a `validation_status` to get the right state would corrupt the field that means + something else. + """ + status = str(result.get("validation_status") or "") + parity = str(getattr(a, "parity", "") or result.get("output_parity") or "").strip().lower() + basis = str(getattr(a, "validation_basis", "") or "").strip().lower() + if not basis or basis == "auto": + # A `validation_status` is only ever set by the Director's validate phase, and that phase is + # the same-session A/B by construction. No status means nothing re-measured this run. + basis = "hot_ab" if status else "unverified" + validated = str(getattr(a, "validated", "") or "auto").strip().lower() + if validated in ("", "auto"): + decided = status == "validated_win" and parity == "pass" + else: + decided = validated in ("1", "true", "yes", "on") + return {"validated": decided, "validation_basis": basis, "parity": parity or "n/a", + "lifecycle": "active" if decided else "candidate", "retained": True} + + +def build_record(a, result: dict, workdir=None) -> dict: + """One run's knowledge document, identical at every rung. + + Both ranking scalars sit flat at the top level because that is the only shape + `sessions/top?metric=` can read. Everything else lives under `value`, including the dimensions + that are already in the canonical id — a record that cannot say what it is once detached from + its address is not auditable. + + `workdir` is a scratch directory whose lifetime must outlast the store write: the synthesized + launch script is created in it and uploaded from it. Passing None asks for the document only, + which is what `retract` and `attest` want when they recompute the content digest — they must + not synthesize files, and must not refuse over a reproducibility rule that only governs new + writes. + """ + identity = identity_of(a) + final = finite_speedup(result.get("final_throughput_tok_s")) + baseline = finite_speedup(result.get("baseline_throughput_tok_s")) + speedup = finite_speedup(result.get("throughput_speedup")) + if speedup is None and final is not None and baseline: + speedup = round(final / baseline, 6) + if final is None: + raise SystemExit("result has no final_throughput_tok_s; refusing to record a run with no " + "measurement — an unranked record is invisible on every page") + kernels, kernel_files = _accepted_kernels(a, result) + state = _record_state(a, result) + value = { + "model": identity["model"], "gpu": identity["gpu"], + "framework": identity["framework"], "framework_version": identity["framework_version"], + "precision": identity["precision"], + # Ints, not the raw argv strings: this is the only copy of the shape once a record is read + # back off a coarse rung, and "1024" sorts and compares differently from 1024 in every + # consumer that touches it. A value that will not parse is dropped, matching counted(). + "workload": {k: int(str(getattr(a, k)).strip()) + for k in ("tp", "isl", "osl", "conc") + if str(getattr(a, k, "") or "").strip().lstrip("-").isdigit()}, + "baseline_throughput_tok_s": baseline, + "final_throughput_tok_s": final, + "direction": str(a.direction or result.get("direction") or ""), + "accepted_config": result.get("accepted_config") if isinstance( + result.get("accepted_config"), dict) else {}, + "accepted_kernels": kernels, + "validation_status": str(result.get("validation_status") or ""), + "upstream": result.get("upstream") if isinstance(result.get("upstream"), dict) else {}, + # GEAK's own comparability keys (schema v2): what basis the pair was measured on, which + # client took the number, which workload points were validated. A stored speedup is only + # meaningful against these, so they ride WITH the number rather than being rediscovered. + "comparability": result.get("comparability") if isinstance( + result.get("comparability"), dict) else {}, + "measured_by": str(a.measured_by or ""), + "recorded_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + } + value.update(state) + files = _artifact_files(a, result) + files.update(kernel_files) + value["artifacts"] = {k: v[0] for k, v in files.items()} + # After `artifacts` (it reads the captured launch/overlay names from there) and before the + # final rebuild (it may ADD the synthesized script and fetched patches to `files`). + value["repro"] = _repro(a, result, value, kernels, files, workdir) + if files: + value["artifacts"] = {k: v[0] for k, v in files.items()} + else: + value.pop("artifacts", None) + document = {"schema": SCHEMA, THROUGHPUT_METRIC: final, "value": value} + if speedup is not None: + document[SPEEDUP_METRIC] = speedup + # Keyed by STORED NAME, not by role: `value.artifacts` maps role -> stored name, and + # materialize() checks every one of those names against what the bundle actually holds. Keying + # the upload by role instead writes `files/patch` while the document promises `final.patch`, + # and the record fails its own integrity check on the way back out. + return {"knowledge": document, "files": {v[0]: v[1] for v in files.values()}, + "speedup": speedup, "throughput": final} + + +def _kernel_records(result: dict): + """Every kernel this run kept, from BOTH tracks, deduped by name. + + e2e_workflow.js banks head-op rewrites into `accepted_heads` and milestone rewrites into + `accepted_kernels`, and its own final view of a run is the union of the two. Reading only + `accepted_kernels` here dropped an entire track — both e2e records already in the remote KB came + from runs whose whole win lived in `accepted_heads`, so they recorded no kernels at all. Union, + not fallback: a run that used both tracks must lose neither. + + MERGE rather than first-wins on a name collision. One op optimized on both tracks is one kernel, + and the two entries know different things about it (a head carries target_callable, a milestone + carries source_path_in_sglang). Dropping the second would reintroduce exactly the silent field + loss this function exists to prevent. + + `short_name` is accepted as a name spelling because that is what the workflow pushed for years + before it also emitted `name`; a record written by an older lane still reads correctly here. + """ + merged, order = {}, [] + for source, raw_list in (("accepted_kernels", result.get("accepted_kernels")), + ("accepted_heads", result.get("accepted_heads"))): + for raw in (raw_list or []): + item = {"name": str(raw)} if not isinstance(raw, dict) else dict(raw) + name = str(item.get("name") or item.get("kernel_name") + or item.get("short_name") or "").strip() + if not name: + continue + item["name"] = name + if source == "accepted_heads": + item.setdefault("from_accepted_heads", True) + if name not in merged: + merged[name] = item + order.append(name) + continue + for key, value in item.items(): + if value in (None, "", 0): + continue + if merged[name].get(key) in (None, "", 0): + merged[name][key] = value + return [merged[name] for name in order] + + +def _accepted_kernels(a, result: dict): + """(entries, {role: (stored_name, local_path)}) for the kernels this run kept. + + A bare list of names was useless to a reader: it said a kernel mattered without saying how to + obtain it, so the next run had to rediscover the same rewrite. Each entry now carries BOTH ways + to get the patch, because each fails differently: + + * `kernel_canonical_id` addresses the kernel lane's own record — the source of truth, with + that kernel's full history, its own champion and its own measurements. Preferred. It can + miss: the kernel page is only there if the kernel lane wrote it, and it is keyed on the + ROCm version this run has to declare. + * `patch` is the diff itself, copied into THIS record under `kernels/.patch`. It costs + duplicated bytes and it goes stale, but it is the only copy that cannot 404 — an e2e record + whose kernel references have rotted still reproduces the configuration it is claiming. + + Accepts either shape from the workflow: plain strings (what FINALIZE_SCHEMA's accepted_kernels + emits today) or objects carrying language/patch/speedup. Strings degrade to a name and a + canonical id, never to an error, so this stays readable by the unmodified e2e_workflow.js. + + Both tracks are read — see _kernel_records for why reading one of them lost half of every run. + """ + entries, files = [], {} + gfx = kbid.segment(a.gfx, kbid.UNKNOWN) + rocm = str(getattr(a, "rocm_version", "") or + (result.get("upstream") or {}).get("rocm_version") or "") + for item in _kernel_records(result): + name = item["name"] + language = str(item.get("language") or item.get("backend") or "").strip() + # Start from what the producer sent and normalize over it, rather than copying a fixed set + # of keys into a fresh dict. A whitelist silently discards everything the workflow knows + # that this function has not been taught about — `kind`, `op_kind`, `e2e_delta_pct`, + # `routed_to` all vanished on the first real run — and the loss is invisible until someone + # reads the record back looking for a field they are sure they wrote. + entry = dict(item) + entry.update({"name": name, "language": language, + # `isolated` is the spelling the workflow used at every push site before + # bankAccepted also emitted `isolated_speedup`; without it an older + # result.json ranks as having no measurement at all. + "isolated_speedup": finite_speedup(item.get("isolated_speedup") + or item.get("speedup") + or item.get("isolated")), + "pct_gpu_time": finite_speedup(item.get("pct_gpu_time"))}) + entry.pop("patch_path", None) # replaced below by the STORED name, if it exists + # Only address the kernel lane when the language is known. Guessing it would mint an id + # that looks authoritative and resolves to nothing, which on a scheme with no search is + # indistinguishable from the kernel never having been optimized. + # + # An env/flag win is NOT a kernel rewrite — it routes the op to an existing implementation + # (`routed_to: aiter`), and the e2e workflow stores that routing target in the same `backend` + # field a real rewrite uses for its LANGUAGE. The kernel lane never wrote a page for it, so + # addressing one here fabricates a permanent dead reference on a store with no delete. + kind = str(item.get("winner_kind") or item.get("kind") or "").strip().lower() + if language and kind not in ("env", "flag"): + entry["kernel_canonical_id"] = kbid.kernel_canonical_ids( + kbid.kernel_identity(gfx, name, language, rocm))[0] + if item.get("session_id"): + entry["kernel_session_id"] = str(item["session_id"]) + patch = str(item.get("patch") or item.get("patch_path") or "") + if patch and os.path.isfile(patch): + stored = "kernels/%s.patch" % kbid.segment(name, "kernel") + files["kernel:" + name] = (stored, patch) + entry["patch"] = stored + entries.append(entry) + return entries, files + + +_ARTIFACT_KEYS = (("patch", "final_patch", "final.patch"), + ("launch", "final_launch_script", "launch.sh"), + ("report", "report_path", "report.md"), + ("overlay", "final_overlay", "overlay.py")) + + +def _artifact_files(a, result: dict) -> dict: + """{role: (stored_name, local_path)} for the run outputs that actually exist on disk. + + A path the result names but the filesystem does not have is dropped here rather than at upload + time, so `value.artifacts` never promises a file the record does not carry — the kernel lane's + materialize() now treats that promise as a hard error, and it should. + """ + found = {} + for role, field, stored in _ARTIFACT_KEYS: + path = str(result.get(field) or "") + if path and os.path.isfile(path): + found[role] = (stored, path) + for extra in (getattr(a, "file", None) or []): # retract recomputes a record but takes no --file + path = str(extra) + if os.path.isfile(path): + found["file:" + os.path.basename(path)] = (os.path.basename(path), path) + return found + + +def _env_pairs(env: str) -> dict: + """`"A=1 B=2"` -> `{"A": "1", "B": "2"}`, which is the shape a reader can act on. + + The workflow carries server env as one flat string because that is what it hands + `bench_e2e.sh`'s EXTRA_ENV, and a reader who wants to know whether a record set + VLLM_USE_AITER should not have to write a parser to find out. Anything that is not a + `KEY=VALUE` token is skipped rather than guessed at; the verbatim string is kept alongside + this dict, so nothing is lost by being strict here. + """ + pairs = {} + for token in str(env or "").split(): + key, sep, val = token.partition("=") + if sep and key and not key[0].isdigit(): + pairs[key] = val + return pairs + + +def _fetch_kernel_patch(root: str, canonical_id: str, name: str, into: str) -> str: + """Pull one kernel's patch out of the kernel lane's local store, or return "". + + The e2e record names the kernel lane's canonical id, but the bytes only ride along if the + workflow happened to still have the patch file on disk at write time — and by the time the + e2e run finalizes, the kernel workflow's scratch is usually gone. This walks over to the + kernel lane's own store and copies its champion's patch in, so the e2e record is complete + without the e2e run having had to hoard files it did not produce. + + Best-effort by construction: a kernel page that does not exist, a bundle that fails its + integrity check, an unreadable root — all of them mean "no patch here", which is a state the + record already knows how to describe (`kernels_without_patch`). Never raises. + """ + try: + from kb.store_local import LocalKBStore + store = LocalKBStore(root, metric=SPEEDUP_METRIC) + found = store.candidates(canonical_id, limit=1) + if not found: + return "" + bundle = store.materialize(canonical_id, found[0].session_id, into) + for candidate in ("patch.diff", "final.patch", "patch"): + path = os.path.join(bundle, "files", candidate) + if os.path.isfile(path): + return path + except Exception: + return "" + return "" + + +def _kernel_patches(a, kernels, files: dict, workdir: str) -> int: + """Fill in the patches the run did not carry, and return how many are STILL missing. + + Mutates `kernels` (setting `patch` on entries it manages to fetch) and `files` (adding the + bytes to upload) in place, because both are already the write path's accumulators and a + third copy would just be a chance for them to disagree. + """ + root = str(getattr(a, "kernel_store", "") or "") + missing = 0 + for entry in kernels: + if entry.get("patch"): + continue + cid = str(entry.get("kernel_canonical_id") or "") + name = str(entry.get("name") or "kernel") + fetched = "" + if root and cid and workdir: + fetched = _fetch_kernel_patch(root, cid, name, + os.path.join(workdir, "kernel_%s" % kbid.segment(name, + "k"))) + if fetched: + stored = "kernels/%s.patch" % kbid.segment(name, "kernel") + files["kernel:" + name] = (stored, fetched) + entry["patch"] = stored + entry["patch_origin"] = "kernel_store" + continue + # Said out loud in the record itself. A reader that sees three accepted kernels and two + # patches has no way to tell whether the third was a no-op or whether its bytes were + # simply lost, and those two readings lead to opposite decisions about whether the + # configuration below is worth trying. + entry["patch_missing"] = True + missing += 1 + return missing + + +def _launch_text(a, result: dict, value: dict, kernels, overlay: str) -> str: + """A runnable `launch.sh` built from the config, for a run that captured no script. + + Written against `e2e_workflow/scripts/bench_e2e.sh`'s env contract, because that script is + how this pipeline actually measured the number being recorded — reproducing the record means + re-entering it at the same door, not inventing a second launcher whose defaults nobody has + compared. Every knob it reads that we know is emitted; the ones we cannot know (MODEL's path, + HOST/PORT) are left as environment overrides with loud defaults. + + The header says SYNTHESIZED in as many words. A reader must be able to tell a script that was + executed and produced the number above from one that was reconstructed afterwards and has + never been run, and no amount of correctness in the body substitutes for saying which it is. + """ + identity = identity_of(a) + workload = value.get("workload") or {} + config = value.get("accepted_config") or {} + flags = str(config.get("flags") or "") + env = str(config.get("env") or "") + model_path = str(result.get("model_path") or (result.get("upstream") or {}).get("model_path") + or "") + lines = [ + "#!/usr/bin/env bash", + "# SYNTHESIZED by e2e_store.py from this record's accepted_config — NOT the script that", + "# produced the number below. It has never been executed as written. Read it before you", + "# run it: the paths it cannot know (MODEL, GEAK_SCRIPTS) are yours to supply.", + "#", + "# identity : %s" % " | ".join( + [identity["model"], identity["gpu"], identity["framework"], + identity["framework_version"], identity["precision"]]), + "# workload : %s" % (json.dumps(workload, sort_keys=True) or "{}"), + "# direction : %s" % (value.get("direction") or "unlabeled"), + "# measured : %s tok/s vs baseline %s" % (value.get("final_throughput_tok_s"), + value.get("baseline_throughput_tok_s")), + "# recorded_at : %s by %s" % (value.get("recorded_at") or "", + value.get("measured_by") or "unknown"), + "", + "set -euo pipefail", + "", + 'HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"', + '# Where GEAK\'s e2e_workflow/scripts lives on THIS box. bench_e2e.sh is the entry point', + '# the recorded number was measured through.', + ': "${GEAK_SCRIPTS:?set GEAK_SCRIPTS to GEAK/e2e_workflow/scripts}"', + ': "${MODEL:=%s}"' % (model_path or ""), + 'if [ -z "${MODEL}" ]; then', + # The caller's spelling, not the canonical segment: `qwen3-397b` is an address, and the + # reader has to type a path or an HF id, which is case-sensitive. + ' echo "set MODEL to the path or HF id of %s (bench_e2e.sh requires it)" >&2; exit 4' + % (str(getattr(a, "model", "") or "") or identity["model"]), + "fi", + "", + ] + if kernels: + patched = [k for k in kernels if k.get("patch")] + lines += ["# Kernel rewrites this configuration depends on. SRC must point at the source", + "# tree the patches were cut against (the serving stack's checkout)."] + if patched: + lines += ['SRC="${SRC:-}"', + 'if [ -n "${SRC}" ]; then', + ' for p in %s; do' % " ".join('"$HERE/%s"' % k["patch"] for k in patched), + ' git -C "$SRC" apply --check "$p" && git -C "$SRC" apply "$p"', + " done", + "else", + ' echo "SRC unset: skipping %d kernel patch(es) — the recorded speedup will ' + 'NOT reproduce without them" >&2' % len(patched), + "fi"] + absent = [k for k in kernels if not k.get("patch")] + if absent: + lines += ["# NO PATCH IN THIS RECORD for: %s" % ", ".join( + str(k.get("name") or "?") for k in absent), + "# Fetch them from the kernel lane (kernel_canonical_id in value.accepted_kernels)", + "# or the number below will not reproduce."] + lines.append("") + if overlay: + lines += ['# Python-level overlay this run served with.', + 'OVERLAY_PYTHONPATH="${OVERLAY_PYTHONPATH:-$HERE/%s}"' % overlay, ""] + pairs = _env_pairs(env) + if pairs: + lines.append("# Server environment, as recorded.") + lines += ["export %s=%s" % (k, _sh_quote(v)) for k, v in sorted(pairs.items())] + lines.append("") + lines += ["exec env \\"] + for key, val in (("BACKEND", identity["framework"]), + ("MODEL", "${MODEL}"), + ("TP", workload.get("tp")), ("ISL", workload.get("isl")), + ("OSL", workload.get("osl")), ("CONC", workload.get("conc")), + ("GPU", "${GPU:-0}"), ("OUT_DIR", "${OUT_DIR:-$PWD/repro_out}")): + if val in (None, "", kbid.UNKNOWN): + continue + lines.append(" %s=%s \\" % (key, _sh_quote(str(val)))) + if flags: + lines.append(" EXTRA_SERVER_ARGS=%s \\" % _sh_quote(flags)) + if env: + lines.append(" EXTRA_ENV=%s \\" % _sh_quote(env)) + if overlay: + lines.append(' OVERLAY_PYTHONPATH="${OVERLAY_PYTHONPATH}" \\') + lines += [' bash "${GEAK_SCRIPTS}/bench_e2e.sh"', ""] + return "\n".join(lines) + + +def _sh_quote(text: str) -> str: + """Single-quote for /bin/sh, leaving `${VAR}` references alone. + + Not shlex.quote: several of the values above are deliberately shell expansions the reader is + meant to be able to override from the environment, and quoting those into literals would turn + a script you can point at your own model into one that tries to open a file called `${MODEL}`. + """ + text = str(text) + if text.startswith("${") and text.endswith("}"): + return '"%s"' % text + return "'%s'" % text.replace("'", "'\"'\"'") + + +def _repro(a, result: dict, value: dict, kernels, files: dict, workdir) -> dict: + """`value.repro`: everything needed to run this configuration again, said twice. + + Once as a script (`launch`) for a reader that wants to run it, and once as structured fields + for a reader that wants to reason about it without parsing shell. Both, not either: the + script is the only artifact that is complete, and the fields are the only form a curation + pass or another agent can compare across records. + + `workdir` is None when the caller only wants the document's shape — `retract`/`attest` + recompute a record purely to re-derive its content digest, and neither should be synthesizing + files or refusing to run because a record it is taking BACK was never reproducible. + """ + artifacts = value.get("artifacts") or {} + captured = str(artifacts.get("launch") or "") + overlay = str(artifacts.get("overlay") or "") + config = value.get("accepted_config") or {} + flags, env = str(config.get("flags") or ""), str(config.get("env") or "") + missing = _kernel_patches(a, kernels, files, workdir) if workdir else sum( + 1 for k in kernels if not k.get("patch")) + if workdir: + # `value.artifacts` is rebuilt by the caller from `files` after this returns, so adding to + # `files` here is enough to make the synthesized script a first-class artifact of the + # record rather than a loose file nobody indexed. + have_patch = any(k.get("patch") for k in kernels) + if not captured and not flags and not env and not have_patch and not overlay: + raise SystemExit( + "result carries no launch script, no accepted_config flags or env, no kernel " + "patch and no overlay; refusing to record a run nobody can reproduce — a record " + "that only says a number was once achieved cannot be acted on, and this store has " + "no delete to take it back with") + if not captured: + path = os.path.join(workdir, "launch.sh") + with open(path, "w") as handle: + handle.write(_launch_text(a, result, value, kernels, overlay)) + os.chmod(path, 0o755) + files["launch"] = ("launch.sh", path) + captured, origin = "launch.sh", "synthesized" + else: + origin = "captured" + else: + origin = "captured" if captured else "" + return { + "launch": captured, + "launch_origin": origin, + "entry_point": "e2e_workflow/scripts/bench_e2e.sh", + "server_args": flags, + "env": env, + "env_pairs": _env_pairs(env), + "model": str(result.get("model_path") + or (result.get("upstream") or {}).get("model_path") or ""), + "backend": identity_of(a)["framework"], + "workload": dict(value.get("workload") or {}), + "overlay": overlay, + "kernels": [{"name": k.get("name") or "", "patch": k.get("patch") or "", + "kernel_canonical_id": k.get("kernel_canonical_id") or ""} + for k in kernels], + "kernels_without_patch": missing, + # The one field a reader can branch on: is what follows enough to re-run, or is it a lead. + "complete": bool(captured) and not missing, + } + + +def cmd_write(a) -> dict: + workdir = tempfile.mkdtemp(prefix="e2e_store_write_") + try: + return _write(a, workdir) + finally: + # Only after every rung has been published: the synthesized launch script and any patches + # fetched from the kernel lane live here, and both planes read them at write time. + shutil.rmtree(workdir, ignore_errors=True) + + +def _write(a, workdir: str) -> dict: + try: + with open(a.result, "r", errors="replace") as handle: + result = json.load(handle) + except (OSError, ValueError) as e: + raise SystemExit("cannot read --result %s: %s" % (a.result, e)) + if not isinstance(result, dict): + raise SystemExit("--result must be a JSON object") + + record = build_record(a, result, workdir) + ladder = ladder_of(a) + sid = kbid.session_id(ladder[0][0], identity_of(a)["model"], + _content_digest(record["knowledge"])) + out = {"applied": bool(a.apply), "session_id": sid, "speedup": record["speedup"], + "throughput_tok_s": record["throughput"], + "files": sorted(record["files"]), "rungs": []} + # A rung ranks on its own metric (throughput on the exact rung, speedup on the coarser ones), so + # each opens its own per-metric store; `publish` writes that one rung, all-or-none. All-or-none + # runs at THIS loop level too: a rung we cannot open or write stops the ladder before a partial + # one is published, and because the exact rung is written first, what lands is never a coarse + # page that outranks the specific one it was meant to summarize. + for cid, tier, metric, floor in ladder: + rung = {"canonical_id": cid, "tier": tier, "metric": metric, "written": False, + "promoted": False, "error": ""} + if not a.apply: + out["rungs"].append(rung) + continue + store, mirror, why = open_plane(a, metric, floor, create=True) + if store is None: + rung["error"] = why + out["rungs"].append(rung) + break + # A re-bench of the same configuration lands on the SAME session id by design (see + # _content_digest) and replaces the document wholesale. Without this, every re-measurement + # would silently reset that record's validation history to zero — and it would look + # perfectly healthy afterwards, because the new document is well-formed and simply says + # nobody has ever tried this. + knowledge = _carrying_ledger(store, cid, sid, record["knowledge"]) + rec = {"canonical_id": cid, "session_id": sid, "knowledge": knowledge} + score_of = lambda r, m=metric: r["knowledge"].get(m) + written, promoted, err = publish(store, [rec], record["files"], score_of) + rung["written"] = bool(written) + rung["promoted"] = bool(promoted) + rung["error"] = err or why # `both` with an unreachable service: recorded, not fatal + if mirror is not None: + # The mirror never gates the primary; its own failure is reported, not raised. It gets + # its own ledger lookup because the two planes drift: a record attested on the remote + # and re-written from a box that only ever wrote locally must not have the remote's + # count overwritten by the local one's. + mrec = dict(rec, knowledge=_carrying_ledger(mirror, cid, sid, record["knowledge"])) + _mw, _mp, merr = publish(mirror, [mrec], record["files"], score_of) + if merr and not rung["error"]: + rung["error"] = merr + out["rungs"].append(rung) + if err: + break + out["ok"] = all(r["written"] for r in out["rungs"]) if a.apply else True + return out + + +def _carrying_ledger(store, cid: str, sid: str, knowledge: dict) -> dict: + """`knowledge` with any attestation ledger the store already holds for this session moved in. + + A store that cannot answer is treated as a store that holds nothing: failing the write over a + lookup would turn a transient service blip into a lost measurement, and the worst case of + guessing wrong here is a reset counter, not a wrong number. + """ + try: + previous = store.get_session(cid, sid) + except Exception: + previous = None + fresh = dict(knowledge) + fresh["value"] = carry_attestations( + previous.get("value") if isinstance(previous, dict) else None, dict(fresh["value"])) + return fresh + + +def _as_number(value): + """A measurement off the command line, or None. Argparse hands these over as strings, and the + caller is usually a shell line built by the workflow, so a stray unit or an empty flag must + drop the evidence rather than abort an attestation that is otherwise perfectly recordable.""" + if value is None or isinstance(value, bool): + return None + try: + return finite_speedup(float(value)) + except (TypeError, ValueError): + return None + + +def cmd_attest(a) -> dict: + """Count one attempt to actually RUN a stored record, at every rung it was written to. + + This is the other half of the read path. A resolve offers records; something downstream takes + one to a box and finds out whether it still holds. Until now that finding-out evaporated, so + the tenth reader of a record that has failed nine times saw exactly what the first reader saw. + + Applied to the whole ladder for the same reason retraction is: all three rungs share one + session id, and counting only on the exact rung leaves the two coarse pages — the ones a + reader on a DIFFERENT workload reads, which is most readers — quoting a stale ledger. + + Deliberately does not touch the ranking scalars or the champion. See kb/attest.py: one failure + on one box is evidence, not a verdict, and burying a record on the strength of it would make + this command too dangerous to run automatically, which would mean it never ran at all. + """ + session_id = str(getattr(a, "session_id", "") or "").strip() + if not session_id: + raise SystemExit("attest needs --session-id (the id the write printed); there is nothing " + "to recompute it from, because the outcome being recorded is not in any " + "result JSON") + evidence = {k: v for k, v in ( + ("measured_tok_s", _as_number(getattr(a, "measured_tok_s", None))), + ("baseline_tok_s", _as_number(getattr(a, "baseline_tok_s", None))), + ("parity", str(getattr(a, "parity", "") or "").strip()), + ("note", str(getattr(a, "note", "") or "").strip()), + ("workload", {k: str(getattr(a, k) or "") for k in ("tp", "isl", "osl", "conc") + if getattr(a, k, None)}), + ) if v not in (None, "", {})} + if evidence.get("measured_tok_s") and evidence.get("baseline_tok_s"): + evidence["delta_pct"] = round( + (evidence["measured_tok_s"] / evidence["baseline_tok_s"] - 1.0) * 100.0, 3) + out = {"applied": bool(a.apply), "session_id": session_id, "outcome": a.outcome, "rungs": []} + for cid, tier, metric, floor in ladder_of(a): + store, mirror, why = open_plane(a, metric, floor) + planes = [p for p in (store, mirror) if p is not None] + if not planes: + out["rungs"].append({"canonical_id": cid, "tier": tier, "error": why, "found": False}) + continue + for plane in planes: + report = attest_session(plane, cid, session_id, a.outcome, + actor=str(a.measured_by or ""), evidence=evidence, + apply=bool(a.apply)) + report.update({"tier": tier, "plane_note": why}) + out["rungs"].append(report) + out["ok"] = attestation_ok(out["rungs"], a.apply) + # Hoisted out of the per-rung reports because every rung holds the identical document, and a + # caller deciding whether to open a curation ticket should not have to notice that. + hints = [r.get("retire_hint") for r in out["rungs"] if r.get("retire_hint")] + out["retire_hint"] = hints[0] if hints else "" + return out + + +def cmd_retract(a) -> dict: + """Take back one already-written record, at every rung it was written to. + + The session id is the SAME at all three rungs — `cmd_write` computes it once from the content + digest and reuses it — so one identity plus one session id addresses the whole ladder, which is + what makes a retraction possible at all without having kept a record of where things landed. + + Two ways to name the session, and the second exists because the first usually is not available: + `--session-id` when the write's output was kept, or `--result` to recompute the digest from the + same JSON the write was fed. The recompute is exact — the digest keys on config, kernel names, + workload and direction, none of which a re-read changes — but it does require the SAME + `--direction`, which is easy to forget and would silently address a session that does not exist. + That case reports `found: false` per rung rather than inventing one. + """ + session_id = str(getattr(a, "session_id", "") or "").strip() + if not session_id: + if not getattr(a, "result", ""): + raise SystemExit("retract needs --session-id, or --result to recompute it") + try: + with open(a.result, "r", errors="replace") as handle: + result = json.load(handle) + except (OSError, ValueError) as e: + raise SystemExit("cannot read --result %s: %s" % (a.result, e)) + record = build_record(a, result) + session_id = kbid.session_id(ladder_of(a)[0][0], identity_of(a)["model"], + _content_digest(record["knowledge"])) + out = {"applied": bool(a.apply), "session_id": session_id, "reason": a.reason, "rungs": []} + for cid, tier, metric, floor in ladder_of(a): + store, mirror, why = open_plane(a, metric, floor) + planes = [p for p in (store, mirror) if p is not None] + if not planes: + out["rungs"].append({"canonical_id": cid, "tier": tier, "error": why, "found": False}) + continue + for plane in planes: + # Both scalars are zeroed on every rung, not just the one this rung ranks on. The + # document is identical at all three rungs by construction, so a rewrite that zeroed + # only the local metric would leave the record still ranked on the other two. + report = retract_session(plane, cid, session_id, a.reason, metric, + extra_metrics=(THROUGHPUT_METRIC, SPEEDUP_METRIC), + actor=str(a.measured_by or ""), scan=int(a.scan), + apply=bool(a.apply)) + report.update({"tier": tier, "metric": metric, "plane_note": why}) + out["rungs"].append(report) + out["ok"] = retraction_ok(out["rungs"], a.apply) + return out + + +def _content_digest(knowledge: dict) -> str: + """Dedup key: the CONFIG, not the measurement. + + Re-benchmarking one config must land on the same session id so `mode="replace"` updates that + record instead of accumulating a page full of near-identical entries that all outrank each + other by noise. So the throughput numbers and the timestamp are deliberately excluded — two + runs of the same config ARE the same candidate, and the later one wins. + """ + value = knowledge.get("value") or {} + payload = json.dumps({"config": value.get("accepted_config") or {}, + "kernels": sorted(str(k.get("name") or "") for k in + (value.get("accepted_kernels") or []) + if isinstance(k, dict)), + "workload": value.get("workload") or {}, + "direction": value.get("direction") or ""}, + sort_keys=True, ensure_ascii=False) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +# -- cli ----------------------------------------------------------------------------------------- + + +def _identity_args(p): + p.add_argument("--model", required=True, help="model name, e.g. Qwen3-397B") + p.add_argument("--gfx", default="", help="gfx target, e.g. gfx950") + p.add_argument("--framework", default="", help="serving stack: vllm | sglang") + p.add_argument("--framework-version", default="", help="serving stack version, e.g. 0.26.0") + p.add_argument("--precision", default="", help="e.g. mxfp8, fp8, bf16") + p.add_argument("--rocm-version", default="", + help="ROCm of the container, used to address the accepted kernels' own records") + p.add_argument("--tp", default=None, help="tensor parallel degree") + p.add_argument("--isl", default=None, help="input sequence length") + p.add_argument("--osl", default=None, help="output sequence length") + p.add_argument("--conc", default=None, help="concurrency") + + +def _state_args(p): + """Overrides for what _record_state would otherwise derive. Shared with `retract` because it + recomputes the content digest, and the digest is computed off a full build_record().""" + p.add_argument("--validated", default="auto", choices=("auto", "true", "false"), + help="auto = validated_win AND parity pass; override when you know better") + p.add_argument("--validation-basis", default="auto", + choices=("auto", "hot_ab", "cold_gate", "unverified"), + help="which gate produced the number (auto: hot_ab if a status was recorded)") + p.add_argument("--parity", default="", help="pass | fail | n/a; defaults to result.output_parity") + + +def _plane_args(p): + p.add_argument("--plane", choices=("local", "remote", "both"), default="local") + p.add_argument("--store", default="", help="on-disk store root (plane local|both)") + p.add_argument("--scan", type=int, default=DEFAULT_SCAN, + help="candidates hydrated per rung before curation") + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + sub = p.add_subparsers(dest="command", required=True) + + q = sub.add_parser("identity", help="print the ladder this deployment reads and writes") + _identity_args(q) + + q = sub.add_parser("resolve", help="offer prior runs for this deployment") + _identity_args(q) + _plane_args(q) + q.add_argument("--top-n", type=int, default=DEFAULT_TOP_N) + q.add_argument("--min-speedup", type=float, default=0.0) + q.add_argument("--sort-by", choices=tuple(SORT_METRICS), default=DEFAULT_SORT_BY, + help="how to order the offer on EVERY rung (default: absolute throughput, " + "high to low). The champion metric per rung is unaffected.") + q.add_argument("--refs-dir", default="", help="write prose references here") + q.add_argument("--cache-dir", default="", help="materialize artifact bundles here") + + q = sub.add_parser("write", help="record one run at every rung") + _identity_args(q) + _plane_args(q) + q.add_argument("--result", required=True, help="JSON from the workflow's report/validate step") + q.add_argument("--direction", default="", help="what this run DID, for the shortlist collapse") + q.add_argument("--measured-by", default="", help="who/what produced the number") + q.add_argument("--file", action="append", default=[], help="extra artifact to attach") + q.add_argument("--kernel-store", default="", + help="kernel lane's on-disk store root; when given, patches this run no longer " + "has on disk are fetched from there by kernel_canonical_id so the record " + "stays reproducible. Off by default: it is real I/O on the write path.") + _state_args(q) + q.add_argument("--apply", action="store_true", help="actually write; default is a dry run") + + q = sub.add_parser("attest", help="count one attempt to RUN a stored record: validated | " + "failed | not_reproduced. Moves no score, no champion.") + _identity_args(q) + _plane_args(q) + q.add_argument("--session-id", required=True, help="the session that was tried") + q.add_argument("--outcome", required=True, choices=OUTCOMES, + help="validated = reproduced a win; failed = ran but did not win; " + "not_reproduced = could not be made to run at all") + q.add_argument("--measured-tok-s", default=None, help="what it did here, for the history entry") + q.add_argument("--baseline-tok-s", default=None, help="what this box does without it") + q.add_argument("--parity", default="", help="pass | fail | n/a on this box") + q.add_argument("--note", default="", help="one line a future reader can act on") + q.add_argument("--measured-by", default="", help="who tried it") + q.add_argument("--apply", action="store_true", help="actually record it; default is a dry run") + + q = sub.add_parser("retract", help="take back a written record (rewrite, since there is no " + "delete): retained=false, scores zeroed, champion re-pointed") + _identity_args(q) + _plane_args(q) + q.add_argument("--session-id", default="", help="the session to retract, from the write output") + q.add_argument("--result", default="", help="recompute the session id from the SAME JSON and " + "--direction the write was given") + q.add_argument("--direction", default="", help="must match the write, or the id will not match") + q.add_argument("--reason", required=True, + help="why this record is wrong; it is all a future reader has to judge by") + q.add_argument("--measured-by", default="", help="who is retracting it") + _state_args(q) + q.add_argument("--apply", action="store_true", help="actually rewrite; default is a dry run") + + a = p.parse_args(argv) + if a.command == "identity": + result = {"identity": identity_of(a), + "ladder": [{"canonical_id": c, "tier": t, "ranked_by": m, "promote_floor": f} + for c, t, m, f in ladder_of(a)]} + elif a.command == "resolve": + result = cmd_resolve(a) + elif a.command == "retract": + result = cmd_retract(a) + elif a.command == "attest": + result = cmd_attest(a) + else: + result = cmd_write(a) + print(json.dumps(result, ensure_ascii=False, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/e2e_workflow/scripts/tests/test_e2e_store.py b/e2e_workflow/scripts/tests/test_e2e_store.py new file mode 100644 index 000000000..e51c627d3 --- /dev/null +++ b/e2e_workflow/scripts/tests/test_e2e_store.py @@ -0,0 +1,542 @@ +#!/usr/bin/env python3 +"""The e2e KB CLI's read/write surface, beyond the retraction round-trips in test_kb_retract. + +Retraction is covered next door; this file drives the paths that one does not touch — the reference +renderer and artifact materializer a `resolve` emits, the kernel/head merge and artifact plumbing a +`write` performs, the two commands' failure modes, and the `identity` echo. Everything runs +`e2e_store.main([...])` in-process (not over a subprocess) so the assertions AND the coverage both +land on the module under test. +""" + +import json +import os +import sys + +import pytest + +_SCRIPTS = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, os.path.dirname(os.path.dirname(_SCRIPTS))) # repo root, for `kb` +sys.path.insert(0, _SCRIPTS) +import e2e_store # noqa: E402 + +IDENTITY = ["--model", "M", "--gfx", "gfx950", "--framework", "vllm", + "--framework-version", "0.26.0", "--precision", "fp8", + "--tp", "8", "--isl", "1024", "--osl", "1024", "--conc", "64"] + + +def _run(command, *args, identity=IDENTITY): + import io + import contextlib + buffer = io.StringIO() + with contextlib.redirect_stdout(buffer): + e2e_store.main([command] + identity + list(args)) + return json.loads(buffer.getvalue()) + + +def _result(path, tput=1000.0, baseline=800.0, **extra): + doc = {"final_throughput_tok_s": tput, "baseline_throughput_tok_s": baseline, + "validation_status": "validated_win", "output_parity": "pass", + "accepted_config": {"env": "A=1"}, "accepted_kernels": []} + doc.update(extra) + path.write_text(json.dumps(doc)) + return str(path) + + +def _write(tmp_path, name, direction, *args, tput=1000.0, **extra): + result = _result(tmp_path / (name + ".json"), tput=tput, **extra) + return _run("write", "--store", str(tmp_path / "store"), "--result", result, + "--direction", direction, "--apply", *args) + + +# -- identity ----------------------------------------------------------------------------------- + + +def test_identity_prints_the_full_ladder(): + out = _run("identity") + assert out["identity"]["model"] == "m" # identity segments are normalized + # Every rung the deployment reads/writes, most specific first, each with the metric it ranks on. + assert out["ladder"][0]["tier"] == "exact" + assert out["ladder"][0]["ranked_by"] == "throughput_tok_s" + assert any(r["ranked_by"] == "speedup" for r in out["ladder"]) + + +# -- resolve: reference prose + artifact bundles ------------------------------------------------ + + +def test_resolve_writes_a_prose_reference(tmp_path): + _write(tmp_path, "a", "tuned") + out = _run("resolve", "--store", str(tmp_path / "store"), + "--refs-dir", str(tmp_path / "refs")) + assert out["read_reason"] == "read" + refs = list((tmp_path / "refs").glob("e2e_reference_*.md")) + assert len(refs) == 1 + text = refs[0].read_text() + assert "# e2e warm start" in text and "throughput" in text and "config:" in text + + +def test_reference_spells_out_each_accepted_kernel(tmp_path): + """The prose kernel line is what the Director reads before opening a patch — it must name the + language, the isolated speedup and the stored path.""" + (tmp_path / "moe.patch").write_text("--- a\n+++ b\n") + _write(tmp_path, "a", "kernelized", + accepted_kernels=[{"name": "moe_stage1", "language": "triton", + "isolated_speedup": 1.84, "patch": str(tmp_path / "moe.patch")}]) + _run("resolve", "--store", str(tmp_path / "store"), "--refs-dir", str(tmp_path / "refs")) + text = list((tmp_path / "refs").glob("*.md"))[0].read_text() + assert "moe_stage1" in text and "triton" in text and "1.84x" in text + assert "kernels/moe_stage1.patch" in text + + +def test_resolve_materializes_artifact_bundles(tmp_path): + (tmp_path / "final.patch").write_text("--- a\n+++ b\n") + _write(tmp_path, "a", "tuned", final_patch=str(tmp_path / "final.patch")) + out = _run("resolve", "--store", str(tmp_path / "store"), + "--cache-dir", str(tmp_path / "cache")) + bundle = out["candidates"][0]["bundle"] + assert "final.patch" in bundle["files"] + assert os.path.isdir(bundle["path"]) + + +def test_a_coarser_rung_match_is_flagged_non_comparable(tmp_path): + """A record written at conc=64 also lands on the workload-agnostic rung; a resolve at conc=128 + misses the exact rung and matches there instead, and the prose must warn the numbers are not + this deployment's.""" + _write(tmp_path, "a", "tuned") + other_conc = [x if x != "64" else "128" for x in IDENTITY] + out = _run("resolve", "--store", str(tmp_path / "store"), + "--refs-dir", str(tmp_path / "refs"), identity=other_conc) + if out["read_reason"] == "read" and out["match_tier"] != "exact": + text = list((tmp_path / "refs").glob("*.md"))[0].read_text() + assert "DIFFERENT workload" in text + + +def test_min_speedup_floors_the_offer(tmp_path): + # baseline 800, tput 1000 => speedup 1.25; a 1.5 floor curates it away. + _write(tmp_path, "a", "tuned", tput=1000.0) + out = _run("resolve", "--store", str(tmp_path / "store"), "--min-speedup", "1.5") + assert out["candidates"] == [] and out["read_reason"] == "e2e_page_not_found" + + +def test_resolve_on_a_missing_store_is_a_miss_not_a_crash(tmp_path): + out = _run("resolve", "--store", str(tmp_path / "nope")) + assert out["candidates"] == [] and out["read_reason"].startswith("no_such_store") + + +def test_a_read_that_raises_is_reported_not_propagated(tmp_path, monkeypatch): + _write(tmp_path, "a", "tuned") + + def boom(self, *a, **k): + raise RuntimeError("disk gone") + + monkeypatch.setattr("kb.store_local.LocalKBStore.candidates", boom) + out = _run("resolve", "--store", str(tmp_path / "store")) + assert out["read_reason"].startswith("read_failed") + + +# -- write: kernel/head merge, artifacts, state, failure modes ---------------------------------- + + +def test_write_merges_accepted_kernels_and_heads(tmp_path): + """Both tracks are read and merged on a name collision — the head entry fills fields the + milestone entry left blank, and neither record is dropped.""" + (tmp_path / "k.patch").write_text("x") + out = _write( + tmp_path, "a", "both-tracks", + accepted_kernels=[{"name": "op1", "language": "triton", "session_id": "kern-sess", + "patch": str(tmp_path / "k.patch"), "isolated_speedup": 1.5}, + {"language": "triton"}], # no name: dropped, not recorded + accepted_heads=[{"name": "op1", "target_callable": "fwd", "language": ""}, # "" won't clobber + {"name": "op2", "language": "triton"}]) + assert out["applied"] is True and out["rungs"][0]["written"] is True + view = _run("resolve", "--store", str(tmp_path / "store"))["candidates"][0] + kernels = {k["name"]: k for k in view["accepted_kernels"]} + assert set(kernels) == {"op1", "op2"} + assert kernels["op1"]["target_callable"] == "fwd" # merged in from the head track + assert kernels["op1"]["patch"] == "kernels/op1.patch" # stored name, from the milestone + assert kernels["op1"].get("kernel_canonical_id") # addressed: language known, real op + assert kernels["op1"]["kernel_session_id"] == "kern-sess" + + +def test_an_env_win_is_not_addressed_as_a_kernel(tmp_path): + """A routed env/flag win is not a rewrite — minting a kernel id for it fabricates a dead + reference on a store with no delete.""" + _write(tmp_path, "a", "routed", + accepted_kernels=[{"name": "moe", "language": "aiter", "winner_kind": "env"}]) + view = _run("resolve", "--store", str(tmp_path / "store"))["candidates"][0] + assert view["accepted_kernels"][0].get("kernel_canonical_id") in (None, "") + + +def test_extra_files_are_attached(tmp_path): + (tmp_path / "notes.txt").write_text("hi") + out = _write(tmp_path, "a", "tuned", "--file", str(tmp_path / "notes.txt")) + assert "notes.txt" in out["files"] + + +def test_a_validated_override_is_honored(tmp_path): + # status/parity would say unvalidated; --validated true overrides to active. + _write(tmp_path, "a", "tuned", "--validated", "true", + validation_status="", output_parity="fail") + view = _run("resolve", "--store", str(tmp_path / "store"))["candidates"][0] + assert view["validated"] is True and view["lifecycle"] == "active" + + +def test_a_dry_run_records_nothing(tmp_path): + result = _result(tmp_path / "r.json") + out = _run("write", "--store", str(tmp_path / "store"), "--result", result, + "--direction", "d") # no --apply + assert out["applied"] is False + assert all(not r["written"] for r in out["rungs"]) + assert not (tmp_path / "store").exists() + + +def test_write_refusing_a_run_with_no_measurement(tmp_path): + (tmp_path / "r.json").write_text(json.dumps({"baseline_throughput_tok_s": 800.0})) + with pytest.raises(SystemExit): + _run("write", "--store", str(tmp_path / "store"), "--result", + str(tmp_path / "r.json"), "--direction", "d", "--apply") + + +def test_write_with_an_unreadable_result(tmp_path): + with pytest.raises(SystemExit): + _run("write", "--store", str(tmp_path / "store"), + "--result", str(tmp_path / "missing.json"), "--direction", "d", "--apply") + + +def test_write_rejects_a_non_object_result(tmp_path): + (tmp_path / "r.json").write_text("[]") + with pytest.raises(SystemExit): + _run("write", "--store", str(tmp_path / "store"), + "--result", str(tmp_path / "r.json"), "--direction", "d", "--apply") + + +def test_a_store_that_cannot_be_opened_stops_the_ladder(tmp_path): + """--store points at a regular file: makedirs fails, the rung cannot open, and the write stops + before publishing a partial ladder.""" + blocker = tmp_path / "blocker" + blocker.write_text("i am a file, not a dir") + result = _result(tmp_path / "r.json") + out = _run("write", "--store", str(blocker), "--result", result, + "--direction", "d", "--apply") + assert out["ok"] is False + assert out["rungs"][0]["written"] is False and out["rungs"][0]["error"] + assert len(out["rungs"]) == 1 # broke, did not try coarser rungs + + +def test_a_mid_ladder_write_failure_stops_and_mirror_is_best_effort(tmp_path, monkeypatch): + """publish erroring on the exact rung stops the ladder; when a mirror is present its own failure + is reported, never raised.""" + real = e2e_store.publish + calls = {"n": 0} + + def flaky(store, recs, files, score_of): + calls["n"] += 1 + return [], [], "boom" # every publish fails + + # Force a two-plane open so the mirror branch is exercised, then fail the write. + local_dir = str(tmp_path / "store") + os.makedirs(local_dir, exist_ok=True) + from kb.store_local import LocalKBStore + prim = LocalKBStore(local_dir, metric="throughput_tok_s", promote_floor=0.0) + mirr = LocalKBStore(str(tmp_path / "mirror"), metric="throughput_tok_s", promote_floor=0.0) + monkeypatch.setattr(e2e_store, "open_plane", lambda a, m, f, create=False: (prim, mirr, "")) + monkeypatch.setattr(e2e_store, "publish", flaky) + result = _result(tmp_path / "r.json") + out = _run("write", "--store", local_dir, "--result", result, + "--direction", "d", "--apply", "--plane", "both") + assert out["ok"] is False and out["rungs"][0]["error"] == "boom" + assert len(out["rungs"]) == 1 # stopped after the exact rung + assert calls["n"] >= 2 # primary AND mirror both attempted + assert e2e_store.publish is flaky and real is not flaky # sanity on the monkeypatch + + +def test_a_materialize_failure_degrades_the_offer(tmp_path, monkeypatch): + """A download problem reports an error on the bundle rather than dropping the candidate — the + config lives in the knowledge doc, so the offer is still usable.""" + (tmp_path / "final.patch").write_text("x") + _write(tmp_path, "a", "tuned", final_patch=str(tmp_path / "final.patch")) + from kb.store_local import KBStoreError + + def boom(self, *a, **k): + raise KBStoreError("cache full") + + monkeypatch.setattr("kb.store_local.LocalKBStore.materialize", boom) + out = _run("resolve", "--store", str(tmp_path / "store"), + "--cache-dir", str(tmp_path / "cache")) + assert out["candidates"][0]["bundle"]["error"] + + +def test_an_unwritable_refs_dir_lets_the_read_stand(tmp_path): + """The prose page is a mirror of the offer, not the offer itself: if refs-dir cannot be made, + resolve still returns the candidate.""" + blocker = tmp_path / "blocker" + blocker.write_text("i am a file") + _write(tmp_path, "a", "tuned") + out = _run("resolve", "--store", str(tmp_path / "store"), + "--refs-dir", str(blocker / "under-a-file")) + assert out["read_reason"] == "read" and out["candidates"] + + +def test_a_mirror_only_failure_is_reported_without_gating_the_primary(tmp_path, monkeypatch): + """The mirror never gates the primary: the exact rung writes locally and its mirror failure is + surfaced on that rung, but the ladder keeps going.""" + from kb.store_local import LocalKBStore + prim = LocalKBStore(str(tmp_path / "store"), metric="throughput_tok_s", promote_floor=0.0) + + class FailingMirror: + def write(self, *a, **k): + raise RuntimeError("mirror down") + + seen = {"n": 0} + + def one_bad_mirror(a, metric, floor, create=False): + seen["n"] += 1 + # A mirror only on the first (exact) rung, so the ladder still advances past it. + return (prim, FailingMirror(), "") if seen["n"] == 1 else (prim, None, "") + + monkeypatch.setattr(e2e_store, "open_plane", one_bad_mirror) + result = _result(tmp_path / "r.json") + out = _run("write", "--store", str(tmp_path / "store"), "--result", result, + "--direction", "d", "--apply", "--plane", "both") + assert out["rungs"][0]["written"] is True # primary landed + assert "mirror down" in out["rungs"][0]["error"] # mirror failure surfaced + assert len(out["rungs"]) > 1 # ladder was NOT stopped by it + + +# -- retract: recompute-from-result and missing-store paths ------------------------------------- + + +def test_retract_recomputes_the_session_from_the_result(tmp_path): + """Given the same --result and --direction, retract recomputes the exact session id the write + minted, without the write's output having been kept.""" + result = _result(tmp_path / "r.json") + written = _run("write", "--store", str(tmp_path / "store"), "--result", result, + "--direction", "d", "--apply")["session_id"] + out = _run("retract", "--store", str(tmp_path / "store"), "--result", result, + "--direction", "d", "--reason", "wrong", "--apply") + assert out["session_id"] == written and out["ok"] is True + + +def test_retract_needs_a_session_or_a_result(tmp_path): + with pytest.raises(SystemExit): + _run("retract", "--store", str(tmp_path / "store"), "--reason", "wrong") + + +def test_retract_with_an_unreadable_result(tmp_path): + with pytest.raises(SystemExit): + _run("retract", "--store", str(tmp_path / "store"), "--result", + str(tmp_path / "missing.json"), "--direction", "d", "--reason", "wrong") + + +def test_retract_on_a_missing_store_reports_not_found(tmp_path): + out = _run("retract", "--store", str(tmp_path / "nope"), + "--session-id", "whatever", "--reason", "wrong", "--apply") + assert out["ok"] is False + assert all(r["found"] is False for r in out["rungs"]) + + +# -- resolve: the offer is ordered by absolute throughput on EVERY rung ------------------------- + + +def _coarse_id(out): + """The tp_any rung's canonical id — the coarse page a reader on another workload lands on.""" + return [r["canonical_id"] for r in out["ladder"] if r["tier"] == "tp_any"][0] + + +def test_a_coarse_rung_is_offered_by_throughput_not_by_its_champion_metric(tmp_path): + """The behaviour the whole read path was changed for. A coarse rung still CROWNS on speedup — + that is what makes its champion comparable across workload points — but the OFFER is ordered by + absolute throughput, because a reader asking "what has run fast on this deployment" is not + asking "what improved most over whatever baseline it happened to have".""" + store = str(tmp_path / "store") + # 1.9x off a slow baseline vs 1.05x off a fast one: opposite orders under the two metrics. + _write(tmp_path, "slowbase", "big-ratio", tput=900.0, baseline=474.0) + _write(tmp_path, "fastbase", "small-ratio", tput=1600.0, baseline=1524.0) + out = _run("resolve", "--store", store, "--tp", "16") # a workload only the coarse rungs hold + assert out["match_tier"] != "exact" + assert out["sorted_by"] == "throughput_tok_s" + assert out["champion_metric"] == "speedup" # the rung's own metric, unchanged + assert [c["direction"] for c in out["candidates"]] == ["small-ratio", "big-ratio"] + + +def test_sort_by_speedup_restores_the_old_order(tmp_path): + _write(tmp_path, "slowbase", "big-ratio", tput=900.0, baseline=474.0) + _write(tmp_path, "fastbase", "small-ratio", tput=1600.0, baseline=1524.0) + out = _run("resolve", "--store", str(tmp_path / "store"), "--tp", "16", + "--sort-by", "speedup") + assert out["sorted_by"] == "speedup" + assert [c["direction"] for c in out["candidates"]] == ["big-ratio", "small-ratio"] + + +def test_the_prose_reference_names_the_metric_it_ordered_by(tmp_path): + _write(tmp_path, "a", "tuned") + _run("resolve", "--store", str(tmp_path / "store"), "--refs-dir", str(tmp_path / "refs")) + text = list((tmp_path / "refs").glob("e2e_reference_*.md"))[0].read_text() + assert "ordered by `throughput_tok_s` (highest first)" in text + + +# -- attest: counting what happened when a record was actually run ------------------------------ + + +def test_attest_counts_on_every_rung_without_moving_anything(tmp_path): + """All three rungs share one session id, so a count that lands only on the exact rung leaves + the two coarse pages — the ones a reader on another workload reads — quoting a stale ledger.""" + store = str(tmp_path / "store") + sid = _write(tmp_path, "a", "tuned")["session_id"] + before = _run("resolve", "--store", store)["candidates"][0] + out = _run("attest", "--store", store, "--session-id", sid, "--outcome", "validated", + "--measured-tok-s", "1100", "--baseline-tok-s", "1000", "--parity", "pass", + "--measured-by", "boxA", "--apply") + assert out["ok"] is True + assert len(out["rungs"]) == 3 and all(r["rewritten"] for r in out["rungs"]) + assert all(r["attestations"]["validations"] == 1 for r in out["rungs"]) + after = _run("resolve", "--store", store)["candidates"][0] + assert after["validations"] == 1 and after["recalls"] == 1 + assert after["throughput_tok_s"] == before["throughput_tok_s"] # no scalar moved + assert after["speedup"] == before["speedup"] + assert after["is_champion"] == before["is_champion"] # no champion re-pointed + # and the evidence rode along, including the delta this command derives rather than trusting + assert out["rungs"][0]["attestations"]["history"][-1]["delta_pct"] == 10.0 + + +def test_repeated_failures_raise_a_retire_hint_but_retire_nothing(tmp_path): + store = str(tmp_path / "store") + sid = _write(tmp_path, "a", "tuned")["session_id"] + for _ in range(2): + _run("attest", "--store", store, "--session-id", sid, + "--outcome", "not_reproduced", "--apply") + view = _run("resolve", "--store", store)["candidates"][0] + assert view["not_reproduced"] == 2 and "could not reproduce" in view["retire_hint"] + assert view["lifecycle"] == "active" # advisory only; still offered, still ranked + assert "**" in list_reference_text(tmp_path, store) # and the prose says so in bold + + +def list_reference_text(tmp_path, store): + _run("resolve", "--store", store, "--refs-dir", str(tmp_path / "refs2")) + return list((tmp_path / "refs2").glob("e2e_reference_*.md"))[0].read_text() + + +def test_re_writing_the_same_config_keeps_its_attestations(tmp_path): + """The bug this carry-forward exists for: _content_digest excludes the measurement, so a + re-bench of one config replaces the SAME session — and a naive write would silently reset the + record's whole validation history while looking perfectly well-formed.""" + store = str(tmp_path / "store") + sid = _write(tmp_path, "a", "tuned", tput=1000.0)["session_id"] + _run("attest", "--store", store, "--session-id", sid, "--outcome", "validated", "--apply") + again = _write(tmp_path, "a", "tuned", tput=1200.0) # same config, new measurement + assert again["session_id"] == sid + view = _run("resolve", "--store", store)["candidates"][0] + assert view["throughput_tok_s"] == 1200.0 # the number DID update + assert view["validations"] == 1 # the ledger did not reset + + +def test_attest_needs_a_session_id(tmp_path): + with pytest.raises(SystemExit): + e2e_store.main(["attest"] + IDENTITY + ["--store", str(tmp_path / "store"), + "--outcome", "validated", "--session-id", ""]) + + +def test_attest_on_a_missing_store_reports_not_found(tmp_path): + out = _run("attest", "--store", str(tmp_path / "nope"), "--session-id", "whatever", + "--outcome", "failed", "--apply") + assert out["ok"] is False and all(r["found"] is False for r in out["rungs"]) + + +# -- repro: every record must carry something you can actually run ------------------------------ + + +def test_a_launch_script_is_synthesized_when_the_run_captured_none(tmp_path): + """The common case: the workflow banked flags and env but no script. Rather than storing a + record nobody can act on, the write builds one against bench_e2e.sh's env contract and says + plainly that it was synthesized.""" + out = _write(tmp_path, "a", "tuned", + accepted_config={"flags": "--max-num-seqs 256", "env": "VLLM_USE_AITER=1"}) + assert "launch.sh" in out["files"] + view = _run("resolve", "--store", str(tmp_path / "store"), + "--cache-dir", str(tmp_path / "cache"))["candidates"][0] + assert view["repro"]["launch"] == "launch.sh" + assert view["repro"]["launch_origin"] == "synthesized" + assert view["repro"]["env_pairs"] == {"VLLM_USE_AITER": "1"} + assert view["repro"]["server_args"] == "--max-num-seqs 256" + text = (tmp_path / "cache" / view["session_id"] / "files" / "launch.sh").read_text() + assert "bench_e2e.sh" in text and "SYNTHESIZED" in text + assert "EXTRA_SERVER_ARGS='--max-num-seqs 256'" in text + assert "export VLLM_USE_AITER='1'" in text + assert "TP='8'" in text and "ISL='1024'" in text # the workload it was measured at + + +def test_a_captured_launch_script_is_stored_verbatim_and_says_so(tmp_path): + (tmp_path / "run.sh").write_text("#!/bin/sh\necho the real thing\n") + _write(tmp_path, "a", "tuned", final_launch_script=str(tmp_path / "run.sh")) + view = _run("resolve", "--store", str(tmp_path / "store"), + "--cache-dir", str(tmp_path / "cache"))["candidates"][0] + assert view["repro"]["launch_origin"] == "captured" + assert (tmp_path / "cache" / view["session_id"] / "files" / "launch.sh").read_text() \ + == "#!/bin/sh\necho the real thing\n" + + +def test_a_kernel_without_its_patch_is_counted_not_hidden(tmp_path): + """A reader seeing three kernels and two patches cannot otherwise tell whether the third was a + no-op or whether its bytes were simply lost, and those read in opposite directions.""" + (tmp_path / "k.patch").write_text("--- a\n+++ b\n") + _write(tmp_path, "a", "tuned", + accepted_kernels=[{"name": "op1", "language": "triton", "patch": str(tmp_path / "k.patch")}, + {"name": "op2", "language": "triton"}]) + view = _run("resolve", "--store", str(tmp_path / "store"))["candidates"][0] + assert view["repro"]["kernels_without_patch"] == 1 + assert view["repro"]["complete"] is False + assert [k["patch"] for k in view["repro"]["kernels"]] == ["kernels/op1.patch", ""] + + +def test_a_kernel_patch_can_be_fetched_from_the_kernel_lanes_own_store(tmp_path): + """--kernel-store closes the usual gap: by the time an e2e run finalizes, the kernel + workflow's scratch is gone, so the bytes only survive where the kernel lane filed them.""" + from kb.store_local import LocalKBStore + kroot = str(tmp_path / "kernelkb") + kstore = LocalKBStore(kroot, metric="speedup") + (tmp_path / "k.patch").write_text("--- a\n+++ b\n") + kid = "geak:kernel:gfx950:op1:triton:rocm:7.2" + kstore.write(kid, "ksess", {"schema": "geak.kernel.v1", "speedup": 1.4, "value": {}}, + {"patch.diff": str(tmp_path / "k.patch")}) + kstore.promote(kid, "ksess", 1.4) + _write(tmp_path, "a", "tuned", "--kernel-store", kroot, "--rocm-version", "7.2", + accepted_kernels=[{"name": "op1", "language": "triton"}]) + view = _run("resolve", "--store", str(tmp_path / "store"), + "--cache-dir", str(tmp_path / "cache"))["candidates"][0] + assert view["repro"]["kernels_without_patch"] == 0 and view["repro"]["complete"] is True + assert (tmp_path / "cache" / view["session_id"] / "files" / "kernels" / "op1.patch") \ + .read_text() == "--- a\n+++ b\n" + + +def test_write_refusing_a_record_nobody_could_reproduce(tmp_path): + """A number with no config, no script, no patch and no overlay is not knowledge — and this + store has no delete to take it back with.""" + (tmp_path / "bare.json").write_text(json.dumps( + {"final_throughput_tok_s": 1000.0, "baseline_throughput_tok_s": 800.0})) + with pytest.raises(SystemExit): + _run("write", "--store", str(tmp_path / "store"), "--result", str(tmp_path / "bare.json"), + "--direction", "d", "--apply") + + +def test_a_kernel_patch_alone_is_enough_to_be_reproducible(tmp_path): + """The gate is "is there anything to act on", not "is there a config" — a run whose whole win + was a kernel rewrite carries the patch and nothing else, and it is perfectly actionable.""" + (tmp_path / "k.patch").write_text("--- a\n+++ b\n") + out = _write(tmp_path, "a", "kernels", accepted_config={}, + accepted_kernels=[{"name": "op1", "language": "triton", + "patch": str(tmp_path / "k.patch")}]) + assert out["rungs"][0]["written"] is True + + +def test_the_prose_reference_points_at_the_launch_script_and_the_patches(tmp_path): + """The Director reads this and then opens a file, so the paths are spelled out — not "see the + bundle", which sends it back to the store to ask a question the page already answered.""" + (tmp_path / "k.patch").write_text("--- a\n+++ b\n") + _write(tmp_path, "a", "tuned", + accepted_kernels=[{"name": "op1", "language": "triton", + "patch": str(tmp_path / "k.patch")}]) + _run("resolve", "--store", str(tmp_path / "store"), "--refs-dir", str(tmp_path / "refs"), + "--cache-dir", str(tmp_path / "cache")) + text = list((tmp_path / "refs").glob("e2e_reference_*.md"))[0].read_text() + assert "- reproduce: " in text and "files/launch.sh" in text + assert "files/kernels/op1.patch" in text + assert "- track record: never benched by anyone since it was recorded" in text diff --git a/e2e_workflow/scripts/tests/test_kb_retract.py b/e2e_workflow/scripts/tests/test_kb_retract.py new file mode 100644 index 000000000..559ae0480 --- /dev/null +++ b/e2e_workflow/scripts/tests/test_kb_retract.py @@ -0,0 +1,202 @@ +#!/usr/bin/env python3 +"""Retraction: taking back a record on a store that has no delete. + +Every test here is about a way retraction can LOOK done and not be. Marking the document while +leaving the ranking scalar, zeroing the scalar while leaving the champion pointing at it, filtering +on read while the writer never set the flag — each of those passes a casual inspection and leaves +the retracted record still steering the next run. +""" + +import json +import os +import sys + +_SCRIPTS = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, os.path.dirname(os.path.dirname(_SCRIPTS))) # repo root, for `kb` +sys.path.insert(0, _SCRIPTS) +import e2e_store # noqa: E402 +from kb.retract import is_retired, retracted_document, retraction_ok # noqa: E402 +from kb.store_local import KBStoreError, LocalKBStore # noqa: E402 + +CID = "geak:e2e:m:gfx950:vllm:0.26.0:fp8:tp_8:isl_1024:osl_1024:conc_64" +IDENTITY = ["--model", "M", "--gfx", "gfx950", "--framework", "vllm", + "--framework-version", "0.26.0", "--precision", "fp8", + "--tp", "8", "--isl", "1024", "--osl", "1024", "--conc", "64"] + + +def _result(path, tput, parity="pass", env="A=1", status="validated_win"): + path.write_text(json.dumps({ + "final_throughput_tok_s": tput, "baseline_throughput_tok_s": 800.0, + "validation_status": status, "output_parity": parity, + "accepted_config": {"env": env}, "accepted_kernels": []})) + return str(path) + + +def _write(tmp_path, name, tput, direction, **kw): + result = _result(tmp_path / (name + ".json"), tput, **kw) + out = _run("write", "--store", str(tmp_path / "store"), "--result", result, + "--direction", direction, "--apply") + return out["session_id"] + + +def _run(command, *args): + import io + import contextlib + buffer = io.StringIO() + with contextlib.redirect_stdout(buffer): + e2e_store.main([command] + IDENTITY + ["--plane", "local"] + list(args)) + return json.loads(buffer.getvalue()) + + +# -- the write-time state fields ------------------------------------------------------------- + + +def test_parity_n_a_is_recorded_as_unvalidated(tmp_path): + """"We did not check" must not read as "we checked and it was fine".""" + _write(tmp_path, "b", 900.0, "unchecked", parity="n/a") + view = _run("resolve", "--store", str(tmp_path / "store"))["candidates"][0] + assert view["validated"] is False + assert view["lifecycle"] == "candidate" + assert view["parity"] == "n/a" + + +def test_a_win_with_parity_is_validated_and_says_which_gate(tmp_path): + _write(tmp_path, "a", 1000.0, "tuned") + view = _run("resolve", "--store", str(tmp_path / "store"))["candidates"][0] + assert view["validated"] is True + assert view["validation_basis"] == "hot_ab" + + +def test_state_does_not_enter_the_dedup_digest(tmp_path): + """Re-recording one config with a corrected verdict must REPLACE it, not mint a second entry. + + The digest keys on the configuration, so the same config written twice is the same candidate. + If `validated` leaked into it, every correction would leave the old verdict live alongside the + new one — on a store with no delete, permanently. + """ + first = _write(tmp_path, "a", 1000.0, "tuned", parity="pass") + second = _write(tmp_path, "a2", 1000.0, "tuned", parity="fail") + assert first == second + assert len(_run("resolve", "--store", str(tmp_path / "store"))["candidates"]) == 1 + + +# -- the rewrite ------------------------------------------------------------------------------ + + +def test_a_retraction_needs_a_reason(): + try: + retracted_document({"speedup": 2.0, "value": {}}, " ", ["speedup"]) + except KBStoreError as e: + assert "reason" in str(e) + else: + raise AssertionError("an unexplained retraction was accepted") + + +def test_the_ranking_scalar_is_zeroed_not_just_flagged(): + """A flag alone is inert: `sessions/top?metric=` reads the scalar, not the flag.""" + document = retracted_document({"speedup": 2.0, "throughput_tok_s": 900.0, "value": {}}, + "wrong bench key", ["speedup", "throughput_tok_s"]) + assert document["speedup"] == 0.0 and document["throughput_tok_s"] == 0.0 + assert is_retired(document["value"]) + # ...and the withdrawn claim survives, or the tombstone cannot be reviewed later. + assert document["value"]["withdrawn_scores"] == {"speedup": 2.0, "throughput_tok_s": 900.0} + + +def test_retracting_the_champion_re_points_it_at_the_survivor(tmp_path): + """Zeroing a document does not un-promote it. The champion is a separate object, and one left + pointing at a retracted record keeps its old score as the bar every future candidate must + clear — a single false record quietly closes the page.""" + champion = _write(tmp_path, "a", 1000.0, "tuned") + survivor = _write(tmp_path, "b", 900.0, "other") + out = _run("retract", "--store", str(tmp_path / "store"), "--session-id", champion, + "--reason", "could not be reproduced", "--apply") + assert out["ok"] is True + store = LocalKBStore(str(tmp_path / "store"), metric="throughput_tok_s", promote_floor=0.0) + assert store.champion(CID)["session_id"] == survivor + + +def test_the_last_record_on_a_page_self_zeroes_its_champion(tmp_path): + session = _write(tmp_path, "a", 1000.0, "tuned") + _run("retract", "--store", str(tmp_path / "store"), "--session-id", session, + "--reason", "sole record, and wrong", "--apply") + store = LocalKBStore(str(tmp_path / "store"), metric="throughput_tok_s", promote_floor=0.0) + champion = store.champion(CID) + # Still pointing somewhere — there is nothing else to point at — but no longer advertising a + # win, and no longer acting as a floor a future candidate has to beat. + assert champion["session_id"] == session and champion["value"] == 0.0 + + +def test_artifacts_survive_the_rewrite(tmp_path): + """The bytes are the evidence. A record has just been called false; that is exactly when + someone wants to read the patch it shipped.""" + (tmp_path / "final.patch").write_text("--- a\n+++ b\n") + (tmp_path / "r.json").write_text(json.dumps({ + "final_throughput_tok_s": 1000.0, "baseline_throughput_tok_s": 800.0, + "validation_status": "validated_win", "output_parity": "pass", + "final_patch": str(tmp_path / "final.patch"), + "accepted_config": {"env": "A=1"}, "accepted_kernels": []})) + session = _run("write", "--store", str(tmp_path / "store"), "--result", + str(tmp_path / "r.json"), "--direction", "d", "--apply")["session_id"] + before = sum(1 for _r, _d, f in os.walk(str(tmp_path / "store")) for _ in f) + _run("retract", "--store", str(tmp_path / "store"), "--session-id", session, + "--reason", "wrong", "--apply") + assert sum(1 for _r, _d, f in os.walk(str(tmp_path / "store")) for _ in f) == before + + +# -- the read ---------------------------------------------------------------------------------- + + +def test_a_retracted_record_is_not_offered(tmp_path): + champion = _write(tmp_path, "a", 1000.0, "tuned") + _write(tmp_path, "b", 900.0, "other") + _run("retract", "--store", str(tmp_path / "store"), "--session-id", champion, + "--reason", "wrong", "--apply") + out = _run("resolve", "--store", str(tmp_path / "store")) + assert [c["direction"] for c in out["candidates"]] == ["other"] + assert out["curation"]["retired"] == 1 + + +def test_retracted_is_dropped_before_the_direction_collapse(tmp_path): + """Order matters. The collapse keeps the best record PER DIRECTION; if a retracted entry is + still in the list when it runs, it wins its direction and evicts the good alternative behind + it — so one false record hides a real one instead of merely removing itself.""" + bad = _write(tmp_path, "a", 1000.0, "tuned") + _write(tmp_path, "b", 950.0, "tuned", env="B=1") + _run("retract", "--store", str(tmp_path / "store"), "--session-id", bad, + "--reason", "wrong", "--apply") + out = _run("resolve", "--store", str(tmp_path / "store")) + assert [c["throughput_tok_s"] for c in out["candidates"]] == [950.0] + + +def test_an_all_retracted_page_reports_why_it_looks_empty(tmp_path): + """"Nobody recorded this" and "everything recorded here was taken back" produce the same + read_reason and mean opposite things about whether to try again.""" + session = _write(tmp_path, "a", 1000.0, "tuned") + _run("retract", "--store", str(tmp_path / "store"), "--session-id", session, + "--reason", "wrong", "--apply") + out = _run("resolve", "--store", str(tmp_path / "store")) + assert out["candidates"] == [] and out["read_reason"] == "e2e_page_not_found" + assert out["curation"]["retired"] == 1 + + +# -- reporting --------------------------------------------------------------------------------- + + +def test_a_page_that_never_held_the_record_is_not_a_failed_retraction(): + assert retraction_ok([{"found": True, "rewritten": True}, + {"found": False, "rewritten": False}], True) is True + + +def test_finding_the_record_nowhere_is_a_failure(): + """Otherwise a typo'd session id reports success while the real record stays live.""" + assert retraction_ok([{"found": False, "rewritten": False}], True) is False + + +def test_a_dry_run_writes_nothing_and_shows_the_document(tmp_path): + session = _write(tmp_path, "a", 1000.0, "tuned") + out = _run("retract", "--store", str(tmp_path / "store"), "--session-id", session, + "--reason", "wrong") + assert out["applied"] is False + assert all(not r["rewritten"] for r in out["rungs"]) + assert out["rungs"][0]["would_write"]["value"]["retired_reason"] == "wrong" + assert _run("resolve", "--store", str(tmp_path / "store"))["candidates"] != [] diff --git a/interface/run_e2e.py b/interface/run_e2e.py index d2b459fa9..20df94c06 100644 --- a/interface/run_e2e.py +++ b/interface/run_e2e.py @@ -2777,6 +2777,12 @@ def _recover_workflow_return(exp_root: Path) -> dict | None: win = _recover_best_intermediate_win(eval_dir) if win is not None: return win + # 1b. no accepted kernel, but the sweep adopted a winning serving config + # (warm-start recall / ck_tune) — recover it as the final floor so a + # validated config gain is not flattened to the default baseline. + cfg_win = _recover_accepted_config_win(eval_dir) + if cfg_win is not None: + return cfg_win return _recover_completed_no_gain(eval_dir) serving = validation.get("serving_config") or {} accepted_config = { @@ -3011,6 +3017,119 @@ def _integrate_candidates(eval_dir: Path) -> list[dict]: return records +def _recover_accepted_serving_config(eval_dir: Path) -> dict | None: + """Recover the SERVING CONFIG the sweep accepted, from ``config/sweep_results.json``. + + The config sweep (warm-start KB recall + ck_tune) writes its winning serving + config here — ``accepted_flags`` / ``accepted_env`` / ``best_throughput_tok_s`` + measured against ``baseline.throughput_tok_s_median`` — and an adopted config + stays live in CURRENT_FLAGS/CURRENT_ENV for the rest of the run (the whole + pipeline compounds ON TOP of it). Neither :func:`_recover_best_intermediate_win` + (reads ``overlay//integrate_result.json``) nor + :func:`_recover_completed_no_gain` (reads ``baseline/`` only) looks here, so a + run that crashed after adopting a config used to be recovered as if the config + never happened — the default baseline was reported and a validated recall gain + (e.g. a +65% warm-start config) was silently discarded. + + Returns ``None`` when there is no sweep file, no measured numbers, or the + accepted config does not actually differ from / beat the baseline (beyond the + sweep's own noise band). Otherwise a compact dict the recovery tiers fold in. + """ + sweep = _read_json(eval_dir / "config" / "sweep_results.json") + if not sweep: + return None + base = sweep.get("baseline") or {} + try: + best_tput = float(sweep.get("best_throughput_tok_s")) + except (TypeError, ValueError): + return None + if best_tput <= 0.0: + return None + try: + sweep_speedup = float(sweep.get("throughput_speedup_vs_baseline")) + except (TypeError, ValueError): + sweep_speedup = 0.0 + # The sweep's own baseline median is the denominator when present, but several + # runs record it as null in this block (only best + speedup survive). Derive it + # from the sweep's own speedup, then fall back to the run's measured default + # baseline — otherwise a real config win (mixtral +6.97%, qwen27b +17.56%) is + # missed just because this one field was null. + baseline_tput = 0.0 + for v in (base.get("throughput_tok_s_median"), + base.get("output_throughput_tok_s_median")): + try: + baseline_tput = float(v) + except (TypeError, ValueError): + continue + if baseline_tput > 0.0: + break + if baseline_tput <= 0.0 and sweep_speedup > 1.0: + baseline_tput = best_tput / sweep_speedup + if baseline_tput <= 0.0: + official = _read_json(eval_dir / "baseline" / "baseline_official.json") + summary = _read_json(eval_dir / "baseline" / "bench_summary.json") + for v in (official.get("baseline_throughput_tok_s"), + official.get("plateau_median_tok_s"), + summary.get("output_throughput_tok_s_median"), + summary.get("throughput_tok_s_median")): + try: + baseline_tput = float(v) + except (TypeError, ValueError): + continue + if baseline_tput > 0.0: + break + if baseline_tput <= 0.0: + return None + flags = str(sweep.get("accepted_flags") or "").strip() + env = str(sweep.get("accepted_env") or "").strip() + base_flags = str(base.get("flags") or "").strip() + base_env = str(base.get("env") or "").strip() + # A sweep that kept nothing writes the baseline config back as "accepted" — that + # is a genuine no-gain, not a config win. Require a real config change AND a real + # gain past the sweep's own acceptance band before crediting it. + changed = (flags and flags != base_flags) or (env and env != base_env) + band = float(sweep.get("noise_band_pct") or 0.5) + if not changed or best_tput <= baseline_tput * (1.0 + band / 100.0): + return None + speedup = sweep_speedup if sweep_speedup > 1.0 else best_tput / baseline_tput + return { + "flags": flags, "env": env, "baseline_tput": baseline_tput, + "best_tput": best_tput, "speedup": speedup, "band_pct": band, + } + + +def _recover_accepted_config_win(eval_dir: Path) -> dict | None: + """Config-only recovery tier: the sweep adopted a winning serving config but no + kernel/head integrate A/B was accepted before the crash. + + Sits between :func:`_recover_best_intermediate_win` (a kernel win, which folds + the config in itself) and :func:`_recover_completed_no_gain` (nothing accepted + at all). Reports the adopted config as the final floor — the served path really + was running it when the run died — so the validated recall/sweep gain survives a + mid-run crash instead of being flattened to the default baseline. + """ + cfg = _recover_accepted_serving_config(eval_dir) + if cfg is None: + return None + return { + "eval_dir": str(eval_dir), + "throughput_speedup": cfg["speedup"], + "baseline_throughput_tok_s": cfg["baseline_tput"], + "final_throughput_tok_s": cfg["best_tput"], + "output_parity": "n/a", + "validation_status": "recovered_intermediate", + # Config-only win: applied through env/flags, so there is no overlay bundle. + "final_overlay": "", + "final_launch_script": "", + "accepted_config": {"flags": cfg["flags"], "env": cfg["env"]}, + "accepted_kernels": [], + "accepted_heads": [], + "recovered_from_disk": True, + "recovered_intermediate": True, + "recovered_config_only": True, + } + + def _recover_best_intermediate_win(eval_dir: Path) -> dict | None: """Salvage the best accepted intermediate win when the run died BEFORE Validate. @@ -3115,11 +3234,40 @@ def _recover_best_intermediate_win(eval_dir: Path) -> dict | None: ): if value and value not in sink: sink.append(value) + # Fold in the sweep-adopted serving config (config/sweep_results.json). It was + # live on the server this kernel A/B ran against, so (a) its flags/env belong in + # accepted_config for a reproducible relaunch, and (b) when the kernel's own + # reference leg was measured ON that config (ref_med at/above the config-applied + # throughput), the honest speedup DENOMINATOR is the DEFAULT baseline, not the + # config-applied ref — otherwise the config's own gain is dropped from the + # headline and the stack (config ⊕ kernel) is under-credited (the DeepSeek / + # gpt-oss recovered_intermediate case). Re-base only when the ref leg is truly at + # the config-applied level, so a kernel A/B run against the raw baseline is never + # double-counted. + baseline_out, speedup_out, config_restacked = ref_med, speedup, False + cfg = _recover_accepted_serving_config(eval_dir) + if cfg is not None: + for value, sink in ((cfg["flags"], flags), (cfg["env"], env)): + if value and value not in sink: + sink.append(value) + # The kernel A/B ran on the config-applied server when its reference leg + # sits on the config-applied SIDE of the gap between the default baseline + # and the config-applied throughput. The midpoint test is robust to ~1% + # measurement drift (the DeepSeek case: ref leg 1481.9 vs config-applied + # 1497) while still refusing to re-base a kernel A/B that ran against the + # RAW baseline (which lands near cfg baseline, far below the midpoint) — + # re-basing that would double-count the config gain. + midpoint = (cfg["baseline_tput"] + cfg["best_tput"]) / 2.0 + if ref_med >= midpoint and cfg["baseline_tput"] < ref_med: + baseline_out = cfg["baseline_tput"] + speedup_out = final_tput / cfg["baseline_tput"] + config_restacked = True return { "eval_dir": str(eval_dir), - "throughput_speedup": speedup, - "baseline_throughput_tok_s": ref_med, + "throughput_speedup": speedup_out, + "baseline_throughput_tok_s": baseline_out, "final_throughput_tok_s": final_tput, + "config_restacked_over_default": config_restacked or None, "output_parity": ir.get("output_parity"), "validation_status": "recovered_intermediate", # Latency from the candidate (accepted) A/B leg when the integrator recorded @@ -3178,6 +3326,11 @@ def _recover_completed_no_gain(eval_dir: Path) -> dict | None: construction (do-no-harm); speedup 1.0 -> :func:`normalize_result` => no_gain. Returns ``None`` only when no baseline throughput was ever measured (the run genuinely produced nothing to keep). + + This is the LAST recovery tier: :func:`_recover_workflow_return` reaches it only + after ruling out an accepted kernel win AND an adopted serving config + (:func:`_recover_accepted_config_win`), so "nothing accepted" is really true + here — the default baseline is the correct floor. """ official = _read_json(eval_dir / "baseline" / "baseline_official.json") summary = _read_json(eval_dir / "baseline" / "bench_summary.json") diff --git a/interface/test_run_e2e_recovery.py b/interface/test_run_e2e_recovery.py index 1e9e6562d..77edbb1de 100644 --- a/interface/test_run_e2e_recovery.py +++ b/interface/test_run_e2e_recovery.py @@ -635,6 +635,92 @@ def test_result_source_director_validation(tmp_path): assert out["result_source"] == "disk_director_validation" +# ── adopted serving-config recovery (config/sweep_results.json) ────────────── +# A run that crashed AFTER the sweep adopted a warm-start / ck_tune config used to +# be recovered from baseline/ only, silently discarding the validated config gain +# (the Qwen3-14B +65.35% case: 5214.3 tok/s adopted -> reported 3153.5, 1.0x). + +def _write_sweep(eval_dir: Path, *, baseline_tput: float, best_tput: float, + speedup: float, flags: str = "--max-model-len 6144", + env: str = "VLLM_ROCM_USE_AITER=1", + base_flags: str = "--max-model-len 6144", base_env: str = "") -> None: + (eval_dir / "config").mkdir(parents=True, exist_ok=True) + (eval_dir / "config" / "sweep_results.json").write_text(json.dumps({ + "phase": "sweep", "backend": "vllm", "noise_band_pct": 0.5, + "baseline": {"throughput_tok_s_median": baseline_tput, + "flags": base_flags, "env": base_env}, + "accepted_flags": flags, "accepted_env": env, + "best_throughput_tok_s": best_tput, + "throughput_speedup_vs_baseline": speedup, + }), encoding="utf-8") + + +def test_recover_config_only_win_from_sweep(tmp_path): + """No accepted kernel, but the sweep adopted a config that beats baseline: the + adopted config is the final floor, NOT the default baseline (speedup > 1).""" + eval_dir = _make_no_gain_eval_dir(tmp_path) # baseline + a REJECTED kernel only + _write_sweep(eval_dir, baseline_tput=3153.524, best_tput=5214.345, speedup=1.6535) + wf = rx._recover_workflow_return(eval_dir.parent) + assert wf["recovered_config_only"] is True + assert wf["recovered_intermediate"] is True + assert not wf.get("recovered_no_gain") + assert wf["baseline_throughput_tok_s"] == pytest.approx(3153.524) + assert wf["final_throughput_tok_s"] == pytest.approx(5214.345) + assert wf["throughput_speedup"] == pytest.approx(1.6535) + assert "VLLM_ROCM_USE_AITER=1" in wf["accepted_config"]["env"] + out = rx.normalize_result(_handoff(eval_dir), wf) + assert out["status"] == "ok", "an adopted config gain must never read as no_gain" + assert out["result_source"] == "disk_intermediate_win" + + +def test_recover_no_gain_when_sweep_kept_nothing(tmp_path): + """A sweep that kept the baseline config (no change / no gain) stays no_gain — + the config tier must not manufacture a win.""" + eval_dir = _make_no_gain_eval_dir(tmp_path) + _write_sweep(eval_dir, baseline_tput=604.8, best_tput=604.8, speedup=1.0, + flags="--trust-remote-code --kv-cache-dtype fp8_e4m3", env="", + base_flags="--trust-remote-code --kv-cache-dtype fp8_e4m3", base_env="") + wf = rx._recover_workflow_return(eval_dir.parent) + assert wf.get("recovered_no_gain") is True + assert not wf.get("recovered_config_only") + out = rx.normalize_result(_handoff(eval_dir), wf) + assert out["status"] == "no_gain" + assert out["result_source"] == "disk_no_gain_synthesis" + + +def test_intermediate_win_restacks_over_default_baseline(tmp_path): + """A kernel win whose ref leg ran ON the adopted config must credit the FULL + stack (config ⊕ kernel) vs the default baseline, and carry the config's + flags/env forward for a reproducible relaunch.""" + eval_dir = _make_eval_dir(tmp_path, accepted=True) # ref_med 461.314, cand 535.352 + # Adopted config: default 400 -> config-applied ~461.3 (== the kernel ref leg). + _write_sweep(eval_dir, baseline_tput=400.0, best_tput=461.314, speedup=1.1533, + flags="--max-model-len 6144", env="VLLM_ROCM_USE_AITER=1") + wf = rx._recover_best_intermediate_win(eval_dir) + assert wf["config_restacked_over_default"] is True + assert wf["baseline_throughput_tok_s"] == pytest.approx(400.0) + assert wf["final_throughput_tok_s"] == pytest.approx(535.352) + assert wf["throughput_speedup"] == pytest.approx(535.352 / 400.0) # full stack + assert "VLLM_ROCM_USE_AITER=1" in wf["accepted_config"]["env"] + assert "--max-model-len 6144" in wf["accepted_config"]["flags"] + + +def test_intermediate_win_no_restack_when_ref_below_config(tmp_path): + """If the kernel A/B ran against the RAW baseline (ref leg well below the + config-applied throughput), do NOT re-base — that would double-count. The + config is still carried in accepted_config as a reproducible lead.""" + eval_dir = _make_eval_dir(tmp_path, accepted=True) # ref_med 461.314 + # Config-applied best (900) is far above the kernel's ref leg (461.3) => the + # kernel was NOT measured on the config; re-basing would fabricate a gain. + _write_sweep(eval_dir, baseline_tput=800.0, best_tput=900.0, speedup=1.125, + flags="--max-model-len 6144", env="VLLM_ROCM_USE_AITER=1") + wf = rx._recover_best_intermediate_win(eval_dir) + assert not wf.get("config_restacked_over_default") + assert wf["baseline_throughput_tok_s"] == pytest.approx(461.314) # unchanged + assert wf["throughput_speedup"] == pytest.approx(535.352 / 461.314) + assert "VLLM_ROCM_USE_AITER=1" in wf["accepted_config"]["env"] + + def test_result_source_live_workflow_return(tmp_path): """A live (scraped) workflow return — no recovery flags — is the canonical source and stamps result_source=workflow_return.""" diff --git a/kb/__init__.py b/kb/__init__.py new file mode 100644 index 000000000..0bef7e7ce --- /dev/null +++ b/kb/__init__.py @@ -0,0 +1,23 @@ +"""The knowledge-base plane, shared by both workflows. + +Nothing in here knows what a kernel is or what an e2e serving run is. It owns the two things the +two workflows have in common — the ADDRESS a record is filed under and the STORE it is filed in — +and leaves the workflow-specific halves (which ladder to build, what to put in `value`) to the +callers in kernel_workflow/ and e2e_workflow/. + + identity.py canonical ids and the fallback ladders, for both domains + store_local.py a KB Store held on disk, in the shape the service uses + store_remote.py the HTTP plane, wearing store_local's interface + store_client.py the standalone HTTP client store_remote is built on + retract.py taking back a record on a store that has no delete + attest.py counting what happened when a record was actually tried on a box + remote_upload.py CLI: push exported records at either plane + +This lives at the repo root rather than under kernel_workflow/scripts/ — where it grew — because +e2e_workflow imports it too, and a workflow should not have to depend on a sibling workflow's +directory being present to read its own knowledge. + +The modules are imported absolutely (`from kb.store_local import ...`), so an entry point that is +executed as a file rather than imported has to put the repo root on sys.path first. Every such +caller does this from its own __file__ and nothing here assumes a cwd. +""" diff --git a/kb/attest.py b/kb/attest.py new file mode 100644 index 000000000..f21c16e90 --- /dev/null +++ b/kb/attest.py @@ -0,0 +1,230 @@ +#!/usr/bin/env python3 +"""Attestation: counting what happened when a stored record was actually TRIED on a box. + +A record's `validated` flag is a judgement its own writer made about its own measurement, once, at +write time. It answers "did the run that produced this number clear a gate", and it can never +answer the question a reader is really asking three months later: **has anyone else since pulled +this record out of the store, run it, and had it work.** That second question is the one that +decides whether a record is worth keeping, and nothing in the schema recorded it — so a config that +has been recalled six times and failed six times looked exactly like one nobody had ever tried. + +This module adds that ledger, under `value.attestations`, in one vocabulary both lanes use: + + recalls how many times this record was pulled and actually PUT ON A BOX + validations of those, how many reproduced a win <- the retire signal + failures of those, how many ran but did not win + not_reproduced of those, how many could not be made to run at all + last_outcome / last_at / last_by the most recent attempt, for a reader in a hurry + history the last HISTORY_LIMIT attempts with their evidence + +`recalls` counts ATTEMPTS ON HARDWARE, not reads. A record that a resolve listed in its top-N and +nobody benched has learned nothing about itself, and counting that as a recall would make the ratio +`validations / recalls` — the only number a retire pass can act on — decay for records that were +never actually doubted. + +WHY THIS IS NOT RETRACTION, even though both rewrite a session document in place. Retraction +declares a record FALSE and therefore has to move it: it zeroes the ranking scalars and re-points +the champion, because a flag alone is inert against a reader that ranks on `knowledge.` +(see kb/retract.py). An attestation declares nothing. One failure on one box is evidence, not a +verdict — the box may have a different ROCm, the flag may have been renamed upstream, the workload +may be off the point the record was tuned for. So this module touches NO ranking scalar and NEVER +re-points the champion. Deciding that an accumulated ledger has become damning, and calling +`retract_session` on the strength of it, is a separate act by a separate caller, which is exactly +the separation that makes the counters trustworthy as input to it. + +Rewrites go through the same `mode="replace"` path both planes' `write()` already use, and reuse +`kb/retract.py:existing_files` so a local rewrite re-supplies the artifacts it would otherwise +delete. +""" + +from __future__ import annotations + +import time + +from kb.retract import existing_files +from kb.store_local import KBStoreError + +# The three things that can happen when a record is taken off the shelf and run. They are counted +# separately because they mean opposite things to a retire pass: `failed` says the claim did not +# hold HERE (the record may still be right elsewhere), while `not_reproduced` says the record could +# not even be applied — a much stronger signal that it is missing something it promised. +VALIDATED = "validated" +FAILED = "failed" +NOT_REPRODUCED = "not_reproduced" +OUTCOMES = (VALIDATED, FAILED, NOT_REPRODUCED) + +# `history` is bounded because it rides inside every knowledge document, and the documents are +# fetched one-per-candidate to rank a page. An unbounded audit log would make every read of a +# popular record slower for a benefit nobody has asked for; the counters are the durable part. +HISTORY_LIMIT = 20 + +# What a history entry may carry, whitelisted so an over-eager caller cannot grow the documents +# without meaning to. The two lanes measure different things — e2e in tokens per second against a +# baseline, kernels in an isolated speedup ratio — and both spellings are here rather than one +# generic `measurement`, because a reader that cannot tell 1.8 tok/s from 1.8x has learned nothing. +_EVIDENCE_KEYS = ("measured_tok_s", "baseline_tok_s", "delta_pct", "measured_speedup", "parity", + "note", "canonical_id", "workload") + + +def empty_attestations() -> dict: + return {"recalls": 0, "validations": 0, "failures": 0, "not_reproduced": 0, + "last_outcome": "", "last_at": "", "last_by": "", "history": []} + + +def _counter(value) -> int: + """A count from a document we did not write. Anything unusable reads as 0, never as a crash.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + return 0 + try: + return max(0, int(value)) + except (TypeError, ValueError, OverflowError): + return 0 + + +def attestations_of(value) -> dict: + """The ledger inside a record's `value`, normalized. Always safe to read fields off.""" + raw = value.get("attestations") if isinstance(value, dict) else None + if not isinstance(raw, dict): + return empty_attestations() + ledger = empty_attestations() + for key in ("recalls", "validations", "failures", "not_reproduced"): + ledger[key] = _counter(raw.get(key)) + for key in ("last_outcome", "last_at", "last_by"): + ledger[key] = str(raw.get(key) or "") + history = raw.get("history") + ledger["history"] = [h for h in history if isinstance(h, dict)][-HISTORY_LIMIT:] \ + if isinstance(history, list) else [] + return ledger + + +def carry_attestations(previous_value, fresh_value: dict) -> dict: + """Move an existing record's ledger onto the document that is about to REPLACE it. + + Both lanes content-address their session ids off the thing being recorded and not off its + measurement, so re-benching one configuration deliberately lands on the SAME session id and + rewrites it. Without this, every re-measurement would silently reset the record's whole + validation history to zero — and the reset would be invisible, because the new document looks + perfectly well-formed. Returns `fresh_value` unchanged when there is nothing to carry. + """ + if not isinstance(previous_value, dict) or "attestations" not in previous_value: + return fresh_value + ledger = attestations_of(previous_value) + if not any(ledger[k] for k in ("recalls", "validations", "failures", "not_reproduced")): + return fresh_value + fresh_value["attestations"] = ledger + return fresh_value + + +def record_attestation(value: dict, outcome: str, *, actor: str = "", evidence=None, + when: str = "") -> dict: + """`value` with one attempt counted onto its ledger. Pure — returns a new dict. + + Split out from the store call so a dry run can show the exact ledger that would land, and so + the kernel lane can apply the identical arithmetic to its on-disk meta.yaml without going + through a KB plane at all. + """ + outcome = str(outcome or "").strip().lower() + if outcome not in OUTCOMES: + raise KBStoreError("unknown attestation outcome %r; expected one of %s" + % (outcome, ", ".join(OUTCOMES))) + ledger = attestations_of(value) + stamp = when or time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + ledger["recalls"] += 1 + ledger[{VALIDATED: "validations", FAILED: "failures", + NOT_REPRODUCED: "not_reproduced"}[outcome]] += 1 + ledger.update({"last_outcome": outcome, "last_at": stamp, "last_by": str(actor or "")}) + entry = {"at": stamp, "outcome": outcome} + if actor: + entry["by"] = str(actor) + for key in _EVIDENCE_KEYS: + item = (evidence or {}).get(key) if isinstance(evidence, dict) else None + if item not in (None, "", [], {}): + entry[key] = item + ledger["history"] = (ledger["history"] + [entry])[-HISTORY_LIMIT:] + updated = dict(value if isinstance(value, dict) else {}) + updated["attestations"] = ledger + return updated + + +def retire_hint(value) -> str: + """Why a curation pass might want to look at this record, or "". Advisory, never enforced. + + Deliberately conservative and deliberately not a boolean: this is read by an agent prompt and + by a human running a curation sweep, and both need to know WHICH pattern fired. Nothing in the + read path filters on it — a record with a hint is still offered, still ranked, still adoptable. + """ + ledger = attestations_of(value) + tried = ledger["recalls"] + if not tried: + return "" + if ledger["not_reproduced"] >= 2 and not ledger["validations"]: + return ("%d attempts could not reproduce it at all and none ever succeeded — the record is " + "probably missing something it promised" % ledger["not_reproduced"]) + if tried >= 3 and not ledger["validations"]: + return ("tried %d times, never reproduced a win" % tried) + return "" + + +def attested_document(knowledge: dict, outcome: str, *, actor: str = "", evidence=None) -> dict: + """A copy of `knowledge` with the attempt counted. Pure — writes nothing. + + Every ranking scalar is left exactly as it was; see the module docstring for why an attestation + must not move a record in the ordering the way a retraction must. + """ + if not isinstance(knowledge, dict): + raise KBStoreError("knowledge is not an object") + document = dict(knowledge) + value = document.get("value") if isinstance(document.get("value"), dict) else {} + document["value"] = record_attestation(dict(value), outcome, actor=actor, evidence=evidence) + return document + + +def attest_session(store, canonical_id: str, session_id: str, outcome: str, *, actor: str = "", + evidence=None, apply: bool = True): + """Count one attempt against one session on one plane. Never raises for a missing record. + + Mirrors `retract_session`'s report shape so a caller that already handles one can handle the + other: a rung the write never reached reports `found: false` rather than failing the run. + """ + report = {"canonical_id": canonical_id, "session_id": session_id, "outcome": outcome, + "found": False, "rewritten": False, "attestations": {}, "error": ""} + knowledge = store.get_session(canonical_id, session_id) + if not isinstance(knowledge, dict): + report["error"] = "no such session on this plane (nothing to attest)" + return report + report["found"] = True + try: + document = attested_document(knowledge, outcome, actor=actor, evidence=evidence) + except KBStoreError as e: + report["error"] = str(e) + return report + report["attestations"] = document["value"]["attestations"] + report["retire_hint"] = retire_hint(document["value"]) + if not apply: + report["would_write"] = document + return report + try: + store.write(canonical_id, session_id, document, + existing_files(store, canonical_id, session_id)) + report["rewritten"] = True + except (KBStoreError, OSError) as e: + report["error"] = "%s: %s" % (type(e).__name__, str(e)[:160]) + return report + + +def attestation_ok(reports, applied: bool) -> bool: + """Did the attestation land, judged over every (page, plane) it visited. + + Same rule as `retraction_ok`: a rung that does not hold the record has nothing to count and is + not a failure, but finding it nowhere at all means the session id is wrong and the caller's + verdict was recorded against nothing. + """ + if not applied: + return True + found = [r for r in reports if r.get("found")] + return bool(found) and all(r.get("rewritten") for r in found) + + +__all__ = ["FAILED", "HISTORY_LIMIT", "NOT_REPRODUCED", "OUTCOMES", "VALIDATED", + "attest_session", "attestation_ok", "attestations_of", "attested_document", + "carry_attestations", "empty_attestations", "record_attestation", "retire_hint"] diff --git a/kb/curate.py b/kb/curate.py new file mode 100644 index 000000000..bb05f6406 --- /dev/null +++ b/kb/curate.py @@ -0,0 +1,26 @@ +"""Curation shared by both lanes' readers: collapsing a ranked candidate list to one rank per idea. + +The store ranks on a scalar; the reader wants breadth. Three implementations of one `direction` +are one idea offered three times, and every candidate a lane adopts costs a full on-box verify, so +the runners-up ride along as `alternates` instead of each taking a slot. An entry with no direction +is its own group — unlabeled is honest, not a group. +""" + + +def collapse_by_direction(ordered, direction_of, unique_of, top_n): + """One rank per IDEA, best first. Input must already be in rank order. + + `direction_of(item)` -> the idea label (case/space-insensitive; empty => its own group). + `unique_of(item)` -> a stable per-item key, used to name an undirected group. + Returns (chosen, alternates-per-chosen, how many were collapsed). + """ + groups, order = {}, [] + for item in ordered: + key = str(direction_of(item) or "").strip().lower() or "__undirected__" + unique_of(item) + if key not in groups: + groups[key] = [] + order.append(key) + groups[key].append(item) + chosen = order[: max(1, int(top_n or 3))] + return ([groups[k][0] for k in chosen], [groups[k][1:] for k in chosen], + sum(len(groups[k]) - 1 for k in chosen)) diff --git a/kb/identity.py b/kb/identity.py new file mode 100644 index 000000000..b42edaad9 --- /dev/null +++ b/kb/identity.py @@ -0,0 +1,249 @@ +#!/usr/bin/env python3 +"""Canonical ids for both workflows, and the fallback ladders that make a miss recoverable. + +One file rather than two because the failure this guards against is silent. The `geak` scheme is +CLIENT-DEFINED: the service declares no dimensions for it, does EXACT canonical-id lookup only +(`POST /v1/kb/search` answers `search_unsupported`), and a GET that misses by one segment returns a +plain 404 that is indistinguishable from "nothing was ever recorded here". The API doc says what +that costs, in as many words: + + 同一个 kernel 用不同维度个数或顺序描述,会落在不同的 rollup 上,服务端无法把这种情况和 + 「两个不同的 kernel」区分开,症状是历史凭空消失。 + +A reader and a writer that disagree by one segment therefore do not raise — the run just cold +starts and nobody finds out. So both sides of both workflows build their address here, and nowhere +else. + +Two schemes, distinguished by the domain segment right after `geak`: + + kernel geak:kernel:{gfx}:{kernel_name}:{backend}:rocm[:{version}] + e2e geak:e2e:{model}:{gfx}:{framework}:{version}:{precision}[:tp_N[:isl_N:osl_N:conc_N]] + +Ordering is by how badly a mismatch hurts, most damaging first, because a canonical id can only be +truncated from the RIGHT. That is what makes the tail droppable and the head mandatory: + + * `gfx` leads because an arch mismatch is the nastiest failure mode available. A gfx942-tuned + patch compiles clean on gfx950 and is simply slower — no error, just a wasted round. + * `kernel_name` and `backend` next: get either wrong and the patch does not apply at all, which + at least fails loudly. + * the ROCm version last, because 7.2 -> 7.3 usually keeps a patch applicable. It is the one + dimension worth being able to drop. + +`producer` is deliberately NOT a dimension, unlike upstream's `kernel:` scheme. It has exactly one +value on our side, and it is not lost by leaving it out: the service stamps it on every record and +the artifact prefix is `kb///`. Keeping it would only cost a segment. The +consequence to accept is that dropping the leading segments no longer yields a native `kernel:` id +verbatim, so widening the credential later means running a conversion rather than a prefix rewrite. + +THE LADDER IS WRITTEN, NOT COMPUTED. The service does no prefix aggregation — "rollup" there means +the session index under one exact canonical id (`rollup 没有自己的 uuid,直接用 canonical_id +寻址`), and no endpoint takes a prefix or depth. A coarse rung answers a read only if somebody +wrote it. So `canonical_ids()` returns every rung, most specific first, and writers must publish to +ALL of them unconditionally. Publishing conditionally is worse than not publishing: the coarse page +becomes a biased subset of runs and ranks worse than an empty one. + +The rungs share one session id, fingerprinted from the MOST SPECIFIC rung, so a measurement is +recognisably one thing at every address it appears. That sharing does NOT make the extra rungs free +on the wire: artifacts live under `kb///` and the storage really is shared, +but the file MANIFEST is per (canonical_id, session_id), so a rung that skips put_files reports +file_count 0 and downloads a bundle with no patch in it. Every rung uploads. What the shared id buys +is that the bytes land on the same keys, not that they are sent once. + +Stdlib only, like its siblings, so a lane agent can reach it over Bash. +""" + +import hashlib +import re + +SCHEME = "geak" +KERNEL_DOMAIN = "kernel" +E2E_DOMAIN = "e2e" + +# kernel-side framework is `rocm` for triton, hip AND ck. Upstream means "the package that owns the +# source being patched" (vllm, sglang); our kernels are standalone extractions, so that reading has +# nothing to say. What actually moves the stack is the container image, and all three languages ship +# in the same one, so its ROCm version is the single number that says whether two speedups were +# measured on the same thing. It is kept as a literal segment rather than dropped so the ladder has +# somewhere to stop: `...:ck:rocm` is a real address, `...:ck` would be a bare prefix. +KERNEL_FRAMEWORK = "rocm" +UNKNOWN_VERSION = "unspecified" # framework known, version not observed. Never guessed. +UNKNOWN = "unknown" + +_DISALLOWED = re.compile(r"[^a-z0-9._+-]+") +_LEADING = re.compile(r"^[^a-z0-9_]+") +_UNSAFE_IN_SESSION = re.compile(r"[^A-Za-z0-9._-]+") +_NAME_BUDGET = 48 # upstream _NAME_BUDGET +_FINGERPRINT_LEN = 12 # upstream _FINGERPRINT_LEN + +# The service's own check, mirrored so a bad segment fails here instead of as a 400 after upload. +# Leading underscores are explicitly legal — Triton kernels are routinely named `_attn_fwd`. +SEGMENT_RE = re.compile(r"^[a-z0-9_][a-z0-9._+-]*$") + + +class IdentityError(ValueError): + """A dimension that cannot be a canonical-id segment.""" + + +def segment(value, fallback: str) -> str: + """Fold a free-form value into one dimension, byte-for-byte as upstream's segment().""" + folded = _DISALLOWED.sub("-", str(value or "").strip().lower()) + folded = _LEADING.sub("", folded).strip("-") + if not folded: + folded = fallback + return folded.encode("ascii", "ignore").decode("ascii")[:256] or fallback + + +def counted(prefix: str, value) -> str: + """`tp_8`, `isl_1024` — a self-describing numeric segment, or "" when there is no number. + + Bare numbers would be legal and unreadable: four of them in a row at the tail of an id (and of + the mirrored directory path) leaves nothing to check a mis-ordered write against. The prefix + costs nothing — there is no search to confuse — and makes a truncated id say where it was cut. + Returning "" rather than a placeholder is what lets the caller drop the whole workload rung: a + run that did not record its shape must not file itself under `isl_unknown` and strand its + result on a page no reader will construct. + """ + try: + number = int(str(value).strip()) + except (TypeError, ValueError): + return "" + if number <= 0: + return "" + return "%s_%d" % (prefix, number) + + +def check(canonical_id: str) -> str: + """Reject an id the service would reject, naming the offending segment.""" + parts = str(canonical_id or "").split(":") + if len(parts) < 2: + raise IdentityError("canonical id needs at least a scheme and one dimension: %r" + % (canonical_id,)) + for index, part in enumerate(parts): + if not SEGMENT_RE.fullmatch(part): + raise IdentityError("canonical_id segment %d=%r is not a valid slug component" + % (index, part)) + return canonical_id + + +# -- kernel scheme ------------------------------------------------------------------------------ + + +def kernel_identity(gfx: str, kernel_name: str, backend: str, rocm_version: str = "") -> dict: + """The four dimensions of a kernel address, already folded. + + `rocm_version` is cut to `.` on purpose. detect_stack() reads whatever + /opt/rocm/.info/version says, which is a full build string on some images (`7.2.0-98765`), + while the recovered backlog only knows `7.2`. Keyed verbatim those land on different identities + and a warm start stops seeing half its own history over a patch release. The exact string still + travels in the record's value, so only the address is coarse. + """ + return { + "gpu": segment(gfx, UNKNOWN), + "kernel_name": segment(kernel_name, UNKNOWN), + "backend": segment(backend, UNKNOWN), + "framework": KERNEL_FRAMEWORK, + "framework_version": _short_version(rocm_version), + } + + +def _short_version(raw: str) -> str: + text = str(raw or "").strip() + if not text: + return UNKNOWN_VERSION + m = re.match(r"\s*(\d+(?:\.\d+)?)", text) + return segment(m.group(1) if m else text, UNKNOWN_VERSION) + + +def kernel_canonical_ids(identity: dict): + """Both rungs, most specific first. + + Rung 2 drops the ROCm version and nothing else. There is no third rung dropping `rocm` itself: + the segment is a constant on the kernel side, so a page without it would hold the same records + as rung 2 under a different name — a second copy of one thing, which is the exact confusion the + doc warns about. + + An entry whose ROCm was never observed keys as `unspecified` and still gets both rungs, so the + coarse page sees it even though the exact page is one nobody will look up. + """ + head = [SCHEME, KERNEL_DOMAIN, identity["gpu"], identity["kernel_name"], + identity["backend"], identity["framework"]] + return [check(":".join(head + [identity["framework_version"]])), check(":".join(head))] + + +# -- e2e scheme --------------------------------------------------------------------------------- + + +def e2e_identity(model: str, gfx: str, framework: str, framework_version: str, precision: str, + tp=None, isl=None, osl=None, conc=None) -> dict: + """The e2e address: what is being served, on what, at what shape. + + `framework` here is the SERVING stack (vllm / sglang) — the opposite of the kernel scheme's + `framework`, which is always `rocm`. The two workflows genuinely key on different things: an + extracted kernel is standalone and only its compile stack matters, while an e2e result is a + statement about a server and is worthless without knowing which one. Note this also collides + with e2e_workflow.js's own `args.backend`, which names the serving adapter, whereas the kernel + scheme's `backend` dimension names the kernel language. Nothing shares a variable across the + two, and this is why. + + `tp` sits AFTER precision and before the workload shape so the ladder can drop the measured + point while keeping the deployment config. Anything unparseable folds to "" and simply removes + the rung that would have carried it. + """ + return { + "model": segment(model, UNKNOWN), + "gpu": segment(gfx, UNKNOWN), + "framework": segment(framework, UNKNOWN), + "framework_version": segment(framework_version, UNKNOWN_VERSION), + "precision": segment(precision, UNKNOWN), + "tp": counted("tp", tp), + "isl": counted("isl", isl), + "osl": counted("osl", osl), + "conc": counted("conc", conc), + } + + +def e2e_canonical_ids(identity: dict): + """Up to three rungs, most specific first: exact workload, TP config, model. + + The last rung is TP-agnostic on purpose, and it is not just a fallback — it is the only page + that can answer "how many ways should I shard this", because that question needs TP4 and TP8 + ranked against each other rather than filed apart. The middle rung answers "given TP=8, how do + I configure it". Rungs that would need a dimension the run did not record are omitted, never + filled with a placeholder. + """ + base = [SCHEME, E2E_DOMAIN, identity["model"], identity["gpu"], identity["framework"], + identity["framework_version"], identity["precision"]] + rungs = [] + tp, isl, osl, conc = identity["tp"], identity["isl"], identity["osl"], identity["conc"] + if tp and isl and osl and conc: + rungs.append(":".join(base + [tp, isl, osl, conc])) + if tp: + rungs.append(":".join(base + [tp])) + rungs.append(":".join(base)) + return [check(r) for r in rungs] + + +# -- session id --------------------------------------------------------------------------------- + + +def session_id(exact_canonical_id: str, name: str, digest: str, producer: str = SCHEME) -> str: + """`---`, upstream's shape. + + Fingerprinted from the MOST SPECIFIC rung and reused verbatim on the coarser ones. Upstream + includes the identity fingerprint because artifacts are partitioned by session id alone, so an + id repeated across two identities makes them share an artifact path. Here that sharing is the + point — the rungs are one measurement filed at several addresses and the bytes are identical — + but it only stays safe while the fingerprint comes from a rung that is unique per measurement. + Fingerprinting each rung separately would instead give the same patch three unrelated ids, and + the coarse pages would stop being reproductions of the exact one. + """ + legible = _UNSAFE_IN_SESSION.sub("-", str(name or "")).strip("-.") + legible = legible[:_NAME_BUDGET].strip("-.") or UNKNOWN + fp = hashlib.sha256(str(exact_canonical_id or "").encode()).hexdigest()[:_FINGERPRINT_LEN] + port = _UNSAFE_IN_SESSION.sub("", str(digest or ""))[:_FINGERPRINT_LEN] + return ("%s-%s-%s-%s" % (segment(producer, SCHEME), legible, fp, port)).strip("-") + + +__all__ = ["E2E_DOMAIN", "IdentityError", "KERNEL_DOMAIN", "KERNEL_FRAMEWORK", "SCHEME", + "SEGMENT_RE", "UNKNOWN", "UNKNOWN_VERSION", "check", "counted", "e2e_canonical_ids", + "e2e_identity", "kernel_canonical_ids", "kernel_identity", "segment", "session_id"] diff --git a/kb/ladder.py b/kb/ladder.py new file mode 100644 index 000000000..c1247151c --- /dev/null +++ b/kb/ladder.py @@ -0,0 +1,31 @@ +"""Publishing one measurement down a ladder of canonical ids — shared by both lanes' writers. + +`publish(store, recs, files, score_of)` writes every rung in `recs` to ONE store, all-or-none, and +returns `(written_cids, promoted_cids, error)`. + +All rungs or none. A partially-filled ladder is the one outcome worth avoiding: the coarse page +would hold whichever runs happened to succeed twice and would rank them as if that were the whole +history — and because the scheme has no search, no reader could ever tell that page was thin. +Stopping at the first failure leaves fewer records than intended but never a page that lies about +its own completeness, since the exact rung is written first. + +`score_of(rec)` returns the scalar the champion pointer ranks on for that rung. The kernel lane +ranks every rung on `speedup`; the e2e lane ranks the exact rung on throughput and the coarser +rungs on speedup, and opens a DIFFERENT store per rung to do it — so it calls this once per rung +with a one-element `recs`, and drives all-or-none at its own loop level (breaking on a non-empty +`error`). The kernel lane, one metric for all rungs, passes the whole ladder in one call. +""" + + +def publish(store, recs, files, score_of): + """Write every rec to `store`, all-or-none. Returns (written, promoted, error).""" + written, promoted = [], [] + for rec in recs: + try: + store.write(rec["canonical_id"], rec["session_id"], rec["knowledge"], files) + written.append(rec["canonical_id"]) + if store.maybe_promote(rec["canonical_id"], rec["session_id"], score_of(rec)): + promoted.append(rec["canonical_id"]) + except Exception as e: + return written, promoted, "%s: %s" % (type(e).__name__, str(e)[:160]) + return written, promoted, "" diff --git a/kb/plane.py b/kb/plane.py new file mode 100644 index 000000000..b29ac93a1 --- /dev/null +++ b/kb/plane.py @@ -0,0 +1,67 @@ +"""Opening a KB plane — the one place both lanes turn a CLI namespace into a store to read or write. + +`open_plane(a, metric, floor, create)` returns `(primary, mirror_or_None, why)` for one plane: + + plane=local -> (LocalKBStore, None, "") reads/writes the on-disk store at a.store + plane=remote -> (RemoteKBStore, None, "") reads/writes the service + plane=both -> (LocalKBStore, RemoteKBStore, why) local is the source of truth; the mirror + is reported, never fatal + +The kernel lane ranks on one metric (`speedup`) and passes `(CHAMPION_METRIC, 1.0)`; the e2e lane +ranks each rung on its own metric (throughput on the exact rung, speedup on the coarser ones) and +passes the rung's `(metric, floor)`. That per-metric parameter is the ONLY thing that used to make +these two openers different functions. + +`--plane both` writes locally FIRST and remotely second, and a remote failure surfaces as `why` +(`remote_unavailable: ...`) rather than as a refusal: the run already spent GPU hours producing the +measurement, and a network blip must not discard it. Without that field an unreachable service +would look exactly like a successful write. +""" + +import os + + +def open_plane(a, metric, floor, create=False): + """(primary, mirror_or_None, why). See module docstring. + + A missing store is a hard miss when reading — a typo'd path must not read as an empty store and + quietly cold-start a run that had experience waiting. Writing creates it, because the first + write into a fresh store is the normal case, not an error. + """ + plane = str(getattr(a, "plane", "local") or "local") + + local = None + if plane in ("local", "both"): + local, why = _open_local(str(getattr(a, "store", "") or ""), metric, floor, create) + if plane == "local": + return local, None, why + if local is None: + return None, None, why + + # plane in ("remote", "both") + try: + from kb.store_remote import RemoteKBStore + except ImportError as e: + return None, None, "store_unavailable: " + str(e)[:120] + scan = getattr(a, "scan", 25) + remote, remote_why = RemoteKBStore.from_env(scan, metric, floor) + if plane == "remote": + return remote, None, remote_why + return local, remote, ("remote_unavailable: " + remote_why if remote is None else "") + + +def _open_local(root, metric, floor, create): + """The on-disk store, or (None, reason). Imported lazily so a box that only has this file can + still open a remote plane.""" + try: + from kb.store_local import LocalKBStore + except ImportError as e: + return None, "store_unavailable: " + str(e)[:120] + if not root or not os.path.isdir(root): + if not create: + return None, "no_such_store: " + root + try: + os.makedirs(root, exist_ok=True) + except OSError as e: + return None, "unusable_store: " + str(e)[:120] + return LocalKBStore(root, metric=metric, promote_floor=floor), "" diff --git a/kb/remote_upload.py b/kb/remote_upload.py new file mode 100755 index 000000000..0e2225173 --- /dev/null +++ b/kb/remote_upload.py @@ -0,0 +1,195 @@ +#!/usr/bin/env python3 +"""Push `experience_store.py export-remote` output to a KernelForge KB Store. + + experience_store.py export-remote --root kb_artifacts --out /tmp/kb.jsonl + kb/remote_upload.py --records /tmp/kb.jsonl # dry run: says what it would do + kb/remote_upload.py --records /tmp/kb.jsonl --local --apply # to the on-disk plane + kb/remote_upload.py --records /tmp/kb.jsonl --apply # to the service + +The two planes take the SAME records, byte for byte. That is the whole point of --local: the +read/apply/optimize/write-back loop can be proven offline, and what proves it is that the service +would receive exactly what the local store received. + +Split from the exporter on purpose. The exporter is pure and offline, so the whole mapping can be +reviewed and diffed before anything leaves the machine; this half is the only code that talks to +the network, and it does nothing until --apply. + +Order per candidate is artifacts, then knowledge, then the champion pointer. That way a record is +never visible referencing bytes the store does not hold yet, and a run interrupted halfway leaves +uploaded-but-unreferenced blobs rather than a record pointing at nothing. + +Needs the upstream client on PYTHONPATH (KernelForge `src/`, or our vendored kb/store_client.py) +plus KB_STORE_URL / KB_STORE_TOKEN. The token is read from the environment and never printed: +--apply logs the canonical id and session id only. +""" + +import argparse +import json +import os +import sys + +# Executed as a file, so the repo root is not on sys.path yet and `kb.` would not resolve. +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +def _load_client(): + """Import the upstream client, or explain precisely what is missing.""" + try: + from kernel_agents.knowledge.remote_exp.kb_store_client import KBStoreClient, KBStoreError + return KBStoreClient, KBStoreError + except ImportError: + pass + try: # our vendored copy of the single file + from kb.store_client import KBStoreClient, KBStoreError # type: ignore + return KBStoreClient, KBStoreError + except ImportError as e: + raise SystemExit( + "cannot import KBStoreClient: " + str(e) + "\n" + " put KernelForge's src/ on PYTHONPATH, or vendor " + "kernel_agents/knowledge/remote_exp/kb_store_client.py at kb/store_client.py" + ) + + +class _LocalBackend: + """The on-disk store behind the client's write surface. + + Same three calls in the same order as the service path, so `upload_one` below does not know + which plane it is writing to. The local store lands a session as one atomic unit, so the + artifacts-then-knowledge ordering is belt and braces here — it matters upstream, where they + are two round trips. + """ + + def __init__(self, root: str): + from kb.store_local import LocalKBStore # noqa: PLC0415 - optional, only for --local + self.store = LocalKBStore(root) + self.root = self.store.root + self._staged = {} + + def put_files(self, cid, sid, entries): + self._staged[(cid, sid)] = {rel: local for (rel, local, _kind, _meta) in entries} + + def put_knowledge(self, cid, knowledge, session_id="", mode="merge"): + files = self._staged.pop((cid, session_id), {}) + return {"session_id": session_id, + "path": self.store.write(cid, session_id, knowledge, files)} + + def set_champion(self, cid, sid, metric="speedup", value=0.0): + self.store.promote(cid, sid, value) + return {"session_id": sid, "metric": metric, "value": value} + + +def read_records(path: str): + out = [] + with open(path, "r", errors="replace") as f: + for n, line in enumerate(f, start=1): + line = line.strip() + if not line: + continue + try: + rec = json.loads(line) + except ValueError as e: + raise SystemExit(f"{path}:{n}: not JSON: {e}") + # A summary line from a stdout capture is not a record; skip it rather than fail, so + # `export-remote ... | tee` output works as an input without hand-editing. + if not isinstance(rec, dict) or not rec.get("canonical_id"): + continue + out.append(rec) + return out + + +def upload_one(store, rec: dict, *, apply: bool, quiet: bool) -> dict: + cid, sid = rec["canonical_id"], rec["session_id"] + files = rec.get("files") or [] + # EVERY rung gets its own put_files, including the coarse ones that share a session id with the + # exact rung. The S3 prefix really is `kb///` and is therefore shared, but + # the file MANIFEST is per (canonical_id, session_id): committing only under the exact id leaves + # the coarse page reporting file_count 0, and a reader that falls back to it materializes a + # bundle with no patch.diff in it. Verified against the service, not assumed — the first cut of + # this skipped rungs > 0 and produced exactly that. + # + # The cost is re-PUTting identical bytes to a key that already holds them, once per extra rung. + # That is the ladder's real price: transfer scales with depth even though storage does not. + missing = [f["local_path"] for f in files if not os.path.isfile(f.get("local_path") or "")] + if missing: + return {"canonical_id": cid, "session_id": sid, "ok": False, + "reason": "missing_local_file: " + missing[0]} + plan = {"canonical_id": cid, "session_id": sid, "files": [f["path"] for f in files], + "bytes": sum(int(f.get("size") or 0) for f in files), + "speedup": (rec.get("knowledge") or {}).get("speedup"), + "champion": bool(rec.get("champion"))} + if not apply: + return dict(plan, ok=True, applied=False) + + if files: + store.put_files(cid, sid, [ + (f["path"], f["local_path"], f.get("kind") or "rewrite", {}) for f in files + ]) + # replace, not merge: this exporter emits the entry's complete current state every time, so a + # merge would keep fields a later curation pass deliberately removed. + store.put_knowledge(cid, rec["knowledge"], session_id=sid, mode="replace") + if rec.get("champion") and rec.get("champion_eligible"): + store.set_champion(cid, sid, metric="speedup", value=float(plan["speedup"] or 0.0)) + if not quiet: + print(json.dumps(dict(plan, ok=True, applied=True), ensure_ascii=False)) + return dict(plan, ok=True, applied=True) + + +def main(argv=None): + p = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--records", required=True, help="JSON lines from export-remote") + p.add_argument("--local", default="", metavar="DIR", + help="write to an on-disk KB store at DIR instead of the service") + p.add_argument("--apply", action="store_true", help="actually upload; default is a dry run") + p.add_argument("--limit", type=int, default=0, help="stop after N candidates (smoke tests)") + p.add_argument("--quiet", action="store_true", help="summary line only") + a = p.parse_args(argv) + + records = read_records(a.records) + if a.limit > 0: + records = records[: a.limit] + + store = None + if a.apply and a.local: + store = _LocalBackend(a.local) + elif a.apply: + KBStoreClient, _KBStoreError = _load_client() + if not (os.environ.get("KB_STORE_URL") or "").strip(): + raise SystemExit("KB_STORE_URL is not set; refusing to --apply") + store = KBStoreClient.from_env() + + ok = failed = sent_bytes = 0 + failures = [] + for rec in records: + try: + result = upload_one(store, rec, apply=a.apply, quiet=a.quiet) + except Exception as e: # one bad candidate must not abandon the rest of the backlog + result = {"canonical_id": rec.get("canonical_id"), "session_id": rec.get("session_id"), + "ok": False, "reason": f"{type(e).__name__}: {str(e)[:200]}"} + if result.get("ok"): + ok += 1 + sent_bytes += int(result.get("bytes") or 0) + if not a.apply and not a.quiet: + print(json.dumps(result, ensure_ascii=False)) + else: + failed += 1 + failures.append(result) + print(json.dumps(result, ensure_ascii=False), file=sys.stderr) + + print(json.dumps({"applied": bool(a.apply), + "plane": "local" if a.local else "service", + "root": getattr(store, "root", "") if a.local else "", + "candidates": len(records), "ok": ok, "failed": failed, + "champions": sum(1 for r in records if r.get("champion")), + "identities": len({r["canonical_id"] for r in records}), + "sessions": len({r.get("session_id") for r in records}), + # What actually goes over the wire, ladder depth included: every rung commits + # its own manifest, so N rungs really do send the artifacts N times. Reported + # from the plans rather than the records so a skipped or failed candidate is + # not counted as transferred. + "bytes": sent_bytes}, ensure_ascii=False)) + return 1 if failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/kb/retract.py b/kb/retract.py new file mode 100644 index 000000000..76d034fd9 --- /dev/null +++ b/kb/retract.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +"""Retraction for the knowledge base: how a record that turned out to be false is taken back. + +The KB Store service exposes no DELETE. Nothing written to it can be removed, and both lanes' +writers say so in their prompts, deliberately. That is not the same as "a wrong record is forever +authoritative", and the difference is this module. + +`put_knowledge(..., mode="replace")` rewrites a whole session document in place, and both planes' +`write()` already use that mode. A session id is a deterministic digest of the record's content, so +any record we wrote can be addressed again later without having kept a note of it. Retraction is +therefore a REWRITE to a tombstone, not a deletion: the page still holds a session with that id, +but the session now says "this was wrong, here is why". + +**A flag alone would be inert.** Ranking on both planes reads a single top-level scalar +(`knowledge.`), and `sessions/top?metric=` is what a reader pages through. A document whose +`retained` is false but whose `speedup` still says 1.9 keeps its place at the head of the list for +every consumer that has not been taught the flag — including the service's own rollup. So a +retraction has to do three things at once, and doing two of them is worse than doing none: + + 1. mark the document (`retained: false` + a `retired_reason` a human can act on), + 2. drop the ranking scalars below every promote floor, so it sinks in the ordering itself, + 3. re-point the identity's champion, because the champion pointer is a separate object that does + not re-derive itself from the sessions it points at. + +Step 3 is the one that surprises. Zeroing a document does not un-promote it; the champion record +keeps the score it was promoted with, and `champion_speedup()` keeps returning that number as the +bar every future candidate must clear. A retracted champion left in place quietly blocks the page. + +What retraction does NOT do: it does not touch the artifacts. The bytes are what let someone audit +the claim afterwards, which is precisely what you want when a record has just been called false. +""" + +from __future__ import annotations + +import os +import time + +from kb.store_local import Candidate, KBStoreError, finite_speedup + + +# Written into `value.lifecycle`. The other two values in circulation are "active" (reproduced) and +# "candidate" (recorded but unreproduced); this is a third state, not a degree of the first two. +RETRACTED = "retracted" + + +def is_retired(value) -> bool: + """Whether a record's `value` says it has been taken back. + + Two spellings because two producers exist. `retained is False` is the kernel lane's local + curation flag (note `is False`, not falsy: a record that never set the key is not retired, and + `None` is "unstated", which is the common case). `retired_reason` is what retraction below + writes, and a non-empty reason is itself the claim. + """ + if not isinstance(value, dict): + return False + return value.get("retained") is False or bool(value.get("retired_reason")) + + +def retracted_document(knowledge: dict, reason: str, metrics, actor: str = "") -> dict: + """A copy of `knowledge` rewritten into a tombstone. Pure — writes nothing. + + Kept separate from the store call so a dry run can show the exact document that would land. + """ + if not isinstance(knowledge, dict): + raise KBStoreError("knowledge is not an object") + reason = str(reason or "").strip() + if not reason: + raise KBStoreError("a retraction needs a reason: it is the only thing a future reader has " + "to judge whether the retraction itself was right") + document = dict(knowledge) + value = dict(document.get("value") if isinstance(document.get("value"), dict) else {}) + # Preserve what the record CLAIMED before we zero the ranking copies. A tombstone that has + # forgotten the number it was retracted for cannot be reviewed, only trusted. + withdrawn = {m: document.get(m) for m in metrics if document.get(m) is not None} + if withdrawn: + value.setdefault("withdrawn_scores", withdrawn) + value.update({"retained": False, "retired_reason": reason, "validated": False, + "lifecycle": RETRACTED, + "retracted_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}) + if actor: + value["retracted_by"] = str(actor) + document["value"] = value + for metric in metrics: + # 0.0, not deleted. A missing scalar makes the record unrankable, and an unrankable record + # on some routes sorts as null-last and on others is dropped from the page entirely — which + # looks like the delete this store does not have, right up until a route that keeps it + # surfaces the record with no score and no explanation. + document[metric] = 0.0 + return document + + +def retraction_ok(reports, applied: bool) -> bool: + """Did the retraction succeed, judged over every (page, plane) it visited. + + A page that does not hold the record is NOT a failure — a ladder rung the write never reached, + or a plane that only ever saw half the ladder, has nothing to take back and reporting it as an + incomplete retraction would train the caller to ignore the field. What IS a failure: finding the + record and not rewriting it, or finding it nowhere at all, which means the session id is wrong + and the record the caller wants to retract is still live somewhere they have not looked. + """ + if not applied: + return True + found = [r for r in reports if r.get("found")] + return bool(found) and all(r.get("rewritten") for r in found) + + +def existing_files(store, canonical_id: str, session_id: str): + """{relative path: absolute source} for a LOCAL session, or None for a remote one. + + The two planes lose artifacts differently on a rewrite and this is the whole reason the caller + cannot just pass `files=None`. `LocalKBStore.write()` stages a fresh directory and swaps it in, + so omitting the files DELETES them. `RemoteKBStore.write()` only calls `put_files` when it is + given some, and the manifest it does not touch survives. So: re-supply on local, stay silent on + remote (re-uploading identical bytes would be the only alternative, and it can fail). + + Public because retraction is not the only rewrite: `kb/attest.py` rewrites the same documents to + count validations, and a second copy of this rule is a second chance to get it wrong. + """ + lister = getattr(store, "session_files", None) + if not callable(lister): + return None + root = os.path.join(store.session_dir(canonical_id, session_id), "files") + return {rel: os.path.join(root, *rel.split("/")) for rel in lister(canonical_id, session_id)} + + +def _replacement_champion(store, canonical_id: str, session_id: str, metric: str, scan: int): + """(session_id, score) of the best surviving candidate, or (None, None). + + Surviving means: not the record being retracted, not already retired, and carrying a real score + under THIS rung's metric. The last clause matters because a coarse rung ranks on `speedup` and + the exact rung on `throughput_tok_s`; promoting a session whose score came from the other + metric would write a champion the store then compares future candidates against incomparably. + """ + try: + found = store.candidates(canonical_id, limit=max(1, int(scan))) + except Exception: + return None, None + for candidate in found: + if candidate.session_id == session_id: + continue + value = candidate.value if isinstance(candidate, Candidate) else {} + if is_retired(value): + continue + score = finite_speedup((candidate.knowledge or {}).get(metric)) + if score is None: + continue + return candidate.session_id, score + return None, None + + +def retract_session(store, canonical_id: str, session_id: str, reason: str, metric: str, + *, extra_metrics=(), actor: str = "", scan: int = 8, apply: bool = True): + """Retract one session on one plane. Returns a report dict; never raises for a missing record. + + `metric` is the rung's ranking metric — the one the champion is compared under. `extra_metrics` + are the other scalars the same document carries (the e2e document holds both `throughput_tok_s` + and `speedup`, and a rewrite that zeroed only one of them would leave the record ranked exactly + where it was on half the ladder). + """ + report = {"canonical_id": canonical_id, "session_id": session_id, "found": False, + "rewritten": False, "champion_was": "", "champion_now": "", "error": ""} + metrics = [metric] + [m for m in extra_metrics if m and m != metric] + knowledge = store.get_session(canonical_id, session_id) + if not isinstance(knowledge, dict): + report["error"] = "no such session on this plane (nothing to retract)" + return report + report["found"] = True + report["was_retired"] = is_retired(knowledge.get("value")) + try: + document = retracted_document(knowledge, reason, metrics, actor=actor) + except KBStoreError as e: + report["error"] = str(e) + return report + report["champion_was"] = str(store.champion(canonical_id).get("session_id") or "") + if not apply: + report["would_write"] = document + return report + try: + store.write(canonical_id, session_id, document, existing_files(store, canonical_id, + session_id)) + report["rewritten"] = True + except (KBStoreError, OSError) as e: + report["error"] = "%s: %s" % (type(e).__name__, str(e)[:160]) + return report + if report["champion_was"] and report["champion_was"] != session_id: + report["champion_now"] = report["champion_was"] # someone else already holds the slot + return report + successor, score = _replacement_champion(store, canonical_id, session_id, metric, scan) + if successor is None: + # Nothing left to crown. Re-point at the tombstone with a zero so the pointer stops + # advertising a win and stops acting as a floor — leaving the old score there would make + # every future candidate have to beat a number we just declared false. + successor, score = session_id, 0.0 + try: + store.promote(canonical_id, successor, float(score)) + report["champion_now"] = successor + except Exception as e: + report["error"] = "champion not re-pointed: %s: %s" % (type(e).__name__, str(e)[:120]) + return report diff --git a/kb/store_client.py b/kb/store_client.py new file mode 100644 index 000000000..8c243b4f3 --- /dev/null +++ b/kb/store_client.py @@ -0,0 +1,879 @@ +"""Standalone client for the KB store. + +Intentionally stdlib-only so producers (Hyperloom orchestrator, agents, +CLI tools) can vendor this single file without pulling in boto3 or an +async HTTP stack. Uploads and downloads go straight to the object store +over presigned URLs; only small JSON control messages touch the service. + +Typical producer flow:: + + store = KBStoreClient.from_env() + store.put_knowledge(cid, {"prs_tested": [...]}) + ref = store.put_file(cid, session_id, "patches/pr-123.patch", + local_path, kind="patch", + meta={"pr_url": url, "outcome": "integrated"}) + +Typical consumer flow:: + + store.download_session(cid, session_id, Path("/tmp/session")) + +Every method raises :class:`KBStoreError` on failure. Producers that treat +the KB as a best-effort side channel should catch it and carry on: losing a +record must never fail the optimization run that produced it. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import urllib.error +import urllib.parse +import urllib.request +import uuid +from collections.abc import Iterable, Mapping +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path, PureWindowsPath +from typing import Any + +DEFAULT_TIMEOUT_SEC = 60.0 +DEFAULT_PARALLELISM = 8 +_READ_CHUNK = 1024 * 1024 + +#: Bundle layout, kept identical to what the archive endpoint emits so the +#: two download routes are interchangeable for consumers. +VALUES_MEMBER = "values.json" +FILES_MEMBER_ROOT = "files" + + +class KBStoreError(RuntimeError): + """Any failure talking to the KB store or the object store.""" + + +#: Must match ``knowledge_base.canonical.RECORD_NAMESPACE``. +RECORD_NAMESPACE = uuid.UUID("0f4d5e6a-8b7c-4d1e-9f2a-3c5b7d9e1f00") + + +def record_id(canonical_id: str, session_id: str) -> str: + """Compute a record's UUID locally, without calling the service. + + The id is derived from the identity rather than allocated, so a + producer can record it (in a session breakdown, a DB row, a log line) + before the record exists and know the value will match. + + The whole identity is hashed, scheme segment included, which is what + keeps an ``inference:`` id from colliding with a ``kernel:`` one. + """ + cid = (canonical_id or "").strip() + scheme, _, dims = cid.partition(":") + if not scheme or not dims: + raise KBStoreError(f"canonical_id {canonical_id!r} is malformed") + sid = (session_id or "").strip() + if not sid: + raise KBStoreError(f"session_id {session_id!r} is malformed") + return str(uuid.uuid5(RECORD_NAMESPACE, f"{cid}|{sid}")) + + +def sha256_of(path: str | Path) -> tuple[str, int]: + """Return ``(hex_digest, size_bytes)`` for a local file.""" + digest = hashlib.sha256() + size = 0 + with open(path, "rb") as handle: + while True: + chunk = handle.read(_READ_CHUNK) + if not chunk: + break + digest.update(chunk) + size += len(chunk) + return digest.hexdigest(), size + + +def _bundle_rel_path(value: Any) -> str: + """Validate a path relative to the bundle's ``files/`` directory.""" + if not isinstance(value, str): + raise KBStoreError(f"artifact path must be a string: {value!r}") + if not value: + raise KBStoreError("artifact path is empty") + if "\0" in value: + raise KBStoreError(f"artifact path contains NUL: {value!r}") + if "\\" in value: + raise KBStoreError(f"artifact path must use forward slashes: {value!r}") + if value.startswith("/") or PureWindowsPath(value).drive: + raise KBStoreError(f"artifact path must be relative: {value!r}") + if any(part in ("", ".", "..") for part in value.split("/")): + raise KBStoreError( + f"artifact path contains an empty or traversal component: {value!r}" + ) + return value + + +def _validated_download_manifest( + listing: Any, +) -> list[tuple[str, str, int, str]]: + """Return strictly validated download entries.""" + if not isinstance(listing, Mapping): + raise KBStoreError("download manifest must be an object") + raw_files = listing.get("files") + if raw_files is None: + return [] + if not isinstance(raw_files, list): + raise KBStoreError("download manifest files must be a list") + + entries: list[tuple[str, str, int, str]] = [] + seen: set[str] = set() + for index, entry in enumerate(raw_files): + if not isinstance(entry, Mapping): + raise KBStoreError(f"download manifest entry {index} must be an object") + rel = _bundle_rel_path(entry.get("path")) + if rel in seen: + raise KBStoreError(f"duplicate artifact path in download manifest: {rel!r}") + seen.add(rel) + + expected_sha = entry.get("sha256") + if ( + not isinstance(expected_sha, str) + or len(expected_sha) != 64 + or any(char not in "0123456789abcdef" for char in expected_sha) + ): + raise KBStoreError( + f"download manifest entry {rel!r} has invalid sha256" + ) + expected_size = entry.get("size") + if ( + isinstance(expected_size, bool) + or not isinstance(expected_size, int) + or expected_size < 0 + ): + raise KBStoreError(f"download manifest entry {rel!r} has invalid size") + url = entry.get("download_url") + if not isinstance(url, str) or not url: + raise KBStoreError( + f"download manifest entry {rel!r} has no download_url" + ) + entries.append((rel, expected_sha, expected_size, url)) + return entries + + +def _checked_download_target(files_root: Path, rel: str) -> Path: + """Build a contained target without following an existing parent symlink.""" + target = files_root.joinpath(*rel.split("/")) + resolved_root = files_root.resolve() + try: + target.resolve(strict=False).relative_to(resolved_root) + except ValueError as exc: + raise KBStoreError( + f"artifact target escapes files directory: {rel!r}" + ) from exc + + current = target.parent + while current != files_root: + if current.is_symlink(): + raise KBStoreError( + f"artifact parent directory may not be a symlink: {current}" + ) + current = current.parent + return target + + +class KBStoreClient: + """Blocking client for the KB store REST surface.""" + + def __init__( + self, + base_url: str, + token: str, + *, + timeout_sec: float = DEFAULT_TIMEOUT_SEC, + parallelism: int = DEFAULT_PARALLELISM, + ) -> None: + if not base_url: + raise KBStoreError("base_url is required") + self._base = base_url.rstrip("/") + self._token = token or "" + self._timeout = timeout_sec + self._parallelism = max(1, parallelism) + + @classmethod + def from_env(cls) -> KBStoreClient: + """Build from ``KB_STORE_URL`` / ``KB_STORE_TOKEN``.""" + base = (os.environ.get("KB_STORE_URL", "") or "").strip() + token = (os.environ.get("KB_STORE_TOKEN", "") or "").strip() + if not base: + raise KBStoreError("KB_STORE_URL is not set") + return cls(base, token) + + @classmethod + def from_env_optional(cls) -> KBStoreClient | None: + """Build from env, or return ``None`` when unconfigured. + + Lets a producer make KB writes opt-in without wrapping every call + site in try/except. + """ + try: + return cls.from_env() + except KBStoreError: + return None + + # -- transport ---------------------------------------------------------- + + def _request(self, method: str, path: str, payload: Any = None) -> Any: + url = self._base + path + data = None + headers = {"Accept": "application/json"} + if self._token: + headers["Authorization"] = f"Bearer {self._token}" + if payload is not None: + data = json.dumps(payload).encode("utf-8") + headers["Content-Type"] = "application/json" + + req = urllib.request.Request(url, data=data, headers=headers, method=method) + try: + with urllib.request.urlopen(req, timeout=self._timeout) as resp: + raw = resp.read().decode("utf-8", errors="replace") + except urllib.error.HTTPError as exc: + body = exc.read().decode("utf-8", errors="replace")[:1024] + raise KBStoreError(f"{method} {path} -> HTTP {exc.code}: {body}") from exc + except Exception as exc: + raise KBStoreError(f"{method} {path} transport error: {exc!r}") from exc + + if not raw.strip(): + return None + try: + return json.loads(raw) + except json.JSONDecodeError as exc: + raise KBStoreError(f"{method} {path}: response was not JSON") from exc + + @staticmethod + def _quote(value: str) -> str: + # Colons are legal in a path segment and the canonical id relies on + # them, so they are explicitly kept unescaped. + return urllib.parse.quote(value, safe=":") + + def _session_base(self, canonical_id: str, session_id: str) -> str: + return ( + f"/v1/kb/{self._quote(canonical_id)}" + f"/sessions/{self._quote(session_id)}" + ) + + # -- knowledge ---------------------------------------------------------- + + def put_knowledge( + self, + canonical_id: str, + knowledge: dict[str, Any], + *, + session_id: str = "", + mode: str = "merge", + ) -> dict[str, Any]: + """Record what this producer knows about an identity. + + ``session_id`` names a candidate under the identity and is optional; + omitting it writes to a slot of this producer's own. Pass it to keep + separate runs comparable — the champion is picked from candidates. + The resolved id comes back as ``session_id`` in the response. + """ + payload: dict[str, Any] = {"knowledge": knowledge, "mode": mode} + if session_id: + payload["session_id"] = session_id + return self._request("POST", f"/v1/kb/{self._quote(canonical_id)}", payload) + + def get_session(self, canonical_id: str, session_id: str) -> dict[str, Any] | None: + """Read a session document, or ``None`` when it does not exist.""" + try: + return self._request("GET", self._session_base(canonical_id, session_id)) + except KBStoreError as exc: + if "HTTP 404" in str(exc): + return None + raise + + def get_record(self, rid: str) -> dict[str, Any] | None: + """Fetch a record by UUID alone, or ``None`` when it does not exist.""" + try: + return self._request("GET", f"/v1/records/{self._quote(rid)}") + except KBStoreError as exc: + if "HTTP 404" in str(exc): + return None + raise + + def get_best_record(self, canonical_id: str) -> dict[str, Any] | None: + """The record to act on for an identity, or ``None`` if there is none. + + Answers from the v1 recipe page when an identity predates this store, + so a caller does not have to know which plane its data lives in. + """ + try: + return self._request("GET", f"/v1/kb/{self._quote(canonical_id)}") + except KBStoreError as exc: + if "HTTP 404" in str(exc): + return None + raise + + def get_rollup(self, canonical_id: str) -> dict[str, Any] | None: + """Read the candidate index, or ``None`` when nothing is recorded.""" + try: + return self._request("GET", f"/v1/kb/{self._quote(canonical_id)}/sessions") + except KBStoreError as exc: + if "HTTP 404" in str(exc): + return None + raise + + def get_top_sessions( + self, + canonical_id: str, + *, + metric: str = "speedup", + limit: int = 3, + offset: int = 0, + ) -> dict[str, Any]: + """Rank scored sessions retained by the identity's rollup index.""" + query = urllib.parse.urlencode( + {"metric": metric, "limit": int(limit), "offset": int(offset)} + ) + path = f"/v1/kb/{self._quote(canonical_id)}/sessions/top?{query}" + return self._request("GET", path) or {} + + def list_identity_files( + self, canonical_id: str, *, kind: str = "" + ) -> list[dict[str, Any]]: + """Artifacts across all sessions of an identity, deduped by digest.""" + path = f"/v1/kb/{self._quote(canonical_id)}/files" + if kind: + path += "?" + urllib.parse.urlencode({"kind": kind}) + result = self._request("GET", path) or {} + return list(result.get("files") or []) + + def set_champion( + self, canonical_id: str, session_id: str, *, metric: str = "throughput", value: float = 0.0 + ) -> dict[str, Any]: + """Promote a session as the identity's best result.""" + return self._request( + "POST", + f"/v1/kb/{self._quote(canonical_id)}/champion", + {"session_id": session_id, "metric": metric, "value": value}, + ) + + # -- upload ------------------------------------------------------------- + + def put_file( + self, + canonical_id: str, + session_id: str, + rel_path: str, + local_path: str | Path, + *, + kind: str = "other", + meta: dict[str, Any] | None = None, + ) -> str: + """Upload one file and return its durable ``kb://`` reference.""" + refs = self.put_files( + canonical_id, + session_id, + [(rel_path, local_path, kind, meta or {})], + ) + return refs[rel_path] + + def put_files( + self, + canonical_id: str, + session_id: str, + items: Iterable[tuple[str, str | Path, str, dict[str, Any]]], + ) -> dict[str, str]: + """Upload a batch of files; returns ``{rel_path: kb:// reference}``. + + Digests are computed locally and declared up front, so the service + can skip bytes it already holds and can pin the uploaded object's + recorded digest into the presigned signature. + """ + validated: list[tuple[str, str | Path, str, dict[str, Any]]] = [] + seen: set[str] = set() + for rel_path, local_path, kind, meta in items: + rel = _bundle_rel_path(rel_path) + if rel in seen: + raise KBStoreError(f"duplicate artifact path for upload: {rel!r}") + seen.add(rel) + validated.append((rel, local_path, kind, meta)) + + entries: list[dict[str, Any]] = [] + sources: dict[str, Path] = {} + for rel_path, local_path, kind, meta in validated: + path = Path(local_path) + if not path.is_file(): + raise KBStoreError(f"not a file: {path}") + digest, size = sha256_of(path) + entries.append( + { + "path": rel_path, + "sha256": digest, + "size": size, + "kind": kind, + "meta": meta or {}, + } + ) + sources[rel_path] = path + if not entries: + return {} + + grant = self._request( + "POST", + self._session_base(canonical_id, session_id) + "/files:grant", + {"files": entries}, + ) + + pending = [ + (u["path"], u["upload_url"]) + for u in (grant.get("uploads") or []) + if not u.get("skip") and u.get("upload_url") + ] + by_path = {e["path"]: e for e in entries} + if pending: + with ThreadPoolExecutor(max_workers=self._parallelism) as pool: + list( + pool.map( + lambda item: self._upload_one( + item[1], sources[item[0]], by_path[item[0]]["sha256"] + ), + pending, + ) + ) + + commit = self._request( + "POST", + self._session_base(canonical_id, session_id) + "/files:commit", + {"files": entries, "verify": True}, + ) + manifest = { + str(f.get("path")): str(f.get("uri") or "") + for f in (commit.get("artifacts") or {}).get("files") or [] + } + return {rel: manifest.get(rel, "") for rel in sources} + + def put_dir( + self, + canonical_id: str, + session_id: str, + local_dir: str | Path, + *, + prefix: str = "", + kind: str = "other", + meta: dict[str, Any] | None = None, + ) -> dict[str, str]: + """Upload a whole directory tree, preserving relative paths.""" + safe_prefix = _bundle_rel_path(prefix) if prefix else "" + root = Path(local_dir) + if not root.is_dir(): + raise KBStoreError(f"not a directory: {root}") + items: list[tuple[str, Path, str, dict[str, Any]]] = [] + for path in sorted(root.rglob("*")): + if not path.is_file(): + continue + rel = path.relative_to(root).as_posix() + if safe_prefix: + rel = f"{safe_prefix}/{rel}" + items.append((rel, path, kind, dict(meta or {}))) + return self.put_files(canonical_id, session_id, items) + + def _upload_one(self, url: str, path: Path, sha256: str) -> None: + with open(path, "rb") as handle: + body = handle.read() + req = urllib.request.Request( + url, + data=body, + method="PUT", + headers={ + "Content-Type": "application/octet-stream", + # Part of the presigned signature; the URL is only valid + # for bytes declared under this digest. + "x-amz-meta-sha256": sha256, + }, + ) + try: + with urllib.request.urlopen(req, timeout=self._timeout) as resp: + if resp.status not in (200, 201, 204): + raise KBStoreError(f"upload of {path} returned HTTP {resp.status}") + except urllib.error.HTTPError as exc: + detail = exc.read().decode("utf-8", errors="replace")[:512] + raise KBStoreError(f"upload of {path} failed: HTTP {exc.code}: {detail}") from exc + except Exception as exc: + raise KBStoreError(f"upload of {path} failed: {exc!r}") from exc + + # -- download ----------------------------------------------------------- + + def list_session_files( + self, canonical_id: str, session_id: str, *, kind: str = "" + ) -> dict[str, Any]: + """Manifest with a short-lived presigned GET URL per file.""" + path = self._session_base(canonical_id, session_id) + "/files" + if kind: + path += "?" + urllib.parse.urlencode({"kind": kind}) + return self._request("GET", path) or {} + + def download_session( + self, + canonical_id: str, + session_id: str, + dest_dir: str | Path, + *, + kind: str = "", + include_values: bool = True, + ) -> list[Path]: + """Download and verify a record in the standard bundle layout. + + Produces the same tree as the archive endpoint:: + + values.json the knowledge payload + files/ every artifact + + so a consumer can read ``values.json`` and resolve any path it + references under ``files/`` without caring which of the two + download routes produced the directory. + + Artifact bytes come straight from the object store over presigned + URLs, concurrently, and never transit the KB store. + """ + listing = self.list_session_files(canonical_id, session_id, kind=kind) + entries = _validated_download_manifest(listing) + root = Path(dest_dir) + if root.is_symlink(): + raise KBStoreError(f"download destination may not be a symlink: {root}") + if root.exists() and not root.is_dir(): + raise KBStoreError(f"download destination is not a directory: {root}") + root.mkdir(parents=True, exist_ok=True) + + files_root = root / FILES_MEMBER_ROOT + targets: dict[str, Path] = {} + if entries: + if files_root.is_symlink(): + raise KBStoreError( + f"files directory may not be a symlink: {files_root}" + ) + if files_root.exists() and not files_root.is_dir(): + raise KBStoreError(f"files path is not a directory: {files_root}") + files_root.mkdir(exist_ok=True) + for rel, _expected_sha, _expected_size, _url in entries: + targets[rel] = _checked_download_target(files_root, rel) + + if include_values: + document = self.get_session(canonical_id, session_id) or {} + values = document.get("knowledge") or {} + values_target = root / VALUES_MEMBER + if values_target.is_symlink(): + raise KBStoreError( + f"values target may not be a symlink: {values_target}" + ) + try: + values_target.resolve(strict=False).relative_to(root.resolve()) + except ValueError as exc: + raise KBStoreError( + "values target escapes download destination" + ) from exc + values_target.write_text( + json.dumps(values, ensure_ascii=False, indent=2, sort_keys=True), + encoding="utf-8", + ) + + def fetch(entry: tuple[str, str, int, str]) -> Path: + rel, expected_sha, expected_size, url = entry + target = targets[rel] + partial: Path | None = None + digest = hashlib.sha256() + size = 0 + try: + target = _checked_download_target(files_root, rel) + target.parent.mkdir(parents=True, exist_ok=True) + target = _checked_download_target(files_root, rel) + partial = target.with_name( + f".{target.name}.{uuid.uuid4().hex}.partial" + ) + with urllib.request.urlopen( + url, timeout=self._timeout + ) as resp, open(partial, "xb") as out: + while True: + chunk = resp.read(_READ_CHUNK) + if not chunk: + break + digest.update(chunk) + size += len(chunk) + out.write(chunk) + actual_sha = digest.hexdigest() + if actual_sha != expected_sha: + raise KBStoreError( + f"download of {rel!r} sha256 mismatch: " + f"expected {expected_sha}, got {actual_sha}" + ) + if size != expected_size: + raise KBStoreError( + f"download of {rel!r} size mismatch: " + f"expected {expected_size}, got {size}" + ) + _checked_download_target(files_root, rel) + os.replace(partial, target) + partial = None + except KBStoreError: + raise + except Exception as exc: + raise KBStoreError(f"download of {rel!r} failed: {exc!r}") from exc + finally: + if partial is not None: + partial.unlink(missing_ok=True) + return target + + if not entries: + return [] + with ThreadPoolExecutor(max_workers=self._parallelism) as pool: + return list(pool.map(fetch, entries)) + + def download_archive( + self, canonical_id: str, session_id: str, dest_file: str | Path + ) -> Path: + """Download the session directory as a single tar.gz.""" + url = self._base + self._session_base(canonical_id, session_id) + "/archive" + headers = {} + if self._token: + headers["Authorization"] = f"Bearer {self._token}" + req = urllib.request.Request(url, headers=headers, method="GET") + target = Path(dest_file) + target.parent.mkdir(parents=True, exist_ok=True) + try: + with urllib.request.urlopen(req, timeout=self._timeout) as resp, open( + target, "wb" + ) as out: + while True: + chunk = resp.read(_READ_CHUNK) + if not chunk: + break + out.write(chunk) + except Exception as exc: + raise KBStoreError(f"archive download failed: {exc!r}") from exc + return target + + +#: Where a sectioned document keeps its per-section maps. The service treats +#: ``knowledge`` as opaque, so this is a producer-side convention rather than +#: part of the record schema; it is fixed because documents already in the +#: store are written under this key. +SECTION_ROOT = "value" + +_DRAFT_DIR_ENV = "KB_DRAFT_DIR" +_WARM_START_DIR_ENV = "KB_WARM_START_DIR" +_SECTIONS_MEMBER = "sections" +_RECIPE_MEMBER = "recipe.json" + + +def _checked_section(name: str) -> str: + """Reject a section name that would escape its subtree or collide oddly.""" + section = str(name or "").strip() + if not section: + raise KBStoreError("section name is required") + if section != section.strip("."): + raise KBStoreError(f"section {name!r} may not start or end with a dot") + bad = set(section) & set("/\\\0") + if bad or section in (".", ".."): + raise KBStoreError(f"section {name!r} may not contain a path separator") + return section + + +class SectionContent: + """One section's knowledge map plus the local files that belong to it.""" + + __slots__ = ("files", "knowledge", "section") + + def __init__( + self, + section: str, + knowledge: dict[str, Any], + files: list[Path] | None = None, + ) -> None: + self.section = section + self.knowledge = knowledge + self.files = list(files or []) + + def __repr__(self) -> str: + return ( + f"SectionContent(section={self.section!r}, " + f"keys={sorted(self.knowledge)!r}, files={len(self.files)})" + ) + + +class KnowledgeSections: + """Section-scoped staging for one knowledge document, backed by a directory. + + A producer is usually several processes: agents that each own one section + and a publisher that uploads once at the end. They cannot share a client + object, so the draft is a directory that both sides open by path. + + Layout under ``root``:: + + sections/
.json one section's staged knowledge map + files/
//... that section's artifacts, ready for put_dir + + ``files`` is laid out exactly as ``put_dir`` expects, so publishing is + ``put_dir(cid, sid, sections.files_dir)`` with no repacking. + """ + + def __init__( + self, + root: str | Path, + *, + warm_start_dir: str | Path | None = None, + ) -> None: + self.root = Path(root) + self.warm_start_dir = Path(warm_start_dir) if warm_start_dir else None + + @classmethod + def from_env(cls) -> KnowledgeSections | None: + """Open the draft an orchestrator prepared, or ``None`` when absent. + + Lets an agent stay agnostic about whether this run publishes at all: + no draft directory means nobody is collecting, so skip the write. + """ + draft = (os.environ.get(_DRAFT_DIR_ENV, "") or "").strip() + if not draft: + return None + warm = (os.environ.get(_WARM_START_DIR_ENV, "") or "").strip() + return cls(draft, warm_start_dir=warm or None) + + @property + def files_dir(self) -> Path: + """The subtree to hand to :meth:`KBStoreClient.put_dir`.""" + return self.root / FILES_MEMBER_ROOT + + # -- write --------------------------------------------------------------- + + def write( + self, + section: str, + knowledge: dict[str, Any], + *, + files: Iterable[str | Path] = (), + kind: str = "artifacts", + mode: str = "merge", + ) -> SectionContent: + """Stage one section's knowledge map and copy its files into the draft. + + ``mode="merge"`` (the default) shallow-merges over what this section + already staged and appends to its file list, so an agent that reports + incrementally does not silently drop its earlier calls. ``"replace"`` + discards the staged map first; staged files always survive because + they may already be referenced by the map being written. + """ + name = _checked_section(section) + if mode not in ("merge", "replace"): + raise KBStoreError(f"mode must be 'merge' or 'replace', got {mode!r}") + if not isinstance(knowledge, dict): + raise KBStoreError(f"knowledge for section {name!r} must be a dict") + try: + json.dumps(knowledge, allow_nan=False) + except (TypeError, ValueError) as exc: + raise KBStoreError(f"section {name!r} is not strict JSON: {exc}") from exc + + staged = self.staged(name) + merged = dict(knowledge) + refs: list[str] = [] + if staged is not None: + refs = [ + path.relative_to(self.files_dir).as_posix() for path in staged.files + ] + if mode == "merge": + merged = {**staged.knowledge, **knowledge} + + added = [self._copy_in(name, source, kind) for source in files] + for ref in added: + if ref and ref not in refs: + refs.append(ref) + + target = self.root / _SECTIONS_MEMBER / f"{name}.json" + target.parent.mkdir(parents=True, exist_ok=True) + payload = {"knowledge": merged, "files": refs} + target.write_text( + json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True), + encoding="utf-8", + ) + return SectionContent(name, merged, [self.files_dir / ref for ref in refs]) + + def _copy_in(self, section: str, source: str | Path, kind: str) -> str: + raw = str(source or "").strip() + if not raw: + return "" + src = Path(raw) + if src.is_symlink(): + raise KBStoreError(f"artifact must not be a symlink: {src}") + if not src.is_file(): + raise KBStoreError(f"artifact is not a readable file: {src}") + safe_kind = _checked_section(kind) + rel = f"{section}/{safe_kind}/{src.name}" + destination = self.files_dir / rel + if destination.exists() and not _same_bytes(src, destination): + digest = hashlib.sha256(str(src.resolve()).encode()).hexdigest()[:10] + rel = f"{section}/{safe_kind}/{src.stem}-{digest}{src.suffix}" + destination = self.files_dir / rel + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(src.read_bytes()) + return rel + + # -- read ---------------------------------------------------------------- + + def staged(self, section: str) -> SectionContent | None: + """Read back what this draft already holds for ``section``.""" + name = _checked_section(section) + target = self.root / _SECTIONS_MEMBER / f"{name}.json" + if not target.is_file(): + return None + try: + payload = json.loads(target.read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + raise KBStoreError(f"staged section {name!r} is unreadable: {exc}") from exc + knowledge = payload.get("knowledge") + refs = payload.get("files") or [] + return SectionContent( + name, + dict(knowledge) if isinstance(knowledge, dict) else {}, + [self.files_dir / str(ref) for ref in refs if str(ref).strip()], + ) + + def read(self, section: str) -> SectionContent | None: + """Return ``section`` from the warm-start record, or ``None``. + + ``None`` means this run has no prior knowledge for the section: either + nothing was downloaded, or the record predates the section. Callers + should treat it as a cold start rather than an error. + """ + name = _checked_section(section) + if self.warm_start_dir is None: + return None + recipe = self.warm_start_dir / _RECIPE_MEMBER + if not recipe.is_file(): + return None + try: + document = json.loads(recipe.read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + raise KBStoreError(f"warm start record is unreadable: {exc}") from exc + value = document.get(SECTION_ROOT) + knowledge = (value or {}).get(name) if isinstance(value, dict) else None + if not isinstance(knowledge, dict): + return None + root = self.warm_start_dir / FILES_MEMBER_ROOT / name + files = ( + sorted(path for path in root.rglob("*") if path.is_file()) + if root.is_dir() + else [] + ) + return SectionContent(name, dict(knowledge), files) + + def sections(self) -> list[str]: + """Every section staged in this draft, in a stable order.""" + root = self.root / _SECTIONS_MEMBER + if not root.is_dir(): + return [] + return sorted(path.stem for path in root.glob("*.json")) + + def document(self) -> dict[str, Any]: + """The staged ``{section: knowledge}`` map to publish under ``value``.""" + return {name: (self.staged(name) or SectionContent(name, {})).knowledge + for name in self.sections()} + + +def _same_bytes(left: Path, right: Path) -> bool: + try: + return left.read_bytes() == right.read_bytes() + except OSError: + return False diff --git a/kb/store_local.py b/kb/store_local.py new file mode 100644 index 000000000..f4f2ce183 --- /dev/null +++ b/kb/store_local.py @@ -0,0 +1,416 @@ +#!/usr/bin/env python3 +"""A KB Store held on disk, in the shape the service uses. + +This is the local plane of the same two-plane design KernelForge already ships +(`kernel_agents/rewrite_by_flydsl/record_store.py`: one `RewriteRecordStore` protocol, +`LocalRewriteRecords` on disk and `KBStoreRewriteRecords` over HTTP, chosen by +`KNOWLEDGE_STORE_MODE`). The layout, the ranking rule and the champion file are copied from +it deliberately, so the whole read/apply/optimize/write-back loop can be exercised offline and +switching to the service later is a change of backend, not of behaviour: + + /geak/kernel/geak/moe_stage1/rocm/7.2/ck/gfx950/ + champion.json {"session_id", "metric": "speedup", "value"} + sessions// + knowledge.json producer / speedup / identity / value + files/patch.diff, files/report.md + +Three properties are load-bearing and must not drift from upstream: + + * ranking is `speedup` descending and nothing else. The store does not know what a bench key + is, so it will happily order a `b:` measurement against a `b2:` one. Comparability is the + caller's job, which is why `resolve-remote` filters on it client-side. + * `candidates()` reads knowledge documents only. Artifacts are fetched by `materialize()`, for + the selected few. Patches here reach 240KB and travel through an agent's tool result. + * every mutation lands by atomic rename, and a session id repeated means overwrite. Session ids + are content-addressed upstream, so re-recording one port updates one candidate rather than + growing a new one per run. + +Stdlib only, like experience_store.py, so a lane agent can call it over Bash. +""" + +import errno +import json +import os +import re +import shutil +import tempfile +import uuid + +try: + import fcntl +except ImportError: # pragma: no cover - POSIX only in practice + fcntl = None + +KNOWLEDGE_FILENAME = "knowledge.json" +CHAMPION_FILENAME = "champion.json" +RECIPE_FILENAME = "recipe.json" +LOCK_FILENAME = ".lock" +CHAMPION_METRIC = "speedup" +ARTIFACT_KIND = "rewrite" + +# Both copied verbatim from upstream. Widening either one produces records the service would +# reject, which is exactly the failure this local plane exists to catch early. +_SESSION_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_SEGMENT_RE = re.compile(r"^[a-z0-9_][a-z0-9._+-]*$") + + +class KBStoreError(RuntimeError): + """Anything that makes a record unusable: a bad id, an unsafe path, a broken document.""" + + +def finite_speedup(value): + """The ranking key, or None when the document does not carry a usable one. + + `True` is an int in Python and would sort as 1.0; a string "1.5" would raise. Both mean the + producer wrote something we cannot rank, so both read as absent rather than as a number. + """ + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + number = float(value) + if number != number or number in (float("inf"), float("-inf")): + return None + return number + + +def validate_session_id(session_id: str) -> str: + """Reject ids that cannot be both a URL segment and a directory name.""" + raw = str(session_id or "").strip() + if not _SESSION_ID_RE.fullmatch(raw): + raise KBStoreError("unusable session id: %r" % (session_id,)) + return raw + + +def safe_rel_path(rel_path: str) -> str: + """Reject artifact paths that could escape the record's files directory.""" + if not isinstance(rel_path, str): + raise KBStoreError("unsafe artifact path: %r" % (rel_path,)) + parts = rel_path.split("/") + if (not rel_path or "\0" in rel_path or "\\" in rel_path or rel_path.startswith("/") + or ":" in rel_path or any(p in ("", ".", "..") for p in parts)): + raise KBStoreError("unsafe artifact path: %r" % (rel_path,)) + return "/".join(parts) + + +def canonical_segments(canonical_id: str): + """Render an identity as nested directory names, scheme first.""" + parts = str(canonical_id or "").split(":") + if len(parts) < 2 or any(not _SEGMENT_RE.fullmatch(p) for p in parts): + raise KBStoreError("unusable canonical id: %r" % (canonical_id,)) + return parts + + +class Candidate(object): + """One recorded port, ranked by the speedup its own document claims.""" + + __slots__ = ("session_id", "knowledge", "speedup", "is_champion") + + def __init__(self, session_id, knowledge, speedup, is_champion): + self.session_id = session_id + self.knowledge = knowledge + self.speedup = speedup + self.is_champion = is_champion + + @property + def value(self): + """The producer-owned half — opaque to the store, everything to the caller.""" + v = self.knowledge.get("value") + return v if isinstance(v, dict) else {} + + +def _write_json(path: str, document) -> None: + _atomic_bytes(path, json.dumps(document, ensure_ascii=False, indent=2, + sort_keys=True).encode("utf-8")) + + +def _atomic_bytes(path: str, payload: bytes) -> None: + directory = os.path.dirname(path) or "." + os.makedirs(directory, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix="." + os.path.basename(path) + ".", dir=directory) + try: + with os.fdopen(fd, "wb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + temporary = "" + finally: + if temporary and os.path.exists(temporary): + os.unlink(temporary) + _fsync_dir(directory) + + +def _fsync_dir(path: str) -> None: + try: + fd = os.open(path, os.O_RDONLY) + except OSError: + return + try: + os.fsync(fd) + except OSError: + pass + finally: + os.close(fd) + + +def _replace_directory(staging: str, destination: str) -> None: + """Install `staging` at `destination`, putting the previous tree back if that fails.""" + parent = os.path.dirname(destination) or "." + backup = "" + if os.path.lexists(destination): + if os.path.islink(destination) or not os.path.isdir(destination): + raise KBStoreError("existing session is not a safe directory: " + destination) + backup = os.path.join(parent, ".%s.backup-%s" % (os.path.basename(destination), + uuid.uuid4().hex)) + os.replace(destination, backup) + try: + os.replace(staging, destination) + except Exception: + if backup: + os.replace(backup, destination) + backup = "" + raise + finally: + if backup: + shutil.rmtree(backup, ignore_errors=True) + _fsync_dir(parent) + + +class LocalKBStore(object): + """Read and write one producer's candidates under a canonical identity, on disk.""" + + def __init__(self, root: str, metric: str = CHAMPION_METRIC, promote_floor: float = 1.0): + self.root = os.path.abspath(os.path.expanduser(str(root))) + # Which flat top-level `knowledge.` scalar ranks this identity, and the value a + # candidate must beat to be worth recording as champion at all. `speedup` above 1.0 is the + # kernel lane's rule and stays the default. The e2e lane's exact-workload rung ranks on + # absolute `throughput_tok_s` instead, where 1.0 would be a nonsense floor: two runs there + # share a workload but not necessarily a baseline, so the higher speedup can easily be the + # slower server. Coarser e2e rungs go back to `speedup`, because their workloads differ and + # absolute numbers across them are not comparable at all. + self.metric = str(metric or CHAMPION_METRIC) + self.promote_floor = float(promote_floor) + + # -- addressing ---------------------------------------------------------------------- + + def identity_dir(self, canonical_id: str) -> str: + return os.path.join(self.root, *canonical_segments(canonical_id)) + + def session_dir(self, canonical_id: str, session_id: str) -> str: + return os.path.join(self.identity_dir(canonical_id), "sessions", + validate_session_id(session_id)) + + def identities(self): + """Every canonical id this store holds. Only useful locally — the service has an index.""" + found = [] + for current, directories, _files in os.walk(self.root, followlinks=False): + if "sessions" not in directories: + continue + rel = os.path.relpath(current, self.root) + if rel == ".": + continue + found.append(":".join(rel.split(os.sep))) + return sorted(found) + + # -- read ---------------------------------------------------------------------------- + + def candidates(self, canonical_id: str, limit: int = 3): + """Rank this identity's candidates. Reads no artifact bytes.""" + sessions_dir = os.path.join(self.identity_dir(canonical_id), "sessions") + if not os.path.isdir(sessions_dir) or os.path.islink(sessions_dir): + return [] + champion_id = str(self.champion(canonical_id).get("session_id") or "") + found = [] + for name in sorted(os.listdir(sessions_dir)): + entry = os.path.join(sessions_dir, name) + if not os.path.isdir(entry) or os.path.islink(entry) or name.startswith("."): + continue + validate_session_id(name) + document = os.path.join(entry, KNOWLEDGE_FILENAME) + if not os.path.isfile(document) or os.path.islink(document): + continue + try: + with open(document, "r", errors="replace") as handle: + knowledge = json.load(handle) + except (OSError, ValueError): + continue # a half-written document is a miss, not a crash + if not isinstance(knowledge, dict): + continue + found.append(Candidate(name, knowledge, finite_speedup(knowledge.get(self.metric)), + name == champion_id)) + # Ties keep the session id order so two runs over one store rank identically. + found.sort(key=lambda c: (-(c.speedup if c.speedup is not None else float("-inf")), + c.session_id)) + return found[: max(0, int(limit))] if limit else found + + def get_session(self, canonical_id: str, session_id: str): + document = os.path.join(self.session_dir(canonical_id, session_id), KNOWLEDGE_FILENAME) + try: + with open(document, "r", errors="replace") as handle: + loaded = json.load(handle) + except (OSError, ValueError): + return None + return loaded if isinstance(loaded, dict) else None + + def session_files(self, canonical_id: str, session_id: str): + """Relative paths of one session's artifacts, without reading them.""" + root = os.path.join(self.session_dir(canonical_id, session_id), "files") + found = [] + for current, _dirs, filenames in os.walk(root, followlinks=False): + for name in filenames: + path = os.path.join(current, name) + if os.path.islink(path) or not os.path.isfile(path): + continue + found.append(os.path.relpath(path, root).replace(os.sep, "/")) + return sorted(found) + + def materialize(self, canonical_id: str, candidate, destination: str) -> str: + """Lay one selected candidate out as the standard bundle: recipe.json + files/.""" + session_id = candidate.session_id if isinstance(candidate, Candidate) else str(candidate) + source = self.session_dir(canonical_id, session_id) + if os.path.islink(source) or not os.path.isdir(source): + raise KBStoreError("candidate session is not a safe directory: " + source) + knowledge = self.get_session(canonical_id, session_id) + if knowledge is None: + raise KBStoreError("candidate knowledge is unreadable: " + source) + if not isinstance(candidate, Candidate): + champion_id = str(self.champion(canonical_id).get("session_id") or "") + candidate = Candidate(session_id, knowledge, + finite_speedup(knowledge.get("speedup")), + session_id == champion_id) + + os.makedirs(destination, exist_ok=True) + bundle = os.path.join(destination, session_id) + staging = tempfile.mkdtemp(prefix="." + session_id + "-", dir=destination) + try: + files_root = os.path.join(staging, "files") + os.makedirs(files_root, exist_ok=True) + for rel in self.session_files(canonical_id, session_id): + target = os.path.join(files_root, *safe_rel_path(rel).split("/")) + os.makedirs(os.path.dirname(target), exist_ok=True) + shutil.copyfile(os.path.join(source, "files", *rel.split("/")), target) + recipe = dict(knowledge) + recipe.update({"canonical_id": canonical_id, "session_id": session_id, + "is_champion": candidate.is_champion, + "champion": candidate.is_champion}) + recipe.setdefault("speedup", candidate.speedup) + _write_json(os.path.join(staging, RECIPE_FILENAME), recipe) + _replace_directory(staging, bundle) + staging = "" + finally: + if staging and os.path.isdir(staging): + shutil.rmtree(staging, ignore_errors=True) + return bundle + + def champion(self, canonical_id: str): + path = os.path.join(self.identity_dir(canonical_id), CHAMPION_FILENAME) + if not os.path.isfile(path) or os.path.islink(path): + return {} + try: + with open(path, "r", errors="replace") as handle: + loaded = json.load(handle) + except (OSError, ValueError): + return {} + return loaded if isinstance(loaded, dict) else {} + + def champion_speedup(self, canonical_id: str): + champion = self.champion(canonical_id) + # A champion recorded under a different metric is not a weaker incumbent, it is an + # incomparable one. Returning None makes the caller treat the slot as empty rather than + # rank tokens-per-second against a ratio. + if str(champion.get("metric") or "") != self.metric: + return None + return finite_speedup(champion.get("value")) + + # -- write --------------------------------------------------------------------------- + + def write(self, canonical_id: str, session_id: str, knowledge, files=None) -> str: + """Record one candidate and its artifacts under an identity, atomically. + + The whole session lands or none of it does: a reader must never see a knowledge document + that references bytes this store does not hold yet. + """ + if not isinstance(knowledge, dict): + raise KBStoreError("knowledge is not an object") + session_id = validate_session_id(session_id) + identity_dir = self.identity_dir(canonical_id) + sessions_dir = os.path.join(identity_dir, "sessions") + os.makedirs(sessions_dir, exist_ok=True) + named = {safe_rel_path(rel): src for rel, src in (files or {}).items()} + + with self._lock(identity_dir): + staging = tempfile.mkdtemp(prefix="." + session_id + ".staging-", dir=sessions_dir) + try: + files_root = os.path.join(staging, "files") + os.makedirs(files_root, exist_ok=True) + for rel in sorted(named): + target = os.path.join(files_root, *rel.split("/")) + os.makedirs(os.path.dirname(target), exist_ok=True) + shutil.copyfile(named[rel], target) + _write_json(os.path.join(staging, KNOWLEDGE_FILENAME), knowledge) + destination = os.path.join(sessions_dir, session_id) + _replace_directory(staging, destination) + staging = "" + finally: + if staging and os.path.isdir(staging): + shutil.rmtree(staging, ignore_errors=True) + return os.path.join(sessions_dir, session_id) + + def promote(self, canonical_id: str, session_id: str, speedup: float) -> None: + """Point the identity's champion at one session. The caller owns the policy.""" + document = {"session_id": validate_session_id(session_id), + "metric": self.metric, "value": float(speedup)} + identity_dir = self.identity_dir(canonical_id) + os.makedirs(identity_dir, exist_ok=True) + with self._lock(identity_dir): + _write_json(os.path.join(identity_dir, CHAMPION_FILENAME), document) + + def maybe_promote(self, canonical_id: str, session_id: str, speedup) -> bool: + """Upstream's gate, verbatim: only a real win, and only over the incumbent.""" + speedup = finite_speedup(speedup) + if speedup is None or speedup <= self.promote_floor: + return False + incumbent = self.champion_speedup(canonical_id) + if incumbent is not None and speedup <= incumbent: + return False + self.promote(canonical_id, session_id, speedup) + return True + + # -- locking ------------------------------------------------------------------------- + + class _Lock(object): + def __init__(self, path): + self.path = path + self.handle = None + + def __enter__(self): + if fcntl is None: + return self + try: + self.handle = open(self.path, "a+") + fcntl.flock(self.handle.fileno(), fcntl.LOCK_EX) + except OSError as error: + # A read-only or lock-less filesystem must not make the store unusable; the + # atomic renames below are still atomic, we just lose writer serialization. + if self.handle is not None: + self.handle.close() + self.handle = None + if error.errno not in (errno.EACCES, errno.EPERM, errno.EROFS, errno.ENOSYS, + errno.ENOLCK): + raise + return self + + def __exit__(self, *exc): + if self.handle is not None: + try: + fcntl.flock(self.handle.fileno(), fcntl.LOCK_UN) + finally: + self.handle.close() + self.handle = None + return False + + def _lock(self, identity_dir: str): + return LocalKBStore._Lock(os.path.join(identity_dir, LOCK_FILENAME)) + + +__all__ = ["ARTIFACT_KIND", "CHAMPION_METRIC", "Candidate", "KBStoreError", "LocalKBStore", + "canonical_segments", "finite_speedup", "safe_rel_path", "validate_session_id"] diff --git a/kb/store_remote.py b/kb/store_remote.py new file mode 100644 index 000000000..b124efee0 --- /dev/null +++ b/kb/store_remote.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 +"""The HTTP plane, wearing LocalKBStore's interface. + +`experience_store.py` and `e2e_store.py` read and write through one small surface — candidates(), +get_session(), materialize(), write(), maybe_promote() — so pointing a lane at the service instead +of at a directory is a constructor change and nothing else. Everything that decides WHAT gets +offered (the bench-key comparability filter, the direction collapse, the min-speedup floor) stays +in the caller, because the service ranks on nothing but the `speedup` number a producer declared +and would happily order a `b:` measurement against a `b2:` one. + +Three places where the service is not a directory, and the adapter has to say so rather than +pretend: + + * `identities()` returns nothing. There is no search on the `geak` scheme + (`POST /v1/kb/search` -> `search_unsupported`), so the pre-ladder near-miss scan that + LocalKBStore supports is simply unavailable here. The ladder itself still works — it constructs + addresses rather than discovering them, which is exactly why it was built that way. + * ranking costs one request per candidate. `GET /sessions/top` returns ids and scores only, no + knowledge, and the caller needs `direction` and `bench_key` to curate. So the top page is + fetched once and then hydrated, bounded by `scan` — a page holds up to 200 sessions and pulling + all of them to offer three would be absurd. + * the ladder is NOT cheap on the wire. Artifacts land under `kb///`, so the + rungs of one write share an S3 prefix — but the file MANIFEST is per (canonical_id, session_id), + and a rung that never called put_files reports file_count 0 and downloads nothing. Committing + once and expecting the coarse page to inherit it produces a bundle with no patch.diff, which is + how this was found. So every rung uploads, and transfer scales with ladder depth even though + storage does not. + +There is no delete. Every write here is permanent — the service exposes no DELETE for `/v1/kb/*`, +so a wrong canonical id is not a mistake that can be cleaned up afterwards, only one that can be +outranked. +""" + +import json +import os + +from kb.store_local import (CHAMPION_METRIC, KBStoreError, Candidate, finite_speedup, + safe_rel_path, validate_session_id) + +DEFAULT_SCAN = 25 # candidates hydrated before curation; well under the 200 rollup cap +ARTIFACT_KIND = "rewrite" + + +class RemoteKBStore(object): + """One producer's candidates under a canonical identity, over HTTP.""" + + def __init__(self, client, scan: int = DEFAULT_SCAN, metric: str = CHAMPION_METRIC, + promote_floor: float = 1.0): + self._client = client + self._scan = max(1, int(scan)) + # See LocalKBStore.__init__ for why this is a parameter. Remotely it is also a hard + # constraint rather than a preference: `sessions/top?metric=` reads one flat top-level + # `knowledge.` scalar and nothing else, so the ranking name has to be decided by + # whoever writes the document, not discovered afterwards. + self.metric = str(metric or CHAMPION_METRIC) + self.promote_floor = float(promote_floor) + self.root = str(getattr(client, "base_url", "") or "kb-store") + + @classmethod + def from_env(cls, scan: int = DEFAULT_SCAN, metric: str = CHAMPION_METRIC, + promote_floor: float = 1.0): + """Build from KB_STORE_URL / KB_STORE_TOKEN, or return (None, reason).""" + try: + from kb.store_client import KBStoreClient + except ImportError as e: + return None, "store_unavailable: " + str(e)[:120] + if not os.environ.get("KB_STORE_URL") or not os.environ.get("KB_STORE_TOKEN"): + return None, "no_credentials: KB_STORE_URL / KB_STORE_TOKEN are not both set" + try: + return cls(KBStoreClient.from_env(), scan, metric, promote_floor), "" + except Exception as e: + return None, "unusable_store: %s: %s" % (type(e).__name__, str(e)[:120]) + + # -- read ---------------------------------------------------------------------------- + + def identities(self): + """Always empty: the scheme supports exact lookup only. See the module docstring.""" + return [] + + def candidates(self, canonical_id: str, limit: int = 3): + """Rank this identity's candidates, hydrating each one's knowledge document. + + A 404 is an empty page, not an error: on a scheme with no search that is the ONLY signal + distinguishing "nothing recorded" from "recorded somewhere else", and the caller's ladder is + what turns it into a next attempt. + """ + try: + top = self._client.get_top_sessions(canonical_id, metric=self.metric, + limit=self._scan, offset=0) + except Exception: + return [] + rows = (top or {}).get("sessions") or [] + found = [] + for row in rows: + sid = str(row.get("session_id") or "") + if not sid: + continue + knowledge = self.get_session(canonical_id, sid) + if not isinstance(knowledge, dict): + continue + # Prefer the score the service computed; fall back to the document so a candidate is + # never dropped just because it was indexed under a different metric name. + speedup = finite_speedup(row.get("score")) + if speedup is None: + speedup = finite_speedup(knowledge.get(self.metric)) + found.append(Candidate(sid, knowledge, speedup, bool(row.get("is_champion")))) + found.sort(key=lambda c: (-(c.speedup if c.speedup is not None else float("-inf")), + c.session_id)) + return found[: max(0, int(limit))] if limit else found + + def get_session(self, canonical_id: str, session_id: str): + try: + record = self._client.get_session(canonical_id, validate_session_id(session_id)) + except Exception: + return None + if not isinstance(record, dict): + return None + # The envelope carries the producer's document under `knowledge`; a record fetched by id + # is already that document on some routes. Accept either rather than guess. + knowledge = record.get("knowledge") + return knowledge if isinstance(knowledge, dict) else record + + def champion(self, canonical_id: str): + """The identity's PROMOTED session, or {} when nothing has been promoted. + + Read from the rollup's `champion` field, NOT from get_best_record(). That endpoint answers + "the record to act on" — it falls back to the most recently selected record when no champion + exists — so using it here reports a brand-new page as already having a champion whose score + is the write that just landed. maybe_promote() then compares a candidate against itself, + never beats it, and no first champion is ever set. LocalKBStore returns {} for an + unpromoted identity and this plane has to mean the same thing by it. + """ + try: + rollup = self._client.get_rollup(canonical_id) + except Exception: + return {} + champion = (rollup or {}).get("champion") + if not isinstance(champion, dict) or not champion.get("session_id"): + return {} + return {"session_id": str(champion.get("session_id") or ""), + "metric": str(champion.get("metric") or self.metric), + "value": champion.get("value")} + + def champion_speedup(self, canonical_id: str): + champion = self.champion(canonical_id) + # Same guard as the local plane: an incumbent recorded under another metric is not a lower + # bar, it is an incomparable number, and ranking tokens-per-second against a ratio would + # either block every promotion or wave through every one. + if str(champion.get("metric") or "") != self.metric: + return None + return finite_speedup(champion.get("value")) + + def materialize(self, canonical_id: str, candidate, destination: str) -> str: + """Download one selected candidate as the standard bundle: recipe.json + files/. + + Same layout LocalKBStore.materialize() produces, so `_render_references` and the lane's + adopt step cannot tell the two planes apart. + """ + session_id = candidate.session_id if isinstance(candidate, Candidate) else str(candidate) + knowledge = (candidate.knowledge if isinstance(candidate, Candidate) + else self.get_session(canonical_id, session_id)) + if not isinstance(knowledge, dict): + raise KBStoreError("candidate knowledge is unreadable: " + session_id) + bundle = os.path.join(destination, validate_session_id(session_id)) + # download_session() lays out `/values.json` + `/files/` itself, so the + # destination is the BUNDLE, not the bundle's files/ directory. Handing it files_root nests + # a second files/ inside and every patch path the caller renders resolves to nothing. + files_root = os.path.join(bundle, "files") + os.makedirs(bundle, exist_ok=True) + try: + self._client.download_session(canonical_id, session_id, bundle) + except Exception as e: + # A knowledge document that names artifacts we could not fetch is worse than a loud + # failure: the lane would hand an agent a patch path that resolves to nothing. + raise KBStoreError("artifact download failed for %s: %s" % (session_id, str(e)[:160])) + # A successful download is not proof the bundle is usable. An identity whose manifest was + # never committed answers with values.json and nothing else, and the caller would only find + # out when an agent opened an empty patch path. Check what the document itself promised. + os.makedirs(files_root, exist_ok=True) # a record with no artifacts still gets the layout + promised = (knowledge.get("value") or {}).get("artifacts") + for name in sorted(set((promised or {}).values())) if isinstance(promised, dict) else []: + if not os.path.isfile(os.path.join(files_root, name)): + raise KBStoreError( + "%s declares %s but %s holds no manifest for this session — the record was " + "written without committing its files" % (session_id, name, canonical_id)) + recipe = dict(knowledge) + recipe.update({"canonical_id": canonical_id, "session_id": session_id, + "is_champion": bool(getattr(candidate, "is_champion", False)), + "champion": bool(getattr(candidate, "is_champion", False))}) + recipe.setdefault(self.metric, getattr(candidate, "speedup", None)) + with open(os.path.join(bundle, "recipe.json"), "w") as handle: + json.dump(recipe, handle, ensure_ascii=False, indent=2, sort_keys=True) + return bundle + + # -- write --------------------------------------------------------------------------- + + def write(self, canonical_id: str, session_id: str, knowledge, files=None) -> str: + """Record one candidate under an identity, uploading its artifacts at most once. + + `mode="replace"` so re-measuring the same patch updates that one session rather than + merging two documents into a chimera — the session id is a digest of the patch, so a repeat + genuinely IS the same candidate. + """ + if not isinstance(knowledge, dict): + raise KBStoreError("knowledge is not an object") + session_id = validate_session_id(session_id) + named = {safe_rel_path(rel): src for rel, src in (files or {}).items()} + try: + self._client.put_knowledge(canonical_id, knowledge, session_id=session_id, + mode="replace") + if named: + self._client.put_files(canonical_id, session_id, [ + (rel, named[rel], ARTIFACT_KIND, {}) for rel in sorted(named)]) + except Exception as e: + raise KBStoreError("%s: %s" % (type(e).__name__, str(e)[:160])) + return "%s/%s" % (canonical_id, session_id) + + def promote(self, canonical_id: str, session_id: str, speedup: float) -> None: + self._client.set_champion(canonical_id, validate_session_id(session_id), + metric=self.metric, value=float(speedup)) + + def maybe_promote(self, canonical_id: str, session_id: str, speedup) -> bool: + """Upstream's gate, verbatim: only a real win, and only over the incumbent.""" + speedup = finite_speedup(speedup) + if speedup is None or speedup <= self.promote_floor: + return False + incumbent = self.champion_speedup(canonical_id) + if incumbent is not None and speedup <= incumbent: + return False + try: + self.promote(canonical_id, session_id, speedup) + except Exception: + return False + return True + + +__all__ = ["ARTIFACT_KIND", "DEFAULT_SCAN", "RemoteKBStore"] diff --git a/kb/tests/test_attest.py b/kb/tests/test_attest.py new file mode 100644 index 000000000..125f23040 --- /dev/null +++ b/kb/tests/test_attest.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 +"""The attestation ledger: counting what happened when a record was actually tried on a box. + +The arithmetic is small; what these tests pin down is the two things that make the counters worth +reading. First, an attestation must NOT move a record — no ranking scalar, no champion — because +one box's failure is evidence and a retraction is a verdict, and collapsing the two would make the +command too dangerous to run automatically. Second, the ledger has to survive a rewrite: session +ids are content-addressed off the config, so re-recording the same configuration replaces the same +document, and a carry-forward that silently drops the history leaves a well-formed record claiming +nobody ever tried it. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +from kb.attest import (HISTORY_LIMIT, attest_session, attestation_ok, # noqa: E402 + attestations_of, attested_document, carry_attestations, + empty_attestations, record_attestation, retire_hint) +from kb.store_local import KBStoreError, LocalKBStore # noqa: E402 + + +def _store(tmp_path): + return LocalKBStore(str(tmp_path / "kb"), metric="speedup") + + +def _doc(speedup=1.5, **value): + body = {"direction": "tuned", "retained": True} + body.update(value) + return {"schema": "geak.e2e.v1", "speedup": speedup, "value": body} + + +def _seed(store, cid="geak:e2e:m", sid="s-1", speedup=1.5): + store.write(cid, sid, _doc(speedup), None) + store.promote(cid, sid, speedup) + return cid, sid + + +# -- the arithmetic ------------------------------------------------------------------------------- + + +def test_an_unattested_value_reads_as_an_empty_ledger(): + """Every field is safe to read off a record written before this module existed.""" + assert attestations_of({}) == empty_attestations() + assert attestations_of(None)["recalls"] == 0 + assert attestations_of({"attestations": "not a dict"})["validations"] == 0 + + +def test_a_counter_that_is_not_a_number_reads_as_zero(): + """A hand-edited document must degrade to 0, not crash the read path that hydrates a page.""" + ledger = attestations_of({"attestations": {"recalls": "seven", "validations": True, + "failures": -3, "not_reproduced": 2.9}}) + assert (ledger["recalls"], ledger["validations"]) == (0, 0) # True is a bool, not a count + assert (ledger["failures"], ledger["not_reproduced"]) == (0, 2) + + +@pytest.mark.parametrize("outcome,field", [("validated", "validations"), ("failed", "failures"), + ("not_reproduced", "not_reproduced")]) +def test_each_outcome_increments_recalls_and_its_own_counter(outcome, field): + value = record_attestation({}, outcome, actor="boxA") + ledger = value["attestations"] + assert ledger["recalls"] == 1 and ledger[field] == 1 + assert ledger["last_outcome"] == outcome and ledger["last_by"] == "boxA" + assert ledger["history"][-1]["outcome"] == outcome + + +def test_an_unknown_outcome_is_refused_rather_than_counted_as_something(): + with pytest.raises(KBStoreError): + record_attestation({}, "worked_i_think") + + +def test_evidence_is_whitelisted_onto_the_history_entry(): + """Both lanes' measurement spellings ride; anything unrecognized is dropped rather than grown + into the document, which is fetched once per candidate to rank a page.""" + value = record_attestation({}, "validated", evidence={ + "measured_tok_s": 1200, "measured_speedup": 1.8, "parity": "pass", + "gpu_serial": "leak me"}) + entry = value["attestations"]["history"][-1] + assert entry["measured_tok_s"] == 1200 and entry["measured_speedup"] == 1.8 + assert entry["parity"] == "pass" and "gpu_serial" not in entry + + +def test_history_is_bounded_but_the_counters_are_not(): + value = {} + for _ in range(HISTORY_LIMIT + 5): + value = record_attestation(value, "failed") + assert len(value["attestations"]["history"]) == HISTORY_LIMIT + assert value["attestations"]["recalls"] == HISTORY_LIMIT + 5 + + +def test_record_attestation_does_not_mutate_its_input(): + original = {"direction": "tuned"} + record_attestation(original, "validated") + assert original == {"direction": "tuned"} + + +# -- the retire hint ------------------------------------------------------------------------------ + + +def test_retire_hint_stays_silent_until_there_is_a_pattern(): + assert retire_hint({}) == "" + assert retire_hint(record_attestation({}, "failed")) == "" # one loss is not a case + twice = record_attestation(record_attestation({}, "failed"), "failed") + assert retire_hint(twice) == "" # losing is not failing + + +def test_two_non_reproductions_with_no_win_is_a_hint(): + value = record_attestation(record_attestation({}, "not_reproduced"), "not_reproduced") + assert "could not reproduce" in retire_hint(value) + + +def test_a_single_validation_clears_the_hint_however_many_failures(): + value = {} + for outcome in ("not_reproduced", "not_reproduced", "failed", "validated"): + value = record_attestation(value, outcome) + assert retire_hint(value) == "" + + +# -- carry-forward across a rewrite --------------------------------------------------------------- + + +def test_a_rewrite_carries_the_ledger_onto_the_replacing_document(): + previous = record_attestation({}, "validated") + fresh = carry_attestations(previous, {"direction": "tuned"}) + assert fresh["attestations"]["validations"] == 1 + + +def test_carrying_from_a_record_nobody_tried_adds_no_ledger(): + """An empty ledger and no ledger mean the same thing, and the shorter one does not imply + somebody looked.""" + assert "attestations" not in carry_attestations({"attestations": empty_attestations()}, {}) + assert "attestations" not in carry_attestations(None, {}) + + +# -- against a real store ------------------------------------------------------------------------- + + +def test_attesting_moves_no_score_and_no_champion(tmp_path): + """The whole reason this is not retraction. A record that failed once on one box keeps its + place in the ordering — the ledger is input to a later human decision, not the decision.""" + store = _store(tmp_path) + cid, sid = _seed(store, speedup=1.9) + report = attest_session(store, cid, sid, "not_reproduced", actor="boxA") + assert report["found"] and report["rewritten"] + after = store.get_session(cid, sid) + assert after["speedup"] == 1.9 # ranking scalar untouched + assert after["value"]["retained"] is True # not a tombstone + assert store.champion(cid)["session_id"] == sid # still crowned + assert after["value"]["attestations"]["not_reproduced"] == 1 + + +def test_attestations_accumulate_across_calls(tmp_path): + store = _store(tmp_path) + cid, sid = _seed(store) + for outcome in ("validated", "failed", "not_reproduced"): + attest_session(store, cid, sid, outcome) + ledger = store.get_session(cid, sid)["value"]["attestations"] + assert (ledger["recalls"], ledger["validations"]) == (3, 1) + assert (ledger["failures"], ledger["not_reproduced"]) == (1, 1) + + +def test_a_local_rewrite_keeps_the_artifacts_it_was_not_given(tmp_path): + """LocalKBStore.write() stages a fresh directory and swaps, so a rewrite that omitted the files + would DELETE them. This is the rule kb/retract.py:existing_files owns and this module reuses.""" + store = _store(tmp_path) + patch = tmp_path / "p.diff" + patch.write_text("--- a\n+++ b\n") + store.write("geak:e2e:m", "s-1", _doc(), {"final.patch": str(patch)}) + attest_session(store, "geak:e2e:m", "s-1", "validated") + assert store.session_files("geak:e2e:m", "s-1") == ["final.patch"] + + +def test_a_dry_run_shows_the_document_and_writes_nothing(tmp_path): + store = _store(tmp_path) + cid, sid = _seed(store) + report = attest_session(store, cid, sid, "validated", apply=False) + assert report["would_write"]["value"]["attestations"]["validations"] == 1 + assert "attestations" not in store.get_session(cid, sid)["value"] + + +def test_a_session_that_is_not_there_reports_it_rather_than_raising(tmp_path): + report = attest_session(_store(tmp_path), "geak:e2e:m", "nope", "validated") + assert report["found"] is False and "no such session" in report["error"] + assert report["rewritten"] is False + + +def test_an_unknown_outcome_against_a_store_reports_rather_than_raises(tmp_path): + store = _store(tmp_path) + cid, sid = _seed(store) + report = attest_session(store, cid, sid, "probably_fine") + assert report["found"] and not report["rewritten"] and "unknown attestation" in report["error"] + + +def test_attestation_ok_needs_the_record_found_somewhere(tmp_path): + """A rung the write never reached has nothing to count; finding it NOWHERE means the caller's + verdict was recorded against nothing at all.""" + assert attestation_ok([{"found": False}], True) is False + assert attestation_ok([{"found": False}, {"found": True, "rewritten": True}], True) is True + assert attestation_ok([{"found": True, "rewritten": False}], True) is False + assert attestation_ok([{"found": False}], False) is True # a dry run cannot fail + + +def test_attested_document_refuses_a_non_object(): + with pytest.raises(KBStoreError): + attested_document("not a document", "validated") diff --git a/kb/tests/test_store_local.py b/kb/tests/test_store_local.py new file mode 100644 index 000000000..e06252c83 --- /dev/null +++ b/kb/tests/test_store_local.py @@ -0,0 +1,283 @@ +"""Tests for the on-disk KB Store (kb/store_local.py). + +This plane exists to be swapped for the KernelForge service without changing behaviour, so what is +pinned here is the contract the service defines, not this implementation's conveniences: + - the address: the canonical id, split on ':', IS the directory path; + - the ranking: `speedup` descending and nothing else, with anything unrankable read as absent; + - the two write outcomes: a repeated session id updates one candidate, a new one appends; + - the champion gate: only a real win, and only over the incumbent; + - the cost model: `candidates()` reads knowledge documents and no artifact bytes. + +The last case cross-checks the tree against upstream's own reader when a KernelForge checkout is +importable, and skips otherwise — the point is to catch drift, not to add a dependency. +""" + +import json +import os +import sys + +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) +from kb.store_local import KBStoreError, LocalKBStore # noqa: E402 + +CID = "geak:kernel:geak:fused_moe_kernel:rocm:7.2:triton:gfx950" + + +def knowledge(speedup=2.0, direction="tile-retune", name="fused_moe_kernel"): + """The four-key document upstream writes; `value` is the producer's own and opaque here.""" + return {"producer": "geak", "speedup": speedup, + "identity": {"producer": "geak", "kernel_name": name, "framework": "rocm", + "framework_version": "7.2", "backend": "triton", "gpu": "gfx950"}, + "value": {"direction": direction, "kernel_name": name}} + + +def artifacts(tmp_path, tag="a", text="patch body\n"): + patch = tmp_path / f"{tag}.diff" + patch.write_text(text) + report = tmp_path / f"{tag}.md" + report.write_text(f"# report {tag}\n") + return {"patch.diff": str(patch), "report.md": str(report)} + + +# --------------------------------------------------------------------------- addressing + +def test_the_canonical_id_is_the_path(tmp_path): + store = LocalKBStore(tmp_path / "store") + store.write(CID, "geak-fused_moe_kernel-aaaa-bbbb", knowledge(), artifacts(tmp_path)) + session = (tmp_path / "store" / "geak" / "kernel" / "geak" / "fused_moe_kernel" / "rocm" / "7.2" + / "triton" / "gfx950" / "sessions" / "geak-fused_moe_kernel-aaaa-bbbb") + assert (session / "knowledge.json").is_file() + assert sorted(os.listdir(session / "files")) == ["patch.diff", "report.md"] + document = json.loads((session / "knowledge.json").read_text()) + assert sorted(document) == ["identity", "producer", "speedup", "value"] + assert store.identities() == [CID] + + +@pytest.mark.parametrize("bad", ["", "kernel", "kernel:GEAK:k", "kernel:geak:k:../etc", "kernel::k"]) +def test_an_unusable_identity_is_refused_not_normalized(tmp_path, bad): + with pytest.raises(KBStoreError): + LocalKBStore(tmp_path).identity_dir(bad) + + +@pytest.mark.parametrize("bad", ["../escape.diff", "/abs.diff", "a/../../b.diff", "a:b.diff", ""]) +def test_an_artifact_cannot_escape_its_session(tmp_path, bad): + source = tmp_path / "src.diff" + source.write_text("x") + with pytest.raises(KBStoreError): + LocalKBStore(tmp_path / "store").write(CID, "sid-1", knowledge(), {bad: str(source)}) + + +@pytest.mark.parametrize("bad", ["", ".hidden", "has space", "a" * 129, "sid/../x"]) +def test_an_unusable_session_id_is_refused(tmp_path, bad): + with pytest.raises(KBStoreError): + LocalKBStore(tmp_path / "store").write(CID, bad, knowledge(), {}) + + +# --------------------------------------------------------------------------- ranking + +def test_candidates_rank_on_speedup_alone(tmp_path): + store = LocalKBStore(tmp_path / "store") + for sid, speedup in [("sid-low", 1.2), ("sid-high", 4.5), ("sid-mid", 2.0)]: + store.write(CID, sid, knowledge(speedup=speedup), {}) + assert [c.session_id for c in store.candidates(CID, limit=0)] == ["sid-high", "sid-mid", "sid-low"] + assert [c.session_id for c in store.candidates(CID, limit=2)] == ["sid-high", "sid-mid"] + + +@pytest.mark.parametrize("unrankable", [None, True, "1.5", {"speedup": 9}, float("nan")]) +def test_a_speedup_the_store_cannot_rank_reads_as_absent(tmp_path, unrankable): + """`True` would sort as 1.0 and "1.5" would raise; both mean the producer wrote something else.""" + store = LocalKBStore(tmp_path / "store") + store.write(CID, "sid-real", knowledge(speedup=1.1), {}) + store.write(CID, "sid-odd", knowledge(speedup=unrankable), {}) + ranked = store.candidates(CID, limit=0) + assert [c.session_id for c in ranked] == ["sid-real", "sid-odd"] + assert ranked[1].speedup is None + + +def test_the_store_does_not_know_what_a_bench_key_is(tmp_path): + """Comparability is the caller's job — `resolve-remote` filters it, the store must not.""" + store = LocalKBStore(tmp_path / "store") + slow = knowledge(speedup=1.5) + slow["value"]["metric"] = {"bench_key": "b:imported"} + fast = knowledge(speedup=40.0) + fast["value"]["metric"] = {"bench_key": "b2:this-box"} + store.write(CID, "sid-imported", slow, {}) + store.write(CID, "sid-onbox", fast, {}) + assert [c.session_id for c in store.candidates(CID, limit=0)] == ["sid-onbox", "sid-imported"] + + +def test_ties_rank_the_same_way_on_every_read(tmp_path): + store = LocalKBStore(tmp_path / "store") + for sid in ("sid-c", "sid-a", "sid-b"): + store.write(CID, sid, knowledge(speedup=2.0), {}) + assert [c.session_id for c in store.candidates(CID, limit=0)] == ["sid-a", "sid-b", "sid-c"] + + +def test_a_cold_identity_is_empty_not_an_error(tmp_path): + store = LocalKBStore(tmp_path / "store") + assert store.candidates(CID, limit=3) == [] + assert store.champion(CID) == {} + assert store.champion_speedup(CID) is None + assert store.get_session(CID, "sid-missing") is None + assert store.session_files(CID, "sid-missing") == [] + + +def test_a_half_written_document_is_a_miss_not_a_crash(tmp_path): + store = LocalKBStore(tmp_path / "store") + store.write(CID, "sid-ok", knowledge(speedup=3.0), {}) + store.write(CID, "sid-broken", knowledge(speedup=9.0), {}) + (tmp_path / "store" / "geak" / "kernel" / "geak" / "fused_moe_kernel" / "rocm" / "7.2" / "triton" + / "gfx950" / "sessions" / "sid-broken" / "knowledge.json").write_text("{not json") + assert [c.session_id for c in store.candidates(CID, limit=0)] == ["sid-ok"] + + +def test_ranking_reads_no_artifact_bytes(tmp_path, monkeypatch): + """A 240KB patch must not be paid for until a candidate is actually selected.""" + store = LocalKBStore(tmp_path / "store") + store.write(CID, "sid-1", knowledge(speedup=2.0), artifacts(tmp_path, "a")) + store.write(CID, "sid-2", knowledge(speedup=3.0), artifacts(tmp_path, "b")) + + real_open = open + + def guard(path, *args, **kw): + assert "/files/" not in str(path), f"candidates() read an artifact: {path}" + return real_open(path, *args, **kw) + + monkeypatch.setattr("builtins.open", guard) + assert len(store.candidates(CID, limit=0)) == 2 + + +# --------------------------------------------------------------------------- write semantics + +def test_the_same_session_id_updates_one_candidate(tmp_path): + """Session ids are content-addressed upstream, so a remeasure must not grow the store.""" + store = LocalKBStore(tmp_path / "store") + store.write(CID, "sid-1", knowledge(speedup=2.0), artifacts(tmp_path, "a", "first\n")) + store.write(CID, "sid-1", knowledge(speedup=2.4), artifacts(tmp_path, "a2", "second\n")) + ranked = store.candidates(CID, limit=0) + assert len(ranked) == 1 and ranked[0].speedup == 2.4 + landed = store.materialize(CID, ranked[0], str(tmp_path / "out")) + with open(os.path.join(landed, "files", "patch.diff")) as handle: + assert handle.read() == "second\n" + + +def test_a_different_session_id_appends_under_the_same_key(tmp_path): + store = LocalKBStore(tmp_path / "store") + store.write(CID, "sid-1", knowledge(speedup=2.0), {}) + store.write(CID, "sid-2", knowledge(speedup=2.1), {}) + assert len(store.candidates(CID, limit=0)) == 2 + assert store.identities() == [CID] + + +def test_a_rewrite_drops_artifacts_the_new_record_no_longer_carries(tmp_path): + """The session lands as one unit, so a stale file must not survive underneath it.""" + store = LocalKBStore(tmp_path / "store") + store.write(CID, "sid-1", knowledge(), artifacts(tmp_path, "a")) + only_patch = {"patch.diff": artifacts(tmp_path, "b")["patch.diff"]} + store.write(CID, "sid-1", knowledge(), only_patch) + assert store.session_files(CID, "sid-1") == ["patch.diff"] + + +def test_knowledge_must_be_a_document(tmp_path): + with pytest.raises(KBStoreError): + LocalKBStore(tmp_path / "store").write(CID, "sid-1", ["not", "a", "document"], {}) + + +# --------------------------------------------------------------------------- champion + +def test_the_champion_moves_only_on_a_real_win_over_the_incumbent(tmp_path): + store = LocalKBStore(tmp_path / "store") + for sid, speedup in [("sid-loss", 0.9), ("sid-tie", 1.0), ("sid-win", 1.8), ("sid-worse", 1.4)]: + store.write(CID, sid, knowledge(speedup=speedup), {}) + assert store.maybe_promote(CID, "sid-loss", 0.9) is False # not faster than its own baseline + assert store.maybe_promote(CID, "sid-tie", 1.0) is False # a tie is not a win + assert store.champion(CID) == {} + assert store.maybe_promote(CID, "sid-win", 1.8) is True + assert store.maybe_promote(CID, "sid-worse", 1.4) is False # loses to the incumbent + assert store.champion_speedup(CID) == 1.8 + assert [c.session_id for c in store.candidates(CID, limit=0) if c.is_champion] == ["sid-win"] + + +@pytest.mark.parametrize("unrankable", [None, True, "2.0"]) +def test_an_unrankable_speedup_never_takes_the_champion_pointer(tmp_path, unrankable): + store = LocalKBStore(tmp_path / "store") + store.write(CID, "sid-1", knowledge(speedup=unrankable), {}) + assert store.maybe_promote(CID, "sid-1", unrankable) is False + assert store.champion(CID) == {} + + +def test_a_champion_written_under_another_metric_is_not_read_as_a_speedup(tmp_path): + store = LocalKBStore(tmp_path / "store") + store.write(CID, "sid-1", knowledge(), {}) + store.promote(CID, "sid-1", 3.0) + path = os.path.join(store.identity_dir(CID), "champion.json") + with open(path, "w") as handle: + json.dump({"session_id": "sid-1", "metric": "latency_ms", "value": 3.0}, handle) + assert store.champion_speedup(CID) is None + + +# --------------------------------------------------------------------------- materialize + +def test_materialize_lays_out_the_bundle_a_caller_can_apply(tmp_path): + store = LocalKBStore(tmp_path / "store") + store.write(CID, "sid-1", knowledge(speedup=2.5), artifacts(tmp_path, "a", "the patch\n")) + store.promote(CID, "sid-1", 2.5) + bundle = store.materialize(CID, store.candidates(CID, limit=1)[0], str(tmp_path / "cache")) + assert sorted(os.listdir(bundle)) == ["files", "recipe.json"] + assert open(os.path.join(bundle, "files", "patch.diff")).read() == "the patch\n" + recipe = json.loads(open(os.path.join(bundle, "recipe.json")).read()) + assert recipe["canonical_id"] == CID and recipe["session_id"] == "sid-1" + assert recipe["is_champion"] is True and recipe["speedup"] == 2.5 + + +def test_materialize_accepts_a_bare_session_id(tmp_path): + """The read path has a session id long before it has a Candidate object.""" + store = LocalKBStore(tmp_path / "store") + store.write(CID, "sid-1", knowledge(), artifacts(tmp_path, "a")) + bundle = store.materialize(CID, "sid-1", str(tmp_path / "cache")) + assert os.path.isfile(os.path.join(bundle, "files", "patch.diff")) + + +def test_materialize_refuses_a_candidate_that_is_not_there(tmp_path): + with pytest.raises(KBStoreError): + LocalKBStore(tmp_path / "store").materialize(CID, "sid-missing", str(tmp_path / "cache")) + + +def test_materializing_twice_replaces_rather_than_merges(tmp_path): + store = LocalKBStore(tmp_path / "store") + store.write(CID, "sid-1", knowledge(), artifacts(tmp_path, "a")) + dest = str(tmp_path / "cache") + store.materialize(CID, "sid-1", dest) + store.write(CID, "sid-1", knowledge(), {"patch.diff": artifacts(tmp_path, "b")["patch.diff"]}) + bundle = store.materialize(CID, "sid-1", dest) + assert os.listdir(os.path.join(bundle, "files")) == ["patch.diff"] + assert not [n for n in os.listdir(dest) if n.startswith(".")], "staging dirs must not be left behind" + + +# --------------------------------------------------------------------------- upstream parity + +def test_upstream_reads_the_tree_we_write(): + """Guards the one thing this file cannot prove on its own: that the shape is still theirs.""" + for candidate_src in ("/tmp/KernelForge/src",): + if os.path.isdir(candidate_src) and candidate_src not in sys.path: + sys.path.insert(0, candidate_src) + record_store = pytest.importorskip("kernel_agents.rewrite_by_flydsl.record_store") + import tempfile + + with tempfile.TemporaryDirectory() as tmp: + ours = LocalKBStore(os.path.join(tmp, "store")) + patch = os.path.join(tmp, "p.diff") + with open(patch, "w") as handle: + handle.write("the patch\n") + ours.write(CID, "sid-low", knowledge(speedup=1.5), {"patch.diff": patch}) + ours.write(CID, "sid-high", knowledge(speedup=3.5), {"patch.diff": patch}) + ours.promote(CID, "sid-high", 3.5) + + theirs = record_store.LocalRewriteRecords(os.path.join(tmp, "store")) + ranked = theirs.candidates(CID, limit=3) + assert [(c.session_id, c.speedup, c.is_champion) for c in ranked] == [ + ("sid-high", 3.5, True), ("sid-low", 1.5, False)] + assert theirs.champion_speedup(CID) == 3.5 + bundle = theirs.materialize(CID, ranked[0], os.path.join(tmp, "cache")) + assert (bundle / "files" / "patch.diff").read_text() == "the patch\n" diff --git a/kernel_workflow/kernel_lane.js b/kernel_workflow/kernel_lane.js index 0eee9a01b..443111910 100644 --- a/kernel_workflow/kernel_lane.js +++ b/kernel_workflow/kernel_lane.js @@ -9,6 +9,7 @@ export const meta = { { title: 'Benchmark', detail: 'benchmark_engineer builds the COMMANDMENT + baseline' }, { title: 'Profile', detail: 'profile_engineer classifies the bottleneck' }, { title: 'Research', detail: 'OPT-IN (args.dra_enabled): researcher fans research questions out in parallel via native WebSearch/WebFetch, writes a ranked-directions brief the planner seeds from' }, + { title: 'WarmStart', detail: 'search the experience KB (remote geak:kernel:* when credentialed, else local kb_artifacts/) for the best curated patch per optimization direction for this (kernel,language,gfx), validate each through the verify gate, adopt the first that passes [warm_start!=off]' }, { title: 'Optimize', detail: 'budget loop: tech_lead plans, specialist OR deep_explore engineers optimize, reprofile' }, { title: 'Verify', detail: 'each candidate patch independently re-benchmarked' }, { title: 'Merge', detail: 'integrator combines the round winners' }, @@ -178,6 +179,76 @@ const DRA_MAX_BLINDSPOTS = (() => { return Number.isFinite(v) && v >= 1 ? v : 4; })(); +// --------------------------------------------------------------------------- +// WARM-START (local experience KB). Before the optimize loop, search the machine-produced +// kb_artifacts/ store for the top-3 best patches for THIS (kernel, language, gfx), validate +// each through the SAME verify_engineer gate, and adopt the first that passes; after Validate, +// write this run's own win back. +// on (default) | read + validate top-3, ADOPT the first that passes. +// reference | read top-3 as prose only, never auto-apply. +// return_after_read | adopt then RETURN before the optimize loop. +// off/false/none, a STATE_DIR resume, or no arch => cold start (byte-identical to pre-feature). +const WARM_START = String(A.warm_start != null ? A.warm_start : 'on').trim().toLowerCase() || 'on'; +const WARM_START_ON = WARM_START !== 'off' && WARM_START !== 'false' && WARM_START !== 'none'; +const WARM_START_REF_ONLY = WARM_START === 'reference'; +const WARM_START_RETURN_AFTER = WARM_START === 'return_after_read'; +const KB_ARTIFACTS_DIR = String(A.kb_artifacts_dir || + (WORKFLOW_DIR ? WORKFLOW_DIR.replace(/\/[^/]*$/, '') + '/kb_artifacts' : '')).replace(/\/+$/, ''); +const EXPERIENCE_STORE = `${WORKFLOW_DIR}/scripts/experience_store.py`; +// Every candidate costs a full on-box verify to reject, so a recorded near-tie is not worth reading. +const WARM_START_MIN_SPEEDUP = Number.isFinite(parseFloat(A.warm_start_min_speedup)) + ? parseFloat(A.warm_start_min_speedup) : 1.05; +// How hard the resolver may work to map this run's kernel name onto a stored page. The name is +// layout-derived (`fused_moe_kernel` standalone, `fused_moe_kernel_task` from an e2e head extraction), +// so `exact` alone would make the head path miss its own history; `fuzzy` also accepts an op_kind. +const WARM_START_MATCH = ['exact', 'normalized', 'fuzzy'].includes(String(A.warm_start_match || '').trim()) + ? String(A.warm_start_match).trim() : 'fuzzy'; +// Which PLANE the experience comes from and goes back to. Same phases, same schemas, same verify +// gate either way — only the two command strings differ, because the store subcommands were built +// to print the same JSON as the directory ones. +// local (default) the curated kb_artifacts/ tree, keyed by slug. +// store a KB Store on disk in the shape the service uses, keyed by canonical id. +// This is the plane that later becomes the remote service, so a run in this +// mode is the rehearsal for it. +const KB_MODE = String(A.kb_mode || 'local').trim().toLowerCase() === 'store' ? 'store' : 'local'; +const KB_STORE_DIR = String(A.kb_store_dir || + (KB_ARTIFACTS_DIR ? KB_ARTIFACTS_DIR.replace(/\/[^/]*$/, '') + '/kb_store_local' : '')).replace(/\/+$/, ''); +// The key carries a rocm .; on a box without /opt/rocm the measured stack is empty and +// every record would file under `unspecified`, which splits one kernel's history across two keys. +const KB_FRAMEWORK_VERSION = String(A.kb_framework_version || '').trim(); +const KB_VERSION_FLAG = KB_FRAMEWORK_VERSION ? ` --framework-version ${JSON.stringify(KB_FRAMEWORK_VERSION)}` : ''; +const KB_ROOT_OK = KB_MODE === 'store' ? !!KB_STORE_DIR : !!KB_ARTIFACTS_DIR; +// Whether this run may also talk to the KB Store SERVICE, on top of whichever local plane KB_MODE +// selected. `auto` (default) uses it when credentials are present and silently does not when they +// are not; `off` restores the directory-only behaviour byte for byte. +const KB_REMOTE = String(A.kb_remote || 'auto').trim().toLowerCase() === 'off' ? 'off' : 'auto'; +// The credentials are NOT visible from this process. They live in the user's profile, and the +// shells these commands run in are non-interactive, so `process.env.KB_STORE_TOKEN` is empty here +// even on a box where the service is configured. Deciding the plane in JS would therefore mean +// deciding it wrong. Instead every emitted command exports its own credentials from the profile's +// two sources and then branches on what it actually got — the test happens where the answer is +// knowable. The token is read from a 0600 file into a variable, never passed in argv, because +// `ps` is world-readable on this box. +// The gateway's internal AMD CA is not in a stock container trust store, so a KB command run inside +// one fails TLS (the old workaround was `curl -k`). DETECT then heal: only when the caller has NOT +// already established trust (SSL_CERT_FILE unset) do we point urllib/requests/curl/node at the first +// readable AMD-root bundle we find. Most callers here run OUTSIDE the warm-start launcher (which sets +// these at `docker run`), so this self-heal is what keeps their KB reachable. Path-only, no CA +// content; overridable with KB_CA_BUNDLE; a no-op when SSL_CERT_FILE is already set or no bundle is +// readable (so CI and already-trusting images are byte-identical). DNS (the host has none +// in-container) is a launch concern, handled with `docker run --add-host`. +const KB_ENV_PRELUDE = + 'export KB_STORE_URL="${KB_STORE_URL:-https://global.primus-safe.amd.com/knowledge-base}"; ' + + 'export KB_STORE_TOKEN="${KB_STORE_TOKEN:-$(cat ~/.geak_kb_token 2>/dev/null)}"; ' + + 'if [ -z "${SSL_CERT_FILE:-}" ]; then for _ca in "${KB_CA_BUNDLE:-}" ' + + '/shared_nfs/hyperloom/ca/amd-ca-combined.pem "$HOME/amd-extra-ca-bundle.pem"; do ' + + '[ -n "$_ca" ] && [ -r "$_ca" ] && { export SSL_CERT_FILE="$_ca" REQUESTS_CA_BUNDLE="$_ca" ' + + 'CURL_CA_BUNDLE="$_ca" NODE_EXTRA_CA_CERTS="$_ca"; break; }; done; fi; '; +// Writing in store mode records BOTH planes in one call, so it needs both roots: the directory tree +// stays the source of truth a curation pass edits, and the store is derived from it. +const KB_WRITE_OK = KB_ROOT_OK && !!KB_ARTIFACTS_DIR; +const kebab = s => String(s || '').toLowerCase().replace(/[^a-z0-9]+/g, '-').replace(/^-+|-+$/g, '').slice(0, 60); + // --------------------------------------------------------------------------- // LEARNED-KNOWLEDGE SINK. Derived from WORKFLOW_DIR and nothing else, so a lane spawned by // e2e_workflow still curates into kernel_workflow/knowledge/learned/ — e2e's learned/ is a @@ -479,6 +550,47 @@ const VALIDATE_SCHEMA = obj({ arbitration_note: { type: 'string' }, final_patch: { type: 'string' }, }, ['director_verified_speedup_geomean', 'validation_status']); +// Warm-start resolver output = experience_store.py `resolve` JSON, verbatim. +const WARMSTART_RESOLVE_SCHEMA = obj({ + read_reason: { type: 'string' }, slug: { type: 'string' }, + // What this run's name derived to vs the page actually served, so a surprising match is visible. + requested_slug: { type: 'string' }, match_tier: { type: 'string' }, + // store mode only: the key the candidates came from. Declared rather than left to + // additionalProperties so the agent relaying this JSON has no reason to drop it. + canonical_id: { type: 'string' }, + // Same kernel, another language: the wrong target_language was passed, not an empty store. + other_language_pages: { type: 'array', items: { type: 'string' } }, + filtered: obj({ + total: { type: 'number' }, retired: { type: 'number' }, below_min_speedup: { type: 'number' }, + same_direction_collapsed: { type: 'number' }, + }), + candidates: { + type: 'array', + items: obj({ + rank: { type: 'number' }, slug: { type: 'string' }, speedup: { type: 'number' }, + exp_dir: { type: 'string' }, arch: { type: 'string' }, + patch_path: { type: 'string' }, prose_path: { type: 'string' }, + strategy: { type: 'string' }, status: { type: 'string' }, + // direction = the IDEA (one rank each). comparable=false: recorded on a different bench key + // than rank 1's, so the order between them is a prior only. + direction: { type: 'string' }, bench_key: { type: 'string' }, comparable: { type: 'boolean' }, + }, ['rank', 'patch_path']), + }, +}, ['read_reason']); +// Warm-start writer output = experience_store.py `write` JSON, verbatim. written=false with reason +// "duplicate_impl" is a success: the store already held this code and counted a reproduction. +const WARMSTART_WRITE_SCHEMA = obj({ + written: { type: 'boolean' }, reason: { type: 'string' }, slug: { type: 'string' }, + dir: { type: 'string' }, speedup: { type: 'number' }, + reproduced: { type: 'string' }, reproductions: { type: 'number' }, + // store mode only: where the same result landed on the KB plane. `replaced` says which of the + // two outcomes it was — a new candidate under the key, or this patch's own one measured again. + remote: obj({ + written: { type: 'boolean' }, reason: { type: 'string' }, canonical_id: { type: 'string' }, + session_id: { type: 'string' }, champion: { type: 'boolean' }, replaced: { type: 'boolean' }, + }), +}, ['written']); + // --------------------------------------------------------------------------- // Prompt helpers. Every agent reads its role file from WORKFLOW_DIR and the // relevant knowledge files itself; the script only passes paths + JSON inputs. @@ -832,7 +944,8 @@ let cumulative = 1.0; // best verified geomean speedup vs the TRUE baseli let bestSeen = 0; // best verified geomean of any candidate, committed or not let noImprove = 0; let bestPerCase = BASELINE_PER_CASE; -let finalWinner = null; // {geomean, arithmetic, per_case, patch, source} +let finalWinner = null; // {geomean, arithmetic, per_case, patch, source} — also set by a warm-start adopt +let roundsCommitted = 0; // rounds this run actually landed; a warm-start adopt is NOT one of them const history = { insights: [], ledger: [], rounds: [], bottleneck_now: profileSummary ? profileSummary.bottleneck : 'unknown', suggest_next: '' }; // One row per KB-seeded direction, joined against what the verifier measured (see the push site in // the round loop). Fed to update_experience and returned to the caller: a driver aggregating these @@ -852,7 +965,174 @@ if (setup.resumed && setup.prior_state) { log(`RESUMED from STATE_DIR: cumulative=${cumulative.toFixed(3)}x, ${history.insights.length} insights, ${history.ledger.length} ledger entries carried forward.`); } -while (dispatched < BUDGET && noImprove < MAX_NO_IMPROVE) { +// =========================================================================== +// PHASE: WarmStart. The recorded speedup only RANKS; adoption is decided by a fresh +// measurement through the verify gate here. gfx comes from the baseline profile's +// on-box `device` string (no extra probe). +// =========================================================================== +const GFX = (String((profileSummary && profileSummary.device) || '').match(/gfx\d+/i) || [''])[0].toLowerCase(); +let warm_start = { adopted: false, read_reason: WARM_START_ON ? 'read' : 'disabled', candidates: [] }; +let skipLoop = false; +if (WARM_START_ON && !setup.resumed && KB_ROOT_OK) { + phase('WarmStart'); + if (!GFX) { + warm_start.read_reason = 'missing_arch'; + log('[kb] warm-start skipped: no gfx detected from the baseline profile device string.'); + } else { + // The two planes take a different root flag and a different name→page rule (the store addresses + // by canonical id, so there is nothing to match fuzzily), and print the same JSON. + const localResolveCmd = KB_MODE === 'store' + ? `resolve-remote --plane local --store ${JSON.stringify(KB_STORE_DIR)}${KB_VERSION_FLAG}` + : `resolve --root ${JSON.stringify(KB_ARTIFACTS_DIR)} --match ${WARM_START_MATCH}`; + const commonArgs = + `--kernel-name ${JSON.stringify(KERNEL_NAME)} --language ${JSON.stringify(TARGET_LANGUAGE)} \\\n` + + ` --gfx ${GFX} --top-n 3 --min-speedup ${WARM_START_MIN_SPEEDUP} \\\n` + + ` --refs-dir ${JSON.stringify(EVAL_DIR + '/kb_references')}`; + // Remote first, local curated tree as the fallback. The service is the shared plane and should + // win when it has an answer, but it is still filling up, while `kb_artifacts/` holds a hand- + // curated history (retired entries, one entry per direction) that a thin remote page must not + // shadow. An empty remote answer is indistinguishable from a 404 on this scheme, so "no + // candidates" — not "no error" — is what triggers the second read. Both reads are seconds and + // no GPU; the thing they protect against is a cold start that costs hours. + const resolveScript = KB_REMOTE === 'off' + ? `python3 ${JSON.stringify(EXPERIENCE_STORE)} ${localResolveCmd} \\\n ${commonArgs}` + : `${KB_ENV_PRELUDE} +REMOTE_OUT='' +if [ -n "$KB_STORE_TOKEN" ]; then + REMOTE_OUT=$(python3 ${JSON.stringify(EXPERIENCE_STORE)} resolve-remote --plane remote${KB_VERSION_FLAG} \\ + ${commonArgs} 2>/dev/null || true) + if printf '%s' "$REMOTE_OUT" | python3 -c 'import json,sys; sys.exit(0 if (json.load(sys.stdin).get("candidates") or []) else 1)' 2>/dev/null; then + printf '%s\\n' "$REMOTE_OUT"; exit 0 + fi +fi +python3 ${JSON.stringify(EXPERIENCE_STORE)} ${localResolveCmd} \\ + ${commonArgs}`; + const resolved = await agentT( + `You are the warm-start resolver. Run EXACTLY this ${KB_REMOTE === 'off' ? 'command' : 'script'} ` + + `and return its single-line JSON stdout verbatim as StructuredOutput — do not add, drop, reorder, ` + + `or reinterpret any field. ` + + (KB_REMOTE === 'off' ? '' : + `It tries the shared KB Store service first and falls back to the on-disk knowledge base by ` + + `itself; run it as one script, do not split it into separate commands. `) + + `An empty candidate list is a VALID answer meaning "nothing recorded for this kernel yet". Do not ` + + `retry with different flags and do not invent candidates: +\`\`\`bash +${resolveScript} +\`\`\``, + { phase: 'WarmStart', label: 'warm_start:resolve', schema: WARMSTART_RESOLVE_SCHEMA }) || {}; + warm_start.read_reason = resolved.read_reason || 'read'; + warm_start.slug = resolved.slug || ''; + warm_start.match_tier = resolved.match_tier || ''; + warm_start.filtered = resolved.filtered || null; + const cands = Array.isArray(resolved.candidates) ? resolved.candidates : []; + warm_start.candidates = cands.map(c => ({ + rank: c.rank, slug: c.slug, speedup: c.speedup, direction: c.direction || '', status: 'read', + })); + const f = resolved.filtered || {}; + // Which plane actually answered. Only the key-addressed subcommand emits `canonical_id`, so its + // presence separates a store read from a slug-tree read; in `local` KB_MODE the only key- + // addressed reader in the script is the remote one, so that is also the remote/fallback tell. + // Worth logging either way: the canonical id is the address, and on a scheme with no search a + // thin answer and a mis-keyed question look identical from the outside. + warm_start.plane = resolved.canonical_id ? (KB_MODE === 'store' ? 'store' : 'remote') : 'local'; + if (resolved.canonical_id) log(`[kb] plane=${warm_start.plane} key=${resolved.canonical_id}`); + else if (KB_REMOTE !== 'off') log('[kb] plane=local (service had no candidates, or no credentials)'); + log(`[kb] experience read: slug=${resolved.slug || '?'} (${resolved.match_tier || 'exact'} match of ` + + `${resolved.requested_slug || KERNEL_NAME}) reason=${warm_start.read_reason} ` + + `candidates=${cands.length}${f.total ? ` of ${f.total} recorded [${f.retired || 0} retired, ` + + `${f.below_min_speedup || 0} below ${WARM_START_MIN_SPEEDUP}x, ` + + `${f.same_direction_collapsed || 0} same-direction]` : ''}`); + const otherLangs = Array.isArray(resolved.other_language_pages) ? resolved.other_language_pages : []; + if (!cands.length && otherLangs.length) { // wrong target_language, not an empty store + log(`[kb] NOTE: no ${TARGET_LANGUAGE} page, but the store holds ${otherLangs.join(', ')} — ` + + `check the target_language this lane was invoked with.`); + } + // reference-only: prose is already mirrored to EVAL_DIR/kb_references by the resolver; do not apply. + if (!WARM_START_REF_ONLY) { + const editableCsv = ((analysis && analysis.modifiable_files) || []).join(','); + for (const c of cands) { // already rank-ordered (fastest first) + const rec = warm_start.candidates.find(x => x.rank === c.rank); + const remapOut = `${EVAL_DIR}/warm_start/cand_${c.rank}/remapped.diff`; + const ver = await agentT( + roleAgent('verify_engineer', 'verify', + 'Validate a HISTORICAL warm-start patch — the SAME gate as any round candidate. The patch was ' + + 'won in the workspace that produced it, whose paths need not exist here, so FIRST run EXACTLY:\n' + + '```bash\n' + + `mkdir -p ${JSON.stringify(`${EVAL_DIR}/warm_start/cand_${c.rank}`)}\n` + + `python3 ${JSON.stringify(EXPERIENCE_STORE)} remap --patch ${JSON.stringify(c.patch_path)} \\\n` + + ` --out ${JSON.stringify(remapOut)} --editable ${JSON.stringify(editableCsv)} \\\n` + + ` --workspace ${JSON.stringify(CANONICAL)}\n` + + '```\n' + + 'It prints one line of JSON. remapped=true: apply the REMAPPED patch (its paths were rewritten ' + + 'onto this workspace). reason "no_change_needed": apply the original PATCH. Any other reason ' + + '(notably "unmapped_paths" — the patch touches files this workspace does not have): return ' + + 'status:"apply_failed", notes "patch_outside_editable_set", and DO NOT apply or benchmark. ' + + 'Apply with `git apply "$P" || git apply --3way "$P"`. On ANY failure restore the working copy ' + + 'fully (git checkout -- . and delete untracked files) before returning so the next candidate ' + + 'starts from a clean tree.', { + CANONICAL, PATCH: c.patch_path, REMAP_OUT: remapOut, + VERIFY_DIR: `${EVAL_DIR}/warm_start/cand_${c.rank}`, + EDITABLE_SET: (analysis && analysis.modifiable_files) || [], + GPU_ID: GPU_LIST[0], SKILL_DIR: WORKFLOW_DIR, COMMANDMENT, BASELINE_PER_CASE, + ...(HARNESS_ADDENDUM ? { HARNESS_ADDENDUM } : {}), + ...(REQUIRE_GRAPH_CAPTURE ? { REQUIRE_GRAPH_CAPTURE: '1' } : {}), + }), + { phase: 'WarmStart', label: `warm_start:verify c${c.rank}`, schema: VERIFY_SCHEMA }); + const sp = primSpeedup(ver); + if (ver && says(ver.status, 'verified') && says(ver.correctness, 'pass') && sp > 1.0) { + const adopt = await agentT( + `You are the TechLead adopting a validated warm-start patch into the canonical workspace. +\`\`\`bash +export GIT_PAGER=cat GIT_TERMINAL_PROMPT=0 GIT_EDITOR=true +cd ${CANONICAL} +git checkout -- . +# the remapped patch when verify produced one (its paths fit THIS workspace), else the stored one +P=${JSON.stringify(remapOut)}; [ -f "$P" ] || P=${JSON.stringify(c.patch_path)} +git apply "$P" || git apply --3way "$P" +git -c user.email=team@workflow -c user.name=team add -A +git -c user.email=team@workflow -c user.name=team commit -q -m "warm-start adopt: ${c.slug} (${sp.toFixed(2)}x)" +git --no-pager diff "$(git rev-list --max-parents=0 HEAD)..HEAD" > ${EVAL_DIR}/current_best.diff +\`\`\` +If BOTH applies fail, apply manually to match intent, then add -A + commit and RE-RUN the COMMANDMENT +correctness check; only report committed=true if it still passes. Return JSON {committed, current_best_diff, note}.`, + { phase: 'WarmStart', label: `warm_start:adopt c${c.rank}`, schema: COMMIT_SCHEMA }); + if (adopt && adopt.committed) { + // The optimize loop now builds ON this patch: cumulative starts at the adopted Nx, so this + // run's own rounds only earn the delta above it (split out as incremental_speedup below). + cumulative = sp; + bestPerCase = (ver.per_case && ver.per_case.length) ? ver.per_case : bestPerCase; + finalWinner = { source: `warm_start:${c.slug}`, geomean: sp, + arithmetic: ver.verified_arithmetic || sp, per_case: bestPerCase, patch: c.patch_path }; + if (bestSeen < sp) bestSeen = sp; + warm_start.adopted = true; + warm_start.adopted_speedup = sp; + warm_start.slug = c.slug; + // Carried into the write-back: the idea this run started from, and the entry it + // descends from — lineage a later curation pass cannot recover from the diff alone. + warm_start.direction = c.direction || ''; + warm_start.exp_dir = c.exp_dir || ''; + if (rec) rec.status = 'adopted'; + log(`[kb] warm-start ADOPTED ${c.slug} @ ${sp.toFixed(2)}x — optimizing from the patched state.`); + profileSummary = await agentT( + roleAgent('profile_engineer', 'reprofile', + 'Re-profile the adopted warm-start state and classify the new bottleneck.', { + WORKSPACE: CANONICAL, EVAL_DIR, SKILL_DIR: WORKFLOW_DIR, GPU_ID: GPU_LIST[0], ROUND: 0, + COMMANDMENT, PREVIOUS_METRICS: profileSummary, + }), + { phase: 'WarmStart', label: 'warm_start:reprofile', schema: PROFILE_SCHEMA }) || profileSummary; + if (history) history.bottleneck_now = profileSummary ? profileSummary.bottleneck : history.bottleneck_now; + if (WARM_START_RETURN_AFTER) { skipLoop = true; warm_start.returned_after_read_kb = true; } + break; + } + } + if (rec && rec.status !== 'adopted') rec.status = (ver && ver.status) ? `rejected:${ver.status}` : 'rejected:apply_failed'; + log(`[kb] warm-start candidate c${c.rank} ${rec ? rec.status : 'rejected'} (${sp ? sp.toFixed(2) + 'x' : 'no measure'}).`); + } + } + } +} + +while (!skipLoop && dispatched < BUDGET && noImprove < MAX_NO_IMPROVE) { // --- (0) HARD STOP: no new round may START past the deadline ---------- // Checked BEFORE round++ so an expired check does not inflate the round count. The in-flight round // is never interrupted — a killed round leaves a half-verified patch and no report. @@ -1129,6 +1409,7 @@ re-check is not required.) Return JSON {committed, current_best_diff, note}.`, cumulative = winner.geomean; bestPerCase = winner.per_case && winner.per_case.length ? winner.per_case : bestPerCase; finalWinner = winner; + roundsCommitted += 1; // --- (f) Re-profile the new best ------------------------------------ profileSummary = await agentT( @@ -1259,6 +1540,89 @@ log(`COMPLETE. ${KERNEL_NAME}: verified ${HAS_WORKLOAD ? 'time-weighted' : 'geom `${HAS_WORKLOAD && Number.isFinite(finalGeomean) ? ` (unweighted geomean ${finalGeomean.toFixed(2)}x)` : ''}` + ` (status ${validation ? validation.validation_status : '?'}). Results in ${EVAL_DIR}`); +// =========================================================================== +// Write this run's outcome back to kb_artifacts/ — the producer half of the loop. +// The script applies its own gate and never fails the run; the `finalPrimary > 1.0` +// pre-check just avoids spending an agent on a run that cannot pass the gate anyway. +// =========================================================================== +let kb_written = null; +if (KB_WRITE_OK && GFX && Number.isFinite(finalPrimary) && finalPrimary > 1.0) { + const kernelClass = (analysis && analysis.kernel_type) || 'unknown'; + const finalPatch = report ? report.final_patch : `${EVAL_DIR}/final_patch.diff`; + const reportPath = report && report.report_path ? report.report_path : `${EVAL_DIR}/tech_lead_report.md`; + // What the store cannot infer from a diff but needs to curate this entry later. `direction` is the + // IDEA, and only a warm start supplies a real one: a round title is free-form per run, so kebabbing + // it would mint a unique label every time and group nothing. Unlabeled is honest — resolve treats + // such an entry as its own direction, and a curation pass assigns the shared label later. + const winnerDirection = kebab(warm_start.direction || ''); + const metricKind = HAS_WORKLOAD ? 'time_weighted' : 'geomean'; + // Commas separate the list, so a case name may not contain one. + const caseNames = (bestPerCase || []).map(c => c && String(c.name || '').replace(/,/g, ';')) + .filter(Boolean).join(','); + // write-remote runs the directory write first and files the same entry in the store, so the two + // planes cannot drift; its extra `remote` field rides along under the schema's open object. + // `--plane both` extends that ordering one hop further: directory tree, then local store, then the + // service. Unlike the read there is no branch on credentials here, because `_open_plane` already + // makes the right call — a `both` without them degrades to the local planes and reports + // `remote_unavailable` rather than failing, which is the behaviour we would have written by hand. + const remoteWriteOn = KB_REMOTE !== 'off' && !!KB_STORE_DIR; + const writeCmd = remoteWriteOn + ? `write-remote --plane both --store ${JSON.stringify(KB_STORE_DIR)}${KB_VERSION_FLAG}` + : KB_MODE === 'store' + ? `write-remote --plane local --store ${JSON.stringify(KB_STORE_DIR)}${KB_VERSION_FLAG}` + : 'write'; + kb_written = await agentT( + `You are the experience writer. Run EXACTLY this command (it applies its own gates and prints a ` + + `single-line JSON) and return that JSON verbatim as StructuredOutput. If the command errors, return ` + + `{"written": false, "reason": "io_error"}` + + (remoteWriteOn + ? ` — do NOT retry it and do NOT adjust its arguments. It may reach the shared KB Store service, ` + + `and every write that service accepts is PERMANENT (it exposes no delete), so a second attempt ` + + `creates a second permanent record instead of fixing the first.` + : `.`) + ` +\`\`\`bash +${remoteWriteOn ? KB_ENV_PRELUDE + '\n' : ''}python3 ${JSON.stringify(EXPERIENCE_STORE)} ${writeCmd} --root ${JSON.stringify(KB_ARTIFACTS_DIR)} \\ + --kernel-name ${JSON.stringify(KERNEL_NAME)} --language ${JSON.stringify(TARGET_LANGUAGE)} \\ + --gfx ${GFX} --kernel-class ${JSON.stringify(kernelClass)} \\ + --speedup ${finalPrimary} --baseline-wall-ms ${BASELINE_GEOMEAN_MS} \\ + --patch ${JSON.stringify(finalPatch)} --eval-dir ${JSON.stringify(EVAL_DIR)} \\ + --report ${JSON.stringify(reportPath)} --metric-kind ${metricKind} \\ + --direction ${JSON.stringify(winnerDirection)} --case-names ${JSON.stringify(caseNames)}\ +${warm_start.exp_dir ? ` \\\n --parent ${JSON.stringify(warm_start.exp_dir)}` : ''} +\`\`\``, + { phase: 'Validate', label: 'kb:write', schema: WARMSTART_WRITE_SCHEMA }); + const remoteWrite = (kb_written && kb_written.remote) || null; + if (remoteWrite) { + log(remoteWrite.written + // `replaced` distinguishes the two write outcomes the store has: a NEW patch appends a session + // under the same key, the SAME patch remeasured lands back on its own content-addressed one. + ? `[kb] plane=store ${remoteWrite.replaced ? 'updated' : 'appended'} ` + + `${remoteWrite.canonical_id}/${remoteWrite.session_id}${remoteWrite.champion ? ' (champion)' : ''}` + : `[kb] plane=store not written: ${remoteWrite.reason || 'unknown'}`); + } + log(kb_written && kb_written.written + ? `[kb] experience written: ${kb_written.slug} (speedup ${finalPrimary.toFixed(2)}, direction ` + + `${winnerDirection || 'unlabeled'})` + : kb_written && kb_written.reason === 'duplicate_impl' + // Not a failure — the expected outcome of adopting a warm start and not beating it. + ? `[kb] experience already known — counted as reproduction #${kb_written.reproductions || '?'} of ` + + `${kb_written.reproduced || kb_written.dir}` + : `[kb] experience not written: ${kb_written ? kb_written.reason : 'writer returned nothing'}`); +} + +// finalPrimary is the total vs the pristine baseline; when a warm-start patch was adopted, split out +// the delta ABOVE it so a KB-derived gain is never reported as this run's own work. +// +// With no round committed after the adopt, the workspace still holds exactly the adopted code, and +// the two numbers are two measurements of the SAME thing taken at different times — the ratio is +// bench noise (~2% on this box; `_gemm_a16_w16_kernel` reported 0.98 and read as a regression). +// Report 1.0 and say why, rather than dressing the noise up as a delta. +const noRoundsAfterAdopt = !!(warm_start.adopted && roundsCommitted === 0); +const incrementalSpeedup = warm_start.adopted && warm_start.adopted_speedup + ? (noRoundsAfterAdopt ? 1.0 + : Number.isFinite(finalPrimary) ? finalPrimary / warm_start.adopted_speedup : null) + : finalPrimary; + // =========================================================================== // PHASE: UpdateExperience — distill ONE reusable card into this workflow's knowledge/learned/ // (see knowledge/learned/README.md). Runs ONCE, at the very end: never per round, because only @@ -1283,7 +1647,6 @@ if (kbGate) log(`[kb] not distilling: ${kbGate}.`); // run did not earn. Reported in review of #411. const kbAccepted = String((validation && validation.validation_status) || '').toLowerCase() === 'accepted'; if (!kbGate && UPDATE_EXPERIENCE_ON && kbAccepted && Number.isFinite(finalPrimary) && finalPrimary > 1.0) { - const GFX = (String((profileSummary && profileSummary.device) || '').match(/gfx\d+/i) || [''])[0].toLowerCase(); try { learned_card = await agentT( roleAgent('update_experience', 'Validate', @@ -1369,7 +1732,27 @@ return { budget_total: BUDGET, report_path: report ? report.report_path : `${EVAL_DIR}/tech_lead_report.md`, final_patch: report ? report.final_patch : `${EVAL_DIR}/final_patch.diff`, - // null when the step was off, or the run earned nothing worth a card. + // Warm-start (local experience KB) outcome. adopted=false + read_reason on a cold start. + warm_start: { + adopted: warm_start.adopted, + read_reason: warm_start.read_reason, + slug: warm_start.slug || '', + // How the store was reached and what it filtered out: tells a genuine cold start from a + // name/language mismatch that only looks like an empty KB. + match_tier: warm_start.match_tier || '', + filtered: warm_start.filtered || null, + adopted_speedup: warm_start.adopted ? warm_start.adopted_speedup : null, + total_speedup: Number.isFinite(finalPrimary) ? finalPrimary : null, + incremental_speedup: incrementalSpeedup, + // true = the run adopted a stored patch and never improved on it; incremental is 1.0 by + // definition, not by measurement. + no_rounds_after_adopt: noRoundsAfterAdopt, + rounds_committed: roundsCommitted, + incremental_improved: !!(warm_start.adopted && Number.isFinite(incrementalSpeedup) && incrementalSpeedup > 1 + MIN_IMPROVE), + returned_after_read_kb: !!warm_start.returned_after_read_kb, + candidates: warm_start.candidates, + }, + kb_written: kb_written && kb_written.written ? { slug: kb_written.slug, dir: kb_written.dir } : null, // Which condition actually ended the loop. RECORDED, not inferred: on the previous campaign this // existed only in the driver's return value, which the caller sees after ALL its kernels finish — // so a parent-process death erased it for every kernel already done, and the realised window had @@ -1398,6 +1781,7 @@ return { } return n ? { distinct: seen.size, issued: n, ratio: +(seen.size / n).toFixed(3) } : null; })(), + // null when the step was off, or the run earned nothing worth a card. learned_card: learned_card && learned_card.card_path ? { action: learned_card.action || '', card_path: learned_card.card_path, key: learned_card.key || '' } : null, diff --git a/kernel_workflow/kernel_workflow.js b/kernel_workflow/kernel_workflow.js index 3c75caf64..b22ff8407 100644 --- a/kernel_workflow/kernel_workflow.js +++ b/kernel_workflow/kernel_workflow.js @@ -76,6 +76,19 @@ const EXPERT_SKILLS_DIR = String(A.expert_skills_dir || (KERNEL_KNOWLEDGE_DIR ? KERNEL_KNOWLEDGE_DIR + '/expert_skills' : '')).replace(/\/+$/, ''); const EXPERT_SKILL_ROLES = new Set(['op_benchmarker']); +// Warm-start experience KB. Passed to each lane explicitly (the bakeoff lane invocation spreads +// specific keys, not ...A) so every language lane reads/writes its own ____ slug. +const WARM_START = String(A.warm_start != null ? A.warm_start : 'on').trim().toLowerCase() || 'on'; +const KB_ARTIFACTS_DIR = String(A.kb_artifacts_dir || + (WORKFLOW_DIR.replace(/\/[^/]*$/, '') + '/kb_artifacts')).replace(/\/+$/, ''); +// Plane selection, forwarded the same way and for the same reason: every bakeoff lane must read and +// write the plane the run was launched with, not each its own default. +const KB_PLANE_ARGS = { + ...(A.kb_mode != null ? { kb_mode: String(A.kb_mode) } : {}), + ...(A.kb_store_dir != null ? { kb_store_dir: String(A.kb_store_dir) } : {}), + ...(A.kb_framework_version != null ? { kb_framework_version: String(A.kb_framework_version) } : {}), +}; + // --------------------------------------------------------------------------- // Schema helpers. // --------------------------------------------------------------------------- @@ -367,6 +380,7 @@ const results = await Promise.all(lanes.map(l => sem.with(1, async ([gpu]) => { exp_root: `${EVAL_DIR}/bakeoff/${l.key}`, use_expert_skills: USE_EXPERT_SKILLS ? 'true' : 'false', expert_skills_dir: EXPERT_SKILLS_DIR, perf_knowledge_dir: KERNEL_KNOWLEDGE_DIR, + warm_start: WARM_START, kb_artifacts_dir: KB_ARTIFACTS_DIR, ...KB_PLANE_ARGS, // Forward the KB switch. This arg object is explicit (the optimize/author path spreads {...A}, // this one does not), so anything omitted here silently reverts to the lane's default — a // caller asking for a KB-off bakeoff would have got eight KB-on lanes and no error. diff --git a/kernel_workflow/roles/tech_lead.md b/kernel_workflow/roles/tech_lead.md index f03c4d214..cba9fa316 100644 --- a/kernel_workflow/roles/tech_lead.md +++ b/kernel_workflow/roles/tech_lead.md @@ -410,7 +410,21 @@ table, `BASELINE_TIMING`, and `BASELINE_GEOMEAN_MS`. - **Final per-test-case table** (baseline ms / optimized ms / speedup; + `count` & weight-share when workload-aligned) + geomean + arithmetic + the time-weighted speedup. - **Key optimizations applied** (what + impact). - - **What didn't work** (dead-ends from the ledger). + - **What didn't work** (dead-ends from the ledger). End this section with the machine-readable + block below, in ADDITION to your prose — it is what the experience store keeps, so the next run + on this kernel does not spend a round re-funding a direction you already closed. One entry per + closed direction; `measured` is the number you actually observed, and if you did not measure it, + say so in `mechanism` instead of inventing a figure. Omit the block entirely if nothing was + closed with evidence — an empty block is worse than none. + + ```` + + ```yaml + - idea: use_buffer_ops=OFF negative control + measured: 0.883x + mechanism: the ambient default is load-bearing, -11.7% + ``` + ```` Return JSON: ```json diff --git a/kernel_workflow/scripts/experience_store.py b/kernel_workflow/scripts/experience_store.py new file mode 100755 index 000000000..a4c1db5ee --- /dev/null +++ b/kernel_workflow/scripts/experience_store.py @@ -0,0 +1,1961 @@ +#!/usr/bin/env python3 +"""Local experience store for the kernel workflow — the machine-produced KB that carries the diff. + +Self-contained, stdlib + PyYAML only, so a lane agent can call it over Bash. On-disk layout +(rooted at --root, default /kb_artifacts): + + ///// + meta.yaml # identity + metric + prose pointers + patch.diff # the winning diff (verbatim copy) + report.md # optional tech_lead report, copied for prose + + slug = ____ # deterministic; read and write derive it identically + +Subcommands: + write Store one measured win behind the gate (missing_arch / no_improvement / empty_diff). + resolve Rank the top-N same-gfx solutions for a slug and mirror their prose into . + remap Rewrite a stored patch's paths onto the calling workspace's layout, or refuse and say why. + languages Which languages a kernel has a page in — the store, not a task_type guess, decides. + backfill-content + Bring imported entries up to the current content shape (dry-run unless --apply). + export-remote + Render entries as KB Store candidates (one JSON line each); uploads nothing. + resolve-remote + `resolve`, but addressed by canonical id against a KB store (kb/store_local.py). + write-remote + `write`, landing the same result in the local store AND under its key. + attest / attest-remote + Count one attempt to USE a stored entry (validated | failed | not_reproduced), so a + later curation pass can retire what nobody can reproduce. Moves no speedup, no rank. + +Speedups only compare within one GPU arch, so resolve drops cross-arch entries outright. Neither +command ever raises: on failure it prints a JSON reason and exits 0 so the caller degrades. + +resolve serves a CURATED top-N, not the raw speedup order: entries the curation retired +(`retained: false`) are never offered, near-ties below `--min-speedup` are not worth a verify slot, +and only one entry per `direction:` is ranked (same-idea runners-up ride along as `alternates`, +since they verify or fail together). A speedup only means something against its own +`metric.bench_key`, so each candidate carries one plus a `comparable` flag against rank 1. +""" + +import argparse +import hashlib +import json +import os +import re +import sys +import tempfile +import time + +# The shared KB plane lives at the repo root as the `kb` package, not beside this file. Executed as +# a CLI from an arbitrary cwd, so the root is derived from __file__ and never from the environment. +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +from kb.attest import OUTCOMES as _OUTCOMES +from kb.curate import collapse_by_direction +from kb.ladder import publish +from kb.plane import open_plane +from kb.store_local import CHAMPION_METRIC + +try: + import yaml +except Exception: # yaml ships in this env; degrade to json-only meta + yaml = None + + +# Identity: read and write MUST derive the slug identically (never via an LLM) or a run +# can never find its own lineage. +def _safe(seg: str) -> str: + """Slug-safe a path segment: keep [A-Za-z0-9._-], collapse the rest to '-'.""" + s = re.sub(r"[^A-Za-z0-9._-]+", "-", str(seg or "").strip()) + s = s.strip("-.") or "x" + return s[:80] + + +def _norm_gfx(gfx: str) -> str: + m = re.search(r"gfx\d+", str(gfx or ""), re.IGNORECASE) + return m.group(0).lower() if m else "" + + +# One kernel is named differently per layout: `fused_moe_kernel` (kernel dir), `fused_moe_kernel_task` +# (e2e head extraction), `triton_fused_moe_kernel.py` (language in the filename). Canonicalizing on +# BOTH sides is what lets a head run find, and extend, its own lineage instead of forking a new page. +_NAME_PREFIXES = ("triton_", "hip_", "ck_", "cuda_", "torch_") +_NAME_SUFFIXES = (".py", ".hip", ".cu", ".cpp", "_task") + + +def canon_name(kernel_name: str) -> str: + """Basename, no language prefix, no task/extension suffix. Case kept for readability; matching + is case-insensitive via _match_key().""" + s = os.path.basename(str(kernel_name or "").strip().rstrip("/")) + changed = True + while changed: + changed = False + for p in _NAME_PREFIXES: + if len(s) > len(p) and s.lower().startswith(p): + s, changed = s[len(p):], True + for suf in _NAME_SUFFIXES: + if len(s) > len(suf) and s.lower().endswith(suf): + s, changed = s[: -len(suf)], True + return s or str(kernel_name or "") + + +def _match_key(kernel_name: str) -> str: + """Comparison key for slug matching: canonical name, case- and separator-insensitive.""" + return re.sub(r"[^a-z0-9]+", "", canon_name(kernel_name).lower()) + + +def make_slug(kernel_name: str, language: str, gfx: str) -> str: + return f"{_safe(canon_name(kernel_name))}__{_safe(language)}__{_norm_gfx(gfx) or 'unknown'}" + + +def _read_meta(meta_path: str): + try: + with open(meta_path, "r") as f: + text = f.read() + if yaml is not None: + return yaml.safe_load(text) or {} + return json.loads(text) + except Exception: + return None + + +def _dump_meta(meta: dict) -> str: + if yaml is not None: + return yaml.safe_dump(meta, sort_keys=False, allow_unicode=True) + return json.dumps(meta, indent=2, ensure_ascii=False) + + +def _atomic_write(path: str, data: str): + """Crash-safe: same-dir temp -> fsync -> os.replace -> dir fsync.""" + d = os.path.dirname(path) or "." + os.makedirs(d, exist_ok=True) + fd, tmp = tempfile.mkstemp(dir=d, prefix=".tmp_", suffix=".swap") + try: + with os.fdopen(fd, "w") as f: + f.write(data) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + finally: + if os.path.exists(tmp): + try: + os.unlink(tmp) + except OSError: + pass + try: + dirfd = os.open(d, os.O_RDONLY) + try: + os.fsync(dirfd) + finally: + os.close(dirfd) + except OSError: + pass + + +def content_signature(patch_text: str) -> str: + """Path-INSENSITIVE identity of a diff: added/removed code lines only, no headers or paths. + + A warm-started run re-emits the patch it adopted as its own `git diff`, from a different + workspace with different path prefixes — byte-different, same code. Without this the store keeps + re-importing its own output as a fresh 'win'; with it, that re-measurement is a REPRODUCTION of + the entry it came from, which is what promotes candidate -> active. + """ + body = [] + for line in (patch_text or "").splitlines(): + if line.startswith(("+++", "---", "diff ", "index ", "@@", "new file", "deleted file", + "similarity ", "rename ", "old mode", "new mode", "Binary files")): + continue + if line[:1] in ("+", "-"): + s = re.sub(r"\s+", " ", line[1:]).strip() + if s: + body.append(line[0] + s) + if not body: + return "" + return "csha:" + hashlib.sha256("\n".join(body).encode("utf-8", "replace")).hexdigest()[:32] + + +def bench_key(metric_kind: str, case_names) -> str: + """Identity of the MEASUREMENT a speedup came from; two speedups compare only when it matches. + Order-insensitive. The `b2:` namespace is deliberate — imported entries carry opaque `b:` keys + from whatever harness produced them, which must never be read as comparable to ours.""" + cases = sorted(c for c in (case_names or []) if c) + if not cases and not metric_kind: + return "" + raw = f"{str(metric_kind or 'unknown')}|{','.join(cases)}" + return "b2:" + hashlib.sha1(raw.encode("utf-8", "replace")).hexdigest()[:12] + + +# --- report prose ------------------------------------------------------------------------------ +# The two sections worth reading first. Heading text varies wildly across the imported backlog +# (`## What didn't work (dead-ends — do not re-fund)`, `(confirmed dead ends)`, ... 20+ suffixes +# over 248 reports), so match the stem only and tolerate the typographic apostrophe. +_SEC_DEAD_ENDS = re.compile(r"^(#{2,3})[^\n]*what\s+didn.?t\s+work[^\n]*$", re.I | re.M) +_SEC_KEY_OPTS = re.compile(r"^(#{2,3})[^\n]*key\s+optimizations[^\n]*$", re.I | re.M) +# Structured dead-ends the tech_lead emits alongside the prose. Absent => we keep the prose only, +# rather than regex-guessing structure out of bullets/tables/paragraphs and inventing empty fields. +_DEAD_ENDS_BLOCK = re.compile( + r"\s*```(?:ya?ml)?\n(.*?)```", re.S | re.I) + + +def _split_section(text: str, pattern): + """(heading+body, text_without_it) for the first match, else ('', text). The body ends at the + next heading of the same or shallower level, so a `###` subsection stays with its parent.""" + m = pattern.search(text or "") + if not m: + return "", text + level = len(m.group(1)) + tail = text[m.end():] + nxt = re.search(r"^#{1,%d} " % level, tail, re.M) + end = m.end() + (nxt.start() if nxt else len(tail)) + return text[m.start():end].rstrip() + "\n", text[:m.start()] + text[end:] + + +def reorder_report(text: str) -> str: + """Hoist 'Key optimizations' and "What didn't work" above everything else. Nothing is dropped — + an agent with room still reads the whole report, one that is tight on context reads the two + sections that change what it does. Returns the text untouched when neither is present.""" + if not text: + return text + key, rest = _split_section(text, _SEC_KEY_OPTS) + dead, rest = _split_section(rest, _SEC_DEAD_ENDS) + if not key and not dead: + return text + return "".join(s for s in (key, dead) if s) + "\n---\n\n" + rest.lstrip("\n") + + +def dead_ends_md(text: str) -> str: + """The "What didn't work" body verbatim, minus any machine-readable block (that is parsed + separately). Kept as text: the 248 imported reports write it as bullets, markdown tables and + plain paragraphs, and no regex turns all three into honest structure.""" + sec, _ = _split_section(text or "", _SEC_DEAD_ENDS) + if not sec: + return "" + body = sec.split("\n", 1)[1] if "\n" in sec else "" + return _DEAD_ENDS_BLOCK.sub("", body).strip() + + +def parse_dead_ends(text: str): + """The tech_lead's machine-readable dead-end list, or []. Each entry keeps whatever keys the + report supplied (idea / measured / mechanism); a malformed block is dropped, never patched up.""" + m = _DEAD_ENDS_BLOCK.search(text or "") + if not m or yaml is None: + return [] + try: + data = yaml.safe_load(m.group(1)) + except Exception: + return [] + if not isinstance(data, list): + return [] + return [{str(k): v for k, v in d.items()} for d in data + if isinstance(d, dict) and str(d.get("idea") or "").strip()] + + +def _techniques(meta: dict): + """The curated one-line summaries of what the patch actually does. Every imported entry has + them and until now nothing read them — they are the densest thing in the store.""" + t = (meta or {}).get("techniques") + if not isinstance(t, list): + return [] + return [str(x).strip() for x in t if str(x).strip()] + + +def _techniques_md(items) -> str: + if not items: + return "" + return "- techniques:\n" + "".join(f" * {i}\n" for i in items) + + +def _stack_str(meta: dict) -> str: + st = (meta or {}).get("verified_stack") + if not isinstance(st, dict) or not st: + return "unrecorded" + return ", ".join(f"{k} {v}" for k, v in sorted(st.items())) + + +def _alternates_md(alts) -> str: + """Same-direction runners-up. They were collapsed out of the ranking because they verify or fail + together, but their techniques are exactly where they differ from rank 1 — so list those.""" + if not alts: + return "- same-direction alternates: 0\n" + lines = [f"- same-direction alternates: {len(alts)}\n"] + for alt in alts: + techs = "; ".join(alt.get("techniques") or []) or "no techniques recorded" + lines.append(f" * {alt['speedup']:.4f}x — {techs}\n") + return "".join(lines) + + +def _prose_body(meta: dict, body: str) -> str: + """The report, with the two load-bearing sections hoisted. meta's dead-ends copy is a FALLBACK + for an entry whose report.md is gone — pasting it next to the report would just duplicate it.""" + if (body or "").strip(): + return reorder_report(body) + out = [] + for d in (meta.get("dead_ends") or []): + if isinstance(d, dict): + bits = [str(d.get(k)) for k in ("measured", "mechanism") if d.get(k)] + out.append(f"- {d.get('idea')}" + (f" — {' — '.join(bits)}" if bits else "")) + md = str(meta.get("dead_ends_md") or "").strip() + if not out and not md: + return "(no report recorded for this entry)" + head = "## What didn't work (from meta; report.md not available)\n\n" + return head + ("\n".join(out) + "\n\n" if out else "") + md + "\n" + + +def _rocm_version() -> str: + try: + with open("/opt/rocm/.info/version", "r", errors="replace") as f: + return f.read().strip().splitlines()[0].strip() + except (OSError, IndexError): + return "" + + +def detect_stack(language: str) -> dict: + """WHAT the speedup was measured on. This runs in the same container as the kernel, so every + value is observed, not inferred; anything unobservable is left out rather than guessed.""" + out = {} + if str(language or "").lower() == "triton": + for mod in ("triton", "torch"): + try: + out[mod] = str(__import__(mod).__version__) + except Exception: + pass + rocm = _rocm_version() + if rocm: + out["rocm"] = rocm + return out + + +def _find_by_content(root: str, gfx: str, slug: str, csig: str): + """(meta, exp_dir) of the entry on this page holding the same code, or None. Hashes patch.diff + for entries written before the signature was recorded (the imported backlog).""" + for meta, exp_dir in _iter_solutions(root, gfx, slug): + known = meta.get("content_signature") + if not known: + try: + with open(os.path.join(exp_dir, "patch.diff"), "r", errors="replace") as f: + known = content_signature(f.read()) + except OSError: + known = "" + if known and known == csig: + return meta, exp_dir + return None + + +def _record_reproduction(dup, csig: str, speedup: float, a) -> dict: + """Count a re-measurement onto the entry that already holds this code. Two of them promote + candidate -> active. The original's metric is NOT overwritten: it was measured on its own bench.""" + meta, exp_dir = dup + try: + reps = int(meta.get("reproductions") or 1) + 1 + except (TypeError, ValueError): + reps = 2 + meta["reproductions"] = reps + meta["content_signature"] = csig + if reps >= 2: + meta["lifecycle"] = "active" + try: + _atomic_write(os.path.join(exp_dir, "meta.yaml"), _dump_meta(meta)) + except OSError as e: + return {"written": False, "reason": "io_error: " + str(e)[:120]} + return { + "written": False, + "reason": "duplicate_impl", + "slug": make_slug(a.kernel_name, a.language, _norm_gfx(a.gfx)), + "dir": exp_dir, + "speedup": round(speedup, 4), + "reproduced": os.path.basename(exp_dir), + "reproductions": reps, + "lifecycle": meta["lifecycle"], + } + + +def cmd_write(a) -> dict: + gfx = _norm_gfx(a.gfx) + if not gfx: + return {"written": False, "reason": "missing_arch"} + + try: + speedup = float(a.speedup) + except (TypeError, ValueError): + return {"written": False, "reason": "invalid_speedup"} + if not (speedup > 1.0): # covers NaN, <=1.0 + return {"written": False, "reason": "no_improvement"} + + patch_text = "" + if a.patch and os.path.isfile(a.patch): + try: + with open(a.patch, "r", errors="replace") as f: + patch_text = f.read() + except OSError: + patch_text = "" + if not patch_text.strip(): + return {"written": False, "reason": "empty_diff"} + + kernel_class = a.kernel_class or "unknown" + case_names = [c.strip() for c in (a.case_names or "").split(",") if c.strip()] + slug = make_slug(a.kernel_name, a.language, gfx) + + # A re-measurement of code the store already holds is a REPRODUCTION, not a new entry. + csig = content_signature(patch_text) + dup = _find_by_content(a.root, gfx, slug, csig) if csig else None + if dup: + return _record_reproduction(dup, csig, speedup, a) + + exp_id = time.strftime("%Y%m%d_%H%M%S") + "_" + hashlib.sha1( + (slug + patch_text[:256] + str(time.time())).encode("utf-8", "replace") + ).hexdigest()[:6] + out_dir = os.path.join(a.root, gfx, _safe(kernel_class), slug, exp_id) + + baseline_ms = None + try: + baseline_ms = float(a.baseline_wall_ms) + except (TypeError, ValueError): + baseline_ms = None + wall_ms = (baseline_ms / speedup) if (baseline_ms and speedup > 0) else None + + meta = { + "lifecycle": "candidate", # earns 'active' only via independent reproduction + "gfx": gfx, + "kernel_class": kernel_class, + "kernel_name": a.kernel_name, + "language": a.language, + "metric": { + "speedup": round(speedup, 6), + "wall_ms": round(wall_ms, 6) if wall_ms is not None else None, + "baseline_wall_ms": round(baseline_ms, 6) if baseline_ms is not None else None, + "gpu_arch": gfx, + # What the speedup was measured against; resolve compares candidates only within one + # bench_key. Empty when the caller does not supply them. + "metric_kind": a.metric_kind or "", + "bench_key": bench_key(a.metric_kind, case_names), + "case_names": case_names, + }, + # The optimization IDEA, not the impl: resolve ranks at most one entry per direction. + "direction": (a.direction or "")[:120], + "content_signature": csig, + "reproductions": 1, + # exp_dir of the warm-start entry this was built on — tells a later curation pass "the same + # idea, one round further" from "an independent second discovery". + "derived_from": a.parent or "", + "verified_on": time.strftime("%Y-%m-%d"), + # Observed here, in the container that took the measurement — a speedup with no stack behind + # it cannot be compared to anything later. + "verified_stack": detect_stack(a.language), + "source_eval_dir": a.eval_dir or "", + } + + # Copy the tech_lead report verbatim as prose; lift its first non-empty line as the strategy, + # and keep its dead-ends so the next run on this kernel does not re-fund a closed direction. + strategy = "" + report_copied = None + if a.report and os.path.isfile(a.report): + try: + with open(a.report, "r", errors="replace") as f: + report_text = f.read() + for line in report_text.splitlines(): + s = line.strip().lstrip("# ").strip() + if s: + strategy = s[:300] + break + report_copied = report_text + structured = parse_dead_ends(report_text) + prose = dead_ends_md(report_text) + if structured: + meta["dead_ends"] = structured + if prose: + meta["dead_ends_md"] = prose + except OSError: + pass + if a.strategy: + strategy = a.strategy[:300] + meta["strategy"] = strategy + + try: + _atomic_write(os.path.join(out_dir, "patch.diff"), patch_text) + _atomic_write(os.path.join(out_dir, "meta.yaml"), _dump_meta(meta)) + if report_copied is not None: + _atomic_write(os.path.join(out_dir, "report.md"), report_copied) + except OSError as e: + return {"written": False, "reason": "io_error: " + str(e)[:120]} + + return { + "written": True, + "reason": "ok", + "slug": slug, + "exp_id": exp_id, + "dir": out_dir, + "speedup": round(speedup, 4), + } + + +def _iter_solutions(root: str, gfx: str, slug: str): + """Yield (meta, exp_dir) for every solution matching (gfx, slug), any kernel_class.""" + base = os.path.join(root, gfx) + if not os.path.isdir(base): + return + for kernel_class in sorted(os.listdir(base)): + slug_dir = os.path.join(base, kernel_class, slug) + if not os.path.isdir(slug_dir): + continue + for exp_id in sorted(os.listdir(slug_dir)): + exp_dir = os.path.join(slug_dir, exp_id) + meta_path = os.path.join(exp_dir, "meta.yaml") + if not os.path.isfile(meta_path): + meta_path = os.path.join(exp_dir, "meta.json") + meta = _read_meta(meta_path) + if isinstance(meta, dict): + yield meta, exp_dir + + +def _list_pages(root: str, gfx: str): + """Yield (slug, match_key, language) for every page under ///. + The slug splits from the RIGHT: a kernel name may itself contain '__' (e.g. `_w8a8__v2`).""" + base = os.path.join(root, gfx) + if not os.path.isdir(base): + return [] + out = {} + for kernel_class in sorted(os.listdir(base)): + kc_dir = os.path.join(base, kernel_class) + if not os.path.isdir(kc_dir): + continue + for slug in sorted(os.listdir(kc_dir)): + if slug in out or not os.path.isdir(os.path.join(kc_dir, slug)): + continue + parts = str(slug).rsplit("__", 2) + name, lang = (parts[0], parts[1]) if len(parts) == 3 else (slug, "") + out[slug] = (slug, _match_key(name), lang.lower()) + return [out[s] for s in sorted(out)] + + +def resolve_slug(root: str, gfx: str, kernel_name: str, language: str, match: str = "fuzzy"): + """Find the kernel page for (kernel_name, language) on this arch, most-specific tier first: + exact the canonical slug is on disk; + normalized same canonical name up to case/separators (`wvsplitk` -> `wvSplitK`); + fuzzy one canonical name contains the other, unambiguously — this is what turns an e2e + op_kind (`fused_moe`) into the `fused_moe_kernel` page. + Returns (slug_or_'', tier, info); info carries the pages NOT served, so a surprising match shows + up in the log instead of silently steering the run. + """ + want_slug = make_slug(kernel_name, language, gfx) + pages = _list_pages(root, gfx) + info = {"other_language_pages": [], "ambiguous": []} + if any(s == want_slug for s, _k, _lg in pages): + return want_slug, "exact", info + + want_key, want_lang = _match_key(kernel_name), str(language or "").strip().lower() + info["other_language_pages"] = [s for s, k, lg in pages if k == want_key and lg != want_lang] + if match == "exact": + return "", "none", info + + same_key = [s for s, k, lg in pages if k == want_key and lg == want_lang] + if same_key: + return same_key[0], "normalized", info + if match != "fuzzy" or len(want_key) < 6: + return "", "none", info + + # Containment, closest name first: `fused_moe` prefers `fused_moe_kernel` over + # `fused_moe_kernel_gptq_awq`. Two equally-close pages are AMBIGUOUS -> serve neither. + cands = [] + for s, k, lg in pages: + if len(k) < 6 or not (want_key in k or k in want_key): + continue + if lg != want_lang: + info["other_language_pages"].append(s) + else: + cands.append((abs(len(k) - len(want_key)), s)) + if not cands: + return "", "none", info + cands.sort() + if len(cands) > 1 and cands[0][0] == cands[1][0]: + info["ambiguous"] = [s for d, s in cands if d == cands[0][0]] + return "", "ambiguous", info + return cands[0][1], "fuzzy", info + + +# --------------------------------------------------------------------------------------------- +# Path remapping. A stored patch was produced in the workspace that won it — an arena checkout +# (`source/triton_fused_moe_kernel.py`, `csrc/...`) — while an e2e head run edits an extracted +# subtree (`kernel_src/.../fused_moe_kernel.py`). Same code, different prefix AND different +# basename, so no `-p` strip depth reaches the file: without rewriting the paths, every warm +# start on the head path fails to apply and the KB is dead weight there. +def _diff_targets(patch_text: str): + """Paths the diff touches: {path: is_new_file}. '' if the diff renames (not remappable).""" + targets, pending_new = {}, False + for line in (patch_text or "").splitlines(): + if line.startswith("rename from ") or line.startswith("rename to "): + return None + if line.startswith("new file mode"): + pending_new = True + elif line.startswith("--- "): + pending_new = pending_new or line[4:].strip() in ("/dev/null", "a//dev/null") + elif line.startswith("+++ "): + p = line[4:].strip().split("\t")[0] + if p != "/dev/null": + targets[re.sub(r"^b/", "", p)] = pending_new + pending_new = False + return targets + + +def _match_path(target: str, editable): + """Best editable path for one patch target, most-specific tier first: identical path, then one + path is the tail of the other, then same basename, then same basename modulo the language + prefix/extension (`triton_fused_moe_kernel.py` -> `fused_moe_kernel.py`). A tier with two + equally good hits is ambiguous -> no mapping, rather than a guess that verify pays to reject.""" + t_base = os.path.basename(target) + tiers = ( + [e for e in editable if e == target], + [e for e in editable if e.endswith("/" + target) or target.endswith("/" + e)], + [e for e in editable if os.path.basename(e) == t_base], + [e for e in editable if _match_key(os.path.basename(e)) == _match_key(t_base) + and os.path.splitext(e)[1] == os.path.splitext(target)[1]], + ) + for tier in tiers: + if len(set(tier)) == 1: + return tier[0] + if tier: + return "" + return "" + + +def _rewrite_paths(patch_text: str, mapping: dict) -> str: + """Rewrite the a//b/ path on every header line; hunks are copied through untouched.""" + out = [] + for line in patch_text.splitlines(): + if line.startswith("diff --git "): + for old, new in mapping.items(): + line = line.replace(f"a/{old} ", f"a/{new} ").replace(f"b/{old}", f"b/{new}") + elif line.startswith("--- a/") or line.startswith("+++ b/"): + head, path = line[:6], line[6:].split("\t")[0] + if path in mapping: + line = head + mapping[path] + out.append(line) + return "\n".join(out) + "\n" + + +def _drop_sections(patch_text: str, drop: set) -> str: + """Remove whole `diff --git` sections for the given target paths.""" + out, keep = [], True + for line in patch_text.splitlines(): + if line.startswith("diff --git "): + tail = line.split(" b/", 1) + keep = not (len(tail) == 2 and tail[1].strip() in drop) + if keep: + out.append(line) + return "\n".join(out) + "\n" + + +# A patch that also touches a non-source file this workspace lacks (a .gitignore line, a README +# note) is still a perfectly good kernel patch. Refusing the whole thing over it wastes the entry; +# the section is dropped and named in `dropped` so the decision is visible. +_SOURCE_EXTS = {".py", ".hip", ".cu", ".cuh", ".cpp", ".cc", ".hpp", ".h", ".c", ".jinja", + ".s", ".asm", ".json", ".yaml", ".yml", ".sh", ".mk", ".txt"} + + +def cmd_remap(a) -> dict: + """Rewrite a stored patch's paths onto THIS workspace's layout, or refuse and say why.""" + try: + with open(a.patch, "r", errors="replace") as f: + patch_text = f.read() + except OSError as e: + return {"remapped": False, "reason": "unreadable_patch: " + str(e)[:80]} + + editable = [p.strip().lstrip("./") for p in (a.editable or "").split(",") if p.strip()] + if not editable and a.workspace: + editable = _walk_workspace(a.workspace) + if not editable: + return {"remapped": False, "reason": "no_editable_set"} + + targets = _diff_targets(patch_text) + if targets is None: + return {"remapped": False, "reason": "rename_not_supported"} + if not targets: + return {"remapped": False, "reason": "no_paths_in_patch"} + + mapping, unmapped, new_files = {}, [], [] + for target, is_new in sorted(targets.items()): + if is_new: + new_files.append(target) + continue + hit = _match_path(target, editable) + if hit and hit != target: + mapping[target] = hit + elif not hit: + unmapped.append(target) + # A file the patch CREATES has nothing to match, so it follows the layout shift its edited + # siblings underwent. No shift (every edited path already fits here) => this workspace has the + # patch's own layout and the new file belongs exactly where the patch puts it. + for target in new_files: + host = next(iter(mapping.values()), None) + if not host and not any(not targets[t] for t in targets): + unmapped.append(target) # a patch of ONLY new files has nothing to anchor to + elif host and os.path.dirname(host) != os.path.dirname(target): + mapping[target] = os.path.join(os.path.dirname(host), os.path.basename(target)) + + dropped = [p for p in unmapped if os.path.splitext(p)[1] not in _SOURCE_EXTS] + unmapped = [p for p in unmapped if p not in dropped] + # All-or-nothing on SOURCE files: applying the mapped half of a patch leaves the workspace + # inconsistent, and verify would pay a full on-box run to discover that. + if unmapped: + return {"remapped": False, "reason": "unmapped_paths", "unmapped": unmapped, + "dropped": dropped, "mapped": mapping} + if not mapping and not dropped: + return {"remapped": False, "reason": "no_change_needed", "mapped": {}} + text = _drop_sections(patch_text, set(dropped)) if dropped else patch_text + try: + _atomic_write(a.out, _rewrite_paths(text, mapping)) + except OSError as e: + return {"remapped": False, "reason": "io_error: " + str(e)[:80]} + return {"remapped": True, "reason": "ok", "out": a.out, "mapped": mapping, "dropped": dropped} + + +_SKIP_DIRS = {".git", "__pycache__", "node_modules", "build", ".venv", "exp"} + + +def _walk_workspace(workspace: str, cap: int = 20000): + """Every file under the workspace, repo-relative, as a fallback editable set. Deliberately NOT + filtered by extension: a whitelist made real targets invisible (a `.cpp.jinja` template that + exists at the patch's exact path) and refused a patch that would have applied verbatim.""" + out = [] + for dirpath, dirnames, filenames in os.walk(workspace): + dirnames[:] = [d for d in dirnames if d not in _SKIP_DIRS and not d.startswith(".")] + for fn in filenames: + out.append(os.path.relpath(os.path.join(dirpath, fn), workspace)) + if len(out) >= cap: + return out + return out + + +def _speedup_of(meta: dict) -> float: + try: + return float((meta.get("metric") or {}).get("speedup")) + except (TypeError, ValueError): + return 0.0 + + +def _is_retired(meta: dict) -> bool: + """The curation's own verdict, as written into meta.yaml by the pass that built the store.""" + return meta.get("retained") is False or bool(meta.get("retired_reason")) + + +def _local_attestations(meta: dict) -> dict: + """This entry's attestation ledger, or {} when nobody has ever tried it. + + Empty rather than a zeroed ledger because `remote_value` drops empty values, and a record + that has never been recalled should carry no ledger at all — four zeroes and no ledger mean + the same thing to a reader, and the shorter one does not imply somebody looked. + """ + from kb.attest import attestations_of + ledger = attestations_of(meta if isinstance(meta, dict) else {}) + counted = any(ledger[k] for k in ("recalls", "validations", "failures", "not_reproduced")) + return ledger if counted else {} + + +def _rank_key(md): + """Recorded speedup, then reproductions, then exp_id for determinism.""" + meta, exp_dir = md + try: + reps = int(meta.get("reproductions") or 0) + except (TypeError, ValueError): + reps = 0 + return (-_speedup_of(meta), -reps, os.path.basename(exp_dir)) + + +def _track_record_md(meta) -> str: + """One line on what happened the last times this patch was adopted, or nothing at all. + + Omitted entirely for an untried entry rather than printed as "0 attempts": the reader is an + agent about to spend a verify slot, and a line that says nothing still costs it a decision. + """ + ledger = _local_attestations(meta) + if not ledger: + return "" + hint = _retire_hint_of(meta) + return ("- track record: adopted %d time(s) — %d reproduced a win, %d did not win, %d would " + "not run%s\n" % (ledger["recalls"], ledger["validations"], ledger["failures"], + ledger["not_reproduced"], " (**%s**)" % hint if hint else "")) + + +def _render_references(refs_dir: str, address: str, summary: str, views): + """Mirror the offered candidates' prose into `refs_dir` and index it, one prose path per view. + + Written up front, before any verdict, so a warm start that is later rejected stays auditable. + Both planes render the same page — a reference reads the same whether the entry came out of a + directory or from behind a KB Store key — so only `address` and each view's `origin` line + differ between them. A page that cannot be written is reported as "" rather than failing the + read: the patch is still adoptable without its prose. + """ + views = list(views) + key = "|".join(v["key"] for v in views).encode("utf-8", "replace") + set_dir = os.path.join(refs_dir, "sets", hashlib.sha1(key).hexdigest()[:7]) + top_bench = views[0]["bench_key"] if views else "" + index_lines = [ + f"# Warm-start references — {address}", "", summary, + f"Speedups compare only within one bench key; rank 1's is `{top_bench or 'none'}`.", "", + ] + paths = [] + for rank, v in enumerate(views, start=1): + meta = v["meta"] + prose_path = os.path.join(set_dir, f"reference_{rank:02d}.md") + try: + body = "" + if os.path.isfile(v["report_path"]): + with open(v["report_path"], "r", errors="replace") as f: + body = f.read() + _atomic_write(prose_path, ( + f"# Reference {rank:02d} — {address}\n\n" + f"- speedup: {v['speedup']:.4f}x ({v['metric_kind'] or 'unknown metric'}, " + f"bench `{v['bench_key'] or 'none'}`)\n" + f"- direction: {v['direction'] or 'unlabeled'}\n" + + _techniques_md(_techniques(meta)) + + f"- strategy: {meta.get('strategy', '')}\n" + + v["origin"] + + f"- verified_on: {meta.get('verified_on', '')}\n" + f"- verified_stack: {_stack_str(meta)}\n" + + _track_record_md(meta) + + _alternates_md(v["alts"]) + + f"\n---\n\n{_prose_body(meta, body)}\n" + )) + except OSError: + prose_path = "" + paths.append(prose_path) + index_lines.append( + f"- Rank {rank}: `{prose_path}` | speedup {v['speedup']:.4f}x | direction " + f"`{v['direction'] or 'unlabeled'}` | bench `{v['bench_key'] or 'none'}` | " + f"patch `{v['patch_path']}` | {len(v['alts'])} alternate(s) | status `read`" + ) + try: + _atomic_write(os.path.join(refs_dir, "index.md"), "\n".join(index_lines) + "\n") + except OSError: + pass + return paths + + +def _candidate(rank: int, v: dict, gfx: str, prose_path: str, top_bench: str) -> dict: + """The candidate record both planes hand the lane. Extra keys ride in `v['extra']`.""" + return dict({ + "rank": rank, + "exp_dir": v["exp_dir"], + "speedup": round(v["speedup"], 4), + "arch": gfx, + "patch_path": v["patch_path"], + "prose_path": prose_path, + "strategy": str(v["meta"].get("strategy") or ""), + "direction": v["direction"], + "techniques": _techniques(v["meta"]), + "bench_key": v["bench_key"], + "metric_kind": v["metric_kind"], + # False = ranked against rank 1 on a DIFFERENT case set, so their ordering is a prior only. + # Adoption is decided by this run's own measurement either way. + "comparable": bool(v["bench_key"]) and v["bench_key"] == top_bench, + "alternates": v["alts"], + # What happened the last times somebody actually adopted this patch, as opposed to the + # speedup its own writer measured once. An entry offered at rank 1 that three lanes have + # since failed to reproduce should not read identically to an untried one, and before this + # it did. Advisory only — nothing here filters on it (see kb/attest.py:retire_hint). + "validations": _local_attestations(v["meta"]).get("validations", 0), + "recalls": _local_attestations(v["meta"]).get("recalls", 0), + "retire_hint": _retire_hint_of(v["meta"]), + "status": "read", + }, **v["extra"]) + + +def _retire_hint_of(meta) -> str: + from kb.attest import retire_hint + return retire_hint(meta if isinstance(meta, dict) else {}) + + +def cmd_resolve(a) -> dict: + gfx = _norm_gfx(a.gfx) + if not gfx: + return {"read_reason": "missing_arch", "candidates": []} + + root = a.root + requested_slug = make_slug(a.kernel_name, a.language, gfx) + if not os.path.isdir(os.path.join(root, gfx)): + return {"read_reason": "kernel_page_not_found", "slug": requested_slug, "candidates": []} + + slug, match_tier, match_info = resolve_slug(root, gfx, a.kernel_name, a.language, a.match) + base_out = { + "slug": slug or requested_slug, + "requested_slug": requested_slug, + "match_tier": match_tier, + "other_language_pages": sorted(set(match_info["other_language_pages"])), + "ambiguous_pages": match_info["ambiguous"], + "candidates": [], + } + if not slug: + reason = ("ambiguous_kernel_page" if match_tier == "ambiguous" + else "no_page_for_language" if match_info["other_language_pages"] + else "kernel_page_not_found") + return dict(base_out, read_reason=reason) + + # The path segment already guarantees same-arch; re-check metric.gpu_arch to catch a mislabeled entry. + found = [(m, d) for (m, d) in _iter_solutions(root, gfx, slug) + if _norm_gfx((m.get("metric") or {}).get("gpu_arch") or m.get("gfx") or gfx) == gfx] + if not found: + return dict(base_out, read_reason="no_same_arch") + + # --- curation gate: what this page may OFFER, before any ranking ------------------------- + total = len(found) + servable = found if a.include_retired else [(m, d) for (m, d) in found if not _is_retired(m)] + retired_n = total - len(servable) + try: + min_speedup = float(a.min_speedup) + except (TypeError, ValueError): + min_speedup = 1.0 + above = [(m, d) for (m, d) in servable if _speedup_of(m) >= min_speedup] + below_n = len(servable) - len(above) + stats = {"total": total, "retired": retired_n, "below_min_speedup": below_n, + "min_speedup": min_speedup} + if not above: + return dict(base_out, filtered=stats, + read_reason="all_retired" if not servable else "below_min_speedup") + + top, alternates, collapsed = collapse_by_direction( + sorted(above, key=_rank_key), lambda md: md[0].get("direction"), lambda md: md[1], a.top_n) + stats["same_direction_collapsed"] = collapsed + + views = [] + for (meta, exp_dir), alt_of in zip(top, alternates): + metric = meta.get("metric") or {} + views.append({ + "key": exp_dir, + "meta": meta, + "exp_dir": exp_dir, + "patch_path": os.path.join(exp_dir, "patch.diff"), + "report_path": os.path.join(exp_dir, "report.md"), + "speedup": _speedup_of(meta), + "direction": str(meta.get("direction") or ""), + "bench_key": str(metric.get("bench_key") or ""), + "metric_kind": str(metric.get("metric_kind") or ""), + "origin": f"- source: {meta.get('source_eval_dir', '')}\n", + "alts": [{ + "exp_dir": d, + "patch_path": os.path.join(d, "patch.diff"), + "speedup": round(_speedup_of(m), 4), + "bench_key": str((m.get("metric") or {}).get("bench_key") or ""), + "techniques": _techniques(m), + } for (m, d) in alt_of], + "extra": {"slug": slug}, + }) + + summary = (f"Matched `{requested_slug}` -> `{slug}` ({match_tier}). {len(top)} direction(s) " + f"offered from {total} recorded run(s): {retired_n} retired by curation, " + f"{below_n} below {min_speedup:g}x, {collapsed} same-direction re-discoveries " + f"moved to `alternates`.") + prose = _render_references(a.refs_dir, f"slug `{slug}` (gfx {gfx})", summary, views) + candidates = [_candidate(rank, v, gfx, p, views[0]["bench_key"]) + for rank, (v, p) in enumerate(zip(views, prose), start=1)] + return dict(base_out, read_reason="read", candidates=candidates, filtered=stats) + + +def cmd_languages(a) -> dict: + """Which languages this kernel actually has a page in. A caller that guesses `triton` for a + kernel the store keeps under `hip`/`ck` gets read_reason=empty and silently loses its history, + so let the store answer instead of a task_type mapping that cannot tell hip from ck.""" + gfx = _norm_gfx(a.gfx) + if not gfx: + return {"languages": [], "reason": "missing_arch"} + pages = _list_pages(a.root, gfx) + want = _match_key(a.kernel_name) + langs = sorted({lg for _s, k, lg in pages if k == want and lg}) + if langs: + return {"gfx": gfx, "languages": langs, "match_tier": "exact", "reason": "ok"} + near = sorted({lg for _s, k, lg in pages if lg and (want in k or k in want)}) + if near: + return {"gfx": gfx, "languages": near, "match_tier": "fuzzy", "reason": "ok"} + return {"gfx": gfx, "languages": [], "match_tier": "none", "reason": "no_page"} + + +# Stacks the imported backlog was measured on, recovered from the campaign's own eval dirs +# (`analysis.json` / `codebase_context.md` device strings). Marked as recovered, not observed — +# a later reader must be able to tell a backfilled stack from one detect_stack() saw first-hand. +_BACKFILL_STACK = { + # rocm is on all three, not just the two that compile against it directly: the whole campaign + # ran in one container image, and rocm is the version the remote identity is keyed on, so a + # triton entry without it exports to a different address than the hip entry beside it. + "triton": {"triton": "3.6.0", "torch": "2.11.0", "rocm": "7.2"}, + "hip": {"rocm": "7.2"}, + "ck": {"rocm": "7.2"}, +} + + +def _backfill_one(meta: dict, exp_dir: str, stacks: dict): + """Fields to add/drop for one entry, as (new_meta, changes) — or (meta, {}) when already done.""" + out = dict(meta) + changes = {"add": [], "drop": [], "fix": []} + + if not str(out.get("dead_ends_md") or "").strip(): + try: + with open(os.path.join(exp_dir, "report.md"), "r", errors="replace") as f: + report = f.read() + except OSError: + report = "" + prose = dead_ends_md(report) + if prose: + out["dead_ends_md"] = prose + changes["add"].append("dead_ends_md") + + # Fill per KEY, not per dict: an early backfill gave triton entries {triton, torch} and no + # rocm, which is exactly the key the remote identity is derived from. Values already present + # are never overwritten — an observed stack always outranks a recovered one. + st = out.get("verified_stack") + st = dict(st) if isinstance(st, dict) else {} + known = stacks.get(str(out.get("language") or "").lower()) or {} + added = [k for k in known if not str(st.get(k) or "").strip()] + if added: + st.update({k: known[k] for k in added}) + st.setdefault("recorded_by", "campaign20_backfill") + out["verified_stack"] = st + changes["add"].append("verified_stack:" + ",".join(sorted(added))) + + # impl_signature is a different hash under a name nothing reads: _find_by_content() falls back + # to re-hashing patch.diff for all 248 entries on every resolve. Recompute under the real name. + if not out.get("content_signature"): + try: + with open(os.path.join(exp_dir, "patch.diff"), "r", errors="replace") as f: + csig = content_signature(f.read()) + except OSError: + csig = "" + if csig: + out["content_signature"] = csig + changes["fix"].append("content_signature") + if "impl_signature" in out and out.get("content_signature"): + out.pop("impl_signature") + changes["drop"].append("impl_signature") + + # Never read, and each is either constant or a duplicate of a field right next to it. + for dead in ("layer", "platforms", "patch_content"): + if dead in out: + out.pop(dead) + changes["drop"].append(dead) + + if not (changes["add"] or changes["drop"] or changes["fix"]): + return meta, {} + return out, changes + + +def cmd_backfill_content(a) -> dict: + """Bring the imported backlog up to the current content shape. Dry-run by default; only ever + adds the fields named above — retained / direction / techniques / metric are never touched.""" + root = a.root + if not os.path.isdir(root): + return {"ok": False, "reason": "no_such_root: " + root} + stacks = dict(_BACKFILL_STACK) + scanned = changed = failed = 0 + for dirpath, _dirs, files in os.walk(root): + if "meta.yaml" not in files: + continue + scanned += 1 + meta_path = os.path.join(dirpath, "meta.yaml") + meta = _read_meta(meta_path) + if not isinstance(meta, dict): + failed += 1 + continue + new_meta, changes = _backfill_one(meta, dirpath, stacks) + if not changes: + continue + changed += 1 + print(json.dumps({"dir": dirpath, **changes}, ensure_ascii=False)) + if a.apply: + try: + _atomic_write(meta_path, _dump_meta(new_meta)) + except OSError as e: + failed += 1 + print(json.dumps({"dir": dirpath, "error": str(e)[:120]})) + return {"ok": True, "applied": bool(a.apply), "scanned": scanned, + "changed": changed, "failed": failed} + + +# --- remote KB export ------------------------------------------------------------------------- +# Record shape mirrors KernelForge's (knowledge/kernel_identity.py and +# rewrite_by_flydsl/{identity,agent_kb,record_store}.py @ baabdae); the ADDRESS does not, and +# kb/identity.py owns it for both workflows and says why. Read and write must both go through it: +# the store finds nothing if the two sides disagree by one segment, and there is no error to notice +# — a mistyped dimension just reads as a cold start. +# +# The scheme is `geak:`, not `kernel:`, because our credential is scoped to `geak` identities and +# 403s on both `kernel:` and `inference:`. That scheme is client-defined and exact-lookup only, so +# every dimension has to be something the READ side can recompute from what it already knows; +# nothing may be derived from run-local state. Two consequences worth having in view here: +# +# * the serving framework (vllm / sglang), its version and the numeric precision are NOT +# dimensions, even though an e2e run knows all three. kernel_lane.js does not — it has no +# upstream awareness at all, and pass-through from e2e forwards only `target_language`. A +# dimension the reader cannot reconstruct is a permanent silent 404. They ride in +# `value.upstream` instead, where a client can filter on them; precision is additionally +# already spelled into most kernel names (`fused_moe_int4_w4a16`, `_w8a8_triton_block_scaled_mm`) +# so keying on it would double-encode and split those pages. +# * every write publishes to BOTH rungs of kernel_canonical_ids(). The service does no prefix +# aggregation, so the version-agnostic page exists only because we put records there. +REMOTE_PRODUCER = "geak" +REMOTE_ARTIFACT_KIND = "rewrite" # upstream ARTIFACT_KIND for a recipe bundle + +try: + from kb import identity as _kbid +except ImportError: # resolve/write stay usable; only the remote pair needs it + _kbid = None + +REMOTE_SCHEME = "geak" +REMOTE_DOMAIN = "kernel" +REMOTE_FRAMEWORK = "rocm" +REMOTE_UNKNOWN_VERSION = "unspecified" + + +def _identity_module(): + if _kbid is None: + raise RuntimeError("kb_identity_unavailable: kb/identity.py must be importable from the repo root") + return _kbid + + +def remote_segment(value, fallback: str) -> str: + """Fold a free-form value into one identity dimension. Delegates so there is one folding rule.""" + return _identity_module().segment(value, fallback) + + +def remote_gpu(gfx: str, override: str = "") -> str: + """The compile target (`gfx950`), NOT the product model, and the LEADING dimension. + + Upstream keys this on the marketing name (`mi355x`) and leaves gfx out of the identity. We + diverge on both counts. gfx is what every producer and consumer on our side already holds — off + the box, out of meta.yaml, out of the e2e harness — whereas the product model exists only as a + lookup table someone has to keep current, and an unmapped arch would file half a kernel's + history under a name nothing looks up. It leads because an arch mismatch is the quietest way to + waste a round: a gfx942-tuned patch compiles clean on gfx950 and is merely slower, where a + wrong kernel_name or language at least fails to apply. + """ + return remote_segment(override or _norm_gfx(gfx), fallback="unknown") + + +def remote_framework_version(meta: dict, override: str = "") -> str: + """The ROCm version this entry was measured on, cut to `.` for the address. + + Coarse on purpose, and now additionally droppable: it is the last segment precisely because + 7.2 -> 7.3 usually keeps a patch applicable, so the second rung of the ladder is the one that + keeps 20 kernels warm through an image upgrade instead of cold-starting all of them at once. + + Never guessed: no rocm key exports as `unspecified`. That entry still gets both rungs, so the + version-agnostic page sees it even though its exact page is one nobody will construct. + """ + stack = meta.get("verified_stack") + raw = str(override or "").strip() + if not raw: + raw = str((stack or {}).get("rocm") or "").strip() if isinstance(stack, dict) else "" + return _identity_module()._short_version(raw) + + +def remote_identity(meta: dict, producer: str = REMOTE_PRODUCER, gpu: str = "", + version: str = "") -> dict: + """The four dimensions of the address. + + `producer` is accepted and recorded but is no longer a dimension: it has one value here, the + service stamps it on every record, and artifacts are already partitioned under + `kb///`. `version` overrides only the key, never the record — a box whose + ROCm this script cannot detect would otherwise file at `:unspecified` and split one kernel's + history in two while `value.verified_stack` keeps saying, correctly, that nothing was observed. + """ + return _identity_module().kernel_identity( + gfx=remote_gpu(meta.get("gfx") or (meta.get("metric") or {}).get("gpu_arch") or "", gpu), + kernel_name=meta.get("kernel_name"), + backend=meta.get("language"), + rocm_version=remote_framework_version(meta, version), + ) + + +def remote_canonical_ids(identity: dict): + """Every address this entry is published at, most specific first. Never a subset.""" + return _identity_module().kernel_canonical_ids(identity) + + +def remote_canonical_id(identity: dict) -> str: + """The exact address — rung 1. The one a fresh write is fingerprinted from.""" + return remote_canonical_ids(identity)[0] + + +def _remote_digest(meta: dict, exp_dir: str) -> str: + """The port fingerprint that names the candidate. Reuses content_signature, which already + dedups this store by patch content, so re-exporting one entry updates one candidate upstream + instead of piling on a new one per run.""" + sig = str(meta.get("content_signature") or "") + if sig: + return sig.split(":", 1)[-1] + try: + with open(os.path.join(exp_dir, "patch.diff"), "r", errors="replace") as f: + return content_signature(f.read()).split(":", 1)[-1] + except OSError: + return "" + + +def remote_session_id(canonical_id: str, kernel_name: str, digest: str) -> str: + """`---`, upstream's shape. + + Pass the EXACT rung: the id is reused verbatim on the coarser one, which is what lets the two + share uploaded artifacts instead of duplicating a 240KB patch. Fingerprinting each rung on its + own would give one measurement two unrelated ids and stop the coarse page from being a + reproduction of the exact one. + """ + return _identity_module().session_id(canonical_id, kernel_name, digest, REMOTE_PRODUCER) + + +def _sha256_file(path: str): + h, size = hashlib.sha256(), 0 + with open(path, "rb") as f: + while True: + chunk = f.read(1024 * 1024) + if not chunk: + break + h.update(chunk) + size += len(chunk) + return h.hexdigest(), size + + +def remote_value(meta: dict, digest: str = "") -> dict: + """The producer-owned half of the record. Upstream treats `value` as opaque — no schema to + satisfy — so this is our own meta.yaml minus what the identity already carries. + + bench_key and metric_kind are not optional here even though nothing upstream reads them: + get_top_sessions ranks purely on the `speedup` number we ourselves declare, so a `b:` entry and + a `b2:` one get ordered against each other as if they were comparable. The reader has to filter + on these client-side, and it cannot do that if we did not send them.""" + metric = meta.get("metric") or {} + value = { + "direction": str(meta.get("direction") or ""), + "techniques": _techniques(meta), + "strategy": str(meta.get("strategy") or ""), + "metric": { + "speedup": metric.get("speedup"), + "wall_ms": metric.get("wall_ms"), + "baseline_wall_ms": metric.get("baseline_wall_ms"), + "metric_kind": str(metric.get("metric_kind") or ""), + "bench_key": str(metric.get("bench_key") or ""), + "case_names": list(metric.get("case_names") or []), + }, + "verified_stack": meta.get("verified_stack") if isinstance(meta.get("verified_stack"), dict) else {}, + "verified_on": str(meta.get("verified_on") or ""), + "measured_by": str(meta.get("measured_by") or ""), + "reproductions": meta.get("reproductions"), + # NOT the same thing as `reproductions`, and the two are easy to conflate into one wrong + # number. `reproductions` counts how many times this lane WROTE the same patch again — a + # measure of how often the optimizer rediscovers an idea, produced entirely by the writer. + # `attestations` counts what happened when somebody READ this record and took it to a box: + # recalls / validations / failures / not_reproduced, in the vocabulary kb/attest.py defines + # and both lanes share. A record can be rediscovered five times and never once survive a + # recall, and only the second number says so. + "attestations": _local_attestations(meta), + "lifecycle": str(meta.get("lifecycle") or ""), + "retained": meta.get("retained"), + # The same two fields the e2e records carry, so one reader can ask "should I believe this" + # of either lane without knowing which one wrote it. Derived, not invented: `active` is + # earned here only by independent reproduction (see the write path), so it already IS the + # validation flag — it was just spelled in a vocabulary nothing outside this file knew. + # The basis is named for what actually produced the number rather than mapped onto the + # e2e taxonomy: a kernel's speedup comes from its own isolated bench harness, and calling + # that a `hot_ab` would claim a serving-level A/B that never ran. + "validated": str(meta.get("lifecycle") or "") == "active", + "validation_basis": "kernel_bench", + # The same digest the session id is built from, so a reader that dedups against its own + # store and the address it was filed under can never disagree about what this patch is. + "content_signature": ("csha:" + digest) if digest else str(meta.get("content_signature") or ""), + "artifacts": {"patch": "patch.diff", "report": "report.md"}, + } + dead = meta.get("dead_ends") + if isinstance(dead, list) and dead: + value["dead_ends"] = dead + # dead_ends_md is deliberately NOT sent: it runs to tens of KB and report.md already carries it + # verbatim as an artifact. Structured dead ends are small enough to ride in the record. + return {k: v for k, v in value.items() if v not in ("", None, [], {})} + + +def remote_records(meta: dict, exp_dir: str, producer: str = REMOTE_PRODUCER, gpu: str = "", + version: str = ""): + """One measurement as upload-ready candidates — one per rung, most specific first. + + All rungs carry the same session id, the same knowledge and the same files. They are not + variants of a result; they are one result filed at every address a reader might construct. That + is why the caller must publish all of them or none: a coarse page fed by only some runs ranks + worse than an empty one, because a reader cannot tell a thin page from a complete one. + + `rung` is stamped on each record so an uploader can skip re-transferring artifacts for rungs + after the first — remotely the bytes are shared via the session id, and re-PUTting them would + only burn the presign window. + """ + identity = remote_identity(meta, producer, gpu, version) + cids = remote_canonical_ids(identity) + digest = _remote_digest(meta, exp_dir) + sid = remote_session_id(cids[0], identity["kernel_name"], digest) + speedup = _speedup_of(meta) + files = [] + for name in ("patch.diff", "report.md"): + path = os.path.join(exp_dir, name) + if not os.path.isfile(path): + continue + file_sha, size = _sha256_file(path) + files.append({"path": name, "local_path": path, "kind": REMOTE_ARTIFACT_KIND, + "sha256": file_sha, "size": size}) + # The knowledge document upstream's own writer produces: four keys, everything else under + # `value`. `speedup` sits at the top because that is the ranking key the service reads — it + # only honours a flat top-level `knowledge.` scalar and rejects a nested path with a 400. + knowledge = { + "producer": remote_segment(producer, REMOTE_PRODUCER), + "speedup": round(speedup, 4) if speedup else None, + "identity": identity, + "value": remote_value(meta, digest), + } + return [{ + "canonical_id": cid, + "session_id": sid, + "exp_dir": exp_dir, + "rung": rung, + "knowledge": knowledge, + "files": files, + # Upstream's own gate: a candidate is always recorded, the pointer moves only on a real win. + # Evaluated per rung, since each address keeps its own champion pointer. + "champion_eligible": speedup > 1.0, + "champion": False, + } for rung, cid in enumerate(cids)] + + +def remote_record(meta: dict, exp_dir: str, producer: str = REMOTE_PRODUCER, gpu: str = "", + version: str = "") -> dict: + """The exact-rung record alone. Kept for callers that only want the address, never for writing — + writing one rung and not the other is the failure mode remote_records() exists to prevent.""" + return remote_records(meta, exp_dir, producer, gpu, version)[0] + + +def cmd_export_remote(a) -> dict: + """Render this store as KB Store candidates, one JSON line each, champion pre-decided. + + Nothing is uploaded here — this only produces what to upload, so the mapping is reviewable and + diffable before anything leaves the machine. kb/remote_upload.py consumes the output. + """ + root = a.root + if not os.path.isdir(root): + return {"ok": False, "reason": "no_such_root: " + root} + want_gfx = _norm_gfx(a.gfx) if a.gfx else "" + want_name = _match_key(a.kernel_name) if a.kernel_name else "" + + records, scanned, skipped = [], 0, {"retired": 0, "no_patch": 0, "unreadable": 0, "filtered": 0} + for dirpath, _dirs, files in sorted(os.walk(root)): + if "meta.yaml" not in files: + continue + scanned += 1 + meta = _read_meta(os.path.join(dirpath, "meta.yaml")) + if not isinstance(meta, dict): + skipped["unreadable"] += 1 + continue + gfx = _norm_gfx(meta.get("gfx") or (meta.get("metric") or {}).get("gpu_arch") or "") + if (want_gfx and gfx != want_gfx) or (want_name and _match_key(meta.get("kernel_name")) != want_name): + skipped["filtered"] += 1 + continue + # Retired entries are dominated duplicates, not negative knowledge, and the service ranks on + # the speedup we declare — offering them would put a retired win in someone's top-N. + if _is_retired(meta) and not a.include_retired: + skipped["retired"] += 1 + continue + if not os.path.isfile(os.path.join(dirpath, "patch.diff")): + skipped["no_patch"] += 1 + continue + records.extend(remote_records(meta, dirpath, a.producer, a.gpu)) + + # One champion per identity, upstream's rule: must beat 1.0x, and highest wins. Ties break on + # session id so two runs of this exporter promote the same candidate. + best = {} + for rec in records: + if not rec["champion_eligible"]: + continue + cur = best.get(rec["canonical_id"]) + key = (rec["knowledge"]["speedup"] or 0.0, rec["session_id"]) + if cur is None or key > cur[0]: + best[rec["canonical_id"]] = (key, rec) + for _key, rec in best.values(): + rec["champion"] = True + + # Byte-identical patches under one identity are ONE candidate upstream, so a collision here is + # the dedup working. Which of them we send still matters: the record is written with + # mode=replace, so keeping the lower of two measurements of the same patch would publish a + # speedup we have already beaten. Highest wins, exp_dir breaks ties, and the dropped rows are + # named in the summary rather than vanishing. + by_id = {} + dropped = [] + for rec in records: + ident = (rec["canonical_id"], rec["session_id"]) + cur = by_id.get(ident) + if cur is None: + by_id[ident] = rec + continue + ranked = sorted((cur, rec), + key=lambda r: (-(r["knowledge"]["speedup"] or 0.0), r["exp_dir"])) + by_id[ident] = ranked[0] + dropped.append(ranked[1]["exp_dir"]) + + emitted = 0 + out = open(a.out, "w") if a.out else None + try: + for rec in records: + if by_id.get((rec["canonical_id"], rec["session_id"])) is not rec: + continue + line = json.dumps(rec, ensure_ascii=False) + (out.write(line + "\n") if out else print(line)) + emitted += 1 + finally: + if out: + out.close() + # `emitted` counts records, not measurements: each entry is published at every rung of its + # ladder, so the honest headline is both numbers. `sessions` is how many distinct measurements + # went out; emitted/sessions should equal the ladder depth for a healthy export. + return {"ok": True, "scanned": scanned, "emitted": emitted, + "sessions": len({r["session_id"] for r in records}), + "identities": len({r["canonical_id"] for r in records}), + "exact_identities": len({r["canonical_id"] for r in records if r["rung"] == 0}), + "champions": len(best), + "deduped": len(dropped), "deduped_dirs": sorted(dropped), + "skipped": skipped, "out": a.out or "-"} + + +def _value_as_meta(value: dict, gfx: str) -> dict: + """Read a record's `value` back as a meta. + + `remote_value()` produced it FROM a meta, minus what the identity already carries, so the + prose helpers below (`_techniques_md`, `_alternates_md`, `_prose_body`) work on it unchanged + and a remote-sourced reference reads exactly like a local one. + """ + meta = dict(value or {}) + metric = dict(meta.get("metric") or {}) + metric.setdefault("gpu_arch", gfx) + meta["metric"] = metric + return meta + + +def _store_ladder(a, gfx: str): + """The addresses to try, most specific first, each paired with the tier it represents. + + framework_version is the one dimension a reader can get wrong without noticing: the store is + keyed on the ROCm an entry was measured on, this box may be on another, and a bare miss looks + exactly like a cold start. The ladder is the answer — but only because the WRITER publishes the + version-agnostic rung too. Nothing here derives a page that was never written; each rung is a + real address that a `write-remote` on this box would also have filled. + + An explicit --canonical-id is taken as given and gets no ladder. A caller that names an address + is usually auditing one page, and silently widening the read would misreport which page + answered. + """ + if a.canonical_id: + return [(a.canonical_id, "exact")] + meta = {"kernel_name": a.kernel_name, "language": a.language, + "verified_stack": detect_stack(a.language)} + identity = remote_identity(meta, a.producer, remote_gpu(gfx, getattr(a, "gpu", "")), + getattr(a, "framework_version", "")) + return list(zip(remote_canonical_ids(identity), ("exact", "any_version"))) + + +def cmd_retract_remote(a) -> dict: + """Take back a key-addressed kernel record. The counterpart to `write-remote`. + + The service has no delete, so this rewrites the session in place: `retained: false`, a reason, + the ranking scalar zeroed, and the identity's champion re-pointed at the best survivor. See + kb/retract.py for why all three are needed and why any two of them is worse than none. + + Both rungs are visited, because `write-remote` filled both with the SAME session id. Retracting + only the exact rung leaves the record live on the version-agnostic page, which is the page a box + on a different ROCm reads — i.e. it would survive exactly where it is least verifiable. + + `--canonical-id` addresses one page only, matching `resolve-remote`'s rule: a caller that names + an address is auditing it, and quietly widening a WRITE beyond what was asked for is not a + behaviour this command should have. + """ + from kb.retract import retract_session, retraction_ok + gfx = _norm_gfx(a.gfx) + if not gfx and not a.canonical_id: + return {"retracted": False, "reason": "missing_arch"} + store, mirror, why = open_plane(a, CHAMPION_METRIC, 1.0) + planes = [p for p in (store, mirror) if p is not None] + if not planes: + return {"retracted": False, "reason": why} + out = {"applied": bool(a.apply), "session_id": a.session_id, "reason": a.reason, + "plane_note": why, "pages": []} + for cid, tier in _store_ladder(a, gfx): + for plane in planes: + report = retract_session(plane, cid, a.session_id, a.reason, CHAMPION_METRIC, + actor=str(getattr(a, "measured_by", "") or ""), + scan=int(a.scan), apply=bool(a.apply)) + out["pages"].append(dict(report, tier=tier)) + out["retracted"] = retraction_ok(out["pages"], a.apply) + return out + + +def _attest_evidence(a) -> dict: + """The one-line record of what this box saw, shared by both attest paths.""" + evidence = {} + for key, raw in (("measured_speedup", getattr(a, "measured_speedup", None)), + ("note", getattr(a, "note", "")), + ("canonical_id", getattr(a, "canonical_id", ""))): + if raw not in (None, ""): + evidence[key] = raw + # Argparse hands the ratio over as a string. Stored as one it would land in meta.yaml quoted, + # and every later reader comparing it against a speedup would be comparing str to float. + if "measured_speedup" in evidence: + try: + evidence["measured_speedup"] = float(evidence["measured_speedup"]) + except (TypeError, ValueError): + evidence.pop("measured_speedup") + return evidence + + +def cmd_attest(a) -> dict: + """Count one attempt to actually USE a stored entry, straight into its meta.yaml. + + The local plane has no session ids and no service — an entry IS a directory — so this is a + read-modify-atomic-write of the same file the write path owns, using the same arithmetic + kb/attest.py applies remotely. Sharing the arithmetic and not the transport is deliberate: the + counters have to mean the same thing on both planes or a curation pass cannot compare them, + but a local store should not need a KB plane to record that a patch did not apply. + + Like the remote one, this moves nothing: the speedup meta declares is left exactly as it was, + and the entry keeps its rank. A patch that failed to apply on one workspace is a fact about + that workspace as much as about the patch. + """ + from kb.attest import record_attestation, retire_hint + exp_dir = str(getattr(a, "exp_dir", "") or "") + meta_path = os.path.join(exp_dir, "meta.yaml") + meta = _read_meta(meta_path) + if not meta: + return {"attested": False, "reason": "no_meta", "exp_dir": exp_dir} + try: + updated = record_attestation(dict(meta), a.outcome, + actor=str(getattr(a, "measured_by", "") or ""), + evidence=_attest_evidence(a)) + except Exception as e: + return {"attested": False, "reason": "bad_outcome: " + str(e)[:120], "exp_dir": exp_dir} + out = {"attested": bool(a.apply), "applied": bool(a.apply), "exp_dir": exp_dir, + "outcome": a.outcome, "attestations": updated["attestations"], + "retire_hint": retire_hint(updated)} + if not a.apply: + return out + try: + _atomic_write(meta_path, _dump_meta(updated)) + except OSError as e: + out.update({"attested": False, "reason": "write_failed: " + str(e)[:120]}) + return out + + +def cmd_attest_remote(a) -> dict: + """The same count, against a key-addressed record on either plane. + + Walks BOTH rungs for the same reason `retract-remote` does: `write-remote` filled them with one + session id, and a box on a different ROCm reads the version-agnostic rung — leaving it with a + stale ledger hides the failures from exactly the readers most likely to hit them. + """ + from kb.attest import attest_session, attestation_ok, retire_hint + gfx = _norm_gfx(a.gfx) + if not gfx and not a.canonical_id: + return {"attested": False, "reason": "missing_arch"} + store, mirror, why = open_plane(a, CHAMPION_METRIC, 1.0) + planes = [p for p in (store, mirror) if p is not None] + if not planes: + return {"attested": False, "reason": why} + out = {"applied": bool(a.apply), "session_id": a.session_id, "outcome": a.outcome, + "plane_note": why, "pages": []} + for cid, tier in _store_ladder(a, gfx): + for plane in planes: + report = attest_session(plane, cid, a.session_id, a.outcome, + actor=str(getattr(a, "measured_by", "") or ""), + evidence=_attest_evidence(a), apply=bool(a.apply)) + out["pages"].append(dict(report, tier=tier)) + out["attested"] = attestation_ok(out["pages"], a.apply) + hints = [p.get("retire_hint") for p in out["pages"] if p.get("retire_hint")] + out["retire_hint"] = hints[0] if hints else "" + return out + + +def _store_near_misses(store, cid: str): + """Identities differing from `cid` only in framework_version, newest-looking last. + + A third tier below the ladder, and only reachable on a store that predates double-writing — + once every write fills the version-agnostic rung, that rung answers first and this never runs. + Kept because the alternative for such a store is a cold start on a kernel that has history. + """ + parts = cid.split(":") + if len(parts) != 7: + return [] + return sorted(other for other in store.identities() + if (lambda s: len(s) == 7 and s[:6] == parts[:6] and s[6] != parts[6]) + (other.split(":"))) + + +def cmd_resolve_remote(a) -> dict: + """Rank the top-N candidates under one canonical id and mirror their prose, like `resolve`. + + Same output shape as `resolve` on purpose: the lane's schema, verify gate and adopt step do not + change when the KB moves behind a key. What differs is where curation happens. The local store + is curated on disk (`retained: false`, one entry per direction); the KB Store ranks on nothing + but the `speedup` a producer declared, so the direction collapse and the bench-key comparability + check have to be redone here, client-side, against the records it hands back. + """ + gfx = _norm_gfx(a.gfx) + if not gfx and not a.canonical_id: + return {"read_reason": "missing_arch", "candidates": []} + # Reading takes ONE plane, never both. Merging two rankings would need a comparability rule + # across planes that nothing here has, and silently preferring one would make a stale local + # mirror shadow the service without saying so. + store, _second, why = open_plane(a, CHAMPION_METRIC, 1.0) + if store is None: + return {"read_reason": why.split(":", 1)[0], "reason": why, "candidates": []} + + ladder = _store_ladder(a, gfx) + cid, match_tier = ladder[0] + segs = cid.split(":") + requested_slug = make_slug(a.kernel_name or (segs[3] if len(segs) > 3 else ""), + a.language or (segs[4] if len(segs) > 4 else ""), gfx) + base_out = {"slug": requested_slug, "requested_slug": requested_slug, "canonical_id": cid, + "match_tier": match_tier, "tried": [c for c, _t in ladder], + "other_language_pages": [], "ambiguous_pages": [], "candidates": []} + + # Descend the ladder, then the pre-ladder near misses. Stopping at the first rung that holds + # anything is deliberate: a coarser page is a superset only if every writer double-wrote, and + # `tried` records the descent so a thin answer can be told apart from a lucky one. + # Retracted records are dropped as the page is read, not after the rung is chosen. The local + # `resolve` has filtered on `_is_retired` since it existed; this path did not, and reported a + # hardcoded `"retired": 0` while serving them — so a record someone had explicitly taken back + # came straight back out of the service. Filtering here rather than below also means a rung + # whose every entry has been retracted correctly reads as EMPTY and the ladder descends, instead + # of stopping on a page that turns out to have nothing to offer. + def live(canonical_id): + rows = store.candidates(canonical_id, limit=0) + kept = [c for c in rows if not _is_retired(c.value)] + return kept, len(rows) - len(kept) + + found, retired = [], 0 + for cid, match_tier in ladder: + found, retired = live(cid) + if found: + break + if not found: + near = _store_near_misses(store, ladder[0][0]) + for other in near: + found, retired = live(other) + if found: + cid, match_tier = other, "other_version" + break + base_out.update({"other_language_pages": near, + "tried": [c for c, _t in ladder] + near}) + if not found: + return dict(base_out, read_reason="kernel_page_not_found") + base_out.update({"canonical_id": cid, "match_tier": match_tier}) + + try: + min_speedup = float(a.min_speedup) + except (TypeError, ValueError): + min_speedup = 1.0 + above = [c for c in found if (c.speedup or 0.0) >= min_speedup] + # `total` counts what the page held, `retired` how many of those were taken back — so the two + # still sum to the page size even though `found` is already the survivors. + stats = {"total": len(found) + retired, "retired": retired, + "below_min_speedup": len(found) - len(above), "min_speedup": min_speedup} + if not above: + return dict(base_out, filtered=stats, read_reason="below_min_speedup") + + # `above` is already speedup-ordered by the store. + top, alternates, collapsed = collapse_by_direction( + above, lambda c: c.value.get("direction"), lambda c: c.session_id, a.top_n) + stats["same_direction_collapsed"] = collapsed + + cache_dir = a.cache_dir or os.path.join(os.path.dirname(os.path.abspath(a.refs_dir)), "kb_cache") + views = [] + for c, alt_of in zip(top, alternates): + meta = _value_as_meta(c.value, gfx) + metric = meta.get("metric") or {} + # Only now do artifact bytes move: the ranking above read knowledge documents alone. + bundle = store.materialize(cid, c, cache_dir) + views.append({ + "key": c.session_id, + "meta": meta, + "exp_dir": bundle, + "patch_path": os.path.join(bundle, "files", "patch.diff"), + "report_path": os.path.join(bundle, "files", "report.md"), + "speedup": c.speedup or 0.0, + "direction": str(meta.get("direction") or ""), + "bench_key": str(metric.get("bench_key") or ""), + "metric_kind": str(metric.get("metric_kind") or ""), + "origin": f"- session: {c.session_id}{' (champion)' if c.is_champion else ''}\n", + # Alternates are materialized too. They are same-direction runners-up, so there are few + # of them, and a candidate listed with a path that resolves to nothing is worse than not + # listing it: the next reader cannot tell a missing file from a broken export. + "alts": [{ + "session_id": alt.session_id, + "patch_path": os.path.join(store.materialize(cid, alt, cache_dir), + "files", "patch.diff"), + "speedup": round(alt.speedup or 0.0, 4), + "bench_key": str((alt.value.get("metric") or {}).get("bench_key") or ""), + "techniques": _techniques(alt.value), + } for alt in alt_of], + "extra": {"slug": requested_slug, "canonical_id": cid, + "session_id": c.session_id, "is_champion": c.is_champion}, + }) + + summary = ( + f"{len(top)} direction(s) offered from {stats['total']} recorded candidate(s): " + f"{stats['below_min_speedup']} below {min_speedup:g}x, " + f"{collapsed} same-direction re-discoveries moved to `alternates`." + + ({"any_version": f" Served from `{cid}` — the version-agnostic page; nothing was" + " recorded under this box's own ROCm.", + "other_version": f" Served from `{cid}` — a DIFFERENT stack version, and not even the" + " version-agnostic page had it."}.get(match_tier, ""))) + prose = _render_references(a.refs_dir, f"`{cid}`", summary, views) + candidates = [_candidate(rank, v, gfx, p, views[0]["bench_key"]) + for rank, (v, p) in enumerate(zip(views, prose), start=1)] + return dict(base_out, read_reason="read", candidates=candidates, filtered=stats) + + +def cmd_write_remote(a) -> dict: + """Store one measured win in BOTH planes, under the same gates as `write`. + + The local entry stays the source of truth — curation, reproductions and dead ends all live + there — and the KB record is derived from it, so the two cannot drift into disagreeing about + what was measured. What lands under the key depends on the patch, not on the caller: + + * a patch the identity has not seen APPENDS a session, because the session id is a digest of + the patch; the champion pointer then moves only if it beat 1.0x and the incumbent. + * the same patch measured again REPLACES that one session in place. It is a reproduction, + not a second candidate, which is exactly what the local plane already calls it. + """ + local = cmd_write(a) + store, also, why = open_plane(a, CHAMPION_METRIC, 1.0, create=True) + if store is None: + return dict(local, remote={"written": False, "reason": why}) + + exp_dir = local.get("dir") or local.get("reproduced") or "" + meta = _read_meta(os.path.join(exp_dir, "meta.yaml")) if exp_dir else None + if not isinstance(meta, dict): + # No local entry means a gate rejected it (no_improvement / empty_diff / duplicate with an + # unreadable target). Nothing measured, nothing to publish. + return dict(local, remote={"written": False, + "reason": local.get("reason") or "no_local_entry"}) + + recs = remote_records(meta, exp_dir, a.producer, + remote_gpu(_norm_gfx(a.gfx), getattr(a, "gpu", "")), + getattr(a, "framework_version", "")) + files = {f["path"]: f["local_path"] for f in recs[0]["files"]} + # Asked BEFORE the write: a session that already exists is this same patch measured again, and + # the caller deserves to know its result replaced one rather than adding one. + replaced = store.get_session(recs[0]["canonical_id"], recs[0]["session_id"]) is not None + written, promoted, error = publish(store, recs, files, + lambda rec: rec["knowledge"].get("speedup")) + if error: # a KB write must not fail a measured result + return dict(local, remote={"written": False, "partial": written, "reason": error}) + out = { + "written": True, "canonical_id": recs[0]["canonical_id"], + "canonical_ids": written, "session_id": recs[0]["session_id"], + "speedup": recs[0]["knowledge"].get("speedup"), "champion": bool(promoted), + "champion_of": promoted, "files": sorted(files), "store": store.root, + # true = this measurement landed on a session that already existed, i.e. the same patch. + "replaced": replaced, + } + if also is not None: + # The second plane never gates the first. It reports its own outcome so an unreachable + # service is visible as a failed mirror rather than as a silent one. + mirrored, mirror_promoted, mirror_error = publish( + also, recs, files, lambda rec: rec["knowledge"].get("speedup")) + out["mirror"] = {"written": not mirror_error, "store": also.root, + "canonical_ids": mirrored, "champion_of": mirror_promoted, + "reason": mirror_error or ""} + elif why: + out["mirror"] = {"written": False, "reason": why} + return dict(local, remote=out) + + +def main(argv=None): + p = argparse.ArgumentParser(description=__doc__) + sub = p.add_subparsers(dest="cmd", required=True) + + def add_write_args(w): + w.add_argument("--root", required=True) + w.add_argument("--kernel-name", dest="kernel_name", required=True) + w.add_argument("--language", required=True) + w.add_argument("--gfx", required=True) + w.add_argument("--kernel-class", dest="kernel_class", default="unknown") + w.add_argument("--speedup", required=True) + w.add_argument("--baseline-wall-ms", dest="baseline_wall_ms", default=None) + w.add_argument("--patch", default="") + w.add_argument("--eval-dir", dest="eval_dir", default="") + w.add_argument("--report", default="") + w.add_argument("--strategy", default="") + # Curation inputs. Without --direction an entry can never be grouped with its own + # re-discoveries; without the bench fields its speedup compares to nothing. + w.add_argument("--direction", default="") + w.add_argument("--metric-kind", dest="metric_kind", default="") + w.add_argument("--case-names", dest="case_names", default="") + w.add_argument("--parent", default="", + help="exp_dir of the warm-start entry this win was built on") + return w + + add_write_args(sub.add_parser("write", help="store one measured win")) + + r = sub.add_parser("resolve", help="enumerate + rank top-N solutions for a slug") + r.add_argument("--root", required=True) + r.add_argument("--kernel-name", dest="kernel_name", required=True) + r.add_argument("--language", required=True) + r.add_argument("--gfx", required=True) + r.add_argument("--top-n", dest="top_n", type=int, default=3, help="max DIRECTIONS to offer") + r.add_argument("--refs-dir", dest="refs_dir", required=True) + r.add_argument("--match", choices=("exact", "normalized", "fuzzy"), default="fuzzy", + help="how hard to try to map the caller's kernel name onto a page (default fuzzy)") + r.add_argument("--min-speedup", dest="min_speedup", type=float, default=1.05, + help="never spend an on-box verify on a recorded win below this (default 1.05)") + r.add_argument("--include-retired", dest="include_retired", action="store_true", + help="also offer entries the curation retired (audit/debug only)") + + lg = sub.add_parser("languages", help="which languages this kernel has a page in") + lg.add_argument("--root", required=True) + lg.add_argument("--kernel-name", dest="kernel_name", required=True) + lg.add_argument("--gfx", required=True) + + bf = sub.add_parser("backfill-content", help="bring imported entries up to the current shape") + bf.add_argument("--root", required=True) + bf.add_argument("--apply", action="store_true", help="write; without it, only report the diff") + + xr = sub.add_parser("export-remote", help="render entries as KB Store candidates (JSON lines)") + xr.add_argument("--root", required=True) + xr.add_argument("--gfx", default="", help="only this arch (default: every arch in the store)") + xr.add_argument("--kernel-name", dest="kernel_name", default="", help="only this kernel") + xr.add_argument("--producer", default=REMOTE_PRODUCER, + help="the system that owns this candidate stream and its champion pointer") + xr.add_argument("--gpu", default="", help="override the gfx dimension; default is the entry's own gfx") + xr.add_argument("--include-retired", dest="include_retired", action="store_true", + help="also export entries the curation retired (they would rank as live wins)") + xr.add_argument("--out", default="", help="write JSON lines here instead of stdout") + + # The key-addressed pair. Same gates, same output shapes as resolve/write — only the plane + # the records live on changes, so the lane can be pointed at either. + def add_plane_args(w): + # `both` writes locally and mirrors to the service; reads always take exactly one plane. + w.add_argument("--plane", choices=("local", "remote", "both"), default="local", + help="local dir, the KB Store service (KB_STORE_URL/KB_STORE_TOKEN), or both") + w.add_argument("--scan", type=int, default=25, + help="remote only: candidates hydrated before curation (page cap is 200)") + return w + + rr = add_plane_args(sub.add_parser("resolve-remote", + help="rank top-N candidates under one canonical id")) + rr.add_argument("--store", default="", help="on-disk KB store root (--plane local/both)") + rr.add_argument("--canonical-id", dest="canonical_id", default="", + help="the key to read; derived from kernel/language/gfx when omitted") + rr.add_argument("--kernel-name", dest="kernel_name", default="") + rr.add_argument("--language", default="") + rr.add_argument("--gfx", default="") + rr.add_argument("--producer", default=REMOTE_PRODUCER) + rr.add_argument("--gpu", default="", help="override the gfx dimension; default is --gfx") + rr.add_argument("--framework-version", dest="framework_version", default="", + help="rocm .; default is detected on this box") + rr.add_argument("--top-n", dest="top_n", type=int, default=3, help="max DIRECTIONS to offer") + rr.add_argument("--refs-dir", dest="refs_dir", required=True) + rr.add_argument("--cache-dir", dest="cache_dir", default="", + help="where selected candidates are materialized (default /../kb_cache)") + rr.add_argument("--min-speedup", dest="min_speedup", type=float, default=1.05) + + wr = add_plane_args(add_write_args( + sub.add_parser("write-remote", help="store one win in the local store AND under its key"))) + wr.add_argument("--store", default="", help="on-disk KB store root (--plane local/both)") + wr.add_argument("--producer", default=REMOTE_PRODUCER) + wr.add_argument("--gpu", default="", help="override the gfx dimension; default is --gfx") + wr.add_argument("--framework-version", dest="framework_version", default="", + help="rocm . for the key; default is the measured stack") + + tr = add_plane_args(sub.add_parser( + "retract-remote", help="take back a written record: retained=false, score zeroed, champion " + "re-pointed (there is no delete — this is a rewrite)")) + tr.add_argument("--store", default="", help="on-disk KB store root (--plane local/both)") + tr.add_argument("--canonical-id", dest="canonical_id", default="", + help="retract on THIS page only; omit to walk both rungs of the ladder") + tr.add_argument("--session-id", dest="session_id", required=True, + help="the session to retract, from the write-remote output") + tr.add_argument("--reason", required=True, + help="why the record is wrong; it is all a future reader has to judge by") + tr.add_argument("--kernel-name", dest="kernel_name", default="") + tr.add_argument("--language", default="") + tr.add_argument("--gfx", default="") + tr.add_argument("--producer", default=REMOTE_PRODUCER) + tr.add_argument("--gpu", default="", help="override the gfx dimension; default is --gfx") + tr.add_argument("--framework-version", dest="framework_version", default="") + tr.add_argument("--measured-by", dest="measured_by", default="", help="who is retracting it") + tr.add_argument("--apply", action="store_true", help="actually rewrite; default is a dry run") + + at = sub.add_parser("attest", help="count one attempt to USE a stored entry (validated | " + "failed | not_reproduced); changes no speedup, no rank") + at.add_argument("--exp-dir", dest="exp_dir", required=True, + help="the entry that was tried, as `resolve` reports it") + at.add_argument("--outcome", required=True, choices=_OUTCOMES, + help="validated = reproduced a win; failed = applied but did not win; " + "not_reproduced = would not apply or would not build") + at.add_argument("--measured-speedup", dest="measured_speedup", default=None, + help="what it did here, for the history entry") + at.add_argument("--note", default="", help="one line a future reader can act on") + at.add_argument("--measured-by", dest="measured_by", default="", help="who tried it") + at.add_argument("--apply", action="store_true", help="actually record it; default is a dry run") + + ar = add_plane_args(sub.add_parser( + "attest-remote", help="the same count, against a key-addressed record on either plane")) + ar.add_argument("--store", default="", help="on-disk KB store root (--plane local/both)") + ar.add_argument("--canonical-id", dest="canonical_id", default="", + help="attest on THIS page only; omit to walk both rungs of the ladder") + ar.add_argument("--session-id", dest="session_id", required=True, + help="the session that was tried, from the write-remote output") + ar.add_argument("--outcome", required=True, choices=_OUTCOMES) + ar.add_argument("--kernel-name", dest="kernel_name", default="") + ar.add_argument("--language", default="") + ar.add_argument("--gfx", default="") + ar.add_argument("--producer", default=REMOTE_PRODUCER) + ar.add_argument("--gpu", default="", help="override the gfx dimension; default is --gfx") + ar.add_argument("--framework-version", dest="framework_version", default="") + ar.add_argument("--measured-speedup", dest="measured_speedup", default=None) + ar.add_argument("--note", default="") + ar.add_argument("--measured-by", dest="measured_by", default="", help="who tried it") + ar.add_argument("--apply", action="store_true", help="actually record it; default is a dry run") + + m = sub.add_parser("remap", help="rewrite a stored patch's paths onto this workspace's layout") + m.add_argument("--patch", required=True) + m.add_argument("--out", required=True) + m.add_argument("--editable", default="", help="comma-separated workspace-relative editable paths") + m.add_argument("--workspace", default="", help="scanned for source files when --editable is empty") + + a = p.parse_args(argv) + try: + if a.cmd == "write": + out = cmd_write(a) + elif a.cmd == "resolve": + out = cmd_resolve(a) + elif a.cmd == "remap": + out = cmd_remap(a) + elif a.cmd == "languages": + out = cmd_languages(a) + elif a.cmd == "backfill-content": + out = cmd_backfill_content(a) + elif a.cmd == "export-remote": + out = cmd_export_remote(a) + elif a.cmd == "resolve-remote": + out = cmd_resolve_remote(a) + elif a.cmd == "write-remote": + out = cmd_write_remote(a) + elif a.cmd == "retract-remote": + out = cmd_retract_remote(a) + elif a.cmd == "attest": + out = cmd_attest(a) + elif a.cmd == "attest-remote": + out = cmd_attest_remote(a) + else: # pragma: no cover + out = {"error": "unknown command"} + except Exception as e: # never crash the caller + err = "exception: " + str(e)[:160] + out = ({"written": False, "reason": err} if a.cmd in ("write", "write-remote") + else {"retracted": False, "reason": err} if a.cmd == "retract-remote" + else {"attested": False, "reason": err} if a.cmd in ("attest", "attest-remote") + else {"remapped": False, "reason": err} if a.cmd == "remap" + else {"read_reason": err, "candidates": []}) + print(json.dumps(out, ensure_ascii=False)) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/kernel_workflow/scripts/tests/test_experience_store.py b/kernel_workflow/scripts/tests/test_experience_store.py new file mode 100644 index 000000000..e388134f0 --- /dev/null +++ b/kernel_workflow/scripts/tests/test_experience_store.py @@ -0,0 +1,1085 @@ +"""Tests for the local experience store (kernel_workflow/scripts/experience_store.py). + +No GPU, no network, no repo state: every case builds its own kb root in tmp_path. The store is the +read/write half of warm start that decides which historical patch a lane spends an on-box verify on, +so what is pinned here is the SELECTION, not the plumbing: + - identity: the same kernel under three different dir layouts is ONE page (read and write agree); + - curation: a retired entry is never offered, one rank per direction, near-ties are not read; + - honesty: a speedup carries the bench it was measured on, and cross-bench ranks say so; + - the loop: re-writing code the store already holds is a reproduction, not a new entry. +""" + +import json +import os +import subprocess +import sys + +import pytest + +STORE = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "experience_store.py") +yaml = pytest.importorskip("yaml") + + +# --------------------------------------------------------------------------- helpers +def run(*args): + """Invoke the store as the lane does — over the CLI — and parse its single-line JSON.""" + p = subprocess.run([sys.executable, STORE] + [str(a) for a in args], + capture_output=True, text=True) + assert p.returncode == 0, f"store must never fail the caller: {p.stderr}" + return json.loads(p.stdout) + + +def patch_text(path="source/k.py", old="BLOCK = 64", new="BLOCK = 128", at=1): + return (f"diff --git a/{path} b/{path}\nindex 111..222 100644\n" + f"--- a/{path}\n+++ b/{path}\n@@ -{at},3 +{at},3 @@\n ctx\n-{old}\n+{new}\n") + + +def write_entry(root, exp_id, *, kernel="fused_moe_kernel", lang="triton", gfx="gfx950", + kclass="triton", speedup=2.0, direction="tile-retune", retired=False, + bench="b:aaa", reproductions=1, lifecycle="candidate", patch=None): + """Seed one entry directly on disk, the way the imported backlog looks.""" + d = os.path.join(root, gfx, kclass, f"{kernel}__{lang}__{gfx}", exp_id) + os.makedirs(d, exist_ok=True) + meta = { + "layer": "artifact", "lifecycle": lifecycle, "gfx": gfx, "kernel_class": kclass, + "kernel_name": kernel, "language": lang, "direction": direction, + "reproductions": reproductions, + "metric": {"speedup": speedup, "gpu_arch": gfx, "bench_key": bench, "metric_kind": "geomean"}, + "strategy": f"strategy for {exp_id}", + } + if retired: + meta["retained"] = False + meta["retired_reason"] = f"duplicate_direction:{direction}" + with open(os.path.join(d, "meta.yaml"), "w") as f: + yaml.safe_dump(meta, f) + with open(os.path.join(d, "patch.diff"), "w") as f: + f.write(patch if patch is not None else patch_text(new=f"BLOCK = {exp_id}")) + return d + + +def resolve(root, refs, kernel="fused_moe_kernel", lang="triton", gfx="gfx950", *extra): + return run("resolve", "--root", root, "--kernel-name", kernel, "--language", lang, + "--gfx", gfx, "--refs-dir", refs, *extra) + + +# --------------------------------------------------------------------------- identity +@pytest.mark.parametrize("name", [ + "fused_moe_kernel", # standalone lane: the kernel dir + "fused_moe_kernel_task", # e2e head path: the EXTRACTED task dir + "triton_fused_moe_kernel.py", # a producer that carries the language in the filename + "/abs/path/to/fused_moe_kernel", # a lane that hands over a path +]) +def test_one_kernel_is_one_page_across_layouts(tmp_path, name): + """If these forked into separate pages, an e2e head run could never find the history of the + kernel it is optimizing.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + write_entry(root, "20260101_000000_aaaaaa", speedup=3.0) + out = resolve(root, refs, kernel=name) + assert out["read_reason"] == "read", out + assert out["slug"] == "fused_moe_kernel__triton__gfx950" + assert out["match_tier"] in ("exact", "normalized") + assert len(out["candidates"]) == 1 + + +def test_write_and_read_derive_the_same_page(tmp_path): + """The docstring's contract: read and write MUST canonicalize identically.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + p = tmp_path / "p.diff" + p.write_text(patch_text()) + w = run("write", "--root", root, "--kernel-name", "triton_demo_kernel.py", "--language", "triton", + "--gfx", "gfx950", "--kernel-class", "triton", "--speedup", "2.5", "--patch", str(p), + "--direction", "vectorize") + assert w["written"] and w["slug"] == "demo_kernel__triton__gfx950" + out = resolve(root, refs, kernel="demo_kernel_task") # the e2e head path's name for it + assert out["match_tier"] == "exact" and len(out["candidates"]) == 1 + + +def test_op_kind_reaches_the_kernel_page_only_when_unambiguous(tmp_path): + """e2e names a head by op_kind (`fused_moe`); the page is `fused_moe_kernel`. Fuzzy bridges that + — but never guesses between two pages that are equally close.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + write_entry(root, "20260101_000000_aaaaaa") + write_entry(root, "20260101_000000_bbbbbb", kernel="fused_moe_kernel_gptq_awq") + out = resolve(root, refs, kernel="fused_moe") + assert out["match_tier"] == "fuzzy" and out["slug"] == "fused_moe_kernel__triton__gfx950", \ + "the CLOSEST containing page wins, not just any page sharing the prefix" + + # two pages the query fits equally well: no basis to choose, so choose neither + write_entry(root, "20260101_000000_cccccc", kernel="fused_moe_alpha") + write_entry(root, "20260101_000000_dddddd", kernel="fused_moe_bravo") + amb = resolve(root, refs, kernel="fused_moe") + assert amb["read_reason"] == "ambiguous_kernel_page" and not amb["candidates"] + assert len(amb["ambiguous_pages"]) == 2 + + # --match exact opts out of all of it: only a page whose slug is literally the requested one. + strict = resolve(root, refs, "fused_moe", "triton", "gfx950", "--match", "exact") + assert strict["read_reason"] == "kernel_page_not_found" and not strict["candidates"] + + +def test_wrong_language_reports_the_page_it_did_find(tmp_path): + """A hip kernel optimized by a lane defaulted to triton must not look like an empty store — + that is the difference between 'no history' and 'this lane was invoked wrong'.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + write_entry(root, "20260101_000000_aaaaaa", kernel="wvSplitK", lang="hip", kclass="hip") + out = resolve(root, refs, kernel="wvsplitk", lang="triton") + assert out["read_reason"] == "no_page_for_language" + assert out["other_language_pages"] == ["wvSplitK__hip__gfx950"] + assert not out["candidates"] + + +# --------------------------------------------------------------------------- curation +def test_retired_entries_are_never_offered(tmp_path): + """`retained: false` is the curation's verdict that a better entry of the same idea is already + served. Ranking by raw speedup would spend verifies re-testing what was rejected.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + write_entry(root, "20260101_000000_hi", speedup=50.0, direction="fp8-bitcast", retired=True) + write_entry(root, "20260101_000001_lo", speedup=4.0, direction="fp8-bitcast") + out = resolve(root, refs) + assert [c["speedup"] for c in out["candidates"]] == [4.0] + assert out["filtered"]["retired"] == 1 + + audit = resolve(root, refs, "fused_moe_kernel", "triton", "gfx950", "--include-retired") + assert audit["candidates"][0]["speedup"] == 50.0 + + +def test_one_rank_per_direction_runners_up_ride_along(tmp_path): + """Three impls of one idea verify or fail together, at one full measurement each. The store + offers the idea once and keeps the rest reachable as alternates.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + write_entry(root, "20260101_000000_a", speedup=9.0, direction="fp8-bitcast") + write_entry(root, "20260101_000001_b", speedup=8.0, direction="fp8-bitcast") + write_entry(root, "20260101_000002_c", speedup=7.0, direction="fp8-bitcast") + write_entry(root, "20260101_000003_d", speedup=2.0, direction="tile-retune") + out = resolve(root, refs, "fused_moe_kernel", "triton", "gfx950", "--top-n", "3") + assert [c["direction"] for c in out["candidates"]] == ["fp8-bitcast", "tile-retune"] + assert [a["speedup"] for a in out["candidates"][0]["alternates"]] == [8.0, 7.0] + assert out["filtered"]["same_direction_collapsed"] == 2 + + +def test_undirected_entries_are_each_their_own_direction(tmp_path): + """An entry written without a direction label must not collapse with every other such entry.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + write_entry(root, "20260101_000000_a", speedup=9.0, direction="") + write_entry(root, "20260101_000001_b", speedup=8.0, direction="") + out = resolve(root, refs) + assert len(out["candidates"]) == 2 + + +def test_near_ties_are_not_worth_a_measurement(tmp_path): + """Reading a recorded 1.02x costs the same full verify as reading a 50x.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + write_entry(root, "20260101_000000_a", speedup=1.02, direction="mask-elision") + out = resolve(root, refs) + assert out["read_reason"] == "below_min_speedup" and not out["candidates"] + assert resolve(root, refs, "fused_moe_kernel", "triton", "gfx950", + "--min-speedup", "1.0")["candidates"][0]["speedup"] == 1.02 + + +def test_rank_order_carries_bench_comparability(tmp_path): + """58x on one case set and 5x on another are not ordered facts. Adoption is decided by a fresh + measurement anyway — but the ranking must not silently claim otherwise.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + write_entry(root, "20260101_000000_a", speedup=58.0, direction="fp8", bench="b:one") + write_entry(root, "20260101_000001_b", speedup=5.0, direction="tiles", bench="b:two") + write_entry(root, "20260101_000002_c", speedup=3.0, direction="loads", bench="b:one") + out = resolve(root, refs) + assert [c["comparable"] for c in out["candidates"]] == [True, False, True] + + +def test_reproduced_entry_outranks_a_one_off_tie(tmp_path): + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + write_entry(root, "20260101_000000_a", speedup=4.0, direction="a", reproductions=1) + write_entry(root, "20260101_000001_b", speedup=4.0, direction="b", reproductions=3) + out = resolve(root, refs) + assert out["candidates"][0]["exp_dir"].endswith("20260101_000001_b") + + +def test_a_fully_retired_page_offers_nothing(tmp_path): + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + write_entry(root, "20260101_000000_a", speedup=4.0, retired=True) + out = resolve(root, refs) + assert out["read_reason"] == "all_retired" and not out["candidates"] + + +# --------------------------------------------------------------------------- the loop +def test_rewriting_known_code_counts_a_reproduction(tmp_path): + """A warm-started run re-emits the patch it adopted as its own diff: different workspace, + byte-different, same code. Otherwise the store keeps re-importing its own output as a fresh win.""" + root = str(tmp_path / "kb") + seeded = write_entry(root, "20260101_000000_a", speedup=3.0, patch=patch_text()) + # same change, e2e head layout + a different hunk offset + reflowed whitespace + p = tmp_path / "again.diff" + p.write_text(patch_text(path="kernel_src/sglang/k.py", old="BLOCK = 64", new="BLOCK = 128", at=87)) + w = run("write", "--root", root, "--kernel-name", "fused_moe_kernel_task", "--language", "triton", + "--gfx", "gfx950", "--kernel-class", "triton", "--speedup", "2.9", "--patch", str(p), + "--direction", "tile-retune", "--eval-dir", "/tmp/run2") + assert w["written"] is False and w["reason"] == "duplicate_impl" + assert w["reproductions"] == 2 and w["lifecycle"] == "active" + meta = yaml.safe_load(open(os.path.join(seeded, "meta.yaml"))) + assert meta["metric"]["speedup"] == 3.0, "the original's own measurement must not be overwritten" + assert meta["reproductions"] == 2 and meta["lifecycle"] == "active" + assert len(os.listdir(os.path.dirname(seeded))) == 1, "no second entry for the same code" + + +def test_genuinely_new_code_is_a_new_entry(tmp_path): + root = str(tmp_path / "kb") + seeded = write_entry(root, "20260101_000000_a", speedup=3.0, patch=patch_text()) + p = tmp_path / "new.diff" + p.write_text(patch_text(new="BLOCK = 256")) + w = run("write", "--root", root, "--kernel-name", "fused_moe_kernel", "--language", "triton", + "--gfx", "gfx950", "--kernel-class", "triton", "--speedup", "3.5", "--patch", str(p), + "--direction", "tile-retune") + assert w["written"] is True + assert len(os.listdir(os.path.dirname(seeded))) == 2 + + +def test_write_records_what_curation_needs(tmp_path): + root = str(tmp_path / "kb") + p = tmp_path / "p.diff" + p.write_text(patch_text()) + w = run("write", "--root", root, "--kernel-name", "k", "--language", "triton", "--gfx", "gfx950", + "--kernel-class", "triton", "--speedup", "2.0", "--patch", str(p), + "--direction", "Vectorize Loads", "--metric-kind", "time_weighted", + "--case-names", "c64,c2", "--parent", "/kb/old/exp") + meta = yaml.safe_load(open(os.path.join(w["dir"], "meta.yaml"))) + assert meta["direction"] == "Vectorize Loads" + assert meta["derived_from"] == "/kb/old/exp" + assert meta["metric"]["case_names"] == ["c64", "c2"] + # bench key is order-insensitive over the case set, and namespaced away from imported `b:` keys + other = run("write", "--root", root, "--kernel-name", "k2", "--language", "triton", "--gfx", + "gfx950", "--kernel-class", "triton", "--speedup", "2.0", "--patch", str(p), + "--metric-kind", "time_weighted", "--case-names", "c2,c64") + bk = yaml.safe_load(open(os.path.join(other["dir"], "meta.yaml")))["metric"]["bench_key"] + assert bk == meta["metric"]["bench_key"] and bk.startswith("b2:") + + +# --------------------------------------------------------------------------- path remapping +def remap(tmp_path, patch, editable="", **kw): + p = tmp_path / "in.diff" + p.write_text(patch) + args = ["remap", "--patch", str(p), "--out", str(tmp_path / "out.diff")] + if editable: + args += ["--editable", editable] + for k, v in kw.items(): + args += ["--" + k.replace("_", "-"), v] + return run(*args), str(tmp_path / "out.diff") + + +def test_a_stored_patch_is_rewritten_onto_the_head_layout(tmp_path): + """The store's patches were won in an arena checkout (`source/triton_.py`); an e2e head + run edits an extracted subtree. Different prefix AND basename, so no -p depth reaches the + file — without this every warm start on the head path fails to apply.""" + out, dest = remap(tmp_path, patch_text(path="source/triton_fused_moe_kernel.py"), + "kernel_src/sglang/layers/moe/fused_moe_kernel.py,kernel_src/other.py") + assert out["remapped"] is True, out + body = open(dest).read() + assert "--- a/kernel_src/sglang/layers/moe/fused_moe_kernel.py" in body + assert "+++ b/kernel_src/sglang/layers/moe/fused_moe_kernel.py" in body + assert "diff --git a/kernel_src/sglang/layers/moe/fused_moe_kernel.py " in body + assert "source/" not in body and "-BLOCK = 64" in body, "hunks must pass through untouched" + + +def test_remap_scans_the_workspace_when_no_editable_set_is_known(tmp_path): + ws = tmp_path / "ws" / "kernel_src" / "moe" + ws.mkdir(parents=True) + (ws / "fused_moe_kernel.py").write_text("x\n") + out, dest = remap(tmp_path, patch_text(path="source/triton_fused_moe_kernel.py"), + workspace=str(tmp_path / "ws")) + assert out["remapped"] is True + assert "+++ b/kernel_src/moe/fused_moe_kernel.py" in open(dest).read() + + +def test_remap_refuses_to_guess_between_two_equally_good_homes(tmp_path): + out, _ = remap(tmp_path, patch_text(path="source/triton_k.py"), "a/k.py,b/k.py") + assert out["remapped"] is False and out["reason"] == "unmapped_paths" + + +def test_remap_is_all_or_nothing(tmp_path): + """Applying the mapped half of a patch leaves a workspace that is neither before nor after.""" + two = (patch_text(path="source/triton_k.py") + patch_text(path="3rdparty/ck/deep.hpp")) + out, dest = remap(tmp_path, two, "kernel_src/k.py") + assert out["remapped"] is False and out["unmapped"] == ["3rdparty/ck/deep.hpp"] + assert not os.path.exists(dest), "a refused remap must not leave a half-rewritten patch behind" + + +def test_remap_reports_when_the_paths_already_fit(tmp_path): + out, _ = remap(tmp_path, patch_text(path="kernel_src/k.py"), "kernel_src/k.py") + assert out["remapped"] is False and out["reason"] == "no_change_needed" + + +def test_a_new_file_lands_beside_the_file_the_patch_edits(tmp_path): + """A created file has nothing to match against, so it follows its edited sibling.""" + add = ("diff --git a/source/helper.py b/source/helper.py\nnew file mode 100644\n" + "--- /dev/null\n+++ b/source/helper.py\n@@ -0,0 +1,1 @@\n+HELPER = 1\n") + out, dest = remap(tmp_path, patch_text(path="source/triton_k.py") + add, "kernel_src/moe/k.py") + assert out["remapped"] is True + assert out["mapped"]["source/helper.py"] == "kernel_src/moe/helper.py" + assert "+++ b/kernel_src/moe/helper.py" in open(dest).read() + + +def test_a_new_file_stays_put_when_the_layout_already_matches(tmp_path): + """The store's own arena layout IS this workspace's layout in the common re-run case. Nothing + shifted, so a root-level file the patch creates must not be dragged into the sibling's dir.""" + add = ("diff --git a/profile_run.py b/profile_run.py\nnew file mode 100644\n" + "--- /dev/null\n+++ b/profile_run.py\n@@ -0,0 +1,1 @@\n+import torch\n") + out, _ = remap(tmp_path, patch_text(path="source/triton_k.py") + add, "source/triton_k.py") + assert out["remapped"] is False and out["reason"] == "no_change_needed" + + +def test_a_non_source_file_this_workspace_lacks_is_dropped_not_refused(tmp_path): + """A `.gitignore` hunk rode along with 4 real store patches; refusing them over it would throw + away the kernel work. The section is dropped, the kernel edit survives, and it is reported.""" + ignore = ("diff --git a/.gitignore b/.gitignore\n--- a/.gitignore\n+++ b/.gitignore\n" + "@@ -5,3 +5,4 @@\n *.o\n+.nfs*\n") + out, dest = remap(tmp_path, ignore + patch_text(path="source/triton_k.py"), + "kernel_src/moe/k.py") + assert out["remapped"] is True and out["dropped"] == [".gitignore"] + body = open(dest).read() + assert ".gitignore" not in body and "+++ b/kernel_src/moe/k.py" in body + + +@pytest.mark.parametrize("patch, reason", [ + ("diff --git a/x b/y\nrename from x\nrename to y\n", "rename_not_supported"), + ("not a diff at all\n", "no_paths_in_patch"), +]) +def test_remap_refuses_what_it_cannot_rewrite(tmp_path, patch, reason): + out, _ = remap(tmp_path, patch, "kernel_src/k.py") + assert out["remapped"] is False and out["reason"] == reason + + +def test_remap_never_raises(tmp_path): + out = run("remap", "--patch", str(tmp_path / "missing.diff"), "--out", str(tmp_path / "o.diff"), + "--editable", "k.py") + assert out["remapped"] is False and out["reason"].startswith("unreadable_patch") + + +# --------------------------------------------------------------------------- degradation +@pytest.mark.parametrize("args, reason", [ + (("--gfx", "cpu"), "missing_arch"), + (("--gfx", "gfx942"), "kernel_page_not_found"), # arch present in the request, absent on disk +]) +def test_resolve_never_raises(tmp_path, args, reason): + """The lane calls this over Bash mid-run: a store problem must degrade to a cold start, never + fail the run.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + write_entry(root, "20260101_000000_a") + out = run("resolve", "--root", root, "--kernel-name", "fused_moe_kernel", "--language", "triton", + "--refs-dir", refs, *args) + assert out["read_reason"] == reason and out["candidates"] == [] + + +def test_resolve_survives_a_corrupt_entry(tmp_path): + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + good = write_entry(root, "20260101_000000_a", speedup=2.0) + bad = os.path.join(os.path.dirname(good), "20260101_000001_b") + os.makedirs(bad) + with open(os.path.join(bad, "meta.yaml"), "w") as f: + f.write("{{{ not yaml") + out = resolve(root, refs) + assert [c["exp_dir"] for c in out["candidates"]] == [good] + + +def test_prose_is_mirrored_for_audit_before_any_verdict(tmp_path): + """A rejected warm start must still be auditable after the run.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + d = write_entry(root, "20260101_000000_a", speedup=2.0) + with open(os.path.join(d, "report.md"), "w") as f: + f.write("# why this worked\nfolded the scale\n") + out = resolve(root, refs) + prose = open(out["candidates"][0]["prose_path"]).read() + assert "folded the scale" in prose and "tile-retune" in prose + assert "direction" in open(os.path.join(refs, "index.md")).read() + + +# --------------------------------------------------------------------------- content +# What an entry CARRIES, and what of it reaches the agent. The store's densest fields (`techniques`, +# the report's dead-ends) were being written and never read; these pin them to the reference. +def run_lines(*args): + """For subcommands that stream one JSON line per record and a summary line last.""" + p = subprocess.run([sys.executable, STORE] + [str(a) for a in args], + capture_output=True, text=True) + assert p.returncode == 0, p.stderr + lines = [json.loads(l) for l in p.stdout.splitlines() if l.strip()] + return lines[:-1], lines[-1] + + +REPORT = """# TechLead Final Report — k + +## Summary + +geomean 3.0x + +## Round-by-round + +### Round 1 + +r1 text + +{heading} + +- `MPerBlock=128`: built, correct, 27% slower. +- occupancy pinning: closed eight times. + +## Stop rationale + +out of budget +""" + + +@pytest.mark.parametrize("heading", [ + "## What didn't work", + "## What didn't work (dead-ends)", + "## What didn't work (dead ends — do not re-fund)", + "## What didn’t work (confirmed dead-ends — do not re-open)", + "### What Didn't Work (dead-ends from the ledger)", +]) +def test_dead_ends_section_survives_every_heading_dialect(tmp_path, heading): + """248 imported reports write this heading 20+ different ways; matching the full title would + silently drop the one section that stops the next run re-funding a closed direction.""" + import importlib.util + spec = importlib.util.spec_from_file_location("es", STORE) + es = importlib.util.module_from_spec(spec) + spec.loader.exec_module(es) + text = REPORT.format(heading=heading) + body = es.dead_ends_md(text) + assert "MPerBlock=128" in body and "occupancy pinning" in body + assert "out of budget" not in body, "the section must end at the next heading" + hoisted = es.reorder_report(text) + assert hoisted.index(heading) < hoisted.index("## Summary") + assert sorted(hoisted.split()) == sorted((text + "\n---\n").split()), "nothing may be dropped" + + +def test_report_without_the_two_sections_is_passed_through_untouched(tmp_path): + import importlib.util + spec = importlib.util.spec_from_file_location("es", STORE) + es = importlib.util.module_from_spec(spec) + spec.loader.exec_module(es) + plain = "# report\n\n## Summary\n\nnothing else\n" + assert es.reorder_report(plain) == plain + assert es.dead_ends_md(plain) == "" + + +def test_techniques_and_dead_ends_reach_the_reference(tmp_path): + """The curated techniques are the densest thing in the store; before this they were written and + never shown to anyone.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + d = write_entry(root, "20260101_000000_a", speedup=3.0) + meta = yaml.safe_load(open(os.path.join(d, "meta.yaml"))) + meta["techniques"] = ["bitcast operands to OCP e4m3fn", "retile BM=256"] + with open(os.path.join(d, "meta.yaml"), "w") as f: + yaml.safe_dump(meta, f) + with open(os.path.join(d, "report.md"), "w") as f: + f.write(REPORT.format(heading="## What didn't work (dead-ends)")) + + out = resolve(root, refs) + cand = out["candidates"][0] + assert cand["techniques"] == ["bitcast operands to OCP e4m3fn", "retile BM=256"] + prose = open(cand["prose_path"]).read() + assert "- techniques:\n * bitcast operands to OCP e4m3fn" in prose + # the two load-bearing sections come before the narrative, so a tight context reads them first + assert prose.index("What didn't work") < prose.index("## Summary") + + +def test_reference_omits_techniques_when_there_are_none(tmp_path): + """An empty heading is worse than no heading: it reads as 'this patch does nothing'.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + write_entry(root, "20260101_000000_a", speedup=3.0) + prose = open(resolve(root, refs)["candidates"][0]["prose_path"]).read() + assert "techniques" not in prose + + +def test_write_keeps_structured_dead_ends_when_the_report_supplies_them(tmp_path): + root = str(tmp_path / "kb") + p = tmp_path / "p.diff" + p.write_text(patch_text()) + rep = tmp_path / "report.md" + rep.write_text(REPORT.format(heading="## What didn't work") + """ + +```yaml +- idea: use_buffer_ops=OFF negative control + measured: 0.883x + mechanism: the ambient default is load-bearing, -11.7% +- not-a-dict +- {mechanism: no idea field, measured: 1.0x} +``` +""") + w = run("write", "--root", root, "--kernel-name", "k", "--language", "triton", "--gfx", "gfx950", + "--kernel-class", "triton", "--speedup", "2.0", "--patch", str(p), "--report", str(rep)) + meta = yaml.safe_load(open(os.path.join(w["dir"], "meta.yaml"))) + assert [d["idea"] for d in meta["dead_ends"]] == ["use_buffer_ops=OFF negative control"], \ + "malformed rows are dropped, never patched up into half-empty structure" + assert meta["dead_ends"][0]["measured"] == "0.883x" + assert "MPerBlock=128" in meta["dead_ends_md"] + assert "dead-ends:yaml" not in meta["dead_ends_md"], "the machine block is not prose" + + +def test_write_falls_back_to_prose_when_there_is_no_structured_block(tmp_path): + root = str(tmp_path / "kb") + p = tmp_path / "p.diff" + p.write_text(patch_text()) + rep = tmp_path / "report.md" + rep.write_text(REPORT.format(heading="## What didn't work (dead ends)")) + w = run("write", "--root", root, "--kernel-name", "k", "--language", "triton", "--gfx", "gfx950", + "--kernel-class", "triton", "--speedup", "2.0", "--patch", str(p), "--report", str(rep)) + meta = yaml.safe_load(open(os.path.join(w["dir"], "meta.yaml"))) + assert "dead_ends" not in meta, "no structure is honest; invented structure is not" + assert "MPerBlock=128" in meta["dead_ends_md"] + + +# --------------------------------------------------------------------------- backfill +def test_backfill_adds_content_without_touching_curation(tmp_path): + root = str(tmp_path / "kb") + d = write_entry(root, "20260101_000000_a", speedup=3.0, direction="tile-retune", retired=True) + with open(os.path.join(d, "report.md"), "w") as f: + f.write(REPORT.format(heading="## What didn't work (dead-ends)")) + before = yaml.safe_load(open(os.path.join(d, "meta.yaml"))) + + _rows, summary = run_lines("backfill-content", "--root", root) + assert summary["changed"] == 1 and summary["applied"] is False + assert yaml.safe_load(open(os.path.join(d, "meta.yaml"))) == before, "dry run writes nothing" + + _rows, summary = run_lines("backfill-content", "--root", root, "--apply") + after = yaml.safe_load(open(os.path.join(d, "meta.yaml"))) + assert "MPerBlock=128" in after["dead_ends_md"] + assert after["verified_stack"]["triton"] == "3.6.0" + assert after["verified_stack"]["recorded_by"] == "campaign20_backfill", \ + "a recovered stack must not read as one we observed" + assert "layer" not in after + for k in ("retained", "retired_reason", "direction", "metric", "reproductions", "lifecycle"): + assert after.get(k) == before.get(k), f"backfill must not touch {k}" + + _rows, summary = run_lines("backfill-content", "--root", root, "--apply") + assert summary["changed"] == 0, "backfill must be idempotent" + + +def test_backfill_recomputes_the_signature_the_dedup_actually_reads(tmp_path): + """The backlog carries `impl_signature`, a different hash under a name nothing reads, so every + resolve re-hashed all 248 patches. Renaming is not enough — the value must be recomputed.""" + root = str(tmp_path / "kb") + d = write_entry(root, "20260101_000000_a", speedup=3.0, patch=patch_text()) + meta = yaml.safe_load(open(os.path.join(d, "meta.yaml"))) + meta["impl_signature"] = "sha256:deadbeef" + with open(os.path.join(d, "meta.yaml"), "w") as f: + yaml.safe_dump(meta, f) + + run_lines("backfill-content", "--root", root, "--apply") + after = yaml.safe_load(open(os.path.join(d, "meta.yaml"))) + assert "impl_signature" not in after + assert after["content_signature"].startswith("csha:") + + p = tmp_path / "again.diff" + p.write_text(patch_text(path="other/layout/k.py", at=99)) + w = run("write", "--root", root, "--kernel-name", "fused_moe_kernel", "--language", "triton", + "--gfx", "gfx950", "--kernel-class", "triton", "--speedup", "3.1", "--patch", str(p)) + assert w["reason"] == "duplicate_impl" and w["reproductions"] == 2 + + +# --------------------------------------------------------------------------- language lookup +def test_languages_lets_the_store_pick_the_language(tmp_path): + """A caller that guesses `triton` for a kernel filed under `hip`/`ck` reads nothing and never + learns why; `task_type: hip2hip` cannot tell those two apart, but the pages can.""" + root = str(tmp_path / "kb") + write_entry(root, "20260101_000000_a", kernel="wvSplitK", lang="hip", kclass="hip") + write_entry(root, "20260101_000000_b", kernel="moe_stage1", lang="ck", kclass="ck") + assert run("languages", "--root", root, "--gfx", "gfx950", + "--kernel-name", "wvSplitK")["languages"] == ["hip"] + assert run("languages", "--root", root, "--gfx", "gfx950", + "--kernel-name", "moe_stage1")["languages"] == ["ck"] + # the name a lane passes is not always the name on disk + assert run("languages", "--root", root, "--gfx", "gfx950", + "--kernel-name", "wvsplitk")["languages"] == ["hip"] + out = run("languages", "--root", root, "--gfx", "gfx950", "--kernel-name", "nothing_here") + assert out["languages"] == [] and out["reason"] == "no_page" + + +# --------------------------------------------------------------------------- remote export +# The mapping onto KernelForge's KB Store. The failure mode being guarded is silent: a wrong +# dimension is not an error, it is a cold start — the write lands at an address nobody reads. +def export(root, *extra): + return run_lines("export-remote", "--root", root, *extra) + + +def stacked(root, exp_id, *, rocm="7.2", **kw): + """An entry with a verified_stack, which is where the remote framework_version comes from.""" + d = write_entry(root, exp_id, **kw) + meta_path = os.path.join(d, "meta.yaml") + meta = yaml.safe_load(open(meta_path)) + meta["verified_stack"] = {"rocm": rocm} if rocm else {} + yaml.safe_dump(meta, open(meta_path, "w")) + return d + + +def test_canonical_id_is_seven_ordered_segments(tmp_path): + """Order and content of the address. Upstream splits on ':' positionally, so a swapped pair is + a different identity that still parses.""" + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_a", kernel="moe_stage1", lang="ck", kclass="ck") + recs, summary = export(root) + assert summary["emitted"] == 1 + assert recs[0]["canonical_id"] == "geak:kernel:geak:moe_stage1:rocm:7.2:ck:gfx950" + # and the echoed identity must reconstruct it, or a reader validating the envelope rejects it + ident = recs[0]["knowledge"]["identity"] + assert ":".join(["kernel", ident["producer"], ident["kernel_name"], ident["framework"], + ident["framework_version"], ident["backend"], ident["gpu"]]) == \ + recs[0]["canonical_id"] + + +@pytest.mark.parametrize("lang,kclass,backend", [ + ("triton", "triton", "triton"), ("hip", "hip", "hip"), ("ck", "ck", "ck")]) +def test_language_is_the_backend_dimension(tmp_path, lang, kclass, backend): + """`backend` is the final implementation type upstream; `framework` stays rocm for all three + because one container image supplies all of them.""" + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_a", kernel="k1", lang=lang, kclass=kclass) + recs, _ = export(root) + ident = recs[0]["knowledge"]["identity"] + assert (ident["backend"], ident["framework"]) == (backend, "rocm") + + +@pytest.mark.parametrize("raw,want", [ + ("7.2", "7.2"), # what the recovered backlog carries + ("7.2.0", "7.2"), # a full release string + ("7.2.0-98765", "7.2"), # what /opt/rocm/.info/version actually holds on some images + ("", "unspecified"), # never guessed +]) +def test_framework_version_is_cut_to_major_minor(tmp_path, raw, want): + """A patch release must not split one kernel's history across two addresses, but the exact + string still has to survive in the payload.""" + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_a", kernel="k1", rocm=raw) + recs, _ = export(root) + assert recs[0]["knowledge"]["identity"]["framework_version"] == want + if raw: + assert recs[0]["knowledge"]["value"]["verified_stack"]["rocm"] == raw + + +def test_session_id_is_content_addressed_and_identity_scoped(tmp_path): + """Re-exporting one entry must update one candidate, not add another; and two identities must + never produce the same id, since artifacts are partitioned by session id alone.""" + root = str(tmp_path / "kb") + same = patch_text(new="BLOCK = 256") + stacked(root, "20260101_000000_a", kernel="k1", lang="triton", kclass="triton", patch=same) + stacked(root, "20260101_000000_b", kernel="k1", lang="hip", kclass="hip", patch=same) + recs, _ = export(root) + assert len(recs) == 2 + a, b = sorted(recs, key=lambda r: r["canonical_id"]) + assert a["session_id"] != b["session_id"], "same patch, different identity -> different id" + assert export(root)[0][0]["session_id"] == recs[0]["session_id"], "not stable across runs" + + +def test_recorded_signature_is_the_one_the_session_id_was_built_from(tmp_path): + """The address and the record must name the same patch. A reader dedups against its own store + by content_signature and fetches by session id; if those two disagree it silently re-downloads + a port it already has, or worse, treats two different patches as one.""" + root = str(tmp_path / "kb") + d = stacked(root, "20260101_000000_a", kernel="k1") + open(os.path.join(d, "report.md"), "w").write(REPORT) + rec = export(root)[0][0] + sig = rec["knowledge"]["value"]["content_signature"] + assert sig.startswith("csha:") + assert sig[len("csha:"):][:12] in rec["session_id"] + # and specifically not an artifact digest, which is what a shadowed local reads as + assert sig[len("csha:"):] not in {f["sha256"] for f in rec["files"]} + + +def test_identical_patches_collapse_to_the_best_measurement(tmp_path): + """Byte-identical patches are one candidate upstream. Which measurement we publish still + matters: the record is written with mode=replace, so sending the lower one publishes a speedup + we have already beaten.""" + root = str(tmp_path / "kb") + same = patch_text(new="BLOCK = 256") + stacked(root, "20260101_000000_a", kernel="k1", speedup=1.20, patch=same) + stacked(root, "20260101_000000_b", kernel="k1", speedup=1.90, patch=same) + recs, summary = export(root) + assert len(recs) == 1 and summary["deduped"] == 1 + assert recs[0]["knowledge"]["speedup"] == 1.9 + assert summary["deduped_dirs"] and "20260101_000000_a" in summary["deduped_dirs"][0] + + +def test_retired_entries_are_not_offered_remotely(tmp_path): + """The service ranks purely on the speedup we declare, so a retired win exported as a live + candidate would surface in someone's top-N with no way to tell.""" + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_a", kernel="k1", speedup=3.0, retired=True, + direction="dead-idea") + stacked(root, "20260101_000000_b", kernel="k1", speedup=1.5, direction="live-idea") + recs, summary = export(root) + assert [r["knowledge"]["value"]["direction"] for r in recs] == ["live-idea"] + assert summary["skipped"]["retired"] == 1 + recs, _ = export(root, "--include-retired") + assert len(recs) == 2 + + +def test_champion_is_one_per_identity_and_must_beat_baseline(tmp_path): + """Upstream's own gate. A candidate is always recorded; only the pointer is earned.""" + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_a", kernel="k1", speedup=1.5, direction="d1") + stacked(root, "20260101_000000_b", kernel="k1", speedup=2.5, direction="d2") + stacked(root, "20260101_000000_c", kernel="k2", speedup=0.9, direction="d3") + recs, summary = export(root) + champs = [r for r in recs if r["champion"]] + assert len(champs) == 1 and champs[0]["knowledge"]["speedup"] == 2.5 + assert summary["identities"] == 2 and summary["champions"] == 1 + losing = [r for r in recs if r["knowledge"]["value"]["direction"] == "d3"][0] + assert losing["champion"] is False and losing["champion_eligible"] is False + + +def test_comparability_fields_travel_with_the_speedup(tmp_path): + """get_top_sessions ranks on the bare `speedup` number and knows nothing about bench keys, so + a reader can only filter incomparable entries out if we sent what to filter on.""" + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_a", kernel="k1", bench="b2:zzz") + value = export(root)[0][0]["knowledge"]["value"] + assert value["metric"]["bench_key"] == "b2:zzz" + assert value["metric"]["metric_kind"] == "geomean" + assert value["direction"] and value["content_signature"] + + +def test_artifacts_are_referenced_not_inlined(tmp_path): + """Patches run to 240KB here and reach an agent through a tool result. The record carries a + path plus a digest; the bytes go up the artifact channel.""" + root = str(tmp_path / "kb") + d = stacked(root, "20260101_000000_a", kernel="k1") + open(os.path.join(d, "report.md"), "w").write(REPORT) + rec = export(root)[0][0] + assert rec["knowledge"]["value"]["artifacts"] == {"patch": "patch.diff", "report": "report.md"} + assert {f["path"] for f in rec["files"]} == {"patch.diff", "report.md"} + for f in rec["files"]: + assert len(f["sha256"]) == 64 and f["size"] > 0 and f["kind"] == "rewrite" + blob = json.dumps(rec) + assert "BLOCK = " not in blob and "Round-by-round" not in blob + + +def test_verbatim_dead_ends_stay_out_of_the_record(tmp_path): + """dead_ends_md runs to tens of KB and report.md already carries it; the structured form is + small enough to ride along and is the half an agent can act on.""" + root = str(tmp_path / "kb") + d = stacked(root, "20260101_000000_a", kernel="k1") + meta_path = os.path.join(d, "meta.yaml") + meta = yaml.safe_load(open(meta_path)) + meta["dead_ends_md"] = "- a very long verbatim section " * 200 + meta["dead_ends"] = [{"idea": "buffer_ops off", "measured": 0.883, "mechanism": "load-bearing"}] + yaml.safe_dump(meta, open(meta_path, "w")) + value = export(root)[0][0]["knowledge"]["value"] + assert "dead_ends_md" not in value + assert value["dead_ends"][0]["measured"] == 0.883 + + +def test_export_filters_and_overrides(tmp_path): + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_a", kernel="k1") + stacked(root, "20260101_000000_b", kernel="k2") + assert len(export(root, "--kernel-name", "k1")[0]) == 1 + assert export(root, "--gfx", "gfx942")[1]["emitted"] == 0 + rec = export(root, "--kernel-name", "k1", "--producer", "forge-loop", "--gpu", "MI300X")[0][0] + assert rec["canonical_id"].startswith("kernel:forge-loop:k1:") + assert rec["canonical_id"].endswith(":gfx942") + assert rec["session_id"].startswith("geak-") # the id prefix is ours, not the producer arg + + +# --------------------------------------------------------------------------- the store plane +# `resolve-remote` / `write-remote` read and write the same experience through a KB Store held on +# disk in the shape the service uses. They exist to make the plane swappable, so what is pinned here +# is that the lane cannot tell the difference: same JSON shape, same curation, same gates — plus the +# two write outcomes the store adds, append vs update, which is what a key-value plane makes visible. + +UPLOADER = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(STORE))), "kb", "remote_upload.py") + + +def seed_store(tmp_path, root, *extra): + """Build a store the way the real one is built: export the directory plane, then load it.""" + jsonl = str(tmp_path / "records.jsonl") + run("export-remote", "--root", root, "--out", jsonl, *extra) + store = str(tmp_path / "store") + p = subprocess.run([sys.executable, UPLOADER, "--records", jsonl, "--local", store, + "--apply", "--quiet"], capture_output=True, text=True) + assert p.returncode == 0, p.stderr + return store + + +def resolve_remote(store, refs, kernel="fused_moe_kernel", lang="triton", gfx="gfx950", *extra): + return run("resolve-remote", "--store", store, "--kernel-name", kernel, "--language", lang, + "--gfx", gfx, "--refs-dir", refs, *extra) + + +def test_both_planes_offer_the_same_candidates(tmp_path): + """The whole premise of the store plane: one curated backlog, two ways to read it.""" + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_a", speedup=4.0, direction="tile-retune") + stacked(root, "20260101_000000_b", speedup=2.0, direction="vectorize") + stacked(root, "20260101_000000_c", speedup=1.01, direction="unroll") # below the floor + store = seed_store(tmp_path, root) + + local = resolve(root, str(tmp_path / "r1"), "fused_moe_kernel", "triton", "gfx950", + "--min-speedup", "1.05") + remote = resolve_remote(store, str(tmp_path / "r2"), "fused_moe_kernel", "triton", "gfx950", + "--min-speedup", "1.05") + assert remote["read_reason"] == local["read_reason"] == "read" + keys = ("rank", "speedup", "direction", "comparable", "kernel_name", "language", "gfx") + assert [{k: c.get(k) for k in keys} for c in remote["candidates"]] == \ + [{k: c.get(k) for k in keys} for c in local["candidates"]] + assert remote["filtered"]["below_min_speedup"] == local["filtered"]["below_min_speedup"] == 1 + assert remote["canonical_id"] == "geak:kernel:geak:fused_moe_kernel:rocm:7.2:triton:gfx950" + + +def test_the_store_plane_curates_what_the_store_itself_cannot(tmp_path): + """The store ranks on speedup alone. Direction collapse and bench comparability are ours.""" + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_a", speedup=40.0, direction="tile-retune", bench="b:imported") + stacked(root, "20260101_000000_b", speedup=6.0, direction="tile-retune", bench="b:imported") + stacked(root, "20260101_000000_c", speedup=3.0, direction="vectorize", bench="b2:onbox") + store = seed_store(tmp_path, root) + out = resolve_remote(store, str(tmp_path / "refs")) + assert [c["direction"] for c in out["candidates"]] == ["tile-retune", "vectorize"] + assert out["filtered"]["same_direction_collapsed"] == 1 + assert out["candidates"][0]["comparable"] is True + assert out["candidates"][1]["comparable"] is False, "a b2: measurement is not a b: one" + assert out["candidates"][0]["alternates"], "the runner-up rides along, it is not discarded" + + +def test_every_offered_patch_resolves_to_real_bytes(tmp_path): + """A candidate listed with a path that resolves to nothing is worse than not listing it.""" + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_a", speedup=4.0, direction="tile-retune") + stacked(root, "20260101_000000_b", speedup=3.0, direction="tile-retune") # an alternate + store = seed_store(tmp_path, root) + out = resolve_remote(store, str(tmp_path / "refs")) + paths = [c["patch_path"] for c in out["candidates"]] + paths += [a["patch_path"] for c in out["candidates"] for a in (c.get("alternates") or [])] + assert paths and all(os.path.getsize(p) > 0 for p in paths) + + +def test_the_store_plane_mirrors_prose_for_the_same_audit(tmp_path): + """The report travels as an artifact, so a store-sourced reference reads like a local one.""" + root = str(tmp_path / "kb") + d = stacked(root, "20260101_000000_a", speedup=4.0) + with open(os.path.join(d, "report.md"), "w") as f: + f.write("# why this worked\nfolded the scale\n") + store = seed_store(tmp_path, root) + refs = str(tmp_path / "refs") + out = resolve_remote(store, refs) + assert out["candidates"], out + prose = open(out["candidates"][0]["prose_path"]).read() + assert "folded the scale" in prose and "tile-retune" in prose + assert "direction" in open(os.path.join(refs, "index.md")).read() + + +def test_a_cold_key_is_a_cold_start_not_a_failure(tmp_path): + """Same vocabulary as the directory plane, so the lane's logging and gates need no branch.""" + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_a", speedup=4.0) + store = seed_store(tmp_path, root) + out = resolve_remote(store, str(tmp_path / "refs"), "some_other_kernel") + local = resolve(root, str(tmp_path / "r2"), "some_other_kernel") + assert out["candidates"] == [] and out["read_reason"] == local["read_reason"] + + +def test_a_store_root_that_is_not_there_is_a_miss_not_an_empty_store(tmp_path): + """A typo'd path must not read as 'no experience' and silently cold-start a warm run.""" + out = resolve_remote(str(tmp_path / "nope"), str(tmp_path / "refs")) + assert out["candidates"] == [] and "no_such_store" in out["reason"] + + +def test_a_key_that_differs_only_in_stack_version_is_reported_not_silently_used(tmp_path): + """`unspecified` vs `7.2` splits one kernel's history in two; the read must say so.""" + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_a", speedup=4.0, rocm="7.2") + store = seed_store(tmp_path, root) + out = resolve_remote(store, str(tmp_path / "refs"), "fused_moe_kernel", "triton", "gfx950", + "--framework-version", "6.4") + assert out["match_tier"] == "other_version" + assert any("rocm:7.2" in p for p in out["other_language_pages"]) + + +def test_a_write_records_both_planes(tmp_path): + root = str(tmp_path / "kb") + store = str(tmp_path / "store") + p = tmp_path / "p.diff" + p.write_text(patch_text()) + w = run("write-remote", "--root", root, "--store", store, "--kernel-name", "fused_moe_kernel", + "--language", "triton", "--gfx", "gfx950", "--kernel-class", "triton", + "--speedup", "2.0", "--patch", str(p), "--direction", "tile-retune", + "--framework-version", "7.2") + assert w["written"] is True and os.path.isfile(os.path.join(w["dir"], "meta.yaml")) + assert w["remote"]["written"] is True + assert w["remote"]["canonical_id"] == "geak:kernel:geak:fused_moe_kernel:rocm:7.2:triton:gfx950" + assert w["remote"]["champion"] is True and w["remote"]["replaced"] is False + out = resolve_remote(store, str(tmp_path / "refs")) + assert [c["speedup"] for c in out["candidates"]] == [2.0] + assert out["candidates"][0]["is_champion"] is True + + +def test_a_new_patch_appends_a_candidate_under_the_same_key(tmp_path): + root, store = str(tmp_path / "kb"), str(tmp_path / "store") + first, second = tmp_path / "a.diff", tmp_path / "b.diff" + first.write_text(patch_text(new="BLOCK = 128")) + second.write_text(patch_text(new="BLOCK = 256")) + a = run("write-remote", "--root", root, "--store", store, "--kernel-name", "k", "--language", + "triton", "--gfx", "gfx950", "--kernel-class", "triton", "--speedup", "2.0", + "--patch", str(first), "--direction", "tile-retune", "--framework-version", "7.2") + b = run("write-remote", "--root", root, "--store", store, "--kernel-name", "k", "--language", + "triton", "--gfx", "gfx950", "--kernel-class", "triton", "--speedup", "3.0", + "--patch", str(second), "--direction", "vectorize", "--framework-version", "7.2") + assert a["remote"]["canonical_id"] == b["remote"]["canonical_id"] + assert b["remote"]["session_id"] != a["remote"]["session_id"] + assert b["remote"]["replaced"] is False and b["remote"]["champion"] is True + out = resolve_remote(store, str(tmp_path / "refs"), "k") + assert [c["speedup"] for c in out["candidates"]] == [3.0, 2.0] + + +def test_remeasuring_one_patch_updates_its_candidate_instead_of_adding_one(tmp_path): + """Session ids are content-addressed, so the store cannot inflate from repeated runs.""" + root, store = str(tmp_path / "kb"), str(tmp_path / "store") + p = tmp_path / "p.diff" + p.write_text(patch_text()) + a = run("write-remote", "--root", root, "--store", store, "--kernel-name", "k", "--language", + "triton", "--gfx", "gfx950", "--kernel-class", "triton", "--speedup", "2.0", + "--patch", str(p), "--direction", "tile-retune", "--framework-version", "7.2") + # the same code re-emitted from a different workspace: shifted hunk, reflowed whitespace + again = tmp_path / "again.diff" + again.write_text(patch_text(path="kernel_src/sglang/k.py", old="BLOCK = 64", + new="BLOCK = 128", at=87)) + b = run("write-remote", "--root", root, "--store", store, "--kernel-name", "k", "--language", + "triton", "--gfx", "gfx950", "--kernel-class", "triton", "--speedup", "2.1", + "--patch", str(again), "--direction", "tile-retune", "--framework-version", "7.2") + assert b["written"] is False and b["reason"] == "duplicate_impl" # the directory plane's view + assert b["remote"]["session_id"] == a["remote"]["session_id"] + assert b["remote"]["replaced"] is True + out = resolve_remote(store, str(tmp_path / "refs"), "k") + assert len(out["candidates"]) == 1, "a reproduction is not a second candidate" + # The confirmation is not lost, it is recorded on the one candidate: same code, measured twice. + document = json.load(open(os.path.join( + store, *b["remote"]["canonical_id"].split(":"), "sessions", + b["remote"]["session_id"], "knowledge.json"))) + assert document["value"]["reproductions"] == 2 + assert document["speedup"] == 2.0, "the original's own measurement is what was recorded" + + +def test_a_store_failure_never_costs_the_measured_result(tmp_path): + """The directory plane is the source of truth; a KB write is bookkeeping on top of it.""" + root = str(tmp_path / "kb") + blocked = tmp_path / "not-a-dir" + blocked.write_text("") + p = tmp_path / "p.diff" + p.write_text(patch_text()) + w = run("write-remote", "--root", root, "--store", str(blocked), "--kernel-name", "k", + "--language", "triton", "--gfx", "gfx950", "--kernel-class", "triton", + "--speedup", "2.0", "--patch", str(p), "--framework-version", "7.2") + assert w["written"] is True, "the local entry must survive a broken store" + assert w["remote"]["written"] is False and w["remote"]["reason"] + + +@pytest.mark.parametrize("args,reason", [ + (("resolve-remote", "--store", "/nonexistent/store"), "a store that is not there"), + (("resolve-remote", "--store", "/tmp", "--canonical-id", "not a valid id"), "a malformed key"), +]) +def test_the_store_read_never_raises(tmp_path, args, reason): + out = run(*args, "--kernel-name", "k", "--language", "triton", "--gfx", "gfx950", + "--refs-dir", str(tmp_path / "refs")) + assert out["candidates"] == [], reason + assert out["read_reason"], "a cold start still has to say why" + + +# --------------------------------------------------------------------------- attestation +# `reproductions` counts the same code being WRITTEN twice. This counts the stored entry being +# READ back out and put on a box — the only signal a later retire pass can act on, and the one +# thing the schema never recorded. +def attest(root, exp_dir, outcome="validated", *extra): + return run("attest", "--exp-dir", exp_dir, "--outcome", outcome, *extra) + + +def test_attesting_counts_the_attempt_and_moves_no_speedup(tmp_path): + root = str(tmp_path / "kb") + d = write_entry(root, "20260101_000000_aaaaaa", speedup=2.0) + out = attest(root, d, "validated", "--measured-speedup", "1.9", "--measured-by", "boxA", + "--note", "held on 7.2", "--apply") + assert out["attested"] is True + assert out["attestations"]["recalls"] == 1 and out["attestations"]["validations"] == 1 + meta = yaml.safe_load(open(os.path.join(d, "meta.yaml"))) + assert meta["metric"]["speedup"] == 2.0, "an attestation is evidence, not a re-measurement" + assert meta["reproductions"] == 1, "not the same counter as a duplicate write" + entry = meta["attestations"]["history"][-1] + assert entry["measured_speedup"] == 1.9 and entry["by"] == "boxA" + + +def test_a_dry_attestation_writes_nothing(tmp_path): + root = str(tmp_path / "kb") + d = write_entry(root, "20260101_000000_aaaaaa") + out = attest(root, d, "failed") + assert out["attested"] is False and out["attestations"]["failures"] == 1 + assert "attestations" not in yaml.safe_load(open(os.path.join(d, "meta.yaml"))) + + +def test_attestations_accumulate_and_raise_a_hint_the_read_surfaces(tmp_path): + """The hint is advisory: the entry is still offered and still ranked. Retiring it is a separate + act by a separate caller, which is what makes the counters safe to write automatically.""" + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + d = write_entry(root, "20260101_000000_aaaaaa") + for _ in range(2): + attest(root, d, "not_reproduced", "--apply") + candidate = resolve(root, refs)["candidates"][0] + assert candidate["recalls"] == 2 and candidate["validations"] == 0 + assert "could not reproduce" in candidate["retire_hint"] + assert "track record" in open(candidate["prose_path"]).read() + + +def test_a_never_tried_entry_reads_as_untried_not_as_failing(tmp_path): + root, refs = str(tmp_path / "kb"), str(tmp_path / "refs") + write_entry(root, "20260101_000000_aaaaaa") + candidate = resolve(root, refs)["candidates"][0] + assert candidate["recalls"] == 0 and candidate["retire_hint"] == "" + + +def test_attest_never_fails_the_caller(tmp_path): + out = run("attest", "--exp-dir", str(tmp_path / "nope"), "--outcome", "validated", "--apply") + assert out["attested"] is False and out["reason"] + out = run("attest", "--exp-dir", str(write_entry(str(tmp_path / "kb"), "20260101_000000_a")), + "--outcome", "validated", "--measured-speedup", "not a number", "--apply") + assert out["attested"] is True, "unusable evidence drops; the attempt still counted" + + +def _seeded(tmp_path, root): + """A store built the way seed_store builds one, plus the session id the export minted.""" + jsonl = str(tmp_path / "records.jsonl") + run("export-remote", "--root", root, "--out", jsonl) + store = str(tmp_path / "store") + p = subprocess.run([sys.executable, UPLOADER, "--records", jsonl, "--local", store, + "--apply", "--quiet"], capture_output=True, text=True) + assert p.returncode == 0, p.stderr + records = [json.loads(line) for line in open(jsonl) if line.strip()] + return store, records[0]["session_id"] + + +def test_the_remote_record_carries_the_local_ledger(tmp_path): + """Both planes read the same page, so a count that only lands on one of them makes the two + disagree about how much anyone has actually tried this.""" + root = str(tmp_path / "kb") + d = stacked(root, "20260101_000000_aaaaaa") + attest(root, d, "validated", "--apply") + store, _ = _seeded(tmp_path, root) + out = resolve_remote(store, str(tmp_path / "refs"), "fused_moe_kernel") + assert out["candidates"][0]["validations"] == 1 + + +def test_attest_remote_counts_against_the_key_addressed_record(tmp_path): + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_aaaaaa") + store, session = _seeded(tmp_path, root) + out = run("attest-remote", "--store", store, "--session-id", session, "--outcome", "validated", + "--kernel-name", "fused_moe_kernel", "--language", "triton", "--gfx", "gfx950", + "--framework-version", "7.2", "--measured-speedup", "2.2", "--apply") + assert out["attested"] is True + found = [r for r in out["pages"] if r["found"]] + assert found and all(r["attestations"]["validations"] == 1 for r in found) + assert resolve_remote(store, str(tmp_path / "refs"), "fused_moe_kernel" + )["candidates"][0]["validations"] == 1 + + +def test_attest_remote_on_an_unknown_session_reports_rather_than_raises(tmp_path): + root = str(tmp_path / "kb") + stacked(root, "20260101_000000_aaaaaa") + store, _ = _seeded(tmp_path, root) + out = run("attest-remote", "--store", store, "--session-id", "nope", "--outcome", "failed", + "--kernel-name", "fused_moe_kernel", "--language", "triton", "--gfx", "gfx950", + "--apply") + assert out["attested"] is False and all(p["found"] is False for p in out["pages"]) diff --git a/kernel_workflow/scripts/tests/test_kb_loop_offline.py b/kernel_workflow/scripts/tests/test_kb_loop_offline.py new file mode 100644 index 000000000..1ffb587a7 --- /dev/null +++ b/kernel_workflow/scripts/tests/test_kb_loop_offline.py @@ -0,0 +1,216 @@ +"""The warm-start loop, end to end, with no GPU and no network. + +A lane's use of the experience KB is a cycle, and every piece of it is tested elsewhere in +isolation. What is only visible when the pieces run in order is whether the cycle CLOSES: + + read a key -> pick the top candidates -> land one on a workspace that does not match the + layout it was recorded from -> optimize further on top of it -> write the result back under + the same key -> read it again and get the improvement, not the starting point. + +The one link this cannot cover is the measurement, which needs a card. Everything on either side +of it is here, so a regression in the plumbing fails in CI rather than half an hour into a run. +""" + +import json +import os +import shutil +import subprocess +import sys + +import pytest + +SCRIPTS = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +STORE = os.path.join(SCRIPTS, "experience_store.py") +UPLOADER = os.path.join(os.path.dirname(os.path.dirname(SCRIPTS)), "kb", "remote_upload.py") +yaml = pytest.importorskip("yaml") +pytestmark = pytest.mark.skipif(shutil.which("git") is None, reason="the loop lands a patch with git") + +KERNEL = "fused_moe_kernel" +CID = "geak:kernel:geak:fused_moe_kernel:rocm:7.2:triton:gfx950" +BASELINE = "import triton\n\nBLOCK = 64\nNUM_WARPS = 4\n" + + +def run(*args): + p = subprocess.run([sys.executable, STORE] + [str(a) for a in args], + capture_output=True, text=True) + assert p.returncode == 0, f"the store must never fail the lane: {p.stderr}" + return json.loads(p.stdout) + + +def git(cwd, *args): + p = subprocess.run(["git"] + list(args), cwd=cwd, capture_output=True, text=True) + assert p.returncode == 0, f"git {' '.join(args)}: {p.stderr}" + return p.stdout + + +def recorded_entry(root, exp_id, *, speedup, direction, patch, report): + """One curated entry, as the imported backlog looks on disk.""" + d = os.path.join(root, "gfx950", "triton", f"{KERNEL}__triton__gfx950", exp_id) + os.makedirs(d, exist_ok=True) + with open(os.path.join(d, "meta.yaml"), "w") as f: + yaml.safe_dump({ + "layer": "artifact", "lifecycle": "candidate", "gfx": "gfx950", + "kernel_class": "triton", "kernel_name": KERNEL, "language": "triton", + "direction": direction, "reproductions": 1, "verified_stack": {"rocm": "7.2"}, + "strategy": f"{direction}: what was tried and why", + "metric": {"speedup": speedup, "gpu_arch": "gfx950", "bench_key": "b:imported", + "metric_kind": "geomean"}, + }, f) + with open(os.path.join(d, "patch.diff"), "w") as f: + f.write(patch) + with open(os.path.join(d, "report.md"), "w") as f: + f.write(report) + return d + + +def recorded_patch(new_value, path="source/triton_fused_moe_kernel.py"): + """A patch as RECORDED: against the layout of the run that produced it, not this workspace.""" + return (f"diff --git a/{path} b/{path}\nindex 1111111..2222222 100644\n" + f"--- a/{path}\n+++ b/{path}\n@@ -1,4 +1,4 @@\n" + f" import triton\n \n-BLOCK = 64\n+BLOCK = {new_value}\n NUM_WARPS = 4\n") + + +def workspace(tmp_path, name="ws"): + """A head-extracted layout: same file, a different place than the patch was recorded against.""" + ws = tmp_path / name + src = ws / "kernel_src" / "sglang" + src.mkdir(parents=True) + (src / "triton_fused_moe_kernel.py").write_text(BASELINE) + git(str(ws), "init", "-q") + git(str(ws), "config", "user.email", "t@t") + git(str(ws), "config", "user.name", "t") + git(str(ws), "add", "-A") + git(str(ws), "commit", "-qm", "baseline") + return ws + + +def build_store(tmp_path, kb_root): + jsonl = str(tmp_path / "records.jsonl") + run("export-remote", "--root", kb_root, "--out", jsonl) + store = str(tmp_path / "store") + p = subprocess.run([sys.executable, UPLOADER, "--records", jsonl, "--local", store, + "--apply", "--quiet"], capture_output=True, text=True) + assert p.returncode == 0, p.stderr + return store + + +def sessions_under(store, cid=CID): + d = os.path.join(store, *cid.split(":"), "sessions") + return sorted(n for n in os.listdir(d) if not n.startswith(".")) if os.path.isdir(d) else [] + + +def champion_of(store, cid=CID): + with open(os.path.join(store, *cid.split(":"), "champion.json")) as f: + return json.load(f) + + +def test_the_warm_start_loop_closes(tmp_path): + kb_root = str(tmp_path / "kb") + recorded_entry(kb_root, "20260101_000000_a", speedup=3.0, direction="tile-retune", + patch=recorded_patch(128), + report="# tile retune\n## Key optimizations\nwidened the tile\n") + recorded_entry(kb_root, "20260101_000000_b", speedup=1.8, direction="vectorize", + patch=recorded_patch(96, path="src/k.py"), + report="# vectorize\n## Key optimizations\nwider loads\n") + recorded_entry(kb_root, "20260101_000000_c", speedup=1.01, direction="unroll", + patch=recorded_patch(72, path="src/k.py"), report="# unroll\n") + store = build_store(tmp_path, kb_root) + + # 1. read the key. Ranked best first, the near-tie left out of the verify budget. + read = run("resolve-remote", "--store", store, "--kernel-name", KERNEL, "--language", "triton", + "--gfx", "gfx950", "--top-n", "3", "--min-speedup", "1.05", + "--framework-version", "7.2", "--refs-dir", str(tmp_path / "refs")) + assert read["read_reason"] == "read" and read["canonical_id"] == CID + assert [c["speedup"] for c in read["candidates"]] == [3.0, 1.8] + assert read["filtered"]["below_min_speedup"] == 1 + top = read["candidates"][0] + assert os.path.getsize(top["patch_path"]) > 0 + assert "widened the tile" in open(top["prose_path"]).read() + + # 2. land it on a workspace whose layout the recorded patch never saw. + ws = workspace(tmp_path) + landed = str(tmp_path / "landed.diff") + remapped = run("remap", "--patch", top["patch_path"], "--out", landed, "--workspace", str(ws), + "--editable", "kernel_src/sglang/triton_fused_moe_kernel.py") + assert remapped["remapped"] is True + git(str(ws), "apply", "--check", landed) + git(str(ws), "apply", landed) + assert "BLOCK = 128" in (ws / "kernel_src" / "sglang" / "triton_fused_moe_kernel.py").read_text() + + # 3. keep optimizing on top of what was adopted. (The measurement is the GPU's job; the number + # here stands in for it.) The patch written back is the FULL diff from the baseline, which is + # what the recorded speedup is measured against. + path = ws / "kernel_src" / "sglang" / "triton_fused_moe_kernel.py" + path.write_text(path.read_text().replace("BLOCK = 128", "BLOCK = 256") + .replace("NUM_WARPS = 4", "NUM_WARPS = 8")) + improved = str(tmp_path / "improved.diff") + with open(improved, "w") as f: + f.write(git(str(ws), "diff")) + + # 4. write it back under the same key, deriving from the entry it was built on. + wrote = run("write-remote", "--root", kb_root, "--store", store, "--kernel-name", KERNEL, + "--language", "triton", "--gfx", "gfx950", "--kernel-class", "triton", + "--speedup", "4.5", "--patch", improved, "--direction", "tile-retune", + "--metric-kind", "geomean", "--framework-version", "7.2", + "--parent", top["exp_dir"]) + assert wrote["written"] is True + assert wrote["remote"]["written"] is True and wrote["remote"]["canonical_id"] == CID + assert wrote["remote"]["replaced"] is False, "new code appends a candidate" + assert wrote["remote"]["champion"] is True + # All three recorded entries are still there, including the near-tie that was filtered out of + # the read: the floor decides what is worth a verify slot, not what the store keeps. + assert len(sessions_under(store)) == 4 + + # 5. read the key again: the loop hands back the improvement, not the starting point. + again = run("resolve-remote", "--store", store, "--kernel-name", KERNEL, "--language", "triton", + "--gfx", "gfx950", "--top-n", "3", "--min-speedup", "1.05", + "--framework-version", "7.2", "--refs-dir", str(tmp_path / "refs2")) + assert [c["speedup"] for c in again["candidates"]] == [4.5, 1.8] + assert again["candidates"][0]["session_id"] == wrote["remote"]["session_id"] + assert again["candidates"][0]["is_champion"] is True + assert champion_of(store)["session_id"] == wrote["remote"]["session_id"] + # 3.0x and 4.5x share a direction, so the one that was superseded rides along as an alternate + # rather than costing a second verify slot. + assert again["filtered"]["same_direction_collapsed"] == 1 + assert [alt["speedup"] for alt in again["candidates"][0]["alternates"]] == [3.0] + + # and the improvement applies to the same baseline the run started from. It was recorded from + # this very layout, so remap has nothing to rewrite and says so rather than emitting a copy. + fresh = workspace(tmp_path, "ws2") + relanded = str(tmp_path / "relanded.diff") + fit = run("remap", "--patch", again["candidates"][0]["patch_path"], "--out", relanded, + "--workspace", str(fresh), "--editable", "kernel_src/sglang/triton_fused_moe_kernel.py") + assert fit["remapped"] is False and fit["reason"] == "no_change_needed" + git(str(fresh), "apply", "--check", again["candidates"][0]["patch_path"]) + git(str(fresh), "apply", again["candidates"][0]["patch_path"]) + assert "BLOCK = 256" in (fresh / "kernel_src" / "sglang" / "triton_fused_moe_kernel.py").read_text() + + +def test_a_second_lap_over_the_same_code_does_not_grow_the_store(tmp_path): + """A run that adopts a warm start and fails to beat it re-emits the patch it adopted. Recording + that as a new candidate every time is how a store fills up with copies of one idea.""" + kb_root = str(tmp_path / "kb") + recorded_entry(kb_root, "20260101_000000_a", speedup=3.0, direction="tile-retune", + patch=recorded_patch(128), report="# tile retune\n") + store = build_store(tmp_path, kb_root) + before = sessions_under(store) + + read = run("resolve-remote", "--store", store, "--kernel-name", KERNEL, "--language", "triton", + "--gfx", "gfx950", "--framework-version", "7.2", "--refs-dir", str(tmp_path / "refs")) + ws = workspace(tmp_path) + landed = str(tmp_path / "landed.diff") + run("remap", "--patch", read["candidates"][0]["patch_path"], "--out", landed, + "--workspace", str(ws), "--editable", "kernel_src/sglang/triton_fused_moe_kernel.py") + git(str(ws), "apply", landed) + reemitted = str(tmp_path / "reemitted.diff") + with open(reemitted, "w") as f: + f.write(git(str(ws), "diff")) + + wrote = run("write-remote", "--root", kb_root, "--store", store, "--kernel-name", KERNEL, + "--language", "triton", "--gfx", "gfx950", "--kernel-class", "triton", + "--speedup", "2.95", "--patch", reemitted, "--direction", "tile-retune", + "--framework-version", "7.2") + assert wrote["written"] is False and wrote["reason"] == "duplicate_impl" + assert wrote["remote"]["replaced"] is True, "the same code lands back on its own candidate" + assert sessions_under(store) == before + assert champion_of(store)["value"] == 3.0, "a slower remeasure does not take the pointer"