From 55b40a94256e7693e7c274a037e5d032854e3fcb Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Mon, 9 Oct 2023 11:04:03 +1100 Subject: [PATCH] feat(events): add batch status and queue status to queue item status changed events The UI will always re-fetch queue and batch status on receiving this event, so we may as well jsut include that data in the event and save the extra network roundtrips. --- invokeai/app/services/events.py | 39 +++++++++++++------ .../session_queue/session_queue_sqlite.py | 24 ++++++++++-- 2 files changed, 48 insertions(+), 15 deletions(-) diff --git a/invokeai/app/services/events.py b/invokeai/app/services/events.py index 3b36ffb917..0a02a03539 100644 --- a/invokeai/app/services/events.py +++ b/invokeai/app/services/events.py @@ -4,7 +4,12 @@ from typing import Any, Optional from invokeai.app.models.image import ProgressImage from invokeai.app.services.model_manager_service import BaseModelType, ModelInfo, ModelType, SubModelType -from invokeai.app.services.session_queue.session_queue_common import EnqueueBatchResult, SessionQueueItem +from invokeai.app.services.session_queue.session_queue_common import ( + BatchStatus, + EnqueueBatchResult, + SessionQueueItem, + SessionQueueStatus, +) from invokeai.app.util.misc import get_timestamp @@ -262,21 +267,31 @@ class EventServiceBase: ), ) - def emit_queue_item_status_changed(self, session_queue_item: SessionQueueItem) -> None: + def emit_queue_item_status_changed( + self, + session_queue_item: SessionQueueItem, + batch_status: BatchStatus, + queue_status: SessionQueueStatus, + ) -> None: """Emitted when a queue item's status changes""" self.__emit_queue_event( event_name="queue_item_status_changed", payload=dict( - queue_id=session_queue_item.queue_id, - queue_item_id=session_queue_item.item_id, - status=session_queue_item.status, - batch_id=session_queue_item.batch_id, - session_id=session_queue_item.session_id, - error=session_queue_item.error, - created_at=str(session_queue_item.created_at) if session_queue_item.created_at else None, - updated_at=str(session_queue_item.updated_at) if session_queue_item.updated_at else None, - started_at=str(session_queue_item.started_at) if session_queue_item.started_at else None, - completed_at=str(session_queue_item.completed_at) if session_queue_item.completed_at else None, + queue_id=queue_status.queue_id, + queue_item=dict( + queue_id=session_queue_item.queue_id, + item_id=session_queue_item.item_id, + status=session_queue_item.status, + batch_id=session_queue_item.batch_id, + session_id=session_queue_item.session_id, + error=session_queue_item.error, + created_at=str(session_queue_item.created_at) if session_queue_item.created_at else None, + updated_at=str(session_queue_item.updated_at) if session_queue_item.updated_at else None, + started_at=str(session_queue_item.started_at) if session_queue_item.started_at else None, + completed_at=str(session_queue_item.completed_at) if session_queue_item.completed_at else None, + ), + batch_status=batch_status.dict(), + queue_status=queue_status.dict(), ), ) diff --git a/invokeai/app/services/session_queue/session_queue_sqlite.py b/invokeai/app/services/session_queue/session_queue_sqlite.py index 19a5346e10..f995576311 100644 --- a/invokeai/app/services/session_queue/session_queue_sqlite.py +++ b/invokeai/app/services/session_queue/session_queue_sqlite.py @@ -427,7 +427,13 @@ class SqliteSessionQueue(SessionQueueBase): finally: self.__lock.release() queue_item = self.get_queue_item(item_id) - self.__invoker.services.events.emit_queue_item_status_changed(queue_item) + batch_status = self.get_batch_status(queue_id=queue_item.queue_id, batch_id=queue_item.batch_id) + queue_status = self.get_queue_status(queue_id=queue_item.queue_id) + self.__invoker.services.events.emit_queue_item_status_changed( + session_queue_item=queue_item, + batch_status=batch_status, + queue_status=queue_status, + ) return queue_item def is_empty(self, queue_id: str) -> IsEmptyResult: @@ -609,7 +615,13 @@ class SqliteSessionQueue(SessionQueueBase): queue_batch_id=current_queue_item.batch_id, graph_execution_state_id=current_queue_item.session_id, ) - self.__invoker.services.events.emit_queue_item_status_changed(current_queue_item) + batch_status = self.get_batch_status(queue_id=queue_id, batch_id=current_queue_item.batch_id) + queue_status = self.get_queue_status(queue_id=queue_id) + self.__invoker.services.events.emit_queue_item_status_changed( + session_queue_item=current_queue_item, + batch_status=batch_status, + queue_status=queue_status, + ) except Exception: self.__conn.rollback() raise @@ -655,7 +667,13 @@ class SqliteSessionQueue(SessionQueueBase): queue_batch_id=current_queue_item.batch_id, graph_execution_state_id=current_queue_item.session_id, ) - self.__invoker.services.events.emit_queue_item_status_changed(current_queue_item) + batch_status = self.get_batch_status(queue_id=queue_id, batch_id=current_queue_item.batch_id) + queue_status = self.get_queue_status(queue_id=queue_id) + self.__invoker.services.events.emit_queue_item_status_changed( + session_queue_item=current_queue_item, + batch_status=batch_status, + queue_status=queue_status, + ) except Exception: self.__conn.rollback() raise