asus
5c7c122105
This commit includes major additions and updates to the frontend and backend architectures, introducing new dataset management, live counting features, batch processing, and triage logic. Includes new UI pages, components, and API routes.
256 lines
8.6 KiB
Python
256 lines
8.6 KiB
Python
"""Job queue: one worker thread, because there is one GPU (REQ-070).
|
|
|
|
Jobs are rows in SQLite, so the list survives a restart (REQ-071). A job that was
|
|
still running when the process died is marked failed at the next startup — it
|
|
cannot be resumed, and pretending otherwise would be worse than saying so.
|
|
|
|
Handlers register themselves by job type:
|
|
|
|
@jobs.handler("extract")
|
|
def _extract(job: Job) -> None:
|
|
...
|
|
|
|
A handler reports progress with `job.progress(i, n, "...")`, writes user-facing
|
|
lines with `job.log("...")`, and checks `job.cancelled` between units of work.
|
|
Raising anything marks the job failed with that exception on it.
|
|
"""
|
|
|
|
import json
|
|
import queue
|
|
import threading
|
|
import time
|
|
import traceback
|
|
from typing import Callable, Dict, List, Optional
|
|
|
|
from backend import db
|
|
|
|
MAX_LOG_LINES = 500
|
|
PROGRESS_FLUSH_SECONDS = 0.5
|
|
|
|
JOB_TYPES = ("extract", "autolabel", "merge", "train", "count", "clock-scan",
|
|
"truck-scan")
|
|
GPU_JOB_TYPES = ("autolabel", "train", "count", "truck-scan")
|
|
"""`extract` is ffmpeg and `merge` is file copying — neither touches the card,
|
|
so neither should be able to block an interactive assist."""
|
|
gpu_lock = threading.Lock()
|
|
"""Held for the duration of any GPU work. The job worker takes it around a
|
|
handler; the interactive assist route takes it around one SAM3 call. One card,
|
|
one holder (REQ-065)."""
|
|
|
|
|
|
class Job:
|
|
"""One queued unit of work. The database row is the source of truth; this
|
|
object is the handle a handler writes through."""
|
|
|
|
def __init__(self, row):
|
|
self.id: int = row["id"]
|
|
self.type: str = row["type"]
|
|
self.project_id: Optional[int] = row["project_id"]
|
|
self.batch_id: Optional[int] = row["batch_id"]
|
|
self.params: dict = json.loads(row["params"])
|
|
self.status: str = row["status"]
|
|
self.current: int = row["progress"]
|
|
self.total: int = row["total"]
|
|
self.message: str = row["message"]
|
|
self.error: str = row["error"]
|
|
self.lines: List[str] = row["log"].splitlines() if row["log"] else []
|
|
self.created_at: float = row["created_at"]
|
|
self.started_at: Optional[float] = row["started_at"]
|
|
self.finished_at: Optional[float] = row["finished_at"]
|
|
self._flushed_at = 0.0
|
|
|
|
# ---- what handlers call ------------------------------------------------
|
|
|
|
@property
|
|
def cancelled(self) -> bool:
|
|
return self.id in _cancelled
|
|
|
|
def progress(self, current: int, total: Optional[int] = None,
|
|
message: Optional[str] = None) -> None:
|
|
self.current = current
|
|
if total is not None:
|
|
self.total = total
|
|
if message is not None:
|
|
self.message = message
|
|
# Throttled: a 3000-frame job would otherwise write 3000 times.
|
|
if time.time() - self._flushed_at >= PROGRESS_FLUSH_SECONDS:
|
|
self.flush()
|
|
|
|
def log(self, message: str) -> None:
|
|
self.lines.append(f"[{time.strftime('%H:%M:%S')}] {message}")
|
|
if len(self.lines) > MAX_LOG_LINES:
|
|
del self.lines[: len(self.lines) - MAX_LOG_LINES]
|
|
self.flush()
|
|
|
|
def flush(self) -> None:
|
|
self._flushed_at = time.time()
|
|
with db.cursor() as cur:
|
|
cur.execute(
|
|
"""UPDATE jobs SET status = ?, progress = ?, total = ?, message = ?,
|
|
error = ?, log = ?, started_at = ?, finished_at = ?
|
|
WHERE id = ?""",
|
|
(self.status, self.current, self.total, self.message, self.error,
|
|
"\n".join(self.lines), self.started_at, self.finished_at, self.id),
|
|
)
|
|
|
|
# ---- serialisation -----------------------------------------------------
|
|
|
|
def to_dict(self) -> dict:
|
|
end = self.finished_at or time.time()
|
|
return {
|
|
"id": self.id,
|
|
"type": self.type,
|
|
"project_id": self.project_id,
|
|
"batch_id": self.batch_id,
|
|
"params": self.params,
|
|
"status": self.status,
|
|
"progress": self.current,
|
|
"total": self.total,
|
|
"message": self.message,
|
|
"error": self.error,
|
|
"log": self.lines,
|
|
"created_at": self.created_at,
|
|
"started_at": self.started_at,
|
|
"finished_at": self.finished_at,
|
|
"elapsed": end - (self.started_at or self.created_at),
|
|
}
|
|
|
|
|
|
_handlers: Dict[str, Callable[[Job], None]] = {}
|
|
_queue: "queue.Queue[int]" = queue.Queue()
|
|
_cancelled: set = set()
|
|
_worker: Optional[threading.Thread] = None
|
|
_worker_lock = threading.Lock()
|
|
|
|
|
|
def handler(job_type: str):
|
|
"""Register the function that runs jobs of this type."""
|
|
if job_type not in JOB_TYPES:
|
|
raise ValueError(f"Unknown job type: {job_type}")
|
|
|
|
def decorate(function: Callable[[Job], None]) -> Callable[[Job], None]:
|
|
_handlers[job_type] = function
|
|
return function
|
|
|
|
return decorate
|
|
|
|
|
|
def create(job_type: str, params: Optional[dict] = None, project_id: Optional[int] = None,
|
|
batch_id: Optional[int] = None, message: str = "") -> Job:
|
|
if job_type not in _handlers:
|
|
raise ValueError(f"No handler registered for job type: {job_type}")
|
|
with db.cursor() as cur:
|
|
cur.execute(
|
|
"""INSERT INTO jobs (project_id, batch_id, type, status, params, message, created_at)
|
|
VALUES (?, ?, ?, 'queued', ?, ?, ?)""",
|
|
(project_id, batch_id, job_type, json.dumps(params or {}), message, time.time()),
|
|
)
|
|
job_id = cur.lastrowid
|
|
job = get(job_id)
|
|
assert job is not None
|
|
_queue.put(job_id)
|
|
_ensure_worker()
|
|
return job
|
|
|
|
|
|
def get(job_id: int) -> Optional[Job]:
|
|
with db.cursor() as cur:
|
|
cur.execute("SELECT * FROM jobs WHERE id = ?", (job_id,))
|
|
row = cur.fetchone()
|
|
return Job(row) if row else None
|
|
|
|
|
|
def listing(project_id: Optional[int] = None, limit: int = 50) -> List[Job]:
|
|
with db.cursor() as cur:
|
|
if project_id is None:
|
|
cur.execute("SELECT * FROM jobs ORDER BY created_at DESC LIMIT ?", (limit,))
|
|
else:
|
|
cur.execute(
|
|
"SELECT * FROM jobs WHERE project_id = ? ORDER BY created_at DESC LIMIT ?",
|
|
(project_id, limit),
|
|
)
|
|
return [Job(row) for row in cur.fetchall()]
|
|
|
|
|
|
def running_types() -> List[str]:
|
|
"""Job types occupying the worker right now — the GPU is shared with the
|
|
interactive SAM3 calls the review editor makes."""
|
|
with db.cursor() as cur:
|
|
cur.execute("SELECT DISTINCT type FROM jobs WHERE status = 'running'")
|
|
return [row[0] for row in cur.fetchall()]
|
|
|
|
|
|
def cancel(job_id: int) -> bool:
|
|
job = get(job_id)
|
|
if job is None or job.status in ("done", "failed", "cancelled"):
|
|
return False
|
|
_cancelled.add(job_id)
|
|
if job.status == "queued":
|
|
# Never started, so no handler will notice the flag — close it out here.
|
|
job.status = "cancelled"
|
|
job.finished_at = time.time()
|
|
job.log("Cancelled before it started")
|
|
else:
|
|
job.log("Cancellation requested")
|
|
return True
|
|
|
|
|
|
def recover() -> int:
|
|
"""Close out jobs left behind by a previous process (REQ-071)."""
|
|
with db.cursor() as cur:
|
|
cur.execute(
|
|
"""UPDATE jobs SET status = 'failed', error = 'interrupted by a server restart',
|
|
finished_at = ?
|
|
WHERE status IN ('queued', 'running')""",
|
|
(time.time(),),
|
|
)
|
|
return cur.rowcount
|
|
|
|
|
|
def _ensure_worker() -> None:
|
|
global _worker
|
|
with _worker_lock:
|
|
if _worker is None or not _worker.is_alive():
|
|
_worker = threading.Thread(target=_worker_loop, name="job-worker", daemon=True)
|
|
_worker.start()
|
|
|
|
|
|
def _worker_loop() -> None:
|
|
while True:
|
|
job_id = _queue.get()
|
|
job = get(job_id)
|
|
if job is None:
|
|
continue
|
|
if job.id in _cancelled:
|
|
_finish(job, "cancelled")
|
|
continue
|
|
_run(job)
|
|
|
|
|
|
def _run(job: Job) -> None:
|
|
job.status = "running"
|
|
job.started_at = time.time()
|
|
job.flush()
|
|
try:
|
|
if job.type in GPU_JOB_TYPES:
|
|
with gpu_lock:
|
|
_handlers[job.type](job)
|
|
else:
|
|
_handlers[job.type](job)
|
|
except Exception as exc:
|
|
job.error = f"{type(exc).__name__}: {exc}"
|
|
job.log(f"FAILED: {job.error}")
|
|
job.log(traceback.format_exc().strip().splitlines()[-1])
|
|
_finish(job, "failed")
|
|
return
|
|
_finish(job, "cancelled" if job.cancelled else "done")
|
|
|
|
|
|
|
|
def _finish(job: Job, status: str) -> None:
|
|
_cancelled.discard(job.id)
|
|
job.status = status
|
|
job.finished_at = time.time()
|
|
job.message = ""
|
|
job.flush()
|