diff --git a/docs/src/content/docs/development/Architecture/model-manager.mdx b/docs/src/content/docs/development/Architecture/model-manager.mdx index dd3a5680b95..970bfeb6454 100644 --- a/docs/src/content/docs/development/Architecture/model-manager.mdx +++ b/docs/src/content/docs/development/Architecture/model-manager.mdx @@ -1020,6 +1020,78 @@ To write such an estimator (`invokeai/backend/util/vae_working_memory.py` has fi - Ship one coefficient entry per measured platform value in `invokeai/backend/working_memory/data/`, either with its own `values` or with `values_from` another estimator's entry. Unmeasured platform values (MPS, XPU, CPU) are priced as `UNMEASURED_PRICED_AS` in `working_memory/platform.py` says. - Ship the measured points next to them: peak reserved bytes (`reserved_peak`), one fresh process per point. `tests/backend/working_memory/` holds each estimate to at least its point, and where a fitted term decides it, to at most 1.3 times the point, or the point's tighter `max_ratio`. The `device_working_mem_gb` floor counts only where an analytic term decides. Every value an entry carries must be decided by at least one point. +#### Measuring estimates: the probe + +`invokeai/backend/working_memory/probe.py` measures what an estimated operation really takes. It is active only inside a `calibration_scope(...)`, which the session runner opens per node of a calibration queue item; without one, `model_on_device()` behaves as above and reads no counter. + +Inside a scope, every `model_on_device(working_memory=...)` on a GPU is a window. The lock reserves the estimate plus calibration headroom (half the estimate, at least 1 GiB), so a low estimate cannot spill. The window runs from the settled lock to the unlock and collects a `MeasuredWindow` in the scope's `windows`; nothing is written while it is open. + +- Counters (`working_memory/counters.py`): torch's reserved and allocated peaks, and this process's usage as the driver counts it. That is D3DKMT on Windows for CUDA and ROCm (the only readers verified on hardware), NVML on Linux CUDA, DRM fdinfo on Linux ROCm, and the MPS driver figure. XPU's Sysman figure covers the whole device, so it only detects another process's load. MPS has no peaks; a sampler thread keeps them while a window is open. +- Demand: per counter, the largest segment peak less the baseline and the bytes co-locks moved onto the device; `demand_bytes` is the largest of those deltas. Each counter misses something: the driver misses HIP's retained pool, and reserved misses cuBLAS and cuDNN workspaces and, under `cudaMallocAsync`, part of the allocated peak. +- Segments: the window opens after a release, so cached blocks and tensors from before the lock form the baseline. Another model locked on the same thread inside the window is a co-lock: its weights go into the offset, and its lock interval is kept out of the peak. An estimated window inside another measures alone while the outer one is frozen. A `TorchDevice.empty_cache()` ends a segment. A direct `torch.cuda.empty_cache()` call (from a library, or the one cuDNN and MIOpen run after an algorithm search) and an offload by the cache's background thread do not: the driver figure can miss what they released, while the reserved and allocated peaks still bound the demand. +- Invalid (`valid=False`, with `reasons`): another thread locked a model on the same device's cache during the window, held one when it opened, or was still locking one then (`foreign_lock`); the allocator failed or retried; Windows still holds more than 256 MiB of the process in shared memory after a short settle; the operation raised or was canceled, or its node failed or was canceled after the window closed; the Krea-2 SDPA benchmark is active; XPU's device-wide usage grew beyond this process's own (`foreign_device_load`); a probe step failed (`probe_error`: a reading, a release, a peak reset or a synchronization, which also voids the scope's later windows when it hit the lock bookkeeping; a window that cannot open is recorded invalid, and its operation runs unmeasured); or no counter could be read. +- Not fit-eligible: a model that was not fully resident when the window opened. The measurement stays valid, but streamed weights are working memory the estimate does not price. + +`working_memory/fingerprint.py` keys measurements. `device_class_for()` gives the same id to identical cards. `environment_for()` splits hard facts, which change an operation's memory and make up its id (the operating system as `windows`, `linux` or `other`, torch, runtime and convolution-library versions, allocator backend and configuration, the ROCm SDPA guard, the memory settings passed in), from soft facts kept for diagnosis (full versions, driver, OS build). + +#### Calibration runs + +A calibration run measures one model's operations on this machine's GPU. Runs only record; no estimate reads them yet. The service is `invokeai/app/services/working_memory_calibration/`, the routes are in `invokeai/app/api/routers/working_memory.py`, and every route is admin-only (403 otherwise). + +1. `POST /api/v2/models/i/{key}/working_memory/calibrate` takes a client-built graph, batch data and an optional `idempotency_key`. The server turns off `use_cache` on every node, sets `runs=1`, sets the graph's id to the run id, and sets the origin to `invokeai:wm-calibration:` and the destination to `invokeai:wm-calibration`. It records the run, enqueues the batch as the requesting admin, then records the run's items. + - The queue takes all of the batch's items or none (`enqueue_batch(..., all_or_none=True)`, checked under its admission lock): a run with only some of its cases would measure only some of them. Without room for all of them nothing is enqueued, the key stays unused, and the run is removed if this request recorded it. + - With a key, the run id is derived from the owner, the key and the request's fingerprint: the model key, the graph without its id, and the batch data. A client that sends no graph id would otherwise get a fresh one each time and could never repeat a request. A repeated request finds its run and gets the same response; nothing new is enqueued, even after the queue collected the key's receipt. Once a reset removed the run, a repeat is refused (404) whether or not the queue still keeps the receipt: the store keeps the ids of removed runs, so the run is never started again and its queued items keep failing at admission. Another request that reuses the key (another model, graph or batch data) gets another run id and so another origin. While the key's receipt exists, the queue refuses that request (409) and the first run stays. Once the queue has collected the acknowledged receipt, the request starts a run of its own. Acknowledge the key through `POST /api/v1/queue/{queue_id}/enqueue_batch/acknowledge`. + - Once its items are queued, the request answers with them, also when the enqueue raised after its commit (the queue announces a batch after committing it). When the queue cannot tell whether it holds them, the run is kept and the enqueue's own error answers. + - If the store fails for a passing reason (an unwritable root, a lock timeout) to record the queued items, a request without a key still succeeds: its items run, and the status finds them by their origin. With a key that is a 503, and a repeat records them: from the key's receipt while the queue keeps it, otherwise from the queue items with the run's origin and the items the runner saw. It never enqueues a second batch for the run. Recording them from the receipt checks neither the model nor the device, so a model removed meanwhile does not fail the repeat; a request that has to enqueue needs the model. A run removed meanwhile (404) or an unusable store file (409) fails the request with or without a key; its queued items then fail at admission. + - Per run id, one creation runs at a time, and it runs to its end even when its request is canceled: the enqueue runs in a worker thread and can commit after the cancellation. Only the server's shutdown cancels the creation with it; the run then stays as far as it got, and with a key a repeat of the request completes it. The lock belongs to the process, so two processes serving one root and its SQLite database can interleave the creations of one run. A run is removed only when this request recorded it and nothing of it was queued. + - The destination only stops clients from routing the results to a gallery or canvas view. Images the graph's nodes save are stored like those of any other session: on the board and with the intermediate flag each node sets. +2. `enqueue_batch` refuses any origin that starts with that prefix, in any case, for every user (400). Only the calibration route creates calibration items. Retries and workflow-call children inherit origin and owner from their source item. +3. `DefaultSessionRunner.run()` checks each item's origin once. For a calibration item it calls `admit()` before any node runs, also for one left pending before a restart: the invoker starts the session processor after every other service. (The queue's start cancels the items that were in progress.) Retries and children are checked here too. Otherwise the item fails with the reason. The checks: + - the owner is an active admin (single-user mode counts as admin); + - the machine has exactly one GPU generation device; + - the run exists for that GPU's device class. +4. Each admitted node runs in `probe.calibration_scope(...)`, keyed by the node's source id. Its windows are persisted after the node completes, fails or is canceled. Unless the node completed, every one of them is persisted invalid (`exception` when it failed, `canceled` otherwise), including a window whose own operation finished before the node raised: `run_node` swallows the failure or cancellation, so the runner takes the outcome from its node hooks. Windows from a device of another class, or of a run that was removed, are dropped with a warning. So are windows of a model reset after the node started, also by another process sharing the root: the runner reads the device class's reset epochs before the node's scope opens, and the write drops the points of every model whose epoch has moved since. Without that, a reset while the node runs (of a VAE inside a transformer's run, say) would be undone by the node's later write. If they cannot be saved, the item fails and the session stops. Every other item gets no scope and runs as before. +5. After the item, the runner records the item's outcome in the run, even when the session raised. An item the queue no longer has (Clear Queue) is recorded as canceled. An item that was still running when the runner raised is recorded as failed with the error. + +The store has one JSON file per device class at `/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; +- the ids of the runs with a retry key that a reset removed (`removed_run_ids`), so that a repeat of the request that started one cannot start it again. They are kept for good, about 40 bytes per removed run. A run without a key gets a fresh id and is never repeated, so its id is not kept; +- per model key, its reset epoch (`model_reset_epochs`): how often the model was reset. The reset raises it in the same file change, so it is only written when the reset is. + +A file without either of these two fields reads them as empty. + +Writes and reads: + +- Every change takes a process lock and an OS lock on the file's `.lock` sibling, re-reads the file, applies the change, and replaces the file atomically. Two processes sharing a root therefore never lose each other's writes. +- On Windows a process that ignores the lock (an antivirus scanner, a search indexer, an editor) can hold the file open and refuse the replace. The write then retries a few times with backoff, about 1.5 s in all (`REPLACE_ATTEMPTS`, `REPLACE_RETRY_S`). +- Readers take the lock as well, because Windows cannot replace a file another process has open. +- Each change rewrites the whole file, which is fine at calibration volumes. + +A file is unusable when it has another schema version, invalid content, an id that does not match its fields or its file name, or a field this version does not know. Every store model forbids unknown fields: ignoring one would drop it at the next rewrite. An unusable file is reported in the status, makes calibration unavailable for its device class, and is never written, moved or deleted. + +A run's status is derived, never stored: + +- The live queue items with the run's origin come first. For items the queue has pruned, the runner's recorded statuses are used. An item the queue no longer has and that the runner never saw finish (the process died mid-item) counts as canceled. +- A retry replaces the item it retried. The queue's own record of the retry links them while it has the retry, also before the runner started it; afterwards, what the runner saw does. A retry the runner never started and the queue no longer has leaves no trace: the run then reads as the outcome of the item it retried, failed for instance. +- The run is `running` while any item is pending or running, `failed` when any failed, and `completed` when all completed. Otherwise it is `partial` if valid points were recorded, or `canceled` if not: an invalid point measured nothing usable. + +Routes, besides the 403 for non-admins and a 503 when the store cannot be read or written (an unwritable root, a lock timeout): + +| Route | Does | Other errors | +| --- | --- | --- | +| `POST .../i/{key}/working_memory/calibrate` | starts a run (above) | 404 unknown model, the run was removed while it was being created, or a reset removed the run the retry key started; 409 no GPU or several generation devices, unusable store file, no room in the queue for all of the run's items (nothing is enqueued; the run is removed if this request recorded it), or a retry key used for another request; the 404 and 409 for the run's store, and with a key the 503, can come after the items were queued (above); 422 a `call_saved_workflow` node, batch data setting `use_cache`, an invalid batch, or fields the server sets; 429 too many unacknowledged retry keys | +| `GET .../working_memory/status` | every device class with its runs, point counts per estimator, version, validity and fit eligibility, and environments with hard and soft facts compared to the current environment; `available` says whether a run can start now (one GPU generation device and a usable store file) | – | +| `DELETE .../working_memory/runs/{run_id}` | removes a run and its points; its items that have not started yet fail at admission, also after a repeat of the request that started it | 404 no usable file holds the run | +| `DELETE .../i/{key}/working_memory` | removes the model's runs with all their points, and every point of the model taken in other runs (a VAE measured inside a transformer's run, say) | 409 no current device class, or an unusable file | +| `DELETE .../working_memory` | removes every run, point and environment of a device class | 409 as above | +| `GET .../working_memory/export` | the file byte for byte: a `CalibrationStoreFile` when this version can read it, otherwise whatever the file holds | 404 no file; 409 no current device class | + +The last three act on the current device class unless `device_class_id` is given. + ### Emitting model loading events When the `context` argument is passed to `load_model_*()`, it will retrieve the invocation event bus from the passed `InvocationContext` object to emit events on the invocation bus. The two events are "model_load_started" and "model_load_completed". Both carry the following payload: diff --git a/docs/src/content/docs/development/Architecture/overview.mdx b/docs/src/content/docs/development/Architecture/overview.mdx index 8cbe18efad5..3d6272db7ff 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()` stops the session processor first: its workers stop taking items and cancel their sessions, and the other services stop only once the workers have exited. The processor waits up to 30 seconds for all its workers together, and after a second still waiting it logs which worker and queue item it waits for. A worker still in a node at that deadline (one without a cancel point) is given up on with an error in the log, so that shutdown finishes; a later `stop()` does not wait for it again. It keeps its GPU, so until it exits no session processor in the same process starts and none replaces the generation device pool: `ApiDependencies.initialize` refuses a restart before it opens the database. + ### Sessions Invocations and links between them form a graph, which is maintained in a session. Sessions can be queued for invocation, which will execute their graph (either the next ready invocation, or all invocations). Sessions also maintain execution history for the graph (including storage of any outputs). An invocation may be added to a session at any time, and there is capability to add and entire graph at once, as well as to automatically link new invocations to previous invocations. Invocations can not be deleted or modified once added. diff --git a/invokeai/app/api/dependencies.py b/invokeai/app/api/dependencies.py index f11b206427e..ee776633b79 100644 --- a/invokeai/app/api/dependencies.py +++ b/invokeai/app/api/dependencies.py @@ -55,6 +55,7 @@ from invokeai.app.services.session_processor.session_processor_default import ( DefaultSessionProcessor, DefaultSessionRunner, + refuse_while_an_abandoned_worker_runs, ) from invokeai.app.services.session_queue.session_queue_default import SessionQueue from invokeai.app.services.shared.database.database import Database @@ -70,6 +71,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 @@ -117,6 +122,10 @@ def initialize( # lazily inside this function for the same reason, one step earlier: it is what fills the # registry, and importing this module is all `scripts/generate_openapi_schema.py` does. validate_architectures() + # A worker the previous invoker's stop gave up on still holds its GPU, and the session processor would refuse + # to start beside it. Refused here, before the database opens (and migrates), so a refused restart leaves + # nothing open behind it. + refuse_while_an_abandoned_worker_runs() logger.info(f"InvokeAI version {__version__}") logger.info(f"Root directory = {str(config.root_path)}") @@ -216,6 +225,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,8 +286,13 @@ def initialize( fonts=fonts, intermediates=intermediates, gallery_maintenance=gallery_maintenance, + 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 2d3c0221abc..25a25c4d8ef 100644 --- a/invokeai/app/api/routers/session_queue.py +++ b/invokeai/app/api/routers/session_queue.py @@ -41,6 +41,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"]) @@ -256,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( @@ -265,6 +273,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..b11ad76277f --- /dev/null +++ b/invokeai/app/api/routers/working_memory.py @@ -0,0 +1,241 @@ +"""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; 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" + }, + 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); with an " + "idempotency key, possibly after the items were queued" + }, + }, +) +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. 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. 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) + except ( + UnknownModelException, + CalibrationRunNotFoundError, + 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 4802132c2b6..7b359a9e02e 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 @@ -786,6 +787,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 f5416b02bb6..618de627229 100644 --- a/invokeai/app/services/invocation_services.py +++ b/invokeai/app/services/invocation_services.py @@ -56,6 +56,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 @@ -117,6 +120,7 @@ def __init__( progress_previews: "ProgressPreviewsBase | None" = None, fonts: "FontService | None" = None, gallery_maintenance: "GalleryMaintenanceService | None" = None, + working_memory_calibration: "WorkingMemoryCalibrationServiceBase | None" = None, board_access: "BoardAccessService | None" = None, ): self.board_images = board_images @@ -170,6 +174,7 @@ def __init__( self.fonts = fonts self.intermediates = intermediates self.gallery_maintenance = gallery_maintenance + self.working_memory_calibration = working_memory_calibration # Stateless over the database, so a real one is the default: no caller should get looser rules by leaving it # out. self.board_access = board_access or BoardAccessService(database) diff --git a/invokeai/app/services/invoker.py b/invokeai/app/services/invoker.py index b504f511ba7..8c01125a05b 100644 --- a/invokeai/app/services/invoker.py +++ b/invokeai/app/services/invoker.py @@ -22,13 +22,25 @@ 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 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 4308e7a4f35..b89b7517f64 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, @@ -45,6 +46,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 +54,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. @@ -64,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 @@ -165,6 +189,11 @@ 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 + # 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 @@ -203,7 +232,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 +278,104 @@ 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; 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, + # 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._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 + try: + service.record_windows(run_id, scope.windows, model_reset_epochs) + except Exception as error: + 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 @@ -573,6 +699,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( @@ -587,6 +714,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, @@ -599,12 +727,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( @@ -708,20 +838,30 @@ def __init__( self._thread_limit = thread_limit self._polling_interval = polling_interval self._workers: list[_SessionWorker] = [] + # 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. - def _resolve_devices(self) -> list[Optional[torch.device]]: - """Determine the per-worker devices from config. + 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). - 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). + `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`). """ - generation_devices = self._invoker.services.configuration.generation_devices - devices = TorchDevice.get_generation_devices(generation_devices) - if not devices: - return [None] - return list(devices) + 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]) + return self._devices def _clone_session_runner(self, template: SessionRunnerBase) -> SessionRunnerBase: """Create an independent runner for an additional worker. @@ -749,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() @@ -760,11 +903,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) @@ -827,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 @@ -838,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 3e7bc4841fb..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, @@ -54,9 +55,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 @@ -290,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 01df20c0052..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, @@ -144,11 +145,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 +160,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 +211,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 +229,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: @@ -1037,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/__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..63207fd6c78 --- /dev/null +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_base.py @@ -0,0 +1,69 @@ +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 +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. + + 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; 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 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: + """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. 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: + """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..e5ee262ca87 --- /dev/null +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_common.py @@ -0,0 +1,280 @@ +"""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) + 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 + +# region API + + +class WorkingMemoryCalibrationRequest(BaseModel): + """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") + + 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 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.", + ) + + +class WorkingMemoryCalibrationRunCreated(BaseModel): + run_id: str + 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" + ) + + +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 " + "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") + 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, 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} {self.reason}" + + +class CalibrationQueueFullError(Exception): + """The queue had no room for all of the run's items, so it took none.""" + + +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..5f59115cefa --- /dev/null +++ b/invokeai/app/services/working_memory_calibration/working_memory_calibration_default.py @@ -0,0 +1,841 @@ +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 +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, + QueueItemLineage, + 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, + CalibrationStoreError, + CalibrationStoreFile, + CalibrationStoreUnusableError, + 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.""" + +_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) + + +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], + 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, 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(kept) + return True, len(kept) + + +def derive_run_status( + 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: `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].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: + 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} + + 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 valid_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] = {} + # 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: + 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 _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): + 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 + 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, 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) + 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 + 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 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) + _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) + + @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. Its key stays unused: a repeat may start it.""" + try: + self._remove_run(run_id, remember=False) + 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) + 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 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( + 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] + 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 {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, QueueItemLineage]]: + """The queue's calibration items, per run, as the queue has them now.""" + 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: + 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, QueueItemLineage]], + ) -> 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] = {} + 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, {}), valid_points_per_run.get(run.run_id, 0) + ) + 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, *, 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 + 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: + 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}, 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 True, self._remove(file, run_ids, model_key)[1] + + 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 + + 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..e11c945e8b1 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,47 @@ 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. + 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 = 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() 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: + probe.locking.remove(self._cache) + 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 +233,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 +282,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/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 3691f374425..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() @@ -1888,6 +1890,28 @@ 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. 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 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 def lock_in_ram(self, cache_entry: CacheRecord) -> None: @@ -1902,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 @@ -1930,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 @@ -1966,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() @@ -2014,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: @@ -2036,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/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/util/devices.py b/invokeai/backend/util/devices.py index b648adc0c8f..8e1c31479a3 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 @@ -25,6 +25,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() @@ -447,14 +460,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..d807541183a --- /dev/null +++ b/invokeai/backend/working_memory/counters.py @@ -0,0 +1,330 @@ +"""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. 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 +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. 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. 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]: + 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..ccdde0a04c1 --- /dev/null +++ b/invokeai/backend/working_memory/fingerprint.py @@ -0,0 +1,173 @@ +"""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 (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 +import ctypes.util +import hashlib +import json +import os +import platform +import sys +from collections.abc import Mapping +from typing import Final, Literal, 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 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] = {"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 "") + 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_build": 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..839e2d785f2 --- /dev/null +++ b/invokeai/backend/working_memory/probe.py @@ -0,0 +1,768 @@ +"""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, 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 +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"}) + + +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") + + +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 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.""" + 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.""" + +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() + + +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, fit_eligible = _validity(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=fit_eligible, + 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 + ) + 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: + _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 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. 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 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: + 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`). 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") + 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 + ) + window.abort() + 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 537d0836d51..068b7b8f779 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. 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": [ + { + "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; 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" + }, + "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); with an idempotency key, possibly after the items were queued" + } + } + } + }, + "/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"], @@ -12557,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": { @@ -24571,6 +24919,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", @@ -116446,6 +116836,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 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, + "type": "object", + "required": ["graph"], + "title": "WorkingMemoryCalibrationRequest", + "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": { + "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. A repeat of the request that finds them unrecorded takes over every queue item carrying the run's origin, retries included" + } + }, + "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 valid 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 a845bdf0a0f..f972a5f9036 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -1195,6 +1195,140 @@ 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. 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. 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; + 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; @@ -7975,6 +8109,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. @@ -52635,6 +52793,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, runs and the + * graph's id, 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 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; + }; + /** 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. A repeat of the request that finds them unrecorded takes over every queue item carrying the run's origin, retries included + */ + 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 valid 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 @@ -55795,6 +56103,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; 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; + }; + 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 (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; + }; + 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); with an idempotency key, possibly after the items were queued */ + 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; @@ -59489,6 +60087,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/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..056b75f71f9 --- /dev/null +++ b/tests/app/routers/test_working_memory_calibration.py @@ -0,0 +1,295 @@ +"""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 + +import pytest +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 +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, + HeldEnqueue, + 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_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: + 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] + + +@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() + + 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..46bd1c6b4e5 --- /dev/null +++ b/tests/app/services/session_processor/test_working_memory_calibration_runner.py @@ -0,0 +1,420 @@ +"""The session runner's calibration hook: calibration items run their nodes inside a probe scope, have the windows +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 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 +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]] = [] +# 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", "hook"] = 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() + if self.behaviour == "hook": + _WHILE_RUNNING[0]() + 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() + _WHILE_RUNNING.clear() + yield + _SEEN.clear() + _ON_CANCEL.clear() + _WHILE_RUNNING.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"], 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", "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) + + 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") + [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 + # Only valid points make a run that ended without completing partial. + 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: + 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, 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 "read-only file system" in (item.error_message or "") + 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 "") + 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, 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( + 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/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/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..7ff7fd0e23a --- /dev/null +++ b/tests/app/services/test_invoker.py @@ -0,0 +1,287 @@ +"""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 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 +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() + + +# 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 new file mode 100644 index 00000000000..72ec62a6fbf --- /dev/null +++ b/tests/app/services/working_memory_calibration/test_working_memory_calibration_service.py @@ -0,0 +1,840 @@ +"""The calibration store and service on a temporary root, with a real SQLite queue.""" + +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.model_records import UnknownModelException +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, + CalibrationRunNotFoundError, + CalibrationStoreError, + CalibrationStoreUnusableError, + WorkingMemoryCalibrationRequest, + WorkingMemoryCalibrationRunCreated, + WorkingMemoryCalibrationStatus, + calibration_run_id, +) +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, + OTHER_CLASS, + HeldEnqueue, + 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,), *, 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, 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") + + +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: + 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 + # 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: + 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" + + # 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" + + # 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 "") + + +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 + 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 186b4659eaf..05e31f2291d 100644 --- a/tests/app/test_api_dependencies.py +++ b/tests/app/test_api_dependencies.py @@ -1,14 +1,26 @@ import asyncio +import threading 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 +30,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 +74,85 @@ 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() + + +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 new file mode 100644 index 00000000000..746c557e834 --- /dev/null +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_locks_held.py @@ -0,0 +1,121 @@ +"""`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 +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 + + +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) 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..92dfd2f491a --- /dev/null +++ b/tests/backend/working_memory/test_counters.py @@ -0,0 +1,234 @@ +"""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" -- and that a peak reset or a synchronization +that fails raises, since it has no "unknown" to fall back on. +""" + +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", + ) + + +@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)) + + 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..84be9d6b1f9 --- /dev/null +++ b/tests/backend/working_memory/test_fingerprint.py @@ -0,0 +1,142 @@ +"""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, os_family + +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) + # 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) + + +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_build", "invokeai"} + assert environment.soft["os_build"] == platform.platform() + 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 + + +@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" diff --git a/tests/backend/working_memory/test_probe.py b/tests/backend/working_memory/test_probe.py new file mode 100644 index 00000000000..ef326b397cd --- /dev/null +++ b/tests/backend/working_memory/test_probe.py @@ -0,0 +1,972 @@ +"""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. 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 Callable, 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 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 + 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.reset_peaks() + + 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 + 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: + 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"]) + + +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) + 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 _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) + + 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..8bcf26d2395 --- /dev/null +++ b/tests/fixtures/working_memory_calibration.py @@ -0,0 +1,107 @@ +"""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. +""" + +import threading +from collections.abc import Mapping +from pathlib import Path +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, +) +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"}) + + +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, + 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), + )