feat(working-memory): measure working memory in admin calibration runs

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.

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

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:

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")

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")

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)