Vito
b1e7abd380
* fix(threads): stop and join the prompt cleanup thread, the image worker and the observer (#302, #370) - PromptTracker: the cleanup thread waits on a stop event instead of time.sleep; new shutdown(timeout=2) joins it and returns whether it ended. - ImageGenerationHandler.shutdown(timeout=5) now joins the worker and returns a bool; new close() is final: schedule() refuses and a sealed run's history re-check is dropped, so no worker restarts afterwards. - ImageMonitor.stop_monitoring(timeout=5) closes the handler first, then stops and joins the watchdog observer within one deadline. - Every test that builds a real tracker, handler or monitor ends with its final stop and an assertion, before its folder and databases go. - #370: the wrapper/session history test drains before dispose(); the held-write bound test uses 30 s ceilings instead of tight wall clocks. * test(history): drain the writer before asserting a bounded close it does not test (#370) * fix(threads): a later stop waits again for a worker or observer an earlier stop left running (#302) Review findings on #372: - ImageGenerationHandler keeps the workers a stop signalled that had not ended; every later shutdown/close joins them again, so a second call no longer returns True while the first worker still runs. - ImageMonitor.stop_monitoring keeps an observer or handler it left running and stops and joins it again on the next call. - A closed handler takes no awaiting history path, and a re-check that raced close() is not kept due. - Tests: second stops of a busy worker and a live observer, the closed handler's re-check paths, the handle test asserts the worker ended and stops before it asserts, and a no-new-thread check instead of set equality. * test(browser): the System reset restores every settings value, safety_backup included The frontend coverage shard of #372 failed in workspace-system-tagging: the review read "No safety backup (the policy is off)". The backups spec saves safety_backup=false, and /__test__/system/reset restored config.json but not PromptManagerConfig.SAFETY_BACKUP (nor INFINITE_SCROLL), so a later spec on the same worker server saw the policy off. The reset now captures and restores every non-tagging value POST /settings can change.
725 lines
28 KiB
Python
725 lines
28 KiB
Python
"""Prompt tracking system for linking generated images to prompts.
|
|
|
|
This module provides thread-safe tracking of active prompt executions to enable
|
|
automatic association of generated images with their source prompts. The system
|
|
maintains prompt context across different execution threads and provides fallback
|
|
mechanisms for reliable image-prompt linking.
|
|
|
|
Key features:
|
|
- Thread-safe prompt tracking using threading.local and locks
|
|
- Automatic cleanup of expired prompt contexts
|
|
- Fallback mechanisms for cross-thread prompt access
|
|
- Support for execution timeouts and extensions
|
|
- Context manager for automatic prompt lifecycle management
|
|
|
|
Typical usage:
|
|
from utils.prompt_tracker import PromptTracker
|
|
|
|
tracker = PromptTracker(db_manager)
|
|
execution_id = tracker.set_current_prompt("beautiful landscape")
|
|
# Images generated after this point will be linked to this prompt
|
|
|
|
Or using the context manager:
|
|
with PromptExecutionContext(tracker, "beautiful landscape"):
|
|
# Generate images here
|
|
pass
|
|
"""
|
|
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from collections import OrderedDict
|
|
from typing import Any, Dict, Optional, Tuple
|
|
|
|
# Import logging system
|
|
try:
|
|
from .hashing import generate_prompt_hash
|
|
from .logging_config import get_logger
|
|
except ImportError:
|
|
import sys
|
|
import os
|
|
|
|
current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, current_dir)
|
|
from utils.hashing import generate_prompt_hash
|
|
from utils.logging_config import get_logger
|
|
|
|
# How long ``PromptTracker.shutdown`` waits for the cleanup thread (#301).
|
|
TRACKER_JOIN_SECONDS = 2.0
|
|
|
|
|
|
class PromptTracker:
|
|
"""Thread-safe tracking of current prompt executions.
|
|
|
|
This class manages active prompt contexts across multiple threads, enabling
|
|
the image monitoring system to correctly associate generated images with their
|
|
source prompts. It uses both thread-local storage and global tracking to handle
|
|
various execution scenarios.
|
|
|
|
The tracker automatically cleans up expired prompts and provides fallback
|
|
mechanisms when prompts are accessed from different threads (e.g., image
|
|
monitoring running in a separate thread from prompt execution).
|
|
|
|
Attributes:
|
|
active_prompts: Global dictionary of active prompt contexts
|
|
cleanup_interval: Seconds between cleanup operations (default: 300)
|
|
prompt_timeout: Seconds before a prompt expires (default: 600)
|
|
"""
|
|
|
|
def __init__(self, db_manager):
|
|
"""
|
|
Initialize the prompt tracker.
|
|
|
|
Args:
|
|
db_manager: Database manager instance for prompt operations
|
|
"""
|
|
self.logger = get_logger("prompt_manager.prompt_tracker")
|
|
self.logger.debug("Initializing PromptTracker")
|
|
|
|
self.db_manager = db_manager
|
|
self._local = threading.local()
|
|
self.active_prompts = {} # Global tracking for multiple threads
|
|
self.lock = threading.Lock()
|
|
# FIFO queue for batch workflows: prompt_ids are pushed during encoding
|
|
# and popped during image linking, preserving correct ordering.
|
|
self._prompt_queue = []
|
|
self._queue_lock = threading.Lock()
|
|
|
|
# Read from GalleryConfig if available, otherwise use defaults
|
|
try:
|
|
from ..py.config import GalleryConfig
|
|
|
|
self.cleanup_interval = GalleryConfig.CLEANUP_INTERVAL
|
|
self.prompt_timeout = GalleryConfig.PROMPT_TIMEOUT
|
|
except Exception:
|
|
self.cleanup_interval = 300 # 5 minutes
|
|
self.prompt_timeout = 600 # 10 minutes
|
|
|
|
# Start cleanup thread; ``shutdown`` sets ``_stop`` to end it.
|
|
self._stop = threading.Event()
|
|
self.cleanup_thread = threading.Thread(
|
|
target=self._cleanup_expired_prompts,
|
|
name="PromptManagerPromptCleanup",
|
|
daemon=True,
|
|
)
|
|
self.cleanup_thread.start()
|
|
self.logger.debug("PromptTracker initialization completed")
|
|
|
|
def set_current_prompt(
|
|
self,
|
|
prompt_text: str,
|
|
additional_data: Optional[Dict[str, Any]] = None,
|
|
push_to_queue: bool = True,
|
|
run_key: Optional[str] = None,
|
|
run_id: Optional[str] = None,
|
|
) -> str:
|
|
"""
|
|
Set the current prompt for this thread for image tracking.
|
|
|
|
This method establishes the prompt context that will be used to link any
|
|
subsequently generated images. The prompt is stored in both thread-local
|
|
storage and global tracking to handle cross-thread access scenarios.
|
|
|
|
Note: Prompt saving is handled by PromptManager to avoid duplicates.
|
|
|
|
Args:
|
|
prompt_text: The prompt text being executed
|
|
additional_data: Additional prompt metadata, should include prompt_id
|
|
from PromptManager
|
|
push_to_queue: Also push onto the FIFO batch queue. Only batch items whose
|
|
text comes from another node need it; typed prompts are linked from
|
|
the image's own metadata instead.
|
|
run_key: The node's run key (``utils.run_keys.prompt_node_key``). The
|
|
image monitor links an image to a tracked prompt only when the
|
|
image's own graph gives the same key (#342).
|
|
run_id: The ComfyUI prompt id executing the node, when the host says.
|
|
|
|
Returns:
|
|
Unique execution ID for this prompt execution
|
|
"""
|
|
execution_id = self.generate_execution_id()
|
|
|
|
# Get prompt_id from additional_data (passed from PromptManager)
|
|
# PromptTracker should NOT save prompts to avoid duplicates
|
|
prompt_id = additional_data.get("prompt_id") if additional_data else None
|
|
|
|
if not prompt_id:
|
|
# Fallback: try to find existing prompt using consistent hash calculation
|
|
try:
|
|
# Same hash as the nodes, so the lookup finds their saved prompt
|
|
prompt_hash = generate_prompt_hash(prompt_text)
|
|
existing_prompt = self.db_manager.get_prompt_by_hash(prompt_hash)
|
|
|
|
if existing_prompt:
|
|
prompt_id = existing_prompt["id"]
|
|
self.logger.debug(f"Found existing prompt ID: {prompt_id}")
|
|
else:
|
|
# Generate a temporary ID for tracking (prompt should be saved
|
|
# by PromptManager)
|
|
prompt_id = f"temp_{int(time.time())}"
|
|
self.logger.debug(f"Using temporary ID for tracking: {prompt_id}")
|
|
except Exception as e:
|
|
self.logger.error(f"Error finding prompt: {e}")
|
|
prompt_id = f"temp_{int(time.time())}"
|
|
|
|
# Create execution context
|
|
execution_context = {
|
|
"id": prompt_id,
|
|
"execution_id": execution_id,
|
|
"text": prompt_text,
|
|
"timestamp": time.time(),
|
|
"thread_id": threading.current_thread().ident,
|
|
"additional_data": additional_data or {},
|
|
"run_key": run_key,
|
|
"run_id": run_id,
|
|
}
|
|
|
|
# Store in thread-local storage
|
|
self._local.current_prompt = execution_context
|
|
|
|
# Store in global tracking with thread safety
|
|
with self.lock:
|
|
self.active_prompts[execution_id] = execution_context
|
|
|
|
self._remember_keyed(execution_context)
|
|
|
|
self.logger.debug(
|
|
f"Set current prompt: {execution_id} -> {prompt_text[:50]}... "
|
|
f"(thread: {threading.current_thread().ident})"
|
|
)
|
|
self.logger.debug(f"Active prompts count: {len(self.active_prompts)}")
|
|
|
|
# Push to batch queue for ordered image linking. Last, with nothing
|
|
# after it that can raise: a caller that marks a failed execution
|
|
# (#342) must never add a second queue slot for the same item.
|
|
if push_to_queue:
|
|
with self._queue_lock:
|
|
self._prompt_queue.append(execution_context)
|
|
return execution_id
|
|
|
|
def get_current_prompt(self) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Get the current prompt context for this thread.
|
|
|
|
Attempts to retrieve the current prompt context, first from thread-local
|
|
storage, then from global tracking as a fallback. This enables image
|
|
monitoring (which may run in a different thread) to access prompt context.
|
|
|
|
Returns:
|
|
Dictionary containing prompt context with keys:
|
|
- id: Prompt ID in database
|
|
- execution_id: Unique execution identifier
|
|
- text: Prompt text
|
|
- timestamp: When the prompt was set
|
|
- thread_id: Thread identifier
|
|
- additional_data: Any additional metadata
|
|
Returns None if no valid prompt context is found
|
|
"""
|
|
current = getattr(self._local, "current_prompt", None)
|
|
|
|
if current:
|
|
# Check if prompt hasn't expired
|
|
if time.time() - current["timestamp"] < self.prompt_timeout:
|
|
return current
|
|
else:
|
|
self.logger.debug(f"Current prompt expired: {current['execution_id']}")
|
|
self.clear_current_prompt()
|
|
|
|
# Fallback: try to find recent prompt from global tracking
|
|
# This is crucial for image monitoring which runs in different threads
|
|
recent_prompt = self._find_recent_prompt()
|
|
if recent_prompt:
|
|
self.logger.debug(
|
|
"Using recent prompt from global tracking: "
|
|
f"{recent_prompt['execution_id']}"
|
|
)
|
|
return recent_prompt
|
|
|
|
self.logger.debug(
|
|
f"No prompt context found (thread: {threading.current_thread().ident})"
|
|
)
|
|
return None
|
|
|
|
def mark_empty(
|
|
self,
|
|
run_key: Optional[str],
|
|
run_id: Optional[str] = None,
|
|
push_to_queue: bool = False,
|
|
) -> None:
|
|
"""Record that a positive node ran with nothing to save (#342).
|
|
|
|
The run key's latest entry becomes a placeholder, so no earlier prompt of
|
|
that key can be linked to this run's images; with ``push_to_queue`` the
|
|
placeholder also holds this item's place in the batch queue, so a later
|
|
item's image never takes the wrong queue entry.
|
|
"""
|
|
if not run_key:
|
|
return
|
|
placeholder = {
|
|
"id": None,
|
|
"placeholder": True,
|
|
"execution_id": self.generate_execution_id(),
|
|
"text": "",
|
|
"timestamp": time.time(),
|
|
"additional_data": {},
|
|
"run_key": run_key,
|
|
"run_id": run_id,
|
|
}
|
|
self._remember_keyed(placeholder)
|
|
if push_to_queue:
|
|
with self._queue_lock:
|
|
self._prompt_queue.append(placeholder)
|
|
|
|
def _keyed_runs(self) -> "KeyedRuns":
|
|
keyed = getattr(self, "_keyed", None)
|
|
if keyed is None:
|
|
with _KEYED_INIT_LOCK:
|
|
keyed = getattr(self, "_keyed", None)
|
|
if keyed is None:
|
|
keyed = self._keyed = KeyedRuns()
|
|
return keyed
|
|
|
|
def _remember_keyed(self, context: Dict[str, Any]) -> None:
|
|
if context.get("run_key"):
|
|
self._keyed_runs().record(context)
|
|
|
|
def keyed_snapshot(self) -> Dict[str, Tuple[Dict[str, Any], Optional[str], int]]:
|
|
"""``{run key: (latest entry, its run id, entries of that run)}`` now.
|
|
|
|
The image monitor captures it when an image appears, before later runs
|
|
can move on.
|
|
"""
|
|
return self._keyed_runs().snapshot(self.prompt_timeout)
|
|
|
|
def pop_for_key(
|
|
self, run_key: Optional[str], run_id: Optional[str] = None
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Pop the oldest queued entry of ``run_key`` (FIFO), from ``run_id`` if set.
|
|
|
|
``run_id`` is the run the key's latest entry came from when the image
|
|
appeared (``keyed_snapshot``): entries of that key queued by earlier
|
|
runs ahead of it are stale and dropped, while entries of later runs
|
|
(queued after the image appeared) stay for their own images. Entries
|
|
of other keys stay queued; expired entries are discarded.
|
|
"""
|
|
if not run_key:
|
|
return None
|
|
now = time.time()
|
|
with self._queue_lock:
|
|
live = [
|
|
ctx
|
|
for ctx in self._prompt_queue
|
|
if now - ctx.get("timestamp", 0) < self.prompt_timeout
|
|
]
|
|
index = next(
|
|
(
|
|
i
|
|
for i, ctx in enumerate(live)
|
|
if ctx.get("run_key") == run_key
|
|
and (run_id is None or ctx.get("run_id") == run_id)
|
|
),
|
|
None,
|
|
)
|
|
if index is None:
|
|
self._prompt_queue = live
|
|
return None
|
|
found = live[index]
|
|
# Same-key entries ahead of it came from earlier runs: stale
|
|
stale_or_found = {
|
|
i for i in range(index + 1) if live[i].get("run_key") == run_key
|
|
}
|
|
self._prompt_queue = [
|
|
ctx for i, ctx in enumerate(live) if i not in stale_or_found
|
|
]
|
|
return found
|
|
|
|
def pop_next_prompt(self) -> Optional[Dict[str, Any]]:
|
|
"""Pop the next prompt from the batch queue (FIFO).
|
|
|
|
In batch workflows, prompts are encoded in order and images are saved
|
|
in the same order. This method pops the oldest queued prompt, ensuring
|
|
each image links to the correct prompt by position.
|
|
|
|
Skips and discards expired entries to prevent stale prompts from
|
|
accumulating across sessions.
|
|
|
|
Returns:
|
|
Prompt context dict, or None if queue is empty or all entries expired
|
|
"""
|
|
now = time.time()
|
|
with self._queue_lock:
|
|
while self._prompt_queue:
|
|
ctx = self._prompt_queue.pop(0)
|
|
if now - ctx.get("timestamp", 0) < self.prompt_timeout:
|
|
self.logger.debug(
|
|
f"Queue pop: prompt {ctx.get('id')} "
|
|
f"({len(self._prompt_queue)} remaining)"
|
|
)
|
|
return ctx
|
|
self.logger.debug(f"Queue discard expired: prompt {ctx.get('id')}")
|
|
return None
|
|
|
|
def _find_recent_prompt(self) -> Optional[Dict[str, Any]]:
|
|
"""Find the most recent prompt that's still valid.
|
|
|
|
Searches the global prompt tracking for the most recently set prompt
|
|
that hasn't expired. This is used as a fallback when thread-local
|
|
storage doesn't contain a current prompt.
|
|
|
|
Returns:
|
|
Most recent valid prompt context, or None if no valid prompts found
|
|
"""
|
|
with self.lock:
|
|
current_time = time.time()
|
|
self.logger.debug(
|
|
f"Searching for recent prompt among {len(self.active_prompts)} "
|
|
"active prompts"
|
|
)
|
|
|
|
recent_prompts = []
|
|
for exec_id, prompt in self.active_prompts.items():
|
|
age_seconds = current_time - prompt["timestamp"]
|
|
self.logger.debug(
|
|
f"Prompt {exec_id}: age={age_seconds:.1f}s, "
|
|
f"timeout={self.prompt_timeout}s"
|
|
)
|
|
|
|
if age_seconds < self.prompt_timeout:
|
|
recent_prompts.append(prompt)
|
|
else:
|
|
self.logger.debug(f"Prompt {exec_id} is expired")
|
|
|
|
if recent_prompts:
|
|
# Return the most recent one
|
|
most_recent = max(recent_prompts, key=lambda p: p["timestamp"])
|
|
self.logger.debug(f"Found recent prompt: {most_recent['execution_id']}")
|
|
return most_recent
|
|
else:
|
|
self.logger.debug("No recent prompts found")
|
|
|
|
return None
|
|
|
|
def clear_current_prompt(self):
|
|
"""Clear the current prompt context for this thread.
|
|
|
|
Removes the prompt context from both thread-local storage and global
|
|
tracking. This should be called when prompt execution is complete,
|
|
though the system also handles automatic cleanup via timeouts.
|
|
"""
|
|
current = getattr(self._local, "current_prompt", None)
|
|
if current:
|
|
execution_id = current["execution_id"]
|
|
self.logger.debug(
|
|
f"Clearing current prompt: {execution_id} "
|
|
f"(thread: {threading.current_thread().ident})"
|
|
)
|
|
|
|
# Clear from thread-local storage
|
|
self._local.current_prompt = None
|
|
|
|
# Remove from global tracking
|
|
with self.lock:
|
|
removed = self.active_prompts.pop(execution_id, None)
|
|
if removed:
|
|
self.logger.debug(
|
|
f"Removed prompt from global tracking: {execution_id}"
|
|
)
|
|
self.logger.debug(
|
|
f"Remaining active prompts: {len(self.active_prompts)}"
|
|
)
|
|
else:
|
|
self.logger.debug(
|
|
f"Prompt {execution_id} was not in global tracking"
|
|
)
|
|
else:
|
|
self.logger.debug(
|
|
"No current prompt to clear "
|
|
f"(thread: {threading.current_thread().ident})"
|
|
)
|
|
|
|
def extend_prompt_timeout(self, execution_id: str, additional_seconds: int = 60):
|
|
"""
|
|
Extend the timeout for a specific prompt execution.
|
|
|
|
This is useful for long-running image generation processes where images
|
|
may continue to be generated well after the initial prompt execution.
|
|
Updates the timestamp to prevent the prompt from being cleaned up.
|
|
|
|
Args:
|
|
execution_id: The execution ID to extend
|
|
additional_seconds: Additional seconds to add to timeout (currently unused,
|
|
method just resets the timestamp to current time)
|
|
"""
|
|
with self.lock:
|
|
if execution_id in self.active_prompts:
|
|
self.active_prompts[execution_id]["timestamp"] = time.time()
|
|
self.logger.debug(f"Extended timeout for prompt: {execution_id}")
|
|
|
|
def generate_execution_id(self) -> str:
|
|
"""Generate a unique execution ID.
|
|
|
|
Creates a unique identifier for this prompt execution using UUID and timestamp.
|
|
|
|
Returns:
|
|
Unique execution ID in format 'exec_{uuid8}_{timestamp}'
|
|
"""
|
|
return f"exec_{uuid.uuid4().hex[:8]}_{int(time.time())}"
|
|
|
|
def _cleanup_expired_prompts(self):
|
|
"""Background thread to clean up expired prompts.
|
|
|
|
Runs continuously as a daemon thread, periodically removing expired
|
|
prompt contexts from the global tracking dictionary. This prevents
|
|
memory leaks from accumulating old prompt data. Every wait ends at
|
|
once when ``shutdown`` sets the stop event.
|
|
"""
|
|
while not self._stop.is_set():
|
|
try:
|
|
current_time = time.time()
|
|
expired_ids = []
|
|
|
|
with self.lock:
|
|
for exec_id, prompt_data in self.active_prompts.items():
|
|
if (
|
|
current_time - prompt_data["timestamp"]
|
|
> self.prompt_timeout
|
|
):
|
|
expired_ids.append(exec_id)
|
|
|
|
for exec_id in expired_ids:
|
|
self.active_prompts.pop(exec_id, None)
|
|
|
|
if expired_ids:
|
|
self.logger.debug(f"Cleaned up {len(expired_ids)} expired prompts")
|
|
|
|
# Also prune expired entries from the batch queue
|
|
with self._queue_lock:
|
|
before = len(self._prompt_queue)
|
|
self._prompt_queue = [
|
|
ctx
|
|
for ctx in self._prompt_queue
|
|
if current_time - ctx.get("timestamp", 0) < self.prompt_timeout
|
|
]
|
|
pruned = before - len(self._prompt_queue)
|
|
if pruned:
|
|
self.logger.debug(f"Pruned {pruned} expired queue entries")
|
|
|
|
if self._stop.wait(self.cleanup_interval):
|
|
break
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error in cleanup thread: {e}")
|
|
if self._stop.wait(60): # Wait a minute before retrying
|
|
break
|
|
|
|
def shutdown(self, timeout: float = TRACKER_JOIN_SECONDS) -> bool:
|
|
"""Stop the cleanup thread and wait for it within ``timeout`` seconds.
|
|
|
|
Returns True when the thread has ended (or never existed). The tracker
|
|
keeps working afterwards; only the background pruning of expired
|
|
entries stops. Safe to call more than once. Called on the cleanup
|
|
thread itself it only signals and returns False.
|
|
"""
|
|
stop = getattr(self, "_stop", None)
|
|
thread = getattr(self, "cleanup_thread", None)
|
|
if stop is None or thread is None:
|
|
return True # built without __init__: no thread was started
|
|
stop.set()
|
|
if thread is threading.current_thread():
|
|
self.logger.debug("shutdown called on the prompt cleanup thread")
|
|
return False
|
|
thread.join(timeout)
|
|
if thread.is_alive():
|
|
self.logger.warning(
|
|
f"prompt cleanup thread still running after {timeout} s"
|
|
)
|
|
return False
|
|
return True
|
|
|
|
def get_active_prompts(self) -> Dict[str, Dict[str, Any]]:
|
|
"""
|
|
Get all currently active prompts.
|
|
|
|
Returns:
|
|
Dictionary mapping execution IDs to prompt context dictionaries.
|
|
This is a copy of the internal tracking dictionary.
|
|
"""
|
|
with self.lock:
|
|
return self.active_prompts.copy()
|
|
|
|
def clear_all_active_prompts(self) -> int:
|
|
"""
|
|
Clear all active prompts.
|
|
|
|
Removes all prompt contexts from both global and thread-local storage.
|
|
This is useful before batch operations or when resetting the tracking state.
|
|
|
|
Returns:
|
|
Number of prompts that were cleared from global tracking
|
|
"""
|
|
with self.lock:
|
|
cleared_count = len(self.active_prompts)
|
|
self.active_prompts.clear()
|
|
|
|
# Clear thread-local storage as well
|
|
self._local.current_prompt = None
|
|
|
|
self.logger.debug(f"Cleared {cleared_count} active prompts")
|
|
return cleared_count
|
|
|
|
def get_status(self) -> Dict[str, Any]:
|
|
"""
|
|
Get tracker status information.
|
|
|
|
Returns:
|
|
Dictionary containing:
|
|
- active_prompts_count: Number of currently active prompts
|
|
- current_prompt_id: ID of current prompt in this thread (if any)
|
|
- current_execution_id: Execution ID of current prompt in this thread
|
|
(if any)
|
|
- thread_id: Current thread identifier
|
|
- prompt_timeout: Configured prompt timeout in seconds
|
|
- cleanup_interval: Configured cleanup interval in seconds
|
|
"""
|
|
with self.lock:
|
|
active_count = len(self.active_prompts)
|
|
|
|
current_prompt = self.get_current_prompt()
|
|
|
|
return {
|
|
"active_prompts_count": active_count,
|
|
"current_prompt_id": current_prompt["id"] if current_prompt else None,
|
|
"current_execution_id": (
|
|
current_prompt["execution_id"] if current_prompt else None
|
|
),
|
|
"thread_id": threading.current_thread().ident,
|
|
"prompt_timeout": self.prompt_timeout,
|
|
"cleanup_interval": self.cleanup_interval,
|
|
}
|
|
|
|
|
|
_KEYED_INIT_LOCK = threading.Lock()
|
|
# Run keys remembered for the image monitor (oldest forgotten first)
|
|
MAX_KEYED_RUNS = 256
|
|
|
|
|
|
class KeyedRuns:
|
|
"""The latest tracked entry per run key, and how many its run produced.
|
|
|
|
More than one entry for a key in one run (a batch fed by a list) means the
|
|
latest one alone cannot say which item an image belongs to.
|
|
"""
|
|
|
|
def __init__(self, max_keys: int = MAX_KEYED_RUNS):
|
|
self._max = max_keys
|
|
self._latest: "OrderedDict[str, Tuple[Dict[str, Any], Optional[str], int]]"
|
|
self._latest = OrderedDict()
|
|
self._lock = threading.Lock()
|
|
|
|
def record(self, context: Dict[str, Any]) -> None:
|
|
key, run_id = context["run_key"], context.get("run_id")
|
|
with self._lock:
|
|
previous = self._latest.pop(key, None)
|
|
same_run = (
|
|
previous is not None and run_id is not None and previous[1] == run_id
|
|
)
|
|
count = previous[2] + 1 if same_run else 1
|
|
self._latest[key] = (context, run_id, count)
|
|
while len(self._latest) > self._max:
|
|
self._latest.popitem(last=False)
|
|
|
|
def snapshot(self, timeout: float) -> Dict[str, Tuple[Dict[str, Any], Any, int]]:
|
|
now = time.time()
|
|
with self._lock:
|
|
return {
|
|
key: entry
|
|
for key, entry in self._latest.items()
|
|
if now - entry[0].get("timestamp", 0) < timeout
|
|
}
|
|
|
|
|
|
class PromptExecutionContext:
|
|
"""Context manager for prompt executions.
|
|
|
|
This context manager provides automatic prompt lifecycle management,
|
|
ensuring prompt context is properly set and cleaned up. While the current
|
|
implementation doesn't automatically clear prompts on exit (to allow for
|
|
images generated after prompt execution), it provides a clean interface
|
|
for prompt management.
|
|
|
|
Example:
|
|
with PromptExecutionContext(tracker, "beautiful landscape") as exec_id:
|
|
# Generate images here
|
|
# Images will be linked to this prompt context
|
|
pass
|
|
# Prompt context remains active for timeout period
|
|
"""
|
|
|
|
def __init__(self, prompt_tracker: PromptTracker, prompt_text: str, **kwargs):
|
|
"""
|
|
Initialize execution context.
|
|
|
|
Args:
|
|
prompt_tracker: PromptTracker instance to use for tracking
|
|
prompt_text: The prompt text to track
|
|
**kwargs: Additional prompt metadata to include in the context
|
|
"""
|
|
self.prompt_tracker = prompt_tracker
|
|
self.prompt_text = prompt_text
|
|
self.additional_data = kwargs
|
|
self.execution_id = None
|
|
|
|
def __enter__(self):
|
|
"""Enter the execution context.
|
|
|
|
Sets the current prompt in the tracker and returns the execution ID.
|
|
|
|
Returns:
|
|
Unique execution ID for this prompt
|
|
"""
|
|
self.execution_id = self.prompt_tracker.set_current_prompt(
|
|
self.prompt_text, self.additional_data
|
|
)
|
|
return self.execution_id
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Exit the execution context.
|
|
|
|
Currently does not clear the prompt context to allow images to be
|
|
generated after prompt execution completes. The prompt will be cleaned
|
|
up automatically after the timeout period.
|
|
|
|
Args:
|
|
exc_type: Exception type (if an exception occurred)
|
|
exc_val: Exception value (if an exception occurred)
|
|
exc_tb: Exception traceback (if an exception occurred)
|
|
"""
|
|
# Don't clear immediately - let the timeout handle it
|
|
# This allows images to be generated after the prompt execution completes
|
|
pass
|
|
|
|
|
|
# Singleton instance management
|
|
_tracker_instance: Optional[PromptTracker] = None
|
|
_tracker_lock = threading.Lock()
|
|
|
|
|
|
def get_prompt_tracker(db_manager) -> PromptTracker:
|
|
"""Get or create the singleton PromptTracker instance.
|
|
|
|
This ensures only one PromptTracker exists across all PromptManager nodes,
|
|
preventing duplicate image linking when multiple nodes are used.
|
|
|
|
Args:
|
|
db_manager: Database manager instance for prompt operations
|
|
|
|
Returns:
|
|
The singleton PromptTracker instance
|
|
"""
|
|
global _tracker_instance
|
|
if _tracker_instance is None:
|
|
with _tracker_lock:
|
|
if _tracker_instance is None:
|
|
_tracker_instance = PromptTracker(db_manager)
|
|
return _tracker_instance
|