f
This commit is contained in:
@@ -11,8 +11,17 @@ from datetime import datetime
|
||||
from threading import Thread, Lock
|
||||
import time
|
||||
|
||||
from API.folder import API_folder_store
|
||||
from SYS.logger import log
|
||||
from SYS.database import (
|
||||
db,
|
||||
insert_worker,
|
||||
update_worker,
|
||||
append_worker_stdout,
|
||||
get_worker_stdout as db_get_worker_stdout,
|
||||
get_active_workers as db_get_active_workers,
|
||||
get_worker as db_get_worker,
|
||||
expire_running_workers as db_expire_running_workers
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -157,7 +166,6 @@ class WorkerLoggingHandler(logging.StreamHandler):
|
||||
def __init__(
|
||||
self,
|
||||
worker_id: str,
|
||||
db: API_folder_store,
|
||||
manager: Optional["WorkerManager"] = None,
|
||||
buffer_size: int = 50,
|
||||
):
|
||||
@@ -165,12 +173,10 @@ class WorkerLoggingHandler(logging.StreamHandler):
|
||||
|
||||
Args:
|
||||
worker_id: ID of the worker to capture logs for
|
||||
db: Reference to LocalLibraryDB for storing logs
|
||||
buffer_size: Number of logs to buffer before flushing to DB
|
||||
"""
|
||||
super().__init__()
|
||||
self.worker_id = worker_id
|
||||
self.db = db
|
||||
self.manager = manager
|
||||
self.buffer_size = buffer_size
|
||||
self.buffer: list[str] = []
|
||||
@@ -232,7 +238,7 @@ class WorkerLoggingHandler(logging.StreamHandler):
|
||||
channel="log"
|
||||
)
|
||||
else:
|
||||
self.db.append_worker_stdout(
|
||||
append_worker_stdout(
|
||||
self.worker_id,
|
||||
log_text,
|
||||
channel="log"
|
||||
@@ -255,29 +261,22 @@ class WorkerLoggingHandler(logging.StreamHandler):
|
||||
|
||||
|
||||
class WorkerManager:
|
||||
"""Manages persistent worker tasks with auto-refresh capability."""
|
||||
"""Manages persistent worker tasks using the central medios.db."""
|
||||
|
||||
def __init__(self, library_root: Path, auto_refresh_interval: float = 2.0):
|
||||
def __init__(self, auto_refresh_interval: float = 2.0):
|
||||
"""Initialize the worker manager.
|
||||
|
||||
Args:
|
||||
library_root: Root directory for the local library database
|
||||
auto_refresh_interval: Seconds between auto-refresh checks (0 = disabled)
|
||||
"""
|
||||
self.library_root = Path(library_root)
|
||||
self.db = API_folder_store(library_root)
|
||||
self.auto_refresh_interval = auto_refresh_interval
|
||||
self.refresh_callbacks: List[Callable] = []
|
||||
self.refresh_thread: Optional[Thread] = None
|
||||
self._stop_refresh = False
|
||||
self._lock = Lock()
|
||||
# Reuse the DB's own lock so there is exactly one lock guarding the
|
||||
# sqlite connection (and it is safe for re-entrant/nested DB usage).
|
||||
self._db_lock = self.db._db_lock
|
||||
self.worker_handlers: Dict[str,
|
||||
WorkerLoggingHandler] = {} # Track active handlers
|
||||
self._worker_last_step: Dict[str,
|
||||
str] = {}
|
||||
self.worker_handlers: Dict[str, WorkerLoggingHandler] = {}
|
||||
self._worker_last_step: Dict[str, str] = {}
|
||||
|
||||
# Buffered stdout/log batching to reduce DB lock contention.
|
||||
self._stdout_buffers: Dict[Tuple[str, str], List[str]] = {}
|
||||
self._stdout_buffer_sizes: Dict[Tuple[str, str], int] = {}
|
||||
@@ -328,13 +327,11 @@ class WorkerManager:
|
||||
Count of workers updated.
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
return self.db.expire_running_workers(
|
||||
older_than_seconds=older_than_seconds,
|
||||
status=status,
|
||||
reason=reason,
|
||||
worker_id_prefix=worker_id_prefix,
|
||||
)
|
||||
return db_expire_running_workers(
|
||||
older_than_seconds=older_than_seconds,
|
||||
status=status,
|
||||
reason=reason or "Stale worker expired"
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(f"Failed to expire stale workers: {exc}", exc_info=True)
|
||||
return 0
|
||||
@@ -362,7 +359,7 @@ class WorkerManager:
|
||||
The logging handler that was created, or None if there was an error
|
||||
"""
|
||||
try:
|
||||
handler = WorkerLoggingHandler(worker_id, self.db, manager=self)
|
||||
handler = WorkerLoggingHandler(worker_id, manager=self)
|
||||
with self._lock:
|
||||
self.worker_handlers[worker_id] = handler
|
||||
|
||||
@@ -437,16 +434,13 @@ class WorkerManager:
|
||||
True if worker was inserted successfully
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
result = self.db.insert_worker(
|
||||
worker_id,
|
||||
worker_type,
|
||||
title,
|
||||
description,
|
||||
total_steps,
|
||||
pipe=pipe
|
||||
)
|
||||
if result > 0:
|
||||
success = insert_worker(
|
||||
worker_id,
|
||||
worker_type,
|
||||
title,
|
||||
description
|
||||
)
|
||||
if success:
|
||||
logger.debug(
|
||||
f"[WorkerManager] Tracking worker: {worker_id} ({worker_type})"
|
||||
)
|
||||
@@ -482,18 +476,16 @@ class WorkerManager:
|
||||
if progress > 0:
|
||||
kwargs["progress"] = progress
|
||||
if current_step:
|
||||
kwargs["current_step"] = current_step
|
||||
kwargs["details"] = current_step
|
||||
if details:
|
||||
kwargs["description"] = details
|
||||
if error:
|
||||
kwargs["error_message"] = error
|
||||
|
||||
if kwargs:
|
||||
kwargs["last_updated"] = datetime.now().isoformat()
|
||||
if "current_step" in kwargs and kwargs["current_step"]:
|
||||
self._worker_last_step[worker_id] = str(kwargs["current_step"])
|
||||
with self._db_lock:
|
||||
return self.db.update_worker(worker_id, **kwargs)
|
||||
if "details" in kwargs and kwargs["details"]:
|
||||
self._worker_last_step[worker_id] = str(kwargs["details"])
|
||||
return update_worker(worker_id, **kwargs)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
@@ -515,7 +507,7 @@ class WorkerManager:
|
||||
worker_id: Unique identifier for the worker
|
||||
result: Result status ('completed', 'error', 'cancelled')
|
||||
error_msg: Error message if any
|
||||
result_data: Result data as JSON string
|
||||
result_data: Result data as JSON string (saved in details)
|
||||
|
||||
Returns:
|
||||
True if update was successful
|
||||
@@ -526,16 +518,15 @@ class WorkerManager:
|
||||
except Exception:
|
||||
pass
|
||||
kwargs = {
|
||||
"status": result,
|
||||
"completed_at": datetime.now().isoformat()
|
||||
"status": "finished",
|
||||
"result": result
|
||||
}
|
||||
if error_msg:
|
||||
kwargs["error_message"] = error_msg
|
||||
if result_data:
|
||||
kwargs["result_data"] = result_data
|
||||
kwargs["details"] = result_data
|
||||
|
||||
with self._db_lock:
|
||||
success = self.db.update_worker(worker_id, **kwargs)
|
||||
success = update_worker(worker_id, **kwargs)
|
||||
logger.info(f"[WorkerManager] Worker finished: {worker_id} ({result})")
|
||||
self._worker_last_step.pop(worker_id, None)
|
||||
return success
|
||||
@@ -553,8 +544,7 @@ class WorkerManager:
|
||||
List of active worker dictionaries
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
return self.db.get_active_workers()
|
||||
return db_get_active_workers()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[WorkerManager] Error getting active workers: {e}",
|
||||
@@ -572,14 +562,9 @@ class WorkerManager:
|
||||
List of finished worker dictionaries
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
all_workers = self.db.get_all_workers(limit=limit)
|
||||
# Filter to only finished workers
|
||||
finished = [
|
||||
w for w in all_workers
|
||||
if w.get("status") in ["completed", "error", "cancelled"]
|
||||
]
|
||||
return finished
|
||||
# We don't have a get_all_workers in database.py yet, but we'll use a local query
|
||||
rows = db.fetchall(f"SELECT * FROM workers WHERE status = 'finished' ORDER BY updated_at DESC LIMIT {limit}")
|
||||
return [dict(row) for row in rows]
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[WorkerManager] Error getting finished workers: {e}",
|
||||
@@ -597,7 +582,13 @@ class WorkerManager:
|
||||
Worker data or None if not found
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
return db_get_worker(worker_id)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[WorkerManager] Error getting worker {worker_id}: {e}",
|
||||
exc_info=True
|
||||
)
|
||||
return None
|
||||
return self.db.get_worker(worker_id)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
@@ -815,14 +806,11 @@ class WorkerManager:
|
||||
ok = True
|
||||
for wid, ch, step, payload in pending_flush:
|
||||
try:
|
||||
with self._db_lock:
|
||||
result = self.db.append_worker_stdout(
|
||||
wid,
|
||||
payload,
|
||||
step=step,
|
||||
channel=ch
|
||||
)
|
||||
ok = ok and result
|
||||
append_worker_stdout(
|
||||
wid,
|
||||
payload,
|
||||
channel=ch
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[WorkerManager] Error flushing stdout for {wid}: {e}",
|
||||
@@ -851,7 +839,6 @@ class WorkerManager:
|
||||
if not chunks:
|
||||
return True
|
||||
text = "".join(chunks)
|
||||
step = self._stdout_buffer_steps.get(key)
|
||||
self._stdout_buffers[key] = []
|
||||
self._stdout_buffer_sizes[key] = 0
|
||||
self._stdout_last_flush[key] = time.monotonic()
|
||||
@@ -860,13 +847,12 @@ class WorkerManager:
|
||||
if not text:
|
||||
return True
|
||||
try:
|
||||
with self._db_lock:
|
||||
return self.db.append_worker_stdout(
|
||||
worker_id,
|
||||
text,
|
||||
step=step,
|
||||
channel=channel,
|
||||
)
|
||||
append_worker_stdout(
|
||||
worker_id,
|
||||
text,
|
||||
channel=channel,
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[WorkerManager] Error flushing stdout for {worker_id}: {e}",
|
||||
@@ -884,8 +870,7 @@ class WorkerManager:
|
||||
Worker's stdout or empty string
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
return self.db.get_worker_stdout(worker_id)
|
||||
return db_get_worker_stdout(worker_id)
|
||||
except Exception as e:
|
||||
logger.error(f"[WorkerManager] Error getting stdout: {e}", exc_info=True)
|
||||
return ""
|
||||
@@ -909,21 +894,20 @@ class WorkerManager:
|
||||
True if clear was successful
|
||||
"""
|
||||
try:
|
||||
with self._db_lock:
|
||||
return self.db.clear_worker_stdout(worker_id)
|
||||
# Not implemented in database.py yet, but we'll add it or skip it
|
||||
db.execute("DELETE FROM worker_stdout WHERE worker_id = ?", (worker_id,))
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"[WorkerManager] Error clearing stdout: {e}", exc_info=True)
|
||||
return False
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the worker manager and database connection."""
|
||||
"""Close the worker manager."""
|
||||
self.stop_auto_refresh()
|
||||
try:
|
||||
self._flush_all_stdout_buffers()
|
||||
except Exception:
|
||||
pass
|
||||
with self._db_lock:
|
||||
self.db.close()
|
||||
logger.info("[WorkerManager] Closed")
|
||||
|
||||
def _flush_all_stdout_buffers(self) -> None:
|
||||
|
||||
Reference in New Issue
Block a user