Skip to content

Commit

Permalink
feat(app): add cancel all except current queue item functionality
Browse files Browse the repository at this point in the history
  • Loading branch information
rikublock committed Dec 6, 2024
1 parent b7132ce commit 0bf431f
Show file tree
Hide file tree
Showing 4 changed files with 59 additions and 0 deletions.
13 changes: 13 additions & 0 deletions invokeai/app/api/routers/session_queue.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
QUEUE_ITEM_STATUS,
Batch,
BatchStatus,
CancelAllExceptCurrentResult,
CancelByBatchIDsResult,
CancelByDestinationResult,
ClearResult,
Expand Down Expand Up @@ -94,6 +95,18 @@ async def Pause(
return ApiDependencies.invoker.services.session_processor.pause()


@session_queue_router.put(
"/{queue_id}/cancel_all_except_current",
operation_id="cancel_all_except_current",
responses={200: {"model": CancelAllExceptCurrentResult}},
)
async def cancel_all_except_current(
queue_id: str = Path(description="The queue id to perform this operation on"),
) -> CancelAllExceptCurrentResult:
"""Immediately cancels all queue items except in-processing items"""
return ApiDependencies.invoker.services.session_queue.cancel_all_except_current(queue_id=queue_id)


@session_queue_router.put(
"/{queue_id}/cancel_by_batch_ids",
operation_id="cancel_by_batch_ids",
Expand Down
6 changes: 6 additions & 0 deletions invokeai/app/services/session_queue/session_queue_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
QUEUE_ITEM_STATUS,
Batch,
BatchStatus,
CancelAllExceptCurrentResult,
CancelByBatchIDsResult,
CancelByDestinationResult,
CancelByQueueIDResult,
Expand Down Expand Up @@ -112,6 +113,11 @@ def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
"""Cancels all queue items with matching queue ID"""
pass

@abstractmethod
def cancel_all_except_current(self, queue_id: str) -> CancelAllExceptCurrentResult:
"""Cancels all queue items except in-progress items"""
pass

@abstractmethod
def list_queue_items(
self,
Expand Down
6 changes: 6 additions & 0 deletions invokeai/app/services/session_queue/session_queue_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -366,6 +366,12 @@ class CancelByQueueIDResult(CancelByBatchIDsResult):
pass


class CancelAllExceptCurrentResult(CancelByBatchIDsResult):
"""Result of canceling all except current"""

pass


class IsEmptyResult(BaseModel):
"""Result of checking if the session queue is empty"""

Expand Down
34 changes: 34 additions & 0 deletions invokeai/app/services/session_queue/session_queue_sqlite.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
QUEUE_ITEM_STATUS,
Batch,
BatchStatus,
CancelAllExceptCurrentResult,
CancelByBatchIDsResult,
CancelByDestinationResult,
CancelByQueueIDResult,
Expand Down Expand Up @@ -510,6 +511,39 @@ def cancel_by_queue_id(self, queue_id: str) -> CancelByQueueIDResult:
self.__lock.release()
return CancelByQueueIDResult(canceled=count)

def cancel_all_except_current(self, queue_id: str) -> CancelAllExceptCurrentResult:
try:
where = """--sql
WHERE
queue_id == ?
AND status == 'pending'
"""
self.__lock.acquire()
self.__cursor.execute(
f"""--sql
SELECT COUNT(*)
FROM session_queue
{where};
""",
(queue_id,),
)
count = self.__cursor.fetchone()[0]
self.__cursor.execute(
f"""--sql
UPDATE session_queue
SET status = 'canceled'
{where};
""",
(queue_id,),
)
self.__conn.commit()
except Exception:
self.__conn.rollback()
raise
finally:
self.__lock.release()
return CancelAllExceptCurrentResult(canceled=count)

def get_queue_item(self, item_id: int) -> SessionQueueItem:
try:
self.__lock.acquire()
Expand Down

0 comments on commit 0bf431f

Please sign in to comment.