[Pfannkuchensack]

Stacked on #9896. Until that merges, the diff also shows its commit; this PR's change is 470db713f5.

Summary

Working-memory estimates can now be checked against what an operation really uses, on the machine that runs it. This PR only records measurements; no estimate changes.

• Probe (invokeai/backend/working_memory/probe.py)
  • During a calibration run, every model_on_device(working_memory=...) window records its demand next to its estimate.
  • Demand is max(own-process driver delta, torch reserved-peak delta, allocated-peak delta), measured from a released baseline.
  • The driver delta comes from D3DKMT on Windows (NVIDIA and AMD), NVML or fdinfo on Linux, and the driver counter on MPS.
  • It locks with headroom, so a too-low estimate cannot spill while it is being measured.
  • It handles co-locked models and nested windows.
  • It marks a window invalid on a foreign lock on the same device, OOM, allocator retries, persistent spill to shared memory, exceptions, or probe errors.
  • Partial residency keeps a window but excludes it from fitting.
  • Without a calibration scope the cost is one ContextVar read per estimated lock.
• Runs (invokeai/app/services/working_memory_calibration/)
  • An admin enqueues a calibration batch through POST /api/v2/models/i/{key}/working_memory/calibrate.
  • Its items carry a reserved origin that enqueue_batch refuses for everyone.
  • At session start the runner re-admits each item: active admin owner, exactly one GPU generation device, run present for this device class.
  • It opens the probe scope per node and stores the windows in a JSON file per device class under databases/working_memory/.
  • The file is written atomically and under a file lock; files it cannot read are kept and never overwritten.
  • Status, reset and export routes are admin-only.

Related Issues / Discussions

Builds on #9896.

QA Instructions

• Tests: affected backend and app suites (working_memory, util, model_manager/load, the calibration service, runner and router tests, session queue/processor): 2079 passed, 115 skipped. Each new file also passes alone.
  • Mutation checks: co-lock offset, nested suspend, foreign lock at open, RAM-pin and unlock bookkeeping, admission refusals, retry lineage and the store's OS lock are each caught.
• Hardware: pytest -m slow tests/backend/working_memory/test_probe.py passed on an RTX 4090 (CUDA 13.0) and an RX 9060 XT (ROCm 10, Windows): the window is valid and covers the reserved growth.
• End to end on the RTX 4090 (server from this branch, cudaMallocAsync):
  • A FLUX VAE encode→decode at 1024² through the calibrate endpoint completed and recorded two valid, verified windows.
    • flux_vae_encode: reserved 2200 MiB, measured 1840 MiB.
    • flux_vae_decode: reserved 4400 MiB, measured 2912 MiB.
  • The status call did not change the process's VRAM.
• API: regenerated openapi.json/schema.ts; pnpm lint and pnpm test in invokeai/frontend/api pass.

Review

Resolved findings

• Runs could report "running" forever after Clear Queue, a crash or a prune.
  • Fix: the outcome is recorded in finally, and an item that is gone counts as canceled.
• Probe errors could leave windows valid with corrupted figures.
  • Fix: every hook, observer or close failure marks the affected windows.
• A foreign lock could slip in between the lock count and window registration at open.
  • Fix: the window is registered before the count.
• calibrate was not safe to retry after a lost response.
  • Fix: it takes an idempotency key.
• Windows os.replace failed transiently while another process held the file.
  • Fix: bounded retry.
• The ROCm attention and conv3d switches were missing from the environment fingerprint.
• Status available and the declared error codes now match what the routes do.

Remaining limits

• The Linux NVML, Linux fdinfo, MPS and XPU readers follow vendor documentation, untested on hardware. They are marked verified=false.
• Direct torch.cuda.empty_cache() calls and background offloads bypass the segment boundary. The torch peaks still bound the demand.
• Multi-GPU calibration is refused by design.

Compatibility / Rollout

• New admin-only routes; regenerated API artifacts.
• No database migration. Calibration data lives in <db_dir>/working_memory/<device_class>.json.
• Nothing reads it for estimates yet.

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)

---

^{Need help on this PR? Tag @codesmith-bot with what you need.}

[JPPhoto]

@Pfannkuchensack Your review has arrived. Merge blockers:

1. Pending calibration work can fail during startup. invocation_services.py:152 starts the session processor before the calibration service. A resumed queue item can reach admit() before the calibration service sets _invoker in working_memory_calibration_default.py:178; the item then fails before measurement. Start the calibration service before session workers, or remove that ordering dependency. If skipped: a restart with pending calibration work can lose it. Likelihood: only when work is pending at startup. Recovery: rerun the failed request. Test: boot with a pending calibration item and verify it is not failed before service initialization.

2. Failed or canceled nodes can produce valid measurements. session_processor_default.py:314 wraps run_node() in the probe scope, but run_node() catches and swallows errors and cancellations at session_processor_default.py:462. The scope therefore closes without an error, so its windows can remain valid. Propagate the node outcome to the scope. If skipped: failed executions can be treated as trustworthy calibration data. Likelihood: whenever a measured node fails or is canceled. Recovery: remove the affected measurements and rerun. Test: fail and cancel an invocation through DefaultSessionRunner and assert its windows are invalid.

3. An idempotency conflict can delete a run another request is using. In working_memory_calibration_default.py:314, the request that first creates a run deletes it on enqueue failure. A concurrent request with the same key but different payload can enqueue successfully first; the creator then deletes the shared run. Cancellation during the threaded enqueue has a similar risk. Protect creation and enqueue as one operation, and only discard a run when no enqueue committed. If skipped: queue items can lose their run and measurement provenance. Likelihood: concurrent or canceled same-key requests. Recovery: remove orphan queue items and rerun. Test: race different payloads under one key and cancel during enqueue.

4. A key does not uniquely identify the calibration request. The run ID uses user and idempotency key, while the queue payload hash omits the model route. Reusing the same key and graph for model B after model A can return model A’s run as a success. Also, queue receipts expire after seven days, so a later retry can enqueue again into the persistent run. Persist and compare a request fingerprint that includes the model key and batch data. If skipped: results may have the wrong model provenance or duplicate samples. Likelihood: key reuse across models or after receipt expiry. Recovery: discard the incorrect run and retry with a new key. Test: reuse a key across model paths and after receipt expiry.

5. Partial batch admission is reported as success. working_memory_calibration_default.py:316 rejects only when zero items enqueue. The queue can admit only the available capacity, so a multi-item request can complete with a partial sample set. Require all items to enqueue or mark the run incomplete and clean up admitted items. If skipped: incomplete calibration data can appear complete. Likelihood: queue near capacity. Recovery: discard the partial run and rerun. Test: submit multiple batch values with capacity for only one.

Other findings/issues:

6. Counter failures can still yield valid windows. counters.py:204 swallows peak-reset and synchronization errors; callers continue measuring in probe.py:402 and probe.py:519. Return failure or raise a probe-local error and mark the window invalid. If skipped: stale peaks or unsettled GPU work can contaminate trusted measurements. Likelihood: transient device/backend errors. Recovery: remove the affected point and rerun. Test: make reset or synchronization fail and assert the window is invalid.

7. A lock can miss probe hooks during an arm race. load_base.py:193 snapshots probe.armed before locking. If it reads false, pauses, then the probe arms and opens a window, the lock proceeds without hooks and can contaminate that window. Coordinate the snapshot and hook registration, or ensure the hooks cover the interleaving. If skipped: concurrent model loading can be included in valid demand. Likelihood: narrow thread interleaving during calibration. Recovery: rerun the measurement. Test: use a barrier to force the interleaving.

8. The generated API contract omits a new HTTP 400. session_queue.py:272 returns 400 for a reserved origin, but the route declaration and generated openapi.json:12884 do not list it. Declare the response and regenerate the API artifact. If skipped: generated clients and docs miss a supported error response. Likelihood: client depends on generated contract. Recovery: regenerate and update clients as needed. Test: verify the route’s OpenAPI response includes 400.

[Pfannkuchensack]

@JPPhoto Thanks for the thorough review — all eight points are addressed in 3ebdb7f (merged with main in 4fbf583).

1. Startup order — Invoker now starts every other service before the session processor and signals the processor to stop first; the generation-device pool is registered before any service starts. test_invoker.py boots with a pending calibration item and forces the worst-case interleaving.
2. Failed/canceled nodes — the runner records each node's outcome; the windows of a node that did not complete are stored invalid (exception / canceled). Tested through DefaultSessionRunner for both.
3. Idempotency conflict — creation is serialized per run id and shielded from request cancellation; a run is removed only if this request recorded it and none of its items are queued. Tests race two payloads under one key, cancel during the enqueue, and raise after the enqueue committed.
4. Request fingerprint — keyed run ids derive from the model key, the graph (its id excluded) and the batch data; a repeat answers from the key's receipt while it exists and adopts the run's queued items after it expired, so it never enqueues a second batch. A different model or data under the same key gets 409.
5. Partial admission — enqueue_batch(..., all_or_none=True) admits every item or none, rechecked under the admission lock; calibration uses it and answers 409 with nothing queued. Covered on every DB backend in TestRaces.
6. Counter failures — failed peak resets and synchronizations raise; the window is stored invalid (probe_error), including windows that cannot open.
7. Arm race — a lock registers itself before reading probe.armed; a window that opens while a lock on its cache is in progress is invalid (foreign_lock). Event-driven tests force both interleavings.
8. API contract — the reserved-origin 400 is declared; OpenAPI and types are regenerated.

Also: a keyless request whose items could not be recorded answers 201 with its queued items (they run) instead of an error; a run whose items were all canceled reads canceled rather than partial.

[JPPhoto]

New merge blockers:

• High - Shutdown can overlap old workers with teardown and restart. invokeai/app/services/invoker.py:39 signals workers, then immediately stops their services. invokeai/app/services/session_processor/session_processor_default.py:934 does not wait for them to exit. A same-process restart can replace the device pool while an old worker still uses the GPU or stopped services. Fix/test: Wait for workers to exit before service teardown and pool replacement. Test stopping and reinitializing while a worker is held in a node.

Other findings/issues:

• Medium - A detached locked model can escape probe validation. invokeai/backend/model_manager/load/model_cache/model_cache.py:1892 counts locks only on cached records, but invokeai/backend/model_manager/load/model_cache/model_cache.py:2993 can detach a still-locked record. A calibration window can miss the foreign lock and record contaminated memory as valid. Fix/test: Count active locks independently of cache membership. Open a window while a detached record remains locked.

• Medium - reset_model can be undone by a pending measurement write. invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py:560 checks the run before writing. If another run measures the reset model, its later write can append the point after reset_model removes it at line 754. Fix/test: Coordinate reset and writes. Block a write across reset and verify the point stays absent.

• Medium - Retry status loses lineage until admission. invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py:593 reads live statuses without retried_from_item_id; derive_run_status can only link retries already observed by the runner. A canceled or pruned retry before admission leaves the original failure effective and may omit the retry from run item IDs, contrary to the documented retry contract. Fix/test: Preserve live retry lineage. Cancel and prune a retry before admission.

• Medium - Keyed recovery can fail if the model was removed. For a keyed run whose queue items were committed but not recorded, _prepare_run at invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py:403 still looks up the model at line 416 before replaying the queue receipt. Removing the model before retry turns a recoverable request into a 404. Fix/test: Replay the matching receipt without requiring the model. Test after removing it.

• Medium - Resetting a run does not invalidate its idempotency key. reset_run at invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py:746 removes the stored run but leaves queue receipts and items. Repeating the same keyed request can recreate the run and allow pending items to pass admission after reset; an expired receipt can also permit duplicate enqueue. Fix/test: Invalidate or tombstone the key on reset. Retry the same request with retained and expired receipts.

[Pfannkuchensack]

@JPPhoto Thanks for the second pass. All six points are addressed in b012444. Each one has a test that fails without its fix.

High: shutdown overlapping old workers

• DefaultSessionProcessor.stop() now joins the workers against one shared 30 s deadline, after it signals stop and cancels their sessions. If it is still waiting after 1 s, it logs which worker and queue item it is waiting for.
• A worker still in its node at the deadline is logged as an error and recorded as abandoned. A repeated stop() does not wait for it again.
• While an abandoned worker is alive:
  • register_generation_devices() and start() refuse with a RuntimeError, so nothing replaces the device pool or starts a second worker on that GPU.
  • ApiDependencies.initialize refuses before it opens the database.
• The database is still disposed after the stop. A late transaction from an abandoned worker is refused by the database's thread claim.
• Why bounded: an unbounded join hangs shutdown forever on a node that never reaches a cancel point. The restart refusal closes the overlap a bounded join would otherwise leave.
• Tests: tests/app/services/test_invoker.py and tests/app/test_api_dependencies.py.

A detached locked record escaping the probe

• ModelCache.locks_held() now reads a counter that every lock and unlock keeps, independent of cache membership. That covers the lock/continue_lock failure paths too.
• CacheRecord.unlock() refuses an unlock without a lock before it decrements, so the counter cannot drift.
• Tests: test_model_cache_locks_held.py, including a dropped-while-locked record and both OOM paths.

reset_model undone by a pending write

• A reset raises a per-model epoch in the store file (model_reset_epochs), in the same file change, so the epoch exists only if the reset committed.
• The runner reads the epochs when the calibration node starts, before the probe scope opens.
• The write drops, under the file lock, every point attributed to a model whose epoch moved since then, including submodel windows such as <key>:vae.
• This also holds for a reset made by another process on the same root.
• Tests: a reset while the node runs, a reset by a second service instance, an in-flight write, and submodel attribution.

Retry lineage before admission

• The run's live read now returns retried_from_item_id (SessionQueueBase.get_queue_item_lineage, one indexed statement whose plan is pinned by a test). A retry stands in for its item before the runner starts it, and a canceled retry counts as the run's outcome.
• Remaining limit, documented: a retry that the queue prunes or clears before the runner ever saw it leaves no trace, so the run reads as the original item's outcome. Covering that would need a hook where the queue creates retries.

Keyed recovery after the model was removed

• When the run is recorded without items and the key's receipt exists, _prepare_run replays the receipt before any model or device lookup.
• A new request still requires the model.

Reset not invalidating the idempotency key

• A reset adds the removed keyed runs to removed_run_ids in the store file. Unkeyed runs are not added, since they cannot be repeated.
• A repeat of the keyed request after a reset is refused with the existing 404, whether the receipt is retained or already collected, on every device class.
• The reset run's pending items keep failing at admission.

Store file: the two new fields (removed_run_ids, model_reset_epochs) have empty defaults. The schema version stays 1, so existing files read unchanged.

Contract: no API shape or status changes. The OpenAPI descriptions are regenerated.

The stack above (#9899, #9910, #9919, #9927, #9946) has these merged in.

[JPPhoto]

Thanks; the six findings from the previous review appear fixed.

Merge blockers

1. Medium - Status can say calibration is available when the store cannot be written. In invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py:671-693, status checks whether the store file can be read, not whether a run can write it. A read-only store can report available=true, then fail run creation with 503. Also, device_class_ids() can let a directory PermissionError escape as 500 (working_memory_calibration_store.py:140-144). Cover read-only and unreadable store directories.

2. Medium - A keyed retry can enqueue a second batch after an ambiguous failure. If queue admission succeeds but recording item IDs fails, the run remains with no IDs. If the key is acknowledged after that 503, its receipt expires, and the queue items are cleared or pruned before observation, a repeat falls through to enqueue at working_memory_calibration_default.py:423-450. Preserve an admission marker or refuse to re-enqueue when admission is uncertain. Test that sequence.

Other issues

3. Low - A completed item can later appear canceled or partial. The queue records completion before the processor persists the terminal calibration observation (session_processor_default.py:660-675,327-330). If the process stops in between and the completed row is later pruned, the saved in_progress observation is treated as canceled (working_memory_calibration_default.py:151-152). Test completion followed by pruning before observation.

4. Low - A conflicting key can return 404 instead of the documented 409. With a different request that names a missing model, _prepare_run() looks up the model before queue admission detects the key conflict (working_memory_calibration_default.py:423-450). The endpoint documents 409 for key reuse with another request (invokeai/app/api/routers/working_memory.py:95-100).

5. Low - Run summaries can omit a retry's source item. derive_run_status() uses retried_from_item_id to determine effective status but does not add that source ID to the returned item_ids (working_memory_calibration_default.py:144-166). If the source is absent from stored, observed, and live IDs, the summary violates its "every queue item" contract.

[Pfannkuchensack]

@JPPhoto Thanks. All five points are addressed in 1370b23, and each has a test that fails without its fix.

1. Status available on a store it cannot write. Status now checks that the current device class's store can be written. It opens an OS-deleted temporary file in the store folder, or in the nearest existing folder above it, and checks that the lock file is writable. On Windows it also checks that the JSON file is not read-only. A store that fails this reports available=false with the reason. device_class_ids() wraps listing errors as a store error, so an unreadable store directory answers 503 instead of 500. Tests cover a read-only lock file, a store root that is a file, and an unlistable directory.

2. A second batch after an ambiguous failure. A run now carries enqueue_attempted. A new run is recorded with it already set, in the same write. Only the attempt that set it clears it, and only when its enqueue is known to have queued nothing. A receipt replay never clears it. A repeat that finds a marked run with no items, no receipt and nothing queued or observed enqueues nothing. It returns the run with empty item_ids, and the run reads canceled. Tests: your sequence; a replay whose enqueue raises before the receipt is collected; and a full queue followed by a later repeat.

3. A completed item later reading canceled. On a clean finish (no error, not canceled, and the queue row is neither failed nor canceled), the runner now records the item as completed before the queue completes it. The final observation still writes the queue's status. Tests:

• the outcome survives a prune;
• the order is the runner's write, then the queue's;
• a failing node, and a failed save, are not recorded as completed.

5. 404 instead of 409. For a keyed request with no run, the reset check runs first (404), then the receipt check (409), and only then the model lookup.

7. Retry sources missing from item_ids. derive_run_status adds every lineage source to item_ids.

Limits, documented:

• If the process stops between the runner's completed record and the queue's completion, and the row is later pruned, the run reads completed. All of that item's nodes did finish.
• If clearing the mark fails, it is only logged. The run then stays marked, so that key cannot start it, though a request with another key can.

Store file: one new field with a default; schema version 1. API: no shape or status changes. The OpenAPI descriptions are regenerated.

The stack above has this merged in.

[JPPhoto]

Findings:

1. Medium - Cross-process requests can return an empty success or enqueue twice. The per-run lock is process-local. If process A records the marked run and pauses before enqueue, process B can return 201 with item_ids=[] at invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py:437-464; A may then enqueue or remove the run. If two prepared attempts overlap and one resumes after the first receipt is collected, both can enqueue, and the later item IDs can replace the earlier ones (:369-382,492-503). Use cross-process coordination or an explicit in-progress response. Add a two-process regression for both sequences.

2. Medium - Final observation can overwrite a saved completion with canceled. The runner records completion before queue completion, but its finally block calls observe() afterward (invokeai/app/services/session_processor/session_processor_default.py:320-341). If the completed queue item is cleared or pruned in between, observe() records canceled over the saved completion (invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py:619-633,650-655). Preserve an existing completed observation and test that sequence.

[Pfannkuchensack]

@JPPhoto Thanks. Both points are addressed in 9f2345c, with tests that fail without the fixes.

1. Cross-process creation

• A keyed run's creation now holds an OS file lock from preparing the run to recording its items. It covers the enqueue and the cleanup when nothing was queued.

• The lock uses the store's own lock primitive (msvcrt.locking / flock), so a holder that dies releases it with its process.

• It is striped over 256 files, working_memory/creations/<xx>.lock, where xx comes from the run id's sha256. That bounds the folder and means no lock file is ever deleted, which would race. The cost is false contention between unrelated runs that share a stripe, which is negligible for this admin action.

• A per-run in-process gate sits in front of it, so concurrent repeats in one process wait on the event loop rather than each holding an executor thread.

• A contender waits up to 30 s and then gets the existing retryable 503. After acquiring the lock it sees what the first attempt did:

  • the recorded items;
  • the receipt to replay;
  • or, once that receipt is collected, the marked run.

  So it can neither return an empty success while the first attempt is enqueuing, nor enqueue a second batch, nor replace the first attempt's item ids.

• Tests: two service instances on one root.

  • A repeat waits while the first instance is held inside its enqueue.
  • A second request cannot prepare while the first is recording the run.
  • Eight concurrent repeats in one process leave one worker on the file lock.

  Separate instances model separate processes here, because the OS arbitrates these locks per handle (LockFileEx) and per open file description (flock).

2. Final observation overwriting a saved completion

• When the queue no longer has the item, the final observation keeps a recorded completion.
• A queue row that still exists stays authoritative, including a failed one. So the run reads the same before and after a prune.
• Tests: completion recorded, item pruned, then the final observation; and a row that is failed in the queue.

Also changed: main's #9855 wake-up test. Merging main brought in #9855's test_worker_wakeup.py. Two of its places assumed that stop() returns while workers are still parked on test gates. Since stop() now joins its workers, the test (62d75d8):

• opens its gates before the cleanup stop;
• runs the stop case on a thread of its own, continuing once every worker's own notification is set.

#9855's guarantee is unchanged. Making the workers share one poll_now_event again fails all three cases.

The stack above has this merged in.