Skip to content

feat(e2e): calibrate working memory and build shipped data from probe runs - #9899

Draft
Pfannkuchensack wants to merge 30 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/family-e2e-calibration
Draft

Pfannkuchensack wants to merge 30 commits into
invoke-ai:mainfrom
Pfannkuchensack:feat/family-e2e-calibration

Conversation

@Pfannkuchensack

@Pfannkuchensack Pfannkuchensack commented Oct 17, 2026 •

Copy link
Copy Markdown
Member

Stacked on #9894 (family E2E tool), #9897 (calibration runs) and #9898 (VAE coefficients per dtype). Until those merge, the diff also shows their commits. This PR's own changes are 7a75080b44 and 9267c3e166; the rest are merges of the three branches.

Summary

Shipped working-memory coefficients can now come from probe runs through the API, without one-off calibration scripts.

  • start --calibrate in the family E2E tool
    • Each compiled case goes through POST /api/v2/models/i/{key}/working_memory/calibrate instead of enqueue_batch.
    • The tool follows the items and saves the store export and status per server lifetime under the run's calibration/. After a crash or stop it copies the store file instead.
    • Time and VRAM are marked as calibration runs and are not rated.
    • A 404/422 for one case becomes that case's error record and the run continues; any other error response stops the run.
  • scripts/working_memory/build_shipped.py
    • Turns exports from one or more cards into probe_demand points. It takes valid, fit-eligible windows, keeps the largest per card, estimator, platform and features, and keeps the larger value when it meets an existing point.
    • Platform keys come from the recorded environment through platform_for_environment, which shares its rules with platform_for.
    • It refits the data, shows each line change and the spread across repetitions, and flags fragmentation: reserved grows ≥10 % while allocated stays flat.
    • --write refuses any result that the shared bound rule rejects. A raise driven by a fragmentation-flagged point needs --accept-fragmentation.
  • Probe points are lower bounds only. They are warm windows in a running server. The estimate must cover them and they may raise a line, but they never set its cap. The fit rates candidate lines only over fresh-process points, and the upper bound (1.3×, listed exceptions, max_ratio) stays with those points, because a cold first run after a server start is the worst case.
  • One bound rule: invokeai/backend/working_memory/bounds.py is used by both the data test and the builder.

Related Issues / Discussions

Builds on #9894, #9897 and #9898.

QA Instructions

  • Tests: the affected suites (working-memory backend, builder, tests/family_e2e, package data, invocations, the calibration service/runner/router) pass: 3438 passed, 8 skipped, 5 xfailed. Each new or changed file also passes alone.
    • Mutation checks are each caught:
      • the calibrate path falling back to enqueue_batch;
      • a missing availability check;
      • the OS taken from the build machine instead of the recording;
      • merging that lets the newer point win, or keeps both;
      • the bound refusal disabled;
      • first instead of largest window;
      • probe points choosing a line;
      • case refusals treated as fatal or as per-case.
  • Hardware, two real calibration runs (SDXL VAE manifests, fp16/fp32, 512–1280², 25 steps) on an RTX 4090 (CUDA, cudaMallocAsync) and an RX 9060 XT (ROCm, Windows). Both completed and exported their stores.
    • The builder on both exports reads 18 points (3 windows not fit-eligible because of partial residency).
    • It flags the RX 9060 XT SDXL fp32 1280² decode as fragmentation: across repetitions in one server, reserved rose from 8020 to 13000 MiB while allocated stayed at 8020 MiB.
    • It refuses the resulting refit (k_decode_fp32 5988 → 7661 B/px), which would push fresh-process points over 1.3×. The shipped data is unchanged.

Review

Resolved findings

  • --write could write data that failed the data test. The upper-bound check only looked at new points, and every new point is a probe point. Now one shared bound rule runs over all points after a refit.
  • The builder skipped the lower bound where an analytic term decides; that check now runs.
  • Lower-bound-only was enforced after the fit instead of in it; the fit now enforces it.
  • A rebuild could replace a shipped point with a lower one. Merging now keeps the larger value.
  • A single case refusal ended the whole run, which resume could not get past.
  • An idempotency key the tool never reused has been removed.
  • The os hard fact moved into feat(working-memory): measure working memory in admin calibration runs #9897, so the environment-id scheme ships once.
  • Tests were added for the availability check, the recorded OS, report history, stop, and the probe-wiring regression.

Remaining limits

  • Calibration runs are Windows-only, like the tool. Other machines' exports come from GET /api/v2/models/working_memory/export.
  • Warm-up windows count toward the maximum, so a cold first window can dominate.
  • fingerprint.os_family and platform._host_os repeat the same sys.platform mapping; a test catches drift between them.

Compatibility / Rollout

  • No shipped values change in this PR.
  • The new point method probe_demand holds lower bounds only.
  • MeasuredPoint gains holds_upper_bound; a max_ratio on a probe point is rejected.

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Meaningful regression coverage added / updated where needed; obsolete tests/code removed
  • Persisted-state and API changes include required migrations / compatibility validation
  • Relevant performance/efficiency opportunities considered; material claims have evidence
  • Material review findings resolved and relevant checks rerun
  • Documentation added / updated (if applicable)
  • Updated What's New copy (if doing a release after this PR)

View with [code]smith Autofix with [code]smith
Need help on this PR? Tag @codesmith-bot with what you need.

Runs each installed model family through a launched scratch server and compares PSNR, time and driver-level VRAM against accepted references.
Graphs come from the frontend builders via an offline graph export; the generate submission plan is shared with the workbench.
VAE estimators return a WorkingMemoryEstimate built from fitted and analytic terms; model_on_device accepts it as working_memory=.
The AutoencoderKL coefficients and their measured points move to backend/working_memory/data, checked by a data test; reservations are unchanged.
A probe records each estimated model_on_device window's device-level demand inside a calibration scope; results go to a per-device-class store under databases/working_memory.
Admin endpoints start, inspect, reset and export runs; the reserved queue origin is refused for ordinary enqueues.
VAE estimators price a measured slope and base per conv backend, allocator and dtype instead of scaling one bf16 constant by element size; CogView4 gets its own lines.
SD/SDXL nodes and their estimate share one run-dtype rule; on ROCm Windows an SDXL fp32 1024² decode reserves 6.2 instead of 14.6 GB and the UNet stays resident.
A lower second reading at a measured size no longer tilts a fitted line, and points that include cast weight growth no longer fit one; two shipped lines move by under 1 %.
…n fact

Environment ids now include windows, linux or other, so calibrations from different operating systems are kept apart; the OS build string moves to soft fact os_build.
… runs

start --calibrate submits each case through the calibration endpoint and saves the store export per server lifetime; scripts/working_memory/build_shipped.py turns exports into probe points and refits.
Probe points bound estimates from below only, and the builder refuses data that the shared bound rule rejects.
# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
@github-actions github-actions Bot added api python PRs that change python files Root invocations PRs that change invocations backend PRs that change backend files services PRs that change app services frontend-deps PRs that change frontend dependencies frontend PRs that change frontend files python-tests PRs that change python tests docs PRs that change docs python-deps PRs that change python dependencies labels Oct 17, 2026
# Conflicts:
#	invokeai/backend/util/vae_working_memory.py
#	tests/app/invocations/test_flux2_working_memory.py
# Conflicts:
#	invokeai/backend/util/vae_working_memory.py
Main added useNsfwChecker and useWatermark to the generation project settings; exported calibration and E2E graphs run without both, as the defaults do.
# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
#	invokeai/backend/model_manager/load/load_base.py
# Conflicts:
#	docs/src/content/docs/development/Architecture/model-manager.mdx
#	invokeai/backend/util/vae_working_memory.py
#	invokeai/backend/working_memory/coefficients.py
#	invokeai/backend/working_memory/data/vae.json
#	invokeai/backend/working_memory/platform.py
#	scripts/calibrate_flux2_working_memory.py
#	tests/app/invocations/test_flux2_working_memory.py
#	tests/app/invocations/test_vae_working_memory_backend.py
#	tests/app/invocations/test_z_image_tiled_decode.py
#	tests/backend/util/test_rocm_sdpa_head_dim_guard.py
#	tests/backend/working_memory/test_shipped_coefficients.py
… tests

CPU wheels on macOS and Windows have no CUDA allocator binding, so the fixture that scripts a CUDA build scripts its backend too.
…s a CUDA environment

The refused-environment test scripted the CUDA build but asked the CPU wheel's missing allocator binding.
@lstein lstein added the 7.0.0 label Oct 19, 2026
@lstein lstein moved this to 7.0 Theme: Modular Design in Invoke - Community Roadmap Oct 19, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

7.0.0 api backend PRs that change backend files docs PRs that change docs frontend PRs that change frontend files frontend-deps PRs that change frontend dependencies invocations PRs that change invocations python PRs that change python files python-deps PRs that change python dependencies python-tests PRs that change python tests Root services PRs that change app services

Projects

Status: 7.0 Theme: Modular Design

Development

Successfully merging this pull request may close these issues.

3 participants