SPB Git forge

spb/llm-api

Public
0commits 0branches 0releases
0 Bsize
maindefault branch
—last push
6.3 KB · 163 lines python
Raw Blame History
1"""Small asyncio job system (downloads, scans, benchmarks, harvests)."""23from __future__ import annotations45import asyncio6import json7import logging8import time9import uuid10from typing import Any, Awaitable, Callable1112from .db import Database13from .events import bus1415log = logging.getLogger("llm_api.jobs")1617QUEUED, RUNNING, COMPLETED, FAILED, CANCELLED = "queued", "running", "completed", "failed", "cancelled"181920class Job:21    def __init__(self, id: str, kind: str, title: str, payload: dict):22        self.id = id23        self.kind = kind24        self.title = title25        self.payload = payload26        self.status = QUEUED27        self.progress = 0.028        self.detail: dict[str, Any] = {}29        self.result: Any = None30        self.error: str | None = None31        self.created_at = time.time()32        self.started_at: float | None = None33        self.finished_at: float | None = None34        self.cancel_event = asyncio.Event()35        self.task: asyncio.Task | None = None36        self._last_flush = 0.03738    def to_dict(self) -> dict:39        return {40            "id": self.id, "kind": self.kind, "title": self.title, "payload": self.payload, "status": self.status,41            "progress": round(self.progress, 4), "detail": self.detail, "result": self.result, "error": self.error,42            "created_at": self.created_at, "started_at": self.started_at, "finished_at": self.finished_at,43        }4445    @property46    def cancelled(self) -> bool:47        return self.cancel_event.is_set()484950class JobRunner:51    def __init__(self, db: Database, max_concurrent: int = 3):52        self.db = db53        self.jobs: dict[str, Job] = {}54        self.sem = asyncio.Semaphore(max_concurrent)55        self._download_lock = asyncio.Semaphore(1)  # downloads are serialized (SSD + bandwidth)5657    async def start(self) -> None:58        # Mark jobs left running by a crashed server as failed59        await self.db.execute("UPDATE jobs SET status=?, error=?, finished_at=? WHERE status IN (?, ?)",60                              (FAILED, "server restarted", time.time(), QUEUED, RUNNING))6162    def submit(self, kind: str, title: str, payload: dict,63               fn: Callable[[Job], Awaitable[Any]], *, exclusive_download: bool = False) -> Job:64        job = Job(uuid.uuid4().hex[:12], kind, title, payload)65        self.jobs[job.id] = job66        asyncio.create_task(self._persist(job))6768        async def _run():69            async with self.sem:70                if exclusive_download:71                    async with self._download_lock:72                        await self._execute(job, fn)73                else:74                    await self._execute(job, fn)7576        job.task = asyncio.create_task(_run(), name=f"job-{kind}-{job.id}")77        bus.publish("job", job.to_dict())78        return job7980    async def _execute(self, job: Job, fn) -> None:81        if job.cancelled:82            job.status = CANCELLED83            job.finished_at = time.time()84            await self._persist(job)85            return86        job.status = RUNNING87        job.started_at = time.time()88        await self._persist(job)89        bus.publish("job", job.to_dict())90        try:91            job.result = await fn(job)92            job.status = CANCELLED if job.cancelled else COMPLETED93            job.progress = 1.0 if job.status == COMPLETED else job.progress94        except asyncio.CancelledError:95            job.status = CANCELLED96        except Exception as e:97            log.exception("job %s failed", job.id)98            job.status = FAILED99            job.error = f"{type(e).__name__}: {e}"100        job.finished_at = time.time()101        await self._persist(job)102        bus.publish("job", job.to_dict())103104    def update(self, job: Job, progress: float | None = None, **detail: Any) -> None:105        if progress is not None:106            job.progress = max(0.0, min(1.0, progress))107        if detail:108            job.detail.update(detail)109        now = time.time()110        if now - job._last_flush > 0.5:111            job._last_flush = now112            bus.publish("job", job.to_dict())113            asyncio.create_task(self._persist(job))114115    async def _persist(self, job: Job) -> None:116        try:117            await self.db.execute(118                "INSERT INTO jobs(id, kind, status, title, payload, progress, detail, result, error, created_at, started_at, finished_at) "119                "VALUES(?,?,?,?,?,?,?,?,?,?,?,?) ON CONFLICT(id) DO UPDATE SET status=excluded.status, progress=excluded.progress, "120                "detail=excluded.detail, result=excluded.result, error=excluded.error, started_at=excluded.started_at, finished_at=excluded.finished_at",121                (job.id, job.kind, job.status, job.title, json.dumps(job.payload), job.progress, json.dumps(job.detail),122                 json.dumps(job.result, default=str) if job.result is not None else None, job.error, job.created_at,123                 job.started_at, job.finished_at))124        except Exception:125            log.exception("persist job failed")126127    async def cancel(self, job_id: str) -> bool:128        job = self.jobs.get(job_id)129        if not job:130            return False131        if job.status in (COMPLETED, FAILED, CANCELLED):132            return False133        job.cancel_event.set()134        if job.status == QUEUED and job.task:135            job.task.cancel()136        return True137138    def list(self, kinds: set[str] | None = None, limit: int = 100) -> list[dict]:139        items = sorted(self.jobs.values(), key=lambda j: -j.created_at)140        if kinds:141            items = [j for j in items if j.kind in kinds]142        return [j.to_dict() for j in items[:limit]]143144    async def history(self, kinds: set[str] | None = None, limit: int = 100) -> list[dict]:145        rows = await self.db.fetchall("SELECT * FROM jobs ORDER BY created_at DESC LIMIT ?", (limit,))146        out = []147        for r in rows:148            if kinds and r["kind"] not in kinds:149                continue150            for k in ("payload", "detail", "result"):151                if r.get(k):152                    try:153                        r[k] = json.loads(r[k])154                    except Exception:155                        pass156            # prefer live state157            live = self.jobs.get(r["id"])158            out.append(live.to_dict() if live else r)159        return out160161    def get(self, job_id: str) -> Job | None:162        return self.jobs.get(job_id)163