feat(aiur): Accurate whole-env sharding based on execution measurements - #537
Draft
samuelburnham wants to merge 7 commits into
Draft
feat(aiur): Accurate whole-env sharding based on execution measurements#537samuelburnham wants to merge 7 commits into
samuelburnham wants to merge 7 commits into
Conversation
…acking Port of sb/measured-ingress's planning layer onto the addr-first kernel (#529), reworked for its lazy fault-in semantics. Profiling (.ixprof v3): the recording kernel's try_get_const consults are drained per constant into a touch graph — the measured set a lazy checker faults in — persisted as a presence-flagged CSR section alongside a new block-level reference graph and a subst_unique counter (splitmix64-keyed dedup of substitution work items: the post-memoization volume a memoizing executor pays). Both graph sections are optional, so older readers' formats stay decodable by version gate alone. Planner: partition_for_aiur_ram bin-packs to a host-RAM cap with faulted-set byte accounting — under #529 the witness always ships the full closure, but only faulted bytes cost ingress trace rows, so a shard's byte feature is owned ∪ touched(owned) (full reference closure as the over-approximating fallback when no touch graph was recorded). The stub/promotion machinery of the original branch is dropped entirely: lazy fault-in computes at run time exactly the set the stubs predicted, with no replay divergence and no repair ladder, and the manifest format stays byte-identical with main (owned-only). Cost models: nlogn(x) = x·log2(x+2) feature form matching Aiur's width·height·log2(height) FFT structure; prove/execute time and RAM coefficients carried over from the pre-#529 aiur bench-suite fits. They are placeholders pending a refit against this kernel — the circuit shrank (deleted positional machinery), the byte feature narrowed to the faulted set, and the recording counters moved under #442 — with provenance noted at the definitions. The costs sidecar (<out>.costs.csv) persists per-shard features incl. subst_unique and nat_arith so the refit has its corpus. CLI: ix shard --backend aiur --max-ram G; ix profile --backend/--top (per-metric block leaderboards with display names) and ix profile sweep (env-wide closure cost sweep: per-root feature CSV + feasibility, min-root-per-hot-block, and diversity reports).
prove --ixe --ixes with no --shard now runs the whole partition through one EnvHandle + one AiurSystem instead of a fire-and-forget loop. Progress persists in the keyed cache ~/.ix/cache/shard-proofs/ (one entry per claim digest, holding the proof's store address): an entry is written only after the proof VERIFIES and binds to the shard's reconstructed CheckEnv claim digest, so a crash or OOM costs only the in-flight shard and re-running the same command resumes. Cached entries are re-verified on resume — the digest pins the claim, not the circuit version. The digest is manifest-independent, so manifests sharing a shard share its proof. Pending shards prove heaviest-predicted-RAM first (from the packer's costs sidecar): if the RAM model's blind spot is going to breach a box it is one of the heaviest shards, and proving those first surfaces the failure in the opening minutes instead of hours in. --jobs N proves shards concurrently (default 1 — each prove peaks at the shard's full predicted RAM). aiur: ExecError::InvalidIOKey now names the channel and key, so a witness that under-ships identifies the constant whose bytes are missing instead of failing generically.
aiur-shard bench backend: profile the env's touch graph, pack at a pinned machine-independent budget (aiurShardBudgetGb = 250, the deliverable prove box — a PR row and its baseline must describe the same partition problem), upload planner metrics (shards, pred-floor-ram, union-bytes) on tight 5% upper bands (profile worker interleaving wobbles touch sets ±2 shards, so no exact pins), then execute the three heaviest-predicted shards through the native kernel and aggregate count-neutral heavy-execute-time / heavy-peak-rss. ix check gains --json/--json-name so a shard check self-reports its bench row with the env parse excluded from the timed window. Init is registered as a benchmark env — it is the shard pipeline's deliverable partition, so its planner trend line matters on its own. Tests/Ix/Kernel/ShardPipeline.lean: end-to-end regression over a small multi-shard env — profile → pack (disjoint cover) → every shard checks → batched prove reaches the composed verdict → second run resumes from the shard-proofs cache with zero pending. Runs under a scratch cache root, never the global ~/.ix/cache; wired into the ixvm suite.
…s preserved ix shard --rebudget K --manifest M.ixes --max-ram G --backend aiur packs only shard K's blocks at the child budget (in the original cut-coherent order); every other shard's block list is reproduced verbatim — asserted, because ShardManifest::build is deterministic in (profile, assignment) and that identity is exactly what keeps the untouched shards' claim digests, and therefore their ~/.ix/cache/shard-proofs entries, valid. The first child keeps id K, the rest append; the aggregation tree's leaf K becomes a left-leaning chain over the children. The costs sidecar is recomputed for every row (identical for untouched shards — the faulted-set union is a function of the owned set alone), so heaviest- first prove ordering picks up the children's predictions. This is the recourse when the RAM model's blind spot packs an outlier shard past what the box can prove: a global repack at a lower budget shifts boundaries from early in the pack order and forfeits most of the proof cache, while a rebudget costs exactly the split shard. Ported from the sb/sp1-shard-agg experiment (where a 148-predicted shard measured ~248 GiB); simplified for the owned-only manifest — the promotion-inheritance step is gone because the witness ships the full closure and a child can never miss an IO key.
…resources fit
Replaces the placeholder coefficients (fit on the pre-addr-first eager
kernel) with a model fit entirely on this kernel, restructured around the
prover's actual work unit.
Stage 1 predicts a run's total FFT cost from profile features:
fft = 2.599e9 + 188.6*nlogn(faulted_bytes) + 71.27*nlogn(subst)
+ 6548*nlogn(def_eq)
def_eq is load-bearing: it appears in every acceptable-tail fit, pricing
the definitional-equality-dense shards that under the old model measured
up to 25% over predicted RAM (reproduced on this kernel before the
refit: the nat-heaviest Init shard proved at 219 GiB against a 175
prediction). nat_arith adds no explanatory power at shard scale on this
corpus and stays a persisted-but-unmodeled feature; the exact path below
covers its tail.
Stage 2 maps FFT cost linearly to resources — physically grounded
(committed LDE volume is proportional to FFT work):
prove_secs = 6.18 + 2.0425 * BFFT RSS_gib = 13.59 + 2.3507 * BFFT
exposed both composed (plan-time, from features) and as *_for_fft
accessors, so a measured FFT cost from an execute-mode stats dump gives
an exact-height prediction — the cheap pre-prove gate for any residual
outlier, which can then be `--rebudget` before burning a prove.
Corpus and fit quality (2026-08-03, this box):
- stage 1: exact FFT of all 34 Init shards at a 250 GiB pack (execute
stats dumps). MAPE 5.0%, worst under -18.5%.
- stage 2: 13 measured Init shard proves spanning 21-81 BFFT. RSS
max|err| 9.9%, wall 12.2%. Composed end-to-end on the proved set:
RAM within 12.4%.
- generalization: 60 InitStd shards the fit never saw. Stage-1 MAPE
5.7%, worst under -15.9% — inside the tail bound.
- execute models refit on the same sweep (wall tracks def_eq alone,
MAPE 9.6%; RSS bytes+def_eq, 3.7%).
AIUR_RAM_USABLE_FRAC rises 0.70 -> 0.75: the composed worst-case
under-prediction (stage-1 tail x stage-2 tail ~ -25%) inverts to 0.797,
leaving the difference as OS/variance margin — headroom that is now
derived from measurement instead of absorbing an unpriced blind spot.
Validation, end to end: repacking Init at 250 GiB yields 32 shards
(stale model: 34, one of which measured unprovable at a projected ~277
GiB); under the new pack no shard's exact-FFT-mapped RAM exceeds the
budget, and the heaviest (88.2 BFFT) proved green at 227.2 GiB measured
vs 220.9 mapped (+2.9%).
Also fixed: both greedy pack loops' tentative cap-test cost accumulated
only bytes+hb before testing — harmless when RAM ignored subst/def_eq,
an under-test under any model with reduction terms. The ShardPipeline
E2E now cuts its tiny env with the fixed-count min-cut partitioner: the
composed model base (~19.7 GiB) exceeds anything a ~200-constant env can
fill, so a RAM budget either degenerates to singletons or one shard;
the claim layer is partition-agnostic and the RAM packer keeps its own
unit coverage.
The witness IOBuffer materialized a shard's full bfs_closure eagerly, 8x-expanded (one G per byte), before the kernel executed a single query — ~2.5 GiB per Init shard buffer, ~30 GiB at Mathlib scale, per in-flight shard. The kernel's laziness was circuit-side only: shipped- but-unfaulted bytes cost zero rows, but the host always paid for the whole closure. IOBuffer gains an optional IOFaultSource backing consulted on a get_info miss: EnvFaultSource resolves ch 2/3/4 keys against the shared mmap'd env and materializes entries on first fault, so host witness RAM scales with the FAULTED set — what the check actually touches — instead of the shipped closure. Soundness-neutral: faulted bytes are blake3-verified against their content-addressed key exactly like eagerly-seeded ones, and the idx/len wires are unconstrained witness either way. Trace generation reads through a frozen lookup (get_info_frozen): it runs circuits in parallel over a shared buffer and never needs to fault — execution already materialized every entry it read. EnvHandle holds Arc<Env> so a backing outlives any handle borrow. The claim-channel seeding (ch 0/1) is factored into seed_shard_check_env_claim, shared by the eager and lazy builders, and all four with-env check/prove FFI paths now use the lazy variants.
Shard boundaries come from executing the env's check schedule through the codegen'd Aiur kernel — thin-frontier CheckEnv claims, the exact unit a prover pays for — with a running cost readout, replacing the counter-prediction path for Aiur. `ix shard env.ixe` runs the scan by default; `--backend zisk` chains profile + guest-cost pack; the .ixprof packer is unchanged for Zisk. - Two-term cut: a shard ends when RAM-per-fft * fft + record bytes reaches the budget's usable headroom — both resources a prove consumes. - One RAM mechanism: a per-worker record-byte cap ((0.70*box - 10)/workers, 8 GiB autoscale floor; --workers 1 grants the whole allowance), enforced inside execution via thread-local accounting in QueryMap::insert and codegen'd record_budget_exceeded() polls (ExecError::RecordBudgetExceeded). A trip mid-segment is an early cut; a block over the cap alone is skipped and named. Fleet RAM is workers x cap by construction — no reactive control. - True costs: every shard is re-priced with one cold CheckEnv — the claim its prove executes — and the merge pass consolidates on true costs to a fixpoint. - Manifest v2: tagged per-shard cost (Unknown / ProfileHeartbeats / ZiskCostUnits / AiurFft, raw fft units), version-checked in both the Rust and Lean parsers; ix prove orders heaviest-first from the manifest; the costs CSV sidecar is report-only. - Parallel scan over chunked schedule with work stealing; --fail-fast (default) / --no-fail-fast uniform across ix shard / check / prove, skipped blocks inventoried in failed.csv; IX_SCAN_LIMIT_BLOCKS slice knob for reproducers. - aiur-shard bench backend gains whole-env rows: scan-time, peak-rss, shards, max-shard-fft, and fft-cost in raw fft units so env totals join the existing per-constant plot. Validated at 250 GiB: Init 40 shards / 1,719 BFFT total (152 s, peak 91 GiB); InitStd 76 / 3,330 (255 s, 127 GiB); Lean 108 / 5,086 (383 s, 123 GiB); zero skipped blocks; deterministic partitions; prove spot check 74.8 GiB real vs 76.1 predicted.
samuelburnham
force-pushed
the
sb/aiur-planner
branch
from
August 4, 2026 23:55
e61a9c8 to
005d9c9
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Aiur sharding: measured scan-and-cut
What this does
Shard boundaries for Aiur proving are now measured, not predicted.
ix shard env.ixeexecutes the env's check schedule through the codegen'd Aiur kernel — the same thin-frontierCheckEnvclaims a prover pays for — with a running cost readout, and cuts a shard wherever the measured cost reaches the RAM budget. The previous counter-based cost model (fit on ingress bytes / substitutions / def-eq counts) mispredicted real content by up to 16× in both directions and OOM'd shards it called safe; execution is the mandatory prefix of proving, so measuring it prices the prove exactly.Design
RAM-per-fft · measured fft + measured record bytesreaches the budget's usable headroom — both resources the prove actually consumes.0.70·box RAM / workers), enforced inside kernel execution. Fleet RAM isworkers × capby construction. A block that outgrows the cap alone is skipped and named infailed.csv, same as a kernel-rejected block — pathological content is excluded by name, not engineered around.CheckEnv— exactly the claim its prove executes — so manifest costs carry no scan-side inflation. The.ixesmanifest (format v2) stores them per shard;ix proveorders heaviest-first from it.IOBufferfaults env bytes on demand, so witness RAM tracks the touched set, not the shipped closure — for scan and prove alike.--fail-fast(default) /--no-fail-fastuniform acrossix shard/check/prove; whole-env rows in theaiur-shardbench backend (scan time, peak RSS, shards, fft totals).Validation (250 GiB budget, 32-core / 250 GB box)
All blocks covered (zero skips), partitions deterministic, every predicted prove RSS under budget. Scan cost is ~2% of the pack's prove cost. Measured prove check: shard predicted 76.1 GiB / 60.5 s proved at 74.8 GiB / 59.4 s.
Usage
The shard-pipeline E2E runs in
lake test -- --ignored ixvm(also in PR CI).Known limits