874 lines
37 KiB
Python
874 lines
37 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from collections.abc import Awaitable, Callable
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
from .database import Database, DatabaseTransientError, utcnow
|
|
|
|
logger = logging.getLogger(__name__)
|
|
JobHandler = Callable[[str, dict], Awaitable[None]]
|
|
|
|
JOB_LANES = {
|
|
"index_video": "ai",
|
|
"transcribe_audio": "ai",
|
|
"suggest_tags": "ai",
|
|
"install_models": "ai",
|
|
"prepare_ai_runtime": "ai",
|
|
"transfer_upload": "transfer",
|
|
"refresh_path": "transfer",
|
|
"aria2_download": "download",
|
|
"scan_source": "scan",
|
|
"backup_export": "transfer",
|
|
}
|
|
LANE_KINDS = {
|
|
lane: frozenset(kind for kind, value in JOB_LANES.items() if value == lane)
|
|
for lane in ("ai", "transfer", "download", "scan")
|
|
}
|
|
JOB_PRIORITIES = {
|
|
"prepare_ai_runtime": -10,
|
|
"install_models": -10,
|
|
"refresh_path": 10,
|
|
# A user-visible upload should start before passive source scans when a
|
|
# caller claims across lanes (for example during recovery or diagnostics).
|
|
"transfer_upload": 0,
|
|
"aria2_download": 10,
|
|
"scan_source": 10,
|
|
# Extract visual/OCR/face metadata before the much longer audio backlog so
|
|
# freshly uploaded videos become useful in the library promptly.
|
|
"index_video": 10,
|
|
"transcribe_audio": 20,
|
|
"suggest_tags": 40,
|
|
}
|
|
|
|
|
|
def job_lane(kind: str) -> str:
|
|
return JOB_LANES.get(kind, "ai")
|
|
|
|
|
|
class JobCancelled(RuntimeError):
|
|
pass
|
|
|
|
|
|
class JobRetry(RuntimeError):
|
|
def __init__(self, message: str, delay_seconds: float):
|
|
super().__init__(message)
|
|
self.delay_seconds = max(0.0, delay_seconds)
|
|
|
|
|
|
class JobQueue:
|
|
def __init__(self, db: Database, poll_seconds: float = 1.0, governor=None, events=None):
|
|
self.db = db
|
|
self.poll_seconds = poll_seconds
|
|
self.governor = governor
|
|
self.events = events
|
|
self.handlers: dict[str, JobHandler] = {}
|
|
self._stopping = asyncio.Event()
|
|
self.owner = f"{os.getpid()}-{uuid.uuid4().hex[:12]}"
|
|
self._progress_guard = threading.Lock()
|
|
self._progress_state: dict[str, tuple[float, float]] = {}
|
|
self._pause_state: dict[str, tuple[str | None, float]] = {}
|
|
|
|
def register(self, kind: str, handler: JobHandler) -> None:
|
|
self.handlers[kind] = handler
|
|
|
|
@staticmethod
|
|
def _resources(payload: dict) -> list[tuple[str, str]]:
|
|
resources: list[tuple[str, str]] = []
|
|
for key, resource_type in (
|
|
("video_id", "video"),
|
|
("source_id", "source"),
|
|
("upload_id", "upload"),
|
|
("collection_id", "collection"),
|
|
("backup_id", "backup"),
|
|
):
|
|
value = payload.get(key)
|
|
if isinstance(value, str) and value:
|
|
resources.append((resource_type, value))
|
|
video_ids = payload.get("video_ids")
|
|
if isinstance(video_ids, list):
|
|
resources.extend(("video", value) for value in video_ids if isinstance(value, str) and value)
|
|
return list(dict.fromkeys(resources))
|
|
|
|
def enqueue(
|
|
self,
|
|
kind: str,
|
|
payload: dict,
|
|
*,
|
|
dedupe_key: str | None = None,
|
|
delay_seconds: float = 0,
|
|
priority: int | None = None,
|
|
connection=None,
|
|
) -> str:
|
|
job_id = str(uuid.uuid4())
|
|
run_after = datetime.now(UTC) + timedelta(seconds=delay_seconds)
|
|
resolved_priority = int(priority if priority is not None else JOB_PRIORITIES.get(kind, 50))
|
|
|
|
def insert(conn) -> str:
|
|
inserted = conn.execute(
|
|
"INSERT INTO jobs(id,kind,payload_json,dedupe_key,priority,run_after,created_at) "
|
|
"VALUES(?,?,?,?,?,?,?) ON CONFLICT DO NOTHING",
|
|
(job_id, kind, json.dumps(payload), dedupe_key, resolved_priority, run_after.isoformat(), utcnow()),
|
|
).rowcount
|
|
if not inserted and dedupe_key:
|
|
row = conn.execute(
|
|
"SELECT id,status,priority FROM jobs WHERE dedupe_key=? AND status IN ('queued','running')",
|
|
(dedupe_key,),
|
|
).fetchone()
|
|
if row:
|
|
# A warm-cache follow-up can intentionally enqueue the same
|
|
# task at a higher priority. Preserve the running worker,
|
|
# but promote a queued duplicate instead of silently losing
|
|
# the caller's scheduling intent.
|
|
if row["status"] == "queued" and int(row["priority"]) > resolved_priority:
|
|
conn.execute(
|
|
"UPDATE jobs SET priority=? WHERE id=? AND status='queued' AND priority>?",
|
|
(resolved_priority, row["id"], resolved_priority),
|
|
)
|
|
return row["id"]
|
|
if not inserted:
|
|
raise sqlite3.IntegrityError("任务写入冲突，且未找到可复用的活动任务")
|
|
resources = self._resources(payload)
|
|
if resources:
|
|
conn.executemany(
|
|
"INSERT INTO job_resources(job_id,resource_type,resource_id) VALUES(?,?,?) "
|
|
"ON CONFLICT DO NOTHING",
|
|
((job_id, resource_type, resource_id) for resource_type, resource_id in resources),
|
|
)
|
|
return job_id
|
|
|
|
if connection is not None:
|
|
result = insert(connection)
|
|
else:
|
|
with self.db.transaction() as conn:
|
|
result = insert(conn)
|
|
if self.events is not None:
|
|
self.events.publish("jobs", job_id=result, action="queued", kind=kind)
|
|
return result
|
|
|
|
@staticmethod
|
|
def _public(row) -> dict:
|
|
item = dict(row)
|
|
item["lane"] = job_lane(item["kind"])
|
|
raw_diagnostics = item.pop("diagnostics_json", "{}")
|
|
try:
|
|
diagnostics = json.loads(raw_diagnostics or "{}")
|
|
except (TypeError, json.JSONDecodeError):
|
|
diagnostics = {}
|
|
item["inference_diagnostics"] = diagnostics if isinstance(diagnostics, dict) else {}
|
|
return item
|
|
|
|
def list(
|
|
self,
|
|
limit: int = 100,
|
|
*,
|
|
lane: str | None = None,
|
|
status: str | None = None,
|
|
) -> list[dict]:
|
|
filters: list[str] = []
|
|
parameters: list[object] = []
|
|
if lane:
|
|
kinds = LANE_KINDS[lane]
|
|
placeholders = ",".join("?" for _ in kinds)
|
|
filters.append(f"kind IN ({placeholders})")
|
|
parameters.extend(sorted(kinds))
|
|
if status:
|
|
filters.append("status=?")
|
|
parameters.append(status)
|
|
where = f" WHERE {' AND '.join(filters)}" if filters else ""
|
|
with self.db.read() as conn:
|
|
rows = conn.execute(
|
|
"SELECT id,kind,status,progress,message,attempts,created_at,started_at,finished_at,error,"
|
|
f"heartbeat_at,pause_reason,cancel_requested,diagnostics_json FROM jobs{where} "
|
|
"ORDER BY created_at DESC LIMIT ?",
|
|
(*parameters, min(max(limit, 1), 500)),
|
|
).fetchall()
|
|
return [self._public(row) for row in rows]
|
|
|
|
def paginate(
|
|
self,
|
|
page: int = 1,
|
|
page_size: int = 10,
|
|
*,
|
|
lane: str | None = None,
|
|
status: str | None = None,
|
|
) -> dict:
|
|
page = max(1, page)
|
|
page_size = min(max(page_size, 1), 50)
|
|
filters: list[str] = []
|
|
parameters: list[str] = []
|
|
if lane:
|
|
kinds = LANE_KINDS[lane]
|
|
placeholders = ",".join("?" for _ in kinds)
|
|
filters.append(f"kind IN ({placeholders})")
|
|
parameters.extend(sorted(kinds))
|
|
if status:
|
|
filters.append("status=?")
|
|
parameters.append(status)
|
|
where = f" WHERE {' AND '.join(filters)}" if filters else ""
|
|
with self.db.read() as conn:
|
|
total = int(conn.execute(f"SELECT count(*) FROM jobs{where}", parameters).fetchone()[0])
|
|
retryable_failed_count = len(self._retryable_failed_rows(conn, lane=lane))
|
|
pages = max(1, (total + page_size - 1) // page_size)
|
|
page = min(page, pages)
|
|
rows = conn.execute(
|
|
"SELECT id,kind,status,progress,message,attempts,created_at,started_at,finished_at,error,"
|
|
f"heartbeat_at,pause_reason,cancel_requested,diagnostics_json FROM jobs{where} "
|
|
"ORDER BY created_at DESC,id DESC LIMIT ? OFFSET ?",
|
|
(*parameters, page_size, (page - 1) * page_size),
|
|
).fetchall()
|
|
return {
|
|
"items": [self._public(row) for row in rows],
|
|
"page": page,
|
|
"page_size": page_size,
|
|
"total": total,
|
|
"pages": pages,
|
|
"retryable_failed_count": retryable_failed_count,
|
|
}
|
|
|
|
def _retryable_failed_rows(self, conn, *, lane: str | None = None, lock: bool = False):
|
|
filters = ["j.status='failed'"]
|
|
parameters: list[object] = []
|
|
if lane:
|
|
kinds = LANE_KINDS[lane]
|
|
placeholders = ",".join("?" for _ in kinds)
|
|
filters.append(f"j.kind IN ({placeholders})")
|
|
parameters.extend(sorted(kinds))
|
|
filters.append(
|
|
"NOT EXISTS(SELECT 1 FROM jobs newer WHERE newer.id<>j.id AND "
|
|
"((j.dedupe_key IS NOT NULL AND newer.dedupe_key=j.dedupe_key) OR "
|
|
"(j.dedupe_key IS NULL AND newer.dedupe_key IS NULL AND newer.kind=j.kind "
|
|
"AND newer.payload_json=j.payload_json)) AND "
|
|
"(newer.status IN ('queued','running') OR newer.created_at>j.created_at))"
|
|
)
|
|
suffix = " FOR UPDATE OF j SKIP LOCKED" if lock else ""
|
|
rows = conn.execute(
|
|
"SELECT j.id,j.kind,j.payload_json,j.dedupe_key,j.priority,j.created_at FROM jobs j WHERE "
|
|
f"{' AND '.join(filters)} ORDER BY j.created_at,j.id{suffix}",
|
|
parameters,
|
|
).fetchall()
|
|
result = []
|
|
for row in rows:
|
|
if row["kind"] not in self.handlers:
|
|
continue
|
|
try:
|
|
payload = json.loads(row["payload_json"])
|
|
except (TypeError, json.JSONDecodeError):
|
|
continue
|
|
if isinstance(payload, dict):
|
|
result.append((row, payload))
|
|
return result
|
|
|
|
def retry_failed(self, *, lane: str | None = None) -> dict:
|
|
"""Retry the latest unresolved failure for each logical task."""
|
|
|
|
if lane is not None and lane not in LANE_KINDS:
|
|
raise ValueError("未知的后台任务通道")
|
|
retried_ids: list[str] = []
|
|
with self.db.transaction() as conn:
|
|
rows = self._retryable_failed_rows(conn, lane=lane, lock=True)
|
|
for row, payload in rows:
|
|
retried_ids.append(
|
|
self.enqueue(
|
|
row["kind"],
|
|
payload,
|
|
dedupe_key=row["dedupe_key"] or f"bulk-retry:{row['id']}",
|
|
priority=int(row["priority"]),
|
|
connection=conn,
|
|
)
|
|
)
|
|
return {
|
|
"lane": lane or "all",
|
|
"retried": len(retried_ids),
|
|
"skipped": 0,
|
|
"job_ids": retried_ids,
|
|
}
|
|
|
|
def retry(self, job_id: str) -> str:
|
|
"""Create a new queued attempt while preserving the failed job record."""
|
|
|
|
with self.db.read() as conn:
|
|
row = conn.execute(
|
|
"SELECT kind,payload_json,dedupe_key,status,priority FROM jobs WHERE id=?",
|
|
(job_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
raise KeyError(job_id)
|
|
if row["status"] != "failed":
|
|
raise ValueError("只有失败的后台任务可以重试")
|
|
if row["kind"] not in self.handlers:
|
|
raise ValueError("该任务类型当前不可重试")
|
|
try:
|
|
payload = json.loads(row["payload_json"])
|
|
except (TypeError, json.JSONDecodeError) as exc:
|
|
raise ValueError("原任务参数损坏，无法重试") from exc
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("原任务参数损坏，无法重试")
|
|
return self.enqueue(
|
|
row["kind"], payload, dedupe_key=row["dedupe_key"], priority=int(row["priority"])
|
|
)
|
|
|
|
def update(self, job_id: str, progress: float, message: str = "") -> None:
|
|
value = min(max(progress, 0), 1)
|
|
now = time.monotonic()
|
|
with self._progress_guard:
|
|
previous_at, previous_value = self._progress_state.get(job_id, (0.0, -1.0))
|
|
if value < 1 and now - previous_at < 2 and abs(value - previous_value) < 0.01:
|
|
return
|
|
try:
|
|
self.db.write_with_retry(
|
|
lambda conn: conn.execute(
|
|
"UPDATE jobs SET progress=?,message=?,heartbeat_at=? WHERE id=?",
|
|
(value, message[:500], utcnow(), job_id),
|
|
),
|
|
timeout_seconds=5,
|
|
)
|
|
except DatabaseTransientError:
|
|
logger.warning("skipped progress update after a transient PostgreSQL conflict for job %s", job_id)
|
|
else:
|
|
with self._progress_guard:
|
|
self._progress_state[job_id] = (now, value)
|
|
if self.events is not None:
|
|
self.events.publish("jobs", job_id=job_id, action="progress", progress=value)
|
|
|
|
def set_diagnostics(self, job_id: str, diagnostics: dict) -> None:
|
|
"""Persist privacy-safe inference metadata for task and UI diagnostics."""
|
|
|
|
safe = {
|
|
str(key): value
|
|
for key, value in dict(diagnostics or {}).items()
|
|
if key
|
|
in {
|
|
"requested_device",
|
|
"actual_device",
|
|
"primary_backend",
|
|
"language_candidates",
|
|
"selected_language",
|
|
"fallback_scope",
|
|
"fallback_reason",
|
|
"gpu_quality_score",
|
|
"cpu_verification_quality",
|
|
}
|
|
}
|
|
payload = json.dumps(safe, ensure_ascii=False, separators=(",", ":"), default=str)
|
|
try:
|
|
self.db.write_with_retry(
|
|
lambda conn: conn.execute(
|
|
"UPDATE jobs SET diagnostics_json=?,heartbeat_at=? WHERE id=?",
|
|
(payload, utcnow(), job_id),
|
|
),
|
|
timeout_seconds=5,
|
|
)
|
|
except DatabaseTransientError:
|
|
# Diagnostics improve observability but are not part of the media
|
|
# result. A short PostgreSQL conflict must not turn an otherwise
|
|
# successful inference into a failed background job.
|
|
logger.warning("skipped diagnostics update after a transient PostgreSQL conflict for job %s", job_id)
|
|
|
|
def set_pause_reason(self, job_id: str, reason: str | None) -> None:
|
|
now = time.monotonic()
|
|
with self._progress_guard:
|
|
previous = self._pause_state.get(job_id)
|
|
if previous and previous[0] == reason and now - previous[1] < 15:
|
|
return
|
|
try:
|
|
self.db.write_with_retry(
|
|
lambda conn: conn.execute(
|
|
"UPDATE jobs SET pause_reason=?,heartbeat_at=?,message=coalesce(?,message) WHERE id=?",
|
|
(reason, utcnow(), reason, job_id),
|
|
),
|
|
timeout_seconds=5,
|
|
)
|
|
except DatabaseTransientError:
|
|
logger.warning("skipped pause heartbeat after a transient PostgreSQL conflict for job %s", job_id)
|
|
else:
|
|
with self._progress_guard:
|
|
self._pause_state[job_id] = (reason, now)
|
|
|
|
def _clear_runtime_state(self, job_id: str) -> None:
|
|
with self._progress_guard:
|
|
self._progress_state.pop(job_id, None)
|
|
self._pause_state.pop(job_id, None)
|
|
|
|
async def _write_critical(self, operation):
|
|
"""Persist a job state transition without blocking the ASGI loop.
|
|
|
|
PostgreSQL can abort a transaction on a deadlock or serialization
|
|
conflict. Keep retrying the complete state transition on a worker
|
|
thread instead of killing a lane worker or blocking the ASGI loop.
|
|
"""
|
|
|
|
while True:
|
|
try:
|
|
return await asyncio.to_thread(
|
|
self.db.write_with_retry,
|
|
operation,
|
|
timeout_seconds=5,
|
|
)
|
|
except DatabaseTransientError:
|
|
if self._stopping.is_set():
|
|
return None
|
|
logger.warning("retrying a job state after a transient PostgreSQL conflict")
|
|
try:
|
|
await asyncio.wait_for(self._stopping.wait(), timeout=self.poll_seconds)
|
|
except TimeoutError:
|
|
pass
|
|
|
|
@staticmethod
|
|
def _requeue_transient_job(conn, job_id: str) -> float:
|
|
row = conn.execute(
|
|
"SELECT attempts,cancel_requested FROM jobs WHERE id=?",
|
|
(job_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
return 0
|
|
if row["cancel_requested"]:
|
|
conn.execute(
|
|
"UPDATE jobs SET status='cancelled',lease_owner=NULL,pause_reason=NULL,"
|
|
"finished_at=?,heartbeat_at=?,message='已取消' WHERE id=?",
|
|
(utcnow(), utcnow(), job_id),
|
|
)
|
|
return 0
|
|
attempts = max(1, int(row["attempts"] or 1))
|
|
delay = float(min(30, 5 * (2 ** min(attempts - 1, 3))))
|
|
run_after = datetime.now(UTC) + timedelta(seconds=delay)
|
|
conn.execute(
|
|
"UPDATE jobs SET status='queued',run_after=?,started_at=NULL,finished_at=NULL,error=NULL,"
|
|
"lease_owner=NULL,heartbeat_at=?,pause_reason=NULL,"
|
|
"message='数据库事务冲突，已自动恢复，稍后重试' WHERE id=?",
|
|
(run_after.isoformat(), utcnow(), job_id),
|
|
)
|
|
return delay
|
|
|
|
@staticmethod
|
|
def _finish_failed(conn, job_id: str, error: str) -> None:
|
|
"""Finish a worker failure without overriding a concurrent cancellation.
|
|
|
|
Video deletion deliberately removes the database row and remote object
|
|
immediately after requesting cancellation. A running media handler can
|
|
therefore observe ``video missing`` before reaching its next checkpoint.
|
|
Resolve the final state atomically so that the already-committed cancel
|
|
request wins over that expected teardown error.
|
|
"""
|
|
|
|
now = utcnow()
|
|
conn.execute(
|
|
"UPDATE jobs SET status=CASE WHEN cancel_requested=1 THEN 'cancelled' ELSE 'failed' END,"
|
|
"error=CASE WHEN cancel_requested=1 THEN NULL ELSE ? END,finished_at=?,pause_reason=NULL,"
|
|
"lease_owner=NULL,heartbeat_at=?,message=CASE WHEN cancel_requested=1 THEN '已取消' ELSE message END "
|
|
"WHERE id=?",
|
|
(error, now, now, job_id),
|
|
)
|
|
|
|
@staticmethod
|
|
def _finish_retry(conn, job_id: str, run_after: str, message: str) -> None:
|
|
"""Requeue a retry unless a concurrent caller has cancelled the job."""
|
|
|
|
now = utcnow()
|
|
conn.execute(
|
|
"UPDATE jobs SET status=CASE WHEN cancel_requested=1 THEN 'cancelled' ELSE 'queued' END,"
|
|
"run_after=?,started_at=NULL,"
|
|
"finished_at=CASE WHEN cancel_requested=1 THEN ? ELSE NULL END,error=NULL,lease_owner=NULL,"
|
|
"heartbeat_at=?,pause_reason=NULL,"
|
|
"message=CASE WHEN cancel_requested=1 THEN '已取消' ELSE ? END WHERE id=?",
|
|
(run_after, now, now, message, job_id),
|
|
)
|
|
|
|
@staticmethod
|
|
def _finish_completed(conn, job_id: str) -> None:
|
|
"""Commit success unless cancellation was requested before finalization."""
|
|
|
|
now = utcnow()
|
|
conn.execute(
|
|
"UPDATE jobs SET status=CASE WHEN cancel_requested=1 THEN 'cancelled' ELSE 'completed' END,"
|
|
"progress=CASE WHEN cancel_requested=1 THEN progress ELSE 1 END,finished_at=?,pause_reason=NULL,"
|
|
"lease_owner=NULL,heartbeat_at=?,error=NULL,"
|
|
"message=CASE WHEN cancel_requested=1 THEN '已取消' ELSE message END WHERE id=?",
|
|
(now, now, job_id),
|
|
)
|
|
|
|
@staticmethod
|
|
def _finish_cancelled(conn, job_id: str) -> None:
|
|
"""Finish a cancellation unless its request was concurrently rolled back.
|
|
|
|
Deletion compensation can clear ``cancel_requested`` after a running
|
|
worker has already left its checkpoint. In that race the payload is
|
|
still valid again, so return the job to the queue instead of committing
|
|
a stale cancelled state after ``restore_cancelled`` has completed.
|
|
"""
|
|
|
|
now = utcnow()
|
|
conn.execute(
|
|
"UPDATE jobs SET status=CASE WHEN cancel_requested=1 THEN 'cancelled' ELSE 'queued' END,"
|
|
"started_at=CASE WHEN cancel_requested=1 THEN started_at ELSE NULL END,"
|
|
"finished_at=CASE WHEN cancel_requested=1 THEN ? ELSE NULL END,lease_owner=NULL,"
|
|
"heartbeat_at=?,pause_reason=NULL,error=NULL,"
|
|
"message=CASE WHEN cancel_requested=1 THEN '已取消' ELSE '取消已回滚，任务重新排队' END "
|
|
"WHERE id=?",
|
|
(now, now, job_id),
|
|
)
|
|
|
|
def _job_state(self, job_id: str):
|
|
with self.db.read() as conn:
|
|
return conn.execute("SELECT kind,status,cancel_requested FROM jobs WHERE id=?", (job_id,)).fetchone()
|
|
|
|
def _cancel_requested(self, job_id: str) -> bool:
|
|
row = self._job_state(job_id)
|
|
# Direct service-level maintenance/tests may use an ephemeral progress
|
|
# id. A missing row is not a cancellation signal; real worker jobs are
|
|
# retained for their full lifecycle.
|
|
return bool(row and row["cancel_requested"])
|
|
|
|
def checkpoint(self, job_id: str, message: str | None = None, *, persist: bool = True) -> None:
|
|
state = self._job_state(job_id)
|
|
if state and state["cancel_requested"]:
|
|
raise JobCancelled("任务已取消")
|
|
|
|
def paused(reason: str) -> None:
|
|
self.set_pause_reason(job_id, reason)
|
|
|
|
def resumed() -> None:
|
|
self.set_pause_reason(job_id, None)
|
|
|
|
# Resource pressure must only park work that has actually been leased
|
|
# by the background worker. Direct maintenance calls and service-level
|
|
# tests can legitimately use a queued or ephemeral progress id and
|
|
# must not wait forever merely because their temporary filesystem is
|
|
# smaller than the NAS reserve policy.
|
|
if self.governor is not None and state is not None and state["status"] == "running":
|
|
if self._should_yield_to_model_maintenance(job_id):
|
|
raise JobCancelled("模型安装优先，音频识别已让路")
|
|
self.governor.wait_sync(
|
|
lane=job_lane(state["kind"]),
|
|
paused=paused,
|
|
resumed=resumed,
|
|
cancelled=lambda: (
|
|
self._cancel_requested(job_id)
|
|
or self._stopping.is_set()
|
|
or self._should_yield_to_model_maintenance(job_id)
|
|
),
|
|
)
|
|
if self._should_yield_to_model_maintenance(job_id):
|
|
raise JobCancelled("模型安装优先，音频识别已让路")
|
|
if self._cancel_requested(job_id):
|
|
raise JobCancelled("任务已取消")
|
|
if persist:
|
|
try:
|
|
self.db.write_with_retry(
|
|
lambda conn: conn.execute(
|
|
"UPDATE jobs SET pause_reason=NULL,heartbeat_at=?,message=coalesce(?,message) WHERE id=?",
|
|
(utcnow(), message, job_id),
|
|
),
|
|
timeout_seconds=5,
|
|
)
|
|
except DatabaseTransientError:
|
|
logger.warning("skipped checkpoint heartbeat after a transient conflict for job %s", job_id)
|
|
|
|
def request_cancel(self, job_id: str) -> None:
|
|
with self.db.transaction() as conn:
|
|
row = conn.execute("SELECT kind,status FROM jobs WHERE id=?", (job_id,)).fetchone()
|
|
if row is None:
|
|
raise KeyError(job_id)
|
|
if row["kind"] == "backup_export" and row["status"] == "running":
|
|
raise ValueError("备份已经开始生成，不能中途取消")
|
|
if row["status"] == "queued":
|
|
now = utcnow()
|
|
conn.execute(
|
|
"UPDATE jobs SET status='cancelled',cancel_requested=1,finished_at=?,message='已取消' WHERE id=?",
|
|
(now, job_id),
|
|
)
|
|
if row["kind"] == "backup_export":
|
|
conn.execute(
|
|
"UPDATE backup_exports SET status='failed',secret_blob=NULL,error='备份任务已取消',"
|
|
"finished_at=? WHERE id IN (SELECT resource_id FROM job_resources "
|
|
"WHERE job_id=? AND resource_type='backup')",
|
|
(now, job_id),
|
|
)
|
|
elif row["status"] == "running":
|
|
conn.execute("UPDATE jobs SET cancel_requested=1,message='正在安全停止' WHERE id=?", (job_id,))
|
|
else:
|
|
raise ValueError("该任务已经结束")
|
|
if self.events is not None:
|
|
self.events.publish("jobs", job_id=job_id, action="cancel")
|
|
|
|
def cancel_by_kind(self, kind: str) -> list[str]:
|
|
with self.db.transaction() as conn:
|
|
rows = conn.execute(
|
|
"SELECT id,status FROM jobs WHERE kind=? AND status IN ('queued','running') ORDER BY created_at",
|
|
(kind,),
|
|
).fetchall()
|
|
cancelled: list[str] = []
|
|
for row in rows:
|
|
if row["status"] == "queued":
|
|
conn.execute(
|
|
"UPDATE jobs SET status='cancelled',cancel_requested=1,finished_at=?,"
|
|
"message='模型安装优先，已取消' WHERE id=?",
|
|
(utcnow(), row["id"]),
|
|
)
|
|
else:
|
|
conn.execute(
|
|
"UPDATE jobs SET cancel_requested=1,message='模型安装优先，正在停止' WHERE id=?",
|
|
(row["id"],),
|
|
)
|
|
cancelled.append(row["id"])
|
|
return cancelled
|
|
|
|
def has_queued_model_maintenance(self) -> bool:
|
|
with self.db.read() as conn:
|
|
return bool(
|
|
conn.execute(
|
|
"SELECT 1 FROM jobs WHERE kind IN ('install_models','prepare_ai_runtime') "
|
|
"AND status='queued' LIMIT 1"
|
|
).fetchone()
|
|
)
|
|
|
|
def _should_yield_to_model_maintenance(self, job_id: str) -> bool:
|
|
with self.db.read() as conn:
|
|
row = conn.execute("SELECT kind,status FROM jobs WHERE id=?", (job_id,)).fetchone()
|
|
return bool(row and row["kind"] == "transcribe_audio" and row["status"] == "running") and (
|
|
self.has_queued_model_maintenance()
|
|
)
|
|
|
|
def cancel_for_video(
|
|
self,
|
|
video_id: str,
|
|
*,
|
|
exclude_job_ids: set[str] | None = None,
|
|
connection=None,
|
|
) -> list[str]:
|
|
excluded = exclude_job_ids or set()
|
|
|
|
def cancel(conn) -> list[str]:
|
|
rows = conn.execute(
|
|
"SELECT DISTINCT j.id,j.status FROM jobs j JOIN job_resources r ON r.job_id=j.id "
|
|
"WHERE r.resource_type='video' AND r.resource_id=? AND j.status IN ('queued','running')",
|
|
(video_id,),
|
|
).fetchall()
|
|
cancelled: list[str] = []
|
|
now = utcnow()
|
|
for row in rows:
|
|
if row["id"] in excluded:
|
|
continue
|
|
if row["status"] == "queued":
|
|
changed = conn.execute(
|
|
"UPDATE jobs SET status='cancelled',cancel_requested=1,finished_at=?,message='已取消' "
|
|
"WHERE id=? AND status='queued'",
|
|
(now, row["id"]),
|
|
).rowcount
|
|
else:
|
|
changed = conn.execute(
|
|
"UPDATE jobs SET cancel_requested=1,message='正在安全停止' "
|
|
"WHERE id=? AND status='running'",
|
|
(row["id"],),
|
|
).rowcount
|
|
if changed:
|
|
cancelled.append(row["id"])
|
|
return cancelled
|
|
|
|
if connection is not None:
|
|
return cancel(connection)
|
|
with self.db.transaction() as conn:
|
|
return cancel(conn)
|
|
|
|
def restore_cancelled(self, job_ids: list[str]) -> int:
|
|
"""Undo a deletion-time cancellation after the delete is compensated.
|
|
|
|
Queued jobs are returned to the queue with their original payload and
|
|
priority. A running worker may already have observed the cancellation,
|
|
so both the still-running and just-cancelled states are handled.
|
|
Completed/failed jobs are intentionally left as historical truth.
|
|
"""
|
|
|
|
identifiers = list(dict.fromkeys(str(value) for value in job_ids if value))
|
|
if not identifiers:
|
|
return 0
|
|
placeholders = ",".join("?" for _ in identifiers)
|
|
with self.db.transaction() as conn:
|
|
restored = conn.execute(
|
|
f"UPDATE jobs SET status=CASE WHEN status='cancelled' THEN 'queued' ELSE status END,"
|
|
"cancel_requested=0,finished_at=CASE WHEN status='cancelled' THEN NULL ELSE finished_at END,"
|
|
"lease_owner=CASE WHEN status='cancelled' THEN NULL ELSE lease_owner END,"
|
|
"pause_reason=NULL,error=CASE WHEN status='cancelled' THEN NULL ELSE error END,"
|
|
"message=CASE WHEN status='cancelled' THEN '删除已回滚，任务重新排队' ELSE message END "
|
|
f"WHERE id IN ({placeholders}) AND status IN ('queued','running','cancelled')",
|
|
identifiers,
|
|
).rowcount
|
|
return int(restored)
|
|
|
|
def cleanup_history(self, *, completed_days: int = 30, failed_days: int = 90) -> dict[str, int]:
|
|
completed_cutoff = (datetime.now(UTC) - timedelta(days=max(1, completed_days))).isoformat()
|
|
failed_cutoff = (datetime.now(UTC) - timedelta(days=max(1, failed_days))).isoformat()
|
|
with self.db.transaction() as conn:
|
|
completed = conn.execute(
|
|
"DELETE FROM jobs WHERE status IN ('completed','cancelled') "
|
|
"AND coalesce(finished_at,created_at)<?",
|
|
(completed_cutoff,),
|
|
).rowcount
|
|
failed = conn.execute(
|
|
"DELETE FROM jobs WHERE status='failed' AND coalesce(finished_at,created_at)<?",
|
|
(failed_cutoff,),
|
|
).rowcount
|
|
return {"completed_cancelled": int(completed), "failed": int(failed)}
|
|
|
|
def recover_stale(self, stale_seconds: int = 300) -> int:
|
|
cutoff = (datetime.now(UTC) - timedelta(seconds=stale_seconds)).isoformat()
|
|
with self.db.transaction() as conn:
|
|
return conn.execute(
|
|
"UPDATE jobs SET status='queued',started_at=NULL,lease_owner=NULL,heartbeat_at=NULL,"
|
|
"pause_reason=NULL,cancel_requested=0,message='应用重启，等待继续执行' "
|
|
"WHERE status='running' AND (lease_owner IS NULL OR lease_owner NOT LIKE ? "
|
|
"OR heartbeat_at IS NULL OR heartbeat_at<?)",
|
|
(f"{self.owner}:%", cutoff),
|
|
).rowcount
|
|
|
|
def _claim(
|
|
self,
|
|
*,
|
|
lane: str | None = None,
|
|
allowed_kinds: frozenset[str] | None = None,
|
|
excluded_kinds: frozenset[str] = frozenset(),
|
|
) -> tuple[str, str, dict] | None:
|
|
lane_name = lane or "ai"
|
|
if lane and allowed_kinds is None:
|
|
allowed_kinds = LANE_KINDS[lane]
|
|
if self.governor is not None and self.governor.pressure_reason(lane=lane_name):
|
|
return None
|
|
filters = ["status='queued'", "cancel_requested=0", "run_after<=?"]
|
|
parameters: list[str] = [utcnow()]
|
|
if allowed_kinds:
|
|
placeholders = ",".join("?" for _ in allowed_kinds)
|
|
filters.append(f"kind IN ({placeholders})")
|
|
parameters.extend(sorted(allowed_kinds))
|
|
if excluded_kinds:
|
|
placeholders = ",".join("?" for _ in excluded_kinds)
|
|
filters.append(f"kind NOT IN ({placeholders})")
|
|
parameters.extend(sorted(excluded_kinds))
|
|
with self.db.transaction() as conn:
|
|
row = conn.execute(
|
|
f"SELECT id,kind,payload_json FROM jobs WHERE {' AND '.join(filters)} "
|
|
"ORDER BY priority,created_at LIMIT 1 FOR UPDATE SKIP LOCKED",
|
|
parameters,
|
|
).fetchone()
|
|
if not row:
|
|
return None
|
|
changed = conn.execute(
|
|
"UPDATE jobs SET status='running',started_at=?,attempts=attempts+1,lease_owner=?,"
|
|
"heartbeat_at=?,pause_reason=NULL WHERE id=? AND status='queued' AND cancel_requested=0",
|
|
(utcnow(), f"{self.owner}:{lane_name}", utcnow(), row["id"]),
|
|
).rowcount
|
|
if not changed:
|
|
return None
|
|
return row["id"], row["kind"], json.loads(row["payload_json"])
|
|
|
|
async def _execute_claimed(self, job_id: str, kind: str, payload: dict) -> None:
|
|
handler = self.handlers.get(kind)
|
|
if handler is None:
|
|
error = f"没有任务处理器：{kind}"
|
|
await self._write_critical(
|
|
lambda conn, error=error, job_id=job_id: self._finish_failed(conn, job_id, error)
|
|
)
|
|
self._clear_runtime_state(job_id)
|
|
return
|
|
try:
|
|
await asyncio.to_thread(self.checkpoint, job_id)
|
|
await handler(job_id, payload)
|
|
except JobCancelled:
|
|
await self._write_critical(
|
|
lambda conn, job_id=job_id: self._finish_cancelled(conn, job_id)
|
|
)
|
|
except JobRetry as exc:
|
|
run_after = datetime.now(UTC) + timedelta(seconds=exc.delay_seconds)
|
|
retry_message = str(exc)[:500]
|
|
await self._write_critical(
|
|
lambda conn, run_after=run_after, retry_message=retry_message, job_id=job_id: self._finish_retry(
|
|
conn, job_id, run_after.isoformat(), retry_message
|
|
)
|
|
)
|
|
except DatabaseTransientError:
|
|
delay = await self._write_critical(
|
|
lambda conn, job_id=job_id: self._requeue_transient_job(conn, job_id)
|
|
)
|
|
logger.warning(
|
|
"job %s hit a transient PostgreSQL conflict and was requeued for %.0fs",
|
|
job_id,
|
|
delay or 0,
|
|
)
|
|
except asyncio.CancelledError:
|
|
|
|
def release(conn, job_id=job_id):
|
|
state = conn.execute("SELECT cancel_requested FROM jobs WHERE id=?", (job_id,)).fetchone()
|
|
if state and state["cancel_requested"]:
|
|
conn.execute(
|
|
"UPDATE jobs SET status='cancelled',lease_owner=NULL,pause_reason=NULL,"
|
|
"finished_at=?,message='已取消' WHERE id=?",
|
|
(utcnow(), job_id),
|
|
)
|
|
else:
|
|
conn.execute(
|
|
"UPDATE jobs SET status='queued',started_at=NULL,lease_owner=NULL,heartbeat_at=NULL,"
|
|
"pause_reason=NULL WHERE id=?",
|
|
(job_id,),
|
|
)
|
|
|
|
await self._write_critical(release)
|
|
raise
|
|
except Exception as exc:
|
|
logger.exception("job %s failed", job_id)
|
|
# Full exception details belong in the service log. API clients
|
|
# receive a short, actionable message without leaking paths,
|
|
# credentials or an implementation traceback.
|
|
error = str(exc).strip()[:1000] or "任务执行失败，请查看服务日志"
|
|
await self._write_critical(
|
|
lambda conn, error=error, job_id=job_id: self._finish_failed(conn, job_id, error)
|
|
)
|
|
else:
|
|
await self._write_critical(
|
|
lambda conn, job_id=job_id: self._finish_completed(conn, job_id)
|
|
)
|
|
finally:
|
|
self._clear_runtime_state(job_id)
|
|
if self.events is not None:
|
|
self.events.publish("jobs", job_id=job_id, action="finished")
|
|
|
|
async def run(
|
|
self,
|
|
*,
|
|
lane: str | None = None,
|
|
allowed_kinds: frozenset[str] | None = None,
|
|
excluded_kinds: frozenset[str] = frozenset(),
|
|
) -> None:
|
|
if lane and allowed_kinds is None:
|
|
allowed_kinds = LANE_KINDS[lane]
|
|
while not self._stopping.is_set():
|
|
try:
|
|
claimed = await asyncio.to_thread(
|
|
self._claim,
|
|
lane=lane,
|
|
allowed_kinds=allowed_kinds,
|
|
excluded_kinds=excluded_kinds,
|
|
)
|
|
except DatabaseTransientError:
|
|
logger.warning("job lane %s is retrying after a PostgreSQL conflict", lane or "ai")
|
|
try:
|
|
await asyncio.wait_for(self._stopping.wait(), timeout=self.poll_seconds)
|
|
except TimeoutError:
|
|
pass
|
|
continue
|
|
if not claimed:
|
|
try:
|
|
await asyncio.wait_for(self._stopping.wait(), timeout=self.poll_seconds)
|
|
except TimeoutError:
|
|
pass
|
|
continue
|
|
await self._execute_claimed(*claimed)
|
|
|
|
def stop(self) -> None:
|
|
self._stopping.set()
|