131 lines
4.5 KiB
Python
131 lines
4.5 KiB
Python
"""UI job bookkeeping for long-running extension processes.
|
|
|
|
When a process is launched by the UI it gets AITK_JOB_ID and a sqlite db
|
|
path; status / step / info are written to the Job row and the UI's stop
|
|
button sets Job.stop. This mirrors what BaseCaptioner does inline, as a
|
|
standalone helper so other UI-driven processes (the inference engine) can
|
|
share it. Outside the UI (no job id or no db file) every call is a no-op.
|
|
"""
|
|
|
|
import os
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
from typing import Callable, Literal, Optional
|
|
|
|
AITK_Status = Literal["running", "stopped", "error", "completed", "queued"]
|
|
|
|
|
|
class UIJobStatus:
|
|
def __init__(self, sqlite_db_path: str = "./aitk_db.db"):
|
|
self.sqlite_db_path = sqlite_db_path
|
|
job_id = os.environ.get("AITK_JOB_ID", None)
|
|
self.job_id = job_id.strip() if job_id else None
|
|
self.is_ui_job = self.job_id is not None and os.path.exists(sqlite_db_path)
|
|
if self.is_ui_job:
|
|
print(f"Using SQLite database at {sqlite_db_path}")
|
|
print(f'Job ID: "{self.job_id}"')
|
|
self._lock = threading.Lock()
|
|
self._watcher = None
|
|
# collapse bursts of identical writes (progress ticks)
|
|
self._last_info = None
|
|
|
|
# ---- db plumbing ----
|
|
def _connect(self):
|
|
conn = sqlite3.connect(self.sqlite_db_path, timeout=30.0)
|
|
conn.isolation_level = None
|
|
return conn
|
|
|
|
def _retry(self, fn, max_retries=3, base_delay=2.0):
|
|
last = None
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
return fn()
|
|
except sqlite3.OperationalError as e:
|
|
if "database is locked" not in str(e):
|
|
raise
|
|
last = e
|
|
if attempt < max_retries:
|
|
time.sleep(base_delay * (2**attempt))
|
|
raise last
|
|
|
|
def _write(self, query: str, params: tuple):
|
|
if not self.is_ui_job:
|
|
return
|
|
|
|
def _do():
|
|
with self._lock, self._connect() as conn:
|
|
cur = conn.cursor()
|
|
cur.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
cur.execute(query, params)
|
|
finally:
|
|
cur.execute("COMMIT")
|
|
|
|
try:
|
|
self._retry(_do)
|
|
except Exception as e:
|
|
print(f"[AITK] Warning: db write failed: {e}")
|
|
|
|
def _read_flag(self, column: str) -> bool:
|
|
if not self.is_ui_job:
|
|
return False
|
|
|
|
def _do():
|
|
with self._lock, self._connect() as conn:
|
|
cur = conn.cursor()
|
|
cur.execute(f"SELECT {column} FROM Job WHERE id = ?", (self.job_id,))
|
|
row = cur.fetchone()
|
|
return False if row is None else row[0] == 1
|
|
|
|
return self._retry(_do)
|
|
|
|
# ---- public api ----
|
|
def update_status(self, status: AITK_Status, info: Optional[str] = None):
|
|
if info is None:
|
|
self._write("UPDATE Job SET status = ? WHERE id = ?", (status, self.job_id))
|
|
else:
|
|
self._last_info = info
|
|
self._write(
|
|
"UPDATE Job SET status = ?, info = ? WHERE id = ?",
|
|
(status, info, self.job_id),
|
|
)
|
|
|
|
def update_info(self, info: str):
|
|
if info == self._last_info:
|
|
return
|
|
self._last_info = info
|
|
self._write("UPDATE Job SET info = ? WHERE id = ?", (info, self.job_id))
|
|
|
|
def update_key(self, key: str, value):
|
|
# key is a column name from our own code, never user input
|
|
self._write(f"UPDATE Job SET {key} = ? WHERE id = ?", (str(value), self.job_id))
|
|
|
|
def should_stop(self) -> bool:
|
|
return self._read_flag("stop")
|
|
|
|
def should_return_to_queue(self) -> bool:
|
|
return self._read_flag("return_to_queue")
|
|
|
|
def start_stop_watcher(self, on_stop: Callable[[str], None], interval_sec: float = 2.0):
|
|
"""Poll the stop / return_to_queue flags on a daemon thread and call
|
|
on_stop("stopped" | "queued") once when one is set."""
|
|
if not self.is_ui_job or self._watcher is not None:
|
|
return
|
|
|
|
def _loop():
|
|
while True:
|
|
try:
|
|
if self.should_stop():
|
|
on_stop("stopped")
|
|
return
|
|
if self.should_return_to_queue():
|
|
on_stop("queued")
|
|
return
|
|
except Exception:
|
|
pass
|
|
time.sleep(interval_sec)
|
|
|
|
self._watcher = threading.Thread(target=_loop, daemon=True)
|
|
self._watcher.start()
|