171 lines
5.3 KiB
Python
171 lines
5.3 KiB
Python
from __future__ import annotations
|
|
|
|
import queue
|
|
import threading
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from .metadata_store import MetadataStore
|
|
from .schemas import JobDetail, JobEvent, JobStatus
|
|
|
|
|
|
@dataclass
|
|
class JobState:
|
|
status: JobStatus
|
|
events: list[JobEvent] = field(default_factory=list)
|
|
subscribers: list[queue.Queue] = field(default_factory=list)
|
|
|
|
|
|
class JobManager:
|
|
def __init__(self, store: Optional[MetadataStore] = None) -> None:
|
|
self._lock = threading.Lock()
|
|
self.store = store
|
|
self._jobs: dict[str, JobState] = {}
|
|
self._restore_from_store()
|
|
|
|
def _restore_from_store(self) -> None:
|
|
if not self.store:
|
|
return
|
|
for status in self.store.load_jobs():
|
|
self._jobs[status.job_id] = JobState(
|
|
status=status,
|
|
events=self.store.load_events(status.job_id, limit=100),
|
|
)
|
|
|
|
def create_job(self) -> str:
|
|
job_id = uuid.uuid4().hex
|
|
now = datetime.now(timezone.utc)
|
|
status = JobStatus(
|
|
job_id=job_id,
|
|
status="queued",
|
|
stage="queued",
|
|
progress_percent=0,
|
|
message="任务已创建",
|
|
created_at=now,
|
|
updated_at=now,
|
|
artifacts={},
|
|
error="",
|
|
)
|
|
with self._lock:
|
|
self._jobs[job_id] = JobState(status=status)
|
|
if self.store:
|
|
self.store.upsert_job(status)
|
|
return job_id
|
|
|
|
def exists(self, job_id: str) -> bool:
|
|
with self._lock:
|
|
return job_id in self._jobs
|
|
|
|
def get_status(self, job_id: str) -> JobStatus:
|
|
with self._lock:
|
|
return self._jobs[job_id].status
|
|
|
|
def get_detail(self, job_id: str) -> JobDetail:
|
|
with self._lock:
|
|
state = self._jobs[job_id]
|
|
return JobDetail(status=state.status, recent_events=state.events[-100:])
|
|
|
|
def delete_job(self, job_id: str) -> None:
|
|
with self._lock:
|
|
self._jobs.pop(job_id, None)
|
|
if self.store:
|
|
self.store.delete_job(job_id)
|
|
|
|
def set_artifacts(self, job_id: str, artifacts: dict[str, str]) -> None:
|
|
with self._lock:
|
|
status = self._jobs[job_id].status
|
|
status.artifacts = artifacts
|
|
status.updated_at = datetime.now(timezone.utc)
|
|
if self.store:
|
|
self.store.upsert_job(status)
|
|
|
|
def set_status(
|
|
self,
|
|
job_id: str,
|
|
*,
|
|
status: str,
|
|
stage: str,
|
|
progress_percent: int,
|
|
message: str,
|
|
error: str | None = None,
|
|
elapsed_seconds: float | None = None,
|
|
) -> None:
|
|
with self._lock:
|
|
s = self._jobs[job_id].status
|
|
s.status = status
|
|
s.stage = stage
|
|
s.progress_percent = progress_percent
|
|
s.message = message
|
|
s.updated_at = datetime.now(timezone.utc)
|
|
if error is not None:
|
|
s.error = error
|
|
if elapsed_seconds is not None:
|
|
s.elapsed_seconds = elapsed_seconds
|
|
if self.store:
|
|
self.store.upsert_job(s)
|
|
|
|
def add_event(
|
|
self,
|
|
job_id: str,
|
|
*,
|
|
kind: str,
|
|
stage: str,
|
|
progress_percent: int,
|
|
message: str,
|
|
elapsed_seconds: float | None = None,
|
|
) -> None:
|
|
event = JobEvent(
|
|
type=kind,
|
|
stage=stage,
|
|
progress_percent=progress_percent,
|
|
message=message,
|
|
timestamp=datetime.now(timezone.utc),
|
|
elapsed_seconds=elapsed_seconds,
|
|
)
|
|
with self._lock:
|
|
state = self._jobs[job_id]
|
|
state.events.append(event)
|
|
state.status.stage = stage
|
|
state.status.progress_percent = progress_percent
|
|
state.status.message = message
|
|
state.status.updated_at = event.timestamp
|
|
if elapsed_seconds is not None:
|
|
state.status.elapsed_seconds = elapsed_seconds
|
|
for sub in state.subscribers:
|
|
sub.put(event)
|
|
if self.store:
|
|
self.store.add_event(job_id, event)
|
|
self.store.upsert_job(state.status)
|
|
|
|
def subscribe(self, job_id: str) -> queue.Queue:
|
|
q: queue.Queue = queue.Queue()
|
|
with self._lock:
|
|
state = self._jobs[job_id]
|
|
# A fast job can emit preview_ready/completed before the browser
|
|
# finishes opening the SSE connection. Replay the existing event
|
|
# history into this subscriber so timing and preview updates are
|
|
# never lost; the lock also makes the snapshot/registration
|
|
# atomic with respect to new events.
|
|
for event in state.events:
|
|
q.put(event)
|
|
state.subscribers.append(q)
|
|
return q
|
|
|
|
def unsubscribe(self, job_id: str, q: queue.Queue) -> None:
|
|
with self._lock:
|
|
subs = self._jobs[job_id].subscribers
|
|
if q in subs:
|
|
subs.remove(q)
|
|
|
|
@staticmethod
|
|
def resolve_artifact_path(status: JobStatus, kind: str) -> Path:
|
|
if kind not in status.artifacts:
|
|
raise KeyError(kind)
|
|
p = status.artifacts[kind]
|
|
if not p:
|
|
raise FileNotFoundError(kind)
|
|
return Path(p)
|