Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 72 additions & 0 deletions docs/src/content/docs/development/Architecture/model-manager.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -1020,6 +1020,78 @@ To write such an estimator (`invokeai/backend/util/vae_working_memory.py` has fi
- Ship one coefficient entry per measured platform value in `invokeai/backend/working_memory/data/`, either with its own `values` or with `values_from` another estimator's entry. Unmeasured platform values (MPS, XPU, CPU) are priced as `UNMEASURED_PRICED_AS` in `working_memory/platform.py` says.
- Ship the measured points next to them: peak reserved bytes (`reserved_peak`), one fresh process per point. `tests/backend/working_memory/` holds each estimate to at least its point, and where a fitted term decides it, to at most 1.3 times the point, or the point's tighter `max_ratio`. The `device_working_mem_gb` floor counts only where an analytic term decides. Every value an entry carries must be decided by at least one point.

#### Measuring estimates: the probe

`invokeai/backend/working_memory/probe.py` measures what an estimated operation really takes. It is active only inside a `calibration_scope(...)`, which the session runner opens per node of a calibration queue item; without one, `model_on_device()` behaves as above and reads no counter.

Inside a scope, every `model_on_device(working_memory=...)` on a GPU is a window. The lock reserves the estimate plus calibration headroom (half the estimate, at least 1 GiB), so a low estimate cannot spill. The window runs from the settled lock to the unlock and collects a `MeasuredWindow` in the scope's `windows`; nothing is written while it is open.

- Counters (`working_memory/counters.py`): torch's reserved and allocated peaks, and this process's usage as the driver counts it. That is D3DKMT on Windows for CUDA and ROCm (the only readers verified on hardware), NVML on Linux CUDA, DRM fdinfo on Linux ROCm, and the MPS driver figure. XPU's Sysman figure covers the whole device, so it only detects another process's load. MPS has no peaks; a sampler thread keeps them while a window is open.
- Demand: per counter, the largest segment peak less the baseline and the bytes co-locks moved onto the device; `demand_bytes` is the largest of those deltas. Each counter misses something: the driver misses HIP's retained pool, and reserved misses cuBLAS and cuDNN workspaces and, under `cudaMallocAsync`, part of the allocated peak.
- Segments: the window opens after a release, so cached blocks and tensors from before the lock form the baseline. Another model locked on the same thread inside the window is a co-lock: its weights go into the offset, and its lock interval is kept out of the peak. An estimated window inside another measures alone while the outer one is frozen. A `TorchDevice.empty_cache()` ends a segment. A direct `torch.cuda.empty_cache()` call (from a library, or the one cuDNN and MIOpen run after an algorithm search) and an offload by the cache's background thread do not: the driver figure can miss what they released, while the reserved and allocated peaks still bound the demand.
- Invalid (`valid=False`, with `reasons`): another thread locked a model on the same device's cache during the window, held one when it opened, or was still locking one then (`foreign_lock`); the allocator failed or retried; Windows still holds more than 256 MiB of the process in shared memory after a short settle; the operation raised or was canceled, or its node failed or was canceled after the window closed; the Krea-2 SDPA benchmark is active; XPU's device-wide usage grew beyond this process's own (`foreign_device_load`); a probe step failed (`probe_error`: a reading, a release, a peak reset or a synchronization, which also voids the scope's later windows when it hit the lock bookkeeping; a window that cannot open is recorded invalid, and its operation runs unmeasured); or no counter could be read.
- Not fit-eligible: a model that was not fully resident when the window opened. The measurement stays valid, but streamed weights are working memory the estimate does not price.

`working_memory/fingerprint.py` keys measurements. `device_class_for()` gives the same id to identical cards. `environment_for()` splits hard facts, which change an operation's memory and make up its id (the operating system as `windows`, `linux` or `other`, torch, runtime and convolution-library versions, allocator backend and configuration, the ROCm SDPA guard, the memory settings passed in), from soft facts kept for diagnosis (full versions, driver, OS build).

#### Calibration runs

A calibration run measures one model's operations on this machine's GPU. Runs only record; no estimate reads them yet. The service is `invokeai/app/services/working_memory_calibration/`, the routes are in `invokeai/app/api/routers/working_memory.py`, and every route is admin-only (403 otherwise).

1. `POST /api/v2/models/i/{key}/working_memory/calibrate` takes a client-built graph, batch data and an optional `idempotency_key`. The server turns off `use_cache` on every node, sets `runs=1`, sets the graph's id to the run id, and sets the origin to `invokeai:wm-calibration:<run_id>` and the destination to `invokeai:wm-calibration`. It records the run, enqueues the batch as the requesting admin, then records the run's items.
- The queue takes all of the batch's items or none (`enqueue_batch(..., all_or_none=True)`, checked under its admission lock): a run with only some of its cases would measure only some of them. Without room for all of them nothing is enqueued, the key stays unused, and the run is removed if this request recorded it.
- With a key, the run id is derived from the owner, the key and the request's fingerprint: the model key, the graph without its id, and the batch data. A client that sends no graph id would otherwise get a fresh one each time and could never repeat a request. A repeated request finds its run and gets the same response; nothing new is enqueued, even after the queue collected the key's receipt. Once a reset removed the run, a repeat is refused (404) whether or not the queue still keeps the receipt: the store keeps the ids of removed runs, so the run is never started again and its queued items keep failing at admission. Another request that reuses the key (another model, graph or batch data) gets another run id and so another origin. While the key's receipt exists, the queue refuses that request (409) and the first run stays. Once the queue has collected the acknowledged receipt, the request starts a run of its own. Acknowledge the key through `POST /api/v1/queue/{queue_id}/enqueue_batch/acknowledge`.
- Once its items are queued, the request answers with them, also when the enqueue raised after its commit (the queue announces a batch after committing it). When the queue cannot tell whether it holds them, the run is kept and the enqueue's own error answers.
- If the store fails for a passing reason (an unwritable root, a lock timeout) to record the queued items, a request without a key still succeeds: its items run, and the status finds them by their origin. With a key that is a 503, and a repeat records them: from the key's receipt while the queue keeps it, otherwise from the queue items with the run's origin and the items the runner saw. It never enqueues a second batch for the run. Recording them from the receipt checks neither the model nor the device, so a model removed meanwhile does not fail the repeat; a request that has to enqueue needs the model. A run removed meanwhile (404) or an unusable store file (409) fails the request with or without a key; its queued items then fail at admission.
- Per run id, one creation runs at a time, and it runs to its end even when its request is canceled: the enqueue runs in a worker thread and can commit after the cancellation. Only the server's shutdown cancels the creation with it; the run then stays as far as it got, and with a key a repeat of the request completes it. The lock belongs to the process, so two processes serving one root and its SQLite database can interleave the creations of one run. A run is removed only when this request recorded it and nothing of it was queued.
- The destination only stops clients from routing the results to a gallery or canvas view. Images the graph's nodes save are stored like those of any other session: on the board and with the intermediate flag each node sets.
2. `enqueue_batch` refuses any origin that starts with that prefix, in any case, for every user (400). Only the calibration route creates calibration items. Retries and workflow-call children inherit origin and owner from their source item.
3. `DefaultSessionRunner.run()` checks each item's origin once. For a calibration item it calls `admit()` before any node runs, also for one left pending before a restart: the invoker starts the session processor after every other service. (The queue's start cancels the items that were in progress.) Retries and children are checked here too. Otherwise the item fails with the reason. The checks:
- the owner is an active admin (single-user mode counts as admin);
- the machine has exactly one GPU generation device;
- the run exists for that GPU's device class.
4. Each admitted node runs in `probe.calibration_scope(...)`, keyed by the node's source id. Its windows are persisted after the node completes, fails or is canceled. Unless the node completed, every one of them is persisted invalid (`exception` when it failed, `canceled` otherwise), including a window whose own operation finished before the node raised: `run_node` swallows the failure or cancellation, so the runner takes the outcome from its node hooks. Windows from a device of another class, or of a run that was removed, are dropped with a warning. So are windows of a model reset after the node started, also by another process sharing the root: the runner reads the device class's reset epochs before the node's scope opens, and the write drops the points of every model whose epoch has moved since. Without that, a reset while the node runs (of a VAE inside a transformer's run, say) would be undone by the node's later write. If they cannot be saved, the item fails and the session stops. Every other item gets no scope and runs as before.
5. After the item, the runner records the item's outcome in the run, even when the session raised. An item the queue no longer has (Clear Queue) is recorded as canceled. An item that was still running when the runner raised is recorded as failed with the error.

The store has one JSON file per device class at `<db_dir>/working_memory/<device_class_id>.json` (`CalibrationStoreFile`, `schema_version` 1). It is machine data, so it lives next to the database rather than inside it. It holds:

- the device class, with its id;
- the runs: model key, model hash as provenance, owner, environment, queue items, and the last status the runner saw for each item;
- the points: `MeasuredWindow` plus `env_id`;
- the environments those reference;
- the ids of the runs with a retry key that a reset removed (`removed_run_ids`), so that a repeat of the request that started one cannot start it again. They are kept for good, about 40 bytes per removed run. A run without a key gets a fresh id and is never repeated, so its id is not kept;
- per model key, its reset epoch (`model_reset_epochs`): how often the model was reset. The reset raises it in the same file change, so it is only written when the reset is.

A file without either of these two fields reads them as empty.

Writes and reads:

- Every change takes a process lock and an OS lock on the file's `.lock` sibling, re-reads the file, applies the change, and replaces the file atomically. Two processes sharing a root therefore never lose each other's writes.
- On Windows a process that ignores the lock (an antivirus scanner, a search indexer, an editor) can hold the file open and refuse the replace. The write then retries a few times with backoff, about 1.5 s in all (`REPLACE_ATTEMPTS`, `REPLACE_RETRY_S`).
- Readers take the lock as well, because Windows cannot replace a file another process has open.
- Each change rewrites the whole file, which is fine at calibration volumes.

A file is unusable when it has another schema version, invalid content, an id that does not match its fields or its file name, or a field this version does not know. Every store model forbids unknown fields: ignoring one would drop it at the next rewrite. An unusable file is reported in the status, makes calibration unavailable for its device class, and is never written, moved or deleted.

A run's status is derived, never stored:

- The live queue items with the run's origin come first. For items the queue has pruned, the runner's recorded statuses are used. An item the queue no longer has and that the runner never saw finish (the process died mid-item) counts as canceled.
- A retry replaces the item it retried. The queue's own record of the retry links them while it has the retry, also before the runner started it; afterwards, what the runner saw does. A retry the runner never started and the queue no longer has leaves no trace: the run then reads as the outcome of the item it retried, failed for instance.
- The run is `running` while any item is pending or running, `failed` when any failed, and `completed` when all completed. Otherwise it is `partial` if valid points were recorded, or `canceled` if not: an invalid point measured nothing usable.

Routes, besides the 403 for non-admins and a 503 when the store cannot be read or written (an unwritable root, a lock timeout):

| Route | Does | Other errors |
| --- | --- | --- |
| `POST .../i/{key}/working_memory/calibrate` | starts a run (above) | 404 unknown model, the run was removed while it was being created, or a reset removed the run the retry key started; 409 no GPU or several generation devices, unusable store file, no room in the queue for all of the run's items (nothing is enqueued; the run is removed if this request recorded it), or a retry key used for another request; the 404 and 409 for the run's store, and with a key the 503, can come after the items were queued (above); 422 a `call_saved_workflow` node, batch data setting `use_cache`, an invalid batch, or fields the server sets; 429 too many unacknowledged retry keys |
| `GET .../working_memory/status` | every device class with its runs, point counts per estimator, version, validity and fit eligibility, and environments with hard and soft facts compared to the current environment; `available` says whether a run can start now (one GPU generation device and a usable store file) | – |
| `DELETE .../working_memory/runs/{run_id}` | removes a run and its points; its items that have not started yet fail at admission, also after a repeat of the request that started it | 404 no usable file holds the run |
| `DELETE .../i/{key}/working_memory` | removes the model's runs with all their points, and every point of the model taken in other runs (a VAE measured inside a transformer's run, say) | 409 no current device class, or an unusable file |
| `DELETE .../working_memory` | removes every run, point and environment of a device class | 409 as above |
| `GET .../working_memory/export` | the file byte for byte: a `CalibrationStoreFile` when this version can read it, otherwise whatever the file holds | 404 no file; 409 no current device class |

The last three act on the current device class unless `device_class_id` is given.

### Emitting model loading events

When the `context` argument is passed to `load_model_*()`, it will retrieve the invocation event bus from the passed `InvocationContext` object to emit events on the invocation bus. The two events are "model_load_started" and "model_load_completed". Both carry the following payload:
Expand Down
2 changes: 2 additions & 0 deletions docs/src/content/docs/development/Architecture/overview.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,8 @@ The invoker (`/invokeai/app/services/invoker.py`) is the primary interface throu
- **invocation services**, which are used by invocations to interact with core functionality.
- **invoker services**, which are used by the invoker to manage sessions and manage the invocation queue.

Creating the invoker starts every service, the session processor last: its workers dequeue at once, and an item left pending before the restart may use any other service. (The queue's start cancels the items that were in progress.) `stop()` stops the session processor first: its workers stop taking items and cancel their sessions, and the other services stop only once the workers have exited. The processor waits up to 30 seconds for all its workers together, and after a second still waiting it logs which worker and queue item it waits for. A worker still in a node at that deadline (one without a cancel point) is given up on with an error in the log, so that shutdown finishes; a later `stop()` does not wait for it again. It keeps its GPU, so until it exits no session processor in the same process starts and none replaces the generation device pool: `ApiDependencies.initialize` refuses a restart before it opens the database.

### Sessions

Invocations and links between them form a graph, which is maintained in a session. Sessions can be queued for invocation, which will execute their graph (either the next ready invocation, or all invocations). Sessions also maintain execution history for the graph (including storage of any outputs). An invocation may be added to a session at any time, and there is capability to add and entire graph at once, as well as to automatically link new invocations to previous invocations. Invocations can not be deleted or modified once added.
Expand Down
Loading
Loading