From 855ebe2c02cd82d0ee81e95fb6f8e6aa73a06990 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 17 Oct 2026 04:57:58 +0200 Subject: [PATCH 1/9] feat(working-memory): ship VAE working-memory coefficients as data 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. --- .../contributing/new-model-integration.mdx | 4 +- .../Architecture/model-manager.mdx | 20 + .../vae/cogview4_image_to_latents.py | 2 +- .../vae/cogview4_latents_to_image.py | 2 +- .../app/invocations/vae/flux2_vae_decode.py | 2 +- .../app/invocations/vae/flux2_vae_encode.py | 2 +- .../app/invocations/vae/flux_vae_decode.py | 4 +- .../app/invocations/vae/flux_vae_encode.py | 2 +- .../vae/ideogram4_latents_to_image.py | 2 +- .../app/invocations/vae/image_to_latents.py | 2 +- .../app/invocations/vae/latents_to_image.py | 2 +- .../invocations/vae/sd3_image_to_latents.py | 2 +- .../invocations/vae/sd3_latents_to_image.py | 2 +- .../vae/z_image_image_to_latents.py | 2 +- .../vae/z_image_latents_to_image.py | 4 +- invokeai/backend/AGENTS.md | 1 + .../flux/extensions/kontext_extension.py | 2 +- invokeai/backend/flux2/ref_image_extension.py | 2 +- .../backend/model_manager/load/load_base.py | 18 +- invokeai/backend/util/vae_working_memory.py | 542 +++++++----- invokeai/backend/working_memory/__init__.py | 6 + .../backend/working_memory/coefficients.py | 191 +++++ invokeai/backend/working_memory/data/vae.json | 769 ++++++++++++++++++ invokeai/backend/working_memory/estimate.py | 52 ++ invokeai/backend/working_memory/platform.py | 45 + invokeai/backend/working_memory/spec.py | 99 +++ pyproject.toml | 2 + scripts/calibrate_flux2_working_memory.py | 37 +- scripts/calibrate_flux_vae_working_memory.py | 6 +- tests/app/invocations/conftest.py | 18 + .../invocations/test_flux2_working_memory.py | 235 ++---- tests/app/invocations/test_flux_vae_decode.py | 2 +- .../test_flux_vae_decode_tiling.py | 18 +- .../test_vae_working_memory_backend.py | 76 +- .../invocations/test_z_image_tiled_decode.py | 135 +-- .../test_z_image_working_memory.py | 12 +- .../model_manager/load/test_loaded_model.py | 64 +- .../util/test_rocm_sdpa_head_dim_guard.py | 30 +- tests/backend/working_memory/__init__.py | 0 .../test_shipped_coefficients.py | 312 +++++++ .../working_memory/test_spec_registry.py | 72 ++ 41 files changed, 2266 insertions(+), 534 deletions(-) create mode 100644 invokeai/backend/working_memory/__init__.py create mode 100644 invokeai/backend/working_memory/coefficients.py create mode 100644 invokeai/backend/working_memory/data/vae.json create mode 100644 invokeai/backend/working_memory/estimate.py create mode 100644 invokeai/backend/working_memory/platform.py create mode 100644 invokeai/backend/working_memory/spec.py create mode 100644 tests/backend/working_memory/__init__.py create mode 100644 tests/backend/working_memory/test_shipped_coefficients.py create mode 100644 tests/backend/working_memory/test_spec_registry.py diff --git a/docs/src/content/docs/contributing/new-model-integration.mdx b/docs/src/content/docs/contributing/new-model-integration.mdx index eda7821880a..04e4530b2d8 100644 --- a/docs/src/content/docs/contributing/new-model-integration.mdx +++ b/docs/src/content/docs/contributing/new-model-integration.mdx @@ -601,11 +601,11 @@ Packages are discovered automatically. A directory of node modules without an `_ ... ``` - Report progress with `context.util.sd_step_callback(state, BaseModelType.NewModel)`; it renders previews through the architecture's `LatentSpaceFacet`. The denoise node is also where these are checked: interrupts (`context.util.is_canceled()`), inpainting, img2img through `denoising_start`, and a working-memory estimate passed to `model_on_device(working_mem_bytes=...)`. + Report progress with `context.util.sd_step_callback(state, BaseModelType.NewModel)`; it renders previews through the architecture's `LatentSpaceFacet`. The denoise node is also where these are checked: interrupts (`context.util.is_canceled()`), inpainting, img2img through `denoising_start`, and a working-memory estimate passed to `model_on_device(working_mem_bytes=)`, or `working_memory=` for an estimator that declares an `EstimatorSpec`. 4. **VAE encode and decode** (`vae/new_model_image_to_latents.py`, `vae/new_model_latents_to_image.py`) - Apply the VAE's latent normalization, scaling and shift or per-channel mean and std, exactly as the reference pipeline does. Handle tiling and give a working-memory estimate. The estimate helpers live in `invokeai/backend/util/vae_working_memory.py`, with calibration scripts in `scripts/calibrate_*_working_memory.py`. Node types conventionally end in `_i2l` and `_l2i`. + Apply the VAE's latent normalization, scaling and shift or per-channel mean and std, exactly as the reference pipeline does. Handle tiling and give a working-memory estimate. The estimate helpers live in `invokeai/backend/util/vae_working_memory.py`, with calibration scripts in `scripts/calibrate_*_working_memory.py`. A new estimator declares an `EstimatorSpec`, returns a `WorkingMemoryEstimate` and takes its fitted coefficients from `invokeai/backend/working_memory/data/`, next to the measured points they were fitted on: the estimate must be at least each point and, where a fitted term decides it, at most 1.3 times it (or the point's `max_ratio`). The [model manager](/development/architecture/model-manager/) lists what such an estimator needs. `working_mem_bytes=` remains for a plain reservation. Node types conventionally end in `_i2l` and `_l2i`. 5. **LoRA loaders** (if LoRA is supported, see [Optional features](#12-optional-features)) diff --git a/docs/src/content/docs/development/Architecture/model-manager.mdx b/docs/src/content/docs/development/Architecture/model-manager.mdx index 846883395b4..2dfa5d93010 100644 --- a/docs/src/content/docs/development/Architecture/model-manager.mdx +++ b/docs/src/content/docs/development/Architecture/model-manager.mdx @@ -972,6 +972,26 @@ with loaded_model.model_on_device() as (state_dict, vae): Since not all models have state dicts, the `state_dict` return value can be None. +An operation that needs more device memory than `device_working_mem_gb` reserves it when it locks the model, so the cache offloads other models instead of running out of memory. Pass the operation's working-memory estimate: + +```py +estimate = estimate_vae_working_memory_flux2( + operation="decode", image_tensor=latents, vae=vae_info.model, device=vae_info.compute_device +) +with vae_info.model_on_device(working_memory=estimate) as (_, vae): + image = vae.decode(latents)[0] +``` + +An estimator that declares an `EstimatorSpec` returns a `WorkingMemoryEstimate` (`invokeai/backend/working_memory/estimate.py`): the reservation, the terms it adds up to, the features it was evaluated from, and the source of its coefficients, which `model_on_device()` logs at debug level. A plain byte count goes in as `working_mem_bytes=` instead; never both. The cache never reserves less than `device_working_mem_gb`. + +To write such an estimator (`invokeai/backend/util/vae_working_memory.py` has five): + +- Register an `EstimatorSpec` at module import, and list the module in `ESTIMATOR_MODULES` (`working_memory/spec.py`); the data test fails on a module that declares a spec without being listed. +- Split the estimator in two. The impure part reads the model and the device and records every runtime decision (tile, element size, score-matrix bytes) as a feature. The pure `terms(features, values)` turns features and coefficient values into terms, building fitted terms with `WorkingMemoryTerm.fitted(...)`. The public function returns `SPEC.estimate(features, device)`. +- Bump `version` whenever `terms` changes the formula. +- 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. + ### 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: diff --git a/invokeai/app/invocations/vae/cogview4_image_to_latents.py b/invokeai/app/invocations/vae/cogview4_image_to_latents.py index 09cde601d32..dd4a3862d98 100644 --- a/invokeai/app/invocations/vae/cogview4_image_to_latents.py +++ b/invokeai/app/invocations/vae/cogview4_image_to_latents.py @@ -43,7 +43,7 @@ def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tenso estimated_working_memory = estimate_vae_working_memory_cogview4( operation="encode", image_tensor=image_tensor, vae=vae_info.model, device=vae_info.compute_device ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) vae.disable_tiling() diff --git a/invokeai/app/invocations/vae/cogview4_latents_to_image.py b/invokeai/app/invocations/vae/cogview4_latents_to_image.py index 99fb2f21261..e8b655b0ce8 100644 --- a/invokeai/app/invocations/vae/cogview4_latents_to_image.py +++ b/invokeai/app/invocations/vae/cogview4_latents_to_image.py @@ -50,7 +50,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: ) with ( SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), - vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae), + vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae), ): context.util.signal_progress("Running VAE") assert isinstance(vae, (AutoencoderKL)) diff --git a/invokeai/app/invocations/vae/flux2_vae_decode.py b/invokeai/app/invocations/vae/flux2_vae_decode.py index f0852f1880f..cb5ed93c10a 100644 --- a/invokeai/app/invocations/vae/flux2_vae_decode.py +++ b/invokeai/app/invocations/vae/flux2_vae_decode.py @@ -57,7 +57,7 @@ def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Ima operation="decode", image_tensor=latents, vae=vae_info.model, device=vae_info.compute_device ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): vae_dtype = next(iter(vae.parameters())).dtype # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from # current param residency: partial loading may have temporarily offloaded all weights to RAM, which would diff --git a/invokeai/app/invocations/vae/flux2_vae_encode.py b/invokeai/app/invocations/vae/flux2_vae_encode.py index 926f3a7c043..4e024158f1c 100644 --- a/invokeai/app/invocations/vae/flux2_vae_encode.py +++ b/invokeai/app/invocations/vae/flux2_vae_encode.py @@ -60,7 +60,7 @@ def _vae_encode(self, vae_info: LoadedModel, image_tensor: torch.Tensor) -> torc device=device, ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): vae_dtype = next(iter(vae.parameters())).dtype image_tensor = image_tensor.to(device=device, dtype=vae_dtype) diff --git a/invokeai/app/invocations/vae/flux_vae_decode.py b/invokeai/app/invocations/vae/flux_vae_decode.py index 170162e8077..f50c2a3272c 100644 --- a/invokeai/app/invocations/vae/flux_vae_decode.py +++ b/invokeai/app/invocations/vae/flux_vae_decode.py @@ -80,7 +80,7 @@ def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents if ( not use_tiling and config.auto_tiled_decode - and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory) + and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory.reservation_bytes) ): use_tiling = True tile_size = self.tile_size @@ -92,7 +92,7 @@ def _vae_decode(self, context: InvocationContext, vae_info: LoadedModel, latents device=vae_info.compute_device, ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) vae_dtype = next(iter(vae.parameters())).dtype # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from diff --git a/invokeai/app/invocations/vae/flux_vae_encode.py b/invokeai/app/invocations/vae/flux_vae_encode.py index 5984fa737ed..a5a10900608 100644 --- a/invokeai/app/invocations/vae/flux_vae_encode.py +++ b/invokeai/app/invocations/vae/flux_vae_encode.py @@ -90,7 +90,7 @@ def vae_encode( # generator held outside would already have been drawn from by the attempt that failed, so # the retry would produce a different latent than a directly tiled run. generator = torch.Generator(device=vae_info.compute_device).manual_seed(0) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) vae_dtype = next(iter(vae.parameters())).dtype # The VAE's own device, as in the decode nodes: a cpu_only VAE (or another worker's GPU) diff --git a/invokeai/app/invocations/vae/ideogram4_latents_to_image.py b/invokeai/app/invocations/vae/ideogram4_latents_to_image.py index f2e72d6e94c..0016a13ccec 100644 --- a/invokeai/app/invocations/vae/ideogram4_latents_to_image.py +++ b/invokeai/app/invocations/vae/ideogram4_latents_to_image.py @@ -81,7 +81,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: operation="decode", image_tensor=z, vae=vae_info.model, device=device ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): context.util.signal_progress("Running VAE") vae_dtype = next(vae.parameters()).dtype TorchDevice.empty_cache() diff --git a/invokeai/app/invocations/vae/image_to_latents.py b/invokeai/app/invocations/vae/image_to_latents.py index 71d3951d79c..d599fadd862 100644 --- a/invokeai/app/invocations/vae/image_to_latents.py +++ b/invokeai/app/invocations/vae/image_to_latents.py @@ -83,7 +83,7 @@ def vae_encode( fp32=upcast, device=vae_info.compute_device, ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): assert isinstance(vae, (AutoencoderKL, AutoencoderTiny)), "VAE must be of type SD-1.5 or SDXL" orig_dtype = vae.dtype if upcast: diff --git a/invokeai/app/invocations/vae/latents_to_image.py b/invokeai/app/invocations/vae/latents_to_image.py index 4fbc87b8659..29c6a1b9db8 100644 --- a/invokeai/app/invocations/vae/latents_to_image.py +++ b/invokeai/app/invocations/vae/latents_to_image.py @@ -65,7 +65,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: ) with ( SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), - vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae), + vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae), ): context.util.signal_progress("Running VAE decoder") assert isinstance(vae, (AutoencoderKL, AutoencoderTiny)) diff --git a/invokeai/app/invocations/vae/sd3_image_to_latents.py b/invokeai/app/invocations/vae/sd3_image_to_latents.py index 5d558eeed52..fe1cb2faaa5 100644 --- a/invokeai/app/invocations/vae/sd3_image_to_latents.py +++ b/invokeai/app/invocations/vae/sd3_image_to_latents.py @@ -39,7 +39,7 @@ def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tenso estimated_working_memory = estimate_vae_working_memory_sd3( operation="encode", image_tensor=image_tensor, vae=vae_info.model, device=vae_info.compute_device ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) vae.disable_tiling() diff --git a/invokeai/app/invocations/vae/sd3_latents_to_image.py b/invokeai/app/invocations/vae/sd3_latents_to_image.py index 442ea0a83b6..87f67d73d8f 100644 --- a/invokeai/app/invocations/vae/sd3_latents_to_image.py +++ b/invokeai/app/invocations/vae/sd3_latents_to_image.py @@ -52,7 +52,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: ) with ( SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), - vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae), + vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae), ): context.util.signal_progress("Running VAE") assert isinstance(vae, (AutoencoderKL)) diff --git a/invokeai/app/invocations/vae/z_image_image_to_latents.py b/invokeai/app/invocations/vae/z_image_image_to_latents.py index de71e54ea0b..149116ae9a0 100644 --- a/invokeai/app/invocations/vae/z_image_image_to_latents.py +++ b/invokeai/app/invocations/vae/z_image_image_to_latents.py @@ -72,7 +72,7 @@ def vae_encode( device=vae_info.compute_device, ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): if not isinstance(vae, AutoencoderKL): raise TypeError( f"Expected AutoencoderKL, got {type(vae).__name__}. VAE model type changed unexpectedly after loading." diff --git a/invokeai/app/invocations/vae/z_image_latents_to_image.py b/invokeai/app/invocations/vae/z_image_latents_to_image.py index 2c5055402f3..af4c4bf07f6 100644 --- a/invokeai/app/invocations/vae/z_image_latents_to_image.py +++ b/invokeai/app/invocations/vae/z_image_latents_to_image.py @@ -74,7 +74,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: if ( not use_tiling and config.auto_tiled_decode - and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory) + and should_pretile_vae_decode(vae_info.compute_device, estimated_working_memory.reservation_bytes) ): use_tiling = True estimated_working_memory = estimate_vae_working_memory_flux( @@ -90,7 +90,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput: # FLUX VAE now honours the axes the UI has always offered for it. seamless_context = SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes) - with seamless_context, vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with seamless_context, vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): context.util.signal_progress("Running VAE") if not isinstance(vae, AutoencoderKL): raise TypeError( diff --git a/invokeai/backend/AGENTS.md b/invokeai/backend/AGENTS.md index 3b30a1c98c3..2f12dc1df2e 100644 --- a/invokeai/backend/AGENTS.md +++ b/invokeai/backend/AGENTS.md @@ -6,6 +6,7 @@ - Use existing device/model-management abstractions, `util/devices.py` included. Keep CPU/MPS/ROCm/XPU/multi-device support. No hardcoded CUDA, GPU zero, or dtype assumptions. - Respect session/thread device ownership. Restore temporary model patches, device state, hooks, allocations on success, cancellation, and exceptions. No request-specific mutable globals. - Define model/cache lifetimes and working-memory estimates. Bound queues, caches, downloads, batches, buffers. Consider peak memory and throughput. +- New or refitted working-memory coefficients go in `working_memory/data/` with the measured points that bound them; their estimators declare an `EstimatorSpec` and return a `WorkingMemoryEstimate`. - Keep model validation/loading safeguards, vendored exclusions, upstream provenance. No unrelated vendored rewrites or formatting. ## Performance diff --git a/invokeai/backend/flux/extensions/kontext_extension.py b/invokeai/backend/flux/extensions/kontext_extension.py index 7e801c54c04..1f31859909f 100644 --- a/invokeai/backend/flux/extensions/kontext_extension.py +++ b/invokeai/backend/flux/extensions/kontext_extension.py @@ -138,7 +138,7 @@ def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]: operation="encode", image_tensor=image_tensor, vae=vae_info.model ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): assert isinstance(vae, AutoencoderKL) vae_dtype = next(iter(vae.parameters())).dtype image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype) diff --git a/invokeai/backend/flux2/ref_image_extension.py b/invokeai/backend/flux2/ref_image_extension.py index 9184b15c1d2..c0e5d36b404 100644 --- a/invokeai/backend/flux2/ref_image_extension.py +++ b/invokeai/backend/flux2/ref_image_extension.py @@ -218,7 +218,7 @@ def _prepare_ref_images(self) -> tuple[torch.Tensor, torch.Tensor]: device=TorchDevice.choose_torch_device(), ) - with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae): vae_dtype = next(iter(vae.parameters())).dtype image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype) diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index 84a8d539516..fddf9a368dc 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -25,6 +25,8 @@ ModelCache, ) from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType +from invokeai.backend.util.logging import InvokeAILogger +from invokeai.backend.working_memory.estimate import WorkingMemoryEstimate class LoadedModelWithoutConfig: @@ -194,13 +196,27 @@ def __exit__(self, *args: Any, **kwargs: Any) -> None: @contextmanager def model_on_device( - self, working_mem_bytes: Optional[int] = None + self, working_mem_bytes: Optional[int] = None, *, working_memory: Optional[WorkingMemoryEstimate] = None ) -> Generator[Tuple[Optional[Dict[str, torch.Tensor]], AnyModel], None, None]: """Return a tuple consisting of the model's state dict (if it exists) and the locked model on execution device. :param working_mem_bytes: The amount of working memory to keep available on the compute device when loading the model. + :param working_memory: A working-memory estimate whose `reservation_bytes` is kept available instead; see + `invokeai.backend.working_memory`. Pass this or `working_mem_bytes`, not both. """ + if isinstance(working_mem_bytes, WorkingMemoryEstimate): + raise TypeError( + "model_on_device() got a WorkingMemoryEstimate as working_mem_bytes; pass it as working_memory=." + ) + if working_memory is not None: + if working_mem_bytes is not None: + raise ValueError("Pass either working_mem_bytes or working_memory to model_on_device(), not both.") + working_mem_bytes = working_memory.reservation_bytes + InvokeAILogger.get_logger(__name__).debug( + f"Working memory for {self._cache_record.key}: {working_memory.estimator} reserves " + f"{working_mem_bytes / 2**20:.0f} MB ({working_memory.source})" + ) self._lock_paced(working_mem_bytes) self._end_first_use_window() try: diff --git a/invokeai/backend/util/vae_working_memory.py b/invokeai/backend/util/vae_working_memory.py index 4c92894f1ee..6462dc1f41a 100644 --- a/invokeai/backend/util/vae_working_memory.py +++ b/invokeai/backend/util/vae_working_memory.py @@ -1,4 +1,6 @@ -from typing import Literal +from collections.abc import Mapping +from dataclasses import replace +from typing import Literal, cast import torch from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL @@ -13,6 +15,8 @@ from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger from invokeai.backend.util.vae_tiling_scope import resolve_tile_size +from invokeai.backend.working_memory.estimate import FeatureValue, WorkingMemoryEstimate, WorkingMemoryTerm +from invokeai.backend.working_memory.spec import EstimatorSpec, register # Share of the device's memory an untiled decode may claim before the decode nodes tile it up front (see # `should_pretile_vae_decode`). High enough that common resolutions on common cards stay untiled -- tiled output is @@ -110,6 +114,96 @@ def _vae_mid_block_score_matrix_bytes( ) +# The SD1/SDXL, SD3, CogView4, FLUX.1 and FLUX.2 VAEs run the same diffusers AutoencoderKL convolution stack, and pay +# the same for it: peak reserved bytes per output pixel per element byte, one fitted coefficient per operation +# (`k_decode`, `k_encode`) and convolution backend. The values ship in `backend/working_memory/data/vae.json`, with the +# points they were fitted on and the per-card table behind them; `scripts/calibrate_flux2_working_memory.py --only vae` +# reproduces that table on any build. +# +# MIOpen's convolution workspaces are simply larger than cuDNN's -- ~1.6x for a decode, ~2.4x for an encode. This is +# not the attention term: it shows up identically on the fused path, and the implied coefficients are flat across +# resolution on both backends, so the linear model itself holds. Nor is "encoding costs half of decoding" architectural: +# it is a cuDNN property (0.49 there, 0.78 on MIOpen), which is why each operation carries its own coefficient instead +# of a ratio. Shipping the MIOpen values everywhere would add ~60% to every cuDNN decode for nothing, so the +# coefficients follow the backend of the device the VAE runs on (`platform_for`). +# +# Caveat for AMD users: `MIOPEN_FIND_MODE=2` (FAST) selects convolution algorithms heuristically instead of by +# benchmark, and measured a uniform 1.28x more memory at every resolution when the MIOpen values were first fitted. +# Invoke has since made FAST its default (`run_app.py`); the RX 9060 XT measurement was taken under it and did not show +# the 1.28x. If you see decodes exceed their reservation on another card, raise `device_working_mem_gb`. +_AUTOENCODER_KL_COEFFICIENTS = frozenset({"k_decode", "k_encode"}) +_CONV_BACKEND_KEYS = ("conv_backend",) + + +def _peak_of_two_phases( + convolution_phase: tuple[WorkingMemoryTerm, ...], score_matrix_bytes: int +) -> tuple[WorkingMemoryTerm, ...]: + """The convolution phase and the mid-block score matrix, of which only the larger is reserved. + + max, not sum: the two peak in different phases of the same forward. The mid-block sits at the 8x-downsampled + bottleneck -- first in the decoder, last in the encoder -- so the full-resolution convolution feature maps that drive + the linear term are not live while the score matrix is, and peak *reserved* is a high-water mark, not a running + total. Measured on the FLUX.2 VAE (the points and their method are in `data/vae.json`): on gfx1100 and gfx1201 + forcing `math` moves the measured peak by nothing at all up to 1024px, and the totals stay flat-linear in area + either way. On CUDA the score matrix only pokes above the convolution peak at 1536px, and then by 2.6GB against the + 21.5GB this term prices standalone, because the attention phase reuses blocks the allocator is already holding. + Summing them reserved 11.1GB for a 1024px gfx1100 decode that measures 6.7GB; taking the max reserves 6.9GB. + + Where a max model is weakest is the crossover, where the two terms are near-equal and whatever overlap exists is + no longer hidden. There is exactly one measured point like that: a 768px encode with cuDNN's linear coefficient and + a materializing kernel measures 1.80GB against a 1.35GB max. It is not reachable as a shortfall, because the cache + floors every reservation at `device_working_mem_gb` (3GB by default, see `ModelCache._get_vram_available`) and the + whole crossover region sits under that floor. On the builds that really do materialize -- gfx1100 and gfx1201 -- + the MIOpen coefficient keeps the linear term above the score term across the measured range, so the crossover does + not arise there at all. + + On a tie the convolutions are the ones reserved. + """ + convolution_bytes: float = 0 + for term in convolution_phase: + convolution_bytes += term.bytes + scores_win = score_matrix_bytes > convolution_bytes + return ( + *(replace(term, included=not scores_win) for term in convolution_phase), + WorkingMemoryTerm("mid_block_scores", score_matrix_bytes, included=scores_win), + ) + + +def _sd_terms(features: Mapping[str, FeatureValue], values: Mapping[str, float]) -> tuple[WorkingMemoryTerm, ...]: + coefficient = f"k_{features['operation']}" + element_size = cast(int, features["element_size"]) + tile = cast(int | None, features["tile"]) + if tile is not None: + area_bytes = tile * tile * element_size + # We add 25% to the working memory estimate when tiling is enabled to account for factors like tile overlap + # and number of tiles. We could make this more precise in the future, but this should be good enough for + # most use cases. + convolutions = WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes * 1.25, values) + else: + area_bytes = cast(int, features["height"]) * cast(int, features["width"]) * element_size + convolutions = WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes, values) + + terms: tuple[WorkingMemoryTerm, ...] = (convolutions,) + # The tiny autoencoder has no mid-block attention, so it is priced without a score matrix. + if features["vae_class"] == "AutoencoderKL": + terms = _peak_of_two_phases(terms, cast(int, features["score_matrix_bytes"])) + if features["fp32"]: + # If we are running in FP32, then we should account for the likely increase in model size (~250MB). + terms += (WorkingMemoryTerm("fp32_weights", 250 * 2**20, True),) + return terms + + +VAE_SD_SPEC = register( + EstimatorSpec( + id="vae.sd", + version=1, + coefficients=_AUTOENCODER_KL_COEFFICIENTS, + platform_keys=_CONV_BACKEND_KEYS, + terms=_sd_terms, + ) +) + + def estimate_vae_working_memory_sd15_sdxl( operation: Literal["encode", "decode"], image_tensor: torch.Tensor, @@ -117,47 +211,75 @@ def estimate_vae_working_memory_sd15_sdxl( tile_size: int | None, fp32: bool, device: torch.device | None = None, -) -> int: - """Estimate the working memory required to encode or decode the given tensor.""" - # It was found experimentally that the peak working memory scales linearly with the number of pixels and the - # element size (precision). This estimate is accurate for both SD1 and SDXL. - element_size = 4 if fp32 else 2 - - # The same AutoencoderKL convolution stack as the FLUX.1 autoencoder, so it takes the same measured constants per - # convolution backend: MIOpen needs ~1.64x what cuDNN does (see `_FLUX_VAE_SCALING_CONSTANTS`). `device` is where - # the VAE will run; it defaults to the session device. - scaling_constant = _flux_vae_scaling_constant(operation, device or TorchDevice.choose_torch_device()) +) -> WorkingMemoryEstimate: + """Estimate the working memory required to encode or decode the given tensor. + It was found experimentally that the peak working memory scales linearly with the number of pixels and the element + size (precision). This estimate is accurate for both SD1 and SDXL. `device` is where the VAE will run; it picks the + convolution backend's coefficients and defaults to the session device. + """ + element_size = 4 if fp32 else 2 + device = device if device is not None else TorchDevice.choose_torch_device() latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 + height = latent_scale_factor_for_operation * image_tensor.shape[-2] + width = latent_scale_factor_for_operation * image_tensor.shape[-1] + + if tile_size == 0: + tile_size = vae.tile_sample_min_size + assert isinstance(tile_size, int) + + features: dict[str, FeatureValue] = { + "operation": operation, + "height": height, + "width": width, + "tile": tile_size, + "element_size": element_size, + "fp32": fp32, + "vae_class": "AutoencoderKL" if isinstance(vae, AutoencoderKL) else "AutoencoderTiny", + } + if isinstance(vae, AutoencoderKL): + score_dtype = torch.float32 if fp32 else next(vae.parameters()).dtype + score_h, score_w = (tile_size, tile_size) if tile_size is not None else (height, width) + features["score_matrix_bytes"] = _vae_mid_block_score_matrix_bytes(score_h, score_w, score_dtype, device=device) - if tile_size is not None: - if tile_size == 0: - tile_size = vae.tile_sample_min_size - assert isinstance(tile_size, int) - h = tile_size - w = tile_size - working_memory = h * w * element_size * scaling_constant + return VAE_SD_SPEC.estimate(features, device) - # We add 25% to the working memory estimate when tiling is enabled to account for factors like tile overlap - # and number of tiles. We could make this more precise in the future, but this should be good enough for - # most use cases. - working_memory = working_memory * 1.25 - else: - h = latent_scale_factor_for_operation * image_tensor.shape[-2] - w = latent_scale_factor_for_operation * image_tensor.shape[-1] - working_memory = h * w * element_size * scaling_constant - if isinstance(vae, AutoencoderKL): - # max, not sum: the mid-block attention and the full-resolution convolutions peak in - # different phases of the forward (see the FLUX.2 estimator for the measurements). - score_dtype = torch.float32 if fp32 else next(vae.parameters()).dtype - working_memory = max(working_memory, _vae_mid_block_score_matrix_bytes(h, w, score_dtype, device=device)) +def _single_pass_terms( + features: Mapping[str, FeatureValue], values: Mapping[str, float] +) -> tuple[WorkingMemoryTerm, ...]: + """An untiled AutoencoderKL pass: the convolutions over the whole image, or the mid-block score matrix.""" + coefficient = f"k_{features['operation']}" + area_bytes = cast(int, features["height"]) * cast(int, features["width"]) * cast(int, features["element_size"]) + convolutions = WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes, values) + return _peak_of_two_phases((convolutions,), cast(int, features["score_matrix_bytes"])) - if fp32: - # If we are running in FP32, then we should account for the likely increase in model size (~250MB). - working_memory += 250 * 2**20 - return int(working_memory) +def _single_pass_features( + operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKL, device: torch.device +) -> dict[str, FeatureValue]: + latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 + height = latent_scale_factor_for_operation * image_tensor.shape[-2] + width = latent_scale_factor_for_operation * image_tensor.shape[-1] + param = next(vae.parameters()) + return { + "operation": operation, + "height": height, + "width": width, + "element_size": param.element_size(), + "score_matrix_bytes": _vae_mid_block_score_matrix_bytes(height, width, param.dtype, device=device), + } + + +VAE_COGVIEW4_SPEC = register( + EstimatorSpec( + id="vae.cogview4", + version=1, + coefficients=_AUTOENCODER_KL_COEFFICIENTS, + platform_keys=_CONV_BACKEND_KEYS, + terms=_single_pass_terms, + ) +) def estimate_vae_working_memory_cogview4( @@ -165,25 +287,14 @@ def estimate_vae_working_memory_cogview4( image_tensor: torch.Tensor, vae: AutoencoderKL, device: torch.device | None = None, -) -> int: - """Estimate the working memory required by the invocation in bytes.""" - latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 - - h = latent_scale_factor_for_operation * image_tensor.shape[-2] - w = latent_scale_factor_for_operation * image_tensor.shape[-1] - element_size = next(vae.parameters()).element_size() - - # The same AutoencoderKL convolution stack as the FLUX.1 autoencoder, so it takes the same measured constants per - # convolution backend: MIOpen needs ~1.64x what cuDNN does (see `_FLUX_VAE_SCALING_CONSTANTS`). `device` is where - # the VAE will run; it defaults to the session device. - scaling_constant = _flux_vae_scaling_constant(operation, device or TorchDevice.choose_torch_device()) - working_memory = h * w * element_size * scaling_constant +) -> WorkingMemoryEstimate: + """Estimate the working memory required by the invocation in bytes. - # max, not sum: the mid-block attention and the full-resolution convolutions peak in different - # phases of the forward (see the FLUX.2 estimator for the measurements). - return int( - max(working_memory, _vae_mid_block_score_matrix_bytes(h, w, next(vae.parameters()).dtype, device=device)) - ) + `device` is where the VAE will run; it picks the convolution backend's coefficients and defaults to the session + device. + """ + device = device if device is not None else TorchDevice.choose_torch_device() + return VAE_COGVIEW4_SPEC.estimate(_single_pass_features(operation, image_tensor, vae, device), device) # What a tiled decode does *not* bound: the assembled image, several times over, per output pixel. @@ -217,13 +328,80 @@ def estimate_vae_working_memory_cogview4( _VAE_MOMENT_CHANNELS_PER_LATENT = 2 +def _flux_terms(features: Mapping[str, FeatureValue], values: Mapping[str, float]) -> tuple[WorkingMemoryTerm, ...]: + operation = features["operation"] + coefficient = f"k_{operation}" + out_h = cast(int, features["height"]) + out_w = cast(int, features["width"]) + element_size = cast(int, features["element_size"]) + tile = cast(int | None, features["tile"]) + + if tile is None: + area_bytes = out_h * out_w * element_size + convolution_phase: tuple[WorkingMemoryTerm, ...] = ( + WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes, values), + ) + return _peak_of_two_phases(convolution_phase, cast(int, features["score_matrix_bytes"])) + + # `calc_tiles_min_overlap` clamps the tile to the image *per axis*, so the largest tile the + # decode builds is this, not `tile` squared. Without the clamp an 8192px tile on a + # 1024x8200 image -- tiled, because one axis exceeds it -- prices 8192x8192. + area_bytes = min(tile, out_h) * min(tile, out_w) * element_size + # A 25% margin for tile overlap and the number of tiles, mirroring the SD1/SDXL estimator. + convolutions = WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes * 1.25, values) + if operation == "decode": + convolution_phase = ( + convolutions, + WorkingMemoryTerm( + "assembled_image", + out_h * out_w * _IMAGE_CHANNELS * (_FLUX_VAE_TILED_IMAGE_COPIES * element_size + 1), + True, + ), + ) + else: + # A tiled encode leaves two things un-bounded: the full-resolution image it slices tiles + # from, and the moments it assembles. + # + # The moments are *not* one copy. `AutoencoderKL._tiled_encode` appends every encoded + # tile to `rows`, then builds `result_rows` from them -- blending against `rows[i-1][j]`, + # so `rows` stays live throughout -- and finally concatenates `enc` on top of both. The + # tiles overlap, so their total area is `1 / (1 - overlap_factor)**2` times the output; + # `result_rows` and `enc` are one output each. Hence the factor below, which is 3.78 at + # the 0.25 every `AutoencoderKL` ships. + live_moment_copies = 1.0 / (1.0 - cast(float, features["tile_overlap_factor"])) ** 2 + 2.0 + moments = ( + (out_h // LATENT_SCALE_FACTOR) + * (out_w // LATENT_SCALE_FACTOR) + * _VAE_MOMENT_CHANNELS_PER_LATENT + * cast(int, features["latent_channels"]) + * element_size + ) + convolution_phase = ( + convolutions, + WorkingMemoryTerm("input_image", out_h * out_w * _IMAGE_CHANNELS * element_size, True), + WorkingMemoryTerm("assembled_moments", moments * live_moment_copies, True), + ) + return _peak_of_two_phases(convolution_phase, cast(int, features["score_matrix_bytes"])) + + +VAE_FLUX_SPEC = register( + EstimatorSpec( + id="vae.flux", + version=1, + coefficients=_AUTOENCODER_KL_COEFFICIENTS, + platform_keys=_CONV_BACKEND_KEYS, + terms=_flux_terms, + ) +) + + def estimate_vae_working_memory_flux( operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKL, tile_size: int | None = None, device: torch.device | None = None, -) -> int: +) -> WorkingMemoryEstimate: """Estimate the working memory required by the invocation in bytes. `tile_size` is in output pixels and defaults to None, i.e. a single-pass decode -- the existing @@ -237,16 +415,11 @@ def estimate_vae_working_memory_flux( for the decode side. `device` is where the VAE will run (a `cpu_only` VAE runs on the CPU); it picks the convolution - backend's constant. Defaults to the session device. + backend's coefficients. Defaults to the session device. """ - latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 - element_size = next(vae.parameters()).element_size() - - # The FLUX.1 autoencoder shares the FLUX.2 VAE's convolution stack and measures the same constants on both - # backends; see `_FLUX_VAE_SCALING_CONSTANTS`. + param = next(vae.parameters()) device = device if device is not None else TorchDevice.choose_torch_device() - scaling_constant = _flux_vae_scaling_constant(operation, device) out_h = latent_scale_factor_for_operation * image_tensor.shape[-2] out_w = latent_scale_factor_for_operation * image_tensor.shape[-1] @@ -261,51 +434,27 @@ def estimate_vae_working_memory_flux( # field carries no upper bound: a 2048px tile on a 1024px image would otherwise reserve ~23GB # for a decode that needs ~4.6GB, evicting the transformer from the cache for nothing. if tile is not None and (tile < out_h or tile < out_w): - # `calc_tiles_min_overlap` clamps the tile to the image *per axis*, so the largest tile the - # decode builds is this, not `tile` squared. Without the clamp an 8192px tile on a - # 1024x8200 image -- tiled, because one axis exceeds it -- prices 8192x8192. - tile_h = min(tile, out_h) - tile_w = min(tile, out_w) - # A 25% margin for tile overlap and the number of tiles, mirroring the SD1/SDXL estimator. - working_memory = tile_h * tile_w * element_size * scaling_constant * 1.25 - if operation == "decode": - working_memory += out_h * out_w * _IMAGE_CHANNELS * (_FLUX_VAE_TILED_IMAGE_COPIES * element_size + 1) - else: - # A tiled encode leaves two things un-bounded: the full-resolution image it slices tiles - # from, and the moments it assembles. - # - # The moments are *not* one copy. `AutoencoderKL._tiled_encode` appends every encoded - # tile to `rows`, then builds `result_rows` from them -- blending against `rows[i-1][j]`, - # so `rows` stays live throughout -- and finally concatenates `enc` on top of both. The - # tiles overlap, so their total area is `1 / (1 - overlap_factor)**2` times the output; - # `result_rows` and `enc` are one output each. Hence the factor below, which is 3.78 at - # the 0.25 every `AutoencoderKL` ships. - # - # Read from the VAE rather than hard-coded: an overlap factor of 0.5 would make the - # tiles alone four times the output, and a fixed 3.78 would then under-reserve. An - # estimate consumed as `free >= estimate` has to err upwards. - config = getattr(vae, "config", None) - latent_channels = getattr(config, "latent_channels", 16) - overlap_factor = getattr(vae, "tile_overlap_factor", 0.25) - live_moment_copies = 1.0 / (1.0 - overlap_factor) ** 2 + 2.0 - moments = ( - (out_h // LATENT_SCALE_FACTOR) - * (out_w // LATENT_SCALE_FACTOR) - * _VAE_MOMENT_CHANNELS_PER_LATENT - * latent_channels - * element_size - ) - working_memory += out_h * out_w * _IMAGE_CHANNELS * element_size - working_memory += moments * live_moment_copies - score_h, score_w = tile_h, tile_w + score_h, score_w = min(tile, out_h), min(tile, out_w) else: - working_memory = out_h * out_w * element_size * scaling_constant + tile = None score_h, score_w = out_h, out_w - # max, not sum: the mid-block attention and the full-resolution convolutions peak in different - # phases of the forward (see the FLUX.2 estimator for the measurements). - dtype = next(vae.parameters()).dtype - return int(max(working_memory, _vae_mid_block_score_matrix_bytes(score_h, score_w, dtype, device=device))) + features: dict[str, FeatureValue] = { + "operation": operation, + "height": out_h, + "width": out_w, + "element_size": param.element_size(), + "tile": tile, + "score_matrix_bytes": _vae_mid_block_score_matrix_bytes(score_h, score_w, param.dtype, device=device), + } + if tile is not None and operation != "decode": + # Read from the VAE rather than hard-coded: an overlap factor of 0.5 would make the tiles + # alone four times the output, and a fixed 3.78 would then under-reserve. An estimate + # consumed as `free >= estimate` has to err upwards. + features["latent_channels"] = getattr(getattr(vae, "config", None), "latent_channels", 16) + features["tile_overlap_factor"] = getattr(vae, "tile_overlap_factor", 0.25) + + return VAE_FLUX_SPEC.estimate(features, device) # The FLUX.2 VAE runs one attention block at the bottom of the encoder and one at the top of the @@ -317,51 +466,32 @@ def estimate_vae_working_memory_flux( _FLUX2_VAE_MID_BLOCK_HEAD_DIM = 512 _FLUX2_VAE_SPATIAL_COMPRESSION = 8 -# Peak reserved bytes per output pixel per element byte, per conv backend and operation, for the FLUX.2 VAE and the -# FLUX.1 autoencoder alike. Fitted with `scripts/calibrate_flux2_working_memory.py --only vae`, which reproduces -# this table on any build: -# -# decode encode encode/decode -# cuDNN RTX 4090, torch 2.7.1 2185 1072 0.49 -# MIOpen RX 9070 XT, torch 2.10 3453 2688 0.78 -# MIOpen PRO W7900, torch 2.10 3525 2688 0.76 -# MIOpen RX 9060 XT, torch 2.12, Windows, FLUX.1 VAE 3451 2688 0.78 -# -# Two AMD generations (RDNA3 and RDNA4), the same numbers: the encode column agrees to the byte and -# the decode column to 2%. So this is MIOpen, not a per-card quirk, and the column below is fitted -# to the larger with ~2% headroom. The FLUX.1 row (bf16, 768-1536px, Invoke's ROCm SDPA guard installed, -# `MIOPEN_FIND_MODE=FAST`) lands on the same numbers: the two VAEs share their convolution stack. The -# cuDNN column is the FLUX.1 constant this estimator shipped with before. -# -# MIOpen's convolution workspaces are simply larger than cuDNN's. This is not the attention term -- -# it shows up identically on the fused path, and the implied constants are flat across resolution on -# both backends, so the linear model itself holds. And the "encoding costs half of decoding" ratio -# the other estimators in this module use turns out to be a cuDNN property rather than an -# architectural one, which is why the two operations carry their own numbers here instead of a ratio. -# -# Shipping the MIOpen numbers everywhere would add ~60% to every cuDNN decode for nothing, so the -# constant follows the backend. -# -# Caveat for AMD users: `MIOPEN_FIND_MODE=2` (FAST) selects convolution algorithms heuristically instead of -# by benchmark, and measured a uniform 1.28x more memory at every resolution when these constants were -# first fitted. Invoke has since made FAST its default (`run_app.py`); the RX 9060 XT row above was taken -# under it and did not show the 1.28x. If you see decodes exceed their reservation on another card, raise -# `device_working_mem_gb`. -_FLUX_VAE_SCALING_CONSTANTS: dict[str, dict[str, int]] = { - "cudnn": {"decode": 2200, "encode": 1100}, - "miopen": {"decode": 3600, "encode": 2750}, -} - -def _flux_vae_scaling_constant(operation: Literal["encode", "decode"], device: torch.device) -> int: - """Pick the pixel-area constant for the convolution backend this device will actually use. - - A HIP build reports ``device.type == "cuda"``, so the torch build -- not the device string -- is - what separates MIOpen from cuDNN. MPS and CPU are unmeasured and take the cuDNN column; on MPS - that pairs with a score-matrix term the probe always charges there, so the total is not thin. - """ - is_rocm = device.type == "cuda" and torch.version.hip is not None - return _FLUX_VAE_SCALING_CONSTANTS["miopen" if is_rocm else "cudnn"][operation] +def _flux2_terms(features: Mapping[str, FeatureValue], values: Mapping[str, float]) -> tuple[WorkingMemoryTerm, ...]: + coefficient = f"k_{features['operation']}" + element_size = cast(int, features["element_size"]) + batch_size = cast(int, features["batch_size"]) + tile = cast(int | None, features["tile"]) + if tile is not None: + area_bytes = tile * tile * element_size + # Add 25% for tile overlap and the blending buffers, mirroring the SD1/SDXL estimate. + convolutions = WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes * 1.25 * batch_size, values) + else: + area_bytes = cast(int, features["height"]) * cast(int, features["width"]) * element_size + convolutions = WorkingMemoryTerm.fitted("convolutions", coefficient, area_bytes * batch_size, values) + # Both terms are per sample; the larger is reserved (see `_peak_of_two_phases` for why, and the measurements). + return _peak_of_two_phases((convolutions,), cast(int, features["score_matrix_bytes"])) + + +VAE_FLUX2_SPEC = register( + EstimatorSpec( + id="vae.flux2", + version=1, + coefficients=_AUTOENCODER_KL_COEFFICIENTS, + platform_keys=_CONV_BACKEND_KEYS, + terms=_flux2_terms, + ) +) def estimate_vae_working_memory_flux2( @@ -370,15 +500,15 @@ def estimate_vae_working_memory_flux2( vae: AutoencoderKLFlux2, tile_size: int | None = None, device: torch.device | None = None, -) -> int: +) -> WorkingMemoryEstimate: """Estimate the working memory required to encode or decode with the FLUX.2 (32-channel) VAE. Peak memory scales linearly with pixel area and element size, as it does for the FLUX.1 VAE, and - the implied constant is flat across 512-1536px on every backend measured. What is *not* constant - is the constant itself: MIOpen's convolution workspaces cost ~1.6x cuDNN's for a decode and ~2.4x - for an encode, so it is looked up per backend -- see ``_FLUX_VAE_SCALING_CONSTANTS`` for the - fitted table and the caveats. Peak *reserved* memory is what is measured throughout, the - conservative quantity that includes allocator overhead. + the implied coefficient is flat across 512-1536px on every backend measured. What is *not* constant + is the coefficient itself: MIOpen's convolution workspaces cost ~1.6x cuDNN's for a decode and ~2.4x + for an encode, so it is shipped per backend -- see `_AUTOENCODER_KL_COEFFICIENTS` for the caveats + and `backend/working_memory/data/vae.json` for the fitted values and their points. Peak *reserved* + memory is what is measured throughout, the conservative quantity that includes allocator overhead. For reference, decoding 1024x1024 peaks at ~4.3GB on cuDNN and ~6.6GB on MIOpen, and 1536x1536 at ~9.6GB -- far above the default ``device_working_mem_gb``, which is why this estimate must be @@ -394,10 +524,10 @@ def estimate_vae_working_memory_flux2( for everything else, so the estimate is right on all of them. The two terms are independent: a build can have a fused kernel and still need the larger - convolution constant (gfx1201 did, before the ROCm rule above). They also do not add -- see the - ``max`` at the end of this function for why, and for the measurements behind it. (Unlike the transformer, this attention does not go through - diffusers' attention dispatcher -- ``AttnProcessor2_0`` calls ``F.scaled_dot_product_attention`` - itself -- so torch's own answer is the whole answer here.) + convolution coefficient (gfx1201 did, before the ROCm rule above). They also do not add -- see + `_peak_of_two_phases` for why, and for the measurements behind it. (Unlike the transformer, this + attention does not go through diffusers' attention dispatcher -- ``AttnProcessor2_0`` calls + ``F.scaled_dot_product_attention`` itself -- so torch's own answer is the whole answer here.) When tiling is enabled the peak is bounded by a single tile instead of the full image (measured ~0.55GB flat at a 512px tile, from 1024px up to the 2024px reference-image cap), and the score @@ -411,53 +541,34 @@ def estimate_vae_working_memory_flux2( call sites all pass batch 1 today; the shared estimator does not assume it. """ param = next(vae.parameters()) - element_size = param.element_size() - device = device if device is not None else TorchDevice.choose_torch_device() - scaling_constant = _flux_vae_scaling_constant(operation, device) batch_size = image_tensor.shape[0] if image_tensor.dim() >= 4 else 1 + latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 + out_h = latent_scale_factor_for_operation * image_tensor.shape[-2] + out_w = latent_scale_factor_for_operation * image_tensor.shape[-1] if tile_size is not None: - # Add 25% for tile overlap and the blending buffers, mirroring the SD1/SDXL estimate. - working_memory = tile_size * tile_size * element_size * scaling_constant * 1.25 mid_block_seq_len = (tile_size // _FLUX2_VAE_SPATIAL_COMPRESSION) ** 2 else: - latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 - out_h = latent_scale_factor_for_operation * image_tensor.shape[-2] - out_w = latent_scale_factor_for_operation * image_tensor.shape[-1] - working_memory = out_h * out_w * element_size * scaling_constant mid_block_seq_len = (out_h // _FLUX2_VAE_SPATIAL_COMPRESSION) * (out_w // _FLUX2_VAE_SPATIAL_COMPRESSION) - working_memory *= batch_size - score_matrix_bytes = sdpa_score_matrix_bytes( - device=device, - dtype=param.dtype, - # The score matrix is (batch, heads, S, S); one head per sample prices the whole batch. - num_heads=_FLUX2_VAE_MID_BLOCK_HEADS * batch_size, - head_dim=_FLUX2_VAE_MID_BLOCK_HEAD_DIM, - seq_len=mid_block_seq_len, - ) - - # max, not sum: the two terms peak in different phases of the same forward. The mid-block sits at - # the 8x-downsampled bottleneck -- first in the decoder, last in the encoder -- so the full- - # resolution convolution feature maps that drive the linear term are not live while the score - # matrix is, and peak *reserved* is a high-water mark, not a running total. Measured (see the - # constants table above for the method): on gfx1100 and gfx1201 forcing `math` moves the measured - # peak by nothing at all up to 1024px, and the totals stay flat-linear in area either way. On - # CUDA the score matrix only pokes above the convolution peak at 1536px, and then by 2.6GB - # against the 21.5GB this term prices standalone, because the attention phase reuses blocks the - # allocator is already holding. Summing them reserved 11.1GB for a 1024px gfx1100 decode that - # measures 6.7GB; taking the max reserves 6.9GB. - # - # Where a max model is weakest is the crossover, where the two terms are near-equal and whatever - # overlap exists is no longer hidden. There is exactly one measured point like that: a 768px - # encode with cuDNN's linear constant and a materializing kernel measures 1.80GB against a - # 1.35GB max. It is not reachable as a shortfall, because the cache floors every reservation at - # `device_working_mem_gb` (3GB by default, see `ModelCache._get_vram_available`) and the whole - # crossover region sits under that floor. On the builds that really do materialize -- gfx1100 and - # gfx1201 -- the MIOpen constant keeps the linear term above the score term across the measured - # range, so the crossover does not arise there at all. - return int(max(working_memory, score_matrix_bytes)) + features: dict[str, FeatureValue] = { + "operation": operation, + "height": out_h, + "width": out_w, + "tile": tile_size, + "batch_size": batch_size, + "element_size": param.element_size(), + "score_matrix_bytes": sdpa_score_matrix_bytes( + device=device, + dtype=param.dtype, + # The score matrix is (batch, heads, S, S); one head per sample prices the whole batch. + num_heads=_FLUX2_VAE_MID_BLOCK_HEADS * batch_size, + head_dim=_FLUX2_VAE_MID_BLOCK_HEAD_DIM, + seq_len=mid_block_seq_len, + ), + } + return VAE_FLUX2_SPEC.estimate(features, device) def estimate_vae_working_memory_anima( @@ -954,30 +1065,27 @@ def estimate_vae_working_memory_qwen_image( return int(working_memory) +VAE_SD3_SPEC = register( + EstimatorSpec( + id="vae.sd3", + version=1, + coefficients=_AUTOENCODER_KL_COEFFICIENTS, + platform_keys=_CONV_BACKEND_KEYS, + terms=_single_pass_terms, + ) +) + + def estimate_vae_working_memory_sd3( operation: Literal["encode", "decode"], image_tensor: torch.Tensor, vae: AutoencoderKL, device: torch.device | None = None, -) -> int: - """Estimate the working memory required by the invocation in bytes.""" - # Encode operations use approximately 50% of the memory required for decode operations - - latent_scale_factor_for_operation = LATENT_SCALE_FACTOR if operation == "decode" else 1 - - h = latent_scale_factor_for_operation * image_tensor.shape[-2] - w = latent_scale_factor_for_operation * image_tensor.shape[-1] - element_size = next(vae.parameters()).element_size() - - # The same AutoencoderKL convolution stack as the FLUX.1 autoencoder, so it takes the same measured constants per - # convolution backend: MIOpen needs ~1.64x what cuDNN does (see `_FLUX_VAE_SCALING_CONSTANTS`). `device` is where - # the VAE will run; it defaults to the session device. - scaling_constant = _flux_vae_scaling_constant(operation, device or TorchDevice.choose_torch_device()) - - working_memory = h * w * element_size * scaling_constant +) -> WorkingMemoryEstimate: + """Estimate the working memory required by the invocation in bytes. - # max, not sum: the mid-block attention and the full-resolution convolutions peak in different - # phases of the forward (see the FLUX.2 estimator for the measurements). - return int( - max(working_memory, _vae_mid_block_score_matrix_bytes(h, w, next(vae.parameters()).dtype, device=device)) - ) + `device` is where the VAE will run; it picks the convolution backend's coefficients and defaults to the session + device. + """ + device = device if device is not None else TorchDevice.choose_torch_device() + return VAE_SD3_SPEC.estimate(_single_pass_features(operation, image_tensor, vae, device), device) diff --git a/invokeai/backend/working_memory/__init__.py b/invokeai/backend/working_memory/__init__.py new file mode 100644 index 00000000000..dd71f347f51 --- /dev/null +++ b/invokeai/backend/working_memory/__init__.py @@ -0,0 +1,6 @@ +"""Working-memory estimates as data: what an operation needs on the compute device, and how that was priced. + +An estimator turns an operation's runtime facts (`features`) and a set of fitted coefficients into a +`WorkingMemoryEstimate` that `LoadedModel.model_on_device(working_memory=...)` reserves. The coefficients ship as data +under `data/`, next to the measured points that bound them; `tests/backend/working_memory/` enforces both. +""" diff --git a/invokeai/backend/working_memory/coefficients.py b/invokeai/backend/working_memory/coefficients.py new file mode 100644 index 00000000000..d8e07b13854 --- /dev/null +++ b/invokeai/backend/working_memory/coefficients.py @@ -0,0 +1,191 @@ +"""The coefficients estimators are evaluated with, and the measured points that bound them. + +Both ship as JSON under `data/`, read through `importlib.resources` so a wheel install finds them. A coefficient entry +belongs to one estimator version on one measured platform value, and either carries its values or borrows another +entry's (`values_from`); a point records the features of an operation that was measured and its measured peak. +`tests/backend/working_memory/test_shipped_coefficients.py` resolves every entry against its `EstimatorSpec` and holds +every point to the estimate the shipped coefficients give it. +""" + +import functools +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from importlib.resources import files +from importlib.resources.abc import Traversable +from types import MappingProxyType +from typing import TYPE_CHECKING, Literal, Self + +from pydantic import BaseModel, ConfigDict, ValidationError, model_validator + +from invokeai.backend.working_memory.estimate import FeatureValue +from invokeai.backend.working_memory.platform import UNMEASURED_PRICED_AS, Platform + +if TYPE_CHECKING: + from invokeai.backend.working_memory.spec import EstimatorSpec + + +class _DataModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class ValuesFrom(_DataModel): + """The entry whose values another entry takes: same platform, this estimator id and version.""" + + estimator: str + version: int + + +class CoefficientEntry(_DataModel): + """The coefficient values of one estimator version on one platform: its own, or another entry's.""" + + estimator: str + version: int + platform: dict[str, str] + values: dict[str, float] | None = None + values_from: ValuesFrom | None = None + """Borrow the values of another estimator's entry on the same platform, so that its measured points bound this + estimator too. One hop: the entry borrowed from carries its own values.""" + provenance: str + """Card, torch build, method and script the values come from, and which points they were fitted on.""" + + @model_validator(mode="after") + def _values_or_values_from(self) -> Self: + if (self.values is None) == (self.values_from is None): + raise ValueError("a coefficient entry carries exactly one of `values` and `values_from`") + return self + + +class MeasuredPoint(_DataModel): + """One measured operation. The estimate for its features must be at least `measured_bytes`, and where a fitted + term decides it, at most 1.3 times that (or `max_ratio`).""" + + estimator: str + version: int + platform: dict[str, str] + features: dict[str, FeatureValue] + measured_bytes: int + method: Literal["reserved_peak"] + """``reserved_peak``: the caching allocator's peak reserved bytes over the operation, one fresh process per + point. Reserved rather than allocated, because reserved is what the model cache's reservation has to cover.""" + device: str + provenance: str + max_ratio: float | None = None + """A tighter upper bound than the default 1.3, as a multiple of `measured_bytes`, for a point the estimate has to + stay closer to. It applies whichever term decides the estimate.""" + + +class WorkingMemoryData(_DataModel): + """One file under `data/`.""" + + schema_version: Literal[1] + coefficients: list[CoefficientEntry] + points: list[MeasuredPoint] + + +@dataclass(frozen=True) +class ResolvedCoefficients: + """The coefficient values an estimate is evaluated with, and where they came from.""" + + values: Mapping[str, float] + source: str + + +def _parse(entry: Traversable) -> WorkingMemoryData: + try: + return WorkingMemoryData.model_validate_json(entry.read_text(encoding="utf-8")) + except ValidationError as e: + raise ValueError(f"Malformed working-memory data file {entry}: {e}") from e + + +def load_data(directory: Traversable) -> Mapping[str, WorkingMemoryData]: + """Every JSON data file in `directory`, parsed, by file name. + + A directory with none raises, naming it: every estimator resolves its coefficients from there, so an install that + lost the files would otherwise fail later as a missing entry for whichever estimator ran first. + """ + entries = ( + sorted((entry for entry in directory.iterdir() if entry.name.endswith(".json")), key=lambda entry: entry.name) + if directory.is_dir() + else [] + ) + if not entries: + raise FileNotFoundError(f"No working-memory data files in {directory}; every EstimatorSpec reads them there.") + return MappingProxyType({entry.name: _parse(entry) for entry in entries}) + + +@functools.cache +def shipped_data() -> Mapping[str, WorkingMemoryData]: + """The data files shipped under `data/`. Read once per process.""" + return load_data(files("invokeai.backend.working_memory") / "data") + + +def _where(keys: Mapping[str, str]) -> str: + return ",".join(f"{key}={value}" for key, value in keys.items()) + + +def _entry( + entries: list[CoefficientEntry], estimator: str, version: int, keys: Mapping[str, str], where: str +) -> CoefficientEntry: + candidates = [entry for entry in entries if entry.estimator == estimator and entry.platform == keys] + matches = [entry for entry in candidates if entry.version == version] + if len(matches) > 1: + raise LookupError( + f"Working-memory coefficients for estimator '{estimator}' version {version} on {where} are declared " + f"{len(matches)} times." + ) + if not matches: + found = f" (found only version {sorted({entry.version for entry in candidates})})" if candidates else "" + raise LookupError( + f"No working-memory coefficients for estimator '{estimator}' version {version} on {where}{found}." + ) + return matches[0] + + +def resolve_coefficients( + data: Iterable[WorkingMemoryData], spec: "EstimatorSpec", platform: Platform, origin: str +) -> ResolvedCoefficients: + """The values `spec` takes on `platform`, and a `source` that says how they were found. + + A platform value nothing is measured for is priced as the measured value `UNMEASURED_PRICED_AS` names. The entry + for the spec's id, version and those platform values either carries the values or borrows them, one hop, from + another estimator's entry on the same platform. No other fallback: an estimate priced with coefficients fitted for + something else is the shortfall this data exists to prevent, so a missing, ambiguous, chained or mismatched entry + raises, naming what is wrong. `origin` prefixes the `source`. + """ + requested = {key: getattr(platform, key) for key in spec.platform_keys} + keys = {key: UNMEASURED_PRICED_AS.get(key, {}).get(value, value) for key, value in requested.items()} + where = _where(keys) if keys == requested else f"{_where(keys)} (priced for {_where(requested)})" + entries = [entry for file in data for entry in file.coefficients] + + entry = _entry(entries, spec.id, spec.version, keys, where) + values, lender = entry.values, None + if entry.values_from is not None: + lender = entry.values_from.estimator + lent = _entry(entries, lender, entry.values_from.version, keys, where) + if lent.values is None: + raise LookupError( + f"Working-memory coefficients for estimator '{spec.id}' on {where} borrow from '{lender}', whose " + "entry borrows in turn; borrow from the entry that carries the values." + ) + values = lent.values + assert values is not None, "a validated entry carries values or values_from" + if set(values) != spec.coefficients: + raise LookupError( + f"Working-memory coefficients for estimator '{spec.id}' on {where} provide {sorted(values)}" + f"{f' (from {lender})' if lender else ''}; the estimator reads {sorted(spec.coefficients)}." + ) + + priced = ",".join(value for key, value in keys.items() if value != requested[key]) + if priced and lender: + note = f" ({priced} values of {lender})" + elif priced: + note = f" ({priced} values)" + else: + note = f" (values of {lender})" if lender else "" + return ResolvedCoefficients(values=MappingProxyType(dict(values)), source=f"{origin}:{_where(requested)}{note}") + + +@functools.cache +def shipped_coefficients(spec: "EstimatorSpec", platform: Platform) -> ResolvedCoefficients: + """The shipped values `spec` takes on `platform`. Resolved once per spec and platform.""" + return resolve_coefficients(shipped_data().values(), spec, platform, origin="shipped") diff --git a/invokeai/backend/working_memory/data/vae.json b/invokeai/backend/working_memory/data/vae.json new file mode 100644 index 00000000000..bae9081384f --- /dev/null +++ b/invokeai/backend/working_memory/data/vae.json @@ -0,0 +1,769 @@ +{ + "schema_version": 1, + "coefficients": [ + { + "estimator": "vae.sd", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the SD1/SDXL VAE runs the same diffusers AutoencoderKL convolution stack. Not measured on this VAE." + }, + { + "estimator": "vae.sd", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the SD1/SDXL VAE runs the same diffusers AutoencoderKL convolution stack. Not measured on this VAE." + }, + { + "estimator": "vae.sd3", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the SD3 VAE runs the same diffusers AutoencoderKL convolution stack. Not measured on this VAE." + }, + { + "estimator": "vae.sd3", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the SD3 VAE runs the same diffusers AutoencoderKL convolution stack. Not measured on this VAE." + }, + { + "estimator": "vae.cogview4", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the CogView4 VAE runs the same diffusers AutoencoderKL convolution stack. Not measured on this VAE." + }, + { + "estimator": "vae.cogview4", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the CogView4 VAE runs the same diffusers AutoencoderKL convolution stack. Not measured on this VAE." + }, + { + "estimator": "vae.flux", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the FLUX.1 VAE runs the same diffusers AutoencoderKL convolution stack. The FLUX.1 point below falls inside them." + }, + { + "estimator": "vae.flux", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "values_from": { + "estimator": "vae.flux2", + "version": 1 + }, + "provenance": "Borrows the vae.flux2 values, so its points bound this estimator too: the FLUX.1 VAE runs the same diffusers AutoencoderKL convolution stack. Confirmed on the FLUX.1 VAE: RX 9060 XT (gfx1200), torch 2.12, Windows, bf16, 768-1536px, Invoke's ROCm SDPA guard installed, MIOPEN_FIND_MODE=FAST: implied 3451 decode / 2688 encode." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "values": { + "k_decode": 2200, + "k_encode": 1100 + }, + "provenance": "Peak reserved bytes per output pixel per element byte of the diffusers AutoencoderKL convolution stack. Fitted on the FLUX.2 VAE (the vae.flux2 cudnn points) with scripts/calibrate_flux2_working_memory.py --only vae: RTX 4090, torch 2.7.1, bf16, untiled, fused SDPA; implied 2180/2185/2165 decode and 1072/1063/1061 encode at 512/768/1024px, flat across resolution. Rounded up to 2200 / 1100, the pair the FLUX.1 estimator already shipped." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "values": { + "k_decode": 3600, + "k_encode": 2750 + }, + "provenance": "Peak reserved bytes per output pixel per element byte of the diffusers AutoencoderKL convolution stack on MIOpen. Fitted on the FLUX.2 VAE (the vae.flux2 miopen points) with scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled: RX 9070 XT (gfx1201) implied 3453/3369/3368 decode, PRO W7900 (gfx1100) 3525 at 512px; encode 2687/2688/2688 on both cards, to the byte. Two AMD generations agreeing that closely make this MIOpen rather than a per-card quirk, so the values are the larger card's plus ~2% headroom: 3600 / 2750. MIOpen's convolution workspaces are larger than cuDNN's (~1.6x decode, ~2.4x encode), identically with fused and math attention, and its encode/decode ratio is 0.78 against cuDNN's 0.49." + } + ], + "points": [ + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 1142461301, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.7.1, bf16, untiled, fused SDPA (score matrix not materialized). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 2576980378, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.7.1, bf16, untiled, fused SDPA (score matrix not materialized). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 4540854174, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.7.1, bf16, untiled, fused SDPA (score matrix not materialized). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 561566974, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.7.1, bf16, untiled, fused SDPA (score matrix not materialized). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 1254130450, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.7.1, bf16, untiled, fused SDPA (score matrix not materialized). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 2224793059, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.7.1, bf16, untiled, fused SDPA (score matrix not materialized). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1280, + "width": 1280, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 7075958620, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, bf16, untiled, fused SDPA: the fused column of the max-versus-sum measurement (#9519). Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 1138753536, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "First FLUX.2 VAE calibration (#9519): CUDA, bf16, untiled, fused SDPA. Recorded in MiB. The decode rows are a separate run from the 512-1024px GiB rows; the 1536px decode is also the fused 9.353 GiB of the max-versus-sum measurement." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 2531262464, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "First FLUX.2 VAE calibration (#9519): CUDA, bf16, untiled, fused SDPA. Recorded in MiB. The decode rows are a separate run from the 512-1024px GiB rows; the 1536px decode is also the fused 9.353 GiB of the max-versus-sum measurement." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 4466933760, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "First FLUX.2 VAE calibration (#9519): CUDA, bf16, untiled, fused SDPA. Recorded in MiB. The decode rows are a separate run from the 512-1024px GiB rows; the 1536px decode is also the fused 9.353 GiB of the max-versus-sum measurement." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1328, + "width": 1328, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 7493124096, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "First FLUX.2 VAE calibration (#9519): CUDA, bf16, untiled, fused SDPA. Recorded in MiB. The decode rows are a separate run from the 512-1024px GiB rows; the 1536px decode is also the fused 9.353 GiB of the max-versus-sum measurement." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1536, + "width": 1536, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 10043260928, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "First FLUX.2 VAE calibration (#9519): CUDA, bf16, untiled, fused SDPA. Recorded in MiB. The decode rows are a separate run from the 512-1024px GiB rows; the 1536px decode is also the fused 9.353 GiB of the max-versus-sum measurement." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 1328, + "width": 1328, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 3168796672, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "First FLUX.2 VAE calibration (#9519): CUDA, bf16, untiled, fused SDPA. Recorded in MiB. The decode rows are a separate run from the 512-1024px GiB rows; the 1536px decode is also the fused 9.353 GiB of the max-versus-sum measurement." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 285212672 + }, + "measured_bytes": 561566974, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae with SDPBackend.MATH forced on CUDA -- the regime MPS lands in; bf16, untiled; score_matrix_bytes is the full math-kernel score matrix. Forcing math measures below the fused path until 1536px. The 768px encode sits at the crossover, where the larger term alone is short and the device_working_mem_gb floor covers it; reproducible to three decimals across runs. Recorded in GiB." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 1443889152 + }, + "measured_bytes": 1933809025, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae with SDPBackend.MATH forced on CUDA -- the regime MPS lands in; bf16, untiled; score_matrix_bytes is the full math-kernel score matrix. Forcing math measures below the fused path until 1536px. The 768px encode sits at the crossover, where the larger term alone is short and the device_working_mem_gb floor covers it; reproducible to three decimals across runs. Recorded in GiB." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 4563402752 + }, + "measured_bytes": 4372276707, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae with SDPBackend.MATH forced on CUDA -- the regime MPS lands in; bf16, untiled; score_matrix_bytes is the full math-kernel score matrix. Forcing math measures below the fused path until 1536px. The 768px encode sits at the crossover, where the larger term alone is short and the device_working_mem_gb floor covers it; reproducible to three decimals across runs. Recorded in GiB." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 4563402752 + }, + "measured_bytes": 3787087413, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae with SDPBackend.MATH forced on CUDA -- the regime MPS lands in; bf16, untiled; score_matrix_bytes is the full math-kernel score matrix. Forcing math measures below the fused path until 1536px. The 768px encode sits at the crossover, where the larger term alone is short and the device_working_mem_gb floor covers it; reproducible to three decimals across runs. Recorded in GiB." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1280, + "width": 1280, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 11141120000 + }, + "measured_bytes": 6360846565, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae with SDPBackend.MATH forced on CUDA -- the regime MPS lands in; bf16, untiled; score_matrix_bytes is the full math-kernel score matrix. Forcing math measures below the fused path until 1536px. The 768px encode sits at the crossover, where the larger term alone is short and the device_working_mem_gb floor covers it; reproducible to three decimals across runs. Recorded in GiB." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1536, + "width": 1536, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 23102226432 + }, + "measured_bytes": 12847320924, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae with SDPBackend.MATH forced on CUDA -- the regime MPS lands in; bf16, untiled; score_matrix_bytes is the full math-kernel score matrix. Forcing math measures below the fused path until 1536px. The 768px encode sits at the crossover, where the larger term alone is short and the device_working_mem_gb floor covers it; reproducible to three decimals across runs. Recorded in GiB." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "decode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 285212672 + }, + "measured_bytes": 1810328715, + "method": "reserved_peak", + "device": "RX 9070 XT (gfx1201)", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "decode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 1443889152 + }, + "measured_bytes": 3973918491, + "method": "reserved_peak", + "device": "RX 9070 XT (gfx1201)", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 4563402752 + }, + "measured_bytes": 7063073718, + "method": "reserved_peak", + "device": "RX 9070 XT (gfx1201)", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "decode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 285212672 + }, + "measured_bytes": 1847909679, + "method": "reserved_peak", + "device": "PRO W7900 (gfx1100)", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals.", + "max_ratio": 1.15 + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "decode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 1443889152 + }, + "measured_bytes": 4059817837, + "method": "reserved_peak", + "device": "PRO W7900 (gfx1100)", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals.", + "max_ratio": 1.15 + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 4563402752 + }, + "measured_bytes": 7197291446, + "method": "reserved_peak", + "device": "PRO W7900 (gfx1100)", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals.", + "max_ratio": 1.15 + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "encode", + "height": 512, + "width": 512, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 285212672 + }, + "measured_bytes": 1408749273, + "method": "reserved_peak", + "device": "RX 9070 XT (gfx1201) and PRO W7900 (gfx1100), identical to the byte", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "encode", + "height": 768, + "width": 768, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 1443889152 + }, + "measured_bytes": 3170759606, + "method": "reserved_peak", + "device": "RX 9070 XT (gfx1201) and PRO W7900 (gfx1100), identical to the byte", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "miopen" + }, + "features": { + "operation": "encode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 4563402752 + }, + "measured_bytes": 5637144576, + "method": "reserved_peak", + "device": "RX 9070 XT (gfx1201) and PRO W7900 (gfx1100), identical to the byte", + "provenance": "scripts/calibrate_flux2_working_memory.py --only vae, torch 2.10, bf16, untiled; a clean run (no MIOPEN_FIND_MODE override, no allocator GC threshold). score_matrix_bytes is the math-kernel score matrix the ROCm head-dim guard prices for the 512-wide head, unchunked; on these cards forcing math moved the peak by nothing up to 1024px. Recorded in GiB to three decimals." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 2, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 8541616210, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "#9519: 1024px decode by latent batch, CUDA, bf16, fused SDPA, one fresh process per point: 4.23 / 7.96 / 11.89 GiB at batch 1 / 2 / 3. The test table wrote them as GiB x 1000 (4229, 7955, 11889); batch 1 is the 4.229 GiB 1024px decode above." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "decode", + "height": 1024, + "width": 1024, + "tile": null, + "batch_size": 3, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 12765716546, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "#9519: 1024px decode by latent batch, CUDA, bf16, fused SDPA, one fresh process per point: 4.23 / 7.96 / 11.89 GiB at batch 1 / 2 / 3. The test table wrote them as GiB x 1000 (4229, 7955, 11889); batch 1 is the 4.229 GiB 1024px decode above." + }, + { + "estimator": "vae.flux2", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 2024, + "width": 2024, + "tile": 512, + "batch_size": 1, + "element_size": 2, + "score_matrix_bytes": 0 + }, + "measured_bytes": 585105408, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "#9519: reference-image encode at the 2024px reference cap with the forced 512px tile, CUDA, bf16, fused SDPA; flat at ~0.55 GB from 1024px up. Recorded in MiB." + }, + { + "estimator": "vae.flux", + "version": 1, + "platform": { + "conv_backend": "cudnn" + }, + "features": { + "operation": "encode", + "height": 3072, + "width": 3072, + "element_size": 2, + "tile": null, + "score_matrix_bytes": 0 + }, + "measured_bytes": 20229295964, + "method": "reserved_peak", + "device": "RTX 4090", + "provenance": "Commit 326bdf1bd0: the FLUX.1 VAE as a diffusers AutoencoderKL, bf16, fused SDPA, untiled 3072px encode: 18.84 GiB. Re-measured 18.57 GiB on torch 2.13+cu130, 2026-10-17, with scripts/calibrate_flux_vae_working_memory.py." + } + ] +} diff --git a/invokeai/backend/working_memory/estimate.py b/invokeai/backend/working_memory/estimate.py new file mode 100644 index 00000000000..1f388687557 --- /dev/null +++ b/invokeai/backend/working_memory/estimate.py @@ -0,0 +1,52 @@ +from collections.abc import Mapping +from dataclasses import dataclass, field + +FeatureValue = int | float | str | bool | None +"""One runtime fact an estimator is evaluated on: a size, an element width, a tile, a byte count it measured.""" + + +@dataclass(frozen=True) +class WorkingMemoryTerm: + """One part of a working-memory estimate. + + A fitted term is proportional to a data coefficient; `fitted` builds it, so its `bytes` are exactly the + coefficient's value times `scale`. An analytic term (a score matrix, an image buffer) carries neither. `included` is + False for a term the estimate priced but does not reserve, such as the smaller of two phases that peak apart. + """ + + name: str + bytes: float + included: bool + coefficient: str | None = None + scale: float | None = None + + @classmethod + def fitted(cls, name: str, coefficient: str, scale: float, values: Mapping[str, float]) -> "WorkingMemoryTerm": + """A term of `scale` units at `values[coefficient]` bytes per unit.""" + return cls(name, values[coefficient] * scale, True, coefficient, scale) + + +@dataclass(frozen=True) +class WorkingMemoryEstimate: + """What an operation is expected to need on its compute device, and the terms that add up to it.""" + + estimator: str + """The id of the `EstimatorSpec` that produced it, e.g. ``vae.flux2``.""" + version: int + features: Mapping[str, FeatureValue] + """Everything the estimate was evaluated from, runtime decisions such as the tile and the score-matrix size + included, so it can be evaluated again from these alone.""" + terms: tuple[WorkingMemoryTerm, ...] + source: str + """Where the coefficients came from, e.g. ``shipped:conv_backend=miopen (values of vae.flux2)``.""" + reservation_bytes: int = field(init=False) + """What the model cache keeps free: the included terms' bytes, truncated. Derived from the terms, never given.""" + + def __post_init__(self) -> None: + # One term at a time in term order rather than `sum()`, which compensates float rounding since Python 3.12, + # so the reservation is exactly the running total the terms describe. + total: float = 0 + for term in self.terms: + if term.included: + total += term.bytes + object.__setattr__(self, "reservation_bytes", int(total)) diff --git a/invokeai/backend/working_memory/platform.py b/invokeai/backend/working_memory/platform.py new file mode 100644 index 00000000000..5dfc2a00997 --- /dev/null +++ b/invokeai/backend/working_memory/platform.py @@ -0,0 +1,45 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal + +import torch + +ConvBackend = Literal["cudnn", "miopen", "mps", "xpu", "cpu"] + + +@dataclass(frozen=True) +class Platform: + """The traits of the machine an operation runs on, as coefficients are keyed. Each field is a platform key an + `EstimatorSpec` can name, and holds the true value: what was measured, and what was not, is `UNMEASURED_PRICED_AS`'s + business.""" + + conv_backend: ConvBackend + + +UNMEASURED_PRICED_AS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType( + { + # MPS, XPU and CPU convolutions are unmeasured and take the cuDNN coefficients. On MPS that pairs with a + # score-matrix term the SDPA probe always charges there, so the total is not thin. + "conv_backend": MappingProxyType({"mps": "cudnn", "xpu": "cudnn", "cpu": "cudnn"}), + } +) +"""Per platform key, the values no coefficients are measured for, and the measured value each is priced as. The one +place that borrowing is declared: the resolver applies it and names it in an estimate's `source`, and shipped data +carries entries for measured values only.""" + + +def platform_for(device: torch.device) -> Platform: + """The platform an operation on `device` runs on. + + A HIP build reports ``device.type == "cuda"``, so the torch build -- not the device string -- is what separates + MIOpen from cuDNN, and a `cpu_only` model on a ROCm build is not on MIOpen. Any device type other than CUDA, MPS and + XPU runs torch's CPU kernels, or none at all, and is priced as the CPU. + """ + if device.type == "cuda": + return Platform(conv_backend="miopen" if torch.version.hip is not None else "cudnn") + if device.type == "mps": + return Platform(conv_backend="mps") + if device.type == "xpu": + return Platform(conv_backend="xpu") + return Platform(conv_backend="cpu") diff --git a/invokeai/backend/working_memory/spec.py b/invokeai/backend/working_memory/spec.py new file mode 100644 index 00000000000..3f0c116a2eb --- /dev/null +++ b/invokeai/backend/working_memory/spec.py @@ -0,0 +1,99 @@ +import importlib +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Final + +import torch + +from invokeai.backend.working_memory.coefficients import ResolvedCoefficients, shipped_coefficients +from invokeai.backend.working_memory.estimate import FeatureValue, WorkingMemoryEstimate, WorkingMemoryTerm +from invokeai.backend.working_memory.platform import platform_for + +TermsFunction = Callable[[Mapping[str, FeatureValue], Mapping[str, float]], tuple[WorkingMemoryTerm, ...]] +"""An estimator's pure part: its features and its coefficient values in, its terms out.""" + +ESTIMATOR_MODULES: Final = ("invokeai.backend.util.vae_working_memory",) +"""Every module that registers an `EstimatorSpec`. The registry holds only the specs of modules imported so far, so a +caller that needs all of them imports these (`import_estimator_modules`); `tests/backend/working_memory/` fails on a +module that declares a spec without being listed here.""" + + +@dataclass(frozen=True) +class EstimatorSpec: + """What a working-memory estimator declares: its id and version, the coefficients its fitted terms read, and the + platform keys those coefficients are shipped per. + + Registered at import of the module that defines the estimator. `terms` is that estimator's pure part, so an + estimate can be evaluated again from recorded features alone -- which is how the measured points under `data/` are + checked. It is behaviour rather than declaration and stays out of equality and hashing. + """ + + id: str + version: int + """Bumped whenever `terms` changes the formula, so coefficients fitted to the old one no longer resolve.""" + coefficients: frozenset[str] + platform_keys: tuple[str, ...] + terms: TermsFunction = field(compare=False, repr=False) + + def estimate(self, features: Mapping[str, FeatureValue], device: torch.device) -> WorkingMemoryEstimate: + """The estimate for an operation on `device`, priced with the shipped coefficients of the platform it runs on.""" + return self.evaluate(features, shipped_coefficients(self, platform_for(device))) + + def evaluate( + self, features: Mapping[str, FeatureValue], coefficients: ResolvedCoefficients + ) -> WorkingMemoryEstimate: + """The estimate for `features` with the given coefficients.""" + return WorkingMemoryEstimate( + estimator=self.id, + version=self.version, + features=features, + terms=self.terms(features, coefficients.values), + source=coefficients.source, + ) + + +_SPECS: Final[dict[str, EstimatorSpec]] = {} + + +def _defined_at(terms: TermsFunction) -> str: + return f"{terms.__module__}.{getattr(terms, '__qualname__', repr(terms))}" + + +def register(spec: EstimatorSpec) -> EstimatorSpec: + """Declare an estimator, once per id, from the module that defines it. + + `importlib.reload` of that module declares the same spec again, from a new `terms` function of the same module and + name: a no-op that adopts the re-executed function. Anything else under an id already taken raises -- differing + declared data, the same data from another function, or the very same function object again (two specs of one module + sharing a `terms` function and, by a copy-paste slip, an id). Each would otherwise silently replace one estimator's + formula or coefficients with another's. + + jurigged (`--dev_reload`) patches function bodies in place and re-runs only the top-level statements that changed: + an edited `terms` body takes effect with nothing re-registered, but an edited `register(...)` statement raises even + when its data is unchanged (a reformat, a comment), because it hands over the same function object. Restart the + server then. + """ + declared = _SPECS.get(spec.id) + if declared is not None: + if declared != spec: + raise ValueError(f"Working-memory estimator '{spec.id}' is already registered as {declared}.") + if declared.terms is spec.terms or _defined_at(declared.terms) != _defined_at(spec.terms): + raise ValueError( + f"Working-memory estimator '{spec.id}' is already registered by {_defined_at(declared.terms)}; " + "ids are unique per estimator." + ) + _SPECS[spec.id] = spec + return spec + + +def registered_specs() -> Mapping[str, EstimatorSpec]: + """Every spec registered so far, by id: those of the estimator modules imported in this process.""" + return MappingProxyType(_SPECS) + + +def import_estimator_modules() -> Mapping[str, EstimatorSpec]: + """Every spec there is: imports `ESTIMATOR_MODULES` first, so the answer does not depend on import order.""" + for name in ESTIMATOR_MODULES: + importlib.import_module(name) + return registered_specs() diff --git a/pyproject.toml b/pyproject.toml index 786a71e56d0..e57bc3dd037 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -548,6 +548,8 @@ version = { attr = "invokeai.version.__version__" } # absent from a wheel install and that loader raises FileNotFoundError. "invokeai.backend.minimax_h3" = ["*.json"] "invokeai.backend.t5" = ["tokenizer/*.json"] +# Fitted working-memory coefficients, read by every estimator that declares an EstimatorSpec. +"invokeai.backend.working_memory" = ["data/*.json"] "invokeai.backend.clip" = ["*.json", "tokenizer/*.json", "tokenizer/*.gz"] "invokeai.app.services.image_index" = ["*.txt"] "invokeai.app.services.workflow_records.default_workflows" = ["*.json"] diff --git a/scripts/calibrate_flux2_working_memory.py b/scripts/calibrate_flux2_working_memory.py index d7d6ca34349..c429c5b9a49 100644 --- a/scripts/calibrate_flux2_working_memory.py +++ b/scripts/calibrate_flux2_working_memory.py @@ -7,8 +7,8 @@ 1. ``SDPA_MATH_BYTES_PER_SCORE_ELEMENT`` (``backend/util/attention.py``) -- bytes per element of the score matrix, charged only where SDPA has no fused kernel and materializes it. -2. ``estimate_vae_working_memory_flux2``'s 2200 / 1100 bytes per pixel per element byte - (``backend/util/vae_working_memory.py``). +2. ``estimate_vae_working_memory_flux2``'s ``k_decode`` / ``k_encode`` bytes per pixel per element byte, + shipped per convolution backend in ``backend/working_memory/data/vae.json``. 3. ``Flux2DenoiseInvocation._estimate_working_memory``'s 0.4 MB per token, defined at the Klein 9B width (4096) and scaled linearly by the loaded variant's width. 4. The dispatch question underneath all of it: does this build's SDPA fuse these shapes, or does it @@ -67,14 +67,19 @@ from invokeai.backend.model_manager.configs.flux2_variant import flux2_hidden_size from invokeai.backend.model_manager.taxonomy import Flux2VariantType from invokeai.backend.util.attention import SDPA_MATH_BYTES_PER_SCORE_ELEMENT, sdpa_score_matrix_bytes -from invokeai.backend.util.vae_working_memory import ( - _flux_vae_scaling_constant, - estimate_vae_working_memory_flux2, -) +from invokeai.backend.util.vae_working_memory import VAE_FLUX2_SPEC, estimate_vae_working_memory_flux2 +from invokeai.backend.working_memory.coefficients import shipped_coefficients +from invokeai.backend.working_memory.platform import platform_for GIB = 1024**3 MIB = 1024**2 + +def _shipped_vae_coefficient(operation: str) -> float: + """The VAE coefficient this build's estimate uses for `operation`: the column its convolution backend selects.""" + return shipped_coefficients(VAE_FLUX2_SPEC, platform_for(torch.device("cuda"))).values[f"k_{operation}"] + + # The FLUX.2 VAE attends on the 8x-downsampled grid with a single 512-wide head. Mirrors # `_FLUX2_VAE_*` in vae_working_memory.py. VAE_SPATIAL_COMPRESSION = 8 @@ -264,7 +269,9 @@ def run() -> None: if peak is None: return row - estimate = estimate_vae_working_memory_flux2(operation=operation, image_tensor=x, vae=vae, device=device) + estimate = estimate_vae_working_memory_flux2( + operation=operation, image_tensor=x, vae=vae, device=device + ).reservation_bytes seq_len = (px // VAE_SPATIAL_COMPRESSION) ** 2 score_bytes = sdpa_score_matrix_bytes( device=device, dtype=param.dtype, num_heads=1, head_dim=VAE_MID_BLOCK_HEAD_DIM, seq_len=seq_len @@ -278,8 +285,8 @@ def run() -> None: "reserved_delta": peak, "estimate": estimate, "score_term": score_bytes, - # The whole measured peak over the pixel area, which is what the 2200/1100-style literals - # name. The score matrix is deliberately NOT backed out: the two terms peak in different + # The whole measured peak over the pixel area, which is what the `k_decode`/`k_encode` + # coefficients name. The score matrix is deliberately NOT backed out: the two terms peak in different # phases of the forward, so the estimate takes the larger of them rather than the sum, and # the convolution phase is what sets the high-water mark until the quadratic term overtakes # it past ~1280px. Where `score_term` exceeds the measured peak, this column is not the @@ -457,10 +464,10 @@ def report_sdpa(shapes: list[tuple[int, int, int]], dtype_name: str) -> list[dic def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[dict]: - decode_k = _flux_vae_scaling_constant("decode", torch.device("cuda")) - encode_k = _flux_vae_scaling_constant("encode", torch.device("cuda")) - print(f"\n=== 3. VAE linear constants (this build selects {decode_k} decode / {encode_k} encode) ===") - print("`implied_k` is the measured peak over pixel area, directly comparable to those literals;") + decode_k = _shipped_vae_coefficient("decode") + encode_k = _shipped_vae_coefficient("encode") + print(f"\n=== 3. VAE linear coefficients (this build selects {decode_k:.0f} decode / {encode_k:.0f} encode) ===") + print("`implied_k` is the measured peak over pixel area, directly comparable to those coefficients;") print("fit on the rows whose `math` column matches what this build really does (section 1).") print("`covered` is the question that matters: is the shipped estimate an upper bound here?") print("Caveat: forcing math on a build that HAS a fused kernel is not equivalent to a build that") @@ -495,7 +502,7 @@ def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[di _flag_non_monotonic(rows) for operation in ("decode", "encode"): # Compare against the column this build actually selects, not against a hard-coded number. - shipped = _flux_vae_scaling_constant(operation, torch.device("cuda")) + shipped = _shipped_vae_coefficient(operation) for force_math in (False, True): ks = [ r["implied_linear_constant"] @@ -511,7 +518,7 @@ def report_vae(pxs: list[int], dtype_name: str, vae_path: str | None) -> list[di continue mode = "math" if force_math else "fused" verdict = "OK" if shipped >= max(ks) else "SHORT" - print(f" {operation} ({mode}): implied_k max = {max(ks):.0f}, shipped = {shipped} -> {verdict}") + print(f" {operation} ({mode}): implied_k max = {max(ks):.0f}, shipped = {shipped:.0f} -> {verdict}") short = [r for r in rows if not r.get("oom") and not r["covered"]] if short: print("\n Points the shipped estimate does NOT cover:") diff --git a/scripts/calibrate_flux_vae_working_memory.py b/scripts/calibrate_flux_vae_working_memory.py index c737005939e..9c62db1c2ef 100644 --- a/scripts/calibrate_flux_vae_working_memory.py +++ b/scripts/calibrate_flux_vae_working_memory.py @@ -93,7 +93,7 @@ def _worker(vae_path: str, operation: str, h: int, w: int, dtype_name: str, tile else: x = torch.randn(1, 3, h, w, device="cuda", dtype=dtype) - estimate = estimate_vae_working_memory_flux(operation, x, vae, tile_size=tile) + estimate = estimate_vae_working_memory_flux(operation, x, vae, tile_size=tile).reservation_bytes torch.cuda.synchronize() torch.cuda.empty_cache() @@ -218,8 +218,8 @@ def main() -> None: print(f" {r['operation']} {r['h']}x{r['w']} tile={r['tile']} headroom {r['headroom']:.2f}x") print() print(" The cache under-reserves here, so it keeps other models resident and the VAE can") - print(" OOM in the last step of a generation. The constants in `vae_working_memory.py` need") - print(" a backend branch, the way `estimate_vae_working_memory_qwen_image` already has one.") + print(" OOM in the last step of a generation. The vae.flux coefficients for this convolution") + print(" backend in `invokeai/backend/working_memory/data/vae.json` need refitting.") else: worst = min(measured, key=lambda r: r["headroom"]) if measured else None if worst: diff --git a/tests/app/invocations/conftest.py b/tests/app/invocations/conftest.py index 6aed50a4190..0151a2d484f 100644 --- a/tests/app/invocations/conftest.py +++ b/tests/app/invocations/conftest.py @@ -5,6 +5,8 @@ import pytest from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL +from invokeai.backend.working_memory.estimate import WorkingMemoryEstimate, WorkingMemoryTerm + def _build_flux_shaped_vae() -> AutoencoderKL: """A real `AutoencoderKL` with the FLUX.1 VAE's *shape* and almost none of its weight. @@ -47,3 +49,19 @@ def _build_flux_shaped_vae() -> AutoencoderKL: def flux_shaped_vae() -> Callable[[], AutoencoderKL]: """Factory, not an instance: a test that needs two independent VAEs should not share one.""" return _build_flux_shaped_vae + + +def _fixed_working_memory(reservation_bytes: int) -> WorkingMemoryEstimate: + return WorkingMemoryEstimate( + estimator="test.fixed", + version=1, + features={}, + terms=(WorkingMemoryTerm("fixed", reservation_bytes, included=True),), + source="test", + ) + + +@pytest.fixture +def fixed_working_memory() -> Callable[[int], WorkingMemoryEstimate]: + """An estimate that reserves exactly the given bytes, for a node whose estimator a test patches out.""" + return _fixed_working_memory diff --git a/tests/app/invocations/test_flux2_working_memory.py b/tests/app/invocations/test_flux2_working_memory.py index 286b6b44511..6eeb2180bb1 100644 --- a/tests/app/invocations/test_flux2_working_memory.py +++ b/tests/app/invocations/test_flux2_working_memory.py @@ -7,7 +7,10 @@ a 1024x1024 generation. See https://github.com/invoke-ai/InvokeAI/issues/9500. The `MEASURED_*` tables below are peak *reserved* memory measured on CUDA in bf16 (the conservative -quantity, including allocator overhead). Every estimate must stay an upper bound on them. +quantity, including allocator overhead). Every estimate must stay an upper bound on them. The VAE's +measured points ship with its coefficients in `invokeai/backend/working_memory/data/vae.json`, where +`tests/backend/working_memory/test_shipped_coefficients.py` holds the estimate to at least each point +and, where a fitted coefficient decides it, at most 1.3x (or the point's `max_ratio`). """ from contextlib import contextmanager @@ -35,14 +38,19 @@ rocm_sdpa_uses_math_kernel, sdpa_score_matrix_bytes, ) -from invokeai.backend.util.vae_working_memory import ( - _FLUX_VAE_SCALING_CONSTANTS, - estimate_vae_working_memory_flux2, -) +from invokeai.backend.util.vae_working_memory import VAE_FLUX2_SPEC, estimate_vae_working_memory_flux2 +from invokeai.backend.working_memory.coefficients import shipped_coefficients +from invokeai.backend.working_memory.platform import Platform MB = 1024**2 GB = 1024**3 + +def _vae_k(conv_backend, operation): + """The shipped FLUX.2 VAE coefficient the estimator resolves for one convolution backend and operation.""" + return shipped_coefficients(VAE_FLUX2_SPEC, Platform(conv_backend)).values[f"k_{operation}"] + + # The estimator's default width, and the head count that follows from it. The suite's pinned # `MEASURED_DENOISE` table was taken on Klein 9B, which is that default, so those rows are unmoved # by the width scaling. @@ -272,18 +280,6 @@ def test_the_lora_margin_scales_with_the_batch(self): class TestFlux2VaeWorkingMemoryEstimate: - # (operation, pixel size, measured peak reserved MB), bf16, untiled. - MEASURED_VAE = [ - ("decode", 512, 1086), - ("decode", 768, 2414), - ("decode", 1024, 4260), - ("decode", 1328, 7146), - ("decode", 1536, 9578), - ("encode", 512, 536), - ("encode", 1024, 2122), - ("encode", 1328, 3022), - ] - def _mock_bf16_vae(self): vae = MagicMock(spec=AutoencoderKLFlux2) vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) # element_size == 2 @@ -293,19 +289,12 @@ def _tensor_for(self, operation, px): # decode receives 32-channel latents at pixels/8; encode receives a pixel image. return torch.zeros(1, 32, px // 8, px // 8) if operation == "decode" else torch.zeros(1, 3, px, px) - @pytest.mark.parametrize("operation, px, measured_mb", MEASURED_VAE) - def test_estimate_is_an_upper_bound_on_measured_peak(self, operation, px, measured_mb): - estimate = estimate_vae_working_memory_flux2( - operation=operation, image_tensor=self._tensor_for(operation, px), vae=self._mock_bf16_vae(), device=FUSED - ) - assert estimate >= measured_mb * MB - - @pytest.mark.parametrize("operation, expected_constant", [("decode", 2200), ("encode", 1100)]) - def test_constant_scales_pixel_area_and_element_size(self, operation, expected_constant): + @pytest.mark.parametrize("operation", ["decode", "encode"]) + def test_constant_scales_pixel_area_and_element_size(self, operation): estimate = estimate_vae_working_memory_flux2( operation=operation, image_tensor=self._tensor_for(operation, 1024), vae=self._mock_bf16_vae(), device=FUSED ) - assert estimate == 1024 * 1024 * 2 * expected_constant + assert estimate.reservation_bytes == 1024 * 1024 * 2 * _vae_k("cudnn", operation) def test_tiled_estimate_is_bounded_by_the_tile_not_the_image(self): """Reference-image encoding forces 512px tiling precisely so the peak stops following the @@ -317,79 +306,26 @@ def test_tiled_estimate_is_bounded_by_the_tile_not_the_image(self): vae=self._mock_bf16_vae(), tile_size=512, device=FUSED, - ) + ).reservation_bytes for px in (1024, 1328, 2024) ] assert len(set(estimates)) == 1 - assert estimates[0] == int(512 * 512 * 2 * 1100 * 1.25) - assert estimates[0] >= 558 * MB # measured tiled peak at 2024px + assert estimates[0] == int(512 * 512 * 2 * _vae_k("cudnn", "encode") * 1.25) # The whole point of tiling: it must shrink the reservation, not just bound the VAE. untiled = estimate_vae_working_memory_flux2( operation="encode", image_tensor=torch.zeros(1, 3, 2024, 2024), vae=self._mock_bf16_vae(), device=FUSED - ) + ).reservation_bytes assert estimates[0] < untiled / 4 class TestVaeConstantsFollowTheConvBackend: - """The pixel-area constant is not one number. MIOpen's convolution workspaces cost far more than - cuDNN's for the same decode, and the gap is not the attention term -- it is identical on the - fused path. Fitted with `scripts/calibrate_flux2_working_memory.py --only vae`, peak reserved, - each point in a fresh process: - - decode encode - cuDNN RTX 4090 2180 2185 2165 1072 1063 1061 - MIOpen RX 9070 XT 3453 3369 3368 2687 2688 2688 - - Flat across 512/768/1024px on both, so the linear model holds; only the coefficient moves. Note - the encode column: cuDNN's is half its decode, MIOpen's is four fifths, so "encoding costs half - of decoding" -- the ratio the other estimators in the module use -- is a cuDNN property, not an - architectural one. + """The pixel-area coefficient is not one number. MIOpen's convolution workspaces cost far more + than cuDNN's for the same decode, and the gap is not the attention term -- it is identical on the + fused path. Both columns, the cards they were fitted on and the points that bound them are in + `invokeai/backend/working_memory/data/vae.json`; what is pinned here is that a device selects the + column of the backend it really runs on. """ - # (backend, operation, px, measured GiB) - MEASURED = [ - ("cudnn", "decode", 512, 1.064), - ("cudnn", "decode", 768, 2.400), - ("cudnn", "decode", 1024, 4.229), - ("cudnn", "encode", 512, 0.523), - ("cudnn", "encode", 768, 1.168), - ("cudnn", "encode", 1024, 2.072), - ("miopen", "decode", 512, 1.686), - ("miopen", "decode", 768, 3.701), - ("miopen", "decode", 1024, 6.578), - # PRO W7900 (gfx1100). Costs slightly more per pixel than gfx1201 and sets the column. - ("miopen", "decode", 512, 1.721), - ("miopen", "decode", 768, 3.781), - ("miopen", "decode", 1024, 6.703), - ("miopen", "encode", 512, 1.312), - ("miopen", "encode", 768, 2.953), - ("miopen", "encode", 1024, 5.250), - ] - - @pytest.mark.parametrize("backend, operation, px, measured_gib", MEASURED) - def test_the_constant_upper_bounds_its_own_backend(self, backend, operation, px, measured_gib): - """Each column has to bound the hardware it was fitted on. Before this the cuDNN column was - used everywhere, leaving a 2.3 GiB shortfall on a 1024px MIOpen decode.""" - constant = _FLUX_VAE_SCALING_CONSTANTS[backend][operation] - estimate = px * px * 2 * constant # bf16 element size - assert estimate >= measured_gib * GB - assert estimate <= 1.3 * measured_gib * GB - - def test_the_cudnn_constant_would_not_cover_miopen(self): - """The regression this guards: a 1024px MIOpen decode needs 6.58 GiB and the cuDNN constant - reserves 4.30. That is the same class of shortfall #9500 reports, one backend over.""" - cudnn = _FLUX_VAE_SCALING_CONSTANTS["cudnn"]["decode"] - assert 1024 * 1024 * 2 * cudnn < 6.578 * GB - - def test_the_encode_ratio_is_not_architectural(self): - """cuDNN's encode is half its decode; MIOpen's is four fifths. A shared ratio cannot express - both, which is why the table carries the two operations separately.""" - ratios = { - backend: consts["encode"] / consts["decode"] for backend, consts in _FLUX_VAE_SCALING_CONSTANTS.items() - } - assert ratios["cudnn"] == pytest.approx(0.50, abs=0.02) - assert ratios["miopen"] == pytest.approx(0.76, abs=0.03) - def test_a_rocm_build_selects_the_miopen_column(self): """A HIP build reports `device.type == "cuda"`, so the torch build is the discriminator.""" vae = MagicMock(spec=AutoencoderKLFlux2) @@ -401,7 +337,7 @@ def estimate(): v.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) return estimate_vae_working_memory_flux2( operation="decode", image_tensor=latents, vae=v, device=torch.device("cuda") - ) + ).reservation_bytes # `_IS_ROCM` is pinned because it is bound at import: on a real HIP build the head-dim # guard short-circuits ahead of the probe patched here and charges a score matrix to both @@ -421,15 +357,15 @@ def estimate(): cuda = estimate() assert rocm > cuda - assert rocm / cuda == pytest.approx(3600 / 2200, rel=1e-3) + assert rocm / cuda == pytest.approx(_vae_k("miopen", "decode") / _vae_k("cudnn", "decode"), rel=1e-3) def test_a_rocm_build_composes_the_miopen_column_with_the_guard_score_matrix(self): """The state every ROCm rig is actually in, driven through the estimator. A HIP build takes the MIOpen column *and* the head-dim guard's score matrix -- the guard is unconditional for this VAE's 512-wide head -- so `max` has to pick between two terms that - are both larger than the cuDNN pair. `TestVaeTermsDoNotAdd` checks the same two terms - against the constants table, which cannot see an estimator that stops combining them: a + are both larger than the cuDNN pair. The shipped points check the same two terms against the + measurements, which cannot see an estimator that stops combining them: a plausible "the score matrix covers it anyway" refactor that drops back to the cuDNN column whenever SDPA materializes under-reserves a 1024px decode by 1.6x and passes every other test in this module. @@ -443,7 +379,7 @@ def estimate(px): image_tensor=torch.zeros(1, 32, px // 8, px // 8), vae=vae, device=torch.device("cuda"), - ) + ).reservation_bytes with ( patch("torch.version.hip", "7.1.25424"), @@ -456,9 +392,8 @@ def estimate(px): at_1024 = estimate(1024) at_1536 = estimate(1536) - miopen = _FLUX_VAE_SCALING_CONSTANTS["miopen"]["decode"] + miopen = _vae_k("miopen", "decode") assert at_1024 == 1024 * 1024 * 2 * miopen, "at 1024px the MIOpen linear term is the larger" - assert at_1024 >= 6.703 * GB, "and it still covers the measured W7900 peak it was fitted to" assert at_1536 == 36864 * 36864 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT, "at 1536px the score matrix is" assert at_1536 > 1536 * 1536 * 2 * miopen @@ -470,13 +405,10 @@ class TestFlux2VaeBatchIsBudgeted: the OOM happens anyway. Measured at 1024px on CUDA/bf16, peak reserved, each point in a fresh process: 4.23GB at batch 1, - 7.96GB at batch 2, 11.89GB at batch 3. Linear, and slightly sub-linear per sample, so scaling the - single-sample estimate is an upper bound rather than a fit. + 7.96GB at batch 2, 11.89GB at batch 3 (shipped as `vae.flux2` points). Linear, and slightly + sub-linear per sample, so scaling the single-sample estimate is an upper bound rather than a fit. """ - # (batch, measured peak reserved MB) for a 1024px decode. - MEASURED_DECODE_BATCH = [(1, 4229), (2, 7955), (3, 11889)] - def _decode_estimate(self, batch, px=1024, tile_size=None, device=FUSED): vae = MagicMock(spec=AutoencoderKLFlux2) vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) @@ -486,13 +418,7 @@ def _decode_estimate(self, batch, px=1024, tile_size=None, device=FUSED): vae=vae, tile_size=tile_size, device=device, - ) - - @pytest.mark.parametrize("batch, measured_mb", MEASURED_DECODE_BATCH) - def test_estimate_is_an_upper_bound_on_the_measured_batch_peak(self, batch, measured_mb): - estimate = self._decode_estimate(batch) - assert estimate >= measured_mb * MB - assert estimate <= 2 * measured_mb * MB + ).reservation_bytes def test_estimate_scales_with_the_batch(self): """The regression in one assertion: before this, all three of these were equal.""" @@ -506,7 +432,7 @@ def test_a_three_dimensional_tensor_is_one_sample(self): vae.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) unbatched = estimate_vae_working_memory_flux2( operation="decode", image_tensor=torch.zeros(32, 128, 128), vae=vae, device=FUSED - ) + ).reservation_bytes assert unbatched == self._decode_estimate(1) def test_tiling_bounds_the_tile_not_the_batch(self): @@ -549,13 +475,13 @@ def _mock_vae_info(self): vae_info.model_on_device = MagicMock(return_value=cm) return vae_info - def test_decode_requests_working_memory(self): + def test_decode_requests_working_memory(self, fixed_working_memory): vae_info = self._mock_vae_info() context = MagicMock() context.models.load.return_value = vae_info context.tensors.load.return_value = torch.zeros(1, 32, 128, 128) - expected = 10 * GB + expected = fixed_working_memory(10 * GB) with patch( "invokeai.app.invocations.vae.flux2_vae_decode.estimate_vae_working_memory_flux2", return_value=expected ) as estimate: @@ -571,14 +497,14 @@ def test_decode_requests_working_memory(self): estimate.assert_called_once() assert estimate.call_args.kwargs["operation"] == "decode" - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected) + vae_info.model_on_device.assert_called_once_with(working_memory=expected) - def test_encode_requests_working_memory(self): + def test_encode_requests_working_memory(self, fixed_working_memory): vae_info = self._mock_vae_info() context = MagicMock() context.models.load.return_value = vae_info - expected = 4 * GB + expected = fixed_working_memory(4 * GB) with ( patch( "invokeai.app.invocations.vae.flux2_vae_encode.estimate_vae_working_memory_flux2", return_value=expected @@ -598,7 +524,7 @@ def test_encode_requests_working_memory(self): estimate.assert_called_once() assert estimate.call_args.kwargs["operation"] == "encode" - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected) + vae_info.model_on_device.assert_called_once_with(working_memory=expected) class _StopBeforeLoad(Exception): @@ -782,9 +708,9 @@ def _pinned_build(): score matrix here", but on a HIP build that same token carries two more decisions with it, and both would make this module's verdict depend on the runner: - * `_flux2_vae_scaling_constant` selects the MIOpen convolution column, 1.6-2.5x the cuDNN one - every expectation guarded by this helper is written against -- so the VAE assertions would - pass on a CUDA runner and fail on a ROCm one. + * `platform_for` selects the MIOpen convolution coefficients, 1.6-2.5x the cuDNN ones every + expectation guarded by this helper is written against -- so the VAE assertions would pass on + a CUDA runner and fail on a ROCm one. * `rocm_sdpa_uses_math_kernel` short-circuits ahead of the probe below, on a threshold a user can move with `INVOKE_ROCM_FUSED_SDPA_MAX_HEAD_DIM`. At the default it happens to agree with `_materializing_probe`; at a lower one it charges a score matrix for the denoise's 128-wide @@ -845,7 +771,7 @@ def _vae_estimate(self, operation, px, device, tile_size=None): vae=self._mock_bf16_vae(), tile_size=tile_size, device=device, - ) + ).reservation_bytes # (operation, pixel size, mid-block tokens). The VAE attends on the 8x-downsampled grid. @pytest.mark.parametrize( @@ -907,20 +833,17 @@ class TestVaeTermsDoNotAdd: decoder, last in the encoder -- so the full-resolution convolution feature maps that drive the linear term are not live while the score matrix is, and peak *reserved* is a high-water mark. - Measured, decode, peak reserved, fresh process per point: - - px cuDNN fused cuDNN forced-math linear(2200) score(17) - 1024 4.229 3.527 4.30 4.25 - 1280 6.590 5.924 6.71 10.38 - 1536 9.353 11.965 9.67 21.52 - - Forcing math measures *below* the fused path until 1536px, and even there it exceeds it by - 2.6GB against the 21.5GB the score term prices standalone -- the attention phase reuses blocks - the allocator is already holding. On gfx1100 and gfx1201, forcing math moves the measured peak - by nothing at all up to 1024px, and the totals stay flat-linear in area either way. + Measured, decode, peak reserved, fresh process per point (the cuDNN fused and forced-math points + in `invokeai/backend/working_memory/data/vae.json`): forcing math measures *below* the fused path + until 1536px, and even there it exceeds it by 2.6GB against the 21.5GB the score term prices + standalone -- the attention phase reuses blocks the allocator is already holding. On gfx1100 and + gfx1201, forcing math moves the measured peak by nothing at all up to 1024px, and the totals stay + flat-linear in area either way. Summing the terms reserved 11.1GB for a 1024px gfx1100 decode that measures 6.7GB. Taking the - max reserves 6.8GB. + max reserves 6.8GB. Where a max model is not an upper bound on its own -- the crossover, a 768px + encode with cuDNN's coefficient and a materializing kernel -- the `device_working_mem_gb` floor + is, and the shipped point there is checked against the floor. """ def _estimate(self, px, device, materializing): @@ -929,64 +852,20 @@ def _estimate(self, px, device, materializing): with _materializing() if materializing else _pinned_build(): return estimate_vae_working_memory_flux2( operation="decode", image_tensor=torch.zeros(1, 32, px // 8, px // 8), vae=vae, device=device - ) - - # (pixel size, measured GiB) for a decode on the PRO W7900. Its 512px point is what sets the - # MIOpen decode constant: implied 3525, where gfx1201 asks for only 3453. - MEASURED_W7900_DECODE = [(1024, 6.703), (768, 3.781), (512, 1.721)] - - @pytest.mark.parametrize("px, measured_gib", MEASURED_W7900_DECODE) - def test_the_max_model_bounds_the_measured_miopen_peak(self, px, measured_gib): - linear = px * px * 2 * _FLUX_VAE_SCALING_CONSTANTS["miopen"]["decode"] - score = (px // 8) ** 4 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - assert max(linear, score) >= measured_gib * GB - assert max(linear, score) <= 1.15 * measured_gib * GB - - def test_the_sum_model_would_have_over_reserved_by_two_thirds(self): - """At 1024px the sum reserves 11.1GiB for a decode that measures 6.7 -- on a 16GB card that - is the difference between the transformer staying resident and being evicted.""" - linear = 1024 * 1024 * 2 * _FLUX_VAE_SCALING_CONSTANTS["miopen"]["decode"] - score = 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - assert (linear + score) / (6.703 * GB) > 1.6 - assert max(linear, score) / (6.703 * GB) < 1.15 + ).reservation_bytes def test_the_estimate_is_the_larger_term_not_the_sum(self): """At 1024px the two are within 2% of each other on the cuDNN column, which makes this the sharpest place to tell the models apart.""" - linear = 1024 * 1024 * 2 * _FLUX_VAE_SCALING_CONSTANTS["cudnn"]["decode"] + linear = 1024 * 1024 * 2 * _vae_k("cudnn", "decode") score = 16384 * 16384 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT assert self._estimate(1024, MATERIALIZING, materializing=True) == max(linear, score) assert self._estimate(1024, MATERIALIZING, materializing=True) < linear + score - # (operation, px, measured GiB) with cuDNN's linear constant and a materializing kernel -- the - # combination MPS lands in, and the only place the max model is not an upper bound on its own. - MEASURED_CUDNN_FORCED_MATH = [ - ("encode", 512, 0.523), - ("encode", 768, 1.801), - ("encode", 1024, 4.072), - ("decode", 1024, 3.527), - ("decode", 1280, 5.924), - ("decode", 1536, 11.965), - ] - - @pytest.mark.parametrize("operation, px, measured_gib", MEASURED_CUDNN_FORCED_MATH) - def test_the_working_memory_floor_covers_the_crossover(self, operation, px, measured_gib): - """A max model is weakest where the two terms are near-equal, and one measured point shows - it: a 768px encode wants 1.80GiB against a 1.35GiB max. It is not reachable as a shortfall, - because the cache floors every reservation at `device_working_mem_gb` and the whole crossover - region sits below that floor. Reproducible to three decimals across runs, so this is the - model's real shape, not noise -- which is why it is pinned rather than rounded away.""" - constant = _FLUX_VAE_SCALING_CONSTANTS["cudnn"][operation] - area = px * px if operation == "decode" else px * px - linear = area * 2 * constant - score = (px // 8) ** 4 * SDPA_MATH_BYTES_PER_SCORE_ELEMENT - floored = max(max(linear, score), 3 * GB) # what ModelCache._get_vram_available reserves - assert floored >= measured_gib * GB - def test_a_fused_build_is_unaffected(self): """The max only ever removes reservation, never adds it: with no score matrix the estimate is the linear term, exactly as before.""" - linear = 1024 * 1024 * 2 * _FLUX_VAE_SCALING_CONSTANTS["cudnn"]["decode"] + linear = 1024 * 1024 * 2 * _vae_k("cudnn", "decode") assert self._estimate(1024, FUSED, materializing=False) == linear def test_regional_prompting_adds_the_score_matrix(self): @@ -1283,7 +1162,7 @@ def estimate(): v.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)]) return estimate_vae_working_memory_flux2( operation="decode", image_tensor=torch.zeros(1, 32, 128, 128), vae=v, device=FUSED - ) + ).reservation_bytes with _diffusers_backend("_native_math"): forced_math = estimate() diff --git a/tests/app/invocations/test_flux_vae_decode.py b/tests/app/invocations/test_flux_vae_decode.py index 4e206cf62b2..8b93e064d3e 100644 --- a/tests/app/invocations/test_flux_vae_decode.py +++ b/tests/app/invocations/test_flux_vae_decode.py @@ -32,7 +32,7 @@ def _loaded_vae(shift_factor: float | None, scaling_factor: float = 0.3611) -> M vae_info.compute_device = torch.device("cpu") @contextmanager - def _on_device(working_mem_bytes=None): + def _on_device(working_memory=None): yield (None, vae) vae_info.model_on_device = _on_device diff --git a/tests/app/invocations/test_flux_vae_decode_tiling.py b/tests/app/invocations/test_flux_vae_decode_tiling.py index 2b6f998067e..1fca3e20687 100644 --- a/tests/app/invocations/test_flux_vae_decode_tiling.py +++ b/tests/app/invocations/test_flux_vae_decode_tiling.py @@ -107,21 +107,25 @@ def test_the_tiling_state_is_restored_afterwards(self, flux_shaped_vae): [(False, False, None), (True, False, 256), (False, True, 0)], ) def test_the_reservation_matches_the_decode_that_will_run( - self, flux_shaped_vae, tiled, force_tiled_decode, expected_tile_size + self, flux_shaped_vae, fixed_working_memory, tiled, force_tiled_decode, expected_tile_size ): """A tiled decode reserved for a single pass is the OOM the retry exists to avoid, and an untiled decode reserved for a tile evicts models it did not need to.""" path = "invokeai.app.invocations.vae.flux_vae_decode.estimate_vae_working_memory_flux" vae = flux_shaped_vae() - _, context, _ = _build_decode_mocks( + vae_info, context, _ = _build_decode_mocks( vae, torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512), force_tiled_decode=force_tiled_decode ) - with patch(path, return_value=1024) as estimate: + estimated = fixed_working_memory(1024) + with patch(path, return_value=estimated) as estimate: _build_invocation(tiled=tiled, tile_size=256 if tiled else 0).invoke(context) assert estimate.call_args.kwargs["tile_size"] == expected_tile_size + vae_info.model_on_device.assert_called_once_with(working_memory=estimated) @pytest.mark.parametrize("auto", [True, False], ids=["auto-tiled-decode-on", "auto-tiled-decode-off"]) - def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(self, auto, flux_shaped_vae): + def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off( + self, auto, flux_shaped_vae, fixed_working_memory + ): """On Windows an allocation that does not fit is paged rather than refused, so the OOM retry never fires and the decision has to be made from the estimate alone.""" module = "invokeai.app.invocations.vae.flux_vae_decode" @@ -129,8 +133,9 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se flux_shaped_vae(), torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512) ) context.config.get.return_value.auto_tiled_decode = auto + untiled, tiled = fixed_working_memory(20 * 2**30), fixed_working_memory(2 * 2**30) with ( - patch(f"{module}.estimate_vae_working_memory_flux", side_effect=[20 * 2**30, 2 * 2**30]) as estimate, + patch(f"{module}.estimate_vae_working_memory_flux", side_effect=[untiled, tiled]) as estimate, patch(f"{module}.should_pretile_vae_decode", return_value=True) as pretile, ): _build_invocation().invoke(context) @@ -138,11 +143,12 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se if auto: pretile.assert_called_once_with(vae_info.compute_device, 20 * 2**30) assert estimate.call_args.kwargs["tile_size"] == 0 - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=2 * 2**30) + vae_info.model_on_device.assert_called_once_with(working_memory=tiled) assert tiling_state[0]["use_tiling"] is True assert tiling_state[0]["tile_sample_min_size"] == DEFAULT_TILE_SAMPLE_MIN_SIZE else: pretile.assert_not_called() + vae_info.model_on_device.assert_called_once_with(working_memory=untiled) assert tiling_state[0]["use_tiling"] is False diff --git a/tests/app/invocations/test_vae_working_memory_backend.py b/tests/app/invocations/test_vae_working_memory_backend.py index 5829d061afe..9dfc4fade93 100644 --- a/tests/app/invocations/test_vae_working_memory_backend.py +++ b/tests/app/invocations/test_vae_working_memory_backend.py @@ -1,7 +1,7 @@ -"""Which convolution-backend constant the diffusers-AutoencoderKL estimators use. +"""Which convolution-backend coefficients the diffusers-AutoencoderKL estimators use, and what the SD estimator adds up. SD1/SDXL, SD3 and CogView4 run the same AutoencoderKL convolution stack as the FLUX.1 autoencoder, so they take the -same measured per-backend constants: MIOpen needs ~1.64x what cuDNN does. The change is large enough to matter -- an +same measured per-backend coefficients: MIOpen needs ~1.64x what cuDNN does. The change is large enough to matter -- an SD1.5 1024px fp32 decode reserves 9.49 GB on cuDNN and 15.36 GB on MIOpen -- and the branch is invisible on a CUDA machine, so it is pinned here rather than left to whichever card the suite runs on. """ @@ -13,10 +13,15 @@ from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL from invokeai.backend.util.vae_working_memory import ( + VAE_COGVIEW4_SPEC, + VAE_SD3_SPEC, + VAE_SD_SPEC, estimate_vae_working_memory_cogview4, estimate_vae_working_memory_sd3, estimate_vae_working_memory_sd15_sdxl, ) +from invokeai.backend.working_memory.coefficients import shipped_coefficients +from invokeai.backend.working_memory.platform import Platform def _vae(dtype: torch.dtype = torch.float16) -> MagicMock: @@ -25,39 +30,74 @@ def _vae(dtype: torch.dtype = torch.float16) -> MagicMock: return vae -def _estimate(estimator, device: torch.device, hip: str | None): +def _estimate(estimator, device: torch.device, hip: str | None) -> int: latents = torch.zeros(1, 4, 64, 64) # 512px output with ( patch("torch.version.hip", hip), # The mid-block score matrix is the other term in these estimates and has its own tests; hold it at zero so - # the assertions below are about the convolution constant alone. + # the assertions below are about the convolution coefficient alone. patch("invokeai.backend.util.vae_working_memory._vae_mid_block_score_matrix_bytes", return_value=0), ): if estimator is estimate_vae_working_memory_sd15_sdxl: - return estimator( + estimate = estimator( operation="decode", image_tensor=latents, vae=_vae(), tile_size=None, fp32=False, device=device ) - return estimator(operation="decode", image_tensor=latents, vae=_vae(), device=device) + else: + estimate = estimator(operation="decode", image_tensor=latents, vae=_vae(), device=device) + return estimate.reservation_bytes + + +def _decode_bytes(spec, conv_backend: str) -> float: + """A 512px fp16 decode at the shipped coefficient of that backend.""" + return 512 * 512 * 2 * shipped_coefficients(spec, Platform(conv_backend)).values["k_decode"] ESTIMATORS = [ - pytest.param(estimate_vae_working_memory_sd15_sdxl, id="sd15-sdxl"), - pytest.param(estimate_vae_working_memory_sd3, id="sd3"), - pytest.param(estimate_vae_working_memory_cogview4, id="cogview4"), + pytest.param(estimate_vae_working_memory_sd15_sdxl, VAE_SD_SPEC, id="sd15-sdxl"), + pytest.param(estimate_vae_working_memory_sd3, VAE_SD3_SPEC, id="sd3"), + pytest.param(estimate_vae_working_memory_cogview4, VAE_COGVIEW4_SPEC, id="cogview4"), ] -@pytest.mark.parametrize("estimator", ESTIMATORS) -def test_a_rocm_vae_reserves_the_miopen_constant(estimator): - assert _estimate(estimator, torch.device("cuda", 0), hip="7.2.0") == 512 * 512 * 2 * 3600 +@pytest.mark.parametrize("estimator, spec", ESTIMATORS) +def test_a_rocm_vae_reserves_the_miopen_coefficient(estimator, spec): + assert _estimate(estimator, torch.device("cuda", 0), hip="7.2.0") == _decode_bytes(spec, "miopen") -@pytest.mark.parametrize("estimator", ESTIMATORS) -def test_a_cuda_vae_reserves_the_cudnn_constant(estimator): - assert _estimate(estimator, torch.device("cuda", 0), hip=None) == 512 * 512 * 2 * 2200 +@pytest.mark.parametrize("estimator, spec", ESTIMATORS) +def test_a_cuda_vae_reserves_the_cudnn_coefficient(estimator, spec): + assert _estimate(estimator, torch.device("cuda", 0), hip=None) == _decode_bytes(spec, "cudnn") -@pytest.mark.parametrize("estimator", ESTIMATORS) -def test_a_cpu_only_vae_reserves_the_cudnn_constant_even_on_a_rocm_build(estimator): +@pytest.mark.parametrize("estimator, spec", ESTIMATORS) +def test_a_cpu_only_vae_reserves_the_cudnn_coefficient_even_on_a_rocm_build(estimator, spec): """A `cpu_only` VAE does not run on MIOpen, so naming its device has to change the answer.""" - assert _estimate(estimator, torch.device("cpu"), hip="7.2.0") == 512 * 512 * 2 * 2200 + assert _estimate(estimator, torch.device("cpu"), hip="7.2.0") == _decode_bytes(spec, "cudnn") + + +def _sd_fp32_decode(tile_size: int | None) -> int: + latents = torch.zeros(1, 4, 96, 128) # a 768x1024 output + with patch("invokeai.backend.util.vae_working_memory._vae_mid_block_score_matrix_bytes", return_value=0): + estimate = estimate_vae_working_memory_sd15_sdxl( + operation="decode", + image_tensor=latents, + vae=_vae(), + tile_size=tile_size, + fp32=True, + device=torch.device("cpu"), + ) + return estimate.reservation_bytes + + +FP32_WEIGHTS = 250 * 2**20 +"""What an fp32 decode adds for the upcast weights.""" + + +def test_a_tiled_fp32_sd_decode_is_one_tile_with_its_margin_plus_the_fp32_weights(): + k = shipped_coefficients(VAE_SD_SPEC, Platform("cudnn")).values["k_decode"] + assert _sd_fp32_decode(512) == int(512 * 512 * 4 * k * 1.25 + FP32_WEIGHTS) + + +def test_an_untiled_fp32_sd_decode_is_the_whole_image_plus_the_fp32_weights(): + k = shipped_coefficients(VAE_SD_SPEC, Platform("cudnn")).values["k_decode"] + assert _sd_fp32_decode(None) == int(768 * 1024 * 4 * k + FP32_WEIGHTS) diff --git a/tests/app/invocations/test_z_image_tiled_decode.py b/tests/app/invocations/test_z_image_tiled_decode.py index 2f8a46a894d..eb9d3a3a9b6 100644 --- a/tests/app/invocations/test_z_image_tiled_decode.py +++ b/tests/app/invocations/test_z_image_tiled_decode.py @@ -10,7 +10,19 @@ from invokeai.app.invocations.vae.z_image_latents_to_image import ZImageLatentsToImageInvocation from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.vae_tiling_scope import DEFAULT_TILE_SAMPLE_MIN_SIZE, MIN_TILE_SAMPLE_SIZE -from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux +from invokeai.backend.util.vae_working_memory import VAE_FLUX_SPEC, estimate_vae_working_memory_flux +from invokeai.backend.working_memory.coefficients import shipped_coefficients +from invokeai.backend.working_memory.platform import Platform + + +def _flux_reservation(*args, **kwargs) -> int: + """What `estimate_vae_working_memory_flux` reserves; the estimator cells are about its size.""" + return estimate_vae_working_memory_flux(*args, **kwargs).reservation_bytes + + +def _k(conv_backend: str, operation: str) -> float: + """The shipped FLUX.1 VAE coefficient the estimator resolves for one convolution backend and operation.""" + return shipped_coefficients(VAE_FLUX_SPEC, Platform(conv_backend)).values[f"k_{operation}"] def _mock_flux_vae(element_size_bytes: int = 2) -> MagicMock: @@ -36,7 +48,7 @@ def _fused_non_rocm_build(self, monkeypatch): score matrix for it. On a ROCm rig that resolves to `cuda`, where the head-dim guard charges 4.25GiB for the 1024px case below -- which the expected value here clears by 1%, so these exact-value assertions turn build-dependent the moment either constant moves. Pin the fused, non-HIP regime they are written for, which - is also the cuDNN column of `_FLUX_VAE_SCALING_CONSTANTS`; `TestClassicVaeEstimators` in + also selects the cuDNN coefficients; `TestClassicVaeEstimators` in `tests/backend/util/test_rocm_sdpa_head_dim_guard.py` owns the ROCm side of this estimator. """ import invokeai.backend.util.attention as attention @@ -46,11 +58,18 @@ def _fused_non_rocm_build(self, monkeypatch): monkeypatch.setattr(vwm.TorchDevice, "choose_torch_device", classmethod(lambda cls: torch.device("cpu"))) @pytest.mark.parametrize( - ("hip", "operation", "constant"), - [(None, "decode", 2200), (None, "encode", 1100), ("7.14.0", "decode", 3600), ("7.14.0", "encode", 2750)], + ("hip", "operation", "conv_backend"), + [ + (None, "decode", "cudnn"), + (None, "encode", "cudnn"), + ("7.14.0", "decode", "miopen"), + ("7.14.0", "encode", "miopen"), + ], ids=["cudnn-decode", "cudnn-encode", "miopen-decode", "miopen-encode"], ) - def test_the_convolution_backend_of_the_vaes_device_picks_the_constant(self, monkeypatch, hip, operation, constant): + def test_the_convolution_backend_of_the_vaes_device_picks_the_constant( + self, monkeypatch, hip, operation, conv_backend + ): """MIOpen's workspaces are larger than cuDNN's: measured 3451 decode / 2688 encode on an RX 9060 XT, the same as the FLUX.2 VAE. A HIP build reports its GPU as "cuda", so the torch build decides, for the VAE's device.""" import invokeai.backend.util.vae_working_memory as vwm @@ -59,22 +78,22 @@ def test_the_convolution_backend_of_the_vaes_device_picks_the_constant(self, mon monkeypatch.setattr(vwm, "_vae_mid_block_score_matrix_bytes", lambda *args, **kwargs: 0) tensor = torch.zeros(1, 16, 64, 64) if operation == "decode" else torch.zeros(1, 3, 512, 512) - estimate = estimate_vae_working_memory_flux( + estimate = _flux_reservation( operation=operation, image_tensor=tensor, vae=_mock_flux_vae(), device=torch.device("cuda", 0) ) - on_cpu = estimate_vae_working_memory_flux( + on_cpu = _flux_reservation( operation=operation, image_tensor=tensor, vae=_mock_flux_vae(), device=torch.device("cpu") ) - assert estimate == 512 * 512 * 2 * constant - assert on_cpu == 512 * 512 * 2 * (2200 if operation == "decode" else 1100), "a cpu_only VAE runs on cuDNN terms" + assert estimate == 512 * 512 * 2 * _k(conv_backend, operation) + assert on_cpu == 512 * 512 * 2 * _k("cudnn", operation), "a cpu_only VAE runs on cuDNN terms" def test_the_default_reproduces_the_untiled_estimate(self): """Regression guard for the six call sites that pass no tile_size at all.""" latents = torch.zeros(1, 16, 128, 128) - # 1024x1024 output px * 2 bytes * 2200 - expected = 1024 * 1024 * 2 * 2200 - actual = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) + # 1024x1024 output px * 2 bytes * the cuDNN decode coefficient + expected = 1024 * 1024 * 2 * _k("cudnn", "decode") + actual = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) assert actual == expected def test_the_tiled_estimate_is_a_tile_bounded_decode_plus_the_assembled_image(self): @@ -86,13 +105,13 @@ def test_the_tiled_estimate_is_a_tile_bounded_decode_plus_the_assembled_image(se under-reservation that no other assertion here can see. """ # 512px tiles, fp16, over a 1536x1536 output. - decode_term = 512 * 512 * 2 * 2200 * 1.25 + decode_term = 512 * 512 * 2 * _k("cudnn", "decode") * 1.25 # 4 RGB buffers at the VAE's element size, plus the 8-bit one the node hands to PIL. Four # is measured, not assumed: the assembled image grows the device peak by ~1.6 copies on a # diffusers VAE and ~0.1 on the host-merging BFL one, and the node adds its in-window # clamp/scale copies on top. See the constant's comment for the measurements. image_term = 1536 * 1536 * 3 * (4 * 2 + 1) - estimate = estimate_vae_working_memory_flux( + estimate = _flux_reservation( operation="decode", image_tensor=torch.zeros(1, 16, 192, 192), vae=_mock_flux_vae(), tile_size=512 ) assert estimate == int(decode_term + image_term) @@ -101,10 +120,10 @@ def test_both_tiled_terms_scale_with_the_vaes_element_size(self): """An fp32 VAE costs twice the bytes per pixel, in the tile and in the assembled image alike. Only the 8-bit buffer is fixed, which is why this is a little under 2x.""" latents = torch.zeros(1, 16, 192, 192) - fp16 = estimate_vae_working_memory_flux( + fp16 = _flux_reservation( operation="decode", image_tensor=latents, vae=_mock_flux_vae(element_size_bytes=2), tile_size=512 ) - fp32 = estimate_vae_working_memory_flux( + fp32 = _flux_reservation( operation="decode", image_tensor=latents, vae=_mock_flux_vae(element_size_bytes=4), tile_size=512 ) assert 1.95 < fp32 / fp16 < 2.0 @@ -120,10 +139,8 @@ def test_a_tiled_estimate_stays_a_fraction_of_the_untiled_one(self, latent_hw): the clamped tile plus the same image term, which can land above the single-pass estimate.) """ latents = torch.zeros(1, 16, *latent_hw) - untiled = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) - tiled = estimate_vae_working_memory_flux( - operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=512 - ) + untiled = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) + tiled = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=512) assert tiled < untiled / 2 def test_a_tile_larger_than_the_image_is_not_estimated_as_tiled(self): @@ -134,10 +151,8 @@ def test_a_tile_larger_than_the_image_is_not_estimated_as_tiled(self): decode needing ~4.6GB, evicting the transformer from the cache to make room for nothing. """ latents = torch.zeros(1, 16, 128, 128) # 1024px - untiled = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) - oversized = estimate_vae_working_memory_flux( - operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=4096 - ) + untiled = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) + oversized = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=4096) assert oversized == untiled def test_a_tile_wider_than_one_axis_is_priced_at_the_clamped_tile(self): @@ -148,10 +163,8 @@ def test_a_tile_wider_than_one_axis_is_priced_at_the_clamped_tile(self): single-pass decode that tiling is there to avoid -- ~92GB against ~37GB. """ latents = torch.zeros(1, 16, 128, 1024) # 1024 x 8192 px: only the long axis exceeds the tile - tiled = estimate_vae_working_memory_flux( - operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=4096 - ) - untiled = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) + tiled = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=4096) + untiled = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae()) assert tiled < untiled def test_a_tiled_estimate_still_covers_the_assembled_image(self): @@ -163,10 +176,10 @@ def test_a_tiled_estimate_still_covers_the_assembled_image(self): second resolution the image term is checked at, so a term that did not scale with area would fail here even if it matched at 1536px. """ - estimate = estimate_vae_working_memory_flux( + estimate = _flux_reservation( operation="decode", image_tensor=torch.zeros(1, 16, 1024, 1024), vae=_mock_flux_vae(), tile_size=512 ) - tile_term = int(512 * 512 * 2 * 2200 * 1.25) + tile_term = int(512 * 512 * 2 * _k("cudnn", "decode") * 1.25) one_full_resolution_buffer = 8192 * 8192 * 3 * 2 # The assembled output plus at least one of the post-processing copies. assert estimate >= tile_term + 2 * one_full_resolution_buffer @@ -178,8 +191,8 @@ def test_the_sentinel_does_not_read_the_size_off_the_vae(self): latents = torch.zeros(1, 16, 192, 192) vae = _mock_flux_vae() vae.tile_sample_min_size = 384 # as if a previous run had set it - estimate = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=vae, tile_size=0) - assert estimate == estimate_vae_working_memory_flux( + estimate = _flux_reservation(operation="decode", image_tensor=latents, vae=vae, tile_size=0) + assert estimate == _flux_reservation( operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=DEFAULT_TILE_SAMPLE_MIN_SIZE ) @@ -190,17 +203,15 @@ def test_the_sentinel_works_on_a_vae_without_that_attribute_at_all(self): vae = MagicMock(spec=AutoencoderKL) vae.parameters.side_effect = lambda: iter([torch.zeros(1, dtype=torch.float16)]) del vae.tile_sample_min_size - estimate = estimate_vae_working_memory_flux(operation="decode", image_tensor=latents, vae=vae, tile_size=0) - assert estimate == estimate_vae_working_memory_flux( + estimate = _flux_reservation(operation="decode", image_tensor=latents, vae=vae, tile_size=0) + assert estimate == _flux_reservation( operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=DEFAULT_TILE_SAMPLE_MIN_SIZE ) def test_a_tile_below_the_cost_floor_is_estimated_at_the_floor(self): latents = torch.zeros(1, 16, 192, 192) - estimate = estimate_vae_working_memory_flux( - operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=8 - ) - assert estimate == estimate_vae_working_memory_flux( + estimate = _flux_reservation(operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=8) + assert estimate == _flux_reservation( operation="decode", image_tensor=latents, vae=_mock_flux_vae(), tile_size=MIN_TILE_SAMPLE_SIZE ) @@ -221,23 +232,21 @@ def _fused_non_rocm_build(self, monkeypatch): monkeypatch.setattr("invokeai.backend.util.vae_working_memory.sdpa_score_matrix_bytes", lambda *a, **k: 0) monkeypatch.setattr(torch.version, "hip", None, raising=False) - def test_the_untiled_estimate_is_unchanged(self): - """Four call sites pass no tile and depend on this number staying what it was.""" + def test_an_untiled_encode_is_the_whole_image_at_the_encode_coefficient(self): + """What the call sites that pass no tile reserve: the full image, priced at the cuDNN encode coefficient.""" image = torch.zeros(1, 3, 1024, 1024) - assert estimate_vae_working_memory_flux("encode", image, _mock_flux_vae()) == 1024 * 1024 * 2 * 1100 + assert _flux_reservation("encode", image, _mock_flux_vae()) == 1024 * 1024 * 2 * _k("cudnn", "encode") def test_a_tiled_estimate_is_flat_in_the_image_size_where_the_untiled_one_is_not(self): vae = _mock_flux_vae() estimates = [ - estimate_vae_working_memory_flux("encode", torch.zeros(1, 3, edge, edge), vae, tile_size=512) + _flux_reservation("encode", torch.zeros(1, 3, edge, edge), vae, tile_size=512) for edge in (1024, 2048, 3072) ] # Not exactly equal: the input image and the assembled latent are not bounded by tiling. # But they are a rounding error beside the tile term, which is what "flat" has to mean. assert max(estimates) < min(estimates) * 1.2 - untiled = [ - estimate_vae_working_memory_flux("encode", torch.zeros(1, 3, edge, edge), vae) for edge in (1024, 3072) - ] + untiled = [_flux_reservation("encode", torch.zeros(1, 3, edge, edge), vae) for edge in (1024, 3072)] assert untiled[1] == untiled[0] * 9 def test_the_tiled_estimate_is_one_tile_plus_the_input_and_every_live_copy_of_the_moments(self): @@ -247,12 +256,10 @@ def test_the_tiled_estimate_is_one_tile_plus_the_input_and_every_live_copy_of_th vae = _mock_flux_vae() edge, tile = 2048, 512 live_moment_copies = 1.0 / (1.0 - 0.25) ** 2 + 2.0 - expected = tile * tile * 2 * 1100 * 1.25 + expected = tile * tile * 2 * _k("cudnn", "encode") * 1.25 expected += edge * edge * 3 * 2 expected += (edge // 8) * (edge // 8) * 2 * 16 * 2 * live_moment_copies - assert estimate_vae_working_memory_flux("encode", torch.zeros(1, 3, edge, edge), vae, tile_size=tile) == int( - expected - ) + assert _flux_reservation("encode", torch.zeros(1, 3, edge, edge), vae, tile_size=tile) == int(expected) def test_it_does_not_inherit_the_decode_residual(self): """The decode adds four element-size copies of the assembled image plus a byte buffer, for @@ -262,8 +269,8 @@ def test_it_does_not_inherit_the_decode_residual(self): image = torch.zeros(1, 3, 2048, 2048) latents = torch.zeros(1, 16, 2048 // 8, 2048 // 8) - encode = estimate_vae_working_memory_flux("encode", image, vae, tile_size=512) - decode = estimate_vae_working_memory_flux("decode", latents, vae, tile_size=512) + encode = _flux_reservation("encode", image, vae, tile_size=512) + decode = _flux_reservation("decode", latents, vae, tile_size=512) # Same tile, same output area, so the whole difference is the residual term. assert decode - encode > 2048 * 2048 * 3 * 2, "the encode is paying for the decode's image copies" @@ -273,9 +280,7 @@ def test_a_tile_larger_than_the_image_is_priced_as_the_untiled_encode(self): oversized tile would otherwise reserve for a pass that never runs.""" vae = _mock_flux_vae() image = torch.zeros(1, 3, 512, 512) - assert estimate_vae_working_memory_flux( - "encode", image, vae, tile_size=1024 - ) == estimate_vae_working_memory_flux("encode", image, vae) + assert _flux_reservation("encode", image, vae, tile_size=1024) == _flux_reservation("encode", image, vae) def _build_decode_mocks(vae, latents: torch.Tensor, decoded: torch.Tensor, force_tiled_decode: bool = False): @@ -364,15 +369,18 @@ def test_the_tiling_state_is_restored_afterwards(self, flux_shaped_vae): assert (vae.use_tiling, vae.tile_sample_min_size, vae.tile_latent_min_size) == before @pytest.mark.parametrize("auto", [True, False], ids=["auto-tiled-decode-on", "auto-tiled-decode-off"]) - def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(self, auto, flux_shaped_vae): + def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off( + self, auto, flux_shaped_vae, fixed_working_memory + ): """Where the driver pages instead of raising an OOM, the retry below never fires: the estimate decides.""" module = "invokeai.app.invocations.vae.z_image_latents_to_image" vae_info, context, state = _build_decode_mocks( flux_shaped_vae(), torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512) ) context.config.get.return_value.auto_tiled_decode = auto + untiled, tiled = fixed_working_memory(20 * 2**30), fixed_working_memory(2 * 2**30) with ( - patch(f"{module}.estimate_vae_working_memory_flux", side_effect=[20 * 2**30, 2 * 2**30]) as estimate, + patch(f"{module}.estimate_vae_working_memory_flux", side_effect=[untiled, tiled]) as estimate, patch(f"{module}.should_pretile_vae_decode", return_value=True) as pretile, ): _build_invocation().invoke(context) @@ -380,18 +388,21 @@ def test_a_decode_too_large_for_its_gpu_is_tiled_up_front_unless_switched_off(se if auto: pretile.assert_called_once_with(vae_info.compute_device, 20 * 2**30) assert estimate.call_args.kwargs["tile_size"] == 0 - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=2 * 2**30) + vae_info.model_on_device.assert_called_once_with(working_memory=tiled) assert state[0]["use_tiling"] is True else: pretile.assert_not_called() assert estimate.call_count == 1 + vae_info.model_on_device.assert_called_once_with(working_memory=untiled) assert state[0]["use_tiling"] is False - def test_requested_tiling_skips_the_untiled_estimate(self, flux_shaped_vae): + def test_requested_tiling_skips_the_untiled_estimate(self, flux_shaped_vae, fixed_working_memory): module = "invokeai.app.invocations.vae.z_image_latents_to_image" _, context, _ = _build_decode_mocks(flux_shaped_vae(), torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) with ( - patch(f"{module}.estimate_vae_working_memory_flux", return_value=2 * 2**30) as estimate, + patch( + f"{module}.estimate_vae_working_memory_flux", return_value=fixed_working_memory(2 * 2**30) + ) as estimate, patch(f"{module}.should_pretile_vae_decode") as pretile, ): _build_invocation(tiled=True).invoke(context) @@ -399,11 +410,13 @@ def test_requested_tiling_skips_the_untiled_estimate(self, flux_shaped_vae): pretile.assert_not_called() @pytest.mark.parametrize("tiled,expected_tile_size", [(False, None), (True, 0)]) - def test_the_estimate_is_tile_bounded_only_when_tiling(self, flux_shaped_vae, tiled, expected_tile_size): + def test_the_estimate_is_tile_bounded_only_when_tiling( + self, flux_shaped_vae, fixed_working_memory, tiled, expected_tile_size + ): path = "invokeai.app.invocations.vae.z_image_latents_to_image.estimate_vae_working_memory_flux" vae = flux_shaped_vae() _, context, _ = _build_decode_mocks(vae, torch.zeros(1, 16, 64, 64), torch.zeros(1, 3, 512, 512)) - with patch(path, return_value=1024) as estimate: + with patch(path, return_value=fixed_working_memory(1024)) as estimate: _build_invocation(tiled=tiled).invoke(context) assert estimate.call_args.kwargs["tile_size"] == expected_tile_size diff --git a/tests/app/invocations/test_z_image_working_memory.py b/tests/app/invocations/test_z_image_working_memory.py index c04e44e2b0d..655714a8c9a 100644 --- a/tests/app/invocations/test_z_image_working_memory.py +++ b/tests/app/invocations/test_z_image_working_memory.py @@ -42,7 +42,7 @@ def _vae_info(vae) -> MagicMock: class TestTheEstimateReachesTheCache: - def test_the_decode_node_reserves_what_it_estimated(self, flux_shaped_vae): + def test_the_decode_node_reserves_what_it_estimated(self, flux_shaped_vae, fixed_working_memory): vae = flux_shaped_vae() vae_info = _vae_info(vae) @@ -53,7 +53,7 @@ def test_the_decode_node_reserves_what_it_estimated(self, flux_shaped_vae): # `ImageOutput.build` validates these, so they cannot be bare MagicMocks. context.images.save.return_value = MagicMock(image_name="test.png", width=512, height=512) - expected_memory = 500 * 1024 * 1024 + expected_memory = fixed_working_memory(500 * 1024 * 1024) path = "invokeai.app.invocations.vae.z_image_latents_to_image.estimate_vae_working_memory_flux" with patch(path, return_value=expected_memory) as estimate: ZImageLatentsToImageInvocation.model_construct( @@ -64,18 +64,18 @@ def test_the_decode_node_reserves_what_it_estimated(self, flux_shaped_vae): ).invoke(context) estimate.assert_called_once() - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected_memory) + vae_info.model_on_device.assert_called_once_with(working_memory=expected_memory) - def test_the_encode_node_reserves_what_it_estimated(self, flux_shaped_vae): + def test_the_encode_node_reserves_what_it_estimated(self, flux_shaped_vae, fixed_working_memory): vae = flux_shaped_vae() vae_info = _vae_info(vae) - expected_memory = 250 * 1024 * 1024 + expected_memory = fixed_working_memory(250 * 1024 * 1024) path = "invokeai.app.invocations.vae.z_image_image_to_latents.estimate_vae_working_memory_flux" with patch(path, return_value=expected_memory) as estimate: latents = ZImageImageToLatentsInvocation.vae_encode(vae_info, torch.zeros(1, 3, 512, 512)) estimate.assert_called_once() - vae_info.model_on_device.assert_called_once_with(working_mem_bytes=expected_memory) + vae_info.model_on_device.assert_called_once_with(working_memory=expected_memory) # The encode really ran: a 512px image through an 8x, 16-channel VAE is a 64px latent. assert latents.shape == (1, 16, 64, 64) diff --git a/tests/backend/model_manager/load/test_loaded_model.py b/tests/backend/model_manager/load/test_loaded_model.py index 04368325bde..9793e2b89e3 100644 --- a/tests/backend/model_manager/load/test_loaded_model.py +++ b/tests/backend/model_manager/load/test_loaded_model.py @@ -1,3 +1,5 @@ +import logging + import pytest import torch @@ -12,6 +14,8 @@ from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import ( apply_custom_layers_to_model, ) +from invokeai.backend.util.logging import InvokeAILogger +from invokeai.backend.working_memory.estimate import WorkingMemoryEstimate, WorkingMemoryTerm class ModelWithRequiredScale(torch.nn.Module): @@ -25,10 +29,12 @@ class FakeCache: def __init__(self): self.lock_calls = 0 self.unlock_calls = 0 + self.reserved: list[int | None] = [] def lock(self, cache_record: CacheRecord, working_mem_bytes: int | None, max_move_bytes: int | None = None) -> bool: - del cache_record, working_mem_bytes, max_move_bytes + del cache_record, max_move_bytes self.lock_calls += 1 + self.reserved.append(working_mem_bytes) return True def unlock(self, cache_record: CacheRecord) -> None: @@ -63,6 +69,62 @@ def test_model_on_device_leaves_full_load_models_unchanged(): assert all(param.device.type == "cpu" for param in model.parameters()) +def _full_load_model(cache: FakeCache) -> LoadedModelWithoutConfig: + cached_model = CachedModelOnlyFullLoad( + model=torch.nn.Linear(4, 4), compute_device=torch.device("meta"), total_bytes=1, keep_ram_copy=False + ) + return LoadedModelWithoutConfig(cache_record=CacheRecord(key="vae-key", cached_model=cached_model), cache=cache) + + +_ESTIMATE = WorkingMemoryEstimate( + estimator="vae.flux2", + version=1, + features={"operation": "decode"}, + terms=(WorkingMemoryTerm.fitted("convolutions", "k_decode", 3 * 2**29 + 0.25, {"k_decode": 2.0}),), + source="shipped:conv_backend=cudnn", +) + + +def test_a_working_memory_estimate_reserves_its_bytes_and_logs_where_they_came_from(caplog): + cache = FakeCache() + logger = "invokeai.backend.model_manager.load.load_base" + # Created before the level is lowered: creating it sets the configured level, which would undo that. + InvokeAILogger.get_logger(logger) + + with ( + caplog.at_level(logging.DEBUG, logger=logger), + _full_load_model(cache).model_on_device(working_memory=_ESTIMATE), + ): + pass + + assert cache.reserved == [3 * 2**30] + assert cache.unlock_calls == 1 + (record,) = [r for r in caplog.records if r.name == logger] + assert "vae-key" in record.message + assert "vae.flux2 reserves 3072 MB (shipped:conv_backend=cudnn)" in record.message + + +def test_an_estimate_and_a_byte_count_together_are_refused_before_locking(): + cache = FakeCache() + + with pytest.raises(ValueError, match="not both"): + with _full_load_model(cache).model_on_device(1024, working_memory=_ESTIMATE): + pass + + assert cache.lock_calls == 0 + + +def test_an_estimate_passed_as_a_byte_count_says_where_it_goes(): + """The positional slot is `working_mem_bytes`; an estimate there would fail deep inside the cache instead.""" + cache = FakeCache() + + with pytest.raises(TypeError, match="working_memory="): + with _full_load_model(cache).model_on_device(_ESTIMATE): + pass + + assert cache.lock_calls == 0 + + def test_supports_partial_loading_reflects_cached_model_type(): partial_model = CachedModelWithPartialLoad( model=torch.nn.Linear(4, 4), compute_device=torch.device("meta"), keep_ram_copy=False diff --git a/tests/backend/util/test_rocm_sdpa_head_dim_guard.py b/tests/backend/util/test_rocm_sdpa_head_dim_guard.py index e82d1ccead1..6710b3afbf3 100644 --- a/tests/backend/util/test_rocm_sdpa_head_dim_guard.py +++ b/tests/backend/util/test_rocm_sdpa_head_dim_guard.py @@ -488,7 +488,7 @@ class TestClassicVaeEstimators: def _rocm_cuda(self, monkeypatch): """The guard's ROCm switch and the convolution backend are separate questions, and this class is about the first: the assertions below carry the cuDNN linear term, so pin `torch.version.hip` (which selects MIOpen's - larger constants) away from it while the attention side stays on ROCm.""" + larger coefficients) away from it while the attention side stays on ROCm.""" import invokeai.backend.util.vae_working_memory as vwm monkeypatch.setattr(attention, "_IS_ROCM", True) @@ -508,16 +508,25 @@ def _score_bytes(px: int) -> int: tokens = (px // 8) ** 2 return tokens * tokens * attention.SDPA_MATH_BYTES_PER_SCORE_ELEMENT + @staticmethod + def _cudnn_decode_bytes(px: int) -> float: + """A bf16 decode's convolution term at the shipped cuDNN coefficient.""" + from invokeai.backend.util.vae_working_memory import VAE_SD_SPEC + from invokeai.backend.working_memory.coefficients import shipped_coefficients + from invokeai.backend.working_memory.platform import Platform + + return px * px * 2 * shipped_coefficients(VAE_SD_SPEC, Platform("cudnn")).values["k_decode"] + def test_sdxl_untiled_decode_reserves_the_score_matrix_when_it_dominates(self, monkeypatch): from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_sd15_sdxl latents = torch.zeros(1, 4, 2048 // 8, 2048 // 8) - linear = 2048 * 2048 * 2 * 2200 + linear = self._cudnn_decode_bytes(2048) with_guard = estimate_vae_working_memory_sd15_sdxl( "decode", latents, self._vae(AutoencoderKL), None, fp32=False - ) + ).reservation_bytes assert with_guard == self._score_bytes(2048) > linear monkeypatch.setattr(attention, "_IS_ROCM", False) @@ -527,7 +536,9 @@ def test_sdxl_untiled_decode_reserves_the_score_matrix_when_it_dominates(self, m real_empty = torch.empty monkeypatch.setattr(torch, "empty", lambda *a, device=None, **kw: real_empty(*a, **kw)) assert ( - estimate_vae_working_memory_sd15_sdxl("decode", latents, self._vae(AutoencoderKL), None, fp32=False) + estimate_vae_working_memory_sd15_sdxl( + "decode", latents, self._vae(AutoencoderKL), None, fp32=False + ).reservation_bytes == linear ) @@ -538,7 +549,8 @@ def test_the_other_autoencoder_kl_estimators_price_it_too(self, estimator): import invokeai.backend.util.vae_working_memory as vwm latents = torch.zeros(1, 16, 2048 // 8, 2048 // 8) - assert getattr(vwm, estimator)("decode", latents, self._vae(AutoencoderKL)) == self._score_bytes(2048) + estimate = getattr(vwm, estimator)("decode", latents, self._vae(AutoencoderKL)) + assert estimate.reservation_bytes == self._score_bytes(2048) def test_a_tiny_autoencoder_has_no_attention_to_price(self): from diffusers.models.autoencoders.autoencoder_tiny import AutoencoderTiny @@ -548,7 +560,7 @@ def test_a_tiny_autoencoder_has_no_attention_to_price(self): latents = torch.zeros(1, 4, 2048 // 8, 2048 // 8) assert estimate_vae_working_memory_sd15_sdxl( "decode", latents, self._vae(AutoencoderTiny), None, fp32=False - ) == (2048 * 2048 * 2 * 2200) + ).reservation_bytes == self._cudnn_decode_bytes(2048) def test_flux1_decode_prices_the_score_matrix_by_tile_when_tiled(self): from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL @@ -556,9 +568,11 @@ def test_flux1_decode_prices_the_score_matrix_by_tile_when_tiled(self): from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux latents = torch.zeros(1, 16, 2048 // 8, 2048 // 8) - untiled = estimate_vae_working_memory_flux("decode", latents, self._vae(AutoencoderKL)) + untiled = estimate_vae_working_memory_flux("decode", latents, self._vae(AutoencoderKL)).reservation_bytes assert untiled == self._score_bytes(2048) - tiled = estimate_vae_working_memory_flux("decode", latents, self._vae(AutoencoderKL), tile_size=1024) + tiled = estimate_vae_working_memory_flux( + "decode", latents, self._vae(AutoencoderKL), tile_size=1024 + ).reservation_bytes assert tiled < self._score_bytes(2048) assert tiled >= self._score_bytes(1024), "a 1024px tile still materializes its own 16384-token score matrix" diff --git a/tests/backend/working_memory/__init__.py b/tests/backend/working_memory/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/backend/working_memory/test_shipped_coefficients.py b/tests/backend/working_memory/test_shipped_coefficients.py new file mode 100644 index 00000000000..d369d03df90 --- /dev/null +++ b/tests/backend/working_memory/test_shipped_coefficients.py @@ -0,0 +1,312 @@ +"""The shipped working-memory data, checked against the estimators that read it. + +Every coefficient entry resolves for the spec it names on a measured platform value, and every value an entry carries +is decided by at least one measured point. A point is peak *reserved* bytes, one fresh process per point (see its +`provenance`), and bounds its estimate from both sides. The model cache consumes an estimate as `free >= estimate`, so +the estimate must be at least the point. Where a fitted term decides it, it must also be at most 1.3 times the point, +or the point's tighter `max_ratio`: a reservation the operation never uses is model weights evicted for nothing. Where +an analytic term (a materialized score matrix) decides it, that term bounds rather than describes, so it gets no upper +bound, and the cache's `device_working_mem_gb` floor counts toward the lower one. +""" + +import itertools +import re +from dataclasses import asdict +from pathlib import Path +from typing import get_args, get_type_hints + +import pytest +import torch +from pydantic import ValidationError + +import invokeai +from invokeai.app.services.config.config_default import InvokeAIAppConfig +from invokeai.backend.working_memory.coefficients import ( + CoefficientEntry, + MeasuredPoint, + ResolvedCoefficients, + ValuesFrom, + WorkingMemoryData, + load_data, + resolve_coefficients, + shipped_coefficients, + shipped_data, +) +from invokeai.backend.working_memory.estimate import WorkingMemoryEstimate +from invokeai.backend.working_memory.platform import UNMEASURED_PRICED_AS, Platform, platform_for +from invokeai.backend.working_memory.spec import ESTIMATOR_MODULES, EstimatorSpec, import_estimator_modules + +MAX_OVER_RESERVATION = 1.3 +"""How far above a measured point a reservation may sit where a fitted coefficient decides it, unless the point +carries a tighter `max_ratio`.""" + +WORKING_MEMORY_FLOOR = int(InvokeAIAppConfig.model_fields["device_working_mem_gb"].default * 2**30) +"""What `ModelCache._working_mem_reserve` floors every reservation at. The field default, not a config instance an +environment variable could move.""" + +# Imported deliberately rather than relying on whichever modules ran first: a spec missing from the registry is a spec +# whose data nobody checks. +SPECS = import_estimator_modules() + + +def _platforms(*, measured: bool) -> list[Platform]: + """Every platform built from measured values only, or every platform with at least one unmeasured value.""" + every = {key: get_args(hint) for key, hint in get_type_hints(Platform).items()} + measured_values = { + key: [v for v in values if v not in UNMEASURED_PRICED_AS.get(key, {})] for key, values in every.items() + } + chosen = measured_values if measured else every + platforms = [Platform(**dict(zip(chosen, values, strict=True))) for values in itertools.product(*chosen.values())] + if measured: + return platforms + return [p for p in platforms if any(value in UNMEASURED_PRICED_AS.get(key, {}) for key, value in asdict(p).items())] + + +def _platform_key(platform: dict[str, str]) -> tuple[tuple[str, str], ...]: + return tuple(sorted(platform.items())) + + +def _point_id(point: MeasuredPoint) -> str: + features = point.features + shape = f"{features.get('height')}x{features.get('width')}" + tile = f"-tile{features['tile']}" if features.get("tile") is not None else "" + batch = f"-b{features['batch_size']}" if features.get("batch_size", 1) != 1 else "" + scores = "-scores" if features.get("score_matrix_bytes") else "" + platform = ",".join(point.platform.values()) + card = point.device.split(" (")[0].replace(" ", "") + return f"{point.estimator}-{platform}-{features.get('operation')}-{shape}{tile}{batch}{scores}-{card}" + + +def _evaluate(point: MeasuredPoint) -> WorkingMemoryEstimate: + spec = SPECS[point.estimator] + assert point.version == spec.version + return spec.evaluate(point.features, shipped_coefficients(spec, Platform(**point.platform))) + + +def _fitted_decides(estimate: WorkingMemoryEstimate) -> bool: + return any(term.included and term.coefficient is not None for term in estimate.terms) + + +POINTS = [pytest.param(point, id=_point_id(point)) for data in shipped_data().values() for point in data.points] + + +def test_every_data_file_parses_and_declares_coefficients() -> None: + for name, data in shipped_data().items(): + assert data.coefficients, f"{name} declares no coefficients" + + +def test_every_coefficient_entry_resolves_for_its_spec_on_a_measured_platform_value() -> None: + """Resolving checks what an entry provides against what its spec reads, the borrowed values included. An entry for + an unmeasured value would never be read: the resolver prices that value as `UNMEASURED_PRICED_AS` says.""" + seen: set[tuple[str, int, tuple[tuple[str, str], ...]]] = set() + for name, data in shipped_data().items(): + for entry in data.coefficients: + where = f"{name}: {entry.estimator} on {entry.platform}" + spec = SPECS.get(entry.estimator) + assert spec is not None, f"{where}: no estimator registers this id" + assert entry.version == spec.version, f"{where}: version {entry.version}, the estimator is {spec.version}" + assert set(entry.platform) == set(spec.platform_keys), f"{where}: the spec keys {spec.platform_keys}" + for key, value in entry.platform.items(): + assert value not in UNMEASURED_PRICED_AS.get(key, {}), f"{where}: {key}={value} is not measured" + key = (entry.estimator, entry.version, _platform_key(entry.platform)) + assert key not in seen, f"{where}: declared twice" + seen.add(key) + shipped_coefficients(spec, Platform(**entry.platform)) + + +@pytest.mark.parametrize("spec", sorted(SPECS.values(), key=lambda spec: spec.id), ids=lambda spec: spec.id) +def test_every_estimator_has_coefficients_on_every_measured_platform_value(spec: EstimatorSpec) -> None: + for platform in _platforms(measured=True): + assert set(shipped_coefficients(spec, platform).values) == spec.coefficients + + +@pytest.mark.parametrize("spec", sorted(SPECS.values(), key=lambda spec: spec.id), ids=lambda spec: spec.id) +def test_an_unmeasured_platform_value_is_priced_as_the_measured_one_and_says_so(spec: EstimatorSpec) -> None: + for platform in _platforms(measured=False): + priced_as = Platform( + **{key: UNMEASURED_PRICED_AS.get(key, {}).get(value, value) for key, value in asdict(platform).items()} + ) + resolved = shipped_coefficients(spec, platform) + assert resolved.values == shipped_coefficients(spec, priced_as).values + assert resolved.source.startswith( + f"shipped:conv_backend={platform.conv_backend} ({priced_as.conv_backend} values" + ) + + +def test_every_carried_value_is_decided_by_a_measured_point() -> None: + """A value no point decides is unguarded: any number would pass. An estimator without points of its own takes a + measured one's values with `values_from`, and is then held to that one's points.""" + entries = { + (entry.estimator, entry.version, _platform_key(entry.platform)): entry + for data in shipped_data().values() + for entry in data.coefficients + } + decided: set[tuple[tuple[str, int, tuple[tuple[str, str], ...]], str]] = set() + for data in shipped_data().values(): + for point in data.points: + key = (point.estimator, point.version, _platform_key(point.platform)) + lender = entries[key].values_from + owner = (lender.estimator, lender.version, key[2]) if lender is not None else key + estimate = _evaluate(point) + decided |= {(owner, term.coefficient) for term in estimate.terms if term.included and term.coefficient} + + undecided = [ + f"{entry.estimator} {entry.platform} {name}" + for key, entry in entries.items() + for name in entry.values or {} + if (key, name) not in decided + ] + assert not undecided, f"no measured point decides these shipped values: {undecided}" + + +@pytest.mark.parametrize("point", POINTS) +def test_every_measured_point_is_bounded_by_its_estimate(point: MeasuredPoint) -> None: + estimate = _evaluate(point) + fitted = _fitted_decides(estimate) + + reserved = estimate.reservation_bytes if fitted else max(estimate.reservation_bytes, WORKING_MEMORY_FLOOR) + assert reserved >= point.measured_bytes, ( + f"reserves {reserved / 2**30:.3f} GiB for a measured {point.measured_bytes / 2**30:.3f} GiB" + ) + if point.max_ratio is not None or fitted: + max_ratio = point.max_ratio if point.max_ratio is not None else MAX_OVER_RESERVATION + assert estimate.reservation_bytes <= max_ratio * point.measured_bytes, ( + f"reserves {estimate.reservation_bytes / point.measured_bytes:.2f}x the measured peak, over {max_ratio}x" + ) + + +def test_a_point_bound_is_never_looser_than_the_default() -> None: + """`max_ratio` tightens the bound for a point the estimate has to stay close to; it never excuses one.""" + for name, data in shipped_data().items(): + for point in data.points: + if point.max_ratio is not None: + assert 1.0 <= point.max_ratio < MAX_OVER_RESERVATION, f"{name}: {_point_id(point)}" + + +@pytest.mark.parametrize( + "device, hip, conv_backend", + [ + (torch.device("cuda", 0), "7.2.0", "miopen"), + (torch.device("cuda", 0), None, "cudnn"), + # A HIP build reports its GPU as "cuda"; anything else on that build is not on MIOpen. + (torch.device("cpu"), "7.2.0", "cpu"), + (torch.device("mps"), "7.2.0", "mps"), + (torch.device("xpu"), None, "xpu"), + (torch.device("meta"), None, "cpu"), + ], +) +def test_the_platform_is_the_convolution_backend_the_device_really_runs(device, hip, conv_backend, monkeypatch) -> None: + monkeypatch.setattr(torch.version, "hip", hip) + assert platform_for(device) == Platform(conv_backend=conv_backend) + + +def _no_terms(features, values): + return () + + +def _spec(spec_id: str, version: int = 1) -> EstimatorSpec: + return EstimatorSpec( + id=spec_id, version=version, coefficients=frozenset({"k"}), platform_keys=("conv_backend",), terms=_no_terms + ) + + +def _entry(estimator: str, *, version: int = 1, values=None, values_from: str | None = None) -> CoefficientEntry: + return CoefficientEntry( + estimator=estimator, + version=version, + platform={"conv_backend": "cudnn"}, + values=values, + values_from=ValuesFrom(estimator=values_from, version=1) if values_from else None, + provenance="-", + ) + + +def _data(*entries: CoefficientEntry) -> list[WorkingMemoryData]: + return [WorkingMemoryData(schema_version=1, coefficients=list(entries), points=[])] + + +def test_a_missing_or_stale_entry_names_the_estimator_and_platform() -> None: + """No fallback to another platform or version: coefficients fitted for something else are the shortfall the data + exists to prevent.""" + spec = _spec("test.estimator", version=2) + + assert resolve_coefficients( + _data(_entry(spec.id, version=2, values={"k": 7.0})), spec, Platform("cudnn"), origin="test" + ) == ResolvedCoefficients(values={"k": 7.0}, source="test:conv_backend=cudnn") + with pytest.raises(LookupError, match=r"'test\.estimator' version 2 on conv_backend=miopen\."): + resolve_coefficients(_data(_entry(spec.id, version=2, values={"k": 7.0})), spec, Platform("miopen"), "test") + with pytest.raises(LookupError, match=r"'test\.estimator' version 2 on conv_backend=cudnn \(found only version"): + resolve_coefficients(_data(_entry(spec.id, version=1, values={"k": 7.0})), spec, Platform("cudnn"), "test") + with pytest.raises(LookupError, match=r"declared 2 times"): + resolve_coefficients( + _data(_entry(spec.id, version=2, values={"k": 7.0}), _entry(spec.id, version=2, values={"k": 8.0})), + spec, + Platform("cudnn"), + "test", + ) + + +def test_borrowed_values_resolve_one_hop_and_the_source_names_both_borrowings() -> None: + data = _data(_entry("test.lender", values={"k": 7.0}), _entry("test.borrower", values_from="test.lender")) + + on_cudnn = resolve_coefficients(data, _spec("test.borrower"), Platform("cudnn"), origin="test") + assert on_cudnn == ResolvedCoefficients(values={"k": 7.0}, source="test:conv_backend=cudnn (values of test.lender)") + on_xpu = resolve_coefficients(data, _spec("test.borrower"), Platform("xpu"), origin="test") + assert on_xpu == ResolvedCoefficients( + values={"k": 7.0}, source="test:conv_backend=xpu (cudnn values of test.lender)" + ) + + +@pytest.mark.parametrize( + "entries, message", + [ + pytest.param( + (_entry("a", values_from="b"), _entry("b", values_from="c"), _entry("c", values={"k": 1.0})), + "borrows in turn", + id="chain", + ), + pytest.param((_entry("a", values_from="b"), _entry("b", values_from="a")), "borrows in turn", id="cycle"), + pytest.param( + (_entry("a", values_from="b"),), r"No working-memory coefficients for estimator 'b'", id="no-lender" + ), + pytest.param( + (_entry("a", values_from="b"), _entry("b", values={"other": 1.0})), r"provide \['other'\]", id="mismatch" + ), + ], +) +def test_a_borrowing_that_is_not_one_hop_to_matching_values_is_refused(entries, message) -> None: + with pytest.raises(LookupError, match=message): + resolve_coefficients(_data(*entries), _spec("a"), Platform("cudnn"), origin="test") + + +@pytest.mark.parametrize("values, values_from", [({"k": 1.0}, "b"), (None, None)], ids=["both", "neither"]) +def test_an_entry_carries_exactly_one_of_values_and_values_from(values, values_from) -> None: + with pytest.raises(ValidationError, match="exactly one of"): + _entry("a", values=values, values_from=values_from) + + +def test_a_data_directory_without_data_files_is_named(tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError, match=re.escape(str(tmp_path))): + load_data(tmp_path) + with pytest.raises(FileNotFoundError, match=re.escape(str(tmp_path / "missing"))): + load_data(tmp_path / "missing") + (tmp_path / "broken.json").write_text('{"schema_version": 1}', encoding="utf-8") + with pytest.raises(ValueError, match="broken.json"): + load_data(tmp_path) + + +@pytest.mark.platform_independent +def test_every_module_that_declares_an_estimator_is_listed() -> None: + """The registry fills as modules import, so a spec declared outside `ESTIMATOR_MODULES` is missing from every + check above until something else happens to import it -- and a ROCm user is the first to find its gap.""" + package = Path(invokeai.__file__).parent + declaring = sorted( + ".".join(path.relative_to(package.parent).with_suffix("").parts) + for owner in ("app", "backend") + for path in (package / owner).rglob("*.py") + if re.search(r"\bEstimatorSpec\(", path.read_text(encoding="utf-8")) + ) + assert declaring, "found no estimator modules -- the scan is broken" + unlisted = sorted(set(declaring) - set(ESTIMATOR_MODULES)) + stale = sorted(set(ESTIMATOR_MODULES) - set(declaring)) + assert not unlisted and not stale, f"add to ESTIMATOR_MODULES: {unlisted}; declares no spec: {stale}" diff --git a/tests/backend/working_memory/test_spec_registry.py b/tests/backend/working_memory/test_spec_registry.py new file mode 100644 index 00000000000..0c323aa179a --- /dev/null +++ b/tests/backend/working_memory/test_spec_registry.py @@ -0,0 +1,72 @@ +"""Registry mechanics, against throwaway specs rather than the shipped ones.""" + +import pytest + +from invokeai.backend.working_memory import spec as spec_module +from invokeai.backend.working_memory.spec import EstimatorSpec, register, registered_specs + + +@pytest.fixture(autouse=True) +def isolated_registry(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(spec_module, "_SPECS", {}) + + +def _module_terms(): + """A fresh `terms` function on every call, with one module and qualified name: what `importlib.reload` of its + defining module produces.""" + + def terms(features, values): + return () + + return terms + + +def _another_estimators_terms(features, values): + return () + + +def _spec(terms, version: int = 1) -> EstimatorSpec: + return EstimatorSpec( + id="test.estimator", + version=version, + coefficients=frozenset({"k"}), + platform_keys=("conv_backend",), + terms=terms, + ) + + +def test_re_executing_the_defining_module_adopts_its_new_terms() -> None: + """`importlib.reload`: the same declaration from the re-executed function.""" + before, after = _module_terms(), _module_terms() + register(_spec(before)) + register(_spec(after)) + + assert registered_specs()["test.estimator"].terms is after + + +def test_a_differing_declaration_under_a_taken_id_raises() -> None: + register(_spec(_module_terms())) + with pytest.raises(ValueError, match="already registered"): + register(_spec(_module_terms(), version=2)) + + assert registered_specs()["test.estimator"].version == 1, "the first one stands" + + +def test_the_same_declaration_from_another_function_raises() -> None: + """`terms` is not part of a spec's equality, so this is the case equality alone would let through.""" + first = _module_terms() + register(_spec(first)) + with pytest.raises(ValueError, match="ids are unique"): + register(_spec(_another_estimators_terms)) + + assert registered_specs()["test.estimator"].terms is first + + +def test_the_same_function_registered_twice_raises() -> None: + """Two specs of one module sharing a `terms` function -- as SD3 and CogView4 do -- under one id. `importlib.reload` + makes new function objects, so it never hands over the same one twice; jurigged re-running an edited `register` + statement does, and raises too (see `register`).""" + terms = _module_terms() + register(_spec(terms)) + with pytest.raises(ValueError, match="ids are unique"): + register(_spec(terms)) From 470db713f551b2826d569a42520a22f55fec2542 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 17 Oct 2026 15:23:25 +0200 Subject: [PATCH 2/9] feat(working-memory): measure working memory in admin calibration runs 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. --- .../Architecture/model-manager.mdx | 64 ++ invokeai/app/api/dependencies.py | 9 + invokeai/app/api/routers/session_queue.py | 12 + invokeai/app/api/routers/working_memory.py | 224 +++++ invokeai/app/api_app.py | 2 + invokeai/app/services/invocation_services.py | 5 + .../session_processor_default.py | 87 +- .../working_memory_calibration/__init__.py | 0 .../working_memory_calibration_base.py | 54 ++ .../working_memory_calibration_common.py | 265 ++++++ .../working_memory_calibration_default.py | 643 +++++++++++++++ .../working_memory_calibration_store.py | 249 ++++++ .../backend/model_manager/load/load_base.py | 76 +- .../load/model_cache/model_cache.py | 9 + invokeai/backend/util/devices.py | 38 +- invokeai/backend/util/nvml.py | 146 ++++ invokeai/backend/util/pci.py | 17 + invokeai/backend/util/wddm.py | 89 +- invokeai/backend/working_memory/counters.py | 333 ++++++++ .../backend/working_memory/fingerprint.py | 162 ++++ invokeai/backend/working_memory/probe.py | 729 ++++++++++++++++ invokeai/frontend/api/openapi.json | 761 +++++++++++++++++ invokeai/frontend/api/schema.ts | 589 +++++++++++++ tests/app/routers/conftest.py | 1 + .../test_working_memory_calibration.py | 244 ++++++ .../test_working_memory_calibration_runner.py | 379 +++++++++ ...test_working_memory_calibration_service.py | 366 +++++++++ .../test_model_cache_locks_held.py | 48 ++ tests/backend/util/test_nvml.py | 111 +++ tests/backend/util/test_wddm.py | 47 +- tests/backend/working_memory/test_counters.py | 203 +++++ .../working_memory/test_fingerprint.py | 119 +++ tests/backend/working_memory/test_probe.py | 775 ++++++++++++++++++ tests/fixtures/working_memory_calibration.py | 86 ++ 34 files changed, 6898 insertions(+), 44 deletions(-) create mode 100644 invokeai/app/api/routers/working_memory.py create mode 100644 invokeai/app/services/working_memory_calibration/__init__.py create mode 100644 invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py create mode 100644 invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py create mode 100644 invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py create mode 100644 invokeai/app/services/working_memory_calibration/working_memory_calibration_store.py create mode 100644 invokeai/backend/util/nvml.py create mode 100644 invokeai/backend/util/pci.py create mode 100644 invokeai/backend/working_memory/counters.py create mode 100644 invokeai/backend/working_memory/fingerprint.py create mode 100644 invokeai/backend/working_memory/probe.py create mode 100644 tests/app/routers/test_working_memory_calibration.py create mode 100644 tests/app/services/session_processor/test_working_memory_calibration_runner.py create mode 100644 tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py create mode 100644 tests/backend/model_manager/load/model_cache/test_model_cache_locks_held.py create mode 100644 tests/backend/util/test_nvml.py create mode 100644 tests/backend/working_memory/test_counters.py create mode 100644 tests/backend/working_memory/test_fingerprint.py create mode 100644 tests/backend/working_memory/test_probe.py create mode 100644 tests/fixtures/working_memory_calibration.py diff --git a/docs/src/content/docs/development/Architecture/model-manager.mdx b/docs/src/content/docs/development/Architecture/model-manager.mdx index 2dfa5d93010..1b8da502c1a 100644 --- a/docs/src/content/docs/development/Architecture/model-manager.mdx +++ b/docs/src/content/docs/development/Architecture/model-manager.mdx @@ -992,6 +992,70 @@ 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 or held one when it opened; 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; 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`, which also voids the scope's later windows when it hit the lock bookkeeping); 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 (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`, and sets the origin to `invokeai:wm-calibration:` and the destination to `invokeai:wm-calibration`. It records the run, then enqueues the batch as the requesting admin. + - With a key, the run id is derived from the owner and the key, so a repeated request rebuilds the identical batch. The queue's enqueue receipt then returns the same items, and the response is the same run; nothing new is enqueued. Reusing the key for a different request is a 409 and leaves the original run in place. Acknowledge the key through `POST /api/v1/queue/{queue_id}/enqueue_batch/acknowledge`. + - 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. 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. Windows from a device of another class, or of a run that was removed, are dropped with a warning. 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 `/working_memory/.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. + +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 run is `running` while any item is pending or running, `failed` when any failed, and `completed` when all completed. Otherwise it is `partial` if points were recorded, or `canceled` if not. + +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; 409 no GPU or several generation devices, unusable store file, full queue (the run is discarded), or a retry key used for another request; 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 | 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: diff --git a/invokeai/app/api/dependencies.py b/invokeai/app/api/dependencies.py index f11b206427e..98d63dbc7c2 100644 --- a/invokeai/app/api/dependencies.py +++ b/invokeai/app/api/dependencies.py @@ -70,6 +70,10 @@ from invokeai.app.services.wildcard_records.wildcard_records_default import WildcardRecordsStorage from invokeai.app.services.workflow_records.workflow_records_default import WorkflowRecordsStorage from invokeai.app.services.workflow_thumbnails.workflow_thumbnails_disk import WorkflowThumbnailFileStorageDisk +from invokeai.app.services.working_memory_calibration.working_memory_calibration_default import ( + WorkingMemoryCalibrationService, +) +from invokeai.app.services.working_memory_calibration.working_memory_calibration_store import CalibrationStore from invokeai.backend.architectures import conditioning_safe_globals from invokeai.backend.architectures import validate as validate_architectures from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData @@ -216,6 +220,10 @@ def initialize( image_index_records = ImageIndexRecords(database) image_index = ImageIndexService() intermediates = IntermediatesService(records=IntermediatesRecords(database), logger=logger) + # Machine data next to the database rather than in it: see working_memory_calibration_store. + working_memory_calibration = WorkingMemoryCalibrationService( + CalibrationStore(configuration.db_path.parent / "working_memory") + ) fonts = FontService( database, fonts_dir=configuration.fonts_path, @@ -273,6 +281,7 @@ def initialize( fonts=fonts, intermediates=intermediates, gallery_maintenance=gallery_maintenance, + working_memory_calibration=working_memory_calibration, ) # Constructing the Invoker starts every service, including the session diff --git a/invokeai/app/api/routers/session_queue.py b/invokeai/app/api/routers/session_queue.py index c64f32623b1..03c60740148 100644 --- a/invokeai/app/api/routers/session_queue.py +++ b/invokeai/app/api/routers/session_queue.py @@ -39,6 +39,10 @@ from invokeai.app.services.shared.graph import Graph, GraphExecutionState from invokeai.app.services.shared.pagination import SQLiteDirection from invokeai.app.services.video_records.video_records_common import VideoRecordNotFoundException +from invokeai.app.services.working_memory_calibration.working_memory_calibration_common import ( + CALIBRATION_ORIGIN_PREFIX, + is_reserved_origin, +) session_queue_router = APIRouter(prefix="/v1/queue", tags=["queue"]) @@ -254,6 +258,14 @@ async def enqueue_batch( prepend: bool = Body(default=False, description="Whether or not to prepend this batch in the queue"), ) -> EnqueueBatchResult: """Processes a batch and enqueues the output graphs for execution for the current user.""" + # Calibration items are measured and admitted as such by the session runner; only the calibration route may + # create them, for administrators. + if is_reserved_origin(batch.origin): + raise HTTPException( + status_code=400, + detail=f"Origins starting with '{CALIBRATION_ORIGIN_PREFIX}' are reserved for working-memory calibration " + "runs, which start through POST /api/v2/models/i/{key}/working_memory/calibrate.", + ) await asyncio.to_thread(assert_image_move_maintenance_inactive) try: diff --git a/invokeai/app/api/routers/working_memory.py b/invokeai/app/api/routers/working_memory.py new file mode 100644 index 00000000000..afed1028611 --- /dev/null +++ b/invokeai/app/api/routers/working_memory.py @@ -0,0 +1,224 @@ +"""Working-memory calibration runs: administrators measure what a model's operations really take on this machine's GPU. + +Mounted under the model manager's prefix: a run calibrates one model (`/i/{key}/working_memory/...`), and its +measurements belong to the machine's device class (`/working_memory/...`). Every route is admin-only: runs spend GPU +time on the instance's behalf, and the store is machine data that names the accounts that created runs. +""" + +import asyncio +from typing import Final, Optional + +from fastapi import Body, HTTPException, Path, Query, Response +from fastapi.routing import APIRouter + +from invokeai.app.api.auth_dependencies import AdminUserOrDefault +from invokeai.app.api.dependencies import ApiDependencies +from invokeai.app.api.routers.image_move_maintenance import assert_image_move_maintenance_inactive +from invokeai.app.services.model_records import UnknownModelException +from invokeai.app.services.session_queue.session_queue_common import ( + EnqueueIdempotencyConflictError, + EnqueueReceiptLimitError, +) +from invokeai.app.services.working_memory_calibration.working_memory_calibration_base import ( + WorkingMemoryCalibrationServiceBase, +) +from invokeai.app.services.working_memory_calibration.working_memory_calibration_common import ( + CalibrationGraphError, + CalibrationQueueFullError, + CalibrationRunNotFoundError, + CalibrationStoreError, + CalibrationStoreUnusableError, + CalibrationUnavailableError, + WorkingMemoryCalibrationRequest, + WorkingMemoryCalibrationResetResult, + WorkingMemoryCalibrationRunCreated, + WorkingMemoryCalibrationStatus, +) + +working_memory_router = APIRouter(prefix="/v2/models", tags=["model_manager"]) + +_DEVICE_CLASS_ID = Query( + default=None, + pattern=r"^[0-9a-f]{16}$", + description="The device class to act on; the current generation device's by default", +) + + +# One description per error, shared by the routes that return it. +_STORE_UNREADABLE: Final = { + "description": "The calibration store cannot be read or written (unwritable root, lock timeout)" +} +_NO_DEVICE_CLASS: Final = { + "description": "No device_class_id was given and there is no current one (no GPU, or several generation " + "devices), or the device class's store file is unusable and kept as it is" +} + + +def _service() -> WorkingMemoryCalibrationServiceBase: + service = ApiDependencies.invoker.services.working_memory_calibration + if service is None: + raise HTTPException(status_code=503, detail="Working-memory calibration is not available") + return service + + +def _http_error(error: Exception) -> HTTPException: + if isinstance(error, (CalibrationRunNotFoundError, UnknownModelException)): + return HTTPException(status_code=404, detail=str(error)) + if isinstance(error, CalibrationGraphError): + return HTTPException(status_code=422, detail=str(error)) + if isinstance( + error, + ( + CalibrationUnavailableError, + CalibrationStoreUnusableError, + CalibrationQueueFullError, + EnqueueIdempotencyConflictError, + ), + ): + return HTTPException(status_code=409, detail=str(error)) + if isinstance(error, EnqueueReceiptLimitError): + return HTTPException(status_code=429, detail=str(error)) + # CalibrationStoreError: the root cannot be read or written (read-only, lock timeout). Also the 503 of `_service`. + return HTTPException(status_code=503, detail=str(error)) + + +@working_memory_router.post( + "/i/{key}/working_memory/calibrate", + operation_id="calibrate_model_working_memory", + status_code=201, + response_model=WorkingMemoryCalibrationRunCreated, + responses={ + 404: {"description": "The model could not be found"}, + 409: { + "description": "This machine cannot calibrate now (no GPU, or several generation devices), the device " + "class's store file is unusable and kept as it is, the queue is full, or the idempotency key was used for " + "another request" + }, + 422: {"description": "The graph cannot be measured as a calibration run"}, + 429: {"description": "Too many unacknowledged idempotency keys; acknowledge or wait, then retry"}, + 503: _STORE_UNREADABLE, + }, +) +async def calibrate_model_working_memory( + current_admin: AdminUserOrDefault, + key: str = Path(description="Key of the model the run calibrates"), + request: WorkingMemoryCalibrationRequest = Body(description="The batch to measure"), +) -> WorkingMemoryCalibrationRunCreated: + """Start a working-memory calibration run: enqueue the batch with every node uncached and measure each node's + model operations on the generation device. Needs exactly one generation device, and it must be a GPU. + + With an `idempotency_key`, repeating the request (after a lost response) returns the run it created and enqueues + nothing new; acknowledge the key through the queue's enqueue_batch/acknowledge route once the response arrived.""" + await asyncio.to_thread(assert_image_move_maintenance_inactive) + try: + return await _service().create_run(key, request, current_admin.user_id) + except ( + UnknownModelException, + CalibrationGraphError, + CalibrationUnavailableError, + CalibrationStoreError, + CalibrationQueueFullError, + EnqueueIdempotencyConflictError, + EnqueueReceiptLimitError, + ) as error: + raise _http_error(error) from error + + +@working_memory_router.get( + "/working_memory/status", + operation_id="get_working_memory_calibration_status", + response_model=WorkingMemoryCalibrationStatus, + responses={503: _STORE_UNREADABLE}, +) +def get_working_memory_calibration_status(current_admin: AdminUserOrDefault) -> WorkingMemoryCalibrationStatus: + """Every device class's calibration runs, with their derived status, point counts and environments.""" + try: + return _service().get_status() + except CalibrationStoreError as error: + raise _http_error(error) from error + + +@working_memory_router.delete( + "/working_memory/runs/{run_id}", + operation_id="delete_working_memory_calibration_run", + response_model=WorkingMemoryCalibrationResetResult, + responses={ + 404: {"description": "No usable store file holds the run"}, + 409: {"description": "The run's store file became unusable after the run was found; it is kept as it is"}, + 503: _STORE_UNREADABLE, + }, +) +def delete_working_memory_calibration_run( + current_admin: AdminUserOrDefault, + run_id: str = Path(description="The calibration run to remove"), +) -> WorkingMemoryCalibrationResetResult: + """Remove a calibration run and its measurements. Its queue items that have not run yet fail when they start.""" + try: + return _service().reset_run(run_id) + except (CalibrationRunNotFoundError, CalibrationStoreError) as error: + raise _http_error(error) from error + + +@working_memory_router.delete( + "/i/{key}/working_memory", + operation_id="reset_model_working_memory_calibration", + response_model=WorkingMemoryCalibrationResetResult, + responses={409: _NO_DEVICE_CLASS, 503: _STORE_UNREADABLE}, +) +def reset_model_working_memory_calibration( + current_admin: AdminUserOrDefault, + key: str = Path(description="Key of the model whose runs and measurements to remove"), + device_class_id: Optional[str] = _DEVICE_CLASS_ID, +) -> WorkingMemoryCalibrationResetResult: + """Remove a model's calibration runs with all their measurements, and every measurement of the model taken in + other runs (a VAE measured inside a transformer's run, say), from one device class.""" + try: + return _service().reset_model(key, device_class_id) + except (CalibrationUnavailableError, CalibrationStoreError) as error: + raise _http_error(error) from error + + +@working_memory_router.delete( + "/working_memory", + operation_id="reset_working_memory_calibration", + response_model=WorkingMemoryCalibrationResetResult, + responses={409: _NO_DEVICE_CLASS, 503: _STORE_UNREADABLE}, +) +def reset_working_memory_calibration( + current_admin: AdminUserOrDefault, + device_class_id: Optional[str] = _DEVICE_CLASS_ID, +) -> WorkingMemoryCalibrationResetResult: + """Remove every calibration run and measurement of one device class.""" + try: + return _service().reset_all(device_class_id) + except (CalibrationUnavailableError, CalibrationStoreError) as error: + raise _http_error(error) from error + + +@working_memory_router.get( + "/working_memory/export", + operation_id="export_working_memory_calibration", + response_class=Response, + responses={ + 200: { + "description": "The device class's store file byte for byte: a CalibrationStoreFile (schema_version 1) " + "when this version can read it, otherwise whatever the file holds", + "content": {"application/json": {}}, + }, + 404: {"description": "The device class has no store file"}, + 409: {"description": "No device_class_id was given and there is no current one"}, + 503: _STORE_UNREADABLE, + }, +) +def export_working_memory_calibration( + current_admin: AdminUserOrDefault, + device_class_id: Optional[str] = _DEVICE_CLASS_ID, +) -> Response: + """The device class's calibration store file as it is on disk, readable by this version or not.""" + try: + raw = _service().export(device_class_id) + except (CalibrationUnavailableError, CalibrationStoreError) as error: + raise _http_error(error) from error + if raw is None: + raise HTTPException(status_code=404, detail="This device class has no calibration store") + return Response(content=raw, media_type="application/json") diff --git a/invokeai/app/api_app.py b/invokeai/app/api_app.py index ee0b722fe0f..42077f88bbe 100644 --- a/invokeai/app/api_app.py +++ b/invokeai/app/api_app.py @@ -52,6 +52,7 @@ virtual_boards, wildcards, workflows, + working_memory, ) from invokeai.app.api.sockets import SocketIO from invokeai.app.services.config.config_default import get_config @@ -783,6 +784,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: app.include_router(utilities.utilities_router, prefix="/api") app.include_router(fonts.fonts_router, prefix="/api") app.include_router(model_manager.model_manager_router, prefix="/api") +app.include_router(working_memory.working_memory_router, prefix="/api") app.include_router(image_moves.image_moves_router, prefix="/api") app.include_router(images.images_router, prefix="/api") app.include_router(image_map.image_map_router, prefix="/api") diff --git a/invokeai/app/services/invocation_services.py b/invokeai/app/services/invocation_services.py index 41de3b8d704..5e4099897f8 100644 --- a/invokeai/app/services/invocation_services.py +++ b/invokeai/app/services/invocation_services.py @@ -55,6 +55,9 @@ from invokeai.app.services.videos.videos_base import VideoServiceABC from invokeai.app.services.workflow_records.workflow_records_base import WorkflowRecordsStorageBase from invokeai.app.services.workflow_thumbnails.workflow_thumbnails_base import WorkflowThumbnailServiceBase + from invokeai.app.services.working_memory_calibration.working_memory_calibration_base import ( + WorkingMemoryCalibrationServiceBase, + ) from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData @@ -116,6 +119,7 @@ def __init__( progress_previews: "ProgressPreviewsBase | None" = None, fonts: "FontService | None" = None, gallery_maintenance: "GalleryMaintenanceService | None" = None, + working_memory_calibration: "WorkingMemoryCalibrationServiceBase | None" = None, ): self.board_images = board_images self.board_image_records = board_image_records @@ -168,3 +172,4 @@ def __init__( self.fonts = fonts self.intermediates = intermediates self.gallery_maintenance = gallery_maintenance + self.working_memory_calibration = working_memory_calibration diff --git a/invokeai/app/services/session_processor/session_processor_default.py b/invokeai/app/services/session_processor/session_processor_default.py index 4308e7a4f35..920b16994a2 100644 --- a/invokeai/app/services/session_processor/session_processor_default.py +++ b/invokeai/app/services/session_processor/session_processor_default.py @@ -45,6 +45,7 @@ ) from invokeai.app.services.shared.graph import CollectInvocation, IterateInvocation, NodeInputError from invokeai.app.services.shared.invocation_context import InvocationContextData, build_invocation_context +from invokeai.app.services.working_memory_calibration.working_memory_calibration_common import calibration_run_id from invokeai.app.util.profiler import Profiler from invokeai.backend.util.device_pool import ( GENERATION_DEVICE_POOL, @@ -52,6 +53,7 @@ ) from invokeai.backend.util.devices import TorchDevice, disable_conv_benchmark_empty_cache from invokeai.backend.util.wddm import paged_bytes +from invokeai.backend.working_memory import probe # Paged-out VRAM worth a warning. Weight uploads page a few hundred MB for an instant (0.25-0.36 GiB measured on an # RX 9060 XT), so the threshold sits above that; a session that ends with more than this paged has lost real VRAM. @@ -165,6 +167,8 @@ def __init__( self.workflow_call_queue_lifecycle = WorkflowCallQueueLifecycle(self) # Per runner, hence per worker and device: each worker remembers the architecture it last ran. self._architecture_switch = ArchitectureSwitchOffload() + # The calibration run of the item this runner is executing, None for every other item. + self._calibration_run_id: Optional[str] = None def start(self, services: InvocationServices, cancel_event: ThreadEvent, profiler: Optional[Profiler] = None): self._services = services @@ -203,7 +207,10 @@ def _run_session_loop(self, queue_item: SessionQueueItem) -> None: if self._cancel_if_owner_revoked(queue_item): break - self.run_node(invocation, queue_item) + if self._calibration_run_id is None: + self.run_node(invocation, queue_item) + elif not self._run_calibration_node(invocation, queue_item, self._calibration_run_id): + break # The session is complete if all invocations have been run or there is an error on the session. # At this time, the queue item may be canceled, but the object itself here won't be updated yet. We must @@ -246,10 +253,88 @@ def _cancel_if_owner_revoked(self, queue_item: SessionQueueItem, *, unreadable_i def run(self, queue_item: SessionQueueItem): # Exceptions raised outside `run_node` are handled by the processor. There is no need to catch them here. + # getattr, like the workflow-call lifecycle's reads: runners also execute stand-in items without an origin. + run_id = calibration_run_id(getattr(queue_item, "origin", None)) + if run_id is not None: + self._run_calibration_item(queue_item, run_id) + return self._on_before_run_session(queue_item=queue_item) self._run_session_loop(queue_item) self._on_after_run_session(queue_item=queue_item) + def _run_calibration_item(self, queue_item: SessionQueueItem, run_id: str) -> None: + """Run a working-memory calibration item, measuring each node (see `working_memory.probe`). + + The item is admitted again here, whatever admitted it to the queue -- a retry and a workflow-call child carry + the origin and owner of the item they came from: its owner must still be an administrator, the machine must + still calibrate on one GPU, and its run must still exist. Otherwise it fails before any node runs. + """ + service = self._services.working_memory_calibration + raised: Optional[BaseException] = None + try: + try: + refusal = ( + "Working-memory calibration is not available in this process." + if service is None + else service.admit(queue_item, run_id) + ) + except Exception as error: + self._services.logger.warning(f"Could not admit calibration item {queue_item.item_id}", exc_info=True) + refusal = f"Could not admit the calibration item: {error}" + if refusal is not None: + self._services.logger.warning(f"Failing calibration item {queue_item.item_id}: {refusal}") + with suppress(SessionQueueItemNotFoundError): + self._services.session_queue.fail_queue_item( + queue_item.item_id, "WorkingMemoryCalibrationError", refusal, refusal + ) + return + self._calibration_run_id = run_id + try: + self._on_before_run_session(queue_item=queue_item) + self._run_session_loop(queue_item) + self._on_after_run_session(queue_item=queue_item) + finally: + self._calibration_run_id = None + except BaseException as error: + raised = error + raise + finally: + # Also when the session raised: `admit` recorded the item as running, and only this replaces that. + if service is not None: + try: + service.observe(queue_item, run_id, raised) + except Exception: + self._services.logger.warning( + f"Could not record the outcome of calibration item {queue_item.item_id}", exc_info=True + ) + + def _run_calibration_node(self, invocation: BaseInvocation, queue_item: SessionQueueItem, run_id: str) -> bool: + """Run one node of a calibration item inside a probe scope and persist its windows, whether the node completed, + failed or was canceled. Returns False when they could not be persisted: the item has failed then.""" + with probe.calibration_scope( + run_id=run_id, + queue_item_id=queue_item.item_id, + # The source id: what the client that built the graph knows the node by. + node_id=queue_item.session.prepared_source_mapping.get(invocation.id, invocation.id), + node_type=invocation.get_type(), + ) as scope: + self.run_node(invocation, queue_item) + if not scope.windows: + return True + service = self._services.working_memory_calibration + try: + assert service is not None + service.record_windows(run_id, scope.windows) + except Exception as error: + message = f"Could not save the working-memory measurements of node {invocation.id}: {error}" + self._services.logger.error(message, exc_info=True) + with suppress(SessionQueueItemNotFoundError): + self._services.session_queue.fail_queue_item( + queue_item.item_id, "WorkingMemoryCalibrationError", message, traceback.format_exc() + ) + return False + return True + def run_node(self, invocation: BaseInvocation, queue_item: SessionQueueItem): try: # Any unhandled exception in this scope is an invocation error & will fail the graph diff --git a/invokeai/app/services/working_memory_calibration/__init__.py b/invokeai/app/services/working_memory_calibration/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py new file mode 100644 index 00000000000..d36e4a0be1c --- /dev/null +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py @@ -0,0 +1,54 @@ +from abc import ABC, abstractmethod +from typing import Optional + +from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem +from invokeai.app.services.working_memory_calibration.working_memory_calibration_common import ( + WorkingMemoryCalibrationRequest, + WorkingMemoryCalibrationResetResult, + WorkingMemoryCalibrationRunCreated, + WorkingMemoryCalibrationStatus, +) +from invokeai.backend.working_memory.probe import MeasuredWindow + + +class WorkingMemoryCalibrationServiceBase(ABC): + """Calibration runs: enqueues them, admits their queue items, and keeps their measurements per device class.""" + + @abstractmethod + async def create_run( + self, model_key: str, request: WorkingMemoryCalibrationRequest, user_id: str + ) -> WorkingMemoryCalibrationRunCreated: + """Validate the request, record the run and enqueue its batch as `user_id`, who must be an administrator.""" + + @abstractmethod + def admit(self, queue_item: SessionQueueItem, run_id: str) -> Optional[str]: + """Called by the session runner before a calibration item runs: None to run it, or why it must fail.""" + + @abstractmethod + def record_windows(self, run_id: str, windows: list[MeasuredWindow]) -> None: + """Persist one node's measured windows. Windows of a run that no longer exists are dropped.""" + + @abstractmethod + def observe(self, queue_item: SessionQueueItem, run_id: str, error: Optional[BaseException] = None) -> None: + """Record the item's queue status in its run, after the session runner is done with it. `error` is what the + runner raised, which fails an item the queue still has as running; a removed item counts as canceled.""" + + @abstractmethod + def get_status(self) -> WorkingMemoryCalibrationStatus: + """Every device class's runs, with derived status, point counts and environments.""" + + @abstractmethod + def reset_run(self, run_id: str) -> WorkingMemoryCalibrationResetResult: + """Remove a run and its points.""" + + @abstractmethod + def reset_model(self, model_key: str, device_class_id: Optional[str] = None) -> WorkingMemoryCalibrationResetResult: + """Remove a model's runs and points from one device class, the current one by default.""" + + @abstractmethod + def reset_all(self, device_class_id: Optional[str] = None) -> WorkingMemoryCalibrationResetResult: + """Remove every run and point of one device class, the current one by default.""" + + @abstractmethod + def export(self, device_class_id: Optional[str] = None) -> Optional[bytes]: + """The device class's store file as it is on disk, the current one by default; None when it has none.""" diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py new file mode 100644 index 00000000000..e7d852aa6c9 --- /dev/null +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py @@ -0,0 +1,265 @@ +"""Working-memory calibration runs: what is stored per device class, and what the API returns about it. + +A calibration run is a batch an administrator enqueues to measure one model's working memory. Its queue items carry the +reserved origin `CALIBRATION_ORIGIN_PREFIX` + run id; the session runner measures every node of such an item with the +working-memory probe and hands the windows to the calibration service, which keeps them in one JSON file per device +class (`CalibrationStoreFile`). Recording only: nothing here changes an estimate. +""" + +import datetime +from typing import Final, Literal, Optional + +from pydantic import BaseModel, ConfigDict, Field, model_validator + +from invokeai.app.services.session_queue.session_queue_common import QUEUE_ITEM_STATUS, BatchDataCollection +from invokeai.app.services.shared.graph import Graph +from invokeai.backend.working_memory.counters import Backend +from invokeai.backend.working_memory.fingerprint import DeviceClass, Environment +from invokeai.backend.working_memory.probe import MeasuredWindow + +CALIBRATION_ORIGIN_PREFIX: Final = "invokeai:wm-calibration:" +"""Queue items whose origin starts with this are calibration items. Only the calibration service may enqueue them.""" +CALIBRATION_DESTINATION: Final = "invokeai:wm-calibration" +"""The destination of calibration items: not a gallery or canvas destination, so no client routes their results.""" +STORE_SCHEMA_VERSION: Final = 1 + + +def calibration_origin(run_id: str) -> str: + return CALIBRATION_ORIGIN_PREFIX + run_id + + +def calibration_run_id(origin: Optional[str]) -> Optional[str]: + """The run id a queue item's origin names, or None for an item that is not a calibration item.""" + if origin is None or not origin.startswith(CALIBRATION_ORIGIN_PREFIX): + return None + return origin[len(CALIBRATION_ORIGIN_PREFIX) :] + + +def is_reserved_origin(origin: Optional[str]) -> bool: + """Whether a client-supplied origin claims the calibration prefix. Case-insensitive, because the queue's origin + filters are: a look-alike origin would otherwise be counted among a run's items.""" + return origin is not None and origin.casefold().startswith(CALIBRATION_ORIGIN_PREFIX) + + +# region Store +# +# One policy for every store model: unknown fields make the file unusable (`extra="forbid"`). Ignoring them instead +# would drop them on the next rewrite, destroying data a newer writer put there. + + +class _StoreModel(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class CalibrationDeviceClass(_StoreModel): + """A `fingerprint.DeviceClass` with its id as a plain field, as it is stored and returned.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + id: str = Field(description="A short stable hash of the other fields") + vendor: str + name: str + total_mib: int + backend: Backend + arch: Optional[str] + + @classmethod + def of(cls, device_class: DeviceClass) -> "CalibrationDeviceClass": + return cls(id=device_class.id, **device_class.model_dump(exclude={"id"})) + + @model_validator(mode="after") + def _id_matches(self) -> "CalibrationDeviceClass": + if DeviceClass(**self.model_dump(exclude={"id"})).id != self.id: + raise ValueError(f"Device class id {self.id} does not match its fields") + return self + + +class CalibrationEnvironment(_StoreModel): + """A `fingerprint.Environment` with its id as a plain field.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + id: str = Field(description="A short stable hash of the hard facts") + hard: dict[str, str] + soft: dict[str, Optional[str]] + + @classmethod + def of(cls, env: Environment) -> "CalibrationEnvironment": + return cls(id=env.id, hard=env.hard, soft=env.soft) + + @model_validator(mode="after") + def _id_matches(self) -> "CalibrationEnvironment": + if Environment(hard=self.hard, soft=self.soft).id != self.id: + raise ValueError(f"Environment id {self.id} does not match its hard facts") + return self + + +class CalibrationPoint(MeasuredWindow): + """One measured window, keyed to the environment it ran under. Forbids unknown fields, as `MeasuredWindow` does.""" + + env_id: str + + +class ObservedItem(_StoreModel): + """What the session runner saw of one of the run's queue items. Kept so that a run's status survives the queue + pruning its items.""" + + status: QUEUE_ITEM_STATUS + error_message: Optional[str] = None + retried_from_item_id: Optional[int] = None + observed_at: datetime.datetime + + +class CalibrationRun(_StoreModel): + run_id: str + model_key: str + model_hash: Optional[str] = Field(description="The model's hash when the run was created; provenance only") + user_id: str + created_at: datetime.datetime + env_id: str = Field(description="The environment the run was created under") + queue_item_ids: list[int] = Field(default_factory=list, description="The items the run was enqueued with") + observed: dict[int, ObservedItem] = Field( + default_factory=dict, description="The last status the session runner saw per item, retries included" + ) + + +class CalibrationStoreFile(_StoreModel): + """The content of `/working_memory/.json`, which the export returns as it is.""" + + schema_version: Literal[1] + device_class: CalibrationDeviceClass + runs: list[CalibrationRun] = Field(default_factory=list) + points: list[CalibrationPoint] = Field(default_factory=list) + envs: list[CalibrationEnvironment] = Field(default_factory=list) + + +# endregion + +# region API + + +class WorkingMemoryCalibrationRequest(BaseModel): + """The batch a calibration run enqueues, built by the client. The server sets origin, destination and runs, and + turns the node cache off.""" + + model_config = ConfigDict(extra="forbid") + + graph: Graph = Field(description="The graph to measure") + data: Optional[BatchDataCollection] = Field(default=None, description="The batch data collection") + idempotency_key: Optional[str] = Field( + default=None, + min_length=1, + max_length=255, + description="A caller-owned retry key. Repeating a request with it returns the run it created and enqueues " + "nothing new; reusing it with a different request is rejected.", + ) + + +class WorkingMemoryCalibrationRunCreated(BaseModel): + run_id: str + device_class_id: str + item_ids: list[int] = Field(description="The queue items the run enqueued") + + +CalibrationRunStatus = Literal["running", "completed", "failed", "partial", "canceled"] + + +class WorkingMemoryCalibrationRunSummary(BaseModel): + run_id: str + model_key: str + model_hash: Optional[str] + user_id: str + created_at: datetime.datetime + env_id: str + status: CalibrationRunStatus = Field( + description="Derived from the run's queue items, a retry standing in for the item it retried: running while " + "any is pending or running; failed when any failed; completed when all completed; otherwise partial when " + "points were recorded, else canceled" + ) + error: Optional[str] = Field(description="The error of a failed item, when the run failed") + item_ids: list[int] = Field(description="Every queue item of the run, retries included") + point_count: int + + +class WorkingMemoryCalibrationPointCount(BaseModel): + estimator: str + version: int + valid: bool + fit_eligible: bool + count: int + + +class WorkingMemoryCalibrationEnvironment(BaseModel): + id: str + hard: dict[str, str] + soft: dict[str, Optional[str]] + hard_matches_current: Optional[bool] = Field( + description="Whether the current environment has the same hard facts; null off the current device class" + ) + soft_matches_current: Optional[bool] = Field( + description="Whether the current environment has the same soft facts; null off the current device class" + ) + point_count: int + + +class WorkingMemoryCalibrationDeviceClass(BaseModel): + device_class_id: str + device_class: Optional[CalibrationDeviceClass] = Field(description="Null when the file is unusable") + usable: bool + problem: Optional[str] = Field(description="Why the file is unusable. It is kept as it is and never written") + is_current: bool + runs: list[WorkingMemoryCalibrationRunSummary] + point_counts: list[WorkingMemoryCalibrationPointCount] + envs: list[WorkingMemoryCalibrationEnvironment] + + +class WorkingMemoryCalibrationStatus(BaseModel): + available: bool = Field( + description="Whether a calibration run can start now: one GPU generation device whose store file is usable" + ) + unavailable_reason: Optional[str] + current_device_class_id: Optional[str] + current_env_id: Optional[str] + device_classes: list[WorkingMemoryCalibrationDeviceClass] + + +class WorkingMemoryCalibrationResetResult(BaseModel): + runs_removed: int + points_removed: int + + +# endregion + +# region Errors + + +class CalibrationUnavailableError(Exception): + """This machine cannot calibrate now: no GPU, or more than one generation device.""" + + +class CalibrationGraphError(ValueError): + """The graph cannot be measured as a calibration run.""" + + +class CalibrationRunNotFoundError(KeyError): + def __init__(self, run_id: str) -> None: + super().__init__(run_id) + self.run_id = run_id + + def __str__(self) -> str: + return f"Calibration run {self.run_id} does not exist" + + +class CalibrationQueueFullError(Exception): + """The queue had no room for the run's items.""" + + +class CalibrationStoreError(Exception): + """The store could not be read or written.""" + + +class CalibrationStoreUnusableError(CalibrationStoreError): + """The device class's file has an unknown schema or is invalid; it is kept untouched and not written.""" + + +# endregion diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py new file mode 100644 index 00000000000..db1a13bdce8 --- /dev/null +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py @@ -0,0 +1,643 @@ +import asyncio +import datetime +import uuid +from collections.abc import Callable, Mapping +from functools import partial +from itertools import chain +from typing import TYPE_CHECKING, Final, Optional + +import torch + +from invokeai.app.invocations.call_saved_workflow import CallSavedWorkflowInvocation +from invokeai.app.services.session_queue.session_queue_common import ( + DEFAULT_QUEUE_ID, + QUEUE_ITEM_STATUS, + Batch, + SessionQueueItem, + SessionQueueItemNotFoundError, +) +from invokeai.app.services.working_memory_calibration.working_memory_calibration_base import ( + WorkingMemoryCalibrationServiceBase, +) +from invokeai.app.services.working_memory_calibration.working_memory_calibration_common import ( + CALIBRATION_DESTINATION, + CALIBRATION_ORIGIN_PREFIX, + STORE_SCHEMA_VERSION, + CalibrationDeviceClass, + CalibrationEnvironment, + CalibrationGraphError, + CalibrationPoint, + CalibrationQueueFullError, + CalibrationRun, + CalibrationRunNotFoundError, + CalibrationRunStatus, + CalibrationStoreFile, + CalibrationUnavailableError, + ObservedItem, + WorkingMemoryCalibrationDeviceClass, + WorkingMemoryCalibrationEnvironment, + WorkingMemoryCalibrationPointCount, + WorkingMemoryCalibrationRequest, + WorkingMemoryCalibrationResetResult, + WorkingMemoryCalibrationRunCreated, + WorkingMemoryCalibrationRunSummary, + WorkingMemoryCalibrationStatus, + calibration_origin, + calibration_run_id, +) +from invokeai.app.services.working_memory_calibration.working_memory_calibration_store import ( + CalibrationStore, + StoreRead, +) +from invokeai.app.util.misc import uuid_string +from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.working_memory import fingerprint +from invokeai.backend.working_memory.counters import backend_for +from invokeai.backend.working_memory.fingerprint import DeviceClass, Environment +from invokeai.backend.working_memory.probe import MeasuredWindow + +if TYPE_CHECKING: + from invokeai.app.services.invocation_services import InvocationServices + from invokeai.app.services.invoker import Invoker + +MEMORY_SETTINGS: Final = ( + "pytorch_cuda_alloc_conf", + "attention_backend", + "attention_type", + "attention_slice_size", + "sequential_guidance", + "pid_memory_optimization", + "wan_memory_optimization", + "rocm_aotriton_experimental", +) +"""The InvokeAI settings that change how much working memory an operation takes; hard facts of a measurement's +environment.""" + +_SUMMARY_CHUNK: Final = 500 +"""Queue items read per summary query when the status collects the calibration items.""" + +_RUN_ID_NAMESPACE: Final = uuid.UUID("3c1f6a52-8d0e-4b6f-9a57-2f0d5c8e41b3") +"""Derives a run id from an owner and a retry key.""" + +_RUNNING: Final[frozenset[QUEUE_ITEM_STATUS]] = frozenset({"pending", "in_progress", "waiting"}) + +SettingValue = str | int | float | bool | None + + +def _now() -> datetime.datetime: + return datetime.datetime.now(datetime.timezone.utc) + + +def _add_env(file: CalibrationStoreFile, env: Environment) -> None: + if all(known.id != env.id for known in file.envs): + file.envs.append(CalibrationEnvironment.of(env)) + + +def _prune_envs(file: CalibrationStoreFile) -> None: + used = {run.env_id for run in file.runs} | {point.env_id for point in file.points} + file.envs = [env for env in file.envs if env.id in used] + + +def _add_points( + file: CalibrationStoreFile, *, run_id: str, env: Environment, points: list[CalibrationPoint] +) -> tuple[bool, bool]: + """Attach `points` to their run; nothing when the file does not hold the run.""" + if all(run.run_id != run_id for run in file.runs): + return False, False + _add_env(file, env) + file.points.extend(points) + return True, True + + +def derive_run_status( + run: CalibrationRun, live: Mapping[int, QUEUE_ITEM_STATUS], point_count: int +) -> tuple[CalibrationRunStatus, Optional[str], list[int]]: + """A run's status, error and items, from its items' statuses: `live` as the queue has them now, the run's own + observations for items the queue no longer has. An item the queue no longer has counts as canceled unless the + runner saw it finish. A retry stands in for the item it retried.""" + item_ids = sorted(set(run.queue_item_ids) | set(run.observed) | set(live)) + + def status_of(item_id: int) -> QUEUE_ITEM_STATUS: + if item_id in live: + return live[item_id] + observed = run.observed.get(item_id) + # Gone from the queue while the runner last saw it running (a crash, a cleared queue): it will not finish. + if observed is None or observed.status in _RUNNING: + return "canceled" + return observed.status + + latest: dict[int, int] = {} + for item_id in item_ids: + observed = run.observed.get(item_id) + lineage = observed.retried_from_item_id if observed and observed.retried_from_item_id else item_id + latest[lineage] = max(latest.get(lineage, item_id), item_id) + effective = sorted(latest.values()) + statuses = {item_id: status_of(item_id) for item_id in effective} + + if not effective: + return "canceled", None, item_ids + if any(status in _RUNNING for status in statuses.values()): + return "running", None, item_ids + failed = [item_id for item_id, status in statuses.items() if status == "failed"] + if failed: + observed = run.observed.get(failed[0]) + error = observed.error_message if observed is not None and observed.error_message else None + return "failed", error or f"Queue item {failed[0]} failed", item_ids + if all(status == "completed" for status in statuses.values()): + return "completed", None, item_ids + return ("partial" if point_count > 0 else "canceled"), None, item_ids + + +class WorkingMemoryCalibrationService(WorkingMemoryCalibrationServiceBase): + """Calibration runs and their store. + + `generation_devices`, `device_class_for` and `environment_for` default to the machine's own; tests pass their own + because a device class can only be read from real hardware. + """ + + def __init__( + self, + store: CalibrationStore, + *, + generation_devices: Optional[Callable[[], list[torch.device]]] = None, + device_class_for: Callable[[torch.device], DeviceClass] = fingerprint.device_class_for, + environment_for: Callable[ + [torch.device, Mapping[str, SettingValue]], Environment + ] = fingerprint.environment_for, + ) -> None: + self._store = store + self._generation_devices = generation_devices + self._device_class_for = device_class_for + self._environment_for = environment_for + # A card's class cannot change while the process runs; reading it touches the driver. + self._device_classes: dict[str, DeviceClass] = {} + # Which file holds a run, so admitting an item and recording its windows do not scan every file. + self._run_files: dict[str, str] = {} + self._invoker: Optional["Invoker"] = None + + def start(self, invoker: "Invoker") -> None: + self._invoker = invoker + + @property + def _services(self) -> "InvocationServices": + assert self._invoker is not None, "The calibration service has not been started" + return self._invoker.services + + # region Devices and environments + + def _devices(self) -> list[torch.device]: + if self._generation_devices is not None: + return list(self._generation_devices()) + configured = TorchDevice.get_generation_devices(self._services.configuration.generation_devices) + return configured or [TorchDevice.choose_torch_device()] + + def _calibration_device(self) -> torch.device: + devices = self._devices() + if len(devices) > 1: + raise CalibrationUnavailableError( + f"Calibration runs on a single generation device, and this machine generates on {len(devices)} " + f"({', '.join(str(device) for device in devices)}). Set generation_devices to one GPU to calibrate it." + ) + if not devices or backend_for(devices[0]) is None: + raise CalibrationUnavailableError( + f"Calibration measures GPU memory, and this machine generates on {devices[0] if devices else 'no device'}." + ) + return devices[0] + + def _device_class(self, device: torch.device) -> DeviceClass: + key = str(device) + if key not in self._device_classes: + self._device_classes[key] = self._device_class_for(device) + return self._device_classes[key] + + def _environment(self, device: torch.device) -> Environment: + config = self._services.configuration + settings: dict[str, SettingValue] = {name: getattr(config, name) for name in MEMORY_SETTINGS} + return self._environment_for(device, settings) + + def _current_device_class_id(self, device_class_id: Optional[str]) -> str: + if device_class_id is not None: + return device_class_id + return self._device_class(self._calibration_device()).id + + # endregion + + # region Runs + + def _find_run(self, run_id: str) -> Optional[tuple[str, CalibrationRun]]: + cached = self._run_files.get(run_id) + candidates = [cached] if cached is not None else [] + candidates += [ + device_class_id for device_class_id in self._store.device_class_ids() if device_class_id != cached + ] + for device_class_id in candidates: + read = self._store.read(device_class_id) + if read is None or read.file is None: + continue + for run in read.file.runs: + if run.run_id == run_id: + self._run_files[run_id] = device_class_id + return device_class_id, run + self._run_files.pop(run_id, None) + return None + + @staticmethod + def _calibration_batch(request: WorkingMemoryCalibrationRequest, run_id: str) -> Batch: + graph = request.graph + for node in graph.nodes.values(): + if isinstance(node, CallSavedWorkflowInvocation): + raise CalibrationGraphError( + f"Node {node.id} calls a saved workflow. A calibration graph cannot: the called workflow would " + "run as separate queue items, outside the measured node." + ) + for datum in chain.from_iterable(request.data or []): + if datum.field_name == "use_cache": + raise CalibrationGraphError( + f"Batch data sets use_cache on node {datum.node_path}; calibration runs every node uncached." + ) + # A cached output would skip the very operation the run measures. + for node in graph.nodes.values(): + node.use_cache = False + try: + return Batch( + graph=graph, + data=request.data, + runs=1, + origin=calibration_origin(run_id), + destination=CALIBRATION_DESTINATION, + idempotency_key=request.idempotency_key, + ) + except ValueError as error: + raise CalibrationGraphError(str(error)) from error + + @staticmethod + def _run_id_for(user_id: str, idempotency_key: Optional[str]) -> str: + """A fresh run id, or with a retry key the same one for every repeat of the request: the run id is part of the + batch's origin, so a replay must rebuild the identical batch for the queue to recognise it.""" + if idempotency_key is None: + return uuid_string() + return str(uuid.uuid5(_RUN_ID_NAMESPACE, f"{user_id}\x00{idempotency_key}")) + + async def create_run( + self, model_key: str, request: WorkingMemoryCalibrationRequest, user_id: str + ) -> WorkingMemoryCalibrationRunCreated: + run_id = self._run_id_for(user_id, request.idempotency_key) + + def prepare() -> tuple[Batch, str, bool]: + """The batch, the run's device class, and whether this call recorded the run (a replay finds it).""" + device = self._calibration_device() + # Raises UnknownModelException for an unknown key. + config = self._services.model_manager.store.get_model(model_key) + batch = self._calibration_batch(request, run_id) + device_class = self._device_class(device) + env = self._environment(device) + run = CalibrationRun( + run_id=run_id, + model_key=model_key, + model_hash=getattr(config, "hash", None), + user_id=user_id, + created_at=_now(), + env_id=env.id, + ) + + def add(file: CalibrationStoreFile) -> tuple[bool, bool]: + if any(known.run_id == run_id for known in file.runs): + return False, False + file.runs.append(run) + _add_env(file, env) + return True, True + + created = self._store.change(device_class.id, add, create=CalibrationDeviceClass.of(device_class)) + self._run_files[run_id] = device_class.id + return batch, device_class.id, bool(created) + + batch, device_class_id, created = await asyncio.to_thread(prepare) + try: + result = await self._services.session_queue.enqueue_batch( + queue_id=DEFAULT_QUEUE_ID, batch=batch, prepend=False, user_id=user_id + ) + except BaseException: + # Only a run this call recorded: a conflicting replay must not remove the original request's run. + if created: + await asyncio.to_thread(self._discard_run, run_id) + raise + if result.enqueued == 0: + if created: + await asyncio.to_thread(self._discard_run, run_id) + raise CalibrationQueueFullError("The queue is full; the calibration run was not started.") + + def remember_items(file: CalibrationStoreFile) -> tuple[bool, None]: + for run in file.runs: + if run.run_id == run_id and run.queue_item_ids != result.item_ids: + run.queue_item_ids = list(result.item_ids) + return True, None + return False, None + + try: + await asyncio.to_thread(self._store.change, device_class_id, remember_items) + except Exception: + # The items are enqueued and carry the run's origin, which is what the status reads them by. + self._services.logger.warning(f"Could not record the items of calibration run {run_id}", exc_info=True) + return WorkingMemoryCalibrationRunCreated( + run_id=run_id, device_class_id=device_class_id, item_ids=list(result.item_ids) + ) + + def _discard_run(self, run_id: str) -> None: + try: + self.reset_run(run_id) + except Exception: + self._services.logger.warning(f"Could not discard calibration run {run_id}", exc_info=True) + + def admit(self, queue_item: SessionQueueItem, run_id: str) -> Optional[str]: + if self._services.configuration.multiuser: + owner = self._services.users.get(queue_item.user_id) + if owner is None or not owner.is_active or not owner.is_admin: + return ( + f"Calibration runs are for administrators only, and the owner of queue item {queue_item.item_id} " + "is not an active administrator." + ) + try: + device = self._calibration_device() + except CalibrationUnavailableError as error: + return str(error) + found = self._find_run(run_id) + if found is None: + return f"Calibration run {run_id} does not exist; it was reset or never created." + device_class_id, _ = found + if device_class_id != self._device_class(device).id: + return ( + f"Calibration run {run_id} was created for device class {device_class_id}, and this machine's GPU " + f"is of class {self._device_class(device).id}." + ) + self._observe( + run_id, + queue_item.item_id, + ObservedItem( + status="in_progress", retried_from_item_id=queue_item.retried_from_item_id, observed_at=_now() + ), + ) + return None + + def _observe(self, run_id: str, item_id: int, observation: ObservedItem) -> None: + found = self._find_run(run_id) + if found is None: + return + + def record(file: CalibrationStoreFile) -> tuple[bool, None]: + for run in file.runs: + if run.run_id == run_id: + previous = run.observed.get(item_id) + if observation.retried_from_item_id is None and previous is not None: + observation_ = observation.model_copy( + update={"retried_from_item_id": previous.retried_from_item_id} + ) + else: + observation_ = observation + run.observed[item_id] = observation_ + return True, None + return False, None + + self._store.change(found[0], record) + + def observe(self, queue_item: SessionQueueItem, run_id: str, error: Optional[BaseException] = None) -> None: + try: + current = self._services.session_queue.get_queue_item(queue_item.item_id) + except SessionQueueItemNotFoundError: + # Removed while it ran (Clear Queue, a deleted item): it will not finish. + self._observe( + run_id, + queue_item.item_id, + ObservedItem( + status="canceled", + error_message="The queue item was removed from the queue.", + retried_from_item_id=queue_item.retried_from_item_id, + observed_at=_now(), + ), + ) + return + status, error_message = current.status, current.error_message + if error is not None and status in _RUNNING: + # The processor fails the item after the runner raised; recorded here, a later prune cannot lose it. + status, error_message = "failed", f"{type(error).__name__}: {error}" + self._observe( + run_id, + current.item_id, + ObservedItem( + status=status, + error_message=error_message, + retried_from_item_id=current.retried_from_item_id, + observed_at=_now(), + ), + ) + + def record_windows(self, run_id: str, windows: list[MeasuredWindow]) -> None: + found = self._find_run(run_id) + if found is None: + self._services.logger.warning( + f"Dropped {len(windows)} working-memory measurements of calibration run {run_id}: the run was reset." + ) + return + run_class_id = found[0] + by_device: dict[str, list[MeasuredWindow]] = {} + for window in windows: + by_device.setdefault(window.device, []).append(window) + for device_name, group in by_device.items(): + device = torch.device(device_name) + device_class = self._device_class(device) + if device_class.id != run_class_id: + # Admission made the run's device class the generation device's; a window elsewhere is a model the + # node ran on another device, which this run does not calibrate. + self._services.logger.warning( + f"Dropped {len(group)} working-memory measurements of calibration run {run_id}: they ran on " + f"{device_name} (device class {device_class.id}), not on the run's device class {run_class_id}." + ) + continue + env = self._environment(device) + points = [CalibrationPoint(**window.model_dump(), env_id=env.id) for window in group] + if not self._store.change(run_class_id, partial(_add_points, run_id=run_id, env=env, points=points)): + self._services.logger.warning( + f"Dropped {len(group)} working-memory measurements of calibration run {run_id}: the run was reset." + ) + + # endregion + + # region Status + + def _live_items(self) -> dict[str, dict[int, QUEUE_ITEM_STATUS]]: + """The queue's calibration items, per run, as the queue has them now.""" + queue = self._services.session_queue + item_ids = queue.get_queue_item_ids(DEFAULT_QUEUE_ID, origin_prefix=CALIBRATION_ORIGIN_PREFIX).item_ids + live: dict[str, dict[int, QUEUE_ITEM_STATUS]] = {} + for start in range(0, len(item_ids), _SUMMARY_CHUNK): + for summary in queue.get_queue_item_summaries_by_ids( + DEFAULT_QUEUE_ID, item_ids[start : start + _SUMMARY_CHUNK] + ): + # The queue's origin filter ignores case; only the exact prefix is the server's. + run_id = calibration_run_id(summary.origin) + if run_id is not None: + live.setdefault(run_id, {})[summary.item_id] = summary.status + return live + + def get_status(self) -> WorkingMemoryCalibrationStatus: + current_class: Optional[DeviceClass] = None + current_env: Optional[Environment] = None + unavailable: Optional[str] = None + try: + device = self._calibration_device() + current_class = self._device_class(device) + current_env = self._environment(device) + except CalibrationUnavailableError as error: + unavailable = str(error) + + live = self._live_items() + device_class_ids = set(self._store.device_class_ids()) + if current_class is not None: + device_class_ids.add(current_class.id) + device_classes = [] + for device_class_id in sorted(device_class_ids): + is_current = current_class is not None and device_class_id == current_class.id + read = self._store.read(device_class_id) + if read is None: + if current_class is None or not is_current: + continue + empty = CalibrationStoreFile( + schema_version=STORE_SCHEMA_VERSION, device_class=CalibrationDeviceClass.of(current_class) + ) + read = StoreRead(device_class_id, empty, None) + if is_current and read.file is None: + # Creating a run would have to write this file, which is kept as it is. + unavailable = f"The calibration store of this GPU's device class cannot be written: {read.problem}" + device_classes.append( + self._device_class_status(read, is_current, current_env if is_current else None, live) + ) + return WorkingMemoryCalibrationStatus( + available=unavailable is None, + unavailable_reason=unavailable, + current_device_class_id=current_class.id if current_class is not None else None, + current_env_id=current_env.id if current_env is not None else None, + device_classes=device_classes, + ) + + @staticmethod + def _device_class_status( + read: StoreRead, + is_current: bool, + current_env: Optional[Environment], + live: Mapping[str, Mapping[int, QUEUE_ITEM_STATUS]], + ) -> WorkingMemoryCalibrationDeviceClass: + file = read.file + if file is None: + return WorkingMemoryCalibrationDeviceClass( + device_class_id=read.device_class_id, + device_class=None, + usable=False, + problem=read.problem, + is_current=is_current, + runs=[], + point_counts=[], + envs=[], + ) + points_per_run: dict[str, int] = {} + points_per_env: dict[str, int] = {} + counts: dict[tuple[str, int, bool, bool], int] = {} + for point in file.points: + points_per_run[point.run_id] = points_per_run.get(point.run_id, 0) + 1 + points_per_env[point.env_id] = points_per_env.get(point.env_id, 0) + 1 + key = (point.estimator, point.version, point.valid, point.fit_eligible) + counts[key] = counts.get(key, 0) + 1 + runs = [] + for run in file.runs: + point_count = points_per_run.get(run.run_id, 0) + status, error, item_ids = derive_run_status(run, live.get(run.run_id, {}), point_count) + runs.append( + WorkingMemoryCalibrationRunSummary( + run_id=run.run_id, + model_key=run.model_key, + model_hash=run.model_hash, + user_id=run.user_id, + created_at=run.created_at, + env_id=run.env_id, + status=status, + error=error, + item_ids=item_ids, + point_count=point_count, + ) + ) + envs = [ + WorkingMemoryCalibrationEnvironment( + id=env.id, + hard=env.hard, + soft=env.soft, + hard_matches_current=None if current_env is None else env.id == current_env.id, + soft_matches_current=None if current_env is None else env.soft == current_env.soft, + point_count=points_per_env.get(env.id, 0), + ) + for env in file.envs + ] + return WorkingMemoryCalibrationDeviceClass( + device_class_id=read.device_class_id, + device_class=file.device_class, + usable=True, + problem=None, + is_current=is_current, + runs=runs, + point_counts=[ + WorkingMemoryCalibrationPointCount( + estimator=estimator, version=version, valid=valid, fit_eligible=fit_eligible, count=count + ) + for (estimator, version, valid, fit_eligible), count in sorted(counts.items()) + ], + envs=envs, + ) + + # endregion + + # region Reset and export + + @staticmethod + def _remove( + file: CalibrationStoreFile, run_ids: set[str], model_key: Optional[str] = None + ) -> tuple[bool, WorkingMemoryCalibrationResetResult]: + runs = [run for run in file.runs if run.run_id not in run_ids] + points = [ + point + for point in file.points + if point.run_id not in run_ids and (model_key is None or point.model_key != model_key) + ] + result = WorkingMemoryCalibrationResetResult( + runs_removed=len(file.runs) - len(runs), points_removed=len(file.points) - len(points) + ) + file.runs, file.points = runs, points + _prune_envs(file) + return result.runs_removed > 0 or result.points_removed > 0, result + + def reset_run(self, run_id: str) -> WorkingMemoryCalibrationResetResult: + found = self._find_run(run_id) + if found is None: + raise CalibrationRunNotFoundError(run_id) + result = self._store.change(found[0], lambda file: self._remove(file, {run_id})) + self._run_files.pop(run_id, None) + return result or WorkingMemoryCalibrationResetResult(runs_removed=0, points_removed=0) + + def reset_model(self, model_key: str, device_class_id: Optional[str] = None) -> WorkingMemoryCalibrationResetResult: + def remove(file: CalibrationStoreFile) -> tuple[bool, WorkingMemoryCalibrationResetResult]: + run_ids = {run.run_id for run in file.runs if run.model_key == model_key} + return self._remove(file, run_ids, model_key) + + result = self._store.change(self._current_device_class_id(device_class_id), remove) + return result or WorkingMemoryCalibrationResetResult(runs_removed=0, points_removed=0) + + def reset_all(self, device_class_id: Optional[str] = None) -> WorkingMemoryCalibrationResetResult: + def remove(file: CalibrationStoreFile) -> tuple[bool, WorkingMemoryCalibrationResetResult]: + result = WorkingMemoryCalibrationResetResult(runs_removed=len(file.runs), points_removed=len(file.points)) + changed = bool(file.runs or file.points or file.envs) + file.runs, file.points, file.envs = [], [], [] + return changed, result + + result = self._store.change(self._current_device_class_id(device_class_id), remove) + return result or WorkingMemoryCalibrationResetResult(runs_removed=0, points_removed=0) + + def export(self, device_class_id: Optional[str] = None) -> Optional[bytes]: + return self._store.read_raw(self._current_device_class_id(device_class_id)) + + # endregion diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_store.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_store.py new file mode 100644 index 00000000000..c669385f78f --- /dev/null +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_store.py @@ -0,0 +1,249 @@ +"""The calibration store: one JSON file per device class under `/working_memory/`. + +Machine data, not account data, so it lives beside the database rather than in it: a device class's measurements stay +valid whichever database the root points at. Every change re-reads the file and is applied to what is on disk under a +process lock and an OS file lock, so two processes sharing a root never lose each other's writes. Files are replaced +atomically. A file this version cannot read -- another schema version, invalid content, a field it does not know, or a device +class that does not match its name -- is reported and never written, moved or deleted. +""" + +import json +import os +import re +import tempfile +import threading +import time +from collections.abc import Callable, Iterator +from contextlib import contextmanager, suppress +from pathlib import Path +from typing import IO, Final, Optional, TypeVar + +from pydantic import ValidationError + +from invokeai.app.services.working_memory_calibration.working_memory_calibration_common import ( + STORE_SCHEMA_VERSION, + CalibrationDeviceClass, + CalibrationStoreError, + CalibrationStoreFile, + CalibrationStoreUnusableError, +) + +LOCK_TIMEOUT_S: float = 30.0 +"""How long a change waits for another process's lock on the same file before it gives up.""" +LOCK_RETRY_S: float = 0.05 +REPLACE_ATTEMPTS: int = 6 +"""How often a write tries to replace the file. On Windows the replace fails with PermissionError while a process that +does not take the store's lock (an antivirus scanner, a search indexer, an editor) has the file open; such a handle is +usually gone within a fraction of a second.""" +REPLACE_RETRY_S: float = 0.05 +"""The first wait between those tries; it doubles each time (about 1.5 s in all).""" + +_FILE_NAME: Final = re.compile(r"^[0-9a-f]{16}\.json$") + +T = TypeVar("T") + + +class StoreRead: + """A device class's file as read: its content, or why it cannot be used.""" + + def __init__(self, device_class_id: str, file: Optional[CalibrationStoreFile], problem: Optional[str]) -> None: + self.device_class_id = device_class_id + self.file = file + self.problem = problem + + +def _parse(raw: bytes, device_class_id: str) -> tuple[Optional[CalibrationStoreFile], Optional[str]]: + try: + data = json.loads(raw) + except ValueError: + return None, "The file is not valid JSON." + if not isinstance(data, dict): + return None, "The file does not hold a calibration store." + version = data.get("schema_version") + if version != STORE_SCHEMA_VERSION: + return None, f"The file has schema version {version!r}; this version of Invoke reads {STORE_SCHEMA_VERSION}." + try: + file = CalibrationStoreFile.model_validate(data) + except ValidationError as error: + first = error.errors()[0] + location = ".".join(str(part) for part in first["loc"]) + return None, f"The file failed validation ({error.error_count()} errors; first at {location}: {first['msg']})." + if file.device_class.id != device_class_id: + return None, f"The file holds device class {file.device_class.id}, not the {device_class_id} its name says." + return file, None + + +@contextmanager +def _os_lock(handle: IO[bytes]) -> Iterator[None]: + if os.name == "nt": + import msvcrt + + deadline = time.monotonic() + LOCK_TIMEOUT_S + while True: + handle.seek(0) + try: + msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1) + break + except OSError: + if time.monotonic() > deadline: + raise CalibrationStoreError(f"Timed out waiting for the lock on {handle.name}") from None + time.sleep(LOCK_RETRY_S) + try: + yield + finally: + handle.seek(0) + msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1) + else: + import fcntl + + deadline = time.monotonic() + LOCK_TIMEOUT_S + while True: + try: + fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + break + except BlockingIOError: + if time.monotonic() > deadline: + raise CalibrationStoreError(f"Timed out waiting for the lock on {handle.name}") from None + time.sleep(LOCK_RETRY_S) + try: + yield + finally: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + + +def _replace(staging: str, path: Path) -> None: + """`os.replace`, retried while another process's handle on the target refuses it (Windows only).""" + delay = REPLACE_RETRY_S + for attempt in range(REPLACE_ATTEMPTS): + try: + os.replace(staging, path) + return + except PermissionError: + if attempt + 1 == REPLACE_ATTEMPTS: + raise + time.sleep(delay) + delay *= 2 + + +class CalibrationStore: + def __init__(self, directory: Path) -> None: + self._dir = directory + self._lock = threading.Lock() + + @property + def directory(self) -> Path: + return self._dir + + def path_for(self, device_class_id: str) -> Path: + return self._dir / f"{device_class_id}.json" + + def device_class_ids(self) -> list[str]: + """The device classes that have a file, readable or not.""" + if not self._dir.is_dir(): + return [] + return sorted(path.stem for path in self._dir.iterdir() if _FILE_NAME.match(path.name) and path.is_file()) + + @contextmanager + def _locked(self, device_class_id: str, *, writing: bool) -> Iterator[None]: + """Hold this process's lock and the OS lock on the file's `.lock` sibling. + + Readers lock too: on Windows a file another process holds open cannot be replaced. A reader that cannot create + the lock file (a read-only root) reads without it, since nobody can write there either. + """ + with self._lock: + lock_path = self._dir / f"{device_class_id}.json.lock" + handle: Optional[IO[bytes]] = None + try: + if writing: + self._dir.mkdir(parents=True, exist_ok=True) + handle = open(lock_path, "a+b") + except OSError as error: + if writing: + raise CalibrationStoreError(f"Cannot open {lock_path}: {error}") from error + if handle is None: + yield + return + with handle, _os_lock(handle): + yield + + def _read_unlocked(self, device_class_id: str) -> Optional[StoreRead]: + path = self.path_for(device_class_id) + try: + raw = path.read_bytes() + except FileNotFoundError: + return None + except OSError as error: + raise CalibrationStoreError(f"Cannot read {path}: {error}") from error + file, problem = _parse(raw, device_class_id) + return StoreRead(device_class_id, file, problem) + + def read(self, device_class_id: str) -> Optional[StoreRead]: + """The device class's file, or None when it has none.""" + if not self.path_for(device_class_id).exists(): + return None + with self._locked(device_class_id, writing=False): + return self._read_unlocked(device_class_id) + + def read_raw(self, device_class_id: str) -> Optional[bytes]: + """The file's bytes as they are, readable or not; None when it has none.""" + path = self.path_for(device_class_id) + if not path.exists(): + return None + with self._locked(device_class_id, writing=False): + try: + return path.read_bytes() + except FileNotFoundError: + return None + except OSError as error: + raise CalibrationStoreError(f"Cannot read {path}: {error}") from error + + def change( + self, + device_class_id: str, + apply: Callable[[CalibrationStoreFile], tuple[bool, T]], + *, + create: Optional[CalibrationDeviceClass] = None, + ) -> Optional[T]: + """Apply `apply` to the file as it is on disk now, and write the file when `apply` says it changed it. + + `apply` returns whether it changed the file, and the result to return. Without a file, `create` names the + device class to start one for; without `create` nothing is written and the result is None. Raises + `CalibrationStoreUnusableError` for a file that cannot be read, which stays as it is. + """ + with self._locked(device_class_id, writing=True): + current = self._read_unlocked(device_class_id) + if current is None: + if create is None: + return None + assert create.id == device_class_id + file = CalibrationStoreFile(schema_version=STORE_SCHEMA_VERSION, device_class=create) + elif current.file is None: + raise CalibrationStoreUnusableError( + f"{self.path_for(device_class_id)} is kept as it is and not written: {current.problem}" + ) + else: + file = current.file + changed, result = apply(file) + if changed: + self._write(self.path_for(device_class_id), file) + return result + + @staticmethod + def _write(path: Path, file: CalibrationStoreFile) -> None: + payload = file.model_dump_json(indent=1).encode("utf-8") + try: + descriptor, staging = tempfile.mkstemp(dir=path.parent, prefix=f"{path.name}.", suffix=".tmp") + except OSError as error: + raise CalibrationStoreError(f"Cannot write {path}: {error}") from error + try: + with os.fdopen(descriptor, "wb") as handle: + handle.write(payload) + handle.flush() + os.fsync(handle.fileno()) + _replace(staging, path) + except BaseException as error: + with suppress(OSError): + os.unlink(staging) + if isinstance(error, OSError): + raise CalibrationStoreError(f"Cannot write {path}: {error}") from error + raise diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index fddf9a368dc..05c75017ada 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -26,6 +26,7 @@ ) from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType from invokeai.backend.util.logging import InvokeAILogger +from invokeai.backend.working_memory import probe from invokeai.backend.working_memory.estimate import WorkingMemoryEstimate @@ -181,18 +182,43 @@ def _lock_paced(self, working_mem_bytes: Optional[int]) -> None: # invariant turning this loop into a spin. assert settled, "uncapped continue_lock pass did not settle" - def __enter__(self) -> AnyModel: - self._lock_paced(None) - self._end_first_use_window() + def _lock_on_device(self, working_mem_bytes: Optional[int]) -> None: + """Lock the model onto its compute device and repair its required tensors there. Nothing stays locked if this + raises. + + While a calibration scope is open anywhere (`probe.armed`), the working-memory probe hears of the lock from its + start until it has settled with the tensors repaired: inside a measured window on this thread that interval is + a co-lock, which moves weights and is kept out of the window's peak; on another thread it voids the window. + """ + observed = probe.armed + if observed: + probe.lock_started(self._cache, self._cache_record) + locked = False try: + self._lock_paced(working_mem_bytes) + locked = True + self._end_first_use_window() self.repair_required_tensors_on_device() - return self.model - except Exception: - self._cache.unlock(self._cache_record) + except BaseException: + if locked: + self._cache.unlock(self._cache_record) + locked = False raise + finally: + if observed: + probe.lock_settled(self._cache, self._cache_record, locked) - def __exit__(self, *args: Any, **kwargs: Any) -> None: + def _unlock(self) -> None: self._cache.unlock(self._cache_record) + if probe.armed: + probe.unlocked(self._cache, self._cache_record) + + def __enter__(self) -> AnyModel: + self._lock_on_device(None) + return self.model + + def __exit__(self, *args: Any, **kwargs: Any) -> None: + self._unlock() @contextmanager def model_on_device( @@ -203,27 +229,45 @@ def model_on_device( :param working_mem_bytes: The amount of working memory to keep available on the compute device when loading the model. :param working_memory: A working-memory estimate whose `reservation_bytes` is kept available instead; see - `invokeai.backend.working_memory`. Pass this or `working_mem_bytes`, not both. + `invokeai.backend.working_memory`. Pass this or `working_mem_bytes`, not both. Inside a calibration scope + (`working_memory.probe`) on a GPU, the lock reserves calibration headroom on top and the probe measures the + operation from the settled lock until the unlock. """ if isinstance(working_mem_bytes, WorkingMemoryEstimate): raise TypeError( "model_on_device() got a WorkingMemoryEstimate as working_mem_bytes; pass it as working_memory=." ) + calibration: Optional[probe.ProbeScope] = None if working_memory is not None: if working_mem_bytes is not None: raise ValueError("Pass either working_mem_bytes or working_memory to model_on_device(), not both.") working_mem_bytes = working_memory.reservation_bytes - InvokeAILogger.get_logger(__name__).debug( + message = ( f"Working memory for {self._cache_record.key}: {working_memory.estimator} reserves " f"{working_mem_bytes / 2**20:.0f} MB ({working_memory.source})" ) - self._lock_paced(working_mem_bytes) - self._end_first_use_window() + scope = probe.current_scope() + if scope is not None and self.compute_device.type != "cpu": + calibration = scope + working_mem_bytes = probe.calibration_reservation(working_mem_bytes) + message += f"; the calibration lock reserves {working_mem_bytes / 2**20:.0f} MB" + InvokeAILogger.get_logger(__name__).debug(message) + self._lock_on_device(working_mem_bytes) + window = None try: - self.repair_required_tensors_on_device() + if calibration is not None: + assert working_memory is not None and working_mem_bytes is not None + window = calibration.open_window(self._cache, self._cache_record, working_memory, working_mem_bytes) yield (self._cache_record.cached_model.get_cpu_state_dict(), self._cache_record.cached_model.model) + except BaseException as error: + if window is not None: + window.close(error) + window = None + raise finally: - self._cache.unlock(self._cache_record) + if window is not None: + window.close(None) + self._unlock() @property def model(self) -> AnyModel: @@ -234,11 +278,13 @@ def model(self) -> AnyModel: def model_in_ram(self) -> Generator[AnyModel, None, None]: """Pin the model's cache record in RAM without moving the model to its execution device.""" self._cache.lock_in_ram(self._cache_record) - self._end_first_use_window() + if probe.armed: + probe.locked_in_ram(self._cache, self._cache_record) try: + self._end_first_use_window() yield self.model finally: - self._cache.unlock(self._cache_record) + self._unlock() @property def compute_device(self) -> torch.device: diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index 3691f374425..59d422ea1d9 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -1888,6 +1888,15 @@ def _claim_first_use(self, cache_entry: CacheRecord) -> Optional[FirstUseClaim]: self.release_first_use_hold(cache_entry, hold_epoch) raise + @synchronized + def locks_held(self) -> int: + """How many locks this cache's records hold right now: VRAM and RAM pins of every thread together. + + The working-memory probe compares it with the locks its own thread took, so that a window opening while + another thread holds a model on this device knows it shares the device. + """ + return sum(record._locks for record in self._cached_models.values()) + @synchronized @record_activity def lock_in_ram(self, cache_entry: CacheRecord) -> None: diff --git a/invokeai/backend/util/devices.py b/invokeai/backend/util/devices.py index 40249246035..276ec2ecd30 100644 --- a/invokeai/backend/util/devices.py +++ b/invokeai/backend/util/devices.py @@ -2,7 +2,7 @@ import threading from collections import Counter, defaultdict from functools import wraps -from typing import Dict, Literal, Optional, Union +from typing import Dict, Literal, Optional, Protocol, Union import torch from deprecated import deprecated @@ -24,6 +24,19 @@ _XPU_MEM_FALLBACK_WARNED: set[str] = set() +class EmptyCacheObserver(Protocol): + def before_empty_cache(self) -> None: ... + + def after_empty_cache(self) -> None: ... + + +EMPTY_CACHE_OBSERVERS: list[EmptyCacheObserver] = [] +"""Told about every `TorchDevice.empty_cache` that releases memory, just before and just after the release, in +registration order. The working-memory probe registers here while a calibration scope is open, so that a release inside +a measured window ends a segment there (`invokeai.backend.working_memory.probe`); otherwise the list is empty and a +release pays one truth test for it. Observers must not raise.""" + + def _xpu_is_available() -> bool: """Return True if a torch XPU (Intel GPU) device is available.""" return hasattr(torch, "xpu") and torch.xpu.is_available() @@ -446,14 +459,21 @@ def empty_cache(cls, force: bool = False) -> bool: # Clear before running: a skip that races in after this point re-sets the flag, so a # request is never lost, only (harmlessly) repeated. cls._empty_cache_deferred.clear() - if torch.backends.mps.is_available(): - torch.mps.empty_cache() - if torch.cuda.is_available(): - # A forced call must get past `install_peer_aware_empty_cache`'s wrapper, which would defer it again. - cuda_empty_cache = torch.cuda.empty_cache - (getattr(cuda_empty_cache, "__wrapped__", cuda_empty_cache) if force else cuda_empty_cache)() - if _xpu_is_available(): - torch.xpu.empty_cache() + observers = tuple(EMPTY_CACHE_OBSERVERS) if EMPTY_CACHE_OBSERVERS else () + for observer in observers: + observer.before_empty_cache() + try: + if torch.backends.mps.is_available(): + torch.mps.empty_cache() + if torch.cuda.is_available(): + # A forced call must get past `install_peer_aware_empty_cache`'s wrapper, which would defer it again. + cuda_empty_cache = torch.cuda.empty_cache + (getattr(cuda_empty_cache, "__wrapped__", cuda_empty_cache) if force else cuda_empty_cache)() + if _xpu_is_available(): + torch.xpu.empty_cache() + finally: + for observer in observers: + observer.after_empty_cache() return True @classmethod diff --git a/invokeai/backend/util/nvml.py b/invokeai/backend/util/nvml.py new file mode 100644 index 00000000000..cffb36efd67 --- /dev/null +++ b/invokeai/backend/util/nvml.py @@ -0,0 +1,146 @@ +"""Minimal NVML queries through ctypes: this process's memory on an NVIDIA device, and the driver version. + +NVML ships with the NVIDIA driver (``libnvidia-ml.so.1`` on Linux, ``nvml.dll`` in System32 on Windows), so this needs +no package. Under WDDM, Windows reports no per-process memory through NVML; the working-memory probe reads it through +D3DKMT there (`wddm.process_local_usage`), and uses NVML on Windows for the driver version only. + +Everything here is best-effort: any failure yields ``None`` ("unknown") rather than raising. +""" + +import ctypes +import functools +import os +import sys +import threading +from typing import Optional + +import torch + +from invokeai.backend.util.logging import InvokeAILogger +from invokeai.backend.util.pci import pci_address + +_NVML_SUCCESS = 0 +_NVML_ERROR_INSUFFICIENT_SIZE = 7 +# `usedGpuMemory` when the driver cannot attribute memory to the process (always so under WDDM). +_NVML_VALUE_NOT_AVAILABLE = 0xFFFFFFFFFFFFFFFF +# NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE. +_DRIVER_VERSION_BUFFER_SIZE = 80 +# Room for processes that start between the count and the fill call. +_SPARE_PROCESS_ENTRIES = 8 + +_LIBRARY_NAMES = ("nvml.dll",) if sys.platform == "win32" else ("libnvidia-ml.so.1", "libnvidia-ml.so") + + +class _ProcessInfo(ctypes.Structure): + """``nvmlProcessInfo_t``, the entry ``nvmlDeviceGetComputeRunningProcesses_v3`` fills.""" + + _fields_ = [ + ("pid", ctypes.c_uint), + ("usedGpuMemory", ctypes.c_ulonglong), + ("gpuInstanceId", ctypes.c_uint), + ("computeInstanceId", ctypes.c_uint), + ] + + +_lock = threading.Lock() +# Per torch device index: the NVML device handle, or None once the lookup has failed. +_handles: dict[int, Optional[ctypes.c_void_p]] = {} + + +def _configure_prototypes(lib: ctypes.CDLL) -> None: + """Declare every entry point used here, so ctypes passes handles at full pointer width.""" + prototypes = { + "nvmlInit_v2": [], + "nvmlDeviceGetHandleByPciBusId_v2": [ctypes.c_char_p, ctypes.POINTER(ctypes.c_void_p)], + "nvmlDeviceGetComputeRunningProcesses_v3": [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_uint), + ctypes.c_void_p, + ], + "nvmlSystemGetDriverVersion": [ctypes.c_char_p, ctypes.c_uint], + } + for name, argtypes in prototypes.items(): + fn = getattr(lib, name) + fn.argtypes = argtypes + fn.restype = ctypes.c_int + + +@functools.lru_cache(maxsize=1) +def _load() -> Optional[ctypes.CDLL]: + """Open and initialise NVML once per process; None where there is no NVIDIA driver.""" + for name in _LIBRARY_NAMES: + try: + lib = ctypes.CDLL(name) + except OSError: + continue + try: + _configure_prototypes(lib) + if lib.nvmlInit_v2() != _NVML_SUCCESS: + return None + except Exception as exc: + InvokeAILogger.get_logger(__name__).debug(f"NVML unusable: {exc}") + return None + return lib + return None + + +def _handle_for(lib: ctypes.CDLL, device: torch.device) -> Optional[ctypes.c_void_p]: + index = device.index if device.index is not None else torch.cuda.current_device() + with _lock: + if index not in _handles: + handle = ctypes.c_void_p() + found = lib.nvmlDeviceGetHandleByPciBusId_v2(pci_address(index).encode(), ctypes.byref(handle)) + _handles[index] = handle if found == _NVML_SUCCESS else None + return _handles[index] + + +def process_used_memory(device: torch.device) -> Optional[int]: + """Bytes NVML attributes to this process on `device`: everything the process holds there, the CUDA context + included. ``None`` off CUDA builds, without NVML, and wherever the driver does not attribute memory to processes.""" + if device.type != "cuda" or torch.version.hip is not None: + return None + try: + lib = _load() + if lib is None: + return None + handle = _handle_for(lib, device) + if handle is None: + return None + count = ctypes.c_uint(0) + status = lib.nvmlDeviceGetComputeRunningProcesses_v3(handle, ctypes.byref(count), None) + if status != _NVML_ERROR_INSUFFICIENT_SIZE: + # Success with an empty buffer means no process at all, so not this one either. + return None + count = ctypes.c_uint(count.value + _SPARE_PROCESS_ENTRIES) + infos = (_ProcessInfo * count.value)() + if lib.nvmlDeviceGetComputeRunningProcesses_v3(handle, ctypes.byref(count), infos) != _NVML_SUCCESS: + return None + pid = os.getpid() + for info in infos[: count.value]: + if info.pid == pid: + return None if info.usedGpuMemory == _NVML_VALUE_NOT_AVAILABLE else int(info.usedGpuMemory) + return None + except Exception: + InvokeAILogger.get_logger(__name__).debug(f"NVML process memory lookup for {device} failed", exc_info=True) + return None + + +def driver_version() -> Optional[str]: + """The NVIDIA driver version, e.g. ``"581.57"``, or None without NVML.""" + try: + lib = _load() + if lib is None: + return None + buffer = ctypes.create_string_buffer(_DRIVER_VERSION_BUFFER_SIZE) + if lib.nvmlSystemGetDriverVersion(buffer, _DRIVER_VERSION_BUFFER_SIZE) != _NVML_SUCCESS: + return None + return buffer.value.decode("ascii", errors="replace") or None + except Exception: + return None + + +def reset_cache() -> None: + """Forget the loader and the device handles (tests only). NVML itself stays initialised.""" + with _lock: + _handles.clear() + _load.cache_clear() diff --git a/invokeai/backend/util/pci.py b/invokeai/backend/util/pci.py new file mode 100644 index 00000000000..765f9431d03 --- /dev/null +++ b/invokeai/backend/util/pci.py @@ -0,0 +1,17 @@ +"""Where a CUDA or HIP device sits on the PCI bus: what links a torch device to the driver interfaces that know it by +location (D3DKMT adapters, NVML, DRM fdinfo).""" + +import torch + + +def pci_location(index: int) -> tuple[int, int, int]: + """``(domain, bus, device)`` of torch's CUDA or HIP device `index`.""" + props = torch.cuda.get_device_properties(index) + return int(props.pci_domain_id), int(props.pci_bus_id), int(props.pci_device_id) + + +def pci_address(index: int) -> str: + """The device's address as ``domain:bus:device.function`` in lower-case hex, e.g. ``0000:03:00.0``: how DRM + fdinfo spells it, and one of the forms NVML parses.""" + domain, bus, device = pci_location(index) + return f"{domain:04x}:{bus:02x}:{device:02x}.0" diff --git a/invokeai/backend/util/wddm.py b/invokeai/backend/util/wddm.py index ec39f7b7389..445cbe1ffd4 100644 --- a/invokeai/backend/util/wddm.py +++ b/invokeai/backend/util/wddm.py @@ -21,6 +21,11 @@ Whether Windows has actually paged this process out is not visible through D3DKMT (``CurrentUsage`` counts paged bytes as local); it is visible in the ``GPU Process Memory`` performance counters Task Manager shows, read here through PDH. +The working-memory probe measures what an operation needs on CUDA and on ROCm alike, so `process_local_usage` (this +process's ``CurrentUsage``) and `process_paged_bytes` answer on either build under Windows; on an RTX 4090 with torch +2.13 the usage follows torch's reserved bytes to three decimals. The ROCm-only readers the model cache decides with keep +their gate. + Everything here is best-effort: any failure yields ``None`` ("unknown") rather than raising, and callers must keep their existing behaviour when the answer is unknown. """ @@ -36,6 +41,7 @@ import torch from invokeai.backend.util.logging import InvokeAILogger +from invokeai.backend.util.pci import pci_location _KMTQAITYPE_ADAPTERADDRESS = 6 _MEMORY_SEGMENT_GROUP_LOCAL = 0 @@ -149,8 +155,12 @@ class _Adapter: _pdh_queries: dict[str, Optional[tuple[ctypes.c_void_p, ctypes.c_void_p]]] = {} +def _windows_gpu(device: torch.device) -> bool: + return sys.platform == "win32" and device.type == "cuda" + + def _supported(device: torch.device) -> bool: - return sys.platform == "win32" and torch.version.hip is not None and device.type == "cuda" + return _windows_gpu(device) and torch.version.hip is not None @functools.lru_cache(maxsize=1) @@ -211,11 +221,11 @@ def _close(lib: ctypes.CDLL, handle: int) -> None: def _resolve_adapter(lib: ctypes.CDLL, index: int) -> Optional[_Adapter]: """Open the WDDM adapter at the torch device's PCI location, closing every other handle the enumeration opened. - ROCm's device properties carry no LUID, so the PCI bus and device numbers are the link. Anything but exactly one - match is treated as unknown rather than guessed. + torch's device properties carry no LUID, on ROCm or CUDA, so the PCI bus and device numbers are the link. Anything + but exactly one match is treated as unknown rather than guessed. """ props = torch.cuda.get_device_properties(index) - location = (int(props.pci_bus_id), int(props.pci_device_id)) + _domain, bus, device_number = pci_location(index) enum = _EnumAdapters2(0, None) if lib.D3DKMTEnumAdapters2(ctypes.byref(enum)) != 0: @@ -241,9 +251,9 @@ def _resolve_adapter(lib: ctypes.CDLL, index: int) -> Optional[_Adapter]: ctypes.cast(ctypes.byref(address), ctypes.c_void_p), ctypes.sizeof(address), ) - if ( - lib.D3DKMTQueryAdapterInfo(ctypes.byref(query)) == 0 - and (address.BusNumber, address.DeviceNumber) == location + if lib.D3DKMTQueryAdapterInfo(ctypes.byref(query)) == 0 and (address.BusNumber, address.DeviceNumber) == ( + bus, + device_number, ): matches.append(info) else: @@ -293,25 +303,61 @@ def local_video_memory(device: torch.device) -> Optional[tuple[int, Optional[int if not _supported(device): return None try: - resolved = _adapter_for(device) - if resolved is None: - return None - lib, adapter = resolved - info = _QueryVideoMemoryInfo(None, adapter.handle, _MEMORY_SEGMENT_GROUP_LOCAL) - if lib.D3DKMTQueryVideoMemoryInfo(ctypes.byref(info)) != 0: + queried = _query_local_memory(device) + if queried is None: return None + info, adapter = queried if not 0 < info.Budget <= adapter.total_bytes: return None - # `CurrentUsage` counts bytes Windows has already paged out, so it exceeds the adapter's own size exactly when - # the process over-commits. It says nothing about residency then, so it is reported as unknown; the budget - # stands, since the model cache's free-VRAM cap is the caller that needs it most in that state. - usage = int(info.CurrentUsage) - return int(info.Budget), usage if 0 <= usage <= adapter.total_bytes else None + # The budget stands when the usage is unknown (see `_plausible_usage`): the model cache's free-VRAM cap is the + # caller that needs it most in exactly that state. + return int(info.Budget), _plausible_usage(info, adapter) except Exception: InvokeAILogger.get_logger(__name__).debug(f"WDDM memory lookup for {device} failed", exc_info=True) return None +def _query_local_memory(device: torch.device) -> Optional[tuple[_QueryVideoMemoryInfo, _Adapter]]: + """This process's local-segment memory info on the device's adapter, with the adapter; None if the driver cannot + answer.""" + resolved = _adapter_for(device) + if resolved is None: + return None + lib, adapter = resolved + info = _QueryVideoMemoryInfo(None, adapter.handle, _MEMORY_SEGMENT_GROUP_LOCAL) + if lib.D3DKMTQueryVideoMemoryInfo(ctypes.byref(info)) != 0: + return None + return info, adapter + + +def _plausible_usage(info: _QueryVideoMemoryInfo, adapter: _Adapter) -> Optional[int]: + """`CurrentUsage`, or None when it exceeds the adapter's own size. + + It counts bytes Windows has already paged out, so it exceeds the adapter exactly when the process over-commits, + and says nothing about residency then. + """ + usage = int(info.CurrentUsage) + return usage if 0 <= usage <= adapter.total_bytes else None + + +def process_local_usage(device: torch.device) -> Optional[int]: + """This process's ``CurrentUsage`` in the local segment of the device's adapter, in bytes, on a CUDA or a ROCm + build under Windows: everything the process has committed there, the allocator's cached blocks, the libraries' own + allocations and any bytes Windows paged out included. + + ``None`` off Windows, for a device that is not a GPU, when the driver cannot answer, and when the usage exceeds the + adapter. + """ + if not _windows_gpu(device): + return None + try: + queried = _query_local_memory(device) + return None if queried is None else _plausible_usage(*queried) + except Exception: + InvokeAILogger.get_logger(__name__).debug(f"WDDM usage lookup for {device} failed", exc_info=True) + return None + + def video_memory_budget(device: torch.device) -> Optional[int]: """This process's WDDM budget on the device's adapter, in bytes: what Windows keeps resident before paging.""" local = local_video_memory(device) @@ -377,6 +423,13 @@ def paged_bytes(device: torch.device) -> Optional[int]: """ if not _supported(device): return None + return process_paged_bytes(device) + + +def process_paged_bytes(device: torch.device) -> Optional[int]: + """`paged_bytes` on a CUDA or a ROCm build under Windows; ``None`` elsewhere.""" + if not _windows_gpu(device): + return None resolved = _adapter_for(device) if resolved is None: return None diff --git a/invokeai/backend/working_memory/counters.py b/invokeai/backend/working_memory/counters.py new file mode 100644 index 00000000000..1089ff5bb33 --- /dev/null +++ b/invokeai/backend/working_memory/counters.py @@ -0,0 +1,333 @@ +"""This process's memory on a compute device, as the working-memory probe reads it at the boundaries of a window. + +Two kinds of counter, and neither alone is an operation's demand: + +- torch's caching allocator: the bytes it reserved from the driver and the bytes live tensors occupy, with exact peaks + since the last reset, on CUDA, ROCm and XPU. MPS keeps current values only; there `PeakSampler` keeps the peaks. +- the driver: what this process holds on the device, including what torch's allocator does not see -- cuBLAS and + cuDNN workspaces, kernel modules, and on ROCm the pool HIP keeps after a free. + +Measured on an RTX 4090 (torch 2.13+cu130, Windows): the driver delta equals torch's reserved-peak delta to three +decimals over VAE decode and encode, matmul, SDPA with and without a mask, and convolutions; first use outside the +allocator adds to the driver alone (+16 MiB at the first matmul, +2 MiB at the first convolution). Under +``backend:cudaMallocAsync`` the reserved peak can sit below the allocated peak (128 vs 136 MiB). On an RX 9060 XT (torch +2.13+rocm10, Windows, expandable segments) the driver keeps HIP's retained pool after ``empty_cache()``, so its delta +reads below reserved (4.3 vs 6.4 GiB for a 1024px decode), except for the first MIOpen decode, which grows it beyond +(2.0 vs 1.72 GiB). So the probe takes the largest delta of the counters it has. + +The driver is read at boundaries only: torch's peaks are exact, so what a boundary read misses is memory outside torch's +allocator that came and went between two boundaries. + +Every reader is best-effort: a failure reads as None, never as an exception in a generation. +""" + +import dataclasses +import sys +import threading +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Literal, Optional + +import torch + +from invokeai.backend.util import nvml, wddm +from invokeai.backend.util.level_zero import xpu_memory_info +from invokeai.backend.util.logging import InvokeAILogger +from invokeai.backend.util.pci import pci_address + +Backend = Literal["cuda", "hip", "mps", "xpu"] + +Method = Literal["windows-d3dkmt", "nvml", "fdinfo", "mps-driver", "xpu-sysman", "torch-only"] +"""Where a reading's driver figures come from. + +- ``windows-d3dkmt``: this process's ``CurrentUsage`` in the adapter's local segment (`wddm.process_local_usage`), on + CUDA and ROCm under Windows, with the paged bytes from PDH. +- ``nvml``: the ``usedGpuMemory`` NVML attributes to this process, on CUDA under Linux. +- ``fdinfo``: ``drm-memory-vram`` summed over this process's DRM clients on the device, on ROCm under Linux. +- ``mps-driver``: ``torch.mps.driver_allocated_memory()``. +- ``xpu-sysman``: the device's used memory from Level Zero Sysman. Device-wide, so it guards against another process + loading the device during a window and is never taken as demand. +- ``torch-only``: no driver figure; the allocator's counters alone. +""" + +VERIFIED_METHODS: Final[frozenset[tuple[Method, Backend]]] = frozenset( + {("windows-d3dkmt", "cuda"), ("windows-d3dkmt", "hip")} +) +"""The (method, backend) pairs whose readings were checked on real hardware: the RTX 4090 and the RX 9060 XT above. The +other readers follow their vendors' documentation, and a window measured with them says so (``verified=False``).""" + +PEAK_SAMPLER_INTERVAL_S: float = 0.003 +"""How often `PeakSampler` samples. A few milliseconds is far below an operation's allocation lifetimes and costs one +getter call per sample.""" + +_FDINFO_DIR = Path("/proc/self/fdinfo") +_FDINFO_UNITS: Final[Mapping[str, int]] = {"": 1, "KiB": 2**10, "MiB": 2**20, "GiB": 2**30} + +# torch.cuda.memory_stats / torch.xpu.memory_stats keys, per CounterReading field. +_STAT_KEYS: Final[Mapping[str, str]] = { + "reserved": "reserved_bytes.all.current", + "allocated": "allocated_bytes.all.current", + "peak_reserved": "reserved_bytes.all.peak", + "peak_allocated": "allocated_bytes.all.peak", + "num_ooms": "num_ooms", + "num_alloc_retries": "num_alloc_retries", +} + + +@dataclass(frozen=True) +class CounterReading: + """One reading of a device's counters for this process, in bytes; None where the platform has no such counter or + its reader failed.""" + + reserved: Optional[int] + allocated: Optional[int] + peak_reserved: Optional[int] + """Since the last `reset_peaks`, like the other peaks.""" + peak_allocated: Optional[int] + num_ooms: Optional[int] + """Allocations that failed, cumulative for the process.""" + num_alloc_retries: Optional[int] + """Allocations that only succeeded after the allocator released its cache and retried, cumulative. A retry + squeezes the peak below what the operation would have taken with room to spare.""" + driver: Optional[int] + """This process's memory on the device as the driver counts it.""" + peak_driver: Optional[int] + """Only where a sampler keeps it (MPS); elsewhere the driver is read at boundaries.""" + paged: Optional[int] + """This process's bytes Windows keeps in shared system memory for the device's adapter.""" + device_used: Optional[int] + """Used memory of the whole device, every process together (XPU).""" + method: Method + + +def backend_for(device: torch.device) -> Optional[Backend]: + """The GPU backend `device` runs on; None for a device that is not a GPU. A HIP build reports ``cuda`` devices.""" + if device.type == "cuda": + return "hip" if torch.version.hip is not None else "cuda" + if device.type in ("mps", "xpu"): + return device.type # type: ignore[return-value] + return None + + +def _allocator_stats(device: torch.device) -> Mapping[str, int]: + try: + if device.type == "cuda": + return torch.cuda.memory_stats(device) + if device.type == "xpu": + return torch.xpu.memory_stats(device) + if device.type == "mps": + return {_STAT_KEYS["allocated"]: torch.mps.current_allocated_memory()} + except Exception: + InvokeAILogger.get_logger(__name__).debug(f"Allocator statistics for {device} unavailable", exc_info=True) + return {} + + +def _fdinfo_size(value: str) -> Optional[int]: + number, _, unit = value.strip().partition(" ") + scale = _FDINFO_UNITS.get(unit.strip()) + return int(number) * scale if number.isdigit() and scale is not None else None + + +def fdinfo_vram_bytes(pdev: str) -> Optional[int]: + """VRAM this process holds on the DRM device at PCI address `pdev`, from ``/proc/self/fdinfo``. + + Each open DRM file reports its client's ``drm-memory-vram``. Duplicated descriptors share a client and report the + same memory, so each ``drm-client-id`` counts once. None when no descriptor of this process is a client of that + device. + """ + try: + entries = list(_FDINFO_DIR.iterdir()) + except OSError: + return None + total: Optional[int] = None + clients: set[str] = set() + for entry in entries: + try: + text = entry.read_text(encoding="utf-8", errors="replace") + except OSError: + # Closed since the listing, or not readable. + continue + fields = {key.strip(): value.strip() for key, _, value in (line.partition(":") for line in text.splitlines())} + if fields.get("drm-pdev", "").lower() != pdev.lower(): + continue + client = fields.get("drm-client-id") + if client is not None: + if client in clients: + continue + clients.add(client) + vram = _fdinfo_size(fields.get("drm-memory-vram", "")) + if vram is not None: + total = vram if total is None else total + vram + return total + + +def _driver_reading(device: torch.device) -> tuple[Optional[int], Optional[int], Optional[int], Method]: + """``(driver, paged, device_used, method)`` for `device` on this platform.""" + backend = backend_for(device) + driver: Optional[int] = None + paged: Optional[int] = None + device_used: Optional[int] = None + method: Method = "torch-only" + try: + if backend in ("cuda", "hip") and sys.platform == "win32": + driver, paged, method = wddm.process_local_usage(device), wddm.process_paged_bytes(device), "windows-d3dkmt" + elif backend == "cuda" and sys.platform.startswith("linux"): + driver, method = nvml.process_used_memory(device), "nvml" + elif backend == "hip" and sys.platform.startswith("linux"): + index = device.index if device.index is not None else torch.cuda.current_device() + driver, method = fdinfo_vram_bytes(pci_address(index)), "fdinfo" + elif backend == "mps": + driver, method = int(torch.mps.driver_allocated_memory()), "mps-driver" + elif backend == "xpu": + info = xpu_memory_info(device) + device_used = None if info is None else info[1] - info[0] + method = "xpu-sysman" + except Exception: + InvokeAILogger.get_logger(__name__).debug(f"Driver memory reading for {device} failed", exc_info=True) + return None, None, None, "torch-only" + if driver is None and device_used is None: + method = "torch-only" + return driver, paged, device_used, method + + +def read_counters(device: torch.device) -> CounterReading: + """This process's counters on `device` now. Peaks are torch's, since the last `reset_peaks`; MPS has none.""" + stats = _allocator_stats(device) + values = {field: (None if key not in stats else int(stats[key])) for field, key in _STAT_KEYS.items()} + driver, paged, device_used, method = _driver_reading(device) + return CounterReading( + **values, driver=driver, peak_driver=None, paged=paged, device_used=device_used, method=method + ) + + +def reset_peaks(device: torch.device) -> None: + """Start torch's peak counters over at the current values. Cumulative counters (OOMs, retries) are kept.""" + try: + if device.type == "cuda": + torch.cuda.reset_peak_memory_stats(device) + elif device.type == "xpu": + torch.xpu.reset_peak_memory_stats(device) + except Exception: + InvokeAILogger.get_logger(__name__).debug(f"Peak reset for {device} failed", exc_info=True) + + +def synchronize(device: torch.device) -> None: + """Wait for the work queued on `device`, so frees it implies have happened.""" + try: + if device.type == "cuda": + torch.cuda.synchronize(device) + elif device.type == "xpu": + torch.xpu.synchronize(device) + elif device.type == "mps": + torch.mps.synchronize() + except Exception: + InvokeAILogger.get_logger(__name__).debug(f"Synchronizing {device} failed", exc_info=True) + + +def _larger(a: Optional[int], b: Optional[int]) -> Optional[int]: + if a is None: + return b + return a if b is None else max(a, b) + + +class PeakSampler: + """The largest values `read` returned since the last `reset`, sampled every `PEAK_SAMPLER_INTERVAL_S` on a daemon + thread between `start` and `stop`. + + For a device whose counters have no peaks: without it, a window would see only the values at its boundaries and + miss everything that came and went in between. `read` returns one value per counter, None where it has none; a + `read` that raises is a sample of Nones. + """ + + def __init__(self, read: Callable[[], tuple[Optional[int], ...]]) -> None: + self._read = read + self._lock = threading.Lock() + self._peaks: tuple[Optional[int], ...] = () + self._stop = threading.Event() + self._thread: Optional[threading.Thread] = None + + def _sample(self) -> tuple[Optional[int], ...]: + try: + return self._read() + except Exception: + return () + + def _fold(self, sample: tuple[Optional[int], ...]) -> tuple[Optional[int], ...]: + with self._lock: + width = max(len(self._peaks), len(sample)) + peaks = self._peaks + (None,) * (width - len(self._peaks)) + sample = sample + (None,) * (width - len(sample)) + self._peaks = tuple(_larger(peak, value) for peak, value in zip(peaks, sample, strict=True)) + return self._peaks + + def _run(self) -> None: + while not self._stop.wait(PEAK_SAMPLER_INTERVAL_S): + self._fold(self._sample()) + + def start(self) -> None: + self.reset() + self._stop.clear() + self._thread = threading.Thread(target=self._run, name="working-memory-peak-sampler", daemon=True) + self._thread.start() + + def stop(self) -> None: + self._stop.set() + if self._thread is not None: + self._thread.join() + self._thread = None + + @property + def running(self) -> bool: + return self._thread is not None and self._thread.is_alive() + + def reset(self) -> None: + """Start the peaks over at the values now.""" + sample = self._sample() + with self._lock: + self._peaks = sample + + def peaks(self) -> tuple[Optional[int], ...]: + """The peaks since the last reset, the values now included.""" + return self._fold(self._sample()) + + +def _mps_values() -> tuple[Optional[int], ...]: + return int(torch.mps.current_allocated_memory()), int(torch.mps.driver_allocated_memory()) + + +class DeviceCounters: + """One device's counters for the length of a window: readings, peak resets and synchronization. + + The peaks are torch's where it keeps them. On MPS, which keeps none, a `PeakSampler` keeps the allocated and the + driver peaks from construction until `close`. + """ + + def __init__(self, device: torch.device) -> None: + self.device = device + self._sampler = PeakSampler(_mps_values) if device.type == "mps" else None + if self._sampler is not None: + self._sampler.start() + + def read(self) -> CounterReading: + reading = read_counters(self.device) + if self._sampler is None: + return reading + peak_allocated, peak_driver = (self._sampler.peaks() + (None, None))[:2] + return dataclasses.replace( + reading, + peak_allocated=_larger(peak_allocated, reading.allocated), + peak_driver=_larger(peak_driver, reading.driver), + ) + + def reset_peaks(self) -> None: + reset_peaks(self.device) + if self._sampler is not None: + self._sampler.reset() + + def synchronize(self) -> None: + synchronize(self.device) + + def close(self) -> None: + if self._sampler is not None: + self._sampler.stop() diff --git a/invokeai/backend/working_memory/fingerprint.py b/invokeai/backend/working_memory/fingerprint.py new file mode 100644 index 00000000000..5f8f057ec7a --- /dev/null +++ b/invokeai/backend/working_memory/fingerprint.py @@ -0,0 +1,162 @@ +"""What a measurement is keyed by: the class of device it ran on, and the software environment that decides how much +memory an operation takes there. + +A device class names a kind of card, not one card: identical cards share it, whatever their index or PCI slot, so +measurements taken on one apply to the other. An environment splits into hard facts, any change of which can change an +operation's memory (torch and runtime versions, the convolution library, the allocator and its configuration, the +attention path, InvokeAI's memory settings), and soft facts recorded for diagnosis only (full versions, the driver, the +OS build). Both are produced deterministically here; which change makes a measurement stale is the caller's decision. +""" + +import ctypes +import ctypes.util +import hashlib +import json +import os +import platform +from collections.abc import Mapping +from typing import Final, Optional + +import psutil +import torch +from pydantic import BaseModel, ConfigDict, computed_field + +from invokeai.backend.util import nvml +from invokeai.backend.util.attention import rocm_sdpa_chunks_math +from invokeai.backend.working_memory.counters import Backend, backend_for +from invokeai.version import __version__ + +_VENDORS: Final[Mapping[Backend, str]] = {"cuda": "nvidia", "hip": "amd", "mps": "apple", "xpu": "intel"} + +HARD_ENVIRONMENT_VARIABLES: Final = ( + "PYTORCH_ALLOC_CONF", + "PYTORCH_CUDA_ALLOC_CONF", + "PYTORCH_HIP_ALLOC_CONF", + "MIOPEN_FIND_MODE", + "PYTORCH_TUNABLEOP_ENABLED", + # The head width above which the ROCm SDPA guard sends attention to the chunked math kernel (`util/attention.py`). + "INVOKE_ROCM_FUSED_SDPA_MAX_HEAD_DIM", + # Whether AOTriton's fused attention runs on GPUs PyTorch marks experimental; without it attention takes the math + # kernel. What `rocm_aotriton_experimental` resolved to (`app/util/rocm_aotriton.py`), or a value set beforehand. + "TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", + # Native or decomposed causal 3D convolutions on ROCm (`backend/wan/rocm_causal_conv3d.py`). + "INVOKEAI_ROCM_CONV3D", +) +"""Environment variables that change how much memory an operation takes: the allocator's configuration, MIOpen's +algorithm search, TunableOp's GEMM selection, and the ROCm attention and convolution kernel choices. Recorded as +``env.``, empty when unset.""" + + +def _short_hash(fields: Mapping[str, object]) -> str: + canonical = json.dumps(fields, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:16] + + +class DeviceClass(BaseModel): + """A kind of compute device. Identical cards share one, whatever their index or PCI location.""" + + model_config = ConfigDict(frozen=True) + + vendor: str + name: str + total_mib: int + backend: Backend + arch: Optional[str] + """``gcnArchName`` on HIP (e.g. ``gfx1200``), ``sm_`` on CUDA, the reported architecture on XPU.""" + + @computed_field # type: ignore[prop-decorator] + @property + def id(self) -> str: + """A short stable hash of the other fields. Derived, never given.""" + return _short_hash(self.model_dump(exclude={"id"})) + + +def _mps_name() -> str: + """The Mac's chip, e.g. ``Apple M2 Max``; MPS names no device.""" + try: + libc = ctypes.CDLL(ctypes.util.find_library("c")) + size = ctypes.c_size_t(0) + if libc.sysctlbyname(b"machdep.cpu.brand_string", None, ctypes.byref(size), None, 0) == 0 and size.value: + buffer = ctypes.create_string_buffer(size.value) + if libc.sysctlbyname(b"machdep.cpu.brand_string", buffer, ctypes.byref(size), None, 0) == 0: + return buffer.value.decode("utf-8", errors="replace") + except Exception: + pass + return "Apple GPU" + + +def device_class_for(device: torch.device) -> DeviceClass: + """The class of `device`. Raises ValueError for a device that is not a GPU.""" + backend = backend_for(device) + if backend is None: + raise ValueError(f"{device} is not a GPU; only GPU devices have a device class.") + if backend in ("cuda", "hip"): + props = torch.cuda.get_device_properties(device) + name, total = torch.cuda.get_device_name(device), int(props.total_memory) + arch = str(props.gcnArchName) if backend == "hip" else f"sm_{props.major}{props.minor}" + elif backend == "xpu": + props = torch.xpu.get_device_properties(device) + name, total = str(props.name), int(props.total_memory) + architecture = getattr(props, "architecture", None) + arch = None if architecture is None else str(architecture) + else: + # Unified memory: the device's share is the machine's RAM. + name, total, arch = _mps_name(), int(psutil.virtual_memory().total), None + return DeviceClass(vendor=_VENDORS[backend], name=name, total_mib=total // 2**20, backend=backend, arch=arch) + + +class Environment(BaseModel): + """The software a measurement ran under: hard facts that change an operation's memory, soft facts that do not.""" + + model_config = ConfigDict(frozen=True) + + hard: dict[str, str] + soft: dict[str, Optional[str]] + + @computed_field # type: ignore[prop-decorator] + @property + def id(self) -> str: + """A short stable hash of the hard facts. Soft facts do not enter it.""" + return _short_hash(self.hard) + + +def _major_minor(version: str) -> str: + return ".".join(version.split("+", 1)[0].split(".")[:2]) + + +def _conv_library_major() -> Optional[str]: + """The major version of cuDNN, or of MIOpen on a HIP build, as loaded at runtime.""" + try: + if torch.backends.cudnn.version() is None: + return None + # (major, minor, patch) on either build; the packed integer `version()` returns is encoded per library. + return str(torch._C._cudnn.getRuntimeVersion()[0]) + except Exception: + return None + + +def environment_for(device: torch.device, settings: Mapping[str, str | int | float | bool | None]) -> Environment: + """The environment an operation on `device` runs under. Every key of `settings` is a hard fact, as + ``setting.``: the caller passes the InvokeAI settings that change memory.""" + backend = backend_for(device) + hard: dict[str, str] = {"torch": _major_minor(torch.__version__)} + if backend in ("cuda", "hip"): + runtime = torch.version.hip if backend == "hip" else torch.version.cuda + hard[backend] = _major_minor(runtime or "") + hard["miopen" if backend == "hip" else "cudnn"] = _conv_library_major() or "" + hard["allocator_backend"] = torch.cuda.get_allocator_backend() + hard["rocm_sdpa_guard"] = str(rocm_sdpa_chunks_math(device.type)).lower() + elif backend == "xpu": + hard["xpu"] = str(getattr(torch.version, "xpu", None) or "") + for name in HARD_ENVIRONMENT_VARIABLES: + hard[f"env.{name}"] = os.environ.get(name, "") + for key, value in settings.items(): + hard[f"setting.{key}"] = str(value) + + soft: dict[str, Optional[str]] = { + "torch_build": torch.__version__, + "driver": nvml.driver_version() if backend == "cuda" else None, + "os": platform.platform(), + "invokeai": __version__, + } + return Environment(hard=hard, soft=soft) diff --git a/invokeai/backend/working_memory/probe.py b/invokeai/backend/working_memory/probe.py new file mode 100644 index 00000000000..bd888590300 --- /dev/null +++ b/invokeai/backend/working_memory/probe.py @@ -0,0 +1,729 @@ +"""The working-memory probe: what an operation really took on its device, recorded next to what was estimated for it. + +Active only inside a `calibration_scope`, which the session runner opens around each node of a calibration queue item. +There, every `LoadedModel.model_on_device(working_memory=...)` on a GPU is a *window*: the model is locked with +calibration headroom (`calibration_reservation`, so that a too-low estimate cannot spill and could never be corrected), +and from the moment the lock has settled until just before the unlock, the probe measures the device. The window's +`MeasuredWindow` lands in the scope's `windows`; nothing is written anywhere while a window is open. + +What a window measures, per counter (driver, torch reserved, torch allocated; see `counters`): + +1. Open: synchronize, release the allocator's cached blocks, read the baseline, reset the peaks. Tensors allocated + before the lock are part of the baseline: in use, as they are in the model cache's own budget. +2. A *co-lock* -- another model locked on the same thread and device while the window is open (a ControlNet, an audio + VAE, the vocoder) -- moves weights, not working memory. The segment ends before its lock (R0, read after a release) + and a new one starts after it settled (R2, after a release); R2 - R0 joins the window's offset, negative when the + lock moved more off the device than on. +3. A *nested window* (an estimated `model_on_device` inside an open one, such as a VAE encode inside a transformer's + window): its lock is a co-lock for the parent, the parent is frozen while it is open, and it measures alone. When + it closes, the parent releases the child's cached blocks, resets the peaks and goes on; the child's lasting outputs + stay in the parent's figures. +4. A `TorchDevice.empty_cache()` inside the window ends a segment: the driver is read before the release. A direct + `torch.cuda.empty_cache()` (a library's, or cuDNN's and MIOpen's after an algorithm search) and an offload on the + cache's background thread do not, so the driver figure can miss what they released; torch's peaks still bound it. +5. Close: synchronize and read the final segment. Per counter, the delta is the largest segment peak less the + baseline and that segment's offset; `demand_bytes` is the largest delta of the counters that could be read. + +A window is invalid (`valid=False`, with its reasons) when another thread locked a model on the same device's cache +while it was open or held one when it opened (Expand Prompt and Image to Prompt run next to sessions), when an +allocation failed or only succeeded after the allocator released its cache (a retry squeezes the peak: an untiled +3072px FLUX VAE decode on an RTX 4090 read 22.5 GiB at the card's limit, ~1280 bytes per pixel instead of ~2150), when +Windows still keeps more than `SPILL_TOLERANCE_BYTES` of the process in shared memory `SPILL_SETTLE_S` after close +(right after a decode the counter spikes transiently, 589 MiB, and is back at 9 MiB at the next read), when the +operation raised or was canceled, when the Krea-2 SDPA benchmark (which resets the peaks itself) is active, when XPU's +device-wide usage grew beyond this process's own, when a probe step failed (`probe_error`; a failed lock hook also voids +the scope's later windows, since its lock bookkeeping is suspect), and when no counter could be read. A model less than fully resident at open keeps the window valid but not `fit_eligible`: streamed +weights are working memory the estimate does not price. + +Without a scope, the probe costs a context-variable read per estimated `model_on_device`, a read of `armed` per lock and +unlock, and a truth test of `devices.EMPTY_CACHE_OBSERVERS` per `empty_cache`. +""" + +import contextvars +import os +import threading +import time +from collections.abc import Iterator +from contextlib import contextmanager +from typing import TYPE_CHECKING, Final, Literal, Optional + +from pydantic import BaseModel, ConfigDict + +from invokeai.app.services.session_processor.session_processor_common import CanceledException +from invokeai.backend.util import devices +from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.logging import InvokeAILogger +from invokeai.backend.working_memory import counters +from invokeai.backend.working_memory.counters import ( + VERIFIED_METHODS, + CounterReading, + DeviceCounters, + Method, + backend_for, +) +from invokeai.backend.working_memory.estimate import FeatureValue, WorkingMemoryEstimate, WorkingMemoryTerm + +if TYPE_CHECKING: + from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord + +GIB = 2**30 +MIB = 2**20 + +CALIBRATION_HEADROOM_MIN_BYTES: Final = 1 * GIB +CALIBRATION_HEADROOM_FRACTION: Final = 0.5 + +SPILL_TOLERANCE_BYTES: Final = 256 * MIB +"""Shared-memory bytes above the window's baseline that count as spill once they persist.""" +SPILL_SETTLE_S: float = 0.5 +"""How long a close waits before re-reading a raised shared-memory counter. Module-level so tests can shorten it.""" +FOREIGN_DEVICE_TOLERANCE_BYTES: Final = 256 * MIB +"""Growth of a device-wide counter (XPU) beyond this process's own reserved growth that marks another process's load.""" + +Reason = Literal[ + "foreign_lock", + "foreign_device_load", + "oom", + "alloc_retry", + "spill", + "exception", + "canceled", + "krea2_sdpa_benchmark", + "counters_unavailable", + "probe_error", + "partial_residency", +] +"""Why a window is invalid, or -- ``partial_residency`` alone -- valid but not fit-eligible.""" + +_NOT_INVALIDATING: Final[frozenset[Reason]] = frozenset({"partial_residency"}) + +CounterName = Literal["driver", "reserved", "allocated"] +_COUNTERS: Final[tuple[CounterName, ...]] = ("driver", "reserved", "allocated") + + +class _Record(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +class CounterDelta(_Record): + """One counter over a window, in bytes.""" + + baseline: int + peak: int + """The largest value any segment reached, offsets not applied.""" + offset: int + """What co-locks moved onto the device by the end of the window; negative when they moved more off.""" + delta: int + """The largest of each segment's peak less the baseline and that segment's offset: the operation's own share.""" + + +class MeasuredWindow(_Record): + """One measured `model_on_device` window: who ran it, what was estimated, and what the device saw.""" + + run_id: str + queue_item_id: int + node_id: str + node_type: str + model_key: str + device: str + + estimator: str + version: int + features: dict[str, FeatureValue] + terms: tuple[WorkingMemoryTerm, ...] + reservation_bytes: int + source: str + reservation_applied_bytes: int + """What the lock reserved: the estimate plus calibration headroom.""" + + driver: Optional[CounterDelta] + reserved: Optional[CounterDelta] + allocated: Optional[CounterDelta] + demand_bytes: Optional[int] + """The largest delta of the counters that could be read.""" + method: Method + """Where the driver figures came from; ``torch-only`` when there were none for the whole window.""" + verified: bool + """Whether that method was verified on this backend on real hardware (`counters.VERIFIED_METHODS`).""" + + resident_fraction: float + """The share of the window's model's weights on the device when the window opened.""" + co_locked_models: list[str] + """The other models this thread held on the device while the window was open.""" + segments: int + duration_s: float + + valid: bool + fit_eligible: bool + reasons: list[Reason] + + +def calibration_reservation(reservation_bytes: int) -> int: + """What a calibration lock reserves for an estimate of `reservation_bytes`: half as much again, at least 1 GiB more.""" + return reservation_bytes + max( + CALIBRATION_HEADROOM_MIN_BYTES, int(CALIBRATION_HEADROOM_FRACTION * reservation_bytes) + ) + + +# All window state, and the registry of open windows, under one lock. Re-entrant because a window's own release runs +# `TorchDevice.empty_cache`, which notifies the probe on the same thread. Lock order: the model cache's lock before +# this one, never the other way round -- `TorchDevice.empty_cache` runs inside cache methods, so nothing here calls into +# a cache while holding it. +_state_lock = threading.RLock() +_open_windows: list["_Window"] = [] +_armed_scopes = 0 + +armed: bool = False +"""True while any thread has a calibration scope open. The model loader checks it before telling the probe about a lock +or an unlock, so without a scope that costs one attribute read.""" + +_boundary = threading.local() + + +def _release() -> None: + """Return the allocator's cached blocks to the driver as a boundary of this thread's windows, which handle it + themselves; other threads' windows see it like any other release.""" + _boundary.active = True + try: + TorchDevice.empty_cache(force=True) + finally: + _boundary.active = False + + +def _value(reading: CounterReading, name: CounterName) -> Optional[int]: + return reading.driver if name == "driver" else reading.reserved if name == "reserved" else reading.allocated + + +def _peak(reading: CounterReading, name: CounterName) -> Optional[int]: + if name == "driver": + return reading.peak_driver if reading.peak_driver is not None else reading.driver + return reading.peak_reserved if name == "reserved" else reading.peak_allocated + + +def _rose(before: Optional[int], after: Optional[int]) -> bool: + return before is not None and after is not None and after > before + + +def _krea2_sdpa_benchmark_active() -> bool: + # Imported here: the Krea-2 attention module pulls in diffusers, which a lock must not. + from invokeai.backend.krea2.attention import KREA2_SDPA_BACKEND_ENV_VAR + + # Any non-empty value runs the benchmark (`_Krea2StepBenchmark.create`), or fails the node if it is invalid. + return bool(os.environ.get(KREA2_SDPA_BACKEND_ENV_VAR, "").strip()) + + +class _Window: + """An open window's measurement. Mutated under `_state_lock` only.""" + + def __init__( + self, + scope: "ProbeScope", + cache: object, + record: "CacheRecord", + estimate: WorkingMemoryEstimate, + reservation_applied_bytes: int, + ) -> None: + self.scope = scope + self.cache = cache + self.record = record + self.estimate = estimate + self.reservation_applied_bytes = reservation_applied_bytes + self.device = record.cached_model.compute_device + self.thread = threading.get_ident() + self.parent: Optional[_Window] = None + self.counters: Optional[DeviceCounters] = None + self.state: Literal["measuring", "co_lock", "suspended"] = "measuring" + self.first: Optional[CounterReading] = None + self.baseline: dict[CounterName, Optional[int]] = {} + self.offset: dict[CounterName, int] = dict.fromkeys(_COUNTERS, 0) + self.peak: dict[CounterName, int] = {} + self.delta: dict[CounterName, int] = {} + self.unavailable: set[CounterName] = set() + self.reasons: list[Reason] = [] + self.segments = 0 + self.pending_reset = False + self.co_lock_record: Optional["CacheRecord"] = None + self.co_lock_start: Optional[CounterReading] = None + self.co_locked: list[str] = [] + self.resident_fraction = 1.0 + self.opened_at = 0.0 + self.recorded = False + + def add_reason(self, reason: Reason) -> None: + if reason not in self.reasons: + self.reasons.append(reason) + + def fail(self) -> None: + """A probe step on this window failed: its figures can no longer be trusted. Mark it, and leave it in a state + the rest of the window can run in.""" + self.add_reason("probe_error") + self.pending_reset = False + if self.state == "co_lock": + self.state = "measuring" + self.co_lock_record = self.co_lock_start = None + + def start(self, reading: CounterReading) -> None: + self.first = reading + self.opened_at = time.perf_counter() + for name in _COUNTERS: + self.baseline[name] = _value(reading, name) + if self.baseline[name] is None: + self.unavailable.add(name) + + def fold(self, reading: CounterReading, ends_segment: bool = True) -> None: + """Take `reading` as the end of the current segment, or with `ends_segment=False` as a sample within it.""" + for name in _COUNTERS: + if name in self.unavailable: + continue + value, base = _peak(reading, name), self.baseline[name] + if value is None or base is None: + self.unavailable.add(name) + continue + self.peak[name] = max(self.peak.get(name, value), value) + share = value - base - self.offset[name] + self.delta[name] = max(self.delta.get(name, share), share) + if ends_segment: + self.segments += 1 + + def begin_co_lock(self, record: "CacheRecord") -> None: + assert self.counters is not None + self.counters.synchronize() + self.fold(self.counters.read()) + _release() + self.co_lock_start = self.counters.read() + self.co_lock_record = record + self.state = "co_lock" + + def end_co_lock(self) -> None: + assert self.counters is not None and self.co_lock_start is not None and self.co_lock_record is not None + self.counters.synchronize() + _release() + settled = self.counters.read() + for name in _COUNTERS: + before, after = _value(self.co_lock_start, name), _value(settled, name) + if before is None or after is None: + self.unavailable.add(name) + else: + self.offset[name] += after - before + if self.co_lock_record.key not in self.co_locked and self.co_lock_record.key != self.record.key: + self.co_locked.append(self.co_lock_record.key) + self.co_lock_record = self.co_lock_start = None + self.counters.reset_peaks() + self.state = "measuring" + + def counter(self, name: CounterName) -> Optional[CounterDelta]: + base = self.baseline.get(name) + if name in self.unavailable or base is None or name not in self.delta: + return None + return CounterDelta(baseline=base, peak=self.peak[name], offset=self.offset[name], delta=self.delta[name]) + + def method(self) -> Method: + if self.first is None: + return "torch-only" + # A driver figure that dropped out mid-window leaves the window to torch's counters. + if self.first.driver is not None and "driver" in self.unavailable: + return "torch-only" + return self.first.method + + def measured(self, duration_s: float) -> MeasuredWindow: + counter_deltas = {name: self.counter(name) for name in _COUNTERS} + deltas = [delta.delta for delta in counter_deltas.values() if delta is not None] + if not deltas: + self.add_reason("counters_unavailable") + valid = not any(reason not in _NOT_INVALIDATING for reason in self.reasons) + method = self.method() + backend = backend_for(self.device) + estimate, scope = self.estimate, self.scope + return MeasuredWindow( + run_id=scope.run_id, + queue_item_id=scope.queue_item_id, + node_id=scope.node_id, + node_type=scope.node_type, + model_key=self.record.key, + device=str(self.device), + estimator=estimate.estimator, + version=estimate.version, + features=dict(estimate.features), + terms=estimate.terms, + reservation_bytes=estimate.reservation_bytes, + source=estimate.source, + reservation_applied_bytes=self.reservation_applied_bytes, + driver=counter_deltas["driver"], + reserved=counter_deltas["reserved"], + allocated=counter_deltas["allocated"], + demand_bytes=max(deltas) if deltas else None, + method=method, + verified=backend is not None and (method, backend) in VERIFIED_METHODS, + resident_fraction=self.resident_fraction, + co_locked_models=list(self.co_locked), + segments=self.segments, + duration_s=duration_s, + valid=valid, + fit_eligible=valid and "partial_residency" not in self.reasons, + reasons=list(self.reasons), + ) + + def close(self, error: Optional[BaseException]) -> None: + """End the window on its own thread, `error` being what the operation raised, and record it in its scope.""" + try: + self._close(error) + except Exception: + InvokeAILogger.get_logger(__name__).warning( + f"Working-memory probe failed to close the window of {self.record.key}", exc_info=True + ) + try: + with _state_lock: + self.add_reason("probe_error") + self._unregister() + parent = self.parent + if parent is not None and parent.state == "suspended": + # Its peaks were not reset after this window, so what it measures from here is suspect. + parent.fail() + parent.state = "measuring" + if not self.recorded: + self.scope.windows.append(self.measured(time.perf_counter() - self.opened_at)) + self.recorded = True + except Exception: + InvokeAILogger.get_logger(__name__).debug("Working-memory probe could not record", exc_info=True) + + def _unregister(self) -> None: + if self in _open_windows: + _open_windows.remove(self) + if self in self.scope._stack: + self.scope._stack.remove(self) + if self.counters is not None: + self.counters.close() + + def _close(self, error: Optional[BaseException]) -> None: + if error is not None: + self.add_reason("canceled" if isinstance(error, CanceledException) else "exception") + final: Optional[CounterReading] = None + with _state_lock: + if self.counters is not None and self.first is not None: + self.counters.synchronize() + final = self.counters.read() + self.fold(final) + first = self.first + if _rose(first.num_ooms, final.num_ooms): + self.add_reason("oom") + if _rose(first.num_alloc_retries, final.num_alloc_retries): + self.add_reason("alloc_retry") + if ( + first.device_used is not None + and final.device_used is not None + and first.reserved is not None + and final.reserved is not None + and (final.device_used - first.device_used) - (final.reserved - first.reserved) + > FOREIGN_DEVICE_TOLERANCE_BYTES + ): + self.add_reason("foreign_device_load") + if self in _open_windows: + _open_windows.remove(self) + if self in self.scope._stack: + self.scope._stack.remove(self) + duration_s = time.perf_counter() - self.opened_at + + if final is not None and self.first is not None and self.counters is not None: + baseline_paged = self.first.paged + if ( + baseline_paged is not None + and final.paged is not None + and final.paged - baseline_paged > SPILL_TOLERANCE_BYTES + ): + # Shared usage spikes for a moment after large transfers; only what is still there later is spill. + time.sleep(SPILL_SETTLE_S) + settled = self.counters.read().paged + if settled is None or settled - baseline_paged > SPILL_TOLERANCE_BYTES: + self.add_reason("spill") + + with _state_lock: + parent = self.parent + if parent is not None and parent in _open_windows and parent.counters is not None: + _release() + parent.counters.reset_peaks() + parent.state = "measuring" + if self.counters is not None: + self.counters.close() + record = self.measured(duration_s) + self.scope.windows.append(record) + self.recorded = True + InvokeAILogger.get_logger(__name__).debug( + f"Measured {record.estimator} on {record.device} for {record.model_key}: " + f"demand {(record.demand_bytes or 0) / MIB:.0f} MB against {record.reservation_bytes / MIB:.0f} MB " + f"estimated ({record.method}{'' if record.valid else ', invalid: ' + ', '.join(record.reasons)})" + ) + + +class ProbeScope: + """One node of a calibration queue item, on the thread that runs it. `windows` collects that node's measurements + in the order they closed; the caller persists them after the node.""" + + def __init__(self, *, run_id: str, queue_item_id: int, node_id: str, node_type: str) -> None: + self.run_id = run_id + self.queue_item_id = queue_item_id + self.node_id = node_id + self.node_type = node_type + self.windows: list[MeasuredWindow] = [] + # This thread's open windows, innermost last. + self._stack: list[_Window] = [] + # The locks this thread holds, per (cache, model key). + self._held: dict[tuple[int, str], int] = {} + # Set when a probe hook failed on this thread: `_held` may be off, so no later window can trust its check. + self._tainted = False + + def _held_on(self, cache: object) -> int: + return sum(count for (cache_id, _key), count in self._held.items() if cache_id == id(cache)) + + def open_window( + self, cache: object, record: "CacheRecord", estimate: WorkingMemoryEstimate, reservation_applied_bytes: int + ) -> Optional[_Window]: + """Start measuring the operation `record`'s model is about to run, once its lock has settled. + + `cache` is the model cache the record is locked in. Never raises: a probe failure leaves the operation alone + and records the window as invalid, or not at all. + """ + window = _Window(self, cache, record, estimate, reservation_applied_bytes) + # Registered before the cache is asked: from here on, a lock another thread takes on this cache voids the + # window (`lock_started`, `lock_settled`), and one it already holds shows in `locks_held()`. Its counters are + # not open yet, so releases pass it by. + with _state_lock: + _open_windows.append(window) + try: + if self._tainted: + window.add_reason("probe_error") + cached_model = record.cached_model + total = cached_model.total_bytes() + resident = cached_model.cur_vram_bytes() + window.resident_fraction = resident / total if total > 0 else 1.0 + if resident < total: + window.add_reason("partial_residency") + window.co_locked = sorted( + key + for (cache_id, key), count in self._held.items() + if cache_id == id(cache) and count > 0 and key != record.key + ) + # Asked outside the probe's lock (see `_state_lock`). Every lock this thread took since the scope opened + # is in `_held`, so any lock beyond those belongs to another thread. + if cache.locks_held() > self._held_on(cache): # type: ignore[attr-defined] + window.add_reason("foreign_lock") + if window.device.type == "cuda" and _krea2_sdpa_benchmark_active(): + window.add_reason("krea2_sdpa_benchmark") + except Exception: + InvokeAILogger.get_logger(__name__).debug("Working-memory probe could not inspect the lock", exc_info=True) + window.add_reason("probe_error") + + try: + window.counters = counters.DeviceCounters(window.device) + with _state_lock: + parent = self._stack[-1] if self._stack else None + window.parent = parent + window.counters.synchronize() + if parent is not None and parent.counters is not None: + if parent.state == "measuring": + parent.fold(parent.counters.read()) + parent.state = "suspended" + _release() + window.start(window.counters.read()) + window.counters.reset_peaks() + self._stack.append(window) + return window + except Exception: + InvokeAILogger.get_logger(__name__).warning( + f"Working-memory probe could not open a window for {record.key}", exc_info=True + ) + with _state_lock: + window._unregister() + parent = window.parent + if parent is not None and parent.state == "suspended": + # Folded and released for a window that never opened. + parent.fail() + parent.state = "measuring" + return None + + def _abandon(self) -> None: + """Drop the windows still open when the scope ends, so no sampler outlives it.""" + with _state_lock: + for window in list(self._stack): + window._unregister() + + +_current_scope: contextvars.ContextVar[Optional[ProbeScope]] = contextvars.ContextVar( + "working_memory_probe_scope", default=None +) + + +def current_scope() -> Optional[ProbeScope]: + """The calibration scope open on this thread, if any.""" + return _current_scope.get() + + +class _EmptyCacheObserver: + """Turns `TorchDevice.empty_cache` into segment boundaries of the open windows.""" + + def before_empty_cache(self) -> None: + try: + me = threading.get_ident() + own_boundary = getattr(_boundary, "active", False) + with _state_lock: + for window in _open_windows: + if window.counters is None or window.state != "measuring": + continue + if window.thread != me: + # Another thread's release (process-global on CUDA) takes this window's cached blocks too: + # read its driver before they go, but leave its peaks alone mid-operation. + window.fold(window.counters.read(), ends_segment=False) + elif not own_boundary: + window.fold(window.counters.read()) + window.pending_reset = True + except Exception: + _all_failed("follow a release") + + def after_empty_cache(self) -> None: + try: + with _state_lock: + for window in _open_windows: + if window.pending_reset and window.counters is not None: + window.counters.reset_peaks() + window.pending_reset = False + except Exception: + _all_failed("follow a release") + + +_observer = _EmptyCacheObserver() + + +def _arm() -> None: + global armed, _armed_scopes + with _state_lock: + _armed_scopes += 1 + if _armed_scopes == 1: + devices.EMPTY_CACHE_OBSERVERS.append(_observer) + armed = True + + +def _disarm() -> None: + global armed, _armed_scopes + with _state_lock: + _armed_scopes -= 1 + if _armed_scopes == 0: + armed = False + if _observer in devices.EMPTY_CACHE_OBSERVERS: + devices.EMPTY_CACHE_OBSERVERS.remove(_observer) + + +@contextmanager +def calibration_scope(*, run_id: str, queue_item_id: int, node_id: str, node_type: str) -> Iterator[ProbeScope]: + """Measure the windows of one node on this thread. Scopes do not nest.""" + if _current_scope.get() is not None: + raise RuntimeError("A calibration scope is already open on this thread; scopes do not nest.") + scope = ProbeScope(run_id=run_id, queue_item_id=queue_item_id, node_id=node_id, node_type=node_type) + token = _current_scope.set(scope) + _arm() + try: + yield scope + finally: + scope._abandon() + _disarm() + _current_scope.reset(token) + + +def _all_failed(what: str) -> None: + """A release boundary failed part-way: any open window may have been folded without its reset.""" + InvokeAILogger.get_logger(__name__).debug(f"Working-memory probe failed to {what}", exc_info=True) + try: + with _state_lock: + for window in _open_windows: + window.fail() + except Exception: + pass + + +def _hook_failed(scope: Optional[ProbeScope], cache: object, what: str) -> None: + """A lock hook failed part-way. On a scope's own thread every window of the scope is suspect, and so is its lock + bookkeeping; on another thread, the windows on that cache which the hook was there to void.""" + InvokeAILogger.get_logger(__name__).debug(f"Working-memory probe failed to {what}", exc_info=True) + try: + with _state_lock: + if scope is not None: + scope._tainted = True + affected = list(scope._stack) + else: + affected = [window for window in _open_windows if window.cache is cache] + for window in affected: + window.fail() + except Exception: + pass + + +def _invalidate_foreign(cache: object) -> None: + me = threading.get_ident() + for window in _open_windows: + if window.cache is cache and window.thread != me: + window.add_reason("foreign_lock") + + +def lock_started(cache: object, record: "CacheRecord") -> None: + """A lock of `record` on its compute device begins on this thread. Called by the model loader while `armed`.""" + scope = _current_scope.get() + try: + with _state_lock: + _invalidate_foreign(cache) + if scope is None or not scope._stack: + return + window = scope._stack[-1] + if ( + window.cache is cache + and window.state == "measuring" + and record.cached_model.compute_device.type != "cpu" + ): + window.begin_co_lock(record) + except Exception: + _hook_failed(scope, cache, "follow a lock") + + +def lock_settled(cache: object, record: "CacheRecord", locked: bool) -> None: + """That lock has settled, its required tensors repaired, or it failed (`locked=False`, nothing stays locked).""" + scope = _current_scope.get() + try: + with _state_lock: + _invalidate_foreign(cache) + if scope is None: + return + if locked: + key = (id(cache), record.key) + scope._held[key] = scope._held.get(key, 0) + 1 + if scope._stack: + window = scope._stack[-1] + if window.state == "co_lock" and window.co_lock_record is record: + window.end_co_lock() + except Exception: + _hook_failed(scope, cache, "follow a lock") + + +def locked_in_ram(cache: object, record: "CacheRecord") -> None: + """`record` was pinned in RAM on this thread: no device memory moved, but the lock counts for `locks_held`.""" + scope = _current_scope.get() + try: + with _state_lock: + _invalidate_foreign(cache) + if scope is not None: + key = (id(cache), record.key) + scope._held[key] = scope._held.get(key, 0) + 1 + except Exception: + _hook_failed(scope, cache, "follow a lock") + + +def unlocked(cache: object, record: "CacheRecord") -> None: + """`record` was unlocked on this thread.""" + scope = _current_scope.get() + try: + with _state_lock: + _invalidate_foreign(cache) + if scope is not None: + key = (id(cache), record.key) + remaining = scope._held.get(key, 0) - 1 + if remaining > 0: + scope._held[key] = remaining + else: + scope._held.pop(key, None) + except Exception: + _hook_failed(scope, cache, "follow an unlock") diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json index a78b1e299b2..f669b2d5e72 100644 --- a/invokeai/frontend/api/openapi.json +++ b/invokeai/frontend/api/openapi.json @@ -5660,6 +5660,351 @@ ] } }, + "/api/v2/models/i/{key}/working_memory/calibrate": { + "post": { + "tags": ["model_manager"], + "summary": "Calibrate Model Working Memory", + "description": "Start a working-memory calibration run: enqueue the batch with every node uncached and measure each node's\nmodel operations on the generation device. Needs exactly one generation device, and it must be a GPU.\n\nWith an `idempotency_key`, repeating the request (after a lost response) returns the run it created and enqueues\nnothing new; acknowledge the key through the queue's enqueue_batch/acknowledge route once the response arrived.", + "operationId": "calibrate_model_working_memory", + "security": [ + { + "HTTPBearer": [] + } + ], + "parameters": [ + { + "name": "key", + "in": "path", + "required": true, + "schema": { + "type": "string", + "description": "Key of the model the run calibrates", + "title": "Key" + }, + "description": "Key of the model the run calibrates" + } + ], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/WorkingMemoryCalibrationRequest", + "description": "The batch to measure" + } + } + } + }, + "responses": { + "201": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/WorkingMemoryCalibrationRunCreated" + } + } + } + }, + "404": { + "description": "The model could not be found" + }, + "409": { + "description": "This machine cannot calibrate now (no GPU, or several generation devices), the device class's store file is unusable and kept as it is, the queue is full, or the idempotency key was used for another request" + }, + "422": { + "description": "The graph cannot be measured as a calibration run" + }, + "429": { + "description": "Too many unacknowledged idempotency keys; acknowledge or wait, then retry" + }, + "503": { + "description": "The calibration store cannot be read or written (unwritable root, lock timeout)" + } + } + } + }, + "/api/v2/models/working_memory/status": { + "get": { + "tags": ["model_manager"], + "summary": "Get Working Memory Calibration Status", + "description": "Every device class's calibration runs, with their derived status, point counts and environments.", + "operationId": "get_working_memory_calibration_status", + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/WorkingMemoryCalibrationStatus" + } + } + } + }, + "503": { + "description": "The calibration store cannot be read or written (unwritable root, lock timeout)" + } + }, + "security": [ + { + "HTTPBearer": [] + } + ] + } + }, + "/api/v2/models/working_memory/runs/{run_id}": { + "delete": { + "tags": ["model_manager"], + "summary": "Delete Working Memory Calibration Run", + "description": "Remove a calibration run and its measurements. Its queue items that have not run yet fail when they start.", + "operationId": "delete_working_memory_calibration_run", + "security": [ + { + "HTTPBearer": [] + } + ], + "parameters": [ + { + "name": "run_id", + "in": "path", + "required": true, + "schema": { + "type": "string", + "description": "The calibration run to remove", + "title": "Run Id" + }, + "description": "The calibration run to remove" + } + ], + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/WorkingMemoryCalibrationResetResult" + } + } + } + }, + "404": { + "description": "No usable store file holds the run" + }, + "409": { + "description": "The run's store file became unusable after the run was found; it is kept as it is" + }, + "503": { + "description": "The calibration store cannot be read or written (unwritable root, lock timeout)" + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + } + } + }, + "/api/v2/models/i/{key}/working_memory": { + "delete": { + "tags": ["model_manager"], + "summary": "Reset Model Working Memory Calibration", + "description": "Remove a model's calibration runs with all their measurements, and every measurement of the model taken in\nother runs (a VAE measured inside a transformer's run, say), from one device class.", + "operationId": "reset_model_working_memory_calibration", + "security": [ + { + "HTTPBearer": [] + } + ], + "parameters": [ + { + "name": "key", + "in": "path", + "required": true, + "schema": { + "type": "string", + "description": "Key of the model whose runs and measurements to remove", + "title": "Key" + }, + "description": "Key of the model whose runs and measurements to remove" + }, + { + "name": "device_class_id", + "in": "query", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string", + "pattern": "^[0-9a-f]{16}$" + }, + { + "type": "null" + } + ], + "description": "The device class to act on; the current generation device's by default", + "title": "Device Class Id" + }, + "description": "The device class to act on; the current generation device's by default" + } + ], + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/WorkingMemoryCalibrationResetResult" + } + } + } + }, + "409": { + "description": "No device_class_id was given and there is no current one (no GPU, or several generation devices), or the device class's store file is unusable and kept as it is" + }, + "503": { + "description": "The calibration store cannot be read or written (unwritable root, lock timeout)" + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + } + } + }, + "/api/v2/models/working_memory": { + "delete": { + "tags": ["model_manager"], + "summary": "Reset Working Memory Calibration", + "description": "Remove every calibration run and measurement of one device class.", + "operationId": "reset_working_memory_calibration", + "security": [ + { + "HTTPBearer": [] + } + ], + "parameters": [ + { + "name": "device_class_id", + "in": "query", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string", + "pattern": "^[0-9a-f]{16}$" + }, + { + "type": "null" + } + ], + "description": "The device class to act on; the current generation device's by default", + "title": "Device Class Id" + }, + "description": "The device class to act on; the current generation device's by default" + } + ], + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/WorkingMemoryCalibrationResetResult" + } + } + } + }, + "409": { + "description": "No device_class_id was given and there is no current one (no GPU, or several generation devices), or the device class's store file is unusable and kept as it is" + }, + "503": { + "description": "The calibration store cannot be read or written (unwritable root, lock timeout)" + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + } + } + }, + "/api/v2/models/working_memory/export": { + "get": { + "tags": ["model_manager"], + "summary": "Export Working Memory Calibration", + "description": "The device class's calibration store file as it is on disk, readable by this version or not.", + "operationId": "export_working_memory_calibration", + "security": [ + { + "HTTPBearer": [] + } + ], + "parameters": [ + { + "name": "device_class_id", + "in": "query", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string", + "pattern": "^[0-9a-f]{16}$" + }, + { + "type": "null" + } + ], + "description": "The device class to act on; the current generation device's by default", + "title": "Device Class Id" + }, + "description": "The device class to act on; the current generation device's by default" + } + ], + "responses": { + "200": { + "description": "The device class's store file byte for byte: a CalibrationStoreFile (schema_version 1) when this version can read it, otherwise whatever the file holds", + "content": { + "application/json": {} + } + }, + "404": { + "description": "The device class has no store file" + }, + "409": { + "description": "No device_class_id was given and there is no current one" + }, + "503": { + "description": "The calibration store cannot be read or written (unwritable root, lock timeout)" + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + } + } + }, "/api/v1/image_moves/start": { "post": { "tags": ["image_moves"], @@ -24174,6 +24519,48 @@ "title": "CalculateImageTilesOutput", "type": "object" }, + "CalibrationDeviceClass": { + "properties": { + "id": { + "type": "string", + "title": "Id", + "description": "A short stable hash of the other fields" + }, + "vendor": { + "type": "string", + "title": "Vendor" + }, + "name": { + "type": "string", + "title": "Name" + }, + "total_mib": { + "type": "integer", + "title": "Total Mib" + }, + "backend": { + "type": "string", + "enum": ["cuda", "hip", "mps", "xpu"], + "title": "Backend" + }, + "arch": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Arch" + } + }, + "additionalProperties": false, + "type": "object", + "required": ["id", "vendor", "name", "total_mib", "backend", "arch"], + "title": "CalibrationDeviceClass", + "description": "A `fingerprint.DeviceClass` with its id as a plain field, as it is stored and returned." + }, "CallSavedWorkflowInvocation": { "category": "workflow", "class": "invocation", @@ -115791,6 +116178,380 @@ ], "title": "WorkflowWithoutID" }, + "WorkingMemoryCalibrationDeviceClass": { + "properties": { + "device_class_id": { + "type": "string", + "title": "Device Class Id" + }, + "device_class": { + "anyOf": [ + { + "$ref": "#/components/schemas/CalibrationDeviceClass" + }, + { + "type": "null" + } + ], + "description": "Null when the file is unusable" + }, + "usable": { + "type": "boolean", + "title": "Usable" + }, + "problem": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Problem", + "description": "Why the file is unusable. It is kept as it is and never written" + }, + "is_current": { + "type": "boolean", + "title": "Is Current" + }, + "runs": { + "items": { + "$ref": "#/components/schemas/WorkingMemoryCalibrationRunSummary" + }, + "type": "array", + "title": "Runs" + }, + "point_counts": { + "items": { + "$ref": "#/components/schemas/WorkingMemoryCalibrationPointCount" + }, + "type": "array", + "title": "Point Counts" + }, + "envs": { + "items": { + "$ref": "#/components/schemas/WorkingMemoryCalibrationEnvironment" + }, + "type": "array", + "title": "Envs" + } + }, + "type": "object", + "required": [ + "device_class_id", + "device_class", + "usable", + "problem", + "is_current", + "runs", + "point_counts", + "envs" + ], + "title": "WorkingMemoryCalibrationDeviceClass" + }, + "WorkingMemoryCalibrationEnvironment": { + "properties": { + "id": { + "type": "string", + "title": "Id" + }, + "hard": { + "additionalProperties": { + "type": "string" + }, + "type": "object", + "title": "Hard" + }, + "soft": { + "additionalProperties": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ] + }, + "type": "object", + "title": "Soft" + }, + "hard_matches_current": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Hard Matches Current", + "description": "Whether the current environment has the same hard facts; null off the current device class" + }, + "soft_matches_current": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "title": "Soft Matches Current", + "description": "Whether the current environment has the same soft facts; null off the current device class" + }, + "point_count": { + "type": "integer", + "title": "Point Count" + } + }, + "type": "object", + "required": ["id", "hard", "soft", "hard_matches_current", "soft_matches_current", "point_count"], + "title": "WorkingMemoryCalibrationEnvironment" + }, + "WorkingMemoryCalibrationPointCount": { + "properties": { + "estimator": { + "type": "string", + "title": "Estimator" + }, + "version": { + "type": "integer", + "title": "Version" + }, + "valid": { + "type": "boolean", + "title": "Valid" + }, + "fit_eligible": { + "type": "boolean", + "title": "Fit Eligible" + }, + "count": { + "type": "integer", + "title": "Count" + } + }, + "type": "object", + "required": ["estimator", "version", "valid", "fit_eligible", "count"], + "title": "WorkingMemoryCalibrationPointCount" + }, + "WorkingMemoryCalibrationRequest": { + "properties": { + "graph": { + "$ref": "#/components/schemas/Graph", + "description": "The graph to measure" + }, + "data": { + "anyOf": [ + { + "items": { + "items": { + "$ref": "#/components/schemas/BatchDatum" + }, + "type": "array" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Data", + "description": "The batch data collection" + }, + "idempotency_key": { + "anyOf": [ + { + "type": "string", + "maxLength": 255, + "minLength": 1 + }, + { + "type": "null" + } + ], + "title": "Idempotency Key", + "description": "A caller-owned retry key. Repeating a request with it returns the run it created and enqueues nothing new; reusing it with a different request is rejected." + } + }, + "additionalProperties": false, + "type": "object", + "required": ["graph"], + "title": "WorkingMemoryCalibrationRequest", + "description": "The batch a calibration run enqueues, built by the client. The server sets origin, destination and runs, and\nturns the node cache off." + }, + "WorkingMemoryCalibrationResetResult": { + "properties": { + "runs_removed": { + "type": "integer", + "title": "Runs Removed" + }, + "points_removed": { + "type": "integer", + "title": "Points Removed" + } + }, + "type": "object", + "required": ["runs_removed", "points_removed"], + "title": "WorkingMemoryCalibrationResetResult" + }, + "WorkingMemoryCalibrationRunCreated": { + "properties": { + "run_id": { + "type": "string", + "title": "Run Id" + }, + "device_class_id": { + "type": "string", + "title": "Device Class Id" + }, + "item_ids": { + "items": { + "type": "integer" + }, + "type": "array", + "title": "Item Ids", + "description": "The queue items the run enqueued" + } + }, + "type": "object", + "required": ["run_id", "device_class_id", "item_ids"], + "title": "WorkingMemoryCalibrationRunCreated" + }, + "WorkingMemoryCalibrationRunSummary": { + "properties": { + "run_id": { + "type": "string", + "title": "Run Id" + }, + "model_key": { + "type": "string", + "title": "Model Key" + }, + "model_hash": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Model Hash" + }, + "user_id": { + "type": "string", + "title": "User Id" + }, + "created_at": { + "type": "string", + "format": "date-time", + "title": "Created At" + }, + "env_id": { + "type": "string", + "title": "Env Id" + }, + "status": { + "type": "string", + "enum": ["running", "completed", "failed", "partial", "canceled"], + "title": "Status", + "description": "Derived from the run's queue items, a retry standing in for the item it retried: running while any is pending or running; failed when any failed; completed when all completed; otherwise partial when points were recorded, else canceled" + }, + "error": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Error", + "description": "The error of a failed item, when the run failed" + }, + "item_ids": { + "items": { + "type": "integer" + }, + "type": "array", + "title": "Item Ids", + "description": "Every queue item of the run, retries included" + }, + "point_count": { + "type": "integer", + "title": "Point Count" + } + }, + "type": "object", + "required": [ + "run_id", + "model_key", + "model_hash", + "user_id", + "created_at", + "env_id", + "status", + "error", + "item_ids", + "point_count" + ], + "title": "WorkingMemoryCalibrationRunSummary" + }, + "WorkingMemoryCalibrationStatus": { + "properties": { + "available": { + "type": "boolean", + "title": "Available", + "description": "Whether a calibration run can start now: one GPU generation device whose store file is usable" + }, + "unavailable_reason": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Unavailable Reason" + }, + "current_device_class_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Current Device Class Id" + }, + "current_env_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Current Env Id" + }, + "device_classes": { + "items": { + "$ref": "#/components/schemas/WorkingMemoryCalibrationDeviceClass" + }, + "type": "array", + "title": "Device Classes" + } + }, + "type": "object", + "required": ["available", "unavailable_reason", "current_device_class_id", "current_env_id", "device_classes"], + "title": "WorkingMemoryCalibrationStatus" + }, "ZImageConditioningField": { "description": "A Z-Image conditioning tensor primitive value", "properties": { diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts index 1c635d6f3c3..0372f0d0f1e 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -1195,6 +1195,131 @@ export type paths = { patch?: never; trace?: never; }; + "/api/v2/models/i/{key}/working_memory/calibrate": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Calibrate Model Working Memory + * @description Start a working-memory calibration run: enqueue the batch with every node uncached and measure each node's + * model operations on the generation device. Needs exactly one generation device, and it must be a GPU. + * + * With an `idempotency_key`, repeating the request (after a lost response) returns the run it created and enqueues + * nothing new; acknowledge the key through the queue's enqueue_batch/acknowledge route once the response arrived. + */ + post: operations["calibrate_model_working_memory"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/api/v2/models/working_memory/status": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Working Memory Calibration Status + * @description Every device class's calibration runs, with their derived status, point counts and environments. + */ + get: operations["get_working_memory_calibration_status"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/api/v2/models/working_memory/runs/{run_id}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + post?: never; + /** + * Delete Working Memory Calibration Run + * @description Remove a calibration run and its measurements. Its queue items that have not run yet fail when they start. + */ + delete: operations["delete_working_memory_calibration_run"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/api/v2/models/i/{key}/working_memory": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + post?: never; + /** + * Reset Model Working Memory Calibration + * @description Remove a model's calibration runs with all their measurements, and every measurement of the model taken in + * other runs (a VAE measured inside a transformer's run, say), from one device class. + */ + delete: operations["reset_model_working_memory_calibration"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/api/v2/models/working_memory": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + post?: never; + /** + * Reset Working Memory Calibration + * @description Remove every calibration run and measurement of one device class. + */ + delete: operations["reset_working_memory_calibration"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/api/v2/models/working_memory/export": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Export Working Memory Calibration + * @description The device class's calibration store file as it is on disk, readable by this version or not. + */ + get: operations["export_working_memory_calibration"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/api/v1/image_moves/start": { parameters: { query?: never; @@ -7891,6 +8016,30 @@ export type components = { */ type: "calculate_image_tiles_output"; }; + /** + * CalibrationDeviceClass + * @description A `fingerprint.DeviceClass` with its id as a plain field, as it is stored and returned. + */ + CalibrationDeviceClass: { + /** + * Id + * @description A short stable hash of the other fields + */ + id: string; + /** Vendor */ + vendor: string; + /** Name */ + name: string; + /** Total Mib */ + total_mib: number; + /** + * Backend + * @enum {string} + */ + backend: "cuda" | "hip" | "mps" | "xpu"; + /** Arch */ + arch: string | null; + }; /** * Call Saved Workflow * @description Displays and later executes against a selected saved workflow. @@ -52359,6 +52508,156 @@ export type components = { [key: string]: components["schemas"]["JsonValue"]; } | null; }; + /** WorkingMemoryCalibrationDeviceClass */ + WorkingMemoryCalibrationDeviceClass: { + /** Device Class Id */ + device_class_id: string; + /** @description Null when the file is unusable */ + device_class: components["schemas"]["CalibrationDeviceClass"] | null; + /** Usable */ + usable: boolean; + /** + * Problem + * @description Why the file is unusable. It is kept as it is and never written + */ + problem: string | null; + /** Is Current */ + is_current: boolean; + /** Runs */ + runs: components["schemas"]["WorkingMemoryCalibrationRunSummary"][]; + /** Point Counts */ + point_counts: components["schemas"]["WorkingMemoryCalibrationPointCount"][]; + /** Envs */ + envs: components["schemas"]["WorkingMemoryCalibrationEnvironment"][]; + }; + /** WorkingMemoryCalibrationEnvironment */ + WorkingMemoryCalibrationEnvironment: { + /** Id */ + id: string; + /** Hard */ + hard: { + [key: string]: string; + }; + /** Soft */ + soft: { + [key: string]: string | null; + }; + /** + * Hard Matches Current + * @description Whether the current environment has the same hard facts; null off the current device class + */ + hard_matches_current: boolean | null; + /** + * Soft Matches Current + * @description Whether the current environment has the same soft facts; null off the current device class + */ + soft_matches_current: boolean | null; + /** Point Count */ + point_count: number; + }; + /** WorkingMemoryCalibrationPointCount */ + WorkingMemoryCalibrationPointCount: { + /** Estimator */ + estimator: string; + /** Version */ + version: number; + /** Valid */ + valid: boolean; + /** Fit Eligible */ + fit_eligible: boolean; + /** Count */ + count: number; + }; + /** + * WorkingMemoryCalibrationRequest + * @description The batch a calibration run enqueues, built by the client. The server sets origin, destination and runs, and + * turns the node cache off. + */ + WorkingMemoryCalibrationRequest: { + /** @description The graph to measure */ + graph: components["schemas"]["Graph"]; + /** + * Data + * @description The batch data collection + */ + data?: components["schemas"]["BatchDatum"][][] | null; + /** + * Idempotency Key + * @description A caller-owned retry key. Repeating a request with it returns the run it created and enqueues nothing new; reusing it with a different request is rejected. + */ + idempotency_key?: string | null; + }; + /** WorkingMemoryCalibrationResetResult */ + WorkingMemoryCalibrationResetResult: { + /** Runs Removed */ + runs_removed: number; + /** Points Removed */ + points_removed: number; + }; + /** WorkingMemoryCalibrationRunCreated */ + WorkingMemoryCalibrationRunCreated: { + /** Run Id */ + run_id: string; + /** Device Class Id */ + device_class_id: string; + /** + * Item Ids + * @description The queue items the run enqueued + */ + item_ids: number[]; + }; + /** WorkingMemoryCalibrationRunSummary */ + WorkingMemoryCalibrationRunSummary: { + /** Run Id */ + run_id: string; + /** Model Key */ + model_key: string; + /** Model Hash */ + model_hash: string | null; + /** User Id */ + user_id: string; + /** + * Created At + * Format: date-time + */ + created_at: string; + /** Env Id */ + env_id: string; + /** + * Status + * @description Derived from the run's queue items, a retry standing in for the item it retried: running while any is pending or running; failed when any failed; completed when all completed; otherwise partial when points were recorded, else canceled + * @enum {string} + */ + status: "running" | "completed" | "failed" | "partial" | "canceled"; + /** + * Error + * @description The error of a failed item, when the run failed + */ + error: string | null; + /** + * Item Ids + * @description Every queue item of the run, retries included + */ + item_ids: number[]; + /** Point Count */ + point_count: number; + }; + /** WorkingMemoryCalibrationStatus */ + WorkingMemoryCalibrationStatus: { + /** + * Available + * @description Whether a calibration run can start now: one GPU generation device whose store file is usable + */ + available: boolean; + /** Unavailable Reason */ + unavailable_reason: string | null; + /** Current Device Class Id */ + current_device_class_id: string | null; + /** Current Env Id */ + current_env_id: string | null; + /** Device Classes */ + device_classes: components["schemas"]["WorkingMemoryCalibrationDeviceClass"][]; + }; /** * ZImageConditioningField * @description A Z-Image conditioning tensor primitive value @@ -55519,6 +55818,296 @@ export interface operations { }; }; }; + calibrate_model_working_memory: { + parameters: { + query?: never; + header?: never; + path: { + /** @description Key of the model the run calibrates */ + key: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["WorkingMemoryCalibrationRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 201: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WorkingMemoryCalibrationRunCreated"]; + }; + }; + /** @description The model could not be found */ + 404: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description This machine cannot calibrate now (no GPU, or several generation devices), the device class's store file is unusable and kept as it is, the queue is full, or the idempotency key was used for another request */ + 409: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description The graph cannot be measured as a calibration run */ + 422: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description Too many unacknowledged idempotency keys; acknowledge or wait, then retry */ + 429: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description The calibration store cannot be read or written (unwritable root, lock timeout) */ + 503: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; + get_working_memory_calibration_status: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WorkingMemoryCalibrationStatus"]; + }; + }; + /** @description The calibration store cannot be read or written (unwritable root, lock timeout) */ + 503: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; + delete_working_memory_calibration_run: { + parameters: { + query?: never; + header?: never; + path: { + /** @description The calibration run to remove */ + run_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WorkingMemoryCalibrationResetResult"]; + }; + }; + /** @description No usable store file holds the run */ + 404: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description The run's store file became unusable after the run was found; it is kept as it is */ + 409: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + /** @description The calibration store cannot be read or written (unwritable root, lock timeout) */ + 503: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; + reset_model_working_memory_calibration: { + parameters: { + query?: { + /** @description The device class to act on; the current generation device's by default */ + device_class_id?: string | null; + }; + header?: never; + path: { + /** @description Key of the model whose runs and measurements to remove */ + key: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WorkingMemoryCalibrationResetResult"]; + }; + }; + /** @description No device_class_id was given and there is no current one (no GPU, or several generation devices), or the device class's store file is unusable and kept as it is */ + 409: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + /** @description The calibration store cannot be read or written (unwritable root, lock timeout) */ + 503: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; + reset_working_memory_calibration: { + parameters: { + query?: { + /** @description The device class to act on; the current generation device's by default */ + device_class_id?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["WorkingMemoryCalibrationResetResult"]; + }; + }; + /** @description No device_class_id was given and there is no current one (no GPU, or several generation devices), or the device class's store file is unusable and kept as it is */ + 409: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + /** @description The calibration store cannot be read or written (unwritable root, lock timeout) */ + 503: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; + export_working_memory_calibration: { + parameters: { + query?: { + /** @description The device class to act on; the current generation device's by default */ + device_class_id?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description The device class's store file byte for byte: a CalibrationStoreFile (schema_version 1) when this version can read it, otherwise whatever the file holds */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description The device class has no store file */ + 404: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description No device_class_id was given and there is no current one */ + 409: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + /** @description The calibration store cannot be read or written (unwritable root, lock timeout) */ + 503: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; + }; + }; start_image_move: { parameters: { query?: never; diff --git a/tests/app/routers/conftest.py b/tests/app/routers/conftest.py index 1221872b69e..994395507ce 100644 --- a/tests/app/routers/conftest.py +++ b/tests/app/routers/conftest.py @@ -63,6 +63,7 @@ def __init__(self, invoker: Invoker) -> None: "invokeai.app.api.routers.workflows", "invokeai.app.api.routers._access", "invokeai.app.api.routers.model_manager", + "invokeai.app.api.routers.working_memory", "invokeai.app.api.routers.app_info", "invokeai.app.api.routers.projects", "invokeai.app.api.routers.boards", diff --git a/tests/app/routers/test_working_memory_calibration.py b/tests/app/routers/test_working_memory_calibration.py new file mode 100644 index 00000000000..fb765c39569 --- /dev/null +++ b/tests/app/routers/test_working_memory_calibration.py @@ -0,0 +1,244 @@ +"""The calibration routes and the reserved origin, against a real SQLite queue and a store on a temporary root.""" + +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock + +import pytest +import torch +from fastapi.testclient import TestClient + +from invokeai.app.invocations.call_saved_workflow import CallSavedWorkflowInvocation +from invokeai.app.invocations.math import AddInvocation +from invokeai.app.services.invoker import Invoker +from invokeai.app.services.model_records import UnknownModelException +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.graph import Graph +from invokeai.app.services.working_memory_calibration.working_memory_calibration_common import ( + CALIBRATION_DESTINATION, + calibration_origin, +) +from tests.app.routers.conftest import MockApiDependencies, _auth +from tests.fixtures.working_memory_calibration import GPU, GPU_CLASS, calibration_service, measured_window + +_DEVICES: list[torch.device] = [GPU] + + +@pytest.fixture +def calibration( + tmp_path: Path, monkeypatch: Any, mock_invoker: Invoker, mock_sqlite_database: Database, enable_multiuser: Any +): + """A started queue, a model store that knows `model-key`, and the calibration service on `tmp_path`.""" + services = mock_invoker.services + queue = SessionQueue(mock_sqlite_database) + services.session_queue = queue + queue.start(mock_invoker) + + def get_model(key: str) -> MagicMock: + if key != "model-key": + raise UnknownModelException(f"unknown model {key}") + return MagicMock(key=key, hash="blake3:abc") + + services.model_manager.store.get_model.side_effect = get_model + _DEVICES[:] = [GPU] + service = calibration_service(tmp_path, devices=lambda: list(_DEVICES)) + service.start(mock_invoker) + services.working_memory_calibration = service + + mock_deps = MockApiDependencies(mock_invoker) + monkeypatch.setattr("invokeai.app.api.routers.session_queue.ApiDependencies", mock_deps) + monkeypatch.setattr("invokeai.app.api.routers.image_move_maintenance.ApiDependencies", mock_deps) + return queue, service + + +def _graph() -> dict: + graph = Graph() + graph.add_node(AddInvocation(id="add", a=1, b=2, use_cache=True)) + return graph.model_dump(mode="json") + + +def _calibrate(client: TestClient, token: str, body: dict | None = None, key: str = "model-key"): + return client.post( + f"/api/v2/models/i/{key}/working_memory/calibrate", + json=body if body is not None else {"graph": _graph()}, + headers=_auth(token), + ) + + +def test_calibrate_enqueues_the_run_uncached_with_the_reserved_origin( + client: TestClient, admin_token: str, calibration +) -> None: + queue, service = calibration + + response = _calibrate(client, admin_token) + + assert response.status_code == 201, response.text + created = response.json() + assert created["device_class_id"] == GPU_CLASS.id + [item_id] = created["item_ids"] + item = queue.get_queue_item(item_id) + assert item.origin == calibration_origin(created["run_id"]) and item.destination == CALIBRATION_DESTINATION + assert [node.use_cache for node in item.session.graph.nodes.values()] == [False] + [run] = service.get_status().device_classes[0].runs + assert (run.run_id, run.model_key, run.item_ids, run.status) == ( + created["run_id"], + "model-key", + [item_id], + "running", + ) + + +def test_non_admins_cannot_calibrate_reset_read_or_export( + client: TestClient, admin_token: str, user1_token: str, calibration +) -> None: + queue, _ = calibration + run_id = _calibrate(client, admin_token).json()["run_id"] + item_ids = queue.get_queue_item_ids("default").item_ids + + responses = [ + _calibrate(client, user1_token), + client.get("/api/v2/models/working_memory/status", headers=_auth(user1_token)), + client.get("/api/v2/models/working_memory/export", headers=_auth(user1_token)), + client.delete(f"/api/v2/models/working_memory/runs/{run_id}", headers=_auth(user1_token)), + client.delete("/api/v2/models/i/model-key/working_memory", headers=_auth(user1_token)), + client.delete("/api/v2/models/working_memory", headers=_auth(user1_token)), + ] + + assert [response.status_code for response in responses] == [403] * 6 + assert queue.get_queue_item_ids("default").item_ids == item_ids + + +@pytest.mark.parametrize("origin", [calibration_origin("forged"), "INVOKEAI:WM-Calibration:forged"]) +def test_enqueue_batch_refuses_the_reserved_origin_even_to_admins( + client: TestClient, admin_token: str, calibration, origin: str +) -> None: + queue, _ = calibration + + response = client.post( + "/api/v1/queue/default/enqueue_batch", + json={"batch": {"graph": _graph(), "origin": origin, "runs": 1}}, + headers=_auth(admin_token), + ) + + assert response.status_code == 400 and "reserved" in response.json()["detail"] + assert queue.get_queue_item_ids("default").item_ids == [] + + +def test_calibrate_refuses_more_than_one_generation_device_and_cpu( + client: TestClient, admin_token: str, calibration +) -> None: + queue, _ = calibration + _DEVICES[:] = [GPU, torch.device("cuda:1")] + assert _calibrate(client, admin_token).status_code == 409 + _DEVICES[:] = [torch.device("cpu")] + assert _calibrate(client, admin_token).status_code == 409 + assert queue.get_queue_item_ids("default").item_ids == [] + + +def test_calibrate_rejects_unknown_models_workflow_calls_and_server_owned_fields( + client: TestClient, admin_token: str, calibration +) -> None: + queue, service = calibration + call_graph = Graph() + call_graph.add_node(CallSavedWorkflowInvocation(id="call", workflow_id="workflow")) + use_cache_data = [[{"node_path": "add", "field_name": "use_cache", "items": [True]}]] + + assert _calibrate(client, admin_token, key="missing").status_code == 404 + assert _calibrate(client, admin_token, {"graph": call_graph.model_dump(mode="json")}).status_code == 422 + assert _calibrate(client, admin_token, {"graph": _graph(), "data": use_cache_data}).status_code == 422 + assert _calibrate(client, admin_token, {"graph": _graph(), "origin": "canvas"}).status_code == 422 + assert queue.get_queue_item_ids("default").item_ids == [] + assert service.get_status().device_classes[0].runs == [] + + +def test_retrying_a_calibration_item_keeps_its_origin_and_owner( + client: TestClient, admin_token: str, calibration +) -> None: + queue, _ = calibration + created = _calibrate(client, admin_token).json() + [item_id] = created["item_ids"] + item = queue.dequeue() + assert item is not None + queue.fail_queue_item(item_id, "ValueError", "failed", "traceback") + + response = client.put("/api/v1/queue/default/retry_items_by_id", json=[item_id], headers=_auth(admin_token)) + + assert response.status_code == 200, response.text + retried = queue.dequeue() + assert retried is not None and retried.retried_from_item_id == item_id + assert (retried.origin, retried.user_id) == (calibration_origin(created["run_id"]), item.user_id) + + +def test_status_reset_and_export_for_admins(client: TestClient, admin_token: str, calibration) -> None: + _, service = calibration + run_id = _calibrate(client, admin_token).json()["run_id"] + service.record_windows(run_id, [measured_window(run_id), measured_window(run_id, model_key="model-key")]) + + status = client.get("/api/v2/models/working_memory/status", headers=_auth(admin_token)).json() + assert status["available"] and status["current_device_class_id"] == GPU_CLASS.id + assert status["device_classes"][0]["runs"][0]["point_count"] == 2 + + export = client.get("/api/v2/models/working_memory/export", headers=_auth(admin_token)) + assert export.status_code == 200 and export.headers["content-type"] == "application/json" + assert [point["run_id"] for point in export.json()["points"]] == [run_id, run_id] + + reset = client.delete("/api/v2/models/i/model-key/working_memory", headers=_auth(admin_token)) + assert reset.json() == {"runs_removed": 1, "points_removed": 2} + assert client.delete(f"/api/v2/models/working_memory/runs/{run_id}", headers=_auth(admin_token)).status_code == 404 + assert client.delete("/api/v2/models/working_memory", headers=_auth(admin_token)).json() == { + "runs_removed": 0, + "points_removed": 0, + } + missing_class = client.get( + "/api/v2/models/working_memory/export", params={"device_class_id": "0" * 16}, headers=_auth(admin_token) + ) + assert missing_class.status_code == 404 + + +def test_a_repeated_calibrate_with_its_retry_key_returns_the_same_run( + client: TestClient, admin_token: str, calibration +) -> None: + queue, service = calibration + body = {"graph": _graph(), "idempotency_key": "lost-response"} + + first = _calibrate(client, admin_token, body) + replay = _calibrate(client, admin_token, body) + + assert first.status_code == replay.status_code == 201 + assert replay.json() == first.json() + assert queue.get_queue_item_ids("default").item_ids == first.json()["item_ids"] + assert [run.run_id for run in service.get_status().device_classes[0].runs] == [first.json()["run_id"]] + + other_graph = Graph() + other_graph.add_node(AddInvocation(id="add", a=5, b=2)) + conflict = _calibrate( + client, admin_token, {"graph": other_graph.model_dump(mode="json"), "idempotency_key": "lost-response"} + ) + + assert conflict.status_code == 409 + # The conflicting request does not take the original run with it. + assert [run.run_id for run in service.get_status().device_classes[0].runs] == [first.json()["run_id"]] + assert queue.get_queue_item_ids("default").item_ids == first.json()["item_ids"] + + +def test_calibrate_into_a_full_queue_is_refused_and_leaves_no_run( + client: TestClient, admin_token: str, mock_invoker: Invoker, calibration +) -> None: + _, service = calibration + mock_invoker.services.configuration.max_queue_size = 1 + kept = _calibrate(client, admin_token).json()["run_id"] + + response = _calibrate(client, admin_token) + + assert response.status_code == 409 and "queue is full" in response.json()["detail"] + assert [run.run_id for run in service.get_status().device_classes[0].runs] == [kept] + + +def test_the_status_carries_the_device_class_id(client: TestClient, admin_token: str, calibration) -> None: + status = client.get("/api/v2/models/working_memory/status", headers=_auth(admin_token)).json() + + assert status["device_classes"][0]["device_class"] == { + "id": GPU_CLASS.id, + **GPU_CLASS.model_dump(exclude={"id"}), + } diff --git a/tests/app/services/session_processor/test_working_memory_calibration_runner.py b/tests/app/services/session_processor/test_working_memory_calibration_runner.py new file mode 100644 index 00000000000..7d15256c7c0 --- /dev/null +++ b/tests/app/services/session_processor/test_working_memory_calibration_runner.py @@ -0,0 +1,379 @@ +"""The session runner's calibration hook: calibration items run their nodes inside a probe scope, have the windows +persisted, and are admitted again at session start; every other item runs exactly as before. + +The probe itself needs a GPU, so a test node plays its part: it appends a window to the scope it runs in, as an +estimated `model_on_device` would. +""" + +import asyncio +import json +import uuid +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from pathlib import Path +from threading import Event as ThreadEvent +from types import SimpleNamespace +from typing import Literal, Optional +from unittest.mock import MagicMock + +import pytest +import torch + +from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation +from invokeai.app.invocations.call_saved_workflow import CallSavedWorkflowInvocation +from invokeai.app.invocations.fields import InputField +from invokeai.app.invocations.primitives import IntegerOutput +from invokeai.app.services.invoker import Invoker +from invokeai.app.services.session_processor.session_processor_common import CanceledException +from invokeai.app.services.session_processor.session_processor_default import DefaultSessionRunner +from invokeai.app.services.session_queue.session_queue_common import Batch, SessionQueueItem +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.graph import Graph, GraphExecutionState +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.app.services.users.users_common import UserCreateRequest, UserUpdateRequest +from invokeai.app.services.working_memory_calibration.working_memory_calibration_common import ( + WorkingMemoryCalibrationRequest, + calibration_origin, +) +from invokeai.app.services.working_memory_calibration.working_memory_calibration_default import ( + WorkingMemoryCalibrationService, +) +from invokeai.backend.working_memory import probe +from tests.fixtures.sqlite_database import sqlite_cursor_of +from tests.fixtures.working_memory_calibration import GPU, GPU_CLASS, OTHER_CLASS, calibration_service, measured_window + +# What each test node saw: the scope it ran in, or None. +_SEEN: list[Optional[tuple[str, int, str, str]]] = [] +_ON_CANCEL: list[Callable[[], None]] = [] + + +@invocation("test_working_memory_calibration_node", version="1.0.0") +class CalibrationProbeNodeInvocation(BaseInvocation): + behaviour: Literal["complete", "fail", "cancel"] = InputField(default="complete") + + def invoke(self, context: InvocationContext) -> IntegerOutput: + scope = probe.current_scope() + if scope is None: + _SEEN.append(None) + else: + _SEEN.append((scope.run_id, scope.queue_item_id, scope.node_id, scope.node_type)) + scope.windows.append(measured_window(scope.run_id, scope.queue_item_id, scope.node_id)) + if self.behaviour == "fail": + raise ValueError("decode failed") + if self.behaviour == "cancel": + _ON_CANCEL[0]() + raise CanceledException() + return IntegerOutput(value=1) + + +class _Stats: + @contextmanager + def collect_stats(self, invocation, graph_execution_state_id) -> Iterator[None]: + yield + + def log_stats(self, graph_execution_state_id) -> None: + pass + + def reset_stats(self, graph_execution_state_id) -> None: + pass + + +@pytest.fixture(autouse=True) +def _reset_seen() -> Iterator[None]: + _SEEN.clear() + _ON_CANCEL.clear() + yield + _SEEN.clear() + _ON_CANCEL.clear() + + +@pytest.fixture +def queue(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: + queue = SessionQueue(mock_sqlite_database) + services = mock_invoker.services + services.session_queue = queue + services.performance_statistics = _Stats() + model_manager = MagicMock() + model_manager.store.get_model.side_effect = lambda key: SimpleNamespace(key=key, hash=None) + services.model_manager = model_manager + services.configuration.offload_on_architecture_switch = False + queue.start(mock_invoker) + return queue + + +@pytest.fixture +def service(tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue) -> WorkingMemoryCalibrationService: + service = calibration_service(tmp_path) + service.start(mock_invoker) + mock_invoker.services.working_memory_calibration = service + return service + + +@pytest.fixture +def runner(mock_invoker: Invoker, queue: SessionQueue) -> tuple[DefaultSessionRunner, ThreadEvent]: + runner = DefaultSessionRunner() + cancel_event = ThreadEvent() + runner.start(services=mock_invoker.services, cancel_event=cancel_event) + return runner, cancel_event + + +def _graph(*behaviours: str) -> Graph: + graph = Graph() + for index, behaviour in enumerate(behaviours): + graph.add_node(CalibrationProbeNodeInvocation(id=f"node{index}", behaviour=behaviour)) + return graph + + +def _create(service: WorkingMemoryCalibrationService, user_id: str, *behaviours: str) -> tuple[str, int]: + created = asyncio.run( + service.create_run("model-key", WorkingMemoryCalibrationRequest(graph=_graph(*behaviours)), user_id) + ) + return created.run_id, created.item_ids[0] + + +def _run_next(queue: SessionQueue, runner: DefaultSessionRunner) -> SessionQueueItem: + item = queue.dequeue() + assert item is not None + runner.workflow_call_queue_lifecycle.run_queue_item(item) + return queue.get_queue_item(item.item_id) + + +def _stored_points(tmp_path: Path) -> list[dict]: + path = tmp_path / "working_memory" / f"{GPU_CLASS.id}.json" + return json.loads(path.read_bytes())["points"] + + +def _user(invoker: Invoker, email: str, is_admin: bool) -> str: + request = UserCreateRequest(email=email, display_name=email, password="TestPass123", is_admin=is_admin) + return invoker.services.users.create(request).user_id + + +def test_each_node_runs_in_its_own_scope_and_its_windows_are_stored( + tmp_path: Path, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner +) -> None: + run_id, item_id = _create(service, "system", "complete", "complete") + + item = _run_next(queue, runner[0]) + + assert item.status == "completed" + node_type = CalibrationProbeNodeInvocation.get_type() + assert sorted(_SEEN) == [(run_id, item_id, "node0", node_type), (run_id, item_id, "node1", node_type)] + assert sorted((point["node_id"], point["run_id"]) for point in _stored_points(tmp_path)) == [ + ("node0", run_id), + ("node1", run_id), + ] + assert not probe.armed + + +@pytest.mark.parametrize("behaviour", ["fail", "cancel"]) +def test_windows_are_stored_when_the_node_fails_or_is_canceled( + tmp_path: Path, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner, behaviour: str +) -> None: + session_runner, cancel_event = runner + run_id, item_id = _create(service, "system", behaviour) + + def cancel() -> None: + queue.cancel_queue_item(item_id) + cancel_event.set() + + _ON_CANCEL.append(cancel) + item = _run_next(queue, session_runner) + + assert item.status == ("failed" if behaviour == "fail" else "canceled") + assert [point["node_id"] for point in _stored_points(tmp_path)] == ["node0"] + [run] = service.get_status().device_classes[0].runs + assert run.status == ("failed" if behaviour == "fail" else "partial") + + +def test_an_item_whose_owner_is_not_an_admin_fails_before_any_node_runs( + tmp_path: Path, mock_invoker: Invoker, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner +) -> None: + mock_invoker.services.configuration.multiuser = True + user_id = _user(mock_invoker, "user@test.com", is_admin=False) + run_id, _ = _create(service, user_id, "complete") + + item = _run_next(queue, runner[0]) + + assert item.status == "failed" and "administrators only" in (item.error_message or "") + assert _SEEN == [] and _stored_points(tmp_path) == [] + [run] = service.get_status().device_classes[0].runs + assert run.status == "failed" and "administrators only" in (run.error or "") + + +def test_an_admin_owned_item_runs_in_multiuser_mode( + mock_invoker: Invoker, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner +) -> None: + mock_invoker.services.configuration.multiuser = True + admin_id = _user(mock_invoker, "admin@test.com", is_admin=True) + _create(service, admin_id, "complete") + + assert _run_next(queue, runner[0]).status == "completed" + assert len(_SEEN) == 1 + + +def test_a_retry_of_a_reset_run_is_refused_at_session_start( + service: WorkingMemoryCalibrationService, queue: SessionQueue, runner +) -> None: + run_id, item_id = _create(service, "system", "fail") + assert _run_next(queue, runner[0]).status == "failed" + + queue.retry_items_by_id("default", [item_id]) + service.reset_run(run_id) + _SEEN.clear() + retried = _run_next(queue, runner[0]) + + assert retried.retried_from_item_id == item_id and retried.origin == calibration_origin(run_id) + assert retried.status == "failed" and "does not exist" in (retried.error_message or "") + assert _SEEN == [] + + +def test_a_workflow_call_child_of_a_calibration_item_is_admitted_again( + mock_invoker: Invoker, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner +) -> None: + mock_invoker.services.configuration.multiuser = True + user_id = _user(mock_invoker, "user@test.com", is_admin=False) + run_id, _ = _create(service, user_id, "complete") + queue.delete_queue_items_by_id(queue.get_queue_item_ids("default").item_ids) + + # A parent waiting on a saved-workflow call, as the queue would hold one: the child inherits origin and owner. + graph = Graph() + graph.add_node(CallSavedWorkflowInvocation(id="call", workflow_id="workflow")) + root = GraphExecutionState(graph=graph) + call = root.next() + assert isinstance(call, CallSavedWorkflowInvocation) + frame = root.build_workflow_call_frame(call.id, "workflow") + root.begin_waiting_on_workflow_call(frame) + child_session = root.create_child_workflow_execution_state(_graph("complete"), frame) + root.attach_waiting_workflow_call_child_sessions([child_session]) + with sqlite_cursor_of(queue) as cursor: + cursor.execute( + "INSERT INTO session_queue (queue_id, session, session_id, batch_id, priority, user_id, status, origin)" + " VALUES ('default', ?, ?, ?, 0, ?, 'in_progress', ?);", + (root.model_dump_json(), root.id, uuid.uuid4().hex, user_id, calibration_origin(run_id)), + ) + root_id = int(cursor.lastrowid or 0) + [child] = queue.enqueue_workflow_call_children(queue.get_queue_item(root_id), [(child_session, None)]) + assert child.origin == calibration_origin(run_id) and child.user_id == user_id + + item = _run_next(queue, runner[0]) + + assert item.item_id == child.item_id + assert item.status == "failed" and "administrators only" in (item.error_message or "") + assert _SEEN == [] + + +def test_an_ordinary_item_runs_without_a_probe_scope( + tmp_path: Path, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner +) -> None: + asyncio.run(queue.enqueue_batch("default", Batch(graph=_graph("complete"), origin="canvas"), prepend=False)) + + item = _run_next(queue, runner[0]) + + assert item.status == "completed" + assert _SEEN == [None] + assert not (tmp_path / "working_memory").exists() + + +def test_a_store_failure_fails_the_item_and_stops_the_session( + tmp_path: Path, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner, monkeypatch +) -> None: + run_id, _ = _create(service, "system", "complete", "complete") + + def refuse(run_id: str, windows) -> None: + raise OSError("read-only file system") + + monkeypatch.setattr(service, "record_windows", refuse) + item = _run_next(queue, runner[0]) + + assert item.status == "failed" and "read-only file system" in (item.error_message or "") + assert len(_SEEN) == 1 + + +def _refused_before_any_node(queue: SessionQueue, runner: DefaultSessionRunner, refusal: str) -> None: + item = _run_next(queue, runner) + assert item.status == "failed" and refusal in (item.error_message or "") + assert _SEEN == [] + + +def test_an_owner_deactivated_after_enqueue_is_refused( + mock_invoker: Invoker, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner +) -> None: + mock_invoker.services.configuration.multiuser = True + _user(mock_invoker, "first-admin@test.com", is_admin=True) + admin_id = _user(mock_invoker, "admin@test.com", is_admin=True) + _create(service, admin_id, "complete") + mock_invoker.services.users.update(admin_id, UserUpdateRequest(is_active=False)) + + _refused_before_any_node(queue, runner[0], "administrators only") + + +def test_an_item_is_refused_once_the_machine_generates_on_two_gpus( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, runner +) -> None: + devices = [GPU] + service = calibration_service(tmp_path, devices=lambda: list(devices)) + service.start(mock_invoker) + mock_invoker.services.working_memory_calibration = service + _create(service, "system", "complete") + devices.append(torch.device("cuda:1")) + + _refused_before_any_node(queue, runner[0], "single generation device") + + +def test_an_item_is_refused_on_a_gpu_of_another_device_class( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, runner +) -> None: + creator = calibration_service(tmp_path) + creator.start(mock_invoker) + run_id, _ = _create(creator, "system", "complete") + # The card was swapped: the same root, another class. + swapped = calibration_service(tmp_path, device_class=OTHER_CLASS) + swapped.start(mock_invoker) + mock_invoker.services.working_memory_calibration = swapped + + _refused_before_any_node(queue, runner[0], f"created for device class {GPU_CLASS.id}") + [run] = next(c for c in creator.get_status().device_classes if c.device_class_id == GPU_CLASS.id).runs + assert run.run_id == run_id and run.status == "failed" + + +def test_an_item_removed_by_clear_queue_mid_run_is_recorded_as_canceled( + tmp_path: Path, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner +) -> None: + session_runner, cancel_event = runner + run_id, item_id = _create(service, "system", "cancel") + + def clear() -> None: + queue.clear("default") + cancel_event.set() + + _ON_CANCEL.append(clear) + item = queue.dequeue() + assert item is not None + session_runner.workflow_call_queue_lifecycle.run_queue_item(item) + + assert queue.get_queue_item_ids("default").item_ids == [] + [run] = service.get_status().device_classes[0].runs + # The node's window was kept; the item, gone from the queue, will not finish. + assert (run.status, run.point_count, run.item_ids) == ("partial", 1, [item_id]) + + +def test_a_runner_failure_outside_the_nodes_is_recorded_as_failed( + service: WorkingMemoryCalibrationService, queue: SessionQueue, runner, monkeypatch: pytest.MonkeyPatch +) -> None: + session_runner, _ = runner + run_id, item_id = _create(service, "system", "complete") + + def explode(queue_item: SessionQueueItem) -> None: + raise RuntimeError("database is locked") + + monkeypatch.setattr(session_runner, "_on_after_run_session", explode) + item = queue.dequeue() + assert item is not None + with pytest.raises(RuntimeError): + session_runner.run(item) + # The processor fails the item next; pruned afterwards, the run keeps the failure the runner recorded. + queue.delete_queue_items_by_id([item_id]) + + [run] = service.get_status().device_classes[0].runs + assert (run.status, run.error) == ("failed", "RuntimeError: database is locked") diff --git a/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py b/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py new file mode 100644 index 00000000000..87e92bea720 --- /dev/null +++ b/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py @@ -0,0 +1,366 @@ +"""The calibration store and service on a temporary root, with a real SQLite queue.""" + +import asyncio +import json +import os +import threading +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from invokeai.app.invocations.math import AddInvocation +from invokeai.app.services.invoker import Invoker +from invokeai.app.services.session_queue.session_queue_common import BatchDatum +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.graph import Graph +from invokeai.app.services.working_memory_calibration import working_memory_calibration_store +from invokeai.app.services.working_memory_calibration.working_memory_calibration_common import ( + CalibrationDeviceClass, + CalibrationStoreError, + CalibrationStoreUnusableError, + WorkingMemoryCalibrationRequest, + WorkingMemoryCalibrationStatus, +) +from invokeai.app.services.working_memory_calibration.working_memory_calibration_default import ( + WorkingMemoryCalibrationService, +) +from tests.fixtures.working_memory_calibration import ( + GPU_CLASS, + OTHER_CLASS, + calibration_service, + measured_window, +) + + +@pytest.fixture +def queue(mock_invoker: Invoker, mock_sqlite_database: Database) -> SessionQueue: + queue = SessionQueue(mock_sqlite_database) + mock_invoker.services.session_queue = queue + model_manager = MagicMock() + model_manager.store.get_model.side_effect = lambda key: SimpleNamespace(key=key, hash=f"blake3:{key}") + mock_invoker.services.model_manager = model_manager + queue.start(mock_invoker) + return queue + + +def _service(root: Path, invoker: Invoker, **kwargs) -> WorkingMemoryCalibrationService: + service = calibration_service(root, **kwargs) + service.start(invoker) + return service + + +def _request(values: tuple[int, ...] = (1,)) -> WorkingMemoryCalibrationRequest: + graph = Graph() + graph.add_node(AddInvocation(id="add", a=1, b=2)) + data = [[BatchDatum(node_path="add", field_name="a", items=list(values))]] if len(values) > 1 else None + return WorkingMemoryCalibrationRequest(graph=graph, data=data) + + +def _create(service: WorkingMemoryCalibrationService, model_key: str = "model-a", values: tuple[int, ...] = (1,)): + return asyncio.run(service.create_run(model_key, _request(values), "system")) + + +def _store_path(root: Path, device_class_id: str = GPU_CLASS.id) -> Path: + return root / "working_memory" / f"{device_class_id}.json" + + +def _run_status(status: WorkingMemoryCalibrationStatus, run_id: str): + return next(run for device_class in status.device_classes for run in device_class.runs if run.run_id == run_id) + + +def test_runs_and_points_round_trip_through_a_fresh_service( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue +) -> None: + service = _service(tmp_path, mock_invoker) + created = _create(service) + service.record_windows( + created.run_id, + [ + measured_window(created.run_id, created.item_ids[0]), + measured_window(created.run_id, created.item_ids[0], reasons=("partial_residency",)), + measured_window(created.run_id, created.item_ids[0], reasons=("spill",)), + ], + ) + + status = _service(tmp_path, mock_invoker).get_status() + + assert status.available and status.current_device_class_id == GPU_CLASS.id + [device_class] = status.device_classes + assert device_class.usable and device_class.is_current + assert device_class.device_class == CalibrationDeviceClass.of(GPU_CLASS) + assert device_class.device_class is not None and device_class.device_class.id == GPU_CLASS.id + [run] = device_class.runs + assert (run.run_id, run.model_key, run.model_hash, run.point_count) == ( + created.run_id, + "model-a", + "blake3:model-a", + 3, + ) + assert run.item_ids == created.item_ids + assert {(count.valid, count.fit_eligible): count.count for count in device_class.point_counts} == { + (True, True): 1, + (True, False): 1, + (False, False): 1, + } + [env] = device_class.envs + assert env.id == run.env_id == status.current_env_id + assert env.hard["setting.attention_backend"] == str(mock_invoker.services.configuration.attention_backend) + assert env.hard_matches_current and env.soft_matches_current and env.point_count == 3 + + +def test_a_failed_write_keeps_the_previous_file_and_leaves_no_staging_file( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + service = _service(tmp_path, mock_invoker) + created = _create(service) + before = _store_path(tmp_path).read_bytes() + + def refuse(source: str, target: str) -> None: + raise OSError("disk full") + + monkeypatch.setattr(os, "replace", refuse) + with pytest.raises(CalibrationStoreError, match="disk full"): + service.record_windows(created.run_id, [measured_window(created.run_id)]) + monkeypatch.undo() + + assert _store_path(tmp_path).read_bytes() == before + assert sorted(path.name for path in (tmp_path / "working_memory").iterdir()) == [ + f"{GPU_CLASS.id}.json", + f"{GPU_CLASS.id}.json.lock", + ] + + +def test_a_replace_refused_by_another_processes_handle_is_retried( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + """On Windows a scanner or indexer holding the file open makes the replace fail with PermissionError for a moment.""" + monkeypatch.setattr(working_memory_calibration_store, "REPLACE_RETRY_S", 0) + service = _service(tmp_path, mock_invoker) + created = _create(service) + real_replace = os.replace + refusals = [PermissionError("in use"), PermissionError("in use")] + + def busy_then_free(source: str, target: str) -> None: + if refusals: + raise refusals.pop() + real_replace(source, target) + + monkeypatch.setattr(os, "replace", busy_then_free) + service.record_windows(created.run_id, [measured_window(created.run_id)]) + assert refusals == [] + assert len(json.loads(_store_path(tmp_path).read_bytes())["points"]) == 1 + + def always_busy(source: str, target: str) -> None: + raise PermissionError("in use") + + monkeypatch.setattr(os, "replace", always_busy) + before = _store_path(tmp_path).read_bytes() + with pytest.raises(CalibrationStoreError, match="in use"): + service.record_windows(created.run_id, [measured_window(created.run_id)]) + monkeypatch.undo() + assert _store_path(tmp_path).read_bytes() == before + assert not any(path.name.endswith(".tmp") for path in (tmp_path / "working_memory").iterdir()) + + +def test_two_services_sharing_a_root_keep_each_others_writes( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue +) -> None: + # Two instances stand for two processes: separate process locks, so only the OS file lock serializes them. + first, second = _service(tmp_path, mock_invoker), _service(tmp_path, mock_invoker) + created = _create(first) + writes_each = 15 + + def write(service: WorkingMemoryCalibrationService, offset: int) -> None: + for index in range(writes_each): + service.record_windows(created.run_id, [measured_window(created.run_id, node_id=f"n{offset + index}")]) + + threads = [ + threading.Thread(target=write, args=(service, offset)) for service, offset in ((first, 0), (second, 100)) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=60) + + stored = json.loads(_store_path(tmp_path).read_bytes()) + assert sorted(point["node_id"] for point in stored["points"]) == sorted( + [f"n{index}" for index in range(writes_each)] + [f"n{100 + index}" for index in range(writes_each)] + ) + + +def _with_extra(field: str): + def mutate(valid: dict) -> bytes: + target = valid["runs"][0] if field == "run" else valid["points"][0] + target["added_by_a_newer_writer"] = True + return json.dumps(valid).encode() + + return mutate + + +@pytest.mark.parametrize( + "mutate", + [ + lambda valid: json.dumps({**valid, "schema_version": 2, "future": []}).encode(), + lambda valid: b'{"schema_version": 1, "device_class": "truncated', + lambda valid: json.dumps({**valid, "device_class": OTHER_CLASS.model_dump()}).encode(), + lambda valid: json.dumps({**valid, "device_class": {**valid["device_class"], "total_mib": 1}}).encode(), + _with_extra("run"), + _with_extra("point"), + ], + ids=[ + "newer-schema", + "invalid-json", + "other-device-class", + "edited-device-class", + "extra-run-field", + "extra-point-field", + ], +) +def test_a_file_this_version_cannot_read_is_reported_and_never_written( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, mutate +) -> None: + """One policy for every store model: an unknown field makes the file unusable rather than being dropped on the + next rewrite.""" + service = _service(tmp_path, mock_invoker) + created = _create(service) + service.record_windows(created.run_id, [measured_window(created.run_id)]) + path = _store_path(tmp_path) + content = mutate(json.loads(path.read_bytes())) + path.write_bytes(content) + item_ids = queue.get_queue_item_ids("default").item_ids + + status = service.get_status() + [device_class] = status.device_classes + assert not device_class.usable and device_class.problem and device_class.runs == [] + assert not status.available and "cannot be written" in (status.unavailable_reason or "") + with pytest.raises(CalibrationStoreUnusableError): + _create(service) + with pytest.raises(CalibrationStoreUnusableError): + service.reset_all() + with pytest.raises(CalibrationStoreUnusableError): + service.reset_model("model-a") + service.record_windows(created.run_id, [measured_window(created.run_id)]) + + assert path.read_bytes() == content + assert service.export() == content + # The failed create enqueued nothing. + assert queue.get_queue_item_ids("default").item_ids == item_ids + + +def test_resets_remove_runs_points_and_unused_environments( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue +) -> None: + service = _service(tmp_path, mock_invoker) + run_a1, run_a2, run_b = _create(service, "model-a"), _create(service, "model-a"), _create(service, "model-b") + for run in (run_a1, run_a2, run_b): + service.record_windows( + run.run_id, [measured_window(run.run_id, model_key="transformer"), measured_window(run.run_id)] + ) + # A window of model-a measured inside model-b's run. + service.record_windows(run_b.run_id, [measured_window(run_b.run_id, model_key="model-a")]) + + removed = service.reset_run(run_a1.run_id) + assert (removed.runs_removed, removed.points_removed) == (1, 2) + + removed = service.reset_model("model-a") + assert (removed.runs_removed, removed.points_removed) == (1, 3) + [device_class] = service.get_status().device_classes + assert [run.run_id for run in device_class.runs] == [run_b.run_id] + assert device_class.runs[0].point_count == 2 + + removed = service.reset_all() + assert (removed.runs_removed, removed.points_removed) == (1, 2) + stored = json.loads(_store_path(tmp_path).read_bytes()) + assert (stored["runs"], stored["points"], stored["envs"]) == ([], [], []) + assert stored["device_class"]["id"] == GPU_CLASS.id + + +def test_windows_of_a_reset_run_are_dropped(tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue) -> None: + service = _service(tmp_path, mock_invoker) + created = _create(service) + service.reset_run(created.run_id) + + service.record_windows(created.run_id, [measured_window(created.run_id)]) + + assert json.loads(_store_path(tmp_path).read_bytes())["points"] == [] + + +def test_run_status_follows_its_items_retries_and_observations( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue +) -> None: + service = _service(tmp_path, mock_invoker) + created = _create(service, values=(1, 2)) + first, second = created.item_ids + + assert _run_status(service.get_status(), created.run_id).status == "running" + + for _ in range(2): + item = queue.dequeue() + assert item is not None + assert service.admit(item, created.run_id) is None + if item.item_id == first: + queue.complete_queue_item(item.item_id) + else: + queue.fail_queue_item(item.item_id, "ValueError", "out of memory", "traceback") + service.observe(item, created.run_id) + run = _run_status(service.get_status(), created.run_id) + assert (run.status, run.error) == ("failed", "out of memory") + + # The retry stands in for the failed item. + queue.retry_items_by_id("default", [second]) + run = _run_status(service.get_status(), created.run_id) + retry = max(run.item_ids) + assert retry not in (first, second) and run.status == "running" + retried = queue.dequeue() + assert retried is not None and retried.item_id == retry + assert service.admit(retried, created.run_id) is None + queue.cancel_queue_item(retry) + service.observe(retried, created.run_id) + assert _run_status(service.get_status(), created.run_id).status == "canceled" + + service.record_windows(created.run_id, [measured_window(created.run_id, first)]) + assert _run_status(service.get_status(), created.run_id).status == "partial" + + # Pruned from the queue, the items keep the status the runner saw. + queue.retry_items_by_id("default", [retry]) + rerun = queue.dequeue() + assert rerun is not None + assert service.admit(rerun, created.run_id) is None + queue.complete_queue_item(rerun.item_id) + service.observe(rerun, created.run_id) + queue.delete_queue_items_by_id([first, second, retry, rerun.item_id]) + run = _run_status(service.get_status(), created.run_id) + assert run.status == "completed" + assert run.item_ids == sorted([first, second, retry, rerun.item_id]) + + +def test_an_item_gone_from_the_queue_while_running_counts_as_canceled( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue +) -> None: + """The process died mid-item, so the runner never recorded an outcome; then the queue history was pruned.""" + service = _service(tmp_path, mock_invoker) + created = _create(service) + item = queue.dequeue() + assert item is not None and service.admit(item, created.run_id) is None + assert _run_status(service.get_status(), created.run_id).status == "running" + + queue.delete_queue_items_by_id([item.item_id]) + + assert _run_status(service.get_status(), created.run_id).status == "canceled" + + +def test_status_reports_machines_that_cannot_calibrate( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue +) -> None: + import torch + + service = _service(tmp_path, mock_invoker, devices=lambda: [torch.device("cuda:0"), torch.device("cuda:1")]) + status = service.get_status() + assert not status.available and "single generation device" in (status.unavailable_reason or "") + assert status.current_device_class_id is None and status.device_classes == [] + + service = _service(tmp_path, mock_invoker, devices=lambda: [torch.device("cpu")]) + assert "GPU memory" in (service.get_status().unavailable_reason or "") diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_locks_held.py b/tests/backend/model_manager/load/model_cache/test_model_cache_locks_held.py new file mode 100644 index 00000000000..dd6ce523bc5 --- /dev/null +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_locks_held.py @@ -0,0 +1,48 @@ +"""`ModelCache.locks_held`: every pin on the cache's records, whichever thread took it and of whichever kind. The +working-memory probe compares it with its own thread's locks to see another thread sharing the device.""" + +import logging +import threading +from unittest.mock import MagicMock + +import pytest +import torch + +from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache + + +@pytest.fixture +def cache(): + logger = MagicMock() + logger.getEffectiveLevel.return_value = logging.INFO + cache = ModelCache( + execution_device_working_mem_gb=1.0, + enable_partial_loading=False, + keep_ram_copy_of_weights=True, + execution_device="cpu", + storage_device="cpu", + logger=logger, + shared_cpu_weights=None, + ) + yield cache + cache.shutdown() + + +def test_counts_device_and_ram_locks_of_every_thread(cache: ModelCache): + cache.put("vae", torch.nn.Linear(2, 2)) + cache.put("lora", torch.nn.Linear(2, 2)) + vae, lora = cache.get("vae"), cache.get("lora") + + cache.lock(vae, None) + cache.lock(vae, None) + other = threading.Thread(target=cache.lock_in_ram, args=(lora,)) + other.start() + other.join() + + assert cache.locks_held() == 3 + + cache.unlock(vae) + cache.unlock(lora) + assert cache.locks_held() == 1 + cache.unlock(vae) + assert cache.locks_held() == 0 diff --git a/tests/backend/util/test_nvml.py b/tests/backend/util/test_nvml.py new file mode 100644 index 00000000000..1de9730fad0 --- /dev/null +++ b/tests/backend/util/test_nvml.py @@ -0,0 +1,111 @@ +"""The NVML readers, against a scripted stand-in for the NVML library. + +The stand-in fills the same ctypes structures the module passes, so the two-call process enumeration, the PCI lookup +and the struct layout are exercised; only the driver behind them is fake. +""" + +import ctypes +import functools +import os +from types import SimpleNamespace +from typing import Optional + +import pytest +import torch + +from invokeai.backend.util import nvml + +GIB = 2**30 +DEVICE = torch.device("cuda", 0) +OUR_BUS_ID = b"0000:03:00.0" + + +class _FakeNvml: + """A device at `OUR_BUS_ID` running `processes` as (pid, used bytes).""" + + def __init__(self, processes: list[tuple[int, int]], driver: str = "581.57"): + self.processes = processes + self.driver = driver + self.bus_ids: list[bytes] = [] + + def nvmlDeviceGetHandleByPciBusId_v2(self, bus_id: bytes, handle_ref) -> int: + self.bus_ids.append(bus_id) + if bus_id != OUR_BUS_ID: + return 6 # NVML_ERROR_NOT_FOUND + handle_ref._obj.value = 0x5EED + return 0 + + def nvmlDeviceGetComputeRunningProcesses_v3(self, handle, count_ref, infos) -> int: + count = count_ref._obj + if infos is None or count.value < len(self.processes): + count.value = len(self.processes) + return 7 if self.processes else 0 # NVML_ERROR_INSUFFICIENT_SIZE + for info, (pid, used) in zip(infos, self.processes, strict=False): + info.pid, info.usedGpuMemory = pid, used + count.value = len(self.processes) + return 0 + + def nvmlSystemGetDriverVersion(self, buffer, length: int) -> int: + buffer.value = self.driver.encode() + return 0 + + +@pytest.fixture +def install(monkeypatch: pytest.MonkeyPatch): + """A CUDA build whose device 0 sits at PCI bus 3. Returns an installer for a fake NVML (or None for no NVML).""" + monkeypatch.setattr(torch.version, "hip", None) + props = SimpleNamespace(pci_domain_id=0, pci_bus_id=3, pci_device_id=0) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda index: props) + nvml.reset_cache() + + def installer(fake: Optional[_FakeNvml]) -> Optional[_FakeNvml]: + monkeypatch.setattr(nvml, "_load", functools.lru_cache(maxsize=1)(lambda: fake)) + return fake + + yield installer + nvml.reset_cache() + + +def test_reads_this_processs_memory_on_the_device_at_its_pci_address(install): + me = os.getpid() + fake = install(_FakeNvml([(me + 1, 9 * GIB), (me, 3 * GIB), (me + 2, GIB)])) + + assert nvml.process_used_memory(DEVICE) == 3 * GIB + assert nvml.process_used_memory(DEVICE) == 3 * GIB + assert fake.bus_ids == [OUR_BUS_ID], "the handle is looked up once" + + +@pytest.mark.parametrize( + "processes", + [[(1, GIB)], [], [(os.getpid(), 0xFFFFFFFFFFFFFFFF)]], + ids=["not-listed", "no-processes", "not-attributed"], +) +def test_unknown_when_nvml_does_not_attribute_memory_to_this_process(install, processes): + install(_FakeNvml(processes)) + assert nvml.process_used_memory(DEVICE) is None + + +def test_unknown_without_nvml_off_cuda_and_when_a_call_raises(install, monkeypatch: pytest.MonkeyPatch): + install(None) + assert nvml.process_used_memory(DEVICE) is None + assert nvml.driver_version() is None + + fake = install(_FakeNvml([(os.getpid(), GIB)])) + monkeypatch.setattr(torch.version, "hip", "7.2.0") + assert nvml.process_used_memory(DEVICE) is None + assert fake.bus_ids == [], "a HIP build never asks NVML" + + monkeypatch.setattr(torch.version, "hip", None) + monkeypatch.setattr(fake, "nvmlDeviceGetComputeRunningProcesses_v3", lambda *args: 1 / 0) + assert nvml.process_used_memory(DEVICE) is None + + +def test_reads_the_driver_version(install): + install(_FakeNvml([], driver="581.57")) + assert nvml.driver_version() == "581.57" + + +def test_the_process_entry_matches_the_nvml_layout(): + """The driver writes entries by offset: a shifted field would read a neighbour's bytes, not fail.""" + assert nvml._ProcessInfo.usedGpuMemory.offset == 8 + assert ctypes.sizeof(nvml._ProcessInfo) == 24 diff --git a/tests/backend/util/test_wddm.py b/tests/backend/util/test_wddm.py index a2c53b65ea3..a2c21a1e2a4 100644 --- a/tests/backend/util/test_wddm.py +++ b/tests/backend/util/test_wddm.py @@ -123,7 +123,7 @@ def windows_rocm(monkeypatch: pytest.MonkeyPatch): """A Windows ROCm build whose device 0 sits at PCI bus 3 with 16 GiB. Returns an installer for a fake gdi32.""" monkeypatch.setattr(sys, "platform", "win32") monkeypatch.setattr(torch.version, "hip", "7.14.0") - props = SimpleNamespace(pci_bus_id=OUR_BUS, pci_device_id=OUR_DEVICE, total_memory=16 * GIB) + props = SimpleNamespace(pci_domain_id=0, pci_bus_id=OUR_BUS, pci_device_id=OUR_DEVICE, total_memory=16 * GIB) monkeypatch.setattr(torch.cuda, "get_device_properties", lambda index: props) wddm.reset_cache() @@ -211,6 +211,40 @@ def test_reset_closes_the_cached_handle(windows_rocm): assert gdi32.closed == [20] +@pytest.mark.parametrize("hip", ["7.14.0", None], ids=["rocm", "cuda"]) +def test_process_usage_is_read_on_either_windows_build(windows_rocm, monkeypatch, hip): + """The working-memory probe measures on CUDA too; the adapter is found by PCI location there just the same.""" + gdi32 = windows_rocm(_FakeGdi32([(10, 7, 0), (20, OUR_BUS, OUR_DEVICE)], usage=5 * GIB)) + monkeypatch.setattr(torch.version, "hip", hip) + + assert wddm.process_local_usage(DEVICE) == 5 * GIB + assert gdi32.closed == [10] + + +@pytest.mark.parametrize( + "kwargs", [{"usage": 17 * GIB}, {"query_status": -1073741823}], ids=["usage-beyond-the-card", "query-fails"] +) +def test_process_usage_is_unknown_rather_than_guessed(windows_rocm, monkeypatch, kwargs): + windows_rocm(_FakeGdi32([(20, OUR_BUS, OUR_DEVICE)], **kwargs)) + monkeypatch.setattr(torch.version, "hip", None) + + assert wddm.process_local_usage(DEVICE) is None + + +def test_the_ungated_readers_elsewhere_never_touch_the_os(monkeypatch): + def loader(): + raise AssertionError("no Windows library may be loaded off Windows") + + monkeypatch.setattr(wddm, "_load_gdi32", functools.lru_cache(maxsize=1)(loader)) + monkeypatch.setattr(wddm, "_load_pdh", functools.lru_cache(maxsize=1)(loader)) + monkeypatch.setattr(sys, "platform", "linux") + + assert wddm.process_local_usage(DEVICE) is None + assert wddm.process_paged_bytes(DEVICE) is None + monkeypatch.setattr(sys, "platform", "win32") + assert wddm.process_local_usage(torch.device("cpu")) is None + + def _instance(pid: int, handle: int) -> str: return f"pid_{pid}_luid_0x00000000_0x{_luid_low(handle):08X}_phys_0" @@ -264,6 +298,17 @@ def test_paged_bytes_is_unknown_before_this_process_used_the_adapter(windows_roc assert wddm.paged_bytes(DEVICE) is None +def test_on_windows_cuda_only_the_ungated_reader_answers_the_paged_amount(windows_rocm, monkeypatch): + """The model cache's spill warning stays a ROCm affair; the probe's spill check needs the figure on CUDA too.""" + windows_rocm(_FakeGdi32([(20, OUR_BUS, OUR_DEVICE)])) + monkeypatch.setattr(torch.version, "hip", None) + pdh = _FakePdh({_instance(os.getpid(), 20): 3 * GIB}) + monkeypatch.setattr(wddm, "_load_pdh", functools.lru_cache(maxsize=1)(lambda: pdh)) + + assert wddm.process_paged_bytes(DEVICE) == 3 * GIB + assert wddm.paged_bytes(DEVICE) is None + + def test_paged_bytes_elsewhere_never_touches_pdh(monkeypatch): def loader(): raise AssertionError("pdh must not be loaded outside Windows ROCm") diff --git a/tests/backend/working_memory/test_counters.py b/tests/backend/working_memory/test_counters.py new file mode 100644 index 00000000000..8e7fe4684c3 --- /dev/null +++ b/tests/backend/working_memory/test_counters.py @@ -0,0 +1,203 @@ +"""The device counters: which driver reader each platform uses, the DRM fdinfo reader, and the MPS peak sampler. + +The OS and torch readers behind them are stand-ins: what is under test is that each platform reads through its own +reader, labels the reading with it, and turns any failure into "unknown". +""" + +import sys +import threading +from pathlib import Path +from typing import Callable + +import pytest +import torch + +from invokeai.backend.util import nvml, wddm +from invokeai.backend.working_memory import counters + +GIB = 2**30 +MIB = 2**20 +PDEV = "0000:03:00.0" +CUDA0 = torch.device("cuda", 0) + +_STATS = { + "reserved_bytes.all.current": 3 * GIB, + "allocated_bytes.all.current": 2 * GIB, + "reserved_bytes.all.peak": 5 * GIB, + "allocated_bytes.all.peak": 4 * GIB, + "num_ooms": 0, + "num_alloc_retries": 1, +} + + +def _fdinfo(pdev: str, client: int, vram: str) -> str: + return ( + "pos:\t0\nflags:\t02100002\nmnt_id:\t26\nino:\t1063\ndrm-driver:\tamdgpu\n" + f"drm-client-id:\t{client}\ndrm-pdev:\t{pdev}\npasid:\t32790\ndrm-memory-vram:\t{vram}\n" + "drm-memory-gtt:\t2048 KiB\ndrm-memory-cpu:\t0 KiB\namd-evicted-vram:\t0 KiB\n" + ) + + +def test_fdinfo_sums_each_client_of_the_device_once(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + (tmp_path / "3").write_text("pos:\t0\nflags:\t0100000\nmnt_id:\t26\nino:\t1063\n") # not a DRM file + (tmp_path / "7").write_text(_fdinfo(PDEV, 11, "1048576 KiB")) + (tmp_path / "8").write_text(_fdinfo(PDEV, 11, "1048576 KiB")) # a duplicated descriptor of the same client + (tmp_path / "9").write_text(_fdinfo(PDEV, 12, "512 MiB")) + (tmp_path / "10").write_text(_fdinfo("0000:0c:00.0", 13, "4194304 KiB")) # the integrated GPU + monkeypatch.setattr(counters, "_FDINFO_DIR", tmp_path) + + assert counters.fdinfo_vram_bytes(PDEV) == 1 * GIB + 512 * MIB + assert counters.fdinfo_vram_bytes("0000:0C:00.0") == 4 * GIB + + +def test_fdinfo_without_a_client_of_the_device_is_unknown(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + (tmp_path / "7").write_text(_fdinfo("0000:0c:00.0", 13, "4194304 KiB")) + monkeypatch.setattr(counters, "_FDINFO_DIR", tmp_path) + assert counters.fdinfo_vram_bytes(PDEV) is None + + monkeypatch.setattr(counters, "_FDINFO_DIR", tmp_path / "missing") + assert counters.fdinfo_vram_bytes(PDEV) is None + + +def _raise(*_args: object) -> None: + raise RuntimeError("the driver did not answer") + + +def _windows(monkeypatch: pytest.MonkeyPatch, reader: Callable) -> None: + monkeypatch.setattr(wddm, "process_local_usage", reader) + monkeypatch.setattr(wddm, "process_paged_bytes", lambda device: 64 * MIB) + + +def _nvml(monkeypatch: pytest.MonkeyPatch, reader: Callable) -> None: + monkeypatch.setattr(nvml, "process_used_memory", reader) + + +def _fdinfo_reader(monkeypatch: pytest.MonkeyPatch, reader: Callable) -> None: + monkeypatch.setattr(counters, "pci_address", lambda index: PDEV) + monkeypatch.setattr(counters, "fdinfo_vram_bytes", reader) + + +_CUDA_PLATFORMS = [ + pytest.param("win32", None, _windows, "windows-d3dkmt", id="windows-cuda"), + pytest.param("win32", "7.14.0", _windows, "windows-d3dkmt", id="windows-rocm"), + pytest.param("linux", None, _nvml, "nvml", id="linux-cuda"), + pytest.param("linux", "7.2.0", _fdinfo_reader, "fdinfo", id="linux-rocm"), +] + + +@pytest.mark.parametrize(("platform", "hip", "install", "method"), _CUDA_PLATFORMS) +def test_each_gpu_platform_reads_its_driver_through_its_own_reader( + monkeypatch: pytest.MonkeyPatch, platform: str, hip: str | None, install: Callable, method: str +): + monkeypatch.setattr(sys, "platform", platform) + monkeypatch.setattr(torch.version, "hip", hip) + monkeypatch.setattr(torch.cuda, "memory_stats", lambda device: dict(_STATS)) + install(monkeypatch, lambda device: 6 * GIB) + + reading = counters.read_counters(CUDA0) + + assert (reading.driver, reading.method) == (6 * GIB, method) + assert (reading.reserved, reading.allocated, reading.peak_reserved, reading.peak_allocated) == ( + 3 * GIB, + 2 * GIB, + 5 * GIB, + 4 * GIB, + ) + assert (reading.num_ooms, reading.num_alloc_retries) == (0, 1) + assert reading.paged == (64 * MIB if platform == "win32" else None) + + +@pytest.mark.parametrize(("platform", "hip", "install", "method"), _CUDA_PLATFORMS) +def test_a_failing_reader_reads_as_unknown_and_never_raises( + monkeypatch: pytest.MonkeyPatch, platform: str, hip: str | None, install: Callable, method: str +): + monkeypatch.setattr(sys, "platform", platform) + monkeypatch.setattr(torch.version, "hip", hip) + monkeypatch.setattr(torch.cuda, "memory_stats", _raise) + install(monkeypatch, _raise) + + reading = counters.read_counters(CUDA0) + + assert (reading.driver, reading.reserved, reading.peak_allocated, reading.method) == ( + None, + None, + None, + "torch-only", + ) + + +def test_xpu_reads_the_device_wide_usage_as_a_guard_not_as_this_processs_driver(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(torch.xpu, "memory_stats", lambda device: dict(_STATS)) + monkeypatch.setattr(counters, "xpu_memory_info", lambda device: (10 * GIB, 16 * GIB)) + + reading = counters.read_counters(torch.device("xpu", 0)) + + assert (reading.driver, reading.device_used, reading.method) == (None, 6 * GIB, "xpu-sysman") + assert reading.reserved == 3 * GIB + + monkeypatch.setattr(counters, "xpu_memory_info", lambda device: None) + assert counters.read_counters(torch.device("xpu", 0)).method == "torch-only" + + +class _FakeMps: + """The two MPS getters, over a value the test moves.""" + + def __init__(self) -> None: + self.allocated, self.driver = 1 * GIB, 2 * GIB + self.spike_reads = 0 + self.spike_seen = threading.Event() + + def current_allocated_memory(self) -> int: + if self.allocated == 5 * GIB and threading.current_thread() is not threading.main_thread(): + # Two sampler reads of the spike: the first has been folded into the peaks by the time of the second. + self.spike_reads += 1 + if self.spike_reads >= 2: + self.spike_seen.set() + return self.allocated + + def driver_allocated_memory(self) -> int: + return self.driver + + +def test_on_mps_a_sampler_keeps_the_peak_that_came_and_went_between_two_reads(monkeypatch: pytest.MonkeyPatch): + fake = _FakeMps() + monkeypatch.setattr(torch.mps, "current_allocated_memory", fake.current_allocated_memory) + monkeypatch.setattr(torch.mps, "driver_allocated_memory", fake.driver_allocated_memory) + monkeypatch.setattr(torch.mps, "synchronize", lambda: None) + monkeypatch.setattr(counters, "PEAK_SAMPLER_INTERVAL_S", 0.0005) + + device_counters = counters.DeviceCounters(torch.device("mps")) + try: + fake.allocated, fake.driver = 5 * GIB, 7 * GIB + assert fake.spike_seen.wait(timeout=30) + fake.allocated, fake.driver = 1 * GIB, 2 * GIB + + reading = device_counters.read() + assert (reading.allocated, reading.driver, reading.method) == (1 * GIB, 2 * GIB, "mps-driver") + assert (reading.peak_allocated, reading.peak_driver) == (5 * GIB, 7 * GIB) + assert reading.peak_reserved is None, "MPS reserves nothing torch reports" + + device_counters.reset_peaks() + assert (device_counters.read().peak_allocated, device_counters.read().peak_driver) == (1 * GIB, 2 * GIB) + finally: + device_counters.close() + assert not device_counters._sampler.running # type: ignore[union-attr] + + +def test_on_mps_failing_getters_read_as_unknown_and_the_sampler_still_stops(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(torch.mps, "current_allocated_memory", _raise) + monkeypatch.setattr(torch.mps, "driver_allocated_memory", _raise) + + device_counters = counters.DeviceCounters(torch.device("mps")) + try: + reading = device_counters.read() + finally: + device_counters.close() + + assert (reading.allocated, reading.peak_allocated, reading.driver, reading.method) == ( + None, + None, + None, + "torch-only", + ) + assert not device_counters._sampler.running # type: ignore[union-attr] diff --git a/tests/backend/working_memory/test_fingerprint.py b/tests/backend/working_memory/test_fingerprint.py new file mode 100644 index 00000000000..d348055545b --- /dev/null +++ b/tests/backend/working_memory/test_fingerprint.py @@ -0,0 +1,119 @@ +"""Device classes and environments: what a measurement is keyed by, from scripted device properties.""" + +import platform +from types import SimpleNamespace + +import pytest +import torch + +from invokeai.backend.util import nvml +from invokeai.backend.working_memory.fingerprint import DeviceClass, device_class_for, environment_for + +GIB = 2**30 +CUDA0 = torch.device("cuda", 0) + + +def _cuda_cards(monkeypatch: pytest.MonkeyPatch, cards: list[tuple[str, SimpleNamespace]], hip: str | None) -> None: + monkeypatch.setattr(torch.version, "hip", hip) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda device: cards[torch.device(device).index][1]) + monkeypatch.setattr(torch.cuda, "get_device_name", lambda device: cards[torch.device(device).index][0]) + + +def _props(total_bytes: int, bus: int, minor: int = 9, gfx: str = "") -> SimpleNamespace: + return SimpleNamespace(total_memory=total_bytes, major=8, minor=minor, pci_bus_id=bus, gcnArchName=gfx) + + +def test_identical_cards_share_a_device_class_and_other_cards_do_not(monkeypatch: pytest.MonkeyPatch): + rtx_4090 = "NVIDIA GeForce RTX 4090" + _cuda_cards( + monkeypatch, + [ + (rtx_4090, _props(24 * GIB, bus=1)), + (rtx_4090, _props(24 * GIB, bus=2)), # the same card in another slot + ("NVIDIA GeForce RTX 4080", _props(16 * GIB, bus=3)), + ], + hip=None, + ) + + first, second, other = (device_class_for(torch.device("cuda", index)) for index in range(3)) + + assert first == second and first.id == second.id + assert other.id != first.id + assert (first.vendor, first.name, first.total_mib, first.backend, first.arch) == ( + "nvidia", + rtx_4090, + 24 * 1024, + "cuda", + "sm_89", + ) + assert DeviceClass.model_validate_json(first.model_dump_json()).id == first.id + + +def test_a_hip_device_class_names_its_gfx_architecture(monkeypatch: pytest.MonkeyPatch): + _cuda_cards(monkeypatch, [("AMD Radeon RX 9060 XT", _props(16 * GIB, bus=3, gfx="gfx1200"))], hip="7.14.0") + + device_class = device_class_for(CUDA0) + + assert (device_class.vendor, device_class.backend, device_class.arch) == ("amd", "hip", "gfx1200") + + +def test_a_cpu_has_no_device_class(): + with pytest.raises(ValueError, match="not a GPU"): + device_class_for(torch.device("cpu")) + + +_SETTINGS = { + "pytorch_cuda_alloc_conf": None, + "attention_type": "sdpa", + "attention_slice_size": "auto", + "sequential_guidance": False, + "device_working_mem_gb": 3.0, +} + + +@pytest.fixture +def cuda_build(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(torch.version, "hip", None) + monkeypatch.setattr(nvml, "driver_version", lambda: "581.57") + monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False) + + +def test_settings_are_hard_facts_and_full_versions_soft_ones(cuda_build): + environment = environment_for(CUDA0, _SETTINGS) + + assert environment == environment_for(CUDA0, dict(_SETTINGS)), "deterministic" + assert {key: environment.hard[f"setting.{key}"] for key in _SETTINGS} == { + "pytorch_cuda_alloc_conf": "None", + "attention_type": "sdpa", + "attention_slice_size": "auto", + "sequential_guidance": "False", + "device_working_mem_gb": "3.0", + } + assert { + "torch", + "cuda", + "cudnn", + "allocator_backend", + "rocm_sdpa_guard", + "env.PYTORCH_CUDA_ALLOC_CONF", + "env.TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", + "env.INVOKEAI_ROCM_CONV3D", + } <= set(environment.hard) + assert environment.hard["torch"].count(".") == 1, "major.minor only" + assert set(environment.soft) == {"torch_build", "driver", "os", "invokeai"} + assert environment.soft["driver"] == "581.57" + + +def test_only_hard_facts_change_the_environment_id(cuda_build, monkeypatch: pytest.MonkeyPatch): + reference = environment_for(CUDA0, _SETTINGS).id + + monkeypatch.setattr(platform, "platform", lambda: "Windows-11-10.0.26100-SP0") + monkeypatch.setattr(nvml, "driver_version", lambda: "590.01") + assert environment_for(CUDA0, _SETTINGS).id == reference + + assert environment_for(CUDA0, {**_SETTINGS, "attention_type": "sliced"}).id != reference + + monkeypatch.setenv("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") + changed = environment_for(CUDA0, _SETTINGS) + assert changed.hard["env.PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" + assert changed.id != reference diff --git a/tests/backend/working_memory/test_probe.py b/tests/backend/working_memory/test_probe.py new file mode 100644 index 00000000000..18f89e62d43 --- /dev/null +++ b/tests/backend/working_memory/test_probe.py @@ -0,0 +1,775 @@ +"""The working-memory probe's windows, measured on a scripted device through the model loader's real lock path. + +`FakeDevice` stands in for a GPU with a caching allocator: live tensors (allocated), the blocks the allocator holds +(reserved = allocated + cached), what the driver counts, peaks since the last reset, and a release that returns the +cached blocks when `TorchDevice.empty_cache` runs -- as a real device does. Each test scripts allocations through it and +reads the `MeasuredWindow` the probe recorded; expected values follow from the script alone. +""" + +import sys +import threading +from contextlib import contextmanager +from typing import Iterator, Optional + +import pytest +import torch + +from invokeai.app.services.session_processor.session_processor_common import CanceledException +from invokeai.backend.model_manager.load.load_base import LoadedModelWithoutConfig +from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord +from invokeai.backend.util import devices +from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.working_memory import counters, probe +from invokeai.backend.working_memory.counters import CounterReading, Method +from invokeai.backend.working_memory.estimate import WorkingMemoryEstimate, WorkingMemoryTerm +from invokeai.backend.working_memory.probe import MeasuredWindow, ProbeScope, calibration_scope + +GIB = 2**30 +MIB = 2**20 + + +class FakeDevice: + """A GPU as the probe reads it. `retained_pool` makes it behave like HIP under Windows: released blocks stay counted + by the driver and serve the next growth of the allocator before the driver grows.""" + + def __init__(self, *, method: Method = "windows-d3dkmt", driver: bool = True, retained_pool: bool = False): + self.method: Method = method + self.has_driver = driver + self.retained_pool = retained_pool + self.readable = True + self.allocated = 0 + self.cached = 0 + self.retained = 0 + self.peak_reserved = 0 + self.peak_allocated = 0 + self.num_ooms = 0 + self.num_alloc_retries = 0 + self.paged: list[int] = [0] + self.device_used_by_others: Optional[int] = None + self.releases = 0 + self.fail_on_release: Optional[int] = None + self.fail_next_read = False + + @property + def reserved(self) -> int: + return self.allocated + self.cached + + def alloc(self, n: int, *, reserve: Optional[int] = None) -> None: + """A live tensor of `n` bytes, from cached blocks first. `reserve` overrides how much the allocator grows.""" + from_cache = min(self.cached, n) + growth = n - from_cache if reserve is None else reserve + self.cached -= from_cache + self.retained -= min(self.retained, growth) + self.allocated += n + # A reserve override (cudaMallocAsync) grows the pool by less than the tensor; cached absorbs the difference. + self.cached += growth - (n - from_cache) + self.peak_allocated = max(self.peak_allocated, self.allocated) + self.peak_reserved = max(self.peak_reserved, self.reserved) + + def free(self, n: int) -> None: + self.allocated -= n + self.cached += n + + def reading(self) -> CounterReading: + if self.fail_next_read: + self.fail_next_read = False + raise RuntimeError("the driver did not answer") + if not self.readable: + return CounterReading(None, None, None, None, None, None, None, None, None, None, "torch-only") + paged = self.paged.pop(0) if len(self.paged) > 1 else self.paged[0] + return CounterReading( + reserved=self.reserved, + allocated=self.allocated, + peak_reserved=self.peak_reserved, + peak_allocated=self.peak_allocated, + num_ooms=self.num_ooms, + num_alloc_retries=self.num_alloc_retries, + driver=self.reserved + self.retained if self.has_driver else None, + peak_driver=None, + paged=paged, + device_used=None if self.device_used_by_others is None else self.reserved + self.device_used_by_others, + method=self.method if self.has_driver or self.device_used_by_others is not None else "torch-only", + ) + + # The caching allocator's side of `TorchDevice.empty_cache`. + def before_empty_cache(self) -> None: + pass + + def after_empty_cache(self) -> None: + self.releases += 1 + if self.releases == self.fail_on_release: + raise RuntimeError("the device context is sick") + if self.retained_pool: + self.retained += self.cached + self.cached = 0 + + +class FakeCounters: + """`counters.DeviceCounters` over a `FakeDevice`.""" + + def __init__(self, device: FakeDevice): + self.device = device + self.closed = False + + def read(self) -> CounterReading: + return self.device.reading() + + def reset_peaks(self) -> None: + self.device.peak_reserved = self.device.reserved + self.device.peak_allocated = self.device.allocated + + def synchronize(self) -> None: + pass + + def close(self) -> None: + self.closed = True + + +class FakeCachedModel: + """A model's weights as the cache sees them: a size, how much of it is on the device, and where it computes.""" + + def __init__(self, total_bytes: int, device: str): + self.total = total_bytes + self.resident = 0 + self.compute_device = torch.device(device) + self.model = object() + + def total_bytes(self) -> int: + return self.total + + def cur_vram_bytes(self) -> int: + return self.resident + + def get_cpu_state_dict(self) -> None: + return None + + +class FakeCache: + """A model cache on a `FakeDevice`: a lock puts the model's weights there, through a staging buffer of + `staging_bytes`, after offloading the models named in `offloads`; `partial` caps a model's residency.""" + + def __init__(self, device: FakeDevice, staging_bytes: int = 0): + self.device = device + self.staging_bytes = staging_bytes + self.reserved: list[Optional[int]] = [] + self.offloads: dict[str, list[CacheRecord]] = {} + self.partial: dict[str, int] = {} + self._locks: dict[str, int] = {} + + def lock(self, record: CacheRecord, working_mem_bytes: Optional[int], max_move_bytes: Optional[int] = None) -> bool: + self.reserved.append(working_mem_bytes) + for victim in self.offloads.get(record.key, []): + self.device.free(victim.cached_model.resident) + victim.cached_model.resident = 0 + model = record.cached_model + if model.compute_device.type != "cpu": + target = self.partial.get(record.key, model.total) + if model.resident < target: + self.device.alloc(self.staging_bytes) + self.device.alloc(target - model.resident) + self.device.free(self.staging_bytes) + model.resident = target + self._locks[record.key] = self._locks.get(record.key, 0) + 1 + return True + + def lock_in_ram(self, record: CacheRecord) -> None: + self._locks[record.key] = self._locks.get(record.key, 0) + 1 + + def unlock(self, record: CacheRecord) -> None: + self._locks[record.key] -= 1 + + def locks_held(self) -> int: + return sum(self._locks.values()) + + +@pytest.fixture +def device(monkeypatch: pytest.MonkeyPatch) -> Iterator[FakeDevice]: + fake = FakeDevice() + opened: list[FakeCounters] = [] + + def device_counters(torch_device: torch.device) -> FakeCounters: + opened.append(FakeCounters(fake)) + return opened[-1] + + # The fake allocator observes releases ahead of the probe, which registers when a scope opens. + monkeypatch.setattr(devices, "EMPTY_CACHE_OBSERVERS", [fake]) + monkeypatch.setattr(counters, "DeviceCounters", device_counters) + monkeypatch.setattr(probe, "SPILL_SETTLE_S", 0.0) + yield fake + assert all(c.closed for c in opened), "every window's counters are closed, so no sampler outlives it" + assert probe.armed is False + + +def _estimate(reservation: int = 2 * GIB) -> WorkingMemoryEstimate: + return WorkingMemoryEstimate( + estimator="vae.test", + version=3, + features={"operation": "decode", "height": 1024, "width": 1024, "fp32": False, "tile": None}, + terms=(WorkingMemoryTerm.fitted("convolutions", "k_decode", 1024 * 1024, {"k_decode": reservation / 2**20}),), + source="shipped:conv_backend=cudnn", + ) + + +def _model(cache: FakeCache, key: str, total_bytes: int = 1 * GIB, device: str = "meta") -> LoadedModelWithoutConfig: + return LoadedModelWithoutConfig( + CacheRecord(key=key, cached_model=FakeCachedModel(total_bytes, device)), cache=cache + ) + + +def _on_another_thread(model: LoadedModelWithoutConfig) -> None: + """Lock and unlock `model` on another thread, as Expand Prompt does next to a session.""" + + def use() -> None: + with model.model_on_device(): + pass + + thread = threading.Thread(target=use) + thread.start() + thread.join() + + +@contextmanager +def _scope() -> Iterator[ProbeScope]: + with calibration_scope(run_id="run-1", queue_item_id=7, node_id="l2i-3", node_type="l2i") as scope: + yield scope + + +def test_a_window_measures_what_the_operation_took_beyond_what_was_in_use(device: FakeDevice): + cache = FakeCache(device) + vae = _model(cache, "vae") + device.alloc(512 * MIB) # latents allocated before the lock: in use, not the operation's + device.alloc(2 * GIB) + device.free(2 * GIB) # an earlier node's cached blocks, released when the window opens + + with _scope() as scope: + with vae.model_on_device(working_memory=_estimate()): + device.alloc(3 * GIB) + device.free(3 * GIB) + + (window,) = scope.windows + assert window.demand_bytes == 3 * GIB + assert window.reserved is not None and window.reserved.delta == 3 * GIB + assert window.allocated is not None and window.allocated.delta == 3 * GIB + assert window.driver is not None and window.driver.delta == 3 * GIB + assert window.reserved.baseline == 512 * MIB + 1 * GIB, "the latents and the VAE's weights, nothing cached" + assert (window.valid, window.fit_eligible, window.reasons) == (True, True, []) + assert (window.run_id, window.queue_item_id, window.node_id, window.node_type) == ("run-1", 7, "l2i-3", "l2i") + assert (window.model_key, window.device, window.estimator, window.version) == ("vae", "meta", "vae.test", 3) + assert window.reservation_bytes == 2 * GIB and window.source == "shipped:conv_backend=cudnn" + assert window.resident_fraction == 1.0 and window.co_locked_models == [] + assert MeasuredWindow.model_validate_json(window.model_dump_json()) == window + + +@pytest.mark.parametrize(("reservation", "applied"), [(1 * GIB, 2 * GIB), (4 * GIB, 6 * GIB)]) +def test_a_calibration_lock_reserves_headroom_beyond_the_estimate(device: FakeDevice, reservation: int, applied: int): + """At least 1 GiB, or half the estimate: an estimate that is too low must not spill, or it could never be fixed.""" + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "vae").model_on_device(working_memory=_estimate(reservation)): + pass + + assert cache.reserved == [applied] + (window,) = scope.windows + assert (window.reservation_bytes, window.reservation_applied_bytes) == (reservation, applied) + + +def test_without_a_scope_the_lock_is_unchanged_and_no_counter_is_read(monkeypatch: pytest.MonkeyPatch): + def refuse(_device: torch.device) -> None: + raise AssertionError("a counter was read without a calibration scope") + + monkeypatch.setattr(counters, "DeviceCounters", refuse) + monkeypatch.setattr(counters, "read_counters", refuse) + cache = FakeCache(FakeDevice()) + vae = _model(cache, "vae") + + with vae.model_on_device(working_memory=_estimate(2 * GIB)): + pass + with vae.model_on_device(512 * MIB): + pass + with vae: + pass + + assert cache.reserved == [2 * GIB, 512 * MIB, None] + assert cache.locks_held() == 0 + + +def test_a_byte_count_or_a_cpu_model_never_opens_a_window(device: FakeDevice): + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "vae").model_on_device(512 * MIB): + pass + with _model(cache, "cpu-vae", device="cpu").model_on_device(working_memory=_estimate(2 * GIB)): + pass + + assert scope.windows == [] + assert cache.reserved == [512 * MIB, 2 * GIB], "no headroom where nothing is measured" + + +def test_a_co_lock_counts_its_weights_as_offset_and_its_lock_outside_the_peak(device: FakeDevice): + """A ControlNet locked inside the window moves 2 GiB of weights through a 3 GiB staging buffer; neither is the + operation's working memory.""" + cache = FakeCache(device, staging_bytes=3 * GIB) + transformer = _model(cache, "transformer") + controlnet = _model(cache, "controlnet", total_bytes=2 * GIB) + + with _scope() as scope: + with transformer.model_on_device(working_memory=_estimate()): + device.alloc(1 * GIB) # lives across the co-lock + with controlnet: + device.alloc(1536 * MIB) + device.free(1536 * MIB) + + (window,) = scope.windows + assert window.demand_bytes == 1 * GIB + 1536 * MIB + assert window.reserved is not None and window.reserved.offset == 2 * GIB + assert window.co_locked_models == ["controlnet"] + assert window.valid + + +def test_a_co_lock_that_offloads_more_than_it_loads_has_a_negative_offset(device: FakeDevice): + cache = FakeCache(device) + resident_before = _model(cache, "text-encoder", total_bytes=1 * GIB) + with resident_before: + pass # resident and unlocked: part of the baseline + vae = _model(cache, "vae") + adapter = _model(cache, "adapter", total_bytes=512 * MIB) + cache.offloads["adapter"] = [resident_before._cache_record] + + with _scope() as scope: + with vae.model_on_device(working_memory=_estimate()): + with adapter: + pass + device.alloc(2 * GIB) + device.free(2 * GIB) + + (window,) = scope.windows + assert window.reserved is not None and window.reserved.offset == -512 * MIB + assert window.demand_bytes == 2 * GIB + + +def test_a_nested_window_measures_alone_and_its_outputs_stay_in_the_parent(device: FakeDevice): + cache = FakeCache(device) + transformer = _model(cache, "transformer", total_bytes=4 * GIB) + vae = _model(cache, "vae", total_bytes=1 * GIB) + + with _scope() as scope: + with transformer.model_on_device(working_memory=_estimate(8 * GIB)): + device.alloc(1 * GIB) + with vae.model_on_device(working_memory=_estimate(2 * GIB)): + device.alloc(2 * GIB) + device.free(2 * GIB) + TorchDevice.empty_cache() # a boundary of the child alone: the parent is frozen + device.alloc(1 * GIB) + device.free(1 * GIB) # still cached when the child closes; the parent releases it + device.alloc(256 * MIB) # the encoded reference latents, kept + device.alloc(512 * MIB) + device.free(512 * MIB) + + child, parent = scope.windows + assert (child.model_key, parent.model_key) == ("vae", "transformer") + assert child.demand_bytes == 2 * GIB and child.segments == 2 + assert child.co_locked_models == ["transformer"] + assert parent.demand_bytes == 1 * GIB + 256 * MIB + 512 * MIB, "not the child's 2 GiB transient" + assert parent.reserved is not None and parent.reserved.offset == 1 * GIB, "the VAE's weights" + assert parent.co_locked_models == ["vae"] + assert child.valid and parent.valid + + +def test_an_empty_cache_inside_a_window_ends_a_segment_so_the_driver_peak_is_kept(device: FakeDevice): + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "vae").model_on_device(working_memory=_estimate()): + device.alloc(3 * GIB) + device.free(3 * GIB) + TorchDevice.empty_cache() + device.alloc(1 * GIB) + device.free(1 * GIB) + + (window,) = scope.windows + assert window.driver is not None and window.driver.delta == 3 * GIB + assert window.reserved is not None and window.reserved.delta == 3 * GIB + assert window.segments == 2 + + +def test_an_allocated_peak_above_the_reserved_peak_decides_the_demand(device: FakeDevice): + """cudaMallocAsync: `max_memory_reserved` measured 128 MiB under an allocated peak of 136 MiB.""" + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "vae").model_on_device(working_memory=_estimate()): + device.alloc(136 * MIB, reserve=128 * MIB) + device.free(136 * MIB) + + (window,) = scope.windows + assert window.reserved is not None and window.reserved.delta == 128 * MIB + assert window.demand_bytes == 136 * MIB + + +def test_a_driver_that_keeps_a_retained_pool_reads_below_reserved_and_reserved_decides(device: FakeDevice): + """HIP under Windows keeps released memory counted and hands it to the next allocation.""" + device.retained_pool = True + device.alloc(3 * GIB) + device.free(3 * GIB) + TorchDevice.empty_cache() # 3 GiB retained; loading the VAE's weights takes 1 GiB of it + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "vae").model_on_device(working_memory=_estimate()): + device.alloc(6 * GIB) + device.free(6 * GIB) + + (window,) = scope.windows + assert window.driver is not None and window.driver.delta == 4 * GIB + assert window.demand_bytes == 6 * GIB + + +def test_without_a_driver_reading_the_window_rests_on_torch_and_is_unverified(device: FakeDevice): + device.has_driver = False + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "vae", device="cuda:0").model_on_device(working_memory=_estimate()): + device.alloc(1 * GIB) + device.free(1 * GIB) + + (window,) = scope.windows + assert window.driver is None and window.demand_bytes == 1 * GIB + assert (window.method, window.verified, window.valid) == ("torch-only", False, True) + + +def test_a_driver_reading_by_a_verified_method_is_marked_verified(device: FakeDevice): + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "vae", device="cuda:0").model_on_device(working_memory=_estimate()): + pass + + (window,) = scope.windows + assert (window.method, window.verified) == ("windows-d3dkmt", True) + + +def test_a_lock_by_another_thread_on_the_same_cache_voids_the_window(device: FakeDevice): + cache = FakeCache(device) + vae = _model(cache, "vae") + prompt_expander = _model(cache, "llm") + elsewhere = _model(FakeCache(device), "other-gpu-model") + + with _scope() as scope: + with vae.model_on_device(working_memory=_estimate()): + _on_another_thread(elsewhere) + with vae.model_on_device(working_memory=_estimate()): + _on_another_thread(prompt_expander) + + other_device, same_device = scope.windows + assert other_device.valid, "another device's cache does not share this one" + assert (same_device.valid, same_device.reasons) == (False, ["foreign_lock"]) + + +def test_a_lock_another_thread_holds_when_the_window_opens_voids_it(device: FakeDevice): + cache = FakeCache(device) + vae = _model(cache, "vae") + prompt_expander = _model(cache, "llm") + holding, done = threading.Event(), threading.Event() + + def hold() -> None: + with prompt_expander.model_on_device(): + holding.set() + done.wait() + + thread = threading.Thread(target=hold) + thread.start() + holding.wait() + try: + with _scope() as scope: + with vae.model_on_device(working_memory=_estimate()): + pass + finally: + done.set() + thread.join() + + (window,) = scope.windows + assert (window.valid, window.reasons) == (False, ["foreign_lock"]) + + +def test_a_lock_another_thread_takes_while_the_window_opens_voids_it( + device: FakeDevice, monkeypatch: pytest.MonkeyPatch +): + """Between the cache's lock count and the baseline lie a synchronization and a release; a lock another thread + takes and settles there must not go unseen.""" + cache = FakeCache(device) + vae = _model(cache, "vae") + prompt_expander = _model(cache, "llm") + device_counters = counters.DeviceCounters + + def opening(torch_device: torch.device) -> FakeCounters: + _on_another_thread(prompt_expander) + return device_counters(torch_device) + + monkeypatch.setattr(counters, "DeviceCounters", opening) + + with _scope() as scope: + with vae.model_on_device(working_memory=_estimate()): + pass + + (window,) = scope.windows + assert (window.valid, window.reasons) == (False, ["foreign_lock"]) + + +def test_a_model_this_thread_pins_in_ram_is_its_own_lock(device: FakeDevice): + """LoRAs are pinned in RAM before the transformer locks; that pin belongs to this thread, not to another.""" + cache = FakeCache(device) + lora = _model(cache, "lora") + + with _scope() as scope: + with lora.model_in_ram(): + with _model(cache, "transformer").model_on_device(working_memory=_estimate()): + pass + + (window,) = scope.windows + assert (window.valid, window.reasons) == (True, []) + assert window.co_locked_models == ["lora"] + + +def test_after_an_unlock_a_lock_another_thread_holds_still_voids_the_next_window(device: FakeDevice): + cache = FakeCache(device) + vae = _model(cache, "vae") + prompt_expander = _model(cache, "llm") + holding, done = threading.Event(), threading.Event() + + def hold() -> None: + with prompt_expander.model_on_device(): + holding.set() + done.wait() + + thread = threading.Thread(target=hold) + with _scope() as scope: + with vae.model_on_device(working_memory=_estimate()): + pass + thread.start() + holding.wait() + try: + with vae.model_on_device(working_memory=_estimate()): + pass + finally: + done.set() + thread.join() + + first, second = scope.windows + assert first.valid + assert (second.valid, second.reasons) == (False, ["foreign_lock"]) + + +@pytest.mark.parametrize("release", [1, 2], ids=["co-lock-start", "co-lock-settled"]) +def test_a_release_that_fails_around_a_co_lock_voids_the_window_and_the_rest_of_the_scope( + device: FakeDevice, release: int +): + cache = FakeCache(device) + transformer = _model(cache, "transformer") + + with _scope() as scope: + with transformer.model_on_device(working_memory=_estimate()): + device.fail_on_release = device.releases + release + with _model(cache, "controlnet"): + pass + device.alloc(1 * GIB) + device.free(1 * GIB) + with transformer.model_on_device(working_memory=_estimate()): + pass + + failed, later = scope.windows + assert (failed.valid, failed.reasons) == (False, ["probe_error"]) + assert (later.valid, later.reasons) == (False, ["probe_error"]), "the scope's lock bookkeeping is suspect" + assert cache.locks_held() == 0 + + +def test_a_segment_boundary_that_fails_voids_the_window(device: FakeDevice): + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "vae").model_on_device(working_memory=_estimate()): + device.fail_next_read = True + TorchDevice.empty_cache() + + (window,) = scope.windows + assert (window.valid, window.reasons) == (False, ["probe_error"]) + + +def test_a_child_window_that_fails_to_close_is_recorded_invalid_and_the_parent_resumes(device: FakeDevice): + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "transformer").model_on_device(working_memory=_estimate()): + with _model(cache, "vae").model_on_device(working_memory=_estimate()): + device.fail_next_read = True + device.alloc(1 * GIB) + device.free(1 * GIB) + + child, parent = scope.windows + assert (child.model_key, child.valid) == ("vae", False) + assert child.reasons == ["probe_error", "counters_unavailable"], "its final segment was never read" + assert (parent.model_key, parent.valid, parent.reasons) == ("transformer", False, ["probe_error"]) + assert parent.segments == 3, "before the child's lock, up to the child, and after it again" + assert cache.locks_held() == 0 + + +@pytest.mark.parametrize(("counter", "reason"), [("num_ooms", "oom"), ("num_alloc_retries", "alloc_retry")]) +def test_an_allocator_that_failed_or_retried_voids_the_window(device: FakeDevice, counter: str, reason: str): + """A retry releases the cache and squeezes the peak below what the operation takes with room to spare.""" + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "vae").model_on_device(working_memory=_estimate()): + setattr(device, counter, getattr(device, counter) + 1) + + (window,) = scope.windows + assert (window.valid, window.fit_eligible, window.reasons) == (False, False, [reason]) + + +@pytest.mark.parametrize(("paged_at_close", "valid"), [([600 * MIB, 9 * MIB], True), ([600 * MIB, 600 * MIB], False)]) +def test_only_spill_that_persists_after_close_voids_the_window(device: FakeDevice, paged_at_close, valid: bool): + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "vae").model_on_device(working_memory=_estimate()): + device.paged = list(paged_at_close) + + (window,) = scope.windows + assert window.valid is valid + assert window.reasons == ([] if valid else ["spill"]) + + +def test_a_device_wide_counter_growing_beyond_this_process_voids_the_window(device: FakeDevice): + """XPU: Sysman counts the whole device, so growth beyond this process's own is another process's load.""" + device.method, device.has_driver, device.device_used_by_others = "xpu-sysman", False, 0 + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "vae").model_on_device(working_memory=_estimate()): + device.alloc(1 * GIB) + device.device_used_by_others = 2 * GIB + + (window,) = scope.windows + assert (window.valid, window.reasons) == (False, ["foreign_device_load"]) + assert window.demand_bytes == 1 * GIB, "the device-wide counter is a guard, never demand" + + +@pytest.mark.parametrize( + ("error", "reason"), [(RuntimeError("decode failed"), "exception"), (CanceledException(), "canceled")] +) +def test_a_window_closed_by_an_error_is_invalid_and_the_error_propagates( + device: FakeDevice, error: BaseException, reason: str +): + cache = FakeCache(device) + + with _scope() as scope: + with pytest.raises(type(error)): + with _model(cache, "vae").model_on_device(working_memory=_estimate()): + device.alloc(1 * GIB) + raise error + + (window,) = scope.windows + assert (window.valid, window.reasons) == (False, [reason]) + assert cache.locks_held() == 0 + + +def test_a_partly_resident_model_is_measured_but_not_fit_eligible(device: FakeDevice): + """Streamed weights are working memory the estimate does not price.""" + cache = FakeCache(device) + cache.partial["transformer"] = 3 * GIB + + with _scope() as scope: + with _model(cache, "transformer", total_bytes=4 * GIB).model_on_device(working_memory=_estimate()): + pass + + (window,) = scope.windows + assert (window.valid, window.fit_eligible, window.reasons) == (True, False, ["partial_residency"]) + assert window.resident_fraction == 0.75 + + +def test_the_krea2_sdpa_benchmark_voids_the_window(device: FakeDevice, monkeypatch: pytest.MonkeyPatch): + """It resets the device's peak statistics itself, under the window's feet.""" + monkeypatch.setenv("INVOKE_KREA2_SDPA_BACKEND", "cudnn") + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "vae", device="cuda:0").model_on_device(working_memory=_estimate()): + pass + + (window,) = scope.windows + assert (window.valid, window.reasons) == (False, ["krea2_sdpa_benchmark"]) + + +def test_a_device_without_readable_counters_gives_an_invalid_window_and_no_demand(device: FakeDevice): + device.readable = False + cache = FakeCache(device) + + with _scope() as scope: + with _model(cache, "vae").model_on_device(working_memory=_estimate()): + pass + + (window,) = scope.windows + assert (window.valid, window.demand_bytes, window.reasons) == (False, None, ["counters_unavailable"]) + + +def test_calibration_scopes_do_not_nest(device: FakeDevice): + with _scope(): + with pytest.raises(RuntimeError, match="do not nest"): + with _scope(): + pass + + +requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA or ROCm GPU") + + +@pytest.mark.slow +@requires_gpu +def test_on_hardware_a_vae_decode_window_is_valid_and_covers_the_reserved_growth(): + from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL + + from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache + from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_sd15_sdxl + + torch.manual_seed(0) + gpu = torch.device("cuda", 0) + vae = AutoencoderKL( + down_block_types=("DownEncoderBlock2D",) * 4, + up_block_types=("UpDecoderBlock2D",) * 4, + block_out_channels=(32, 64, 128, 128), + layers_per_block=1, + latent_channels=4, + ).to(torch.float16) + cache = ModelCache( + execution_device_working_mem_gb=1.0, + enable_partial_loading=True, + keep_ram_copy_of_weights=True, + execution_device=gpu, + storage_device="cpu", + shared_cpu_weights=None, + ) + cache.put("vae", vae) + loaded = LoadedModelWithoutConfig(cache.get("vae"), cache) + latents = torch.randn(1, 4, 96, 96, dtype=torch.float16) + estimate = estimate_vae_working_memory_sd15_sdxl( + operation="decode", image_tensor=latents, vae=vae, tile_size=None, fp32=False, device=gpu + ) + + with calibration_scope(run_id="hw", queue_item_id=1, node_id="l2i", node_type="l2i") as scope: + with loaded.model_on_device(working_memory=estimate) as (_, model): + torch.cuda.synchronize(gpu) + reserved_at_open = torch.cuda.memory_reserved(gpu) + with torch.inference_mode(): + model.decode(latents.to(gpu)) + torch.cuda.synchronize(gpu) + reserved_growth = torch.cuda.max_memory_reserved(gpu) - reserved_at_open + + (window,) = scope.windows + assert window.valid, window.reasons + assert reserved_growth > 0 + assert window.demand_bytes is not None and window.demand_bytes >= reserved_growth + assert window.reserved is not None and window.reserved.delta == reserved_growth + if sys.platform == "win32": + assert (window.method, window.verified) == ("windows-d3dkmt", True) + assert window.driver is not None diff --git a/tests/fixtures/working_memory_calibration.py b/tests/fixtures/working_memory_calibration.py new file mode 100644 index 00000000000..0f4c5d1d317 --- /dev/null +++ b/tests/fixtures/working_memory_calibration.py @@ -0,0 +1,86 @@ +"""A working-memory calibration service on a temporary root, with a faked GPU. + +A device class and an environment can only be read from real hardware, so the service gets fixed ones: one class for +``cuda:0``, an environment whose hard facts are the settings passed in. Windows are built the way the probe reports +them, without measuring anything. +""" + +from collections.abc import Mapping +from pathlib import Path +from typing import Callable, Optional + +import torch + +from invokeai.app.services.working_memory_calibration.working_memory_calibration_default import ( + SettingValue, + WorkingMemoryCalibrationService, +) +from invokeai.app.services.working_memory_calibration.working_memory_calibration_store import CalibrationStore +from invokeai.backend.working_memory.estimate import WorkingMemoryTerm +from invokeai.backend.working_memory.fingerprint import DeviceClass, Environment +from invokeai.backend.working_memory.probe import MeasuredWindow, Reason + +GPU = torch.device("cuda:0") +GPU_CLASS = DeviceClass(vendor="nvidia", name="Test GPU 24GB", total_mib=24564, backend="cuda", arch="sm_89") +OTHER_CLASS = DeviceClass(vendor="amd", name="Test GPU 16GB", total_mib=16304, backend="hip", arch="gfx1200") + + +def fake_environment(device: torch.device, settings: Mapping[str, SettingValue]) -> Environment: + hard = {"torch": "2.7", **{f"setting.{key}": str(value) for key, value in settings.items()}} + return Environment(hard=hard, soft={"driver": "1.0", "os": "test"}) + + +def calibration_service( + root: Path, + devices: Optional[Callable[[], list[torch.device]]] = None, + device_class: DeviceClass = GPU_CLASS, +) -> WorkingMemoryCalibrationService: + """A service storing under `root`; not started (call `start(invoker)`).""" + return WorkingMemoryCalibrationService( + CalibrationStore(root / "working_memory"), + generation_devices=devices or (lambda: [GPU]), + device_class_for=lambda device: device_class, + environment_for=fake_environment, + ) + + +def measured_window( + run_id: str, + queue_item_id: int = 1, + node_id: str = "decode", + *, + model_key: str = "vae-key", + estimator: str = "vae_decode", + version: int = 1, + demand_bytes: int = 3 * 2**30, + reasons: tuple[Reason, ...] = (), +) -> MeasuredWindow: + valid = not any(reason != "partial_residency" for reason in reasons) + return MeasuredWindow( + run_id=run_id, + queue_item_id=queue_item_id, + node_id=node_id, + node_type="test_calibration_node", + model_key=model_key, + device=str(GPU), + estimator=estimator, + version=version, + features={"pixels": 1024 * 1024}, + terms=(WorkingMemoryTerm(name="activations", bytes=2.5 * 2**30, included=True),), + reservation_bytes=int(2.5 * 2**30), + source="data", + reservation_applied_bytes=int(3.75 * 2**30), + driver=None, + reserved=None, + allocated=None, + demand_bytes=demand_bytes, + method="torch-only", + verified=False, + resident_fraction=1.0, + co_locked_models=[], + segments=1, + duration_s=0.5, + valid=valid, + fit_eligible=valid and "partial_residency" not in reasons, + reasons=list(reasons), + ) From fab1519974ec072291d644972976037d47bc4038 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 17 Oct 2026 18:37:02 +0200 Subject: [PATCH 3/9] fix(working-memory): record the operating system as a hard calibration 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. --- .../Architecture/model-manager.mdx | 2 +- .../backend/working_memory/fingerprint.py | 23 ++++++++++++----- .../working_memory/test_fingerprint.py | 25 +++++++++++++++++-- 3 files changed, 41 insertions(+), 9 deletions(-) diff --git a/docs/src/content/docs/development/Architecture/model-manager.mdx b/docs/src/content/docs/development/Architecture/model-manager.mdx index 1b8da502c1a..4e839daf307 100644 --- a/docs/src/content/docs/development/Architecture/model-manager.mdx +++ b/docs/src/content/docs/development/Architecture/model-manager.mdx @@ -1004,7 +1004,7 @@ Inside a scope, every `model_on_device(working_memory=...)` on a GPU is a window - Invalid (`valid=False`, with `reasons`): another thread locked a model on the same device's cache during the window or held one when it opened; 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; 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`, which also voids the scope's later windows when it hit the lock bookkeeping); 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 (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). +`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 diff --git a/invokeai/backend/working_memory/fingerprint.py b/invokeai/backend/working_memory/fingerprint.py index 5f8f057ec7a..ccdde0a04c1 100644 --- a/invokeai/backend/working_memory/fingerprint.py +++ b/invokeai/backend/working_memory/fingerprint.py @@ -3,9 +3,10 @@ A device class names a kind of card, not one card: identical cards share it, whatever their index or PCI slot, so measurements taken on one apply to the other. An environment splits into hard facts, any change of which can change an -operation's memory (torch and runtime versions, the convolution library, the allocator and its configuration, the -attention path, InvokeAI's memory settings), and soft facts recorded for diagnosis only (full versions, the driver, the -OS build). Both are produced deterministically here; which change makes a measurement stale is the caller's decision. +operation's memory (the operating system, torch and runtime versions, the convolution library, the allocator and its +configuration, the attention path, InvokeAI's memory settings), and soft facts recorded for diagnosis only (full +versions, the driver, the OS build). Both are produced deterministically here; which change makes a measurement stale +is the caller's decision. """ import ctypes @@ -14,8 +15,9 @@ import json import os import platform +import sys from collections.abc import Mapping -from typing import Final, Optional +from typing import Final, Literal, Optional import psutil import torch @@ -135,11 +137,20 @@ def _conv_library_major() -> Optional[str]: return None +def os_family(platform_name: str) -> Literal["windows", "linux", "other"]: + """The operating system family of a ``sys.platform`` value, as a hard fact: the driver model differs (WDDM pages + an allocation that does not fit into shared memory instead of failing it), so the same card and torch build can + take different memory on Windows and on Linux.""" + if platform_name == "win32": + return "windows" + return "linux" if platform_name.startswith("linux") else "other" + + def environment_for(device: torch.device, settings: Mapping[str, str | int | float | bool | None]) -> Environment: """The environment an operation on `device` runs under. Every key of `settings` is a hard fact, as ``setting.``: the caller passes the InvokeAI settings that change memory.""" backend = backend_for(device) - hard: dict[str, str] = {"torch": _major_minor(torch.__version__)} + hard: dict[str, str] = {"os": os_family(sys.platform), "torch": _major_minor(torch.__version__)} if backend in ("cuda", "hip"): runtime = torch.version.hip if backend == "hip" else torch.version.cuda hard[backend] = _major_minor(runtime or "") @@ -156,7 +167,7 @@ def environment_for(device: torch.device, settings: Mapping[str, str | int | flo soft: dict[str, Optional[str]] = { "torch_build": torch.__version__, "driver": nvml.driver_version() if backend == "cuda" else None, - "os": platform.platform(), + "os_build": platform.platform(), "invokeai": __version__, } return Environment(hard=hard, soft=soft) diff --git a/tests/backend/working_memory/test_fingerprint.py b/tests/backend/working_memory/test_fingerprint.py index d348055545b..5956254c286 100644 --- a/tests/backend/working_memory/test_fingerprint.py +++ b/tests/backend/working_memory/test_fingerprint.py @@ -1,13 +1,14 @@ """Device classes and environments: what a measurement is keyed by, from scripted device properties.""" import platform +import sys from types import SimpleNamespace import pytest import torch from invokeai.backend.util import nvml -from invokeai.backend.working_memory.fingerprint import DeviceClass, device_class_for, environment_for +from invokeai.backend.working_memory.fingerprint import DeviceClass, device_class_for, environment_for, os_family GIB = 2**30 CUDA0 = torch.device("cuda", 0) @@ -100,7 +101,8 @@ def test_settings_are_hard_facts_and_full_versions_soft_ones(cuda_build): "env.INVOKEAI_ROCM_CONV3D", } <= set(environment.hard) assert environment.hard["torch"].count(".") == 1, "major.minor only" - assert set(environment.soft) == {"torch_build", "driver", "os", "invokeai"} + assert set(environment.soft) == {"torch_build", "driver", "os_build", "invokeai"} + assert environment.soft["os_build"] == platform.platform() assert environment.soft["driver"] == "581.57" @@ -117,3 +119,22 @@ def test_only_hard_facts_change_the_environment_id(cuda_build, monkeypatch: pyte changed = environment_for(CUDA0, _SETTINGS) assert changed.hard["env.PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True" assert changed.id != reference + + +@pytest.mark.parametrize( + ("platform_name", "family"), + [("win32", "windows"), ("linux", "linux"), ("linux2", "linux"), ("darwin", "other"), ("freebsd14", "other")], +) +def test_the_os_family_is_named_from_sys_platform(platform_name: str, family: str): + assert os_family(platform_name) == family + + +def test_the_operating_system_is_a_hard_fact(cuda_build, monkeypatch: pytest.MonkeyPatch): + ids = {} + for platform_name, family in (("win32", "windows"), ("linux", "linux"), ("darwin", "other")): + monkeypatch.setattr(sys, "platform", platform_name) + environment = environment_for(CUDA0, _SETTINGS) + assert environment.hard["os"] == family + ids[family] = environment.id + + assert len(set(ids.values())) == 3, "the same card and software on another OS is another environment" From f004b1ade8355cd19e65d3cee7e8c911975b6800 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Mon, 19 Oct 2026 20:12:09 +0200 Subject: [PATCH 4/9] test(working-memory): script the allocator backend in the fingerprint tests CPU wheels on macOS and Windows have no CUDA allocator binding, so the fixture that scripts a CUDA build scripts its backend too. --- tests/backend/working_memory/test_fingerprint.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/backend/working_memory/test_fingerprint.py b/tests/backend/working_memory/test_fingerprint.py index 5956254c286..84be9d6b1f9 100644 --- a/tests/backend/working_memory/test_fingerprint.py +++ b/tests/backend/working_memory/test_fingerprint.py @@ -75,6 +75,8 @@ def test_a_cpu_has_no_device_class(): @pytest.fixture def cuda_build(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(torch.version, "hip", None) + # A CPU wheel (macOS, Windows) has no allocator binding to ask. + monkeypatch.setattr(torch.cuda, "get_allocator_backend", lambda: "native") monkeypatch.setattr(nvml, "driver_version", lambda: "581.57") monkeypatch.delenv("PYTORCH_CUDA_ALLOC_CONF", raising=False) From 3ebdb7f8a347037209136c8f3e54ae676ec14040 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 21 Oct 2026 00:31:01 +0200 Subject: [PATCH 5/9] fix(working-memory): keep calibration runs and their measurements trustworthy The session processor starts after every other service; windows of failed or canceled nodes, failed counter resets and locks racing a window are stored invalid. Keyed requests are identified by model, graph and batch data, are created once per run and admit all of their items or none; the reserved-origin 400 is declared. --- .../Architecture/model-manager.mdx | 18 +- .../development/Architecture/overview.mdx | 2 + invokeai/app/api/dependencies.py | 4 + invokeai/app/api/routers/session_queue.py | 4 + invokeai/app/api/routers/working_memory.py | 32 +- invokeai/app/services/invoker.py | 23 +- .../session_processor_default.py | 60 ++- .../session_queue/session_queue_base.py | 6 +- .../session_queue/session_queue_default.py | 14 +- .../working_memory_calibration_base.py | 5 +- .../working_memory_calibration_common.py | 18 +- .../working_memory_calibration_default.py | 288 ++++++++++---- .../backend/model_manager/load/load_base.py | 12 +- invokeai/backend/util/device_pool.py | 5 +- invokeai/backend/working_memory/counters.py | 35 +- invokeai/backend/working_memory/probe.py | 123 ++++-- invokeai/frontend/api/openapi.json | 19 +- invokeai/frontend/api/schema.ts | 37 +- .../test_working_memory_calibration.py | 53 ++- .../test_working_memory_calibration_runner.py | 32 +- .../test_session_queue_writes.py | 27 ++ .../test_image_move_startup_safety.py | 65 ---- tests/app/services/test_invoker.py | 108 ++++++ ...test_working_memory_calibration_service.py | 364 +++++++++++++++++- tests/app/test_api_dependencies.py | 68 +++- tests/backend/working_memory/test_counters.py | 33 +- tests/backend/working_memory/test_probe.py | 205 +++++++++- tests/fixtures/working_memory_calibration.py | 23 +- 28 files changed, 1372 insertions(+), 311 deletions(-) create mode 100644 tests/app/services/test_invoker.py diff --git a/docs/src/content/docs/development/Architecture/model-manager.mdx b/docs/src/content/docs/development/Architecture/model-manager.mdx index 7520f263aa2..1cf3041c696 100644 --- a/docs/src/content/docs/development/Architecture/model-manager.mdx +++ b/docs/src/content/docs/development/Architecture/model-manager.mdx @@ -1029,7 +1029,7 @@ Inside a scope, every `model_on_device(working_memory=...)` on a GPU is a window - 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 or held one when it opened; 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; 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`, which also voids the scope's later windows when it hit the lock bookkeeping); or no counter could be read. +- 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). @@ -1038,15 +1038,19 @@ Inside a scope, every `model_on_device(working_memory=...)` on a GPU is a window 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`, and sets the origin to `invokeai:wm-calibration:` and the destination to `invokeai:wm-calibration`. It records the run, then enqueues the batch as the requesting admin. - - With a key, the run id is derived from the owner and the key, so a repeated request rebuilds the identical batch. The queue's enqueue receipt then returns the same items, and the response is the same run; nothing new is enqueued. Reusing the key for a different request is a 409 and leaves the original run in place. Acknowledge the key through `POST /api/v1/queue/{queue_id}/enqueue_batch/acknowledge`. +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:` 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. 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. 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. Retries and children are checked here too. Otherwise the item fails with the reason. The checks: +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. Windows from a device of another class, or of a run that was removed, are dropped with a warning. If they cannot be saved, the item fails and the session stops. Every other item gets no scope and runs as before. +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. 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 `/working_memory/.json` (`CalibrationStoreFile`, `schema_version` 1). It is machine data, so it lives next to the database rather than inside it. It holds: @@ -1069,13 +1073,13 @@ 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 run is `running` while any item is pending or running, `failed` when any failed, and `completed` when all completed. Otherwise it is `partial` if points were recorded, or `canceled` if not. +- 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; 409 no GPU or several generation devices, unusable store file, full queue (the run is discarded), or a retry key used for another request; 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 | +| `POST .../i/{key}/working_memory/calibrate` | starts a run (above) | 404 unknown model, or the run was removed while it was being created; 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 | 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 | diff --git a/docs/src/content/docs/development/Architecture/overview.mdx b/docs/src/content/docs/development/Architecture/overview.mdx index 8cbe18efad5..0bb0b3bf812 100644 --- a/docs/src/content/docs/development/Architecture/overview.mdx +++ b/docs/src/content/docs/development/Architecture/overview.mdx @@ -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()` signals the session processor first, so that its workers stop taking items and cancel their sessions before the other services stop; it does not wait for the workers. + ### 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. diff --git a/invokeai/app/api/dependencies.py b/invokeai/app/api/dependencies.py index 98d63dbc7c2..84cb7c3354a 100644 --- a/invokeai/app/api/dependencies.py +++ b/invokeai/app/api/dependencies.py @@ -284,6 +284,10 @@ def initialize( working_memory_calibration=working_memory_calibration, ) + # Before the Invoker starts any service: work outside the session queue, such as the image index's backfill, + # borrows idle GPUs from the generation device pool as soon as its service starts. + session_processor.register_generation_devices(configuration) + # Constructing the Invoker starts every service, including the session # processor thread (which may immediately resume queue items that were # pending at shutdown). Trigger torch attention's lazy, non-thread-safe diff --git a/invokeai/app/api/routers/session_queue.py b/invokeai/app/api/routers/session_queue.py index 55ade472e6e..25a25c4d8ef 100644 --- a/invokeai/app/api/routers/session_queue.py +++ b/invokeai/app/api/routers/session_queue.py @@ -260,6 +260,10 @@ def get_queue_item_for_mutation(queue_id: str, item_id: int, current_user: Curre operation_id="enqueue_batch", responses={ 201: {"model": EnqueueBatchResult}, + 400: { + "description": f"The batch's origin starts with '{CALIBRATION_ORIGIN_PREFIX}' (in any case), which is " + "reserved for working-memory calibration runs" + }, }, ) async def enqueue_batch( diff --git a/invokeai/app/api/routers/working_memory.py b/invokeai/app/api/routers/working_memory.py index afed1028611..10517161298 100644 --- a/invokeai/app/api/routers/working_memory.py +++ b/invokeai/app/api/routers/working_memory.py @@ -88,15 +88,22 @@ def _http_error(error: Exception) -> HTTPException: status_code=201, response_model=WorkingMemoryCalibrationRunCreated, responses={ - 404: {"description": "The model could not be found"}, + 404: { + "description": "The model could not be found, or the run was removed while it was being created; its " + "items are queued then and fail when they start" + }, 409: { "description": "This machine cannot calibrate now (no GPU, or several generation devices), the device " - "class's store file is unusable and kept as it is, the queue is full, or the idempotency key was used for " - "another request" + "class's store file is unusable and kept as it is (possibly after the items were queued, which then fail " + "when they start), the queue has no room for all of the run's items (it then takes none), or the " + "idempotency key was used for another request" }, 422: {"description": "The graph cannot be measured as a calibration run"}, 429: {"description": "Too many unacknowledged idempotency keys; acknowledge or wait, then retry"}, - 503: _STORE_UNREADABLE, + 503: { + "description": "The calibration store cannot be read or written (unwritable root, lock timeout); with an " + "idempotency key, possibly after the items were queued" + }, }, ) async def calibrate_model_working_memory( @@ -105,15 +112,24 @@ async def calibrate_model_working_memory( request: WorkingMemoryCalibrationRequest = Body(description="The batch to measure"), ) -> WorkingMemoryCalibrationRunCreated: """Start a working-memory calibration run: enqueue the batch with every node uncached and measure each node's - model operations on the generation device. Needs exactly one generation device, and it must be a GPU. - - With an `idempotency_key`, repeating the request (after a lost response) returns the run it created and enqueues - nothing new; acknowledge the key through the queue's enqueue_batch/acknowledge route once the response arrived.""" + model operations on the generation device. Needs exactly one generation device, and it must be a GPU. The queue + takes all of the batch's items or none. + + Once the items are queued, a request without an `idempotency_key` succeeds even when the store fails for a passing + reason (an unwritable root, a lock timeout) to record them: they run all the same, and the status finds them by + their origin. With a key, that failure is a 503 although the items are queued. A run removed meanwhile (404) or an + unusable store file (409) fails the request with or without a key; its queued items then fail when they start. + + With a key, repeating the request (also after a lost response) records the queued items and returns the run; it + enqueues nothing new. The items come from the queue's enqueue receipt for the key while it exists, otherwise from + the queue items that carry the run's origin and those the runner saw. Acknowledge the key through the queue's + enqueue_batch/acknowledge route once the response arrived.""" await asyncio.to_thread(assert_image_move_maintenance_inactive) try: return await _service().create_run(key, request, current_admin.user_id) except ( UnknownModelException, + CalibrationRunNotFoundError, CalibrationGraphError, CalibrationUnavailableError, CalibrationStoreError, diff --git a/invokeai/app/services/invoker.py b/invokeai/app/services/invoker.py index b504f511ba7..17594fb266e 100644 --- a/invokeai/app/services/invoker.py +++ b/invokeai/app/services/invoker.py @@ -22,13 +22,24 @@ def __stop_service(self, service) -> None: if callable(stop_op): stop_op(self) + def _services_but_the_session_processor(self) -> list[object]: + return [getattr(self.services, name) for name in vars(self.services) if name != "session_processor"] + def _start(self) -> None: - """Starts the invoker. This is called automatically when the invoker is created.""" - for service in vars(self.services): - self.__start_service(getattr(self.services, service)) + """Starts the invoker. This is called automatically when the invoker is created. + + The session processor starts after every other service: its workers dequeue at once, and an item left pending + before the restart (the queue's start cancels those that were in progress) may reach any service, which must + have been started by then. + """ + for service in self._services_but_the_session_processor(): + self.__start_service(service) + self.__start_service(self.services.session_processor) def stop(self) -> None: """Stops the invoker. A new invoker will have to be created to execute further.""" - # First stop all services - for service in vars(self.services): - self.__stop_service(getattr(self.services, service)) + # The session processor is signalled first, so that its workers stop taking queue items and cancel the sessions + # they run before the services those sessions use stop. Its stop() does not wait for the workers. + self.__stop_service(self.services.session_processor) + for service in self._services_but_the_session_processor(): + self.__stop_service(service) diff --git a/invokeai/app/services/session_processor/session_processor_default.py b/invokeai/app/services/session_processor/session_processor_default.py index 920b16994a2..b8d0328e4ba 100644 --- a/invokeai/app/services/session_processor/session_processor_default.py +++ b/invokeai/app/services/session_processor/session_processor_default.py @@ -4,13 +4,14 @@ from contextlib import contextmanager, nullcontext, suppress from threading import BoundedSemaphore, Thread from threading import Event as ThreadEvent -from typing import Iterator, Optional +from typing import Iterator, Literal, Optional import torch from starlette.concurrency import run_in_threadpool from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput from invokeai.app.invocations.call_saved_workflow import CallSavedWorkflowInvocation +from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.events.events_common import ( BatchEnqueuedEvent, FastAPIEvent, @@ -169,6 +170,9 @@ def __init__( self._architecture_switch = ArchitectureSwitchOffload() # The calibration run of the item this runner is executing, None for every other item. self._calibration_run_id: Optional[str] = None + # What the outcome hooks last reported: `run_node` swallows a node's failure or cancellation, and a calibration + # node must not keep valid measurements then. + self._node_outcome: Optional[Literal["completed", "failed"]] = None def start(self, services: InvocationServices, cancel_event: ThreadEvent, profiler: Optional[Profiler] = None): self._services = services @@ -310,7 +314,8 @@ def _run_calibration_item(self, queue_item: SessionQueueItem, run_id: str) -> No def _run_calibration_node(self, invocation: BaseInvocation, queue_item: SessionQueueItem, run_id: str) -> bool: """Run one node of a calibration item inside a probe scope and persist its windows, whether the node completed, - failed or was canceled. Returns False when they could not be persisted: the item has failed then.""" + failed or was canceled; unless it completed, they are persisted invalid. Returns False when they could not be + persisted: the item has failed then.""" with probe.calibration_scope( run_id=run_id, queue_item_id=queue_item.item_id, @@ -318,7 +323,11 @@ def _run_calibration_node(self, invocation: BaseInvocation, queue_item: SessionQ node_id=queue_item.session.prepared_source_mapping.get(invocation.id, invocation.id), node_type=invocation.get_type(), ) as scope: + self._node_outcome = None self.run_node(invocation, queue_item) + if self._node_outcome != "completed": + # A node that neither completed nor failed was canceled: `run_node` passes over a CanceledException. + scope.invalidate("exception" if self._node_outcome == "failed" else "canceled") if not scope.windows: return True service = self._services.working_memory_calibration @@ -658,6 +667,7 @@ def _on_after_run_node( - Emits an invocation complete event. - Run any callbacks registered for this event. + - Records that the node completed (see `_run_calibration_node`). """ self._services.logger.debug( @@ -672,6 +682,7 @@ def _on_after_run_node( for callback in self._on_after_run_node_callbacks: callback(invocation=invocation, queue_item=queue_item, output=output) + self._node_outcome = "completed" def _on_node_error( self, @@ -684,12 +695,14 @@ def _on_node_error( ): """Called when a node errors. Node errors may occur when running or preparing the node.. + - Records that the node failed, ahead of anything that can raise (see `_run_calibration_node`). - Set the node error on the session object. - Log the error. - Fail the queue item. - Emits an invocation error event. - Run any callbacks registered for this event. """ + self._node_outcome = "failed" self._services.progress_previews.clear_node(queue_item.item_id, invocation.id) self._services.logger.debug( @@ -793,20 +806,27 @@ def __init__( self._thread_limit = thread_limit self._polling_interval = polling_interval self._workers: list[_SessionWorker] = [] - - def _resolve_devices(self) -> list[Optional[torch.device]]: - """Determine the per-worker devices from config. - - Resolves `generation_devices` (which defaults to `"auto"` — every available GPU) into one - normalized device per worker. Returns a single `None` (legacy single-worker, device chosen by - the global config) only if the resolution is empty (e.g. `generation_devices` set to an empty - list). + # One per worker, once `register_generation_devices` resolved them. + self._devices: Optional[list[Optional[torch.device]]] = None + + def register_generation_devices(self, configuration: InvokeAIAppConfig) -> list[Optional[torch.device]]: + """Resolve the workers' devices and register them in `GENERATION_DEVICE_POOL`, so the model loader can find + idle GPUs to borrow. Returns the devices, one per worker. + + Work outside the session queue borrows from the pool from the moment its service starts (the image index's + backfill, say), so the composition root calls this before the invoker starts any service + (`ApiDependencies.initialize`). `start` calls it only when nothing did: registering again would drop a borrow in + flight. Registering creates no CUDA context: the pool holds device objects, and a worker pins its GPU in CUDA + only once it claims its first queue item (#9413). + + `generation_devices` (by default `"auto"`, every available GPU) resolves to one normalized device per worker, + or to a single `None` (one worker on the device the global config chooses) when the resolution is empty, as + with an empty list. """ - generation_devices = self._invoker.services.configuration.generation_devices - devices = TorchDevice.get_generation_devices(generation_devices) - if not devices: - return [None] - return list(devices) + devices = TorchDevice.get_generation_devices(configuration.generation_devices) + self._devices = list(devices) if devices else [None] + GENERATION_DEVICE_POOL.set_generation_devices([d for d in self._devices if d is not None]) + return self._devices def _clone_session_runner(self, template: SessionRunnerBase) -> SessionRunnerBase: """Create an independent runner for an additional worker. @@ -845,11 +865,11 @@ def start(self, invoker: Invoker) -> None: register_events(QueueItemStatusChangedEvent, self._on_queue_item_status_changed) register_events(UserAccessChangedEvent, self._on_user_access_changed) - devices = self._resolve_devices() - - # Register the generation devices so the model loader can discover idle GPUs to host text - # encoders on (see offload_text_encoders_to_idle_gpus). None means legacy single-device mode. - GENERATION_DEVICE_POOL.set_generation_devices([d for d in devices if d is not None]) + devices = ( + self._devices + if self._devices is not None + else self.register_generation_devices(self._invoker.services.configuration) + ) # A worker that stood aside for a lent GPU re-polls as soon as that borrow ends. GENERATION_DEVICE_POOL.set_release_listener(self._poll_now_event.set) diff --git a/invokeai/app/services/session_queue/session_queue_base.py b/invokeai/app/services/session_queue/session_queue_base.py index 3e7bc4841fb..2414e846fdf 100644 --- a/invokeai/app/services/session_queue/session_queue_base.py +++ b/invokeai/app/services/session_queue/session_queue_base.py @@ -54,9 +54,11 @@ def dequeue(self, device: Optional[str] = None) -> Optional[SessionQueueItem]: @abstractmethod def enqueue_batch( - self, queue_id: str, batch: Batch, prepend: bool, user_id: str = "system" + self, queue_id: str, batch: Batch, prepend: bool, user_id: str = "system", *, all_or_none: bool = False ) -> Coroutine[Any, Any, EnqueueBatchResult]: - """Enqueues all permutations of a batch for execution for a specific user.""" + """Enqueues all permutations of a batch for execution for a specific user. A queue without room for all of + them takes as many as fit, or with `all_or_none` none of them. A repeat with the batch's idempotency key + returns the key's settled receipt as it is: `all_or_none` is not part of the payload the receipt matches.""" pass @abstractmethod diff --git a/invokeai/app/services/session_queue/session_queue_default.py b/invokeai/app/services/session_queue/session_queue_default.py index 01df20c0052..c33dfbdfc70 100644 --- a/invokeai/app/services/session_queue/session_queue_default.py +++ b/invokeai/app/services/session_queue/session_queue_default.py @@ -144,11 +144,11 @@ def _prune_terminal_to_limit(self, queue_id: str, keep: int) -> int: return self._queries.session_queue.prune(queue_id, user_id=None, keep=keep) async def enqueue_batch( - self, queue_id: str, batch: Batch, prepend: bool, user_id: str = "system" + self, queue_id: str, batch: Batch, prepend: bool, user_id: str = "system", *, all_or_none: bool = False ) -> EnqueueBatchResult: # The route awaits this method, but every operation below is synchronous database/CPU work. Keep the # complete enqueue in one worker operation, off the event loop. - return await asyncio.to_thread(self._enqueue_batch, queue_id, batch, prepend, user_id) + return await asyncio.to_thread(self._enqueue_batch, queue_id, batch, prepend, user_id, all_or_none) def acknowledge_enqueue(self, queue_id: str, idempotency_key: str, user_id: str = "system") -> None: self._queries.session_queue.acknowledge_receipt(queue_id, user_id, idempotency_key) @@ -159,7 +159,9 @@ def get_enqueue_receipt( receipt = self._queries.session_queue.enqueue_receipt(queue_id, user_id, idempotency_key) return _decode_enqueue_receipt(*receipt) if receipt is not None else None - def _enqueue_batch(self, queue_id: str, batch: Batch, prepend: bool, user_id: str) -> EnqueueBatchResult: + def _enqueue_batch( + self, queue_id: str, batch: Batch, prepend: bool, user_id: str, all_or_none: bool = False + ) -> EnqueueBatchResult: requested_count = calc_session_count(batch=batch) payload_hash = ( hashlib.sha256( @@ -208,11 +210,12 @@ def survey(q: Queries) -> SettledReceipt | int: # Preparing the values is CPU work, done outside a transaction; admission below takes what still fits. max_queue_size = self.__invoker.services.configuration.max_queue_size + room = max(0, max_queue_size - surveyed) prepared_values = prepare_values_to_insert( queue_id=queue_id, batch=batch, priority=0, - max_new_queue_items=max(0, max_queue_size - surveyed), + max_new_queue_items=0 if all_or_none and requested_count > room else room, user_id=user_id, ) @@ -225,6 +228,9 @@ def admit(q: Queries) -> SettledReceipt | EnqueueBatchResult: require_project(q) max_new_queue_items = max(0, max_queue_size - q.session_queue.pending_count(queue_id)) + # Checked under the admission lock: the room the survey saw may have been taken meanwhile. + if all_or_none and min(len(prepared_values), max_new_queue_items) < requested_count: + max_new_queue_items = 0 priority = q.session_queue.top_pending_priority(queue_id) + 1 if prepend else 0 values_to_insert = prepared_values[:max_new_queue_items] if priority != 0: diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py index d36e4a0be1c..68a52bafbb9 100644 --- a/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py @@ -18,7 +18,10 @@ class WorkingMemoryCalibrationServiceBase(ABC): async def create_run( self, model_key: str, request: WorkingMemoryCalibrationRequest, user_id: str ) -> WorkingMemoryCalibrationRunCreated: - """Validate the request, record the run and enqueue its batch as `user_id`, who must be an administrator.""" + """Validate the request, record the run and enqueue its batch as `user_id`, who must be an administrator. + + The queue takes all of the batch's items or none (`CalibrationQueueFullError`). With a retry key, a repeat of + the request returns the run it created and enqueues nothing.""" @abstractmethod def admit(self, queue_item: SessionQueueItem, run_id: str) -> Optional[str]: diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py index e7d852aa6c9..ed4ad8a8b16 100644 --- a/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py @@ -139,8 +139,8 @@ class CalibrationStoreFile(_StoreModel): class WorkingMemoryCalibrationRequest(BaseModel): - """The batch a calibration run enqueues, built by the client. The server sets origin, destination and runs, and - turns the node cache off.""" + """The batch a calibration run enqueues, built by the client. The server sets origin, destination, runs and the + graph's id, and turns the node cache off.""" model_config = ConfigDict(extra="forbid") @@ -150,15 +150,19 @@ class WorkingMemoryCalibrationRequest(BaseModel): default=None, min_length=1, max_length=255, - description="A caller-owned retry key. Repeating a request with it returns the run it created and enqueues " - "nothing new; reusing it with a different request is rejected.", + description="A caller-owned retry key. Repeating the request with it returns the run it created and enqueues " + "nothing new. Reusing it for another request (another model, graph or batch data) is rejected while the " + "queue keeps the key's enqueue receipt.", ) class WorkingMemoryCalibrationRunCreated(BaseModel): run_id: str device_class_id: str - item_ids: list[int] = Field(description="The queue items the run enqueued") + item_ids: list[int] = Field( + description="The queue items the run enqueued. A repeat of the request that finds them unrecorded takes over " + "every queue item carrying the run's origin, retries included" + ) CalibrationRunStatus = Literal["running", "completed", "failed", "partial", "canceled"] @@ -174,7 +178,7 @@ class WorkingMemoryCalibrationRunSummary(BaseModel): status: CalibrationRunStatus = Field( description="Derived from the run's queue items, a retry standing in for the item it retried: running while " "any is pending or running; failed when any failed; completed when all completed; otherwise partial when " - "points were recorded, else canceled" + "valid points were recorded, else canceled" ) error: Optional[str] = Field(description="The error of a failed item, when the run failed") item_ids: list[int] = Field(description="Every queue item of the run, retries included") @@ -251,7 +255,7 @@ def __str__(self) -> str: class CalibrationQueueFullError(Exception): - """The queue had no room for the run's items.""" + """The queue had no room for all of the run's items, so it took none.""" class CalibrationStoreError(Exception): diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py index db1a13bdce8..6138fb3f086 100644 --- a/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py @@ -1,6 +1,9 @@ import asyncio import datetime +import hashlib +import json import uuid +import weakref from collections.abc import Callable, Mapping from functools import partial from itertools import chain @@ -31,7 +34,9 @@ CalibrationRun, CalibrationRunNotFoundError, CalibrationRunStatus, + CalibrationStoreError, CalibrationStoreFile, + CalibrationStoreUnusableError, CalibrationUnavailableError, ObservedItem, WorkingMemoryCalibrationDeviceClass, @@ -77,7 +82,7 @@ """Queue items read per summary query when the status collects the calibration items.""" _RUN_ID_NAMESPACE: Final = uuid.UUID("3c1f6a52-8d0e-4b6f-9a57-2f0d5c8e41b3") -"""Derives a run id from an owner and a retry key.""" +"""Derives a run id from an owner, a retry key and the request's fingerprint.""" _RUNNING: Final[frozenset[QUEUE_ITEM_STATUS]] = frozenset({"pending", "in_progress", "waiting"}) @@ -110,11 +115,12 @@ def _add_points( def derive_run_status( - run: CalibrationRun, live: Mapping[int, QUEUE_ITEM_STATUS], point_count: int + run: CalibrationRun, live: Mapping[int, QUEUE_ITEM_STATUS], valid_point_count: int ) -> tuple[CalibrationRunStatus, Optional[str], list[int]]: """A run's status, error and items, from its items' statuses: `live` as the queue has them now, the run's own observations for items the queue no longer has. An item the queue no longer has counts as canceled unless the - runner saw it finish. A retry stands in for the item it retried.""" + runner saw it finish. A retry stands in for the item it retried. A run that ended without completing is partial + only with valid points: an invalid one (a canceled or failed operation, a spill) measured nothing usable.""" item_ids = sorted(set(run.queue_item_ids) | set(run.observed) | set(live)) def status_of(item_id: int) -> QUEUE_ITEM_STATUS: @@ -145,7 +151,7 @@ def status_of(item_id: int) -> QUEUE_ITEM_STATUS: return "failed", error or f"Queue item {failed[0]} failed", item_ids if all(status == "completed" for status in statuses.values()): return "completed", None, item_ids - return ("partial" if point_count > 0 else "canceled"), None, item_ids + return ("partial" if valid_point_count > 0 else "canceled"), None, item_ids class WorkingMemoryCalibrationService(WorkingMemoryCalibrationServiceBase): @@ -173,6 +179,10 @@ def __init__( self._device_classes: dict[str, DeviceClass] = {} # Which file holds a run, so admitting an item and recording its windows do not scan every file. self._run_files: dict[str, str] = {} + # Creations in flight, referenced here so that one whose request was canceled still runs to its end. + self._creations: set[asyncio.Task[WorkingMemoryCalibrationRunCreated]] = set() + # Per run id while a creation uses it: a repeat of a request waits for the attempt still in flight. + self._run_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() self._invoker: Optional["Invoker"] = None def start(self, invoker: "Invoker") -> None: @@ -242,7 +252,16 @@ def _find_run(self, run_id: str) -> Optional[tuple[str, CalibrationRun]]: return None @staticmethod - def _calibration_batch(request: WorkingMemoryCalibrationRequest, run_id: str) -> Batch: + def _identify(model_key: str, request: WorkingMemoryCalibrationRequest, user_id: str) -> str: + """Check the request's graph and turn its node cache off, then return the run's id, which also becomes the + graph's id. + + Without a retry key the id is fresh. With one it is derived from the owner, the key and the request's + fingerprint: the model and the batch the server builds (graph and batch data). A repeat of the request finds + its run by that id, and another request reusing the key gets another id. The id is part of the batch's + origin, so the queue's receipt for the key refuses that other request. The client's graph id is not part of + the request: a graph sent without one gets a fresh id each time, so a repeat could never match. + """ graph = request.graph for node in graph.nodes.values(): if isinstance(node, CallSavedWorkflowInvocation): @@ -258,9 +277,146 @@ def _calibration_batch(request: WorkingMemoryCalibrationRequest, run_id: str) -> # A cached output would skip the very operation the run measures. for node in graph.nodes.values(): node.use_cache = False + if request.idempotency_key is None: + run_id = uuid_string() + else: + fingerprint = { + "model_key": model_key, + "graph": graph.model_dump(mode="json", exclude={"id"}), + "data": None + if request.data is None + else [[datum.model_dump(mode="json") for datum in group] for group in request.data], + } + digest = hashlib.sha256( + json.dumps(fingerprint, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode("utf-8") + ).hexdigest() + run_id = str(uuid.uuid5(_RUN_ID_NAMESPACE, f"{user_id}\x00{request.idempotency_key}\x00{digest}")) + # The queue's receipt hashes the batch, graph id included: a repeat must build the same one. + graph.id = run_id + return run_id + + def _run_lock(self, run_id: str) -> asyncio.Lock: + lock = self._run_locks.get(run_id) + if lock is None: + lock = self._run_locks[run_id] = asyncio.Lock() + return lock + + async def create_run( + self, model_key: str, request: WorkingMemoryCalibrationRequest, user_id: str + ) -> WorkingMemoryCalibrationRunCreated: + # The enqueue runs in a worker thread and can commit after the request was canceled. So the creation runs to + # its end whatever happens to the request: the run keeps the items that were queued, or is removed. Only the + # server's shutdown cancels the creation as well; the run stays as far as it got, for a keyed repeat to finish. + creation = asyncio.create_task(self._create_run(model_key, request, user_id)) + self._creations.add(creation) + creation.add_done_callback(self._creations.discard) + try: + return await asyncio.shield(creation) + except asyncio.CancelledError: + creation.add_done_callback(self._log_unawaited_failure) + raise + + def _log_unawaited_failure(self, creation: "asyncio.Task[WorkingMemoryCalibrationRunCreated]") -> None: + """Nobody awaits the creation of a canceled request: its failure would go unseen.""" + if not creation.cancelled() and (error := creation.exception()) is not None: + self._services.logger.warning( + "The creation of a calibration run whose request was canceled failed", exc_info=error + ) + + async def _create_run( + self, model_key: str, request: WorkingMemoryCalibrationRequest, user_id: str + ) -> WorkingMemoryCalibrationRunCreated: + run_id = await asyncio.to_thread(self._identify, model_key, request, user_id) + # Recording the run, enqueuing its batch and recording its items or removing it: one step for each run id. The + # lock is this process's; two processes serving one root and its SQLite database are not kept apart. + async with self._run_lock(run_id): + prepared = await asyncio.to_thread(self._prepare_run, model_key, request, user_id, run_id) + if isinstance(prepared, WorkingMemoryCalibrationRunCreated): + return prepared + batch, device_class_id, created = prepared + try: + result = await self._services.session_queue.enqueue_batch( + queue_id=DEFAULT_QUEUE_ID, batch=batch, prepend=False, user_id=user_id, all_or_none=True + ) + except Exception: + # The queue announces a batch after committing it, so an enqueue can raise once its items are queued. + # Items of an earlier attempt that a reset left behind carry the run's origin too, and answer a repeat + # of the request as the key's receipt would. + try: + queued: Optional[list[int]] = await asyncio.to_thread(self._queued_item_ids, run_id) + except Exception: + self._services.logger.warning( + f"Could not tell whether calibration run {run_id} has queue items; it is kept", exc_info=True + ) + queued = None + if queued is None: + # The enqueue's own error answers the request. + raise + if not queued: + if created: + await asyncio.to_thread(self._discard, run_id) + raise + self._services.logger.warning( + f"Enqueuing calibration run {run_id} raised, and the queue holds its items {queued}", + exc_info=True, + ) + item_ids = queued + else: + if result.enqueued == 0: + if created: + await asyncio.to_thread(self._discard, run_id) + # A partial run would measure some of the batch's cases, so the queue takes all of them or none. + raise CalibrationQueueFullError( + f"The queue is full: it has no room for all {result.requested} items of the calibration run, " + "so the run was not started." + ) + item_ids = list(result.item_ids) + try: + await asyncio.to_thread(self._record_items, device_class_id, run_id, item_ids) + except CalibrationStoreError as error: + # An unusable file fails the run's items at admission, as a removed run (not caught here) does: an error + # with a key or without. With a key, a repeat records the items, from the key's receipt or, once the + # queue has collected that, from their origin. Without one the request cannot be repeated, only sent + # again, which would queue the batch twice: a passing store failure answers with the items. + if isinstance(error, CalibrationStoreUnusableError) or request.idempotency_key is not None: + raise + self._services.logger.warning( + f"Could not record the queue items {item_ids} of calibration run {run_id}; they run all the same, " + "and the status finds them by their origin", + exc_info=True, + ) + return WorkingMemoryCalibrationRunCreated(run_id=run_id, device_class_id=device_class_id, item_ids=item_ids) + + def _prepare_run( + self, model_key: str, request: WorkingMemoryCalibrationRequest, user_id: str, run_id: str + ) -> WorkingMemoryCalibrationRunCreated | tuple[Batch, str, bool]: + """The run a repeated request already enqueued; or the batch to enqueue, the run's device class and whether + this call recorded the run.""" + key = request.idempotency_key + found: Optional[tuple[str, CalibrationRun]] = None + if key is not None and (found := self._find_run(run_id)) is not None: + device_class_id, run = found + if run.queue_item_ids: + return WorkingMemoryCalibrationRunCreated( + run_id=run_id, device_class_id=device_class_id, item_ids=list(run.queue_item_ids) + ) + # An earlier attempt may have queued the batch without recording its items. While the queue keeps the + # key's receipt, the enqueue below answers with the items it names, all of them. Without the receipt, + # enqueuing again would queue the batch a second time: the items carrying the run's origin, and those the + # runner saw, are taken over instead. + if self._services.session_queue.get_enqueue_receipt(DEFAULT_QUEUE_ID, key, user_id) is None: + item_ids = sorted(set(self._queued_item_ids(run_id)) | set(run.observed)) + if item_ids: + self._record_items(device_class_id, run_id, item_ids) + return WorkingMemoryCalibrationRunCreated( + run_id=run_id, device_class_id=device_class_id, item_ids=item_ids + ) + device = self._calibration_device() + # Raises UnknownModelException for an unknown key. + config = self._services.model_manager.store.get_model(model_key) try: - return Batch( - graph=graph, + batch = Batch( + graph=request.graph, data=request.data, runs=1, origin=calibration_origin(run_id), @@ -269,85 +425,56 @@ def _calibration_batch(request: WorkingMemoryCalibrationRequest, run_id: str) -> ) except ValueError as error: raise CalibrationGraphError(str(error)) from error + if found is not None: + # The run is recorded without items: enqueue its batch, which the key's receipt answers while it exists. + return batch, found[0], False + device_class = self._device_class(device) + env = self._environment(device) + run = CalibrationRun( + run_id=run_id, + model_key=model_key, + model_hash=getattr(config, "hash", None), + user_id=user_id, + created_at=_now(), + env_id=env.id, + ) - @staticmethod - def _run_id_for(user_id: str, idempotency_key: Optional[str]) -> str: - """A fresh run id, or with a retry key the same one for every repeat of the request: the run id is part of the - batch's origin, so a replay must rebuild the identical batch for the queue to recognise it.""" - if idempotency_key is None: - return uuid_string() - return str(uuid.uuid5(_RUN_ID_NAMESPACE, f"{user_id}\x00{idempotency_key}")) - - async def create_run( - self, model_key: str, request: WorkingMemoryCalibrationRequest, user_id: str - ) -> WorkingMemoryCalibrationRunCreated: - run_id = self._run_id_for(user_id, request.idempotency_key) - - def prepare() -> tuple[Batch, str, bool]: - """The batch, the run's device class, and whether this call recorded the run (a replay finds it).""" - device = self._calibration_device() - # Raises UnknownModelException for an unknown key. - config = self._services.model_manager.store.get_model(model_key) - batch = self._calibration_batch(request, run_id) - device_class = self._device_class(device) - env = self._environment(device) - run = CalibrationRun( - run_id=run_id, - model_key=model_key, - model_hash=getattr(config, "hash", None), - user_id=user_id, - created_at=_now(), - env_id=env.id, - ) - - def add(file: CalibrationStoreFile) -> tuple[bool, bool]: - if any(known.run_id == run_id for known in file.runs): - return False, False - file.runs.append(run) - _add_env(file, env) - return True, True - - created = self._store.change(device_class.id, add, create=CalibrationDeviceClass.of(device_class)) - self._run_files[run_id] = device_class.id - return batch, device_class.id, bool(created) - - batch, device_class_id, created = await asyncio.to_thread(prepare) - try: - result = await self._services.session_queue.enqueue_batch( - queue_id=DEFAULT_QUEUE_ID, batch=batch, prepend=False, user_id=user_id - ) - except BaseException: - # Only a run this call recorded: a conflicting replay must not remove the original request's run. - if created: - await asyncio.to_thread(self._discard_run, run_id) - raise - if result.enqueued == 0: - if created: - await asyncio.to_thread(self._discard_run, run_id) - raise CalibrationQueueFullError("The queue is full; the calibration run was not started.") + def add(file: CalibrationStoreFile) -> tuple[bool, bool]: + if any(known.run_id == run_id for known in file.runs): + return False, False + file.runs.append(run) + _add_env(file, env) + return True, True - def remember_items(file: CalibrationStoreFile) -> tuple[bool, None]: - for run in file.runs: - if run.run_id == run_id and run.queue_item_ids != result.item_ids: - run.queue_item_ids = list(result.item_ids) - return True, None - return False, None + created = self._store.change(device_class.id, add, create=CalibrationDeviceClass.of(device_class)) + self._run_files[run_id] = device_class.id + return batch, device_class.id, bool(created) - try: - await asyncio.to_thread(self._store.change, device_class_id, remember_items) - except Exception: - # The items are enqueued and carry the run's origin, which is what the status reads them by. - self._services.logger.warning(f"Could not record the items of calibration run {run_id}", exc_info=True) - return WorkingMemoryCalibrationRunCreated( - run_id=run_id, device_class_id=device_class_id, item_ids=list(result.item_ids) - ) + def _queued_item_ids(self, run_id: str) -> list[int]: + """The queue items carrying the run's origin: its batch's, with their retries and workflow-call children.""" + queue = self._services.session_queue + return sorted(queue.get_queue_item_ids(DEFAULT_QUEUE_ID, origin_prefix=calibration_origin(run_id)).item_ids) - def _discard_run(self, run_id: str) -> None: + def _discard(self, run_id: str) -> None: + """Remove a run this request recorded and queued nothing for.""" try: self.reset_run(run_id) except Exception: self._services.logger.warning(f"Could not discard calibration run {run_id}", exc_info=True) + def _record_items(self, device_class_id: str, run_id: str, item_ids: list[int]) -> None: + def remember(file: CalibrationStoreFile) -> tuple[bool, bool]: + for run in file.runs: + if run.run_id == run_id: + changed = run.queue_item_ids != item_ids + run.queue_item_ids = list(item_ids) + return changed, True + return False, False + + if not self._store.change(device_class_id, remember): + # Reset while it was being created: its items fail at admission. + raise CalibrationRunNotFoundError(run_id) + def admit(self, queue_item: SessionQueueItem, run_id: str) -> Optional[str]: if self._services.configuration.multiuser: owner = self._services.users.get(queue_item.user_id) @@ -538,17 +665,22 @@ def _device_class_status( envs=[], ) points_per_run: dict[str, int] = {} + valid_points_per_run: dict[str, int] = {} points_per_env: dict[str, int] = {} counts: dict[tuple[str, int, bool, bool], int] = {} for point in file.points: points_per_run[point.run_id] = points_per_run.get(point.run_id, 0) + 1 + if point.valid: + valid_points_per_run[point.run_id] = valid_points_per_run.get(point.run_id, 0) + 1 points_per_env[point.env_id] = points_per_env.get(point.env_id, 0) + 1 key = (point.estimator, point.version, point.valid, point.fit_eligible) counts[key] = counts.get(key, 0) + 1 runs = [] for run in file.runs: point_count = points_per_run.get(run.run_id, 0) - status, error, item_ids = derive_run_status(run, live.get(run.run_id, {}), point_count) + status, error, item_ids = derive_run_status( + run, live.get(run.run_id, {}), valid_points_per_run.get(run.run_id, 0) + ) runs.append( WorkingMemoryCalibrationRunSummary( run_id=run.run_id, diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index 05c75017ada..e11c945e8b1 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -189,12 +189,15 @@ def _lock_on_device(self, working_mem_bytes: Optional[int]) -> None: While a calibration scope is open anywhere (`probe.armed`), the working-memory probe hears of the lock from its start until it has settled with the tensors repaired: inside a measured window on this thread that interval is a co-lock, which moves weights and is kept out of the window's peak; on another thread it voids the window. + The lock is in `probe.locking` from before `armed` is read until it has settled or failed, so a window that + opens on this cache meanwhile is void even if this read missed the scope arming. """ - observed = probe.armed - if observed: - probe.lock_started(self._cache, self._cache_record) - locked = False + observed = locked = False + probe.locking.append(self._cache) try: + observed = probe.armed + if observed: + probe.lock_started(self._cache, self._cache_record) self._lock_paced(working_mem_bytes) locked = True self._end_first_use_window() @@ -205,6 +208,7 @@ def _lock_on_device(self, working_mem_bytes: Optional[int]) -> None: locked = False raise finally: + probe.locking.remove(self._cache) if observed: probe.lock_settled(self._cache, self._cache_record, locked) diff --git a/invokeai/backend/util/device_pool.py b/invokeai/backend/util/device_pool.py index e9044160689..fbd68f7af16 100644 --- a/invokeai/backend/util/device_pool.py +++ b/invokeai/backend/util/device_pool.py @@ -75,7 +75,10 @@ def set_release_listener(self, listener: Optional[Callable[[], None]]) -> None: self._on_release = listener def set_generation_devices(self, devices: list[torch.device]) -> None: - """Register the full set of generation devices (called once at processor startup). + """Register the full set of generation devices, replacing those registered before (a borrow in flight is + dropped). The session processor's `register_generation_devices` calls it once: from the composition root + (`ApiDependencies.initialize`) before the invoker starts any service, or from the processor's `start` when + nothing did. Only GPU devices participate in idle-offload; others are ignored. """ diff --git a/invokeai/backend/working_memory/counters.py b/invokeai/backend/working_memory/counters.py index 1089ff5bb33..d807541183a 100644 --- a/invokeai/backend/working_memory/counters.py +++ b/invokeai/backend/working_memory/counters.py @@ -18,7 +18,9 @@ The driver is read at boundaries only: torch's peaks are exact, so what a boundary read misses is memory outside torch's allocator that came and went between two boundaries. -Every reader is best-effort: a failure reads as None, never as an exception in a generation. +Every reader is best-effort: a failure reads as None. A peak reset or a synchronization that fails raises instead: the +window would go on with peaks from before its boundary or with work still in flight, so the probe, which never lets an +exception reach a generation, marks the window invalid. """ import dataclasses @@ -202,27 +204,22 @@ def read_counters(device: torch.device) -> CounterReading: def reset_peaks(device: torch.device) -> None: - """Start torch's peak counters over at the current values. Cumulative counters (OOMs, retries) are kept.""" - try: - if device.type == "cuda": - torch.cuda.reset_peak_memory_stats(device) - elif device.type == "xpu": - torch.xpu.reset_peak_memory_stats(device) - except Exception: - InvokeAILogger.get_logger(__name__).debug(f"Peak reset for {device} failed", exc_info=True) + """Start torch's peak counters over at the current values. Cumulative counters (OOMs, retries) are kept. Raises + what torch raises.""" + if device.type == "cuda": + torch.cuda.reset_peak_memory_stats(device) + elif device.type == "xpu": + torch.xpu.reset_peak_memory_stats(device) def synchronize(device: torch.device) -> None: - """Wait for the work queued on `device`, so frees it implies have happened.""" - try: - if device.type == "cuda": - torch.cuda.synchronize(device) - elif device.type == "xpu": - torch.xpu.synchronize(device) - elif device.type == "mps": - torch.mps.synchronize() - except Exception: - InvokeAILogger.get_logger(__name__).debug(f"Synchronizing {device} failed", exc_info=True) + """Wait for the work queued on `device`, so frees it implies have happened. Raises what torch raises.""" + if device.type == "cuda": + torch.cuda.synchronize(device) + elif device.type == "xpu": + torch.xpu.synchronize(device) + elif device.type == "mps": + torch.mps.synchronize() def _larger(a: Optional[int], b: Optional[int]) -> Optional[int]: diff --git a/invokeai/backend/working_memory/probe.py b/invokeai/backend/working_memory/probe.py index bd888590300..839e2d785f2 100644 --- a/invokeai/backend/working_memory/probe.py +++ b/invokeai/backend/working_memory/probe.py @@ -25,18 +25,21 @@ baseline and that segment's offset; `demand_bytes` is the largest delta of the counters that could be read. A window is invalid (`valid=False`, with its reasons) when another thread locked a model on the same device's cache -while it was open or held one when it opened (Expand Prompt and Image to Prompt run next to sessions), when an -allocation failed or only succeeded after the allocator released its cache (a retry squeezes the peak: an untiled -3072px FLUX VAE decode on an RTX 4090 read 22.5 GiB at the card's limit, ~1280 bytes per pixel instead of ~2150), when -Windows still keeps more than `SPILL_TOLERANCE_BYTES` of the process in shared memory `SPILL_SETTLE_S` after close -(right after a decode the counter spikes transiently, 589 MiB, and is back at 9 MiB at the next read), when the -operation raised or was canceled, when the Krea-2 SDPA benchmark (which resets the peaks itself) is active, when XPU's -device-wide usage grew beyond this process's own, when a probe step failed (`probe_error`; a failed lock hook also voids -the scope's later windows, since its lock bookkeeping is suspect), and when no counter could be read. A model less than fully resident at open keeps the window valid but not `fit_eligible`: streamed -weights are working memory the estimate does not price. - -Without a scope, the probe costs a context-variable read per estimated `model_on_device`, a read of `armed` per lock and -unlock, and a truth test of `devices.EMPTY_CACHE_OBSERVERS` per `empty_cache`. +while it was open, held one when it opened or was still locking one then (Expand Prompt and Image to Prompt run next to +sessions; see `locking`), when an allocation failed or only succeeded after the allocator released its cache (a retry +squeezes the peak: an untiled 3072px FLUX VAE decode on an RTX 4090 read 22.5 GiB at the card's limit, ~1280 bytes per +pixel instead of ~2150), when Windows still keeps more than `SPILL_TOLERANCE_BYTES` of the process in shared memory +`SPILL_SETTLE_S` after close (right after a decode the counter spikes transiently, 589 MiB, and is back at 9 MiB at the +next read), when the operation raised or was canceled, or its node failed or was canceled after the window closed +(`ProbeScope.invalidate`), when the Krea-2 SDPA benchmark (which resets the peaks itself) is active, when XPU's +device-wide usage grew beyond this process's own, when a probe step failed (`probe_error`: a reading, a release, a peak +reset or a synchronization, at the window's open as much as later; a failed lock hook also voids the scope's later +windows, since its lock bookkeeping is suspect), and when no counter could be read. A model less than fully resident at +open keeps the window valid but not `fit_eligible`: streamed weights are working memory the estimate does not price. + +Without a scope, the probe costs a context-variable read per estimated `model_on_device`, an entry in `locking` and a +read of `armed` per lock, a read of `armed` per unlock, and a truth test of `devices.EMPTY_CACHE_OBSERVERS` per +`empty_cache`. """ import contextvars @@ -96,6 +99,13 @@ _NOT_INVALIDATING: Final[frozenset[Reason]] = frozenset({"partial_residency"}) + +def _validity(reasons: list[Reason]) -> tuple[bool, bool]: + """``(valid, fit_eligible)`` of a window with `reasons`.""" + valid = all(reason in _NOT_INVALIDATING for reason in reasons) + return valid, valid and "partial_residency" not in reasons + + CounterName = Literal["driver", "reserved", "allocated"] _COUNTERS: Final[tuple[CounterName, ...]] = ("driver", "reserved", "allocated") @@ -156,6 +166,14 @@ class MeasuredWindow(_Record): fit_eligible: bool reasons: list[Reason] + def invalidated(self, reason: Reason) -> "MeasuredWindow": + """This window with `reason` among its reasons, and valid and fit-eligible as they then are.""" + if reason in self.reasons: + return self + reasons = [*self.reasons, reason] + valid, fit_eligible = _validity(reasons) + return self.model_copy(update={"reasons": reasons, "valid": valid, "fit_eligible": fit_eligible}) + def calibration_reservation(reservation_bytes: int) -> int: """What a calibration lock reserves for an estimate of `reservation_bytes`: half as much again, at least 1 GiB more.""" @@ -176,6 +194,14 @@ def calibration_reservation(reservation_bytes: int) -> int: """True while any thread has a calibration scope open. The model loader checks it before telling the probe about a lock or an unlock, so without a scope that costs one attribute read.""" +locking: list[object] = [] +"""The cache of every model lock in progress, once per lock, whichever thread takes it. The model loader adds it before +it reads `armed` and removes it once the lock has settled or failed. A lock that read `armed` before a scope armed never +reports to the probe, and one that reported its start before a window was registered reports again only once it has +settled, perhaps after the window closed. Both move weights while the window measures, so a window that opens while +another thread's lock is in progress on its cache is void. A list, so that the loader's side is an append and a remove, +each atomic under the GIL, and takes no lock.""" + _boundary = threading.local() @@ -329,7 +355,7 @@ def measured(self, duration_s: float) -> MeasuredWindow: deltas = [delta.delta for delta in counter_deltas.values() if delta is not None] if not deltas: self.add_reason("counters_unavailable") - valid = not any(reason not in _NOT_INVALIDATING for reason in self.reasons) + valid, fit_eligible = _validity(self.reasons) method = self.method() backend = backend_for(self.device) estimate, scope = self.estimate, self.scope @@ -358,7 +384,7 @@ def measured(self, duration_s: float) -> MeasuredWindow: segments=self.segments, duration_s=duration_s, valid=valid, - fit_eligible=valid and "partial_residency" not in self.reasons, + fit_eligible=fit_eligible, reasons=list(self.reasons), ) @@ -370,20 +396,26 @@ def close(self, error: Optional[BaseException]) -> None: InvokeAILogger.get_logger(__name__).warning( f"Working-memory probe failed to close the window of {self.record.key}", exc_info=True ) - try: - with _state_lock: - self.add_reason("probe_error") - self._unregister() - parent = self.parent - if parent is not None and parent.state == "suspended": - # Its peaks were not reset after this window, so what it measures from here is suspect. - parent.fail() - parent.state = "measuring" - if not self.recorded: - self.scope.windows.append(self.measured(time.perf_counter() - self.opened_at)) - self.recorded = True - except Exception: - InvokeAILogger.get_logger(__name__).debug("Working-memory probe could not record", exc_info=True) + self.abort() + + def abort(self) -> None: + """End the window after a probe step failed as it opened or closed; record it as invalid. Never raises.""" + try: + with _state_lock: + self.add_reason("probe_error") + self._unregister() + parent = self.parent + if parent is not None and parent.state == "suspended": + # Folded and released for this window, and not reset after it: what it measures from here is + # suspect. + parent.fail() + parent.state = "measuring" + if not self.recorded: + duration_s = time.perf_counter() - self.opened_at if self.first is not None else 0.0 + self.scope.windows.append(self.measured(duration_s)) + self.recorded = True + except Exception: + InvokeAILogger.get_logger(__name__).debug("Working-memory probe could not record", exc_info=True) def _unregister(self) -> None: if self in _open_windows: @@ -473,18 +505,28 @@ def __init__(self, *, run_id: str, queue_item_id: int, node_id: str, node_type: def _held_on(self, cache: object) -> int: return sum(count for (cache_id, _key), count in self._held.items() if cache_id == id(cache)) + def invalidate(self, reason: Literal["exception", "canceled"]) -> None: + """Void every window recorded so far: the node failed or was canceled. + + The session runner swallows a node's failure or cancellation, and calls this instead. It covers the windows that + closed before the node raised: their own operation finished, but the node did not, and no measurement of a node + that did not complete is trusted. + """ + self.windows = [window.invalidated(reason) for window in self.windows] + def open_window( self, cache: object, record: "CacheRecord", estimate: WorkingMemoryEstimate, reservation_applied_bytes: int ) -> Optional[_Window]: """Start measuring the operation `record`'s model is about to run, once its lock has settled. `cache` is the model cache the record is locked in. Never raises: a probe failure leaves the operation alone - and records the window as invalid, or not at all. + and records the window as invalid. A window that cannot be opened is recorded at once, and the operation runs + unmeasured. """ window = _Window(self, cache, record, estimate, reservation_applied_bytes) - # Registered before the cache is asked: from here on, a lock another thread takes on this cache voids the - # window (`lock_started`, `lock_settled`), and one it already holds shows in `locks_held()`. Its counters are - # not open yet, so releases pass it by. + # Registered before the cache is asked: from here on, a lock another thread starts or settles on this cache + # voids the window (`lock_started`, `lock_settled`), and one in progress or held shows in `locking` and + # `locks_held()`. Its counters are not open yet, so releases pass it by. with _state_lock: _open_windows.append(window) try: @@ -501,9 +543,12 @@ def open_window( for (cache_id, key), count in self._held.items() if cache_id == id(cache) and count > 0 and key != record.key ) - # Asked outside the probe's lock (see `_state_lock`). Every lock this thread took since the scope opened - # is in `_held`, so any lock beyond those belongs to another thread. - if cache.locks_held() > self._held_on(cache): # type: ignore[attr-defined] + # Asked outside the probe's lock (see `_state_lock`). A lock in progress is another thread's: this thread + # opens a window only once its own lock has settled. Every lock this thread took since the scope opened is + # in `_held`, so any lock held beyond those is another thread's too. `locking` is read first: a lock that + # leaves it before `locks_held()` is read still holds its record, or is over before the baseline. + in_progress = any(entry is cache for entry in tuple(locking)) + if in_progress or cache.locks_held() > self._held_on(cache): # type: ignore[attr-defined] window.add_reason("foreign_lock") if window.device.type == "cuda" and _krea2_sdpa_benchmark_active(): window.add_reason("krea2_sdpa_benchmark") @@ -530,13 +575,7 @@ def open_window( InvokeAILogger.get_logger(__name__).warning( f"Working-memory probe could not open a window for {record.key}", exc_info=True ) - with _state_lock: - window._unregister() - parent = window.parent - if parent is not None and parent.state == "suspended": - # Folded and released for a window that never opened. - parent.fail() - parent.state = "measuring" + window.abort() return None def _abandon(self) -> None: diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json index 7787a41e476..ee6add31301 100644 --- a/invokeai/frontend/api/openapi.json +++ b/invokeai/frontend/api/openapi.json @@ -5664,7 +5664,7 @@ "post": { "tags": ["model_manager"], "summary": "Calibrate Model Working Memory", - "description": "Start a working-memory calibration run: enqueue the batch with every node uncached and measure each node's\nmodel operations on the generation device. Needs exactly one generation device, and it must be a GPU.\n\nWith an `idempotency_key`, repeating the request (after a lost response) returns the run it created and enqueues\nnothing new; acknowledge the key through the queue's enqueue_batch/acknowledge route once the response arrived.", + "description": "Start a working-memory calibration run: enqueue the batch with every node uncached and measure each node's\nmodel operations on the generation device. Needs exactly one generation device, and it must be a GPU. The queue\ntakes all of the batch's items or none.\n\nOnce the items are queued, a request without an `idempotency_key` succeeds even when the store fails for a passing\nreason (an unwritable root, a lock timeout) to record them: they run all the same, and the status finds them by\ntheir origin. With a key, that failure is a 503 although the items are queued. A run removed meanwhile (404) or an\nunusable store file (409) fails the request with or without a key; its queued items then fail when they start.\n\nWith a key, repeating the request (also after a lost response) records the queued items and returns the run; it\nenqueues nothing new. The items come from the queue's enqueue receipt for the key while it exists, otherwise from\nthe queue items that carry the run's origin and those the runner saw. Acknowledge the key through the queue's\nenqueue_batch/acknowledge route once the response arrived.", "operationId": "calibrate_model_working_memory", "security": [ { @@ -5707,10 +5707,10 @@ } }, "404": { - "description": "The model could not be found" + "description": "The model could not be found, or the run was removed while it was being created; its items are queued then and fail when they start" }, "409": { - "description": "This machine cannot calibrate now (no GPU, or several generation devices), the device class's store file is unusable and kept as it is, the queue is full, or the idempotency key was used for another request" + "description": "This machine cannot calibrate now (no GPU, or several generation devices), the device class's store file is unusable and kept as it is (possibly after the items were queued, which then fail when they start), the queue has no room for all of the run's items (it then takes none), or the idempotency key was used for another request" }, "422": { "description": "The graph cannot be measured as a calibration run" @@ -5719,7 +5719,7 @@ "description": "Too many unacknowledged idempotency keys; acknowledge or wait, then retry" }, "503": { - "description": "The calibration store cannot be read or written (unwritable root, lock timeout)" + "description": "The calibration store cannot be read or written (unwritable root, lock timeout); with an idempotency key, possibly after the items were queued" } } } @@ -12902,6 +12902,9 @@ }, "description": "Created" }, + "400": { + "description": "The batch's origin starts with 'invokeai:wm-calibration:' (in any case), which is reserved for working-memory calibration runs" + }, "422": { "description": "Validation Error", "content": { @@ -117028,14 +117031,14 @@ } ], "title": "Idempotency Key", - "description": "A caller-owned retry key. Repeating a request with it returns the run it created and enqueues nothing new; reusing it with a different request is rejected." + "description": "A caller-owned retry key. Repeating the request with it returns the run it created and enqueues nothing new. Reusing it for another request (another model, graph or batch data) is rejected while the queue keeps the key's enqueue receipt." } }, "additionalProperties": false, "type": "object", "required": ["graph"], "title": "WorkingMemoryCalibrationRequest", - "description": "The batch a calibration run enqueues, built by the client. The server sets origin, destination and runs, and\nturns the node cache off." + "description": "The batch a calibration run enqueues, built by the client. The server sets origin, destination, runs and the\ngraph's id, and turns the node cache off." }, "WorkingMemoryCalibrationResetResult": { "properties": { @@ -117068,7 +117071,7 @@ }, "type": "array", "title": "Item Ids", - "description": "The queue items the run enqueued" + "description": "The queue items the run enqueued. A repeat of the request that finds them unrecorded takes over every queue item carrying the run's origin, retries included" } }, "type": "object", @@ -117113,7 +117116,7 @@ "type": "string", "enum": ["running", "completed", "failed", "partial", "canceled"], "title": "Status", - "description": "Derived from the run's queue items, a retry standing in for the item it retried: running while any is pending or running; failed when any failed; completed when all completed; otherwise partial when points were recorded, else canceled" + "description": "Derived from the run's queue items, a retry standing in for the item it retried: running while any is pending or running; failed when any failed; completed when all completed; otherwise partial when valid points were recorded, else canceled" }, "error": { "anyOf": [ diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts index 92a48342ddf..68f521a5ae9 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -1207,10 +1207,18 @@ export type paths = { /** * Calibrate Model Working Memory * @description Start a working-memory calibration run: enqueue the batch with every node uncached and measure each node's - * model operations on the generation device. Needs exactly one generation device, and it must be a GPU. + * model operations on the generation device. Needs exactly one generation device, and it must be a GPU. The queue + * takes all of the batch's items or none. * - * With an `idempotency_key`, repeating the request (after a lost response) returns the run it created and enqueues - * nothing new; acknowledge the key through the queue's enqueue_batch/acknowledge route once the response arrived. + * Once the items are queued, a request without an `idempotency_key` succeeds even when the store fails for a passing + * reason (an unwritable root, a lock timeout) to record them: they run all the same, and the status finds them by + * their origin. With a key, that failure is a 503 although the items are queued. A run removed meanwhile (404) or an + * unusable store file (409) fails the request with or without a key; its queued items then fail when they start. + * + * With a key, repeating the request (also after a lost response) records the queued items and returns the run; it + * enqueues nothing new. The items come from the queue's enqueue receipt for the key while it exists, otherwise from + * the queue items that carry the run's origin and those the runner saw. Acknowledge the key through the queue's + * enqueue_batch/acknowledge route once the response arrived. */ post: operations["calibrate_model_working_memory"]; delete?: never; @@ -52846,8 +52854,8 @@ export type components = { }; /** * WorkingMemoryCalibrationRequest - * @description The batch a calibration run enqueues, built by the client. The server sets origin, destination and runs, and - * turns the node cache off. + * @description The batch a calibration run enqueues, built by the client. The server sets origin, destination, runs and the + * graph's id, and turns the node cache off. */ WorkingMemoryCalibrationRequest: { /** @description The graph to measure */ @@ -52859,7 +52867,7 @@ export type components = { data?: components["schemas"]["BatchDatum"][][] | null; /** * Idempotency Key - * @description A caller-owned retry key. Repeating a request with it returns the run it created and enqueues nothing new; reusing it with a different request is rejected. + * @description A caller-owned retry key. Repeating the request with it returns the run it created and enqueues nothing new. Reusing it for another request (another model, graph or batch data) is rejected while the queue keeps the key's enqueue receipt. */ idempotency_key?: string | null; }; @@ -52878,7 +52886,7 @@ export type components = { device_class_id: string; /** * Item Ids - * @description The queue items the run enqueued + * @description The queue items the run enqueued. A repeat of the request that finds them unrecorded takes over every queue item carrying the run's origin, retries included */ item_ids: number[]; }; @@ -52901,7 +52909,7 @@ export type components = { env_id: string; /** * Status - * @description Derived from the run's queue items, a retry standing in for the item it retried: running while any is pending or running; failed when any failed; completed when all completed; otherwise partial when points were recorded, else canceled + * @description Derived from the run's queue items, a retry standing in for the item it retried: running while any is pending or running; failed when any failed; completed when all completed; otherwise partial when valid points were recorded, else canceled * @enum {string} */ status: "running" | "completed" | "failed" | "partial" | "canceled"; @@ -56119,14 +56127,14 @@ export interface operations { "application/json": components["schemas"]["WorkingMemoryCalibrationRunCreated"]; }; }; - /** @description The model could not be found */ + /** @description The model could not be found, or the run was removed while it was being created; its items are queued then and fail when they start */ 404: { headers: { [name: string]: unknown; }; content?: never; }; - /** @description This machine cannot calibrate now (no GPU, or several generation devices), the device class's store file is unusable and kept as it is, the queue is full, or the idempotency key was used for another request */ + /** @description This machine cannot calibrate now (no GPU, or several generation devices), the device class's store file is unusable and kept as it is (possibly after the items were queued, which then fail when they start), the queue has no room for all of the run's items (it then takes none), or the idempotency key was used for another request */ 409: { headers: { [name: string]: unknown; @@ -56147,7 +56155,7 @@ export interface operations { }; content?: never; }; - /** @description The calibration store cannot be read or written (unwritable root, lock timeout) */ + /** @description The calibration store cannot be read or written (unwritable root, lock timeout); with an idempotency key, possibly after the items were queued */ 503: { headers: { [name: string]: unknown; @@ -60078,6 +60086,13 @@ export interface operations { "application/json": components["schemas"]["EnqueueBatchResult"]; }; }; + /** @description The batch's origin starts with 'invokeai:wm-calibration:' (in any case), which is reserved for working-memory calibration runs */ + 400: { + headers: { + [name: string]: unknown; + }; + content?: never; + }; /** @description Validation Error */ 422: { headers: { diff --git a/tests/app/routers/test_working_memory_calibration.py b/tests/app/routers/test_working_memory_calibration.py index fb765c39569..056b75f71f9 100644 --- a/tests/app/routers/test_working_memory_calibration.py +++ b/tests/app/routers/test_working_memory_calibration.py @@ -1,5 +1,6 @@ """The calibration routes and the reserved origin, against a real SQLite queue and a store on a temporary root.""" +import threading from pathlib import Path from typing import Any from unittest.mock import MagicMock @@ -8,6 +9,7 @@ import torch from fastapi.testclient import TestClient +from invokeai.app.api_app import app from invokeai.app.invocations.call_saved_workflow import CallSavedWorkflowInvocation from invokeai.app.invocations.math import AddInvocation from invokeai.app.services.invoker import Invoker @@ -20,7 +22,13 @@ calibration_origin, ) from tests.app.routers.conftest import MockApiDependencies, _auth -from tests.fixtures.working_memory_calibration import GPU, GPU_CLASS, calibration_service, measured_window +from tests.fixtures.working_memory_calibration import ( + GPU, + GPU_CLASS, + HeldEnqueue, + calibration_service, + measured_window, +) _DEVICES: list[torch.device] = [GPU] @@ -125,6 +133,12 @@ def test_enqueue_batch_refuses_the_reserved_origin_even_to_admins( assert queue.get_queue_item_ids("default").item_ids == [] +def test_enqueue_batch_documents_the_reserved_origin_refusal() -> None: + responses = app.openapi()["paths"]["/api/v1/queue/{queue_id}/enqueue_batch"]["post"]["responses"] + + assert "reserved for working-memory calibration" in responses["400"]["description"] + + def test_calibrate_refuses_more_than_one_generation_device_and_cpu( client: TestClient, admin_token: str, calibration ) -> None: @@ -235,6 +249,43 @@ def test_calibrate_into_a_full_queue_is_refused_and_leaves_no_run( assert [run.run_id for run in service.get_status().device_classes[0].runs] == [kept] +@pytest.mark.parametrize( + ("idempotency_key", "lost", "status_code"), + [(None, "run", 404), ("lost-midway", "run", 404), (None, "file", 409)], + ids=["run-removed-without-key", "run-removed-with-key", "file-unusable-without-key"], +) +def test_a_run_lost_while_it_is_being_created_fails_the_request_with_its_items_queued( + client: TestClient, + admin_token: str, + calibration, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + idempotency_key: str | None, + lost: str, + status_code: int, +) -> None: + """Admitted, the items fail when they start: without a key as well, the request is no success.""" + queue, service = calibration + held = HeldEnqueue(monkeypatch) + responses = [] + body = {"graph": _graph(), "idempotency_key": idempotency_key} + request = threading.Thread(target=lambda: responses.append(_calibrate(client, admin_token, body))) + request.start() + assert held.entered.wait(10) + [run] = service.get_status().device_classes[0].runs + if lost == "run": + service.reset_run(run.run_id) + else: + (tmp_path / "working_memory" / f"{GPU_CLASS.id}.json").write_bytes(b'{"schema_version": 2}') + held.release.set() + request.join(timeout=30) + + [response] = responses + assert response.status_code == status_code, response.text + [item_id] = queue.get_queue_item_ids("default").item_ids + assert queue.get_queue_item(item_id).origin == calibration_origin(run.run_id) + + def test_the_status_carries_the_device_class_id(client: TestClient, admin_token: str, calibration) -> None: status = client.get("/api/v2/models/working_memory/status", headers=_auth(admin_token)).json() diff --git a/tests/app/services/session_processor/test_working_memory_calibration_runner.py b/tests/app/services/session_processor/test_working_memory_calibration_runner.py index 7d15256c7c0..3fe38688d65 100644 --- a/tests/app/services/session_processor/test_working_memory_calibration_runner.py +++ b/tests/app/services/session_processor/test_working_memory_calibration_runner.py @@ -1,8 +1,9 @@ """The session runner's calibration hook: calibration items run their nodes inside a probe scope, have the windows -persisted, and are admitted again at session start; every other item runs exactly as before. +persisted -- invalid unless the node completed -- and are admitted again at session start; every other item runs exactly +as before. -The probe itself needs a GPU, so a test node plays its part: it appends a window to the scope it runs in, as an -estimated `model_on_device` would. +The probe itself needs a GPU, so a test node plays its part: it appends a valid window to the scope it runs in, as an +estimated `model_on_device` would when its operation finishes, and then completes, fails or is canceled. """ import asyncio @@ -159,17 +160,19 @@ def test_each_node_runs_in_its_own_scope_and_its_windows_are_stored( assert item.status == "completed" node_type = CalibrationProbeNodeInvocation.get_type() assert sorted(_SEEN) == [(run_id, item_id, "node0", node_type), (run_id, item_id, "node1", node_type)] - assert sorted((point["node_id"], point["run_id"]) for point in _stored_points(tmp_path)) == [ - ("node0", run_id), - ("node1", run_id), + assert sorted((point["node_id"], point["run_id"], point["valid"]) for point in _stored_points(tmp_path)) == [ + ("node0", run_id, True), + ("node1", run_id, True), ] assert not probe.armed -@pytest.mark.parametrize("behaviour", ["fail", "cancel"]) -def test_windows_are_stored_when_the_node_fails_or_is_canceled( - tmp_path: Path, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner, behaviour: str +@pytest.mark.parametrize(("behaviour", "reason"), [("fail", "exception"), ("cancel", "canceled")]) +def test_the_windows_of_a_node_that_fails_or_is_canceled_are_stored_invalid( + tmp_path: Path, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner, behaviour: str, reason: str ) -> None: + """The window closed before the node raised, so the probe saw nothing wrong; `run_node` swallows the failure or the + cancellation, and only the runner can tell the scope.""" session_runner, cancel_event = runner run_id, item_id = _create(service, "system", behaviour) @@ -181,9 +184,12 @@ def cancel() -> None: item = _run_next(queue, session_runner) assert item.status == ("failed" if behaviour == "fail" else "canceled") - assert [point["node_id"] for point in _stored_points(tmp_path)] == ["node0"] + [point] = _stored_points(tmp_path) + assert point["node_id"] == "node0" + assert (point["valid"], point["fit_eligible"], point["reasons"]) == (False, False, [reason]) [run] = service.get_status().device_classes[0].runs - assert run.status == ("failed" if behaviour == "fail" else "partial") + # Only valid points make a run that ended without completing partial. + assert run.status == ("failed" if behaviour == "fail" else "canceled") def test_an_item_whose_owner_is_not_an_admin_fails_before_any_node_runs( @@ -354,8 +360,8 @@ def clear() -> None: assert queue.get_queue_item_ids("default").item_ids == [] [run] = service.get_status().device_classes[0].runs - # The node's window was kept; the item, gone from the queue, will not finish. - assert (run.status, run.point_count, run.item_ids) == ("partial", 1, [item_id]) + # The node's window was kept, invalid since the node was canceled; the item, gone from the queue, will not finish. + assert (run.status, run.point_count, run.item_ids) == ("canceled", 1, [item_id]) def test_a_runner_failure_outside_the_nodes_is_recorded_as_failed( diff --git a/tests/app/services/session_queue/test_session_queue_writes.py b/tests/app/services/session_queue/test_session_queue_writes.py index dadb0cd5e5a..7accd02fcef 100644 --- a/tests/app/services/session_queue/test_session_queue_writes.py +++ b/tests/app/services/session_queue/test_session_queue_writes.py @@ -24,6 +24,7 @@ QueueItemStatusChangedEvent, ) from invokeai.app.services.invoker import Invoker +from invokeai.app.services.session_queue import session_queue_default from invokeai.app.services.session_queue.session_queue_common import Batch, EnqueueProjectNotFoundError from invokeai.app.services.session_queue.session_queue_default import SessionQueue from invokeai.app.services.shared.database.database import Database @@ -503,6 +504,32 @@ def fill(q: Queries) -> None: assert errors == [] assert queue.get_queue_status(QUEUE).pending == 2 + def test_an_all_or_none_enqueue_takes_no_place_once_its_items_no_longer_fit( + self, database: Database, events: TestEventService, accounts: None, monkeypatch: pytest.MonkeyPatch + ) -> None: + # Another admission takes two of three places between the enqueue's survey and its own admission, where one of + # its two items would still fit. Committed in between, the fill is seen only by the recount, on every backend. + queue = _service(database, events, max_queue_size=3) + prepare = session_queue_default.prepare_values_to_insert + + def fill(q: Queries) -> None: + q.locks.acquire(DatabaseLock.SESSION_QUEUE_ADMISSION) + for _ in range(2): + q.session_queue.insert_item(_pending_values()) + + def prepare_while_another_admits(*args: Any, **kwargs: Any) -> Any: + database.queries.run(fill) + return prepare(*args, **kwargs) + + monkeypatch.setattr(session_queue_default, "prepare_values_to_insert", prepare_while_another_admits) + batch = _batch(runs=2, key="whole") + + result = asyncio.run(queue.enqueue_batch(QUEUE, batch, False, "alice", all_or_none=True)) + + assert (result.requested, result.enqueued, result.item_ids) == (2, 0, []) + assert queue.get_queue_status(QUEUE).pending == 2 + assert queue.get_enqueue_receipt(QUEUE, "whole", "alice") is None + @pytest.mark.parametrize("lock", [DatabaseLock.SESSION_QUEUE_ADMISSION, DatabaseLock.MEDIA_PROTECTION]) def test_a_retry_waits_for_the_admission_and_the_cleanup( self, database: Database, queue: SessionQueue, lock: DatabaseLock diff --git a/tests/app/services/test_image_move_startup_safety.py b/tests/app/services/test_image_move_startup_safety.py index ba9213fd683..066ba352adb 100644 --- a/tests/app/services/test_image_move_startup_safety.py +++ b/tests/app/services/test_image_move_startup_safety.py @@ -1,73 +1,8 @@ from unittest.mock import MagicMock -from invokeai.app.services.invocation_services import InvocationServices -from invokeai.app.services.invoker import Invoker from invokeai.app.services.session_processor.session_processor_default import DefaultSessionProcessor -def _services(**overrides): - services = { - "board_image_records": object(), - "board_images": object(), - "board_records": object(), - "boards": object(), - "bulk_download": object(), - "configuration": object(), - "database": object(), - "events": object(), - "images": object(), - "image_files": object(), - "image_records": object(), - "logger": object(), - "model_images": object(), - "model_manager": object(), - "model_relationships": object(), - "model_relationship_records": object(), - "download_queue": object(), - "external_generation": object(), - "performance_statistics": object(), - "session_queue": object(), - "session_processor": object(), - "invocation_cache": object(), - "names": object(), - "urls": object(), - "workflow_records": object(), - "tensors": object(), - "conditioning": object(), - "style_preset_records": object(), - "style_preset_image_files": object(), - "workflow_thumbnails": object(), - "client_state_persistence": object(), - "project_records": object(), - "users": object(), - "wildcard_records": object(), - "videos": object(), - "video_files": object(), - "video_records": object(), - "board_video_records": object(), - "gallery": object(), - "image_index_records": object(), - "image_index": object(), - "intermediates": object(), - "image_moves": None, - "system_prompt_records": object(), - } - services.update(overrides) - return InvocationServices(**services) - - -def test_image_moves_start_before_session_processor() -> None: - started: list[str] = [] - image_moves = MagicMock() - image_moves.start.side_effect = lambda _invoker: started.append("image_moves") - session_processor = MagicMock() - session_processor.start.side_effect = lambda _invoker: started.append("session_processor") - - Invoker(_services(image_moves=image_moves, session_processor=session_processor)) - - assert started == ["image_moves", "session_processor"] - - def test_session_processor_detects_active_image_move_maintenance() -> None: image_moves = MagicMock() image_moves.is_maintenance_active.return_value = True diff --git a/tests/app/services/test_invoker.py b/tests/app/services/test_invoker.py new file mode 100644 index 00000000000..0f040f765af --- /dev/null +++ b/tests/app/services/test_invoker.py @@ -0,0 +1,108 @@ +"""The invoker's service lifecycle: the session processor starts after every other service and stops before them.""" + +import asyncio +import inspect +import time +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from invokeai.app.invocations.math import AddInvocation +from invokeai.app.services.invocation_services import InvocationServices +from invokeai.app.services.invoker import Invoker +from invokeai.app.services.session_processor.session_processor_default import DefaultSessionProcessor +from invokeai.app.services.session_queue.session_queue_default import SessionQueue +from invokeai.app.services.shared.database.database import Database +from invokeai.app.services.shared.graph import Graph +from invokeai.app.services.working_memory_calibration.working_memory_calibration_common import ( + WorkingMemoryCalibrationRequest, +) +from invokeai.backend.util.device_pool import GENERATION_DEVICE_POOL +from tests.fixtures.working_memory_calibration import calibration_service + + +def test_the_session_processor_starts_after_and_stops_before_every_other_service() -> None: + calls: list[tuple[str, str]] = [] + + class Service: + def __init__(self, name: str) -> None: + self.name = name + + def start(self, invoker: Invoker) -> None: + calls.append(("start", self.name)) + + def stop(self, invoker: Invoker) -> None: + calls.append(("stop", self.name)) + + names = list(inspect.signature(InvocationServices).parameters) + invoker = Invoker(InvocationServices(**{name: Service(name) for name in names})) + started = [name for action, name in calls if action == "start"] + calls.clear() + invoker.stop() + stopped = [name for action, name in calls if action == "stop"] + + assert sorted(started) == sorted(stopped) == sorted(names) + assert started[-1] == stopped[0] == "session_processor" + + +class _ProcessorThatRunsTheResumedItemFirst(DefaultSessionProcessor): + """Starts like the real processor, then returns only once the item pending at startup has run: a worker that is + quicker than the rest of startup, which services started after the processor would not keep up with.""" + + def __init__(self, queue: SessionQueue, item_id: int) -> None: + super().__init__() + self._resumed_queue = queue + self._resumed_item_id = item_id + + def start(self, invoker: Invoker) -> None: + super().start(invoker) + deadline = time.monotonic() + 30 + while self._resumed_queue.get_queue_item(self._resumed_item_id).status in ("pending", "in_progress"): + assert time.monotonic() < deadline, "the item pending at startup did not run" + time.sleep(0.01) + + +def test_a_calibration_item_pending_at_startup_runs_once_every_service_started( + tmp_path: Path, mock_invoker: Invoker, mock_sqlite_database: Database, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + "invokeai.app.services.session_processor.session_processor_default.paged_bytes", lambda device: None + ) + services = mock_invoker.services + services.configuration.generation_devices = ["cpu"] + services.configuration.offload_on_architecture_switch = False + services.model_manager = MagicMock() + services.model_manager.store.get_model.side_effect = lambda key: SimpleNamespace(key=key, hash=None) + + # The previous process enqueued a calibration run and stopped before its item ran. + previous_queue = SessionQueue(mock_sqlite_database) + services.session_queue = previous_queue + previous_queue.start(mock_invoker) + creator = calibration_service(tmp_path) + creator.start(mock_invoker) + graph = Graph() + graph.add_node(AddInvocation(id="add", a=1, b=2)) + created = asyncio.run(creator.create_run("model-key", WorkingMemoryCalibrationRequest(graph=graph), "system")) + [item_id] = created.item_ids + + # This process: new services over the same database and root, started by its invoker. + queue = SessionQueue(mock_sqlite_database) + calibration = calibration_service(tmp_path) + processor = _ProcessorThatRunsTheResumedItemFirst(queue, item_id) + services.session_queue, services.working_memory_calibration = queue, calibration + services.session_processor = processor + invoker = Invoker(services) + try: + item = queue.get_queue_item(item_id) + assert item.status == "completed", item.error_message + [run] = calibration.get_status().device_classes[0].runs + assert (run.run_id, run.status) == (created.run_id, "completed") + finally: + invoker.stop() + for worker in processor._workers: + assert worker.thread is not None + worker.thread.join(timeout=5) + # The processor registered its devices and a release listener in the process-wide pool. + GENERATION_DEVICE_POOL.reset() diff --git a/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py b/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py index 87e92bea720..4402002be7d 100644 --- a/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py +++ b/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py @@ -2,34 +2,48 @@ import asyncio import json +import logging import os import threading +import time from pathlib import Path from types import SimpleNamespace +from typing import Optional from unittest.mock import MagicMock import pytest from invokeai.app.invocations.math import AddInvocation from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_queue.session_queue_common import BatchDatum +from invokeai.app.services.session_queue.session_queue_common import ( + Batch, + BatchDatum, + EnqueueIdempotencyConflictError, + SessionQueueItem, +) from invokeai.app.services.session_queue.session_queue_default import SessionQueue from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph from invokeai.app.services.working_memory_calibration import working_memory_calibration_store from invokeai.app.services.working_memory_calibration.working_memory_calibration_common import ( + CALIBRATION_ORIGIN_PREFIX, CalibrationDeviceClass, + CalibrationQueueFullError, CalibrationStoreError, CalibrationStoreUnusableError, WorkingMemoryCalibrationRequest, + WorkingMemoryCalibrationRunCreated, WorkingMemoryCalibrationStatus, + calibration_run_id, ) from invokeai.app.services.working_memory_calibration.working_memory_calibration_default import ( WorkingMemoryCalibrationService, ) +from tests.fixtures.sqlite_database import sqlite_cursor from tests.fixtures.working_memory_calibration import ( GPU_CLASS, OTHER_CLASS, + HeldEnqueue, calibration_service, measured_window, ) @@ -52,15 +66,41 @@ def _service(root: Path, invoker: Invoker, **kwargs) -> WorkingMemoryCalibration return service -def _request(values: tuple[int, ...] = (1,)) -> WorkingMemoryCalibrationRequest: - graph = Graph() - graph.add_node(AddInvocation(id="add", a=1, b=2)) +def _request( + values: tuple[int, ...] = (1,), *, key: Optional[str] = None, b: int = 2, graph_id: Optional[str] = None +) -> WorkingMemoryCalibrationRequest: + # Without an id the graph gets a fresh one, as a client that sends none does. + graph = Graph() if graph_id is None else Graph(id=graph_id) + graph.add_node(AddInvocation(id="add", a=1, b=b)) data = [[BatchDatum(node_path="add", field_name="a", items=list(values))]] if len(values) > 1 else None - return WorkingMemoryCalibrationRequest(graph=graph, data=data) + return WorkingMemoryCalibrationRequest(graph=graph, data=data, idempotency_key=key) + + +def _create( + service: WorkingMemoryCalibrationService, + model_key: str = "model-a", + values: tuple[int, ...] = (1,), + *, + key: Optional[str] = None, +): + return asyncio.run(service.create_run(model_key, _request(values, key=key), "system")) + + +def _calibration_item_ids(queue: SessionQueue) -> list[int]: + return sorted(queue.get_queue_item_ids("default", origin_prefix=CALIBRATION_ORIGIN_PREFIX).item_ids) + +def _refuse_writes(source: str, target: str) -> None: + raise OSError("disk full") -def _create(service: WorkingMemoryCalibrationService, model_key: str = "model-a", values: tuple[int, ...] = (1,)): - return asyncio.run(service.create_run(model_key, _request(values), "system")) + +async def _dequeued(queue: SessionQueue) -> SessionQueueItem: + """The next item, once the queue has one: what the session processor claims as soon as an enqueue commits.""" + deadline = time.monotonic() + 10 + while (item := queue.dequeue()) is None: + assert time.monotonic() < deadline, "nothing was enqueued" + await asyncio.sleep(0.01) + return item def _store_path(root: Path, device_class_id: str = GPU_CLASS.id) -> Path: @@ -321,6 +361,9 @@ def test_run_status_follows_its_items_retries_and_observations( service.observe(retried, created.run_id) assert _run_status(service.get_status(), created.run_id).status == "canceled" + # Only a valid point measured something: a canceled operation's own window does not make the run partial. + service.record_windows(created.run_id, [measured_window(created.run_id, retry, reasons=("canceled",))]) + assert _run_status(service.get_status(), created.run_id).status == "canceled" service.record_windows(created.run_id, [measured_window(created.run_id, first)]) assert _run_status(service.get_status(), created.run_id).status == "partial" @@ -364,3 +407,310 @@ def test_status_reports_machines_that_cannot_calibrate( service = _service(tmp_path, mock_invoker, devices=lambda: [torch.device("cpu")]) assert "GPU memory" in (service.get_status().unavailable_reason or "") + + +def test_a_request_reusing_a_key_mid_enqueue_is_refused_without_taking_the_other_run( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + """Two requests share a retry key with different graphs. The second commits while the first is still enqueuing: + the first is refused and removes only the run it recorded.""" + service = _service(tmp_path, mock_invoker) + held = HeldEnqueue(monkeypatch) + + async def race(): + first = asyncio.create_task(service.create_run("model-a", _request(key="shared", b=2), "system")) + assert await asyncio.to_thread(held.entered.wait, 10) + second = await service.create_run("model-a", _request(key="shared", b=5), "system") + held.release.set() + with pytest.raises(EnqueueIdempotencyConflictError): + await first + return second + + second = asyncio.run(race()) + + [item_id] = second.item_ids + assert _calibration_item_ids(queue) == [item_id] + [run] = service.get_status().device_classes[0].runs + assert (run.run_id, run.item_ids) == (second.run_id, [item_id]) + item = queue.dequeue() + assert item is not None and service.admit(item, second.run_id) is None + + +def test_a_request_canceled_while_enqueuing_keeps_the_run_of_the_items_it_enqueued( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + """The client went away while the enqueue ran in its worker thread, which committed afterwards.""" + service = _service(tmp_path, mock_invoker) + held = HeldEnqueue(monkeypatch) + + async def cancel_then_retry(): + creating = asyncio.create_task(service.create_run("model-a", _request(key="flaky-connection"), "system")) + assert await asyncio.to_thread(held.entered.wait, 10) + creating.cancel() + with pytest.raises(asyncio.CancelledError): + await creating + held.release.set() + item = await _dequeued(queue) + refusal = await asyncio.to_thread(service.admit, item, calibration_run_id(item.origin)) + # The client retries the lost request. + retry = await service.create_run("model-a", _request(key="flaky-connection"), "system") + return item, refusal, retry + + item, refusal, retry = asyncio.run(cancel_then_retry()) + + assert refusal is None + assert (retry.run_id, retry.item_ids) == (calibration_run_id(item.origin), [item.item_id]) + assert _calibration_item_ids(queue) == [item.item_id] + [run] = service.get_status().device_classes[0].runs + assert (run.run_id, run.item_ids, run.status) == (retry.run_id, [item.item_id], "running") + + +def _collect_receipt(queue: SessionQueue, database: Database, key: str) -> None: + """The client acknowledged the key, the retention passed, and an admission collected the receipt.""" + queue.acknowledge_enqueue("default", key, "system") + with sqlite_cursor(database) as cursor: + cursor.execute( + "UPDATE session_queue_enqueue_receipts SET acknowledged_at = '2000-01-01' WHERE idempotency_key = ?;", + (key,), + ) + collect = Batch(graph=_request().graph, origin="canvas", idempotency_key=f"collects-{key}") + asyncio.run(queue.enqueue_batch("default", collect, False, "system")) + assert queue.get_enqueue_receipt("default", key, "system") is None + + +@pytest.mark.parametrize( + ("other_model", "other_values"), [("model-b", (1, 2)), ("model-a", (1, 3))], ids=["another-model", "other-data"] +) +def test_a_retry_key_names_one_request_across_models_batch_data_and_receipt_expiry( + tmp_path: Path, + mock_invoker: Invoker, + queue: SessionQueue, + mock_sqlite_database: Database, + other_model: str, + other_values: tuple[int, ...], +) -> None: + service = _service(tmp_path, mock_invoker) + first = _create(service, "model-a", (1, 2), key="calibrate") + + # The same key for another model or other batch data is another request, which the key's receipt refuses. + with pytest.raises(EnqueueIdempotencyConflictError): + _create(service, other_model, other_values, key="calibrate") + assert [run.run_id for run in service.get_status().device_classes[0].runs] == [first.run_id] + + _collect_receipt(queue, mock_sqlite_database, "calibrate") + + # A late repeat of the first request still answers with its run and enqueues nothing. + assert _create(service, "model-a", (1, 2), key="calibrate") == first + assert _calibration_item_ids(queue) == first.item_ids + # The other request now starts a run of its own. + other = _create(service, other_model, other_values, key="calibrate") + runs = {run.run_id: run for run in service.get_status().device_classes[0].runs} + assert other.run_id != first.run_id and runs.keys() == {first.run_id, other.run_id} + assert (runs[first.run_id].model_key, runs[first.run_id].item_ids) == ("model-a", first.item_ids) + assert (runs[other.run_id].model_key, runs[other.run_id].item_ids) == (other_model, other.item_ids) + assert _calibration_item_ids(queue) == sorted(first.item_ids + other.item_ids) + + +def test_a_batch_the_queue_cannot_take_whole_starts_no_run( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue +) -> None: + """A run with only some of its cases would measure only some of them, so it gets all of its items or none.""" + service = _service(tmp_path, mock_invoker) + mock_invoker.services.configuration.max_queue_size = 1 + + with pytest.raises(CalibrationQueueFullError, match="all 2 items"): + _create(service, values=(1, 2), key="two-sizes") + + assert queue.get_queue_item_ids("default").item_ids == [] + assert service.get_status().device_classes[0].runs == [] + # The refusal left the key unused: with room for both items, the same request starts the run. + mock_invoker.services.configuration.max_queue_size = 2 + created = _create(service, values=(1, 2), key="two-sizes") + assert len(created.item_ids) == 2 and _calibration_item_ids(queue) == created.item_ids + + +def test_a_request_canceled_while_enqueuing_into_a_queue_that_filled_up_leaves_no_run( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + """The creation goes on after the cancellation, finds no room left, and removes the run it recorded.""" + service = _service(tmp_path, mock_invoker) + held = HeldEnqueue(monkeypatch) + mock_invoker.services.configuration.max_queue_size = 1 + + async def cancel_fill_then_repeat(): + creating = asyncio.create_task(service.create_run("model-a", _request(key="lost"), "system")) + assert await asyncio.to_thread(held.entered.wait, 10) + creating.cancel() + with pytest.raises(asyncio.CancelledError): + await creating + # Another run takes the last place while the canceled one waits between survey and admission. + other = await service.create_run("model-b", _request(), "system") + held.release.set() + # A repeat waits for the creation still in flight, then finds no room either. + with pytest.raises(CalibrationQueueFullError): + await service.create_run("model-a", _request(key="lost"), "system") + return other + + other = asyncio.run(cancel_fill_then_repeat()) + + assert [run.run_id for run in service.get_status().device_classes[0].runs] == [other.run_id] + assert _calibration_item_ids(queue) == other.item_ids + + +def test_an_enqueue_that_raised_after_its_commit_answers_with_the_queued_items( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + """The queue announces a batch after committing it: a failed announcement leaves the batch queued.""" + service = _service(tmp_path, mock_invoker) + + def announce(*args: object, **kwargs: object) -> None: + raise RuntimeError("the event bus is gone") + + monkeypatch.setattr(mock_invoker.services.events, "emit_batch_enqueued", announce) + created = _create(service) + + [item_id] = created.item_ids + assert _calibration_item_ids(queue) == [item_id] + [run] = service.get_status().device_classes[0].runs + assert (run.run_id, run.item_ids) == (created.run_id, [item_id]) + item = queue.dequeue() + assert item is not None and service.admit(item, created.run_id) is None + + +def test_a_repeat_of_a_request_waits_for_the_attempt_in_flight( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + service = _service(tmp_path, mock_invoker) + held = HeldEnqueue(monkeypatch) + + async def repeat_while_in_flight(): + first = asyncio.create_task(service.create_run("model-a", _request(key="slow"), "system")) + assert await asyncio.to_thread(held.entered.wait, 10) + repeat = asyncio.create_task(service.create_run("model-a", _request(key="slow"), "system")) + # Half a second, as `tests/fixtures/races.waiting` gives a change that must wait. + done, _ = await asyncio.wait({repeat}, timeout=0.5) + held.release.set() + return bool(done), await first, await repeat + + repeat_finished_early, first, repeat = asyncio.run(repeat_while_in_flight()) + + assert not repeat_finished_early, "the repeat did not wait for the attempt in flight" + assert repeat == first and _calibration_item_ids(queue) == first.item_ids + + +def _create_while_recording_fails( + service: WorkingMemoryCalibrationService, monkeypatch: pytest.MonkeyPatch, request: WorkingMemoryCalibrationRequest +) -> WorkingMemoryCalibrationRunCreated: + """Create a run whose store starts refusing writes once the run is recorded: recording its items fails.""" + held = HeldEnqueue(monkeypatch) + replace = os.replace + + async def create(): + creating = asyncio.create_task(service.create_run("model-a", request, "system")) + assert await asyncio.to_thread(held.entered.wait, 10) + monkeypatch.setattr(os, "replace", _refuse_writes) + held.release.set() + try: + return await creating + finally: + monkeypatch.setattr(os, "replace", replace) + + return asyncio.run(create()) + + +def test_a_request_without_a_key_answers_with_its_queued_items_when_the_store_fails_to_record_them( + tmp_path: Path, + mock_invoker: Invoker, + queue: SessionQueue, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """The items run all the same; sending the request again would queue the batch a second time.""" + service = _service(tmp_path, mock_invoker) + + with caplog.at_level(logging.WARNING): + created = _create_while_recording_fails(service, monkeypatch, _request()) + + [item_id] = created.item_ids + assert _calibration_item_ids(queue) == [item_id] + assert any(record.levelno == logging.WARNING and created.run_id in record.getMessage() for record in caplog.records) + [run] = service.get_status().device_classes[0].runs + assert (run.run_id, run.item_ids, run.status) == (created.run_id, [item_id], "running") + + +def test_a_repeat_answers_with_every_item_the_keys_receipt_names( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + """One of the unrecorded items left the queue meanwhile. The receipt still names it, so the run cannot read as + completed once the others have run.""" + service = _service(tmp_path, mock_invoker) + with pytest.raises(CalibrationStoreError): + _create_while_recording_fails(service, monkeypatch, _request((1, 2, 3), key="three-sizes")) + queued = _calibration_item_ids(queue) + queue.delete_queue_items_by_id(queued[:1]) + + repeat = _create(service, values=(1, 2, 3), key="three-sizes") + + assert len(queued) == 3 and repeat.item_ids == queued + assert _calibration_item_ids(queue) == queued[1:] + [stored] = json.loads(_store_path(tmp_path).read_bytes())["runs"] + assert stored["queue_item_ids"] == queued + + +def test_a_repeat_takes_over_the_queued_items_once_the_keys_receipt_is_collected( + tmp_path: Path, + mock_invoker: Invoker, + queue: SessionQueue, + mock_sqlite_database: Database, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Enqueuing the batch again would queue it a second time.""" + service = _service(tmp_path, mock_invoker) + with pytest.raises(CalibrationStoreError): + _create_while_recording_fails(service, monkeypatch, _request(key="unrecorded")) + [item_id] = _calibration_item_ids(queue) + _collect_receipt(queue, mock_sqlite_database, "unrecorded") + + repeat = _create(service, key="unrecorded") + + assert repeat.item_ids == [item_id] and _calibration_item_ids(queue) == [item_id] + [stored] = json.loads(_store_path(tmp_path).read_bytes())["runs"] + assert (stored["run_id"], stored["queue_item_ids"]) == (repeat.run_id, [item_id]) + + +def test_a_failed_enqueue_answers_with_its_own_error_when_the_queue_cannot_be_read( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + """Whether anything of the run was queued cannot be told then, so the run is kept.""" + service = _service(tmp_path, mock_invoker) + first = _create(service, key="taken") + + def unreadable(*args: object, **kwargs: object) -> None: + raise RuntimeError("the database is gone") + + with monkeypatch.context() as patch: + patch.setattr(queue, "get_queue_item_ids", unreadable) + with pytest.raises(EnqueueIdempotencyConflictError): + _create(service, "model-b", key="taken") + + runs = {run.run_id: run.model_key for run in service.get_status().device_classes[0].runs} + assert len(runs) == 2 and runs[first.run_id] == "model-a" and "model-b" in runs.values() + + +def test_requests_that_differ_only_in_their_graph_id_are_one_request( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue +) -> None: + """A client that sends no graph id gets a fresh one with each request.""" + service = _service(tmp_path, mock_invoker) + + def create(graph_id: Optional[str]): + return asyncio.run(service.create_run("model-a", _request(key="one-request", graph_id=graph_id), "system")) + + first = create(None) + assert create("chosen-by-the-client") == first + # With the run reset and its items gone from the queue, only the key's receipt knows the request. It matches the + # whole batch, graph id included. + service.reset_run(first.run_id) + queue.delete_queue_items_by_id(first.item_ids) + assert create(None) == first + assert _calibration_item_ids(queue) == [] diff --git a/tests/app/test_api_dependencies.py b/tests/app/test_api_dependencies.py index 186b4659eaf..3a2fcef2b54 100644 --- a/tests/app/test_api_dependencies.py +++ b/tests/app/test_api_dependencies.py @@ -1,14 +1,25 @@ import asyncio from logging import getLogger +from typing import Optional from unittest.mock import MagicMock +import pytest +import torch + from invokeai.app.api import dependencies from invokeai.app.api.dependencies import ApiDependencies from invokeai.app.services.config.config_default import InvokeAIAppConfig +from invokeai.app.services.image_index.image_index_default import ImageIndexService +from invokeai.app.services.invoker import Invoker +from invokeai.app.services.session_processor import session_processor_default from invokeai.app.services.shared.database.database import Database +from invokeai.backend.util.device_pool import GENERATION_DEVICE_POOL +from invokeai.backend.util.devices import TorchDevice -def test_initialize_does_not_vacuum_database(monkeypatch, tmp_path) -> None: +def _skip_process_wide_setup(monkeypatch: pytest.MonkeyPatch) -> None: + """The architecture registry, the attention warm-up, the JWT secret and the model manager: process-wide or heavy, + and not what these tests check.""" monkeypatch.setattr(dependencies, "validate_architectures", lambda: None) monkeypatch.setattr(dependencies, "warm_up_attention", lambda *_: None) monkeypatch.setattr(dependencies, "set_jwt_secret", lambda _: None) @@ -18,6 +29,10 @@ def test_initialize_does_not_vacuum_database(monkeypatch, tmp_path) -> None: MagicMock(return_value=MagicMock()), ) + +def test_initialize_does_not_vacuum_database(monkeypatch, tmp_path) -> None: + _skip_process_wide_setup(monkeypatch) + original_init_database = dependencies.init_database initialized_databases: list[Database] = [] @@ -58,3 +73,54 @@ def initialize_database(*args, **kwargs) -> Database: database.dispose() ApiDependencies.database = None loop.close() + + +def test_the_generation_devices_are_registered_before_any_service_starts(monkeypatch, tmp_path) -> None: + """Work outside the session queue borrows an idle GPU from the moment its service starts -- the image index's first + embedding batch runs before the session processor, which starts last -- and keeps it across the processor's start.""" + _skip_process_wide_setup(monkeypatch) + # Two GPUs, on a machine that may have none: a worker pins its GPU in CUDA only once it claims a queue item, and + # the empty queue gives it none. + gpus = [torch.device("cuda:0"), torch.device("cuda:1")] + monkeypatch.setattr(TorchDevice, "get_generation_devices", classmethod(lambda cls, generation_devices: list(gpus))) + # Process-wide torch state, which the processor sets for more than one CUDA worker. + monkeypatch.setattr(session_processor_default, "disable_conv_benchmark_empty_cache", lambda: None) + borrowed: list[Optional[torch.device]] = [] + image_index_start = ImageIndexService.start + + def start_and_borrow(self: ImageIndexService, invoker: Invoker) -> None: + image_index_start(self, invoker) + borrowed.append(GENERATION_DEVICE_POOL.try_borrow_off_queue("cuda")) + + monkeypatch.setattr(ImageIndexService, "start", start_and_borrow) + config = InvokeAIAppConfig( + use_memory_db=True, + image_index_enabled=False, + node_cache_size=0, + scan_models_on_startup=False, + ) + config._root = tmp_path + loop = asyncio.new_event_loop() + original_invoker = getattr(ApiDependencies, "invoker", None) + + try: + ApiDependencies.initialize(config=config, event_handler_id=0, loop=loop, logger=getLogger("test")) + + assert borrowed == [gpus[0]] + assert GENERATION_DEVICE_POOL.is_lent(gpus[0]), "the processor's start kept the borrow" + GENERATION_DEVICE_POOL.release_borrow(gpus[0]) + finally: + invoker = getattr(ApiDependencies, "invoker", None) + if invoker is not None and invoker is not original_invoker: + ApiDependencies.shutdown() + for worker in invoker.services.session_processor._workers: + assert worker.thread is not None + worker.thread.join(timeout=5) + if original_invoker is None: + delattr(ApiDependencies, "invoker") + else: + ApiDependencies.invoker = original_invoker + ApiDependencies.database = None + # The processor registered the devices and a release listener in the process-wide pool. + GENERATION_DEVICE_POOL.reset() + loop.close() diff --git a/tests/backend/working_memory/test_counters.py b/tests/backend/working_memory/test_counters.py index 8e7fe4684c3..92dfd2f491a 100644 --- a/tests/backend/working_memory/test_counters.py +++ b/tests/backend/working_memory/test_counters.py @@ -1,7 +1,8 @@ """The device counters: which driver reader each platform uses, the DRM fdinfo reader, and the MPS peak sampler. The OS and torch readers behind them are stand-ins: what is under test is that each platform reads through its own -reader, labels the reading with it, and turns any failure into "unknown". +reader, labels the reading with it, and turns any failure into "unknown" -- and that a peak reset or a synchronization +that fails raises, since it has no "unknown" to fall back on. """ import sys @@ -126,6 +127,36 @@ def test_a_failing_reader_reads_as_unknown_and_never_raises( ) +@pytest.mark.parametrize( + ("device", "module", "function"), + [ + pytest.param(CUDA0, torch.cuda, "reset_peak_memory_stats", id="cuda-reset"), + pytest.param(CUDA0, torch.cuda, "synchronize", id="cuda-synchronize"), + pytest.param(torch.device("xpu", 0), torch.xpu, "reset_peak_memory_stats", id="xpu-reset"), + pytest.param(torch.device("xpu", 0), torch.xpu, "synchronize", id="xpu-synchronize"), + pytest.param(torch.device("mps"), torch.mps, "synchronize", id="mps-synchronize"), + ], +) +def test_a_peak_reset_or_synchronization_that_fails_raises( + monkeypatch: pytest.MonkeyPatch, device: torch.device, module: object, function: str +): + """A window that went on would measure from an earlier operation's peaks, or with that work still in flight.""" + monkeypatch.setattr(module, function, _raise) + # What the MPS peak sampler reads while it runs. + monkeypatch.setattr(torch.mps, "current_allocated_memory", lambda: 1 * GIB) + monkeypatch.setattr(torch.mps, "driver_allocated_memory", lambda: 2 * GIB) + + device_counters = counters.DeviceCounters(device) + try: + with pytest.raises(RuntimeError, match="the driver did not answer"): + if function == "synchronize": + device_counters.synchronize() + else: + device_counters.reset_peaks() + finally: + device_counters.close() + + def test_xpu_reads_the_device_wide_usage_as_a_guard_not_as_this_processs_driver(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(torch.xpu, "memory_stats", lambda device: dict(_STATS)) monkeypatch.setattr(counters, "xpu_memory_info", lambda device: (10 * GIB, 16 * GIB)) diff --git a/tests/backend/working_memory/test_probe.py b/tests/backend/working_memory/test_probe.py index 18f89e62d43..ef326b397cd 100644 --- a/tests/backend/working_memory/test_probe.py +++ b/tests/backend/working_memory/test_probe.py @@ -3,13 +3,14 @@ `FakeDevice` stands in for a GPU with a caching allocator: live tensors (allocated), the blocks the allocator holds (reserved = allocated + cached), what the driver counts, peaks since the last reset, and a release that returns the cached blocks when `TorchDevice.empty_cache` runs -- as a real device does. Each test scripts allocations through it and -reads the `MeasuredWindow` the probe recorded; expected values follow from the script alone. +reads the `MeasuredWindow` the probe recorded; expected values follow from the script alone. The probe reads it through +`FakeCounters`, or with `cuda_device` through the real `counters.DeviceCounters` and torch's CUDA functions. """ import sys import threading from contextlib import contextmanager -from typing import Iterator, Optional +from typing import Callable, Iterator, Optional import pytest import torch @@ -70,6 +71,21 @@ def free(self, n: int) -> None: self.allocated -= n self.cached += n + def reset_peaks(self) -> None: + self.peak_reserved = self.reserved + self.peak_allocated = self.allocated + + def memory_stats(self) -> dict[str, int]: + """What `torch.cuda.memory_stats` reports for this device.""" + return { + "reserved_bytes.all.current": self.reserved, + "allocated_bytes.all.current": self.allocated, + "reserved_bytes.all.peak": self.peak_reserved, + "allocated_bytes.all.peak": self.peak_allocated, + "num_ooms": self.num_ooms, + "num_alloc_retries": self.num_alloc_retries, + } + def reading(self) -> CounterReading: if self.fail_next_read: self.fail_next_read = False @@ -115,8 +131,7 @@ def read(self) -> CounterReading: return self.device.reading() def reset_peaks(self) -> None: - self.device.peak_reserved = self.device.reserved - self.device.peak_allocated = self.device.allocated + self.device.reset_peaks() def synchronize(self) -> None: pass @@ -198,6 +213,24 @@ def device_counters(torch_device: torch.device) -> FakeCounters: yield fake assert all(c.closed for c in opened), "every window's counters are closed, so no sampler outlives it" assert probe.armed is False + assert probe.locking == [], "every lock leaves `locking` once it has settled or failed" + + +_DEVICE_COUNTERS = counters.DeviceCounters + + +@pytest.fixture +def cuda_device(device: FakeDevice, monkeypatch: pytest.MonkeyPatch) -> FakeDevice: + """`device` as cuda:0 behind the real `counters.DeviceCounters`: only torch's allocator statistics, peak reset and + synchronization, and the driver's figure, are the fake's.""" + monkeypatch.setattr(counters, "DeviceCounters", _DEVICE_COUNTERS) + monkeypatch.setattr(torch.cuda, "memory_stats", lambda _device: device.memory_stats()) + monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", lambda _device: device.reset_peaks()) + monkeypatch.setattr(torch.cuda, "synchronize", lambda _device: None) + monkeypatch.setattr( + counters, "_driver_reading", lambda _device: (device.reserved + device.retained, 0, None, "windows-d3dkmt") + ) + return device def _estimate(reservation: int = 2 * GIB) -> WorkingMemoryEstimate: @@ -518,6 +551,123 @@ def opening(torch_device: torch.device) -> FakeCounters: assert (window.valid, window.reasons) == (False, ["foreign_lock"]) +class GatedCache(FakeCache): + """A cache whose lock of the model `gated` holds still twice: before it moves anything, as a lock that waits for + the cache's lock or for MODEL_LOAD_LOCK does, and once it moved the weights, before it settles, as a paced lock does + between passes.""" + + def __init__(self, device: FakeDevice, gated: str): + super().__init__(device) + self.gated = gated + self.waiting, self.go, self.moved, self.settle = (threading.Event() for _ in range(4)) + + def lock(self, record: CacheRecord, working_mem_bytes: Optional[int], max_move_bytes: Optional[int] = None) -> bool: + if record.key != self.gated: + return super().lock(record, working_mem_bytes, max_move_bytes) + self.waiting.set() + if not self.go.wait(timeout=30): + raise TimeoutError("the test never let the lock go on") + settled = super().lock(record, working_mem_bytes, max_move_bytes) + self.moved.set() + if not self.settle.wait(timeout=30): + raise TimeoutError("the test never let the lock settle") + return settled + + +@pytest.mark.parametrize("lock_read_armed", ["before-the-scope-armed", "before-the-window-opened"]) +def test_a_lock_another_thread_has_in_progress_when_the_window_opens_voids_it(device: FakeDevice, lock_read_armed: str): + """Read before the scope armed, `armed` says nothing is measured, and the lock never reports. Read after it, the lock + reports its start while no window is registered and its end once the window has closed. Either way its weights + move while the window measures.""" + cache = GatedCache(device, gated="llm") + vae = _model(cache, "vae") + llm = _model(cache, "llm", total_bytes=2 * GIB) + + def expand_prompt() -> None: + with llm.model_on_device(): + pass + + thread = threading.Thread(target=expand_prompt) + + def lock_until_it_waits() -> None: + thread.start() + assert cache.waiting.wait(timeout=30), "the lock read `armed` and waits for the cache" + + try: + if lock_read_armed == "before-the-scope-armed": + lock_until_it_waits() + with _scope() as scope: + if lock_read_armed == "before-the-window-opened": + lock_until_it_waits() + with vae.model_on_device(working_memory=_estimate()): + cache.go.set() + assert cache.moved.wait(timeout=30), "the weights moved onto the device inside the window" + cache.settle.set() + thread.join(timeout=30) + assert not thread.is_alive() + finally: + cache.go.set() + cache.settle.set() + if thread.ident is not None: + thread.join() + + (window,) = scope.windows + assert (window.valid, window.reasons) == (False, ["foreign_lock"]) + assert cache.locks_held() == 0 + + +class GatedAppend(list): + """`probe.locking`, whose append on the thread named `gated` waits until the test lets it go on.""" + + def __init__(self, gated: str): + super().__init__() + self.gated = gated + self.waiting, self.go = threading.Event(), threading.Event() + + def append(self, cache: object) -> None: + if threading.current_thread().name == self.gated: + self.waiting.set() + if not self.go.wait(timeout=30): + raise TimeoutError("the test never let the lock go on") + super().append(cache) + + +def test_a_lock_enters_locking_before_it_reads_armed(device: FakeDevice, monkeypatch: pytest.MonkeyPatch): + """The other way round, a lock could read `armed` before the scope armed and enter `locking` only after the window + opened: neither would see the other, and its weights would move while the window measures.""" + locking = GatedAppend(gated="expand-prompt") + monkeypatch.setattr(probe, "locking", locking) + cache = FakeCache(device) + vae = _model(cache, "vae") + llm = _model(cache, "llm", total_bytes=2 * GIB) + locked, release = threading.Event(), threading.Event() + + def expand_prompt() -> None: + with llm.model_on_device(): + locked.set() + release.wait(timeout=30) + + thread = threading.Thread(target=expand_prompt, name="expand-prompt") + try: + thread.start() + assert locking.waiting.wait(timeout=30), "the lock is entering `locking`" + with _scope() as scope: + with vae.model_on_device(working_memory=_estimate()): + locking.go.set() + assert locked.wait(timeout=30), "the weights moved onto the device inside the window" + # Held until the window has closed, so that its unlock tells the window nothing. + release.set() + thread.join(timeout=30) + assert not thread.is_alive() + finally: + locking.go.set() + release.set() + thread.join() + + (window,) = scope.windows + assert (window.valid, window.reasons) == (False, ["foreign_lock"]) + + def test_a_model_this_thread_pins_in_ram_is_its_own_lock(device: FakeDevice): """LoRAs are pinned in RAM before the transformer locks; that pin belongs to this thread, not to another.""" cache = FakeCache(device) @@ -597,6 +747,53 @@ def test_a_segment_boundary_that_fails_voids_the_window(device: FakeDevice): assert (window.valid, window.reasons) == (False, ["probe_error"]) +def _works(times: int, then: Callable[[], None]) -> Callable[..., None]: + """A torch call that does `then` `times` times and raises from then on, as on a device context gone bad.""" + left = [times] + + def call(*_args: object) -> None: + if left[0] == 0: + raise RuntimeError("CUDA error: an illegal memory access was encountered") + left[0] -= 1 + then() + + return call + + +_UNOPENED = ["probe_error", "counters_unavailable"] +"""A window that could not open: the operation ran unmeasured.""" + + +@pytest.mark.parametrize( + ("function", "works", "reasons"), + [ + pytest.param("synchronize", 0, _UNOPENED, id="synchronize-at-open"), + pytest.param("reset_peak_memory_stats", 0, _UNOPENED, id="reset-at-open"), + pytest.param("reset_peak_memory_stats", 1, ["probe_error"], id="reset-after-a-release"), + pytest.param("synchronize", 1, ["probe_error"], id="synchronize-at-close"), + ], +) +def test_a_window_whose_peaks_cannot_be_reset_or_device_synchronized_is_invalid( + cuda_device: FakeDevice, monkeypatch: pytest.MonkeyPatch, function: str, works: int, reasons: list[str] +): + """Peaks not reset carry what came before into the window; unsynchronized work lands in the wrong segment.""" + then = cuda_device.reset_peaks if function == "reset_peak_memory_stats" else (lambda: None) + monkeypatch.setattr(torch.cuda, function, _works(works, then)) + cache = FakeCache(cuda_device) + + with _scope() as scope: + with _model(cache, "vae", device="cuda:0").model_on_device(working_memory=_estimate()): + cuda_device.alloc(2 * GIB) + cuda_device.free(2 * GIB) + TorchDevice.empty_cache() + cuda_device.alloc(1 * GIB) + cuda_device.free(1 * GIB) + + (window,) = scope.windows + assert (window.valid, window.fit_eligible, window.reasons) == (False, False, reasons) + assert cache.locks_held() == 0 + + def test_a_child_window_that_fails_to_close_is_recorded_invalid_and_the_parent_resumes(device: FakeDevice): cache = FakeCache(device) diff --git a/tests/fixtures/working_memory_calibration.py b/tests/fixtures/working_memory_calibration.py index 0f4c5d1d317..8bcf26d2395 100644 --- a/tests/fixtures/working_memory_calibration.py +++ b/tests/fixtures/working_memory_calibration.py @@ -5,12 +5,15 @@ them, without measuring anything. """ +import threading from collections.abc import Mapping from pathlib import Path -from typing import Callable, Optional +from typing import Any, Callable, Optional +import pytest import torch +from invokeai.app.services.session_queue import session_queue_default from invokeai.app.services.working_memory_calibration.working_memory_calibration_default import ( SettingValue, WorkingMemoryCalibrationService, @@ -30,6 +33,24 @@ def fake_environment(device: torch.device, settings: Mapping[str, SettingValue]) return Environment(hard=hard, soft={"driver": "1.0", "os": "test"}) +class HeldEnqueue: + """Holds the next enqueue between its survey of the queue and its admission until released, as a busy worker + thread would. Later enqueues pass.""" + + def __init__(self, monkeypatch: pytest.MonkeyPatch) -> None: + self.entered = threading.Event() + self.release = threading.Event() + prepare = session_queue_default.prepare_values_to_insert + + def held(*args: Any, **kwargs: Any): + monkeypatch.setattr(session_queue_default, "prepare_values_to_insert", prepare) + self.entered.set() + assert self.release.wait(timeout=10) + return prepare(*args, **kwargs) + + monkeypatch.setattr(session_queue_default, "prepare_values_to_insert", held) + + def calibration_service( root: Path, devices: Optional[Callable[[], list[torch.device]]] = None, From b01244400b26a79f6106cc6e8988b13fa3c497e6 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 21 Oct 2026 06:16:30 +0200 Subject: [PATCH 6/9] fix(working-memory): wait for session workers on stop and fence calibration writes behind resets Stop joins the workers within a shared 30 s deadline and refuses a restart while one still runs; the cache counts locks apart from its records. A model reset raises a per-model epoch the runner reads at node start, keyed runs leave a tombstone on reset, retries keep their lineage before admission, and a recorded receipt replays without the model. --- .../Architecture/model-manager.mdx | 18 +- .../development/Architecture/overview.mdx | 2 +- invokeai/app/api/dependencies.py | 5 + invokeai/app/api/routers/working_memory.py | 9 +- invokeai/app/services/invoker.py | 5 +- .../session_processor_default.py | 93 ++++++++- .../session_queue/session_queue_base.py | 7 + .../session_queue/session_queue_common.py | 9 + .../session_queue/session_queue_default.py | 11 +- .../shared/database/queries/session_queue.py | 14 ++ .../working_memory_calibration_base.py | 20 +- .../working_memory_calibration_common.py | 19 +- .../working_memory_calibration_default.py | 188 ++++++++++++------ .../load/model_cache/cache_record.py | 6 +- .../load/model_cache/model_cache.py | 33 ++- invokeai/frontend/api/openapi.json | 6 +- invokeai/frontend/api/schema.ts | 9 +- .../test_working_memory_calibration_runner.py | 39 +++- .../test_session_queue_query_plans.py | 11 + tests/app/services/test_invoker.py | 183 ++++++++++++++++- ...test_working_memory_calibration_service.py | 136 ++++++++++++- tests/app/test_api_dependencies.py | 32 +++ .../test_model_cache_locks_held.py | 77 ++++++- 23 files changed, 808 insertions(+), 124 deletions(-) diff --git a/docs/src/content/docs/development/Architecture/model-manager.mdx b/docs/src/content/docs/development/Architecture/model-manager.mdx index 1cf3041c696..970bfeb6454 100644 --- a/docs/src/content/docs/development/Architecture/model-manager.mdx +++ b/docs/src/content/docs/development/Architecture/model-manager.mdx @@ -1040,9 +1040,9 @@ A calibration run measures one model's operations on this machine's GPU. Runs on 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:` 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. 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`. + - 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. 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. + - 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. @@ -1050,7 +1050,7 @@ A calibration run measures one model's operations on this machine's GPU. Runs on - 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. If they cannot be saved, the item fails and the session stops. Every other item gets no scope and runs as before. +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 `/working_memory/.json` (`CalibrationStoreFile`, `schema_version` 1). It is machine data, so it lives next to the database rather than inside it. It holds: @@ -1058,7 +1058,11 @@ The store has one JSON file per device class at `/working_memory/ HTTPException: response_model=WorkingMemoryCalibrationRunCreated, responses={ 404: { - "description": "The model could not be found, or the run was removed while it was being created; its " - "items are queued then and fail when they start" + "description": "The model could not be found; the run was removed while it was being created, whose " + "items are queued then and fail when they start; or a reset removed the run the idempotency key started" }, 409: { "description": "This machine cannot calibrate now (no GPU, or several generation devices), the device " @@ -122,8 +122,9 @@ async def calibrate_model_working_memory( With a key, repeating the request (also after a lost response) records the queued items and returns the run; it enqueues nothing new. The items come from the queue's enqueue receipt for the key while it exists, otherwise from - the queue items that carry the run's origin and those the runner saw. Acknowledge the key through the queue's - enqueue_batch/acknowledge route once the response arrived.""" + the queue items that carry the run's origin and those the runner saw. Once a reset removed the run, a repeat is + refused (404). Acknowledge the key through the queue's enqueue_batch/acknowledge route once the response + arrived.""" await asyncio.to_thread(assert_image_move_maintenance_inactive) try: return await _service().create_run(key, request, current_admin.user_id) diff --git a/invokeai/app/services/invoker.py b/invokeai/app/services/invoker.py index 17594fb266e..8c01125a05b 100644 --- a/invokeai/app/services/invoker.py +++ b/invokeai/app/services/invoker.py @@ -38,8 +38,9 @@ def _start(self) -> None: def stop(self) -> None: """Stops the invoker. A new invoker will have to be created to execute further.""" - # The session processor is signalled first, so that its workers stop taking queue items and cancel the sessions - # they run before the services those sessions use stop. Its stop() does not wait for the workers. + # The session processor stops first: its workers stop taking queue items and cancel the sessions they run, and + # its stop() waits for them to exit (bounded, see DefaultSessionProcessor.stop) before the services those + # sessions use stop. self.__stop_service(self.services.session_processor) for service in self._services_but_the_session_processor(): self.__stop_service(service) diff --git a/invokeai/app/services/session_processor/session_processor_default.py b/invokeai/app/services/session_processor/session_processor_default.py index b8d0328e4ba..b89b7517f64 100644 --- a/invokeai/app/services/session_processor/session_processor_default.py +++ b/invokeai/app/services/session_processor/session_processor_default.py @@ -67,6 +67,27 @@ OWNER_LOOKUP_ATTEMPTS = 3 OWNER_LOOKUP_RETRY_SECONDS = 0.25 +# How long `DefaultSessionProcessor.stop` waits, in all, for its workers to leave their sessions. A node that checks for +# cancellation stops within one step; the bound also covers most nodes without such a check (a large VAE decode), and +# keeps a node that never returns from holding up shutdown for good. +WORKER_STOP_TIMEOUT_SECONDS = 30.0 +# After this long, `stop` says which worker it is waiting for: the wait runs on the event-loop thread at shutdown. +WORKER_STOP_NOTICE_SECONDS = 1.0 + +# Workers a stop gave up on. Process-wide, like the device pool whose locks they may still hold: no processor of this +# process starts, and none replaces the pool, while one of them runs. +_abandoned_workers: list[Thread] = [] + + +def refuse_while_an_abandoned_worker_runs() -> None: + """Raise RuntimeError while a worker an earlier `DefaultSessionProcessor.stop` gave up on still runs.""" + _abandoned_workers[:] = [thread for thread in _abandoned_workers if thread.is_alive()] + if _abandoned_workers: + raise RuntimeError( + f"Session worker {', '.join(thread.name for thread in _abandoned_workers)} of the previous session " + "processor is still running a node; it would share its GPU with the new workers. Restart the process." + ) + def queue_owner_is_active( services: InvocationServices, queue_item: SessionQueueItem, *, unreadable_is_active: bool = False @@ -316,6 +337,13 @@ def _run_calibration_node(self, invocation: BaseInvocation, queue_item: SessionQ """Run one node of a calibration item inside a probe scope and persist its windows, whether the node completed, failed or was canceled; unless it completed, they are persisted invalid. Returns False when they could not be persisted: the item has failed then.""" + service = self._services.working_memory_calibration + try: + assert service is not None + # Read before the node runs: a model reset while it runs removes points these windows must not bring back. + model_reset_epochs = service.model_reset_epochs(run_id) + except Exception as error: + return self._fail_calibration_node(invocation, queue_item, error) with probe.calibration_scope( run_id=run_id, queue_item_id=queue_item.item_id, @@ -330,20 +358,24 @@ def _run_calibration_node(self, invocation: BaseInvocation, queue_item: SessionQ scope.invalidate("exception" if self._node_outcome == "failed" else "canceled") if not scope.windows: return True - service = self._services.working_memory_calibration try: - assert service is not None - service.record_windows(run_id, scope.windows) + service.record_windows(run_id, scope.windows, model_reset_epochs) except Exception as error: - message = f"Could not save the working-memory measurements of node {invocation.id}: {error}" - self._services.logger.error(message, exc_info=True) - with suppress(SessionQueueItemNotFoundError): - self._services.session_queue.fail_queue_item( - queue_item.item_id, "WorkingMemoryCalibrationError", message, traceback.format_exc() - ) - return False + return self._fail_calibration_node(invocation, queue_item, error) return True + def _fail_calibration_node( + self, invocation: BaseInvocation, queue_item: SessionQueueItem, error: Exception + ) -> bool: + """Fail the item whose node's measurements cannot be saved. Call from the `except` block (for the traceback).""" + message = f"Could not save the working-memory measurements of node {invocation.id}: {error}" + self._services.logger.error(message, exc_info=True) + with suppress(SessionQueueItemNotFoundError): + self._services.session_queue.fail_queue_item( + queue_item.item_id, "WorkingMemoryCalibrationError", message, traceback.format_exc() + ) + return False + def run_node(self, invocation: BaseInvocation, queue_item: SessionQueueItem): try: # Any unhandled exception in this scope is an invocation error & will fail the graph @@ -822,7 +854,10 @@ def register_generation_devices(self, configuration: InvokeAIAppConfig) -> list[ `generation_devices` (by default `"auto"`, every available GPU) resolves to one normalized device per worker, or to a single `None` (one worker on the device the global config chooses) when the resolution is empty, as with an empty list. + + Raises RuntimeError while a worker an earlier `stop` gave up on still runs (see `stop`). """ + refuse_while_an_abandoned_worker_runs() devices = TorchDevice.get_generation_devices(configuration.generation_devices) self._devices = list(devices) if devices else [None] GENERATION_DEVICE_POOL.set_generation_devices([d for d in self._devices if d is not None]) @@ -854,6 +889,9 @@ def _clone_session_runner(self, template: SessionRunnerBase) -> SessionRunnerBas ) def start(self, invoker: Invoker) -> None: + """Start one worker per generation device. Raises RuntimeError while a worker an earlier `stop` gave up on + still runs (see `stop`).""" + refuse_while_an_abandoned_worker_runs() self._invoker: Invoker = invoker self._resume_event = ThreadEvent() @@ -932,6 +970,15 @@ def start(self, invoker: Invoker) -> None: worker.thread.start() def stop(self, *args, **kwargs) -> None: + """Stop the workers and wait for them to exit, up to `WORKER_STOP_TIMEOUT_SECONDS` for all of them together. + + The invoker stops every other service once this returns, so a worker that exits in time never runs beside a + stopped service. A wait longer than `WORKER_STOP_NOTICE_SECONDS` names the worker and its queue item in the + log. A worker still in a node at the deadline is given up on, with an error in the log, so that shutdown + finishes: a node without a cancel point may never return. A later stop does not wait for it again. It keeps + its GPU, so until it exits no session processor of this process starts and none replaces the generation device + pool. The services it uses stop under it; once the database is disposed, its writes are refused. + """ self._stop_event.set() # Cancel any in-progress generation so that long-running nodes (e.g. denoising) stop at # the next step boundary instead of running to completion. Without this, a generation @@ -943,6 +990,32 @@ def stop(self, *args, **kwargs) -> None: self._poll_now_event.set() self._resume_event.set() + deadline = time.monotonic() + WORKER_STOP_TIMEOUT_SECONDS + for worker in self._workers: + # A worker given up on by an earlier stop gets no second wait. + if worker.thread is None or worker.thread in _abandoned_workers: + continue + worker.thread.join(timeout=max(0.0, min(deadline - time.monotonic(), WORKER_STOP_NOTICE_SECONDS))) + remaining = deadline - time.monotonic() + if worker.thread.is_alive() and remaining > 0: + self._invoker.services.logger.info( + f"Waiting up to {remaining:.0f}s for the {self._describe_running(worker)} to stop" + ) + worker.thread.join(timeout=remaining) + if worker.thread.is_alive(): + _abandoned_workers.append(worker.thread) + self._invoker.services.logger.error( + f"The {self._describe_running(worker)} did not stop within {WORKER_STOP_TIMEOUT_SECONDS:.0f}s; " + "shutting down without it. Restart the process before starting Invoke again." + ) + + @staticmethod + def _describe_running(worker: _SessionWorker) -> str: + assert worker.thread is not None + queue_item = worker.queue_item + item = f" (queue item {queue_item.item_id})" if queue_item is not None else "" + return f"session worker {worker.thread.name} on {worker.label}{item}" + def _poll_now(self) -> None: self._poll_now_event.set() diff --git a/invokeai/app/services/session_queue/session_queue_base.py b/invokeai/app/services/session_queue/session_queue_base.py index 2414e846fdf..edf6dc85836 100644 --- a/invokeai/app/services/session_queue/session_queue_base.py +++ b/invokeai/app/services/session_queue/session_queue_base.py @@ -21,6 +21,7 @@ NodeFieldValue, PruneResult, QueueItemDeliveryResult, + QueueItemLineage, QueueItemOwner, RetryItemsResult, SessionQueueCountsByDestination, @@ -292,6 +293,12 @@ def get_queue_item_ids( matching items only when no limit is given.""" pass + @abstractmethod + def get_queue_item_lineage(self, queue_id: str, origin_prefix: str) -> list[QueueItemLineage]: + """The status, origin and retry source of every item whose origin starts with `origin_prefix` (ignoring + case), in no order.""" + pass + @abstractmethod def get_queue_item_summaries_by_ids(self, queue_id: str, item_ids: list[int]) -> list[SessionQueueItemSummary]: """Gets lightweight queue item summaries in the requested item ID order.""" diff --git a/invokeai/app/services/session_queue/session_queue_common.py b/invokeai/app/services/session_queue/session_queue_common.py index 31fc0ada27c..458f043af4d 100644 --- a/invokeai/app/services/session_queue/session_queue_common.py +++ b/invokeai/app/services/session_queue/session_queue_common.py @@ -216,6 +216,15 @@ def validate_graph(cls, v: Graph): QUEUE_ITEM_STATUS = Literal["pending", "in_progress", "waiting", "completed", "failed", "canceled"] +class QueueItemLineage(NamedTuple): + """An item's status and origin, and the item it retried: a retry of a retry names the first item as well.""" + + item_id: int + status: QUEUE_ITEM_STATUS + origin: Optional[str] + retried_from_item_id: Optional[int] + + class ItemIdsResult(BaseModel): """Response containing ordered item ids with metadata for optimistic updates.""" diff --git a/invokeai/app/services/session_queue/session_queue_default.py b/invokeai/app/services/session_queue/session_queue_default.py index c33dfbdfc70..33f8968838b 100644 --- a/invokeai/app/services/session_queue/session_queue_default.py +++ b/invokeai/app/services/session_queue/session_queue_default.py @@ -3,7 +3,7 @@ import json import math from datetime import datetime, timedelta, timezone -from typing import Any, Literal, Optional +from typing import Any, Literal, Optional, cast from pydantic_core import to_jsonable_python @@ -32,6 +32,7 @@ NodeFieldValue, PruneResult, QueueItemDeliveryResult, + QueueItemLineage, QueueItemOwner, RetriedQueueItem, RetryItemsResult, @@ -1043,6 +1044,14 @@ def get_queue_item_ids( def has_active_queue_work(self) -> bool: return self._queries.session_queue.has_active_work() + def get_queue_item_lineage(self, queue_id: str, origin_prefix: str) -> list[QueueItemLineage]: + return [ + QueueItemLineage(item_id, cast(QUEUE_ITEM_STATUS, status), origin, retried_from_item_id) + for item_id, status, origin, retried_from_item_id in self._queries.session_queue.lineage( + queue_id, origin_prefix + ) + ] + def get_queue_item_summaries_by_ids(self, queue_id: str, item_ids: list[int]) -> list[SessionQueueItemSummary]: if not item_ids: return [] diff --git a/invokeai/app/services/shared/database/queries/session_queue.py b/invokeai/app/services/shared/database/queries/session_queue.py index 120ee0c38af..0e8056d09aa 100644 --- a/invokeai/app/services/shared/database/queries/session_queue.py +++ b/invokeai/app/services/shared/database/queries/session_queue.py @@ -114,6 +114,10 @@ def _origin_pattern(origin_prefix: Optional[str]) -> Optional[str]: return like_prefix(origin_prefix) if origin_prefix is not None else None +# The origin is tested on the listing index; only the matching items' rows are read. +_LINEAGE = select(_Q.item_id, _Q.status, _Q.origin, _Q.retried_from_item_id).where(_IN_QUEUE, _origin_matches()) + + @functools.cache def _item_by_status(status: str, by_origin: bool, by_owner: bool) -> Select[Any]: """The pending item to run next, or an item in progress (the earliest enqueued of several).""" @@ -249,6 +253,10 @@ def _summaries(rows: Sequence[Sequence[Any]]) -> list[dict[str, Any]]: return [dict(zip(_SUMMARY_NAMES, row, strict=True)) for row in rows] +def _lineage(rows: Sequence[Sequence[Any]]) -> list[tuple[int, str, Optional[str], Optional[int]]]: + return [(row[0], row[1], row[2], row[3]) for row in rows] + + def _counts(rows: Sequence[Sequence[Any]]) -> dict[str, int]: return {str(row[0]): int(row[1]) for row in rows} @@ -841,6 +849,12 @@ def item_ids( } return list(conn.execute(statement, parameters).scalars().all()) + @mapped(_lineage) + @read + def lineage(self, conn: Connection, queue_id: str, origin_prefix: str) -> Sequence[Row[Any]]: + """Id, status, origin and retry source of the queue's items with the origin, in no order.""" + return conn.execute(_LINEAGE, {"queue_id": queue_id, "origin_pattern": like_prefix(origin_prefix)}).all() + @mapped(_summaries) @read def summaries(self, conn: Connection, queue_id: str, item_ids: Sequence[int]) -> Sequence[Row[Any]]: diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py index 68a52bafbb9..63207fd6c78 100644 --- a/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py @@ -1,4 +1,5 @@ from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import Optional from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem @@ -21,15 +22,25 @@ async def create_run( """Validate the request, record the run and enqueue its batch as `user_id`, who must be an administrator. The queue takes all of the batch's items or none (`CalibrationQueueFullError`). With a retry key, a repeat of - the request returns the run it created and enqueues nothing.""" + the request returns the run it created and enqueues nothing; once a reset removed the run, it raises + `CalibrationRunNotFoundError`.""" @abstractmethod def admit(self, queue_item: SessionQueueItem, run_id: str) -> Optional[str]: """Called by the session runner before a calibration item runs: None to run it, or why it must fail.""" @abstractmethod - def record_windows(self, run_id: str, windows: list[MeasuredWindow]) -> None: - """Persist one node's measured windows. Windows of a run that no longer exists are dropped.""" + def model_reset_epochs(self, run_id: str) -> dict[str, int]: + """The reset epoch of every model reset on the run's device class; empty when the run does not exist. Read + when a node starts and passed to `record_windows`.""" + + @abstractmethod + def record_windows( + self, run_id: str, windows: list[MeasuredWindow], model_reset_epochs: Optional[Mapping[str, int]] = None + ) -> None: + """Persist one node's measured windows. Windows of a run that no longer exists are dropped, and so are those + of a model reset since `model_reset_epochs` was read (by default, at this call): the reset removed the + model's points, these among them. A reset by another process sharing the root counts as well.""" @abstractmethod def observe(self, queue_item: SessionQueueItem, run_id: str, error: Optional[BaseException] = None) -> None: @@ -42,7 +53,8 @@ def get_status(self) -> WorkingMemoryCalibrationStatus: @abstractmethod def reset_run(self, run_id: str) -> WorkingMemoryCalibrationResetResult: - """Remove a run and its points.""" + """Remove a run and its points. Every reset keeps the ids of the runs it removed: a repeat of the request that + started one cannot start it again.""" @abstractmethod def reset_model(self, model_key: str, device_class_id: Optional[str] = None) -> WorkingMemoryCalibrationResetResult: diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py index ed4ad8a8b16..e5ee262ca87 100644 --- a/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py @@ -131,6 +131,16 @@ class CalibrationStoreFile(_StoreModel): runs: list[CalibrationRun] = Field(default_factory=list) points: list[CalibrationPoint] = Field(default_factory=list) envs: list[CalibrationEnvironment] = Field(default_factory=list) + removed_run_ids: list[str] = Field( + default_factory=list, + description="The runs with a retry key that a reset removed, kept for good. A repeat of the request that " + "started one cannot start it again, so its queue items keep failing at admission", + ) + model_reset_epochs: dict[str, int] = Field( + default_factory=dict, + description="Per model key, how often the model was reset. A node's points of a model reset after the node " + "started are dropped", + ) # endregion @@ -151,8 +161,8 @@ class WorkingMemoryCalibrationRequest(BaseModel): min_length=1, max_length=255, description="A caller-owned retry key. Repeating the request with it returns the run it created and enqueues " - "nothing new. Reusing it for another request (another model, graph or batch data) is rejected while the " - "queue keeps the key's enqueue receipt.", + "nothing new; once a reset removed that run, the repeat is refused. Reusing the key for another request " + "(another model, graph or batch data) is rejected while the queue keeps the key's enqueue receipt.", ) @@ -246,12 +256,13 @@ class CalibrationGraphError(ValueError): class CalibrationRunNotFoundError(KeyError): - def __init__(self, run_id: str) -> None: + def __init__(self, run_id: str, reason: str = "does not exist") -> None: super().__init__(run_id) self.run_id = run_id + self.reason = reason def __str__(self) -> str: - return f"Calibration run {self.run_id} does not exist" + return f"Calibration run {self.run_id} {self.reason}" class CalibrationQueueFullError(Exception): diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py index 6138fb3f086..5f59115cefa 100644 --- a/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py @@ -16,6 +16,7 @@ DEFAULT_QUEUE_ID, QUEUE_ITEM_STATUS, Batch, + QueueItemLineage, SessionQueueItem, SessionQueueItemNotFoundError, ) @@ -78,17 +79,25 @@ """The InvokeAI settings that change how much working memory an operation takes; hard facts of a measurement's environment.""" -_SUMMARY_CHUNK: Final = 500 -"""Queue items read per summary query when the status collects the calibration items.""" - _RUN_ID_NAMESPACE: Final = uuid.UUID("3c1f6a52-8d0e-4b6f-9a57-2f0d5c8e41b3") """Derives a run id from an owner, a retry key and the request's fingerprint.""" _RUNNING: Final[frozenset[QUEUE_ITEM_STATUS]] = frozenset({"pending", "in_progress", "waiting"}) +_RESET: Final = "was reset; a repeat of the request that started it cannot start it again" + SettingValue = str | int | float | bool | None +def _is_keyed(run_id: str) -> bool: + """Whether a retry key named the run: its id is then derived (a UUID 5, see `_identify`), otherwise fresh + (`uuid_string`, a UUID 4) and never repeated.""" + try: + return uuid.UUID(run_id).version == 5 + except ValueError: + return False + + def _now() -> datetime.datetime: return datetime.datetime.now(datetime.timezone.utc) @@ -104,38 +113,55 @@ def _prune_envs(file: CalibrationStoreFile) -> None: def _add_points( - file: CalibrationStoreFile, *, run_id: str, env: Environment, points: list[CalibrationPoint] -) -> tuple[bool, bool]: - """Attach `points` to their run; nothing when the file does not hold the run.""" + file: CalibrationStoreFile, + *, + run_id: str, + env: Environment, + points: list[CalibrationPoint], + epochs: Mapping[str, int], +) -> tuple[bool, int]: + """Attach `points` to their run, except those of a model whose reset epoch moved past `epochs`: a reset removed + the model's points after the points were measured. How many it attached; none when the file lacks the run.""" if all(run.run_id != run_id for run in file.runs): - return False, False + return False, 0 + current = file.model_reset_epochs + kept = [point for point in points if current.get(point.model_key, 0) <= epochs.get(point.model_key, 0)] + if not kept: + return False, 0 _add_env(file, env) - file.points.extend(points) - return True, True + file.points.extend(kept) + return True, len(kept) def derive_run_status( - run: CalibrationRun, live: Mapping[int, QUEUE_ITEM_STATUS], valid_point_count: int + run: CalibrationRun, live: Mapping[int, QueueItemLineage], valid_point_count: int ) -> tuple[CalibrationRunStatus, Optional[str], list[int]]: - """A run's status, error and items, from its items' statuses: `live` as the queue has them now, the run's own - observations for items the queue no longer has. An item the queue no longer has counts as canceled unless the - runner saw it finish. A retry stands in for the item it retried. A run that ended without completing is partial - only with valid points: an invalid one (a canceled or failed operation, a spill) measured nothing usable.""" + """A run's status, error and items, from its items: `live` as the queue has them now, the run's own observations + for items the queue no longer has. An item the queue no longer has counts as canceled unless the runner saw it + finish. A retry stands in for the item it retried, also before the runner started it. A run that ended without + completing is partial only with valid points: an invalid one (a canceled or failed operation, a spill) measured + nothing usable.""" item_ids = sorted(set(run.queue_item_ids) | set(run.observed) | set(live)) def status_of(item_id: int) -> QUEUE_ITEM_STATUS: if item_id in live: - return live[item_id] + return live[item_id].status observed = run.observed.get(item_id) # Gone from the queue while the runner last saw it running (a crash, a cleared queue): it will not finish. if observed is None or observed.status in _RUNNING: return "canceled" return observed.status + def retried_from(item_id: int) -> Optional[int]: + if item_id in live: + return live[item_id].retried_from_item_id + observed = run.observed.get(item_id) + return observed.retried_from_item_id if observed is not None else None + latest: dict[int, int] = {} for item_id in item_ids: - observed = run.observed.get(item_id) - lineage = observed.retried_from_item_id if observed and observed.retried_from_item_id else item_id + source = retried_from(item_id) + lineage = item_id if source is None else source latest[lineage] = max(latest.get(lineage, item_id), item_id) effective = sorted(latest.values()) statuses = {item_id: status_of(item_id) for item_id in effective} @@ -401,32 +427,26 @@ def _prepare_run( run_id=run_id, device_class_id=device_class_id, item_ids=list(run.queue_item_ids) ) # An earlier attempt may have queued the batch without recording its items. While the queue keeps the - # key's receipt, the enqueue below answers with the items it names, all of them. Without the receipt, - # enqueuing again would queue the batch a second time: the items carrying the run's origin, and those the - # runner saw, are taken over instead. - if self._services.session_queue.get_enqueue_receipt(DEFAULT_QUEUE_ID, key, user_id) is None: - item_ids = sorted(set(self._queued_item_ids(run_id)) | set(run.observed)) - if item_ids: - self._record_items(device_class_id, run_id, item_ids) - return WorkingMemoryCalibrationRunCreated( - run_id=run_id, device_class_id=device_class_id, item_ids=item_ids - ) + # key's receipt, enqueuing the batch again answers with the items it names, all of them. That needs + # neither the model, which may have been removed since, nor the device. Without the receipt, enqueuing + # again would queue the batch a second time: the items carrying the run's origin, and those the runner + # saw, are taken over instead. + if self._services.session_queue.get_enqueue_receipt(DEFAULT_QUEUE_ID, key, user_id) is not None: + return self._batch(request, run_id), device_class_id, False + item_ids = sorted(set(self._queued_item_ids(run_id)) | set(run.observed)) + if item_ids: + self._record_items(device_class_id, run_id, item_ids) + return WorkingMemoryCalibrationRunCreated( + run_id=run_id, device_class_id=device_class_id, item_ids=item_ids + ) + elif key is not None and self._was_reset(run_id): + raise CalibrationRunNotFoundError(run_id, _RESET) device = self._calibration_device() # Raises UnknownModelException for an unknown key. config = self._services.model_manager.store.get_model(model_key) - try: - batch = Batch( - graph=request.graph, - data=request.data, - runs=1, - origin=calibration_origin(run_id), - destination=CALIBRATION_DESTINATION, - idempotency_key=request.idempotency_key, - ) - except ValueError as error: - raise CalibrationGraphError(str(error)) from error + batch = self._batch(request, run_id) if found is not None: - # The run is recorded without items: enqueue its batch, which the key's receipt answers while it exists. + # The run is recorded without items, and nothing of it is queued: enqueue its batch. return batch, found[0], False device_class = self._device_class(device) env = self._environment(device) @@ -440,6 +460,8 @@ def _prepare_run( ) def add(file: CalibrationStoreFile) -> tuple[bool, bool]: + if run_id in file.removed_run_ids: + raise CalibrationRunNotFoundError(run_id, _RESET) if any(known.run_id == run_id for known in file.runs): return False, False file.runs.append(run) @@ -450,15 +472,36 @@ def add(file: CalibrationStoreFile) -> tuple[bool, bool]: self._run_files[run_id] = device_class.id return batch, device_class.id, bool(created) + @staticmethod + def _batch(request: WorkingMemoryCalibrationRequest, run_id: str) -> Batch: + try: + return Batch( + graph=request.graph, + data=request.data, + runs=1, + origin=calibration_origin(run_id), + destination=CALIBRATION_DESTINATION, + idempotency_key=request.idempotency_key, + ) + except ValueError as error: + raise CalibrationGraphError(str(error)) from error + + def _was_reset(self, run_id: str) -> bool: + for device_class_id in self._store.device_class_ids(): + read = self._store.read(device_class_id) + if read is not None and read.file is not None and run_id in read.file.removed_run_ids: + return True + return False + def _queued_item_ids(self, run_id: str) -> list[int]: """The queue items carrying the run's origin: its batch's, with their retries and workflow-call children.""" queue = self._services.session_queue return sorted(queue.get_queue_item_ids(DEFAULT_QUEUE_ID, origin_prefix=calibration_origin(run_id)).item_ids) def _discard(self, run_id: str) -> None: - """Remove a run this request recorded and queued nothing for.""" + """Remove a run this request recorded and queued nothing for. Its key stays unused: a repeat may start it.""" try: - self.reset_run(run_id) + self._remove_run(run_id, remember=False) except Exception: self._services.logger.warning(f"Could not discard calibration run {run_id}", exc_info=True) @@ -557,7 +600,18 @@ def observe(self, queue_item: SessionQueueItem, run_id: str, error: Optional[Bas ), ) - def record_windows(self, run_id: str, windows: list[MeasuredWindow]) -> None: + def model_reset_epochs(self, run_id: str) -> dict[str, int]: + found = self._find_run(run_id) + if found is None: + return {} + read = self._store.read(found[0]) + return dict(read.file.model_reset_epochs) if read is not None and read.file is not None else {} + + def record_windows( + self, run_id: str, windows: list[MeasuredWindow], model_reset_epochs: Optional[Mapping[str, int]] = None + ) -> None: + if model_reset_epochs is None: + model_reset_epochs = self.model_reset_epochs(run_id) found = self._find_run(run_id) if found is None: self._services.logger.warning( @@ -581,28 +635,26 @@ def record_windows(self, run_id: str, windows: list[MeasuredWindow]) -> None: continue env = self._environment(device) points = [CalibrationPoint(**window.model_dump(), env_id=env.id) for window in group] - if not self._store.change(run_class_id, partial(_add_points, run_id=run_id, env=env, points=points)): + add = partial(_add_points, run_id=run_id, env=env, points=points, epochs=model_reset_epochs) + dropped = len(points) - (self._store.change(run_class_id, add) or 0) + if dropped: self._services.logger.warning( - f"Dropped {len(group)} working-memory measurements of calibration run {run_id}: the run was reset." + f"Dropped {dropped} working-memory measurements of calibration run {run_id}: the run, or the " + "model they measured, was reset." ) # endregion # region Status - def _live_items(self) -> dict[str, dict[int, QUEUE_ITEM_STATUS]]: + def _live_items(self) -> dict[str, dict[int, QueueItemLineage]]: """The queue's calibration items, per run, as the queue has them now.""" - queue = self._services.session_queue - item_ids = queue.get_queue_item_ids(DEFAULT_QUEUE_ID, origin_prefix=CALIBRATION_ORIGIN_PREFIX).item_ids - live: dict[str, dict[int, QUEUE_ITEM_STATUS]] = {} - for start in range(0, len(item_ids), _SUMMARY_CHUNK): - for summary in queue.get_queue_item_summaries_by_ids( - DEFAULT_QUEUE_ID, item_ids[start : start + _SUMMARY_CHUNK] - ): - # The queue's origin filter ignores case; only the exact prefix is the server's. - run_id = calibration_run_id(summary.origin) - if run_id is not None: - live.setdefault(run_id, {})[summary.item_id] = summary.status + live: dict[str, dict[int, QueueItemLineage]] = {} + for item in self._services.session_queue.get_queue_item_lineage(DEFAULT_QUEUE_ID, CALIBRATION_ORIGIN_PREFIX): + # The queue's origin filter ignores case; only the exact prefix is the server's. + run_id = calibration_run_id(item.origin) + if run_id is not None: + live.setdefault(run_id, {})[item.item_id] = item return live def get_status(self) -> WorkingMemoryCalibrationStatus: @@ -650,7 +702,7 @@ def _device_class_status( read: StoreRead, is_current: bool, current_env: Optional[Environment], - live: Mapping[str, Mapping[int, QUEUE_ITEM_STATUS]], + live: Mapping[str, Mapping[int, QueueItemLineage]], ) -> WorkingMemoryCalibrationDeviceClass: file = read.file if file is None: @@ -728,8 +780,14 @@ def _device_class_status( @staticmethod def _remove( - file: CalibrationStoreFile, run_ids: set[str], model_key: Optional[str] = None + file: CalibrationStoreFile, run_ids: set[str], model_key: Optional[str] = None, *, remember: bool = True ) -> tuple[bool, WorkingMemoryCalibrationResetResult]: + """Remove the runs and their points, and with `model_key` every point of that model. `remember` keeps the + removed runs' ids, so that a repeat of the request that started one cannot start it again.""" + if remember: + file.removed_run_ids.extend( + run.run_id for run in file.runs if run.run_id in run_ids and _is_keyed(run.run_id) + ) runs = [run for run in file.runs if run.run_id not in run_ids] points = [ point @@ -744,25 +802,33 @@ def _remove( return result.runs_removed > 0 or result.points_removed > 0, result def reset_run(self, run_id: str) -> WorkingMemoryCalibrationResetResult: + return self._remove_run(run_id, remember=True) + + def _remove_run(self, run_id: str, *, remember: bool) -> WorkingMemoryCalibrationResetResult: found = self._find_run(run_id) if found is None: raise CalibrationRunNotFoundError(run_id) - result = self._store.change(found[0], lambda file: self._remove(file, {run_id})) + result = self._store.change(found[0], lambda file: self._remove(file, {run_id}, remember=remember)) self._run_files.pop(run_id, None) return result or WorkingMemoryCalibrationResetResult(runs_removed=0, points_removed=0) def reset_model(self, model_key: str, device_class_id: Optional[str] = None) -> WorkingMemoryCalibrationResetResult: + class_id = self._current_device_class_id(device_class_id) + def remove(file: CalibrationStoreFile) -> tuple[bool, WorkingMemoryCalibrationResetResult]: + # Written with the removal or not at all: a node that read the older epoch drops its points of the model. + file.model_reset_epochs[model_key] = file.model_reset_epochs.get(model_key, 0) + 1 run_ids = {run.run_id for run in file.runs if run.model_key == model_key} - return self._remove(file, run_ids, model_key) + return True, self._remove(file, run_ids, model_key)[1] - result = self._store.change(self._current_device_class_id(device_class_id), remove) + result = self._store.change(class_id, remove) return result or WorkingMemoryCalibrationResetResult(runs_removed=0, points_removed=0) def reset_all(self, device_class_id: Optional[str] = None) -> WorkingMemoryCalibrationResetResult: def remove(file: CalibrationStoreFile) -> tuple[bool, WorkingMemoryCalibrationResetResult]: result = WorkingMemoryCalibrationResetResult(runs_removed=len(file.runs), points_removed=len(file.points)) changed = bool(file.runs or file.points or file.envs) + file.removed_run_ids.extend(run.run_id for run in file.runs if _is_keyed(run.run_id)) file.runs, file.points, file.envs = [], [], [] return changed, result diff --git a/invokeai/backend/model_manager/load/model_cache/cache_record.py b/invokeai/backend/model_manager/load/model_cache/cache_record.py index 14fb4879af4..5a30eb74712 100644 --- a/invokeai/backend/model_manager/load/model_cache/cache_record.py +++ b/invokeai/backend/model_manager/load/model_cache/cache_record.py @@ -114,9 +114,11 @@ def lock(self) -> None: self._locks += 1 def unlock(self) -> None: - """Unlock this record.""" + """Unlock this record. An unlock without a lock to release raises and changes nothing, so the count stays + right for the next lock and for the cache's `locks_held`.""" + if self._locks == 0: + raise RuntimeError(f"Model cache entry {self.key} was unlocked more often than it was locked") self._locks -= 1 - assert self._locks >= 0 @property def is_locked(self) -> bool: diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index 59d422ea1d9..8e2097095ec 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -630,6 +630,8 @@ def __init__( self._cached_models: Dict[str, CacheRecord] = {} self._cache_stack: List[str] = [] + # Locks on this cache's records, detached ones included (see `locks_held`); changed only with `_lock` held. + self._locks_held = 0 self._ram_cache_size_bytes = self._calc_ram_available_to_model_cache() @@ -1893,9 +1895,22 @@ def locks_held(self) -> int: """How many locks this cache's records hold right now: VRAM and RAM pins of every thread together. The working-memory probe compares it with the locks its own thread took, so that a window opening while - another thread holds a model on this device knows it shares the device. + another thread holds a model on this device knows it shares the device. Counted as locks are taken and + released, not over the cached records: a record dropped while locked (the VRAM-move error paths) still holds + its model on the device until its holder unlocks it. """ - return sum(record._locks for record in self._cached_models.values()) + return self._locks_held + + def _pin(self, cache_entry: CacheRecord) -> None: + """Lock `cache_entry` and count the lock in `locks_held`. Caller holds the lock.""" + cache_entry.lock() + self._locks_held += 1 + + def _unpin(self, cache_entry: CacheRecord) -> None: + """Release a lock `_pin` took. Caller holds the lock. The record refuses an unlock it has no lock for before + either count changes, so a double unlock leaves `locks_held` right.""" + cache_entry.unlock() + self._locks_held -= 1 @synchronized @record_activity @@ -1911,7 +1926,7 @@ def lock_in_ram(self, cache_entry: CacheRecord) -> None: "the RAM cache. This is a sign that the model loading order is non-optimal in the invocation code " "(See https://github.com/invoke-ai/InvokeAI/issues/7513)." ) - cache_entry.lock() + self._pin(cache_entry) cache_entry.awaiting_first_use = False @synchronized @@ -1939,7 +1954,7 @@ def lock( "(See https://github.com/invoke-ai/InvokeAI/issues/7513)." ) # cache_entry = self._cached_models[key] - cache_entry.lock() + self._pin(cache_entry) # End of the post-admission grace: from here the entry is pinned by its lock count, and # after the final unlock it is ordinary evictable cache content. cache_entry.awaiting_first_use = False @@ -1975,10 +1990,10 @@ def lock( ) except torch.OutOfMemoryError: self._logger.warning("Insufficient GPU memory to load model. Aborting") - cache_entry.unlock() + self._unpin(cache_entry) raise except Exception: - cache_entry.unlock() + self._unpin(cache_entry) raise self._log_cache_state() @@ -2023,10 +2038,10 @@ def continue_lock( ) except torch.OutOfMemoryError: self._logger.warning("Insufficient GPU memory to load model. Aborting") - cache_entry.unlock() + self._unpin(cache_entry) raise except Exception: - cache_entry.unlock() + self._unpin(cache_entry) raise if settled: @@ -2045,7 +2060,7 @@ def unlock(self, cache_entry: CacheRecord) -> None: "(See https://github.com/invoke-ai/InvokeAI/issues/7513)." ) # cache_entry = self._cached_models[key] - cache_entry.unlock() + self._unpin(cache_entry) self._logger.debug( f"Unlocked model {cache_entry.key} (Type: {cache_entry.cached_model.model.__class__.__name__})" ) diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json index ee6add31301..068b7b8f779 100644 --- a/invokeai/frontend/api/openapi.json +++ b/invokeai/frontend/api/openapi.json @@ -5664,7 +5664,7 @@ "post": { "tags": ["model_manager"], "summary": "Calibrate Model Working Memory", - "description": "Start a working-memory calibration run: enqueue the batch with every node uncached and measure each node's\nmodel operations on the generation device. Needs exactly one generation device, and it must be a GPU. The queue\ntakes all of the batch's items or none.\n\nOnce the items are queued, a request without an `idempotency_key` succeeds even when the store fails for a passing\nreason (an unwritable root, a lock timeout) to record them: they run all the same, and the status finds them by\ntheir origin. With a key, that failure is a 503 although the items are queued. A run removed meanwhile (404) or an\nunusable store file (409) fails the request with or without a key; its queued items then fail when they start.\n\nWith a key, repeating the request (also after a lost response) records the queued items and returns the run; it\nenqueues nothing new. The items come from the queue's enqueue receipt for the key while it exists, otherwise from\nthe queue items that carry the run's origin and those the runner saw. Acknowledge the key through the queue's\nenqueue_batch/acknowledge route once the response arrived.", + "description": "Start a working-memory calibration run: enqueue the batch with every node uncached and measure each node's\nmodel operations on the generation device. Needs exactly one generation device, and it must be a GPU. The queue\ntakes all of the batch's items or none.\n\nOnce the items are queued, a request without an `idempotency_key` succeeds even when the store fails for a passing\nreason (an unwritable root, a lock timeout) to record them: they run all the same, and the status finds them by\ntheir origin. With a key, that failure is a 503 although the items are queued. A run removed meanwhile (404) or an\nunusable store file (409) fails the request with or without a key; its queued items then fail when they start.\n\nWith a key, repeating the request (also after a lost response) records the queued items and returns the run; it\nenqueues nothing new. The items come from the queue's enqueue receipt for the key while it exists, otherwise from\nthe queue items that carry the run's origin and those the runner saw. Once a reset removed the run, a repeat is\nrefused (404). Acknowledge the key through the queue's enqueue_batch/acknowledge route once the response\narrived.", "operationId": "calibrate_model_working_memory", "security": [ { @@ -5707,7 +5707,7 @@ } }, "404": { - "description": "The model could not be found, or the run was removed while it was being created; its items are queued then and fail when they start" + "description": "The model could not be found; the run was removed while it was being created, whose items are queued then and fail when they start; or a reset removed the run the idempotency key started" }, "409": { "description": "This machine cannot calibrate now (no GPU, or several generation devices), the device class's store file is unusable and kept as it is (possibly after the items were queued, which then fail when they start), the queue has no room for all of the run's items (it then takes none), or the idempotency key was used for another request" @@ -117031,7 +117031,7 @@ } ], "title": "Idempotency Key", - "description": "A caller-owned retry key. Repeating the request with it returns the run it created and enqueues nothing new. Reusing it for another request (another model, graph or batch data) is rejected while the queue keeps the key's enqueue receipt." + "description": "A caller-owned retry key. Repeating the request with it returns the run it created and enqueues nothing new; once a reset removed that run, the repeat is refused. Reusing the key for another request (another model, graph or batch data) is rejected while the queue keeps the key's enqueue receipt." } }, "additionalProperties": false, diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts index 68f521a5ae9..f972a5f9036 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -1217,8 +1217,9 @@ export type paths = { * * With a key, repeating the request (also after a lost response) records the queued items and returns the run; it * enqueues nothing new. The items come from the queue's enqueue receipt for the key while it exists, otherwise from - * the queue items that carry the run's origin and those the runner saw. Acknowledge the key through the queue's - * enqueue_batch/acknowledge route once the response arrived. + * the queue items that carry the run's origin and those the runner saw. Once a reset removed the run, a repeat is + * refused (404). Acknowledge the key through the queue's enqueue_batch/acknowledge route once the response + * arrived. */ post: operations["calibrate_model_working_memory"]; delete?: never; @@ -52867,7 +52868,7 @@ export type components = { data?: components["schemas"]["BatchDatum"][][] | null; /** * Idempotency Key - * @description A caller-owned retry key. Repeating the request with it returns the run it created and enqueues nothing new. Reusing it for another request (another model, graph or batch data) is rejected while the queue keeps the key's enqueue receipt. + * @description A caller-owned retry key. Repeating the request with it returns the run it created and enqueues nothing new; once a reset removed that run, the repeat is refused. Reusing the key for another request (another model, graph or batch data) is rejected while the queue keeps the key's enqueue receipt. */ idempotency_key?: string | null; }; @@ -56127,7 +56128,7 @@ export interface operations { "application/json": components["schemas"]["WorkingMemoryCalibrationRunCreated"]; }; }; - /** @description The model could not be found, or the run was removed while it was being created; its items are queued then and fail when they start */ + /** @description The model could not be found; the run was removed while it was being created, whose items are queued then and fail when they start; or a reset removed the run the idempotency key started */ 404: { headers: { [name: string]: unknown; diff --git a/tests/app/services/session_processor/test_working_memory_calibration_runner.py b/tests/app/services/session_processor/test_working_memory_calibration_runner.py index 3fe38688d65..46bd1c6b4e5 100644 --- a/tests/app/services/session_processor/test_working_memory_calibration_runner.py +++ b/tests/app/services/session_processor/test_working_memory_calibration_runner.py @@ -47,11 +47,13 @@ # What each test node saw: the scope it ran in, or None. _SEEN: list[Optional[tuple[str, int, str, str]]] = [] _ON_CANCEL: list[Callable[[], None]] = [] +# What a "hook" node does after its window closed, before it completes. +_WHILE_RUNNING: list[Callable[[], None]] = [] @invocation("test_working_memory_calibration_node", version="1.0.0") class CalibrationProbeNodeInvocation(BaseInvocation): - behaviour: Literal["complete", "fail", "cancel"] = InputField(default="complete") + behaviour: Literal["complete", "fail", "cancel", "hook"] = InputField(default="complete") def invoke(self, context: InvocationContext) -> IntegerOutput: scope = probe.current_scope() @@ -65,6 +67,8 @@ def invoke(self, context: InvocationContext) -> IntegerOutput: if self.behaviour == "cancel": _ON_CANCEL[0]() raise CanceledException() + if self.behaviour == "hook": + _WHILE_RUNNING[0]() return IntegerOutput(value=1) @@ -84,9 +88,11 @@ def reset_stats(self, graph_execution_state_id) -> None: def _reset_seen() -> Iterator[None]: _SEEN.clear() _ON_CANCEL.clear() + _WHILE_RUNNING.clear() yield _SEEN.clear() _ON_CANCEL.clear() + _WHILE_RUNNING.clear() @pytest.fixture @@ -192,6 +198,19 @@ def cancel() -> None: assert run.status == ("failed" if behaviour == "fail" else "canceled") +def test_a_model_reset_while_the_node_runs_keeps_its_points_of_the_model_out( + tmp_path: Path, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner +) -> None: + """The node's window of the VAE closed, then the VAE was reset before the node ended: its windows are persisted + only after that, and must not bring the VAE's points back. The next node measures after the reset.""" + _create(service, "system", "hook", "complete") + _WHILE_RUNNING.append(lambda: service.reset_model("vae-key")) + + assert _run_next(queue, runner[0]).status == "completed" + + assert [point["node_id"] for point in _stored_points(tmp_path)] == ["node1"] + + def test_an_item_whose_owner_is_not_an_admin_fails_before_any_node_runs( tmp_path: Path, mock_invoker: Invoker, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner ) -> None: @@ -286,7 +305,7 @@ def test_a_store_failure_fails_the_item_and_stops_the_session( ) -> None: run_id, _ = _create(service, "system", "complete", "complete") - def refuse(run_id: str, windows) -> None: + def refuse(run_id: str, windows, model_reset_epochs) -> None: raise OSError("read-only file system") monkeypatch.setattr(service, "record_windows", refuse) @@ -296,6 +315,22 @@ def refuse(run_id: str, windows) -> None: assert len(_SEEN) == 1 +def test_a_store_that_cannot_be_read_at_node_start_fails_the_item_before_the_node_runs( + tmp_path: Path, service: WorkingMemoryCalibrationService, queue: SessionQueue, runner, monkeypatch +) -> None: + """Without the reset epochs the node's windows could undo a reset, so the node does not run.""" + _create(service, "system", "complete", "complete") + + def unreadable(run_id: str) -> dict[str, int]: + raise OSError("lock timed out") + + monkeypatch.setattr(service, "model_reset_epochs", unreadable) + item = _run_next(queue, runner[0]) + + assert item.status == "failed" and "lock timed out" in (item.error_message or "") + assert _SEEN == [] and _stored_points(tmp_path) == [] + + def _refused_before_any_node(queue: SessionQueue, runner: DefaultSessionRunner, refusal: str) -> None: item = _run_next(queue, runner) assert item.status == "failed" and refusal in (item.error_message or "") diff --git a/tests/app/services/session_queue/test_session_queue_query_plans.py b/tests/app/services/session_queue/test_session_queue_query_plans.py index ac581d526f9..7bb6f36cb33 100644 --- a/tests/app/services/session_queue/test_session_queue_query_plans.py +++ b/tests/app/services/session_queue/test_session_queue_query_plans.py @@ -129,6 +129,17 @@ def test_queue_listing_and_counts_read_only_the_listing_index(session_queue: Ses assert counts[0].count("SEARCH session_queue") == 1, counts[0] +def test_the_lineage_read_tests_the_origin_on_the_listing_index( + session_queue: SessionQueue, database: Database +) -> None: + """The calibration status reads the retry source of its own items: the origin is tested on the listing index, + and no sort follows.""" + _enqueue(session_queue) + [lineage] = _statements(database, lambda: session_queue.get_queue_item_lineage("default", ORIGIN_PREFIX)) + assert f"USING INDEX {LISTING_INDEX} (queue_id=?)" in lineage.plan + assert "TEMP B-TREE" not in lineage.plan + + def test_a_non_admins_project_reads_keep_their_plans(session_queue: SessionQueue, database: Database) -> None: """Limited to the caller's items, a project's listing and counts still read the listing index alone, and its current and next items are still found by their status: the account's own index would read its whole history.""" diff --git a/tests/app/services/test_invoker.py b/tests/app/services/test_invoker.py index 0f040f765af..7ff7fd0e23a 100644 --- a/tests/app/services/test_invoker.py +++ b/tests/app/services/test_invoker.py @@ -1,18 +1,28 @@ -"""The invoker's service lifecycle: the session processor starts after every other service and stops before them.""" +"""The invoker's service lifecycle: the session processor starts after every other service and stops before them, and +its workers leave their sessions before any other service stops.""" import asyncio import inspect +import threading import time +from collections.abc import Iterator from pathlib import Path from types import SimpleNamespace +from typing import Optional from unittest.mock import MagicMock import pytest from invokeai.app.invocations.math import AddInvocation +from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.invocation_services import InvocationServices from invokeai.app.services.invoker import Invoker -from invokeai.app.services.session_processor.session_processor_default import DefaultSessionProcessor +from invokeai.app.services.session_processor import session_processor_default +from invokeai.app.services.session_processor.session_processor_default import ( + DefaultSessionProcessor, + DefaultSessionRunner, +) +from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem from invokeai.app.services.session_queue.session_queue_default import SessionQueue from invokeai.app.services.shared.database.database import Database from invokeai.app.services.shared.graph import Graph @@ -106,3 +116,172 @@ def test_a_calibration_item_pending_at_startup_runs_once_every_service_started( worker.thread.join(timeout=5) # The processor registered its devices and a release listener in the process-wide pool. GENERATION_DEVICE_POOL.reset() + + +# Bounds a hang only: every wait below is released by an event. +_HANG_S = 30 + + +class _RunnerHeldInANode(DefaultSessionRunner): + """Runs every queue item as one node with no cancel point, like a long decode: it returns once `release` is set.""" + + def __init__(self) -> None: + super().__init__() + self.entered = threading.Event() + self.release = threading.Event() + + def run(self, queue_item: SessionQueueItem) -> None: + self.entered.set() + assert self.release.wait(timeout=_HANG_S), "the test never released the node" + + +class _OneItemQueue: + def __init__(self) -> None: + self._item: Optional[SimpleNamespace] = SimpleNamespace( + item_id=1, session_id="session", queue_id="default", user_id="system" + ) + + def dequeue(self, device: Optional[str] = None) -> Optional[SimpleNamespace]: + item, self._item = self._item, None + return item + + def get_queue_item(self, item_id: int) -> SimpleNamespace: + return SimpleNamespace(item_id=item_id, status="in_progress") + + +def _start_invoker(processor: DefaultSessionProcessor, teardown: list[tuple[str, bool]]) -> Invoker: + """An invoker of `processor` and stand-in services. Each service's stop records whether a worker was alive then.""" + + class Service: + def __init__(self, name: str) -> None: + self.name = name + + def stop(self, invoker: Invoker) -> None: + alive = any(w.thread is not None and w.thread.is_alive() for w in processor._workers) + teardown.append((self.name, alive)) + + services = InvocationServices(**{name: Service(name) for name in inspect.signature(InvocationServices).parameters}) + services.configuration = InvokeAIAppConfig(generation_devices=["cpu"], offload_on_architecture_switch=False) + services.logger = MagicMock() + services.session_queue = _OneItemQueue() # type: ignore[assignment] + services.image_moves = None # type: ignore[assignment] + services.session_processor = processor + return Invoker(services) + + +@pytest.fixture +def no_paging_reads(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr( + "invokeai.app.services.session_processor.session_processor_default.paged_bytes", lambda device: None + ) + yield + # The processors registered their devices and a release listener in the process-wide pool. + GENERATION_DEVICE_POOL.reset() + + +def test_stopping_waits_for_a_worker_to_leave_its_node_before_any_other_service_stops( + no_paging_reads: None, monkeypatch: pytest.MonkeyPatch +) -> None: + runner = _RunnerHeldInANode() + processor = DefaultSessionProcessor(session_runner=runner) + teardown: list[tuple[str, bool]] = [] + invoker = _start_invoker(processor, teardown) + [worker] = processor._workers + assert worker.thread is not None + try: + assert runner.entered.wait(timeout=_HANG_S) + # The node returns only once stop() waits for the worker: released any earlier, it could return before a stop + # that does not wait reached the other services. + join = worker.thread.join + + def release_then_join(timeout: Optional[float] = None) -> None: + runner.release.set() + join(timeout) + + monkeypatch.setattr(worker.thread, "join", release_then_join) + invoker.stop() + + assert not worker.thread.is_alive() + assert teardown and [name for name, alive in teardown if alive] == [] + finally: + runner.release.set() + worker.thread.join(timeout=_HANG_S) + + +def test_a_stop_that_waits_says_which_worker_and_queue_item_it_waits_for( + no_paging_reads: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(session_processor_default, "WORKER_STOP_NOTICE_SECONDS", 0.0) + runner = _RunnerHeldInANode() + processor = DefaultSessionProcessor(session_runner=runner) + invoker = _start_invoker(processor, []) + [worker] = processor._workers + assert worker.thread is not None + notices: list[str] = [] + + def info(message: str) -> None: + # The node returns once the stop has said it is waiting: a stop that waits in silence waits out its deadline. + if message.startswith("Waiting"): + notices.append(message) + runner.release.set() + + invoker.services.logger.info.side_effect = info + try: + assert runner.entered.wait(timeout=_HANG_S) + invoker.stop() + + assert not worker.thread.is_alive() + assert len(notices) == 1 + assert "session_processor_0" in notices[0] and "queue item 1" in notices[0] and "up to 30s" in notices[0] + invoker.services.logger.error.assert_not_called() + finally: + runner.release.set() + worker.thread.join(timeout=_HANG_S) + + +def test_a_worker_still_in_its_node_after_the_stop_keeps_a_new_processor_from_starting( + no_paging_reads: None, monkeypatch: pytest.MonkeyPatch +) -> None: + stop_timeout = session_processor_default.WORKER_STOP_TIMEOUT_SECONDS + # A stop that gives up on the worker at once: the node outlasts any bound. + monkeypatch.setattr(session_processor_default, "WORKER_STOP_TIMEOUT_SECONDS", 0.0) + runner = _RunnerHeldInANode() + processor = DefaultSessionProcessor(session_runner=runner) + invoker = _start_invoker(processor, []) + [worker] = processor._workers + assert worker.thread is not None + # Registered before the stop, so that its start is the one to refuse. + successor = DefaultSessionProcessor(session_runner=DefaultSessionRunner()) + successor.register_generation_devices(InvokeAIAppConfig(generation_devices=["cpu"])) + join = worker.thread.join + try: + assert runner.entered.wait(timeout=_HANG_S) + invoker.stop() + assert worker.thread.is_alive() + [(message,), _] = invoker.services.logger.error.call_args + assert "did not stop" in message and "queue item 1" in message + + # Ctrl+C stops the invoker a second time, with the full timeout: the abandoned worker gets no second wait. + monkeypatch.setattr(session_processor_default, "WORKER_STOP_TIMEOUT_SECONDS", stop_timeout) + joins: list[Optional[float]] = [] + monkeypatch.setattr(worker.thread, "join", lambda timeout=None: joins.append(timeout)) + invoker.stop() + assert joins == [] + assert invoker.services.logger.error.call_count == 1 + + # A restart in this process: neither a second worker nor a new device pool may replace the one still running. + with pytest.raises(RuntimeError, match="still running"): + _start_invoker(successor, []) + assert successor._workers == [] + with pytest.raises(RuntimeError, match="still running"): + DefaultSessionProcessor().register_generation_devices(InvokeAIAppConfig(generation_devices=["cpu"])) + finally: + runner.release.set() + join(timeout=_HANG_S) + + # Once the old worker has exited, the process starts a processor again. + restarted_runner = _RunnerHeldInANode() + restarted_runner.release.set() + restarted = DefaultSessionProcessor(session_runner=restarted_runner) + _start_invoker(restarted, []).stop() + assert [w.thread is not None and not w.thread.is_alive() for w in restarted._workers] == [True] diff --git a/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py b/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py index 4402002be7d..72ec62a6fbf 100644 --- a/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py +++ b/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py @@ -15,6 +15,7 @@ from invokeai.app.invocations.math import AddInvocation from invokeai.app.services.invoker import Invoker +from invokeai.app.services.model_records import UnknownModelException from invokeai.app.services.session_queue.session_queue_common import ( Batch, BatchDatum, @@ -29,6 +30,7 @@ CALIBRATION_ORIGIN_PREFIX, CalibrationDeviceClass, CalibrationQueueFullError, + CalibrationRunNotFoundError, CalibrationStoreError, CalibrationStoreUnusableError, WorkingMemoryCalibrationRequest, @@ -39,6 +41,7 @@ from invokeai.app.services.working_memory_calibration.working_memory_calibration_default import ( WorkingMemoryCalibrationService, ) +from invokeai.app.services.working_memory_calibration.working_memory_calibration_store import CalibrationStore from tests.fixtures.sqlite_database import sqlite_cursor from tests.fixtures.working_memory_calibration import ( GPU_CLASS, @@ -316,6 +319,8 @@ def test_resets_remove_runs_points_and_unused_environments( stored = json.loads(_store_path(tmp_path).read_bytes()) assert (stored["runs"], stored["points"], stored["envs"]) == ([], [], []) assert stored["device_class"]["id"] == GPU_CLASS.id + # Runs without a retry key can never be started again, so no reset keeps their ids. + assert stored["removed_run_ids"] == [] def test_windows_of_a_reset_run_are_dropped(tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue) -> None: @@ -708,9 +713,128 @@ def create(graph_id: Optional[str]): first = create(None) assert create("chosen-by-the-client") == first - # With the run reset and its items gone from the queue, only the key's receipt knows the request. It matches the - # whole batch, graph id included. - service.reset_run(first.run_id) - queue.delete_queue_items_by_id(first.item_ids) - assert create(None) == first - assert _calibration_item_ids(queue) == [] + assert _calibration_item_ids(queue) == first.item_ids + + +def test_a_model_reset_removes_its_points_that_a_write_in_flight_brings( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + """A window of model-a measured in model-b's run reaches the store only after model-a was reset.""" + service = _service(tmp_path, mock_invoker) + run_b = _create(service, "model-b") + entered, release = threading.Event(), threading.Event() + change = CalibrationStore.change + + def held(self: CalibrationStore, *args, **kwargs): + if threading.current_thread() is writer: + entered.set() + assert release.wait(timeout=10) + return change(self, *args, **kwargs) + + monkeypatch.setattr(CalibrationStore, "change", held) + windows = [measured_window(run_b.run_id, model_key="model-a"), measured_window(run_b.run_id)] + writer = threading.Thread(target=service.record_windows, args=(run_b.run_id, windows)) + writer.start() + assert entered.wait(timeout=10) + service.reset_model("model-a") + release.set() + writer.join(timeout=30) + + assert not writer.is_alive() + assert [point["model_key"] for point in json.loads(_store_path(tmp_path).read_bytes())["points"]] == ["vae-key"] + + +def test_a_model_reset_by_another_process_keeps_out_the_points_a_node_measured_before_it( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue +) -> None: + """Two service instances stand for two processes sharing the root.""" + service, other = _service(tmp_path, mock_invoker), _service(tmp_path, mock_invoker) + run_b = _create(service, "model-b") + node_start = service.model_reset_epochs(run_b.run_id) + + other.reset_model("model-a") + windows = [measured_window(run_b.run_id, model_key="model-a"), measured_window(run_b.run_id)] + service.record_windows(run_b.run_id, windows, node_start) + # A node that starts after the reset keeps its points of the model. + service.record_windows(run_b.run_id, [measured_window(run_b.run_id, model_key="model-a")]) + + points = json.loads(_store_path(tmp_path).read_bytes())["points"] + assert [point["model_key"] for point in points] == ["vae-key", "model-a"] + + +def test_a_retry_stands_in_for_its_item_before_the_runner_starts_it( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue +) -> None: + """The failed item was pruned from the queue, then its retry canceled before it ran: the run did not fail.""" + service = _service(tmp_path, mock_invoker) + created = _create(service) + [item_id] = created.item_ids + item = queue.dequeue() + assert item is not None and service.admit(item, created.run_id) is None + queue.fail_queue_item(item_id, "ValueError", "out of memory", "traceback") + service.observe(item, created.run_id) + queue.retry_items_by_id("default", [item_id]) + [retry] = [queued for queued in _calibration_item_ids(queue) if queued != item_id] + + queue.prune("default") + queue.cancel_queue_item(retry) + + run = _run_status(service.get_status(), created.run_id) + assert (run.status, run.error, run.item_ids) == ("canceled", None, [item_id, retry]) + + +def test_a_repeat_records_the_queued_items_of_a_model_removed_meanwhile( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + """The key's receipt names the items; only a new request needs the model.""" + service = _service(tmp_path, mock_invoker) + with pytest.raises(CalibrationStoreError): + _create_while_recording_fails(service, monkeypatch, _request(key="model-removed")) + [item_id] = _calibration_item_ids(queue) + + def removed(key: str) -> None: + raise UnknownModelException(key) + + mock_invoker.services.model_manager.store.get_model.side_effect = removed + repeat = _create(service, key="model-removed") + + assert repeat.item_ids == [item_id] and _calibration_item_ids(queue) == [item_id] + [stored] = json.loads(_store_path(tmp_path).read_bytes())["runs"] + assert (stored["run_id"], stored["queue_item_ids"]) == (repeat.run_id, [item_id]) + with pytest.raises(UnknownModelException): + _create(service, key="another-request") + + +@pytest.mark.parametrize("receipt", ["retained", "collected"]) +@pytest.mark.parametrize("reset", ["run", "model", "all"]) +def test_a_reset_run_cannot_be_started_again_by_its_retry_key( + tmp_path: Path, + mock_invoker: Invoker, + queue: SessionQueue, + mock_sqlite_database: Database, + reset: str, + receipt: str, +) -> None: + """Its items stay queued and fail at admission; repeating the request neither revives the run nor enqueues.""" + service = _service(tmp_path, mock_invoker) + created = _create(service, key="reset-me") + if reset == "run": + service.reset_run(created.run_id) + elif reset == "model": + service.reset_model("model-a") + else: + service.reset_all() + if receipt == "collected": + _collect_receipt(queue, mock_sqlite_database, "reset-me") + + with pytest.raises(CalibrationRunNotFoundError, match="reset"): + _create(service, key="reset-me") + # Nor on a GPU of another class, whose file would otherwise take the run. + with pytest.raises(CalibrationRunNotFoundError, match="reset"): + _create(_service(tmp_path, mock_invoker, device_class=OTHER_CLASS), key="reset-me") + + assert _calibration_item_ids(queue) == created.item_ids + assert service.get_status().device_classes[0].runs == [] + item = queue.dequeue() + assert item is not None and item.item_id in created.item_ids + assert "does not exist" in (service.admit(item, created.run_id) or "") diff --git a/tests/app/test_api_dependencies.py b/tests/app/test_api_dependencies.py index 3a2fcef2b54..05e31f2291d 100644 --- a/tests/app/test_api_dependencies.py +++ b/tests/app/test_api_dependencies.py @@ -1,4 +1,5 @@ import asyncio +import threading from logging import getLogger from typing import Optional from unittest.mock import MagicMock @@ -124,3 +125,34 @@ def start_and_borrow(self: ImageIndexService, invoker: Invoker) -> None: # The processor registered the devices and a release listener in the process-wide pool. GENERATION_DEVICE_POOL.reset() loop.close() + + +def test_a_restart_refused_for_a_worker_the_last_stop_gave_up_on_opens_no_database(monkeypatch, tmp_path) -> None: + _skip_process_wide_setup(monkeypatch) + opened: list[Database] = [] + original_init_database = dependencies.init_database + + def init_database(*args, **kwargs) -> Database: + opened.append(original_init_database(*args, **kwargs)) + return opened[-1] + + monkeypatch.setattr(dependencies, "init_database", init_database) + # A session worker the previous invoker's stop gave up on: still in its node. + release = threading.Event() + stuck = threading.Thread(target=release.wait, name="session_processor_0", daemon=True) + stuck.start() + monkeypatch.setattr(session_processor_default, "_abandoned_workers", [stuck]) + config = InvokeAIAppConfig(use_memory_db=True, image_index_enabled=False, scan_models_on_startup=False) + config._root = tmp_path + loop = asyncio.new_event_loop() + try: + with pytest.raises(RuntimeError, match="still running"): + ApiDependencies.initialize(config=config, event_handler_id=0, loop=loop, logger=getLogger("test")) + assert opened == [] + finally: + release.set() + stuck.join(timeout=30) + for database in opened: + database.dispose() + GENERATION_DEVICE_POOL.reset() + loop.close() diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_locks_held.py b/tests/backend/model_manager/load/model_cache/test_model_cache_locks_held.py index dd6ce523bc5..746c557e834 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_locks_held.py +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_locks_held.py @@ -1,5 +1,6 @@ -"""`ModelCache.locks_held`: every pin on the cache's records, whichever thread took it and of whichever kind. The -working-memory probe compares it with its own thread's locks to see another thread sharing the device.""" +"""`ModelCache.locks_held`: every pin on the cache's records, those it dropped included, whichever thread took it and +of whichever kind. The working-memory probe compares it with its own thread's locks to see another thread sharing the +device.""" import logging import threading @@ -46,3 +47,75 @@ def test_counts_device_and_ram_locks_of_every_thread(cache: ModelCache): assert cache.locks_held() == 1 cache.unlock(vae) assert cache.locks_held() == 0 + + +def test_counts_the_locks_on_a_record_the_cache_dropped_while_it_was_locked(cache: ModelCache): + cache.put("vae", torch.nn.Linear(2, 2)) + dropped = cache.get("vae") + cache.lock(dropped, None) + # The VRAM-move error paths drop a record that is still locked; reaching them takes a GPU. + cache._delete_cache_entry(dropped) + cache.put("vae", torch.nn.Linear(2, 2)) + readmitted = cache.get("vae") + cache.lock_in_ram(readmitted) + + assert cache.locks_held() == 2 + + cache.unlock(dropped) + assert cache.locks_held() == 1 + cache.unlock(readmitted) + assert cache.locks_held() == 0 + + +def test_an_unlock_without_a_lock_is_refused_and_leaves_the_count_right(cache: ModelCache): + cache.put("vae", torch.nn.Linear(2, 2)) + vae = cache.get("vae") + cache.lock(vae, None) + cache.unlock(vae) + + with pytest.raises(RuntimeError, match="unlocked more often than it was locked"): + cache.unlock(vae) + + assert cache.locks_held() == 0 + cache.lock(vae, None) + assert (cache.locks_held(), vae.is_locked) == (1, True) + cache.unlock(vae) + assert (cache.locks_held(), vae.is_locked) == (0, False) + + +def _failing_vram_load(cache: ModelCache, monkeypatch: pytest.MonkeyPatch, settles_first: bool): + """A record on a GPU whose VRAM move fails: on its first pass, or on the pass after a paced first one.""" + cache.put("vae", torch.nn.Linear(2, 2)) + vae = cache.get("vae") + # Only the device the record names: the move it would start is the stub below. + monkeypatch.setattr(vae.cached_model, "_compute_device", torch.device("cuda:0")) + passes = iter([False] if settles_first else []) + + def load_locked_model(cache_entry, working_mem_bytes=None, max_move_bytes=None, stream_started_at=None) -> bool: + for settled in passes: + return settled + raise torch.OutOfMemoryError("no room") + + monkeypatch.setattr(cache, "_load_locked_model", load_locked_model) + monkeypatch.setattr(cache, "_log_cache_state", lambda *args, **kwargs: None) + return vae + + +def test_a_lock_whose_vram_move_fails_holds_no_lock(cache: ModelCache, monkeypatch: pytest.MonkeyPatch): + vae = _failing_vram_load(cache, monkeypatch, settles_first=False) + + with pytest.raises(torch.OutOfMemoryError): + cache.lock(vae, None) + + assert (cache.locks_held(), vae.is_locked) == (0, False) + + +def test_a_paced_lock_whose_next_pass_fails_holds_no_lock(cache: ModelCache, monkeypatch: pytest.MonkeyPatch): + vae = _failing_vram_load(cache, monkeypatch, settles_first=True) + assert cache.lock(vae, None, max_move_bytes=1) is False + assert cache.locks_held() == 1 + + with pytest.raises(torch.OutOfMemoryError): + cache.continue_lock(vae, None, max_move_bytes=1) + + assert (cache.locks_held(), vae.is_locked) == (0, False) From 1370b2347f7b944c6fa8b199700d05ff892035d9 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Wed, 21 Oct 2026 21:33:32 +0200 Subject: [PATCH 7/9] fix(working-memory): refuse unwritable stores and duplicate enqueues in calibration runs Status reports a store it cannot write as unavailable, and a run is marked before its batch is enqueued so a repeat never enqueues it twice. A clean finish is recorded before the queue completes the item, a key conflict answers 409 before the model lookup, and run summaries list a retry's source. --- .../Architecture/model-manager.mdx | 12 +- .../session_processor_default.py | 20 +++ .../working_memory_calibration_base.py | 6 + .../working_memory_calibration_common.py | 14 +- .../working_memory_calibration_default.py | 100 +++++++++--- .../working_memory_calibration_store.py | 29 +++- invokeai/frontend/api/openapi.json | 6 +- invokeai/frontend/api/schema.ts | 6 +- .../test_working_memory_calibration.py | 30 ++++ .../test_working_memory_calibration_runner.py | 52 +++++++ ...test_working_memory_calibration_service.py | 144 +++++++++++++++++- 11 files changed, 376 insertions(+), 43 deletions(-) diff --git a/docs/src/content/docs/development/Architecture/model-manager.mdx b/docs/src/content/docs/development/Architecture/model-manager.mdx index 970bfeb6454..e34df87dd26 100644 --- a/docs/src/content/docs/development/Architecture/model-manager.mdx +++ b/docs/src/content/docs/development/Architecture/model-manager.mdx @@ -1040,9 +1040,9 @@ A calibration run measures one model's operations on this machine's GPU. Runs on 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:` 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`. + - 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, that request is refused (409) and the first run stays; the receipt is checked before the request's model, so an unknown model does not turn the refusal into a 404. 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. + - 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: the run is marked before its batch is enqueued (`enqueue_attempted`, set as the run is recorded), and only the attempt that set the mark clears it again, when its own enqueue queued nothing; a replay of the key's receipt never clears it, since the attempt that queued the batch may have queued it. A repeat that finds a marked run without items, without the receipt and with no item in the queue or seen by the runner answers with the run and no items, and the run reads as canceled. That is the case when the items were cleared or pruned before the runner saw them, when the server shut down between the mark and the enqueue, or when an enqueue raised, the queue could not be read to tell what it had queued, and it had queued nothing. Should clearing the mark fail (it is logged), the run stays marked although nothing was queued, and the key cannot start it; a request with another key can. 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. @@ -1051,12 +1051,12 @@ A calibration run measures one model's operations on this machine's GPU. Runs on - 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. +5. After the item, the runner records the item's outcome in the run, even when the session raised. A session that ran every node without an error, and that nothing canceled or failed (a failed save of its measurements fails the item), is recorded as completed before the queue completes the item, so that a process stopping in between cannot leave the item to read as canceled once the queue has pruned it; the queue's status then replaces that record. 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 `/working_memory/.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 runs: model key, model hash as provenance, owner, environment, queue items, whether an enqueue of the run's batch may have committed, 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; @@ -1076,7 +1076,7 @@ A file is unusable when it has another schema version, invalid content, an id th 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. +- 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's items include each item a known retry came from, even when nothing else knows that item any more. - 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): @@ -1084,7 +1084,7 @@ Routes, besides the 403 for non-admins and a 503 when the store cannot be read o | 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) | – | +| `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 that can be written. The check stages a temporary file in the store's folder (or the nearest folder above it that exists) and removes it at once, checks that the lock file is writable, and on Windows that the store file is not read-only. | 503 when the store folder cannot be listed | | `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 | diff --git a/invokeai/app/services/session_processor/session_processor_default.py b/invokeai/app/services/session_processor/session_processor_default.py index b89b7517f64..68716b86591 100644 --- a/invokeai/app/services/session_processor/session_processor_default.py +++ b/invokeai/app/services/session_processor/session_processor_default.py @@ -317,6 +317,17 @@ def _run_calibration_item(self, queue_item: SessionQueueItem, run_id: str) -> No try: self._on_before_run_session(queue_item=queue_item) self._run_session_loop(queue_item) + if service is not None: + # Before the queue completes the item: the runner's own record survives a prune of the item. + try: + if self._calibration_session_completed(queue_item): + service.observe_completed(queue_item, run_id) + except SessionQueueItemNotFoundError: + pass + except Exception: + self._services.logger.warning( + f"Could not record the completion of calibration item {queue_item.item_id}", exc_info=True + ) self._on_after_run_session(queue_item=queue_item) finally: self._calibration_run_id = None @@ -333,6 +344,15 @@ def _run_calibration_item(self, queue_item: SessionQueueItem, run_id: str) -> No f"Could not record the outcome of calibration item {queue_item.item_id}", exc_info=True ) + def _calibration_session_completed(self, queue_item: SessionQueueItem) -> bool: + """Whether the queue is about to complete the item: its session ran every node without an error, and nothing + canceled or failed it. A save that failed and an owner revoked after the last node have failed or canceled + the item in the queue, while its session reads as complete.""" + session = queue_item.session + if not session.is_complete() or session.has_error() or self._is_canceled(): + return False + return self._services.session_queue.get_queue_item(queue_item.item_id).status not in ("canceled", "failed") + def _run_calibration_node(self, invocation: BaseInvocation, queue_item: SessionQueueItem, run_id: str) -> bool: """Run one node of a calibration item inside a probe scope and persist its windows, whether the node completed, failed or was canceled; unless it completed, they are persisted invalid. Returns False when they could not be diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py index 63207fd6c78..8362cbf8220 100644 --- a/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py @@ -29,6 +29,12 @@ async def create_run( def admit(self, queue_item: SessionQueueItem, run_id: str) -> Optional[str]: """Called by the session runner before a calibration item runs: None to run it, or why it must fail.""" + @abstractmethod + def observe_completed(self, queue_item: SessionQueueItem, run_id: str) -> None: + """Record that the item's session ran every node without an error and nothing canceled or failed it, before + the queue completes the item: should the process stop in between and the queue later prune the item, its + outcome is not lost. `observe` replaces it with the queue's status.""" + @abstractmethod def model_reset_epochs(self, run_id: str) -> dict[str, int]: """The reset epoch of every model reset on the run's device class; empty when the run does not exist. Read diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py index e5ee262ca87..74fa75f9e6e 100644 --- a/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py @@ -118,6 +118,12 @@ class CalibrationRun(_StoreModel): created_at: datetime.datetime env_id: str = Field(description="The environment the run was created under") queue_item_ids: list[int] = Field(default_factory=list, description="The items the run was enqueued with") + enqueue_attempted: bool = Field( + default=False, + description="Set before the run's batch is enqueued, and cleared only by the attempt that set it, when its " + "enqueue queued nothing. A repeat of the request that finds no items of the run then enqueues nothing: they " + "may have been queued", + ) observed: dict[int, ObservedItem] = Field( default_factory=dict, description="The last status the session runner saw per item, retries included" ) @@ -171,7 +177,8 @@ class WorkingMemoryCalibrationRunCreated(BaseModel): device_class_id: str item_ids: list[int] = Field( description="The queue items the run enqueued. A repeat of the request that finds them unrecorded takes over " - "every queue item carrying the run's origin, retries included" + "every queue item carrying the run's origin, retries included; none when the queue no longer has them and the " + "runner never saw them" ) @@ -191,7 +198,7 @@ class WorkingMemoryCalibrationRunSummary(BaseModel): "valid points were recorded, else canceled" ) error: Optional[str] = Field(description="The error of a failed item, when the run failed") - item_ids: list[int] = Field(description="Every queue item of the run, retries included") + item_ids: list[int] = Field(description="Every queue item of the run known, retries and the items they retried") point_count: int @@ -229,7 +236,8 @@ class WorkingMemoryCalibrationDeviceClass(BaseModel): class WorkingMemoryCalibrationStatus(BaseModel): available: bool = Field( - description="Whether a calibration run can start now: one GPU generation device whose store file is usable" + description="Whether a calibration run can start now: one GPU generation device whose store file is usable " + "and can be written" ) unavailable_reason: Optional[str] current_device_class_id: Optional[str] diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py index 5f59115cefa..b6bbcaf19f6 100644 --- a/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py @@ -7,7 +7,7 @@ from collections.abc import Callable, Mapping from functools import partial from itertools import chain -from typing import TYPE_CHECKING, Final, Optional +from typing import TYPE_CHECKING, Final, NamedTuple, Optional import torch @@ -16,6 +16,7 @@ DEFAULT_QUEUE_ID, QUEUE_ITEM_STATUS, Batch, + EnqueueIdempotencyConflictError, QueueItemLineage, SessionQueueItem, SessionQueueItemNotFoundError, @@ -163,6 +164,8 @@ def retried_from(item_id: int) -> Optional[int]: source = retried_from(item_id) lineage = item_id if source is None else source latest[lineage] = max(latest.get(lineage, item_id), item_id) + # A retry names the item it came from, which may be known by nothing else any more. + item_ids = sorted(set(item_ids) | set(latest)) effective = sorted(latest.values()) statuses = {item_id: status_of(item_id) for item_id in effective} @@ -180,6 +183,16 @@ def retried_from(item_id: int) -> Optional[int]: return ("partial" if valid_point_count > 0 else "canceled"), None, item_ids +class _Enqueue(NamedTuple): + """What a creation enqueues: the batch, the run's device class, whether this request recorded the run, and + whether it marks the run's enqueue, and so may clear the mark again when nothing was queued.""" + + batch: Batch + device_class_id: str + created: bool + marks: bool + + class WorkingMemoryCalibrationService(WorkingMemoryCalibrationServiceBase): """Calibration runs and their store. @@ -359,7 +372,11 @@ async def _create_run( prepared = await asyncio.to_thread(self._prepare_run, model_key, request, user_id, run_id) if isinstance(prepared, WorkingMemoryCalibrationRunCreated): return prepared - batch, device_class_id, created = prepared + batch, device_class_id, created, marks = prepared + if marks and not created: + # Before the enqueue: once it may have committed, a repeat that finds no items must not enqueue again. + # A run this request recorded was marked as it was recorded. + await asyncio.to_thread(self._set_enqueue_attempted, device_class_id, run_id, True) try: result = await self._services.session_queue.enqueue_batch( queue_id=DEFAULT_QUEUE_ID, batch=batch, prepend=False, user_id=user_id, all_or_none=True @@ -379,8 +396,7 @@ async def _create_run( # The enqueue's own error answers the request. raise if not queued: - if created: - await asyncio.to_thread(self._discard, run_id) + await asyncio.to_thread(self._enqueued_nothing, device_class_id, run_id, created, marks) raise self._services.logger.warning( f"Enqueuing calibration run {run_id} raised, and the queue holds its items {queued}", @@ -389,8 +405,7 @@ async def _create_run( item_ids = queued else: if result.enqueued == 0: - if created: - await asyncio.to_thread(self._discard, run_id) + await asyncio.to_thread(self._enqueued_nothing, device_class_id, run_id, created, marks) # A partial run would measure some of the batch's cases, so the queue takes all of them or none. raise CalibrationQueueFullError( f"The queue is full: it has no room for all {result.requested} items of the calibration run, " @@ -415,9 +430,8 @@ async def _create_run( def _prepare_run( self, model_key: str, request: WorkingMemoryCalibrationRequest, user_id: str, run_id: str - ) -> WorkingMemoryCalibrationRunCreated | tuple[Batch, str, bool]: - """The run a repeated request already enqueued; or the batch to enqueue, the run's device class and whether - this call recorded the run.""" + ) -> "WorkingMemoryCalibrationRunCreated | _Enqueue": + """The run a repeated request already enqueued, or what to enqueue.""" key = request.idempotency_key found: Optional[tuple[str, CalibrationRun]] = None if key is not None and (found := self._find_run(run_id)) is not None: @@ -432,22 +446,36 @@ def _prepare_run( # again would queue the batch a second time: the items carrying the run's origin, and those the runner # saw, are taken over instead. if self._services.session_queue.get_enqueue_receipt(DEFAULT_QUEUE_ID, key, user_id) is not None: - return self._batch(request, run_id), device_class_id, False + # The attempt that queued them marked the run, and the mark stays whatever this replay finds. + return _Enqueue(self._batch(request, run_id), device_class_id, created=False, marks=False) item_ids = sorted(set(self._queued_item_ids(run_id)) | set(run.observed)) if item_ids: self._record_items(device_class_id, run_id, item_ids) return WorkingMemoryCalibrationRunCreated( run_id=run_id, device_class_id=device_class_id, item_ids=item_ids ) - elif key is not None and self._was_reset(run_id): - raise CalibrationRunNotFoundError(run_id, _RESET) + if run.enqueue_attempted: + # The batch may have been queued, and its items have left the queue unseen: enqueuing it again could + # run it twice. The run stays without items and reads as canceled. + self._services.logger.warning( + f"Calibration run {run_id} may have been enqueued, and the queue no longer has its items; a repeat " + "of its request enqueues nothing" + ) + return WorkingMemoryCalibrationRunCreated(run_id=run_id, device_class_id=device_class_id, item_ids=[]) + elif key is not None: + if self._was_reset(run_id): + raise CalibrationRunNotFoundError(run_id, _RESET) + # The key's receipt was written by an enqueue of a recorded run, and no run has this request's id: the + # key belongs to another request. Refused here, before the model of this one is looked up. + if self._services.session_queue.get_enqueue_receipt(DEFAULT_QUEUE_ID, key, user_id) is not None: + raise EnqueueIdempotencyConflictError(f"Idempotency key {key} is already used by another submission") device = self._calibration_device() # Raises UnknownModelException for an unknown key. config = self._services.model_manager.store.get_model(model_key) batch = self._batch(request, run_id) if found is not None: - # The run is recorded without items, and nothing of it is queued: enqueue its batch. - return batch, found[0], False + # The run is recorded without items and unmarked, and nothing of it is queued: enqueue its batch. + return _Enqueue(batch, found[0], created=False, marks=True) device_class = self._device_class(device) env = self._environment(device) run = CalibrationRun( @@ -457,6 +485,8 @@ def _prepare_run( user_id=user_id, created_at=_now(), env_id=env.id, + # Marked as it is recorded: the enqueue follows, and nothing else is written in between. + enqueue_attempted=True, ) def add(file: CalibrationStoreFile) -> tuple[bool, bool]: @@ -468,9 +498,9 @@ def add(file: CalibrationStoreFile) -> tuple[bool, bool]: _add_env(file, env) return True, True - created = self._store.change(device_class.id, add, create=CalibrationDeviceClass.of(device_class)) + created = bool(self._store.change(device_class.id, add, create=CalibrationDeviceClass.of(device_class))) self._run_files[run_id] = device_class.id - return batch, device_class.id, bool(created) + return _Enqueue(batch, device_class.id, created=created, marks=created) @staticmethod def _batch(request: WorkingMemoryCalibrationRequest, run_id: str) -> Batch: @@ -498,13 +528,30 @@ def _queued_item_ids(self, run_id: str) -> list[int]: queue = self._services.session_queue return sorted(queue.get_queue_item_ids(DEFAULT_QUEUE_ID, origin_prefix=calibration_origin(run_id)).item_ids) - def _discard(self, run_id: str) -> None: - """Remove a run this request recorded and queued nothing for. Its key stays unused: a repeat may start it.""" + def _enqueued_nothing(self, device_class_id: str, run_id: str, created: bool, marks: bool) -> None: + """The enqueue queued nothing: remove the run this request recorded, or clear the mark this request set on + one it found, so that a repeat may enqueue the batch. A mark an earlier attempt set stays: that attempt may + have queued the batch.""" try: - self._remove_run(run_id, remember=False) + if created: + self._remove_run(run_id, remember=False) + elif marks: + self._set_enqueue_attempted(device_class_id, run_id, False) except Exception: self._services.logger.warning(f"Could not discard calibration run {run_id}", exc_info=True) + def _set_enqueue_attempted(self, device_class_id: str, run_id: str, attempted: bool) -> None: + def mark(file: CalibrationStoreFile) -> tuple[bool, bool]: + for run in file.runs: + if run.run_id == run_id: + changed = run.enqueue_attempted != attempted + run.enqueue_attempted = attempted + return changed, True + return False, False + + if not self._store.change(device_class_id, mark): + raise CalibrationRunNotFoundError(run_id) + def _record_items(self, device_class_id: str, run_id: str, item_ids: list[int]) -> None: def remember(file: CalibrationStoreFile) -> tuple[bool, bool]: for run in file.runs: @@ -600,6 +647,13 @@ def observe(self, queue_item: SessionQueueItem, run_id: str, error: Optional[Bas ), ) + def observe_completed(self, queue_item: SessionQueueItem, run_id: str) -> None: + self._observe( + run_id, + queue_item.item_id, + ObservedItem(status="completed", retried_from_item_id=queue_item.retried_from_item_id, observed_at=_now()), + ) + def model_reset_epochs(self, run_id: str) -> dict[str, int]: found = self._find_run(run_id) if found is None: @@ -683,9 +737,11 @@ def get_status(self) -> WorkingMemoryCalibrationStatus: schema_version=STORE_SCHEMA_VERSION, device_class=CalibrationDeviceClass.of(current_class) ) read = StoreRead(device_class_id, empty, None) - if is_current and read.file is None: - # Creating a run would have to write this file, which is kept as it is. - unavailable = f"The calibration store of this GPU's device class cannot be written: {read.problem}" + if is_current: + # Creating a run writes this file: an unusable one is kept as it is, and the root may refuse writes. + problem = read.problem if read.file is None else self._store.write_problem(device_class_id) + if problem is not None: + unavailable = f"The calibration store of this GPU's device class cannot be written: {problem}" device_classes.append( self._device_class_status(read, is_current, current_env if is_current else None, live) ) diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_store.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_store.py index c669385f78f..f865b16571a 100644 --- a/invokeai/app/services/working_memory_calibration/working_memory_calibration_store.py +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_store.py @@ -141,7 +141,34 @@ def device_class_ids(self) -> list[str]: """The device classes that have a file, readable or not.""" if not self._dir.is_dir(): return [] - return sorted(path.stem for path in self._dir.iterdir() if _FILE_NAME.match(path.name) and path.is_file()) + try: + return sorted(path.stem for path in self._dir.iterdir() if _FILE_NAME.match(path.name) and path.is_file()) + except OSError as error: + raise CalibrationStoreError(f"Cannot list {self._dir}: {error}") from error + + def write_problem(self, device_class_id: str) -> Optional[str]: + """Why a change of the device class's file could not be written now, or None. + + Checks what a write needs without changing anything that lasts: a file can be created in the directory (or, + while there is none, in the nearest folder that exists, where the write would create it); the lock file can be + opened for writing; and, on Windows, the file is not read-only, which refuses the replace there. The file it + creates is a temporary one the system deletes, also should the process stop. + """ + folder = self._dir + while not folder.exists() and folder.parent != folder: + folder = folder.parent + try: + with tempfile.TemporaryFile(dir=folder, prefix=".invokeai-write-check."): + pass + except OSError as error: + return f"Cannot write in {folder}: {error}" + lock_path = self._dir / f"{device_class_id}.json.lock" + if lock_path.exists() and not os.access(lock_path, os.W_OK): + return f"{lock_path} is read-only" + path = self.path_for(device_class_id) + if os.name == "nt" and path.exists() and not os.access(path, os.W_OK): + return f"{path} is read-only" + return None @contextmanager def _locked(self, device_class_id: str, *, writing: bool) -> Iterator[None]: diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json index 068b7b8f779..2eec3dbdfb6 100644 --- a/invokeai/frontend/api/openapi.json +++ b/invokeai/frontend/api/openapi.json @@ -117071,7 +117071,7 @@ }, "type": "array", "title": "Item Ids", - "description": "The queue items the run enqueued. A repeat of the request that finds them unrecorded takes over every queue item carrying the run's origin, retries included" + "description": "The queue items the run enqueued. A repeat of the request that finds them unrecorded takes over every queue item carrying the run's origin, retries included; none when the queue no longer has them and the runner never saw them" } }, "type": "object", @@ -117136,7 +117136,7 @@ }, "type": "array", "title": "Item Ids", - "description": "Every queue item of the run, retries included" + "description": "Every queue item of the run known, retries and the items they retried" }, "point_count": { "type": "integer", @@ -117163,7 +117163,7 @@ "available": { "type": "boolean", "title": "Available", - "description": "Whether a calibration run can start now: one GPU generation device whose store file is usable" + "description": "Whether a calibration run can start now: one GPU generation device whose store file is usable and can be written" }, "unavailable_reason": { "anyOf": [ diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts index f972a5f9036..eeb7d31420d 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -52887,7 +52887,7 @@ export type components = { device_class_id: string; /** * Item Ids - * @description The queue items the run enqueued. A repeat of the request that finds them unrecorded takes over every queue item carrying the run's origin, retries included + * @description The queue items the run enqueued. A repeat of the request that finds them unrecorded takes over every queue item carrying the run's origin, retries included; none when the queue no longer has them and the runner never saw them */ item_ids: number[]; }; @@ -52921,7 +52921,7 @@ export type components = { error: string | null; /** * Item Ids - * @description Every queue item of the run, retries included + * @description Every queue item of the run known, retries and the items they retried */ item_ids: number[]; /** Point Count */ @@ -52931,7 +52931,7 @@ export type components = { WorkingMemoryCalibrationStatus: { /** * Available - * @description Whether a calibration run can start now: one GPU generation device whose store file is usable + * @description Whether a calibration run can start now: one GPU generation device whose store file is usable and can be written */ available: boolean; /** Unavailable Reason */ diff --git a/tests/app/routers/test_working_memory_calibration.py b/tests/app/routers/test_working_memory_calibration.py index 056b75f71f9..35cc2de322f 100644 --- a/tests/app/routers/test_working_memory_calibration.py +++ b/tests/app/routers/test_working_memory_calibration.py @@ -293,3 +293,33 @@ def test_the_status_carries_the_device_class_id(client: TestClient, admin_token: "id": GPU_CLASS.id, **GPU_CLASS.model_dump(exclude={"id"}), } + + +def test_a_store_directory_that_cannot_be_listed_answers_the_status_with_503( + client: TestClient, admin_token: str, calibration, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + store = tmp_path / "working_memory" + store.mkdir() + iterdir = Path.iterdir + + def denied(path: Path): + if path == store: + raise PermissionError(13, "Permission denied", str(path)) + return iterdir(path) + + monkeypatch.setattr(Path, "iterdir", denied) + response = client.get("/api/v2/models/working_memory/status", headers=_auth(admin_token)) + + assert response.status_code == 503 and "Permission denied" in response.json()["detail"] + + +def test_a_retry_key_reused_for_an_unknown_model_is_a_conflict( + client: TestClient, admin_token: str, calibration +) -> None: + """Key reuse is the documented 409, whichever model the other request names.""" + body = {"graph": _graph(), "idempotency_key": "taken"} + assert _calibrate(client, admin_token, body).status_code == 201 + + response = _calibrate(client, admin_token, body, key="no-such-model") + + assert response.status_code == 409, response.text diff --git a/tests/app/services/session_processor/test_working_memory_calibration_runner.py b/tests/app/services/session_processor/test_working_memory_calibration_runner.py index 46bd1c6b4e5..2ff3073bc97 100644 --- a/tests/app/services/session_processor/test_working_memory_calibration_runner.py +++ b/tests/app/services/session_processor/test_working_memory_calibration_runner.py @@ -331,6 +331,58 @@ def unreadable(run_id: str) -> dict[str, int]: assert _SEEN == [] and _stored_points(tmp_path) == [] +def test_an_item_that_completed_keeps_its_outcome_when_the_queue_prunes_it_before_the_runner_recorded_it( + service: WorkingMemoryCalibrationService, queue: SessionQueue, runner, monkeypatch +) -> None: + """The process stopped after the queue completed the item and before the runner recorded the outcome: the runner + recorded the completion before the queue did.""" + run_id, _ = _create(service, "system", "complete") + monkeypatch.setattr(service, "observe", lambda *args, **kwargs: None) + order: list[str] = [] + observe_completed, complete_queue_item = service.observe_completed, queue.complete_queue_item + + def recorded(*args, **kwargs) -> None: + order.append("runner") + observe_completed(*args, **kwargs) + + def completed(*args, **kwargs): + order.append("queue") + return complete_queue_item(*args, **kwargs) + + monkeypatch.setattr(service, "observe_completed", recorded) + monkeypatch.setattr(queue, "complete_queue_item", completed) + + assert _run_next(queue, runner[0]).status == "completed" + queue.prune("default") + + assert order == ["runner", "queue"] + [run] = service.get_status().device_classes[0].runs + assert (run.run_id, run.status) == (run_id, "completed") + + +@pytest.mark.parametrize("ending", ["node-failed", "save-failed"]) +def test_an_item_that_did_not_complete_is_not_recorded_as_completed( + service: WorkingMemoryCalibrationService, queue: SessionQueue, runner, monkeypatch, ending: str +) -> None: + """Its session reads as complete all the same. With the runner's last record lost and the item pruned, the run + must not read as a success.""" + run_id, _ = _create(service, "system", "fail" if ending == "node-failed" else "complete") + monkeypatch.setattr(service, "observe", lambda *args, **kwargs: None) + if ending == "save-failed": + + def refuse(run_id: str, windows, model_reset_epochs) -> None: + raise OSError("read-only file system") + + monkeypatch.setattr(service, "record_windows", refuse) + + item = _run_next(queue, runner[0]) + assert item.status == "failed" and item.session.is_complete() + queue.prune("default") + + [run] = service.get_status().device_classes[0].runs + assert (run.run_id, run.status) == (run_id, "canceled") + + def _refused_before_any_node(queue: SessionQueue, runner: DefaultSessionRunner, refusal: str) -> None: item = _run_next(queue, runner) assert item.status == "failed" and refusal in (item.error_message or "") diff --git a/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py b/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py index 72ec62a6fbf..2374026c343 100644 --- a/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py +++ b/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py @@ -4,6 +4,7 @@ import json import logging import os +import stat import threading import time from pathlib import Path @@ -688,18 +689,21 @@ def test_a_failed_enqueue_answers_with_its_own_error_when_the_queue_cannot_be_re ) -> None: """Whether anything of the run was queued cannot be told then, so the run is kept.""" service = _service(tmp_path, mock_invoker) - first = _create(service, key="taken") def unreadable(*args: object, **kwargs: object) -> None: raise RuntimeError("the database is gone") + async def failing_enqueue(*args: object, **kwargs: object) -> None: + raise RuntimeError("the connection dropped") + with monkeypatch.context() as patch: + patch.setattr(queue, "enqueue_batch", failing_enqueue) patch.setattr(queue, "get_queue_item_ids", unreadable) - with pytest.raises(EnqueueIdempotencyConflictError): - _create(service, "model-b", key="taken") + with pytest.raises(RuntimeError, match="the connection dropped"): + _create(service, key="unknown-outcome") - runs = {run.run_id: run.model_key for run in service.get_status().device_classes[0].runs} - assert len(runs) == 2 and runs[first.run_id] == "model-a" and "model-b" in runs.values() + [run] = service.get_status().device_classes[0].runs + assert run.model_key == "model-a" def test_requests_that_differ_only_in_their_graph_id_are_one_request( @@ -838,3 +842,133 @@ def test_a_reset_run_cannot_be_started_again_by_its_retry_key( item = queue.dequeue() assert item is not None and item.item_id in created.item_ids assert "does not exist" in (service.admit(item, created.run_id) or "") + + +def _read_only(path: Path) -> None: + os.chmod(path, stat.S_IREAD) + + +@pytest.mark.parametrize("unwritable", ["root-is-a-file", "read-only-lock-file"]) +def test_a_store_that_cannot_be_written_makes_calibration_unavailable( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, unwritable: str +) -> None: + """The status says so, and a run cannot start: the two agree.""" + if unwritable == "root-is-a-file": + root = tmp_path / "root" + root.write_bytes(b"") + service = _service(root, mock_invoker) + else: + if hasattr(os, "geteuid") and os.geteuid() == 0: + pytest.skip("root writes read-only files") + service = _service(tmp_path, mock_invoker) + _create(service) + lock = _store_path(tmp_path).with_name(f"{GPU_CLASS.id}.json.lock") + _read_only(lock) + try: + status = service.get_status() + + assert not status.available and "cannot be written" in (status.unavailable_reason or "") + with pytest.raises(CalibrationStoreError): + _create(service, key="refused") + finally: + if unwritable != "root-is-a-file": + os.chmod(lock, stat.S_IREAD | stat.S_IWRITE) + + +def test_a_repeat_enqueues_nothing_once_the_items_of_an_uncertain_admission_are_gone( + tmp_path: Path, + mock_invoker: Invoker, + queue: SessionQueue, + mock_sqlite_database: Database, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The items were queued but not recorded (a 503); the client acknowledged the key, the queue collected the + receipt and the items left the queue before the runner saw them. Enqueuing again would run the batch twice.""" + service = _service(tmp_path, mock_invoker) + with pytest.raises(CalibrationStoreError): + _create_while_recording_fails(service, monkeypatch, _request(key="uncertain")) + [item_id] = _calibration_item_ids(queue) + _collect_receipt(queue, mock_sqlite_database, "uncertain") + queue.delete_queue_items_by_id([item_id]) + + repeat = _create(service, key="uncertain") + + assert repeat.item_ids == [] and _calibration_item_ids(queue) == [] + [run] = service.get_status().device_classes[0].runs + assert (run.run_id, run.status) == (repeat.run_id, "canceled") + + +def test_a_run_lists_the_item_its_retry_came_from( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + """The run never recorded the item (its store failed, without a key), the runner never saw it, and the queue + deleted it: only its retry still names it.""" + service = _service(tmp_path, mock_invoker) + created = _create_while_recording_fails(service, monkeypatch, _request()) + [item_id] = created.item_ids + item = queue.dequeue() + assert item is not None + queue.fail_queue_item(item_id, "ValueError", "out of memory", "traceback") + queue.retry_items_by_id("default", [item_id]) + queue.delete_queue_items_by_id([item_id]) + [retry] = _calibration_item_ids(queue) + + run = _run_status(service.get_status(), created.run_id) + + assert run.item_ids == [item_id, retry] + + +def test_a_failed_replay_of_the_keys_receipt_keeps_the_runs_mark( + tmp_path: Path, + mock_invoker: Invoker, + queue: SessionQueue, + mock_sqlite_database: Database, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The first attempt queued the batch and could not record its items, which were pruned. A repeat replays the + receipt and its enqueue raises: that attempt queued nothing, but the first may have.""" + service = _service(tmp_path, mock_invoker) + with pytest.raises(CalibrationStoreError): + _create_while_recording_fails(service, monkeypatch, _request(key="replayed")) + queue.delete_queue_items_by_id(_calibration_item_ids(queue)) + + async def failing_enqueue(*args: object, **kwargs: object) -> None: + raise RuntimeError("the connection dropped") + + with monkeypatch.context() as patch: + patch.setattr(queue, "enqueue_batch", failing_enqueue) + with pytest.raises(RuntimeError): + _create(service, key="replayed") + _collect_receipt(queue, mock_sqlite_database, "replayed") + + repeat = _create(service, key="replayed") + + assert repeat.item_ids == [] and _calibration_item_ids(queue) == [] + + +def test_a_found_run_whose_enqueue_found_no_room_is_enqueued_by_a_later_repeat( + tmp_path: Path, + mock_invoker: Invoker, + queue: SessionQueue, + mock_sqlite_database: Database, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A run recorded without its items and without the mark (as a store written before the mark existed has it), + with nothing queued: the repeat that marks it and finds the queue full clears the mark again.""" + service = _service(tmp_path, mock_invoker) + with pytest.raises(CalibrationStoreError): + _create_while_recording_fails(service, monkeypatch, _request((1, 2), key="unmarked")) + queue.delete_queue_items_by_id(_calibration_item_ids(queue)) + _collect_receipt(queue, mock_sqlite_database, "unmarked") + path = _store_path(tmp_path) + stored = json.loads(path.read_bytes()) + del stored["runs"][0]["enqueue_attempted"] + path.write_bytes(json.dumps(stored).encode()) + + mock_invoker.services.configuration.max_queue_size = 1 + with pytest.raises(CalibrationQueueFullError): + _create(service, values=(1, 2), key="unmarked") + mock_invoker.services.configuration.max_queue_size = 10 + repeat = _create(service, values=(1, 2), key="unmarked") + + assert len(repeat.item_ids) == 2 and _calibration_item_ids(queue) == repeat.item_ids From 9f2345c5db8a1fd51cce4896a47ceef9f85491c9 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 22 Oct 2026 00:13:13 +0200 Subject: [PATCH 8/9] fix(working-memory): serialize calibration run creation across processes A keyed run's creation holds an OS file lock, striped over 256 files, from preparing the run to recording its items, behind a per-run in-process gate. The final observation keeps a recorded completion of an item the queue no longer has. --- .../Architecture/model-manager.mdx | 4 +- .../working_memory_calibration_default.py | 56 +++++-- .../working_memory_calibration_store.py | 22 +++ ...test_working_memory_calibration_service.py | 137 ++++++++++++++++++ 4 files changed, 204 insertions(+), 15 deletions(-) diff --git a/docs/src/content/docs/development/Architecture/model-manager.mdx b/docs/src/content/docs/development/Architecture/model-manager.mdx index e34df87dd26..e454b546f0f 100644 --- a/docs/src/content/docs/development/Architecture/model-manager.mdx +++ b/docs/src/content/docs/development/Architecture/model-manager.mdx @@ -1043,7 +1043,7 @@ A calibration run measures one model's operations on this machine's GPU. Runs on - 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, that request is refused (409) and the first run stays; the receipt is checked before the request's model, so an unknown model does not turn the refusal into a 404. 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: the run is marked before its batch is enqueued (`enqueue_attempted`, set as the run is recorded), and only the attempt that set the mark clears it again, when its own enqueue queued nothing; a replay of the key's receipt never clears it, since the attempt that queued the batch may have queued it. A repeat that finds a marked run without items, without the receipt and with no item in the queue or seen by the runner answers with the run and no items, and the run reads as canceled. That is the case when the items were cleared or pruned before the runner saw them, when the server shut down between the mark and the enqueue, or when an enqueue raised, the queue could not be read to tell what it had queued, and it had queued nothing. Should clearing the mark fail (it is logged), the run stays marked although nothing was queued, and the key cannot start it; a request with another key can. 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. + - 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. For a run with a retry key, "one at a time" holds across every process serving the root: the creation holds an OS file lock (the store's own lock primitive) from finding or recording the run to recording its items, so another process's repeat waits and then finds what the first attempt did. A holder that dies releases the lock with its process. A repeat waits up to 30 s (`LOCK_TIMEOUT_S`) and then answers 503, which with a key can be retried. The run ids share 256 empty lock files, `/working_memory/creations/.lock` by the first two hex digits of the id's SHA-256. No lock file is ever removed (removing one could let two processes each lock a file of that name), and the folder never holds more than 256. Unrelated runs on one file wait for each other, which costs little for a short administrator's action. Within one process the repeats of a run first wait on an in-process lock per run, on the event loop, so only one of them at a time takes a worker thread to wait for the file lock: a burst of repeats cannot fill the default executor. A run without a key has a fresh id that no other request names, so it takes no lock. 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: @@ -1051,7 +1051,7 @@ A calibration run measures one model's operations on this machine's GPU. Runs on - 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. A session that ran every node without an error, and that nothing canceled or failed (a failed save of its measurements fails the item), is recorded as completed before the queue completes the item, so that a process stopping in between cannot leave the item to read as canceled once the queue has pruned it; the queue's status then replaces that record. 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. +5. After the item, the runner records the item's outcome in the run, even when the session raised. A session that ran every node without an error, and that nothing canceled or failed (a failed save of its measurements fails the item), is recorded as completed before the queue completes the item, so that a process stopping in between cannot leave the item to read as canceled once the queue has pruned it; the queue's status then replaces that record. Only a queue row replaces it: once the queue has pruned the item, the final record leaves the completion as it is instead of recording the missing item as canceled. A row that still exists and says failed or canceled wins. 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 `/working_memory/.json` (`CalibrationStoreFile`, `schema_version` 1). It is machine data, so it lives next to the database rather than inside it. It holds: diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py index b6bbcaf19f6..055f8bc061e 100644 --- a/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py @@ -4,7 +4,8 @@ import json import uuid import weakref -from collections.abc import Callable, Mapping +from collections.abc import AsyncIterator, Callable, Mapping +from contextlib import ExitStack, asynccontextmanager from functools import partial from itertools import chain from typing import TYPE_CHECKING, Final, NamedTuple, Optional @@ -220,8 +221,8 @@ def __init__( self._run_files: dict[str, str] = {} # Creations in flight, referenced here so that one whose request was canceled still runs to its end. self._creations: set[asyncio.Task[WorkingMemoryCalibrationRunCreated]] = set() - # Per run id while a creation uses it: a repeat of a request waits for the attempt still in flight. - self._run_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() + # Per run id while a creation of this process uses it, in front of the file lock (see `_creation_lock`). + self._run_gates: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() self._invoker: Optional["Invoker"] = None def start(self, invoker: "Invoker") -> None: @@ -334,11 +335,35 @@ def _identify(model_key: str, request: WorkingMemoryCalibrationRequest, user_id: graph.id = run_id return run_id - def _run_lock(self, run_id: str) -> asyncio.Lock: - lock = self._run_locks.get(run_id) - if lock is None: - lock = self._run_locks[run_id] = asyncio.Lock() - return lock + @asynccontextmanager + async def _creation_lock(self, run_id: str) -> AsyncIterator[None]: + """Hold the run's creation lock, which every process serving the root takes (`CalibrationStore.creation_lock`): + a repeat of the request waits for the attempt in flight, in this process or another, and then finds what it + did. Only a retry key's run id can be contended; a fresh one is not locked. + + In front of it, the repeats of one process wait on the run's gate, on the event loop: only one of them at a + time takes a worker thread to wait for the file lock. A burst of repeats would otherwise fill the default + executor, stalling the holder's own worker-thread steps and a shutdown that waits for the executor.""" + if not _is_keyed(run_id): + yield + return + gate = self._run_gates.get(run_id) + if gate is None: + gate = self._run_gates[run_id] = asyncio.Lock() + async with gate: + held = ExitStack() + # A file lock: taken in a worker thread, which waits for another holder without blocking the loop. + acquiring = asyncio.ensure_future(asyncio.to_thread(held.enter_context, self._store.creation_lock(run_id))) + try: + await asyncio.shield(acquiring) + except asyncio.CancelledError: + # The worker may still take the lock: give it back once it has. + acquiring.add_done_callback(lambda _: held.close()) + raise + try: + yield + finally: + held.close() async def create_run( self, model_key: str, request: WorkingMemoryCalibrationRequest, user_id: str @@ -366,9 +391,9 @@ async def _create_run( self, model_key: str, request: WorkingMemoryCalibrationRequest, user_id: str ) -> WorkingMemoryCalibrationRunCreated: run_id = await asyncio.to_thread(self._identify, model_key, request, user_id) - # Recording the run, enqueuing its batch and recording its items or removing it: one step for each run id. The - # lock is this process's; two processes serving one root and its SQLite database are not kept apart. - async with self._run_lock(run_id): + # Recording the run, enqueuing its batch and recording its items or removing it: one step for each run id, in + # every process serving the root. + async with self._creation_lock(run_id): prepared = await asyncio.to_thread(self._prepare_run, model_key, request, user_id, run_id) if isinstance(prepared, WorkingMemoryCalibrationRunCreated): return prepared @@ -595,7 +620,8 @@ def admit(self, queue_item: SessionQueueItem, run_id: str) -> Optional[str]: ) return None - def _observe(self, run_id: str, item_id: int, observation: ObservedItem) -> None: + def _observe(self, run_id: str, item_id: int, observation: ObservedItem, *, over_completed: bool = True) -> None: + """Record what the runner saw of the item; without `over_completed`, not over a completion it recorded.""" found = self._find_run(run_id) if found is None: return @@ -604,6 +630,8 @@ def record(file: CalibrationStoreFile) -> tuple[bool, None]: for run in file.runs: if run.run_id == run_id: previous = run.observed.get(item_id) + if not over_completed and previous is not None and previous.status == "completed": + return False, None if observation.retried_from_item_id is None and previous is not None: observation_ = observation.model_copy( update={"retried_from_item_id": previous.retried_from_item_id} @@ -620,7 +648,8 @@ def observe(self, queue_item: SessionQueueItem, run_id: str, error: Optional[Bas try: current = self._services.session_queue.get_queue_item(queue_item.item_id) except SessionQueueItemNotFoundError: - # Removed while it ran (Clear Queue, a deleted item): it will not finish. + # Removed while it ran (Clear Queue, a deleted item): it will not finish. Unless the runner recorded its + # completion: the queue then completed it and pruned it since. A row that still exists has the last word. self._observe( run_id, queue_item.item_id, @@ -630,6 +659,7 @@ def observe(self, queue_item: SessionQueueItem, run_id: str, error: Optional[Bas retried_from_item_id=queue_item.retried_from_item_id, observed_at=_now(), ), + over_completed=False, ) return status, error_message = current.status, current.error_message diff --git a/invokeai/app/services/working_memory_calibration/working_memory_calibration_store.py b/invokeai/app/services/working_memory_calibration/working_memory_calibration_store.py index f865b16571a..1e7d6a8d24f 100644 --- a/invokeai/app/services/working_memory_calibration/working_memory_calibration_store.py +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_store.py @@ -7,6 +7,7 @@ class that does not match its name -- is reported and never written, moved or deleted. """ +import hashlib import json import os import re @@ -193,6 +194,27 @@ def _locked(self, device_class_id: str, *, writing: bool) -> Iterator[None]: with handle, _os_lock(handle): yield + @contextmanager + def creation_lock(self, run_id: str) -> Iterator[None]: + """Hold the OS lock of one run's creation, which every process serving this root takes on the same file. A + holder that dies releases it with its process. Waits `LOCK_TIMEOUT_S` for another holder, then raises + `CalibrationStoreError`. + + The run ids share 256 lock files, `creations/.lock` by the first two hex digits of the id's SHA-256: the + folder never holds more, so no file is ever removed (removing one could let a process that opened it before the + removal and one that creates it anew each hold a lock of their own). Unrelated runs on one file wait for each + other, which costs little: a creation is a short administrator's action. No store file has such a name. + """ + stripe = hashlib.sha256(run_id.encode("utf-8")).hexdigest()[:2] + lock_path = self._dir / "creations" / f"{stripe}.lock" + try: + lock_path.parent.mkdir(parents=True, exist_ok=True) + handle = open(lock_path, "a+b") + except OSError as error: + raise CalibrationStoreError(f"Cannot open {lock_path}: {error}") from error + with handle, _os_lock(handle): + yield + def _read_unlocked(self, device_class_id: str) -> Optional[StoreRead]: path = self.path_for(device_class_id) try: diff --git a/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py b/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py index 2374026c343..e2623297791 100644 --- a/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py +++ b/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py @@ -7,6 +7,7 @@ import stat import threading import time +from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace from typing import Optional @@ -972,3 +973,139 @@ def test_a_found_run_whose_enqueue_found_no_room_is_enqueued_by_a_later_repeat( repeat = _create(service, values=(1, 2), key="unmarked") assert len(repeat.item_ids) == 2 and _calibration_item_ids(queue) == repeat.item_ids + + +def _in_thread(work) -> tuple[threading.Thread, list]: + """Starts `work` on a thread of its own; the list receives what it returned or raised.""" + outcome: list = [] + + def run() -> None: + try: + outcome.append(work()) + except BaseException as error: # noqa: BLE001 - returned to the test + outcome.append(error) + + thread = threading.Thread(target=run) + thread.start() + return thread, outcome + + +def test_a_repeat_from_another_process_waits_for_the_attempt_that_is_enqueuing( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + """The first process recorded and marked the run and is enqueuing it. Two service instances on one root stand for + two processes: they share nothing in memory, only the root's files and the database.""" + first, second = _service(tmp_path, mock_invoker), _service(tmp_path, mock_invoker) + held = HeldEnqueue(monkeypatch) + creating, created = _in_thread(lambda: _create(first, key="two-processes")) + assert held.entered.wait(timeout=10) + + repeating, repeated = _in_thread(lambda: _create(second, key="two-processes")) + repeating.join(timeout=0.5) + waited = repeating.is_alive() + held.release.set() + creating.join(timeout=30) + repeating.join(timeout=30) + + assert waited, f"the repeat did not wait for the attempt in flight; it answered {repeated!r}" + assert repeated == created and len(created[0].item_ids) == 1 + assert _calibration_item_ids(queue) == created[0].item_ids + + +def test_a_request_from_another_process_cannot_prepare_while_the_first_attempt_records_the_run( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + """Prepared side by side, both attempts would enqueue: the one resuming after the queue collected the other's + receipt queues the batch again, and its items replace the other's.""" + first, second = _service(tmp_path, mock_invoker), _service(tmp_path, mock_invoker) + first_store = first._store + entered, release = threading.Event(), threading.Event() + change = CalibrationStore.change + + def held(self: CalibrationStore, *args, **kwargs): + # The first change of the first process's store records the run. + if self is first_store and not entered.is_set(): + entered.set() + assert release.wait(timeout=10) + return change(self, *args, **kwargs) + + monkeypatch.setattr(CalibrationStore, "change", held) + creating, created = _in_thread(lambda: _create(first, key="side-by-side")) + assert entered.wait(timeout=10) + + repeating, repeated = _in_thread(lambda: _create(second, key="side-by-side")) + repeating.join(timeout=0.5) + waited = repeating.is_alive() + release.set() + creating.join(timeout=30) + repeating.join(timeout=30) + + assert waited, f"the other request prepared while the first recorded the run; it answered {repeated!r}" + assert repeated == created and len(created[0].item_ids) == 1 + assert _calibration_item_ids(queue) == created[0].item_ids + [stored] = json.loads(_store_path(tmp_path).read_bytes())["runs"] + assert stored["queue_item_ids"] == created[0].item_ids + + +@pytest.mark.parametrize(("afterwards", "outcome"), [("pruned", "completed"), ("failed-in-queue", "failed")]) +def test_the_final_observation_keeps_a_recorded_completion_of_an_item_the_queue_no_longer_has( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, afterwards: str, outcome: str +) -> None: + """The runner recorded the completion, then the queue completed the item and pruned it before the runner's last + record. A queue row that still exists and failed the item has the last word.""" + service = _service(tmp_path, mock_invoker) + created = _create(service) + item = queue.dequeue() + assert item is not None and service.admit(item, created.run_id) is None + service.observe_completed(item, created.run_id) + if afterwards == "pruned": + queue.complete_queue_item(item.item_id) + queue.prune("default") + else: + queue.fail_queue_item(item.item_id, "RuntimeError", "the session ended badly", "traceback") + + service.observe(item, created.run_id) + + assert _run_status(service.get_status(), created.run_id).status == outcome + + +def test_repeats_in_one_process_leave_one_worker_waiting_for_the_file_lock( + tmp_path: Path, mock_invoker: Invoker, queue: SessionQueue, monkeypatch: pytest.MonkeyPatch +) -> None: + """The other repeats wait on the event loop: a burst would otherwise hold a worker thread each, polling the lock.""" + service = _service(tmp_path, mock_invoker) + held = HeldEnqueue(monkeypatch) + count = threading.Lock() + inside, most = [0], [0] + creation_lock = service._store.creation_lock + + @contextmanager + def counted(run_id: str): + with count: + inside[0] += 1 + most[0] = max(most[0], inside[0]) + try: + with creation_lock(run_id): + yield + finally: + with count: + inside[0] -= 1 + + monkeypatch.setattr(service._store, "creation_lock", counted) + + async def burst(): + first = asyncio.create_task(service.create_run("model-a", _request(key="burst"), "system")) + assert await asyncio.to_thread(held.entered.wait, 10) + repeats = [ + asyncio.create_task(service.create_run("model-a", _request(key="burst"), "system")) for _ in range(8) + ] + # Half a second, as `tests/fixtures/races.waiting` gives a change that must wait. + await asyncio.wait(repeats, timeout=0.5) + most_while_held = most[0] + held.release.set() + return most_while_held, await first, await asyncio.gather(*repeats) + + most_while_held, first, repeats = asyncio.run(burst()) + + assert most_while_held == 1 + assert all(repeat == first for repeat in repeats) and _calibration_item_ids(queue) == first.item_ids From 62d75d83eb2af9d0d108569e35394d99c9740020 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Thu, 22 Oct 2026 01:03:57 +0200 Subject: [PATCH 9/9] test(processor): release the wake-up test's gates before a stop that waits for its workers Stop now joins the workers, so the wake-up test opens its gates before stopping and runs the stop case on a thread of its own. --- .../session_processor/test_worker_wakeup.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/tests/app/services/session_processor/test_worker_wakeup.py b/tests/app/services/session_processor/test_worker_wakeup.py index 0eaa31a3a25..13a6111299b 100644 --- a/tests/app/services/session_processor/test_worker_wakeup.py +++ b/tests/app/services/session_processor/test_worker_wakeup.py @@ -50,10 +50,11 @@ def _poll_events(processor): def _stop_and_join(processor, gates=()): try: - if hasattr(processor, "_stop_event"): - processor.stop() + # Opened first: stop() waits for the workers, and a gated worker exits only once its gate opens. for gate in gates: gate.set() + if hasattr(processor, "_stop_event"): + processor.stop() for event in _poll_events(processor).values(): event.set() for worker in processor._workers: @@ -130,6 +131,8 @@ def run_item(_lifecycle, _item): monkeypatch.setattr(WorkflowCallQueueLifecycle, "run_queue_item", run_item) + stopper = threading.Thread(target=processor.stop, name="stopper", daemon=True) + def notify(): nonlocal available if producer == "release": @@ -139,7 +142,11 @@ def notify(): available = True asyncio.run(processor._on_batch_enqueued(("batch_enqueued", None))) else: - processor.stop() + # stop() waits for the workers, which these gates hold: it runs beside the test, which goes on once every + # worker has been notified, as a stop() that returned at once had done before B's clear. + stopper.start() + for name, event in _poll_events(processor).items(): + assert event.wait(20), f"stop did not notify {name}" try: processor.start(invoker) @@ -167,7 +174,10 @@ def notify(): assert claims == [str(devices[0])] finally: _stop_and_join(processor, gates.values()) + if stopper.ident is not None: + stopper.join(20) + assert not stopper.is_alive() assert all(worker.thread is not None and not worker.thread.is_alive() for worker in processor._workers) assert errors == [], f"worker coordination failed: {errors}" invoker.services.logger.error.assert_not_called()