157 lines
7.5 KiB
Python
157 lines
7.5 KiB
Python
import time
|
|
import heapq
|
|
import logging
|
|
from typing import Optional, Any, Callable
|
|
|
|
|
|
class QueueHookManager:
|
|
"""Manages installation/uninstallation of queue hooks.
|
|
|
|
Responsible for wrapping prompt queue methods to add persistence and pause behavior.
|
|
"""
|
|
|
|
def __init__(self, *, is_paused_fn: Callable[[], bool], on_job_started: Callable[[str], None], on_task_done: Callable[[Any, Any, Any], None], should_run_when_paused: Optional[Callable[[str], bool]] = None):
|
|
self._original_queue_get = None
|
|
self._original_task_done = None
|
|
self._installed = False
|
|
self._is_paused = is_paused_fn
|
|
self._on_job_started = on_job_started
|
|
self._on_task_done = on_task_done
|
|
self._should_run_when_paused = should_run_when_paused
|
|
|
|
def install(self) -> None:
|
|
"""Install hooks into execution.PromptQueue if not already installed."""
|
|
if self._installed:
|
|
return
|
|
import execution
|
|
|
|
if self._original_queue_get is None:
|
|
self._original_queue_get = execution.PromptQueue.get
|
|
|
|
def get_wrapper(q_self, timeout=None):
|
|
# If paused, only allow items explicitly permitted by the manager (run-selected mode)
|
|
try:
|
|
if self._is_paused():
|
|
allowed = False
|
|
if callable(self._should_run_when_paused):
|
|
try:
|
|
top = q_self.queue[0] if getattr(q_self, 'queue', None) else None
|
|
except Exception:
|
|
top = None
|
|
if top is not None:
|
|
try:
|
|
top_pid = str(top[1])
|
|
allowed = bool(self._should_run_when_paused(top_pid))
|
|
except Exception:
|
|
allowed = False
|
|
if not allowed:
|
|
time.sleep(0.1)
|
|
return None
|
|
except Exception:
|
|
# If anything goes wrong during checks, be safe and respect pause
|
|
if self._is_paused():
|
|
time.sleep(0.1)
|
|
return None
|
|
|
|
result = self._original_queue_get(q_self, timeout=timeout)
|
|
if result is not None:
|
|
try:
|
|
item, _item_id = result
|
|
prompt_id = item[1]
|
|
# Helper: sanitize prompt by stripping non-node metadata keys
|
|
def _sanitize_item(it):
|
|
try:
|
|
n, pid, prompt, extra, outs = it
|
|
if isinstance(prompt, dict):
|
|
try:
|
|
prompt_clean = {k: v for k, v in prompt.items() if k not in ('workflow', 'name')}
|
|
except Exception:
|
|
prompt_clean = prompt
|
|
if prompt_clean is not prompt:
|
|
return (n, pid, prompt_clean, extra, outs)
|
|
except Exception:
|
|
pass
|
|
return it
|
|
# If paused, ensure the popped item is allowed; otherwise, reinsert and yield None
|
|
if self._is_paused() and callable(self._should_run_when_paused):
|
|
try:
|
|
if not self._should_run_when_paused(str(prompt_id)):
|
|
# Undo the side effects of the pop from the original get():
|
|
# 1) Remove from currently_running
|
|
# 2) Reinsert the item into the heap queue
|
|
# 3) Notify the server that the queue/running set changed
|
|
try:
|
|
with q_self.mutex:
|
|
try:
|
|
q_self.currently_running.pop(_item_id, None)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
safe_item = _sanitize_item(item)
|
|
q_self.queue.append(safe_item)
|
|
heapq.heapify(q_self.queue)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
q_self.server.queue_updated()
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
time.sleep(0.05)
|
|
return None
|
|
except Exception:
|
|
time.sleep(0.05)
|
|
return None
|
|
# Sanitize both the running copy and the returned item to prevent crashes
|
|
try:
|
|
with q_self.mutex:
|
|
if _item_id in q_self.currently_running:
|
|
try:
|
|
q_self.currently_running[_item_id] = _sanitize_item(q_self.currently_running[_item_id])
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
# Replace the returned item with sanitized version
|
|
try:
|
|
item = _sanitize_item(item)
|
|
result = (item, _item_id)
|
|
except Exception:
|
|
pass
|
|
self._on_job_started(prompt_id)
|
|
except Exception as e:
|
|
logging.debug(f"QueueHookManager get_wrapper failed: {e}")
|
|
return result
|
|
|
|
execution.PromptQueue.get = get_wrapper
|
|
|
|
if self._original_task_done is None:
|
|
self._original_task_done = execution.PromptQueue.task_done
|
|
|
|
def task_done_wrapper(q_self, item_id, history_result, status):
|
|
try:
|
|
self._on_task_done((q_self, item_id, history_result, status))
|
|
except Exception as e:
|
|
logging.debug(f"QueueHookManager task_done on_task_done failed: {e}")
|
|
self._original_task_done(q_self, item_id, history_result, status)
|
|
|
|
execution.PromptQueue.task_done = task_done_wrapper
|
|
|
|
self._installed = True
|
|
|
|
def uninstall(self) -> None:
|
|
"""Restore original methods."""
|
|
if not self._installed:
|
|
return
|
|
import execution
|
|
if self._original_queue_get is not None:
|
|
execution.PromptQueue.get = self._original_queue_get
|
|
self._original_queue_get = None
|
|
if self._original_task_done is not None:
|
|
execution.PromptQueue.task_done = self._original_task_done
|
|
self._original_task_done = None
|
|
self._installed = False
|
|
|
|
|