256 lines
8.5 KiB
Python
256 lines
8.5 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")
|
|
GPU_JOB_TYPES = ("autolabel", "train")
|
|
"""`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()
|