spb/ai-atlas

1"""Postgres-backed job queue (`jobs` table, `FOR UPDATE SKIP LOCKED`): priorities, retries with backoff, dead letters, batches.2Local-first: no extra broker. Redis is used only for cross-process locks and caches."""3from __future__ import annotations45import asyncio6import logging7import socket8import traceback9from collections.abc import Awaitable, Callable10from datetime import UTC, datetime, timedelta11from typing import Any1213from sqlalchemy.ext.asyncio import AsyncConnection1415from aiatlas.config import settings16from aiatlas.db import execute, fetch_all, fetch_one, jsonb, transaction17from aiatlas.ids import new_id1819log = logging.getLogger(__name__)2021Handler = Callable[[dict[str, Any], dict[str, Any]], Awaitable[dict[str, Any] | None]]22_HANDLERS: dict[str, Handler] = {}232425def handler(kind: str) -> Callable[[Handler], Handler]:26    def deco(fn: Handler) -> Handler:27        _HANDLERS[kind] = fn28        return fn29    return deco303132async def enqueue(conn: AsyncConnection, kind: str, payload: dict[str, Any], *, priority: int = 5, run_after: datetime | None = None,33                  max_attempts: int = 3, batch_id: str | None = None, dedupe_key: str | None = None) -> str | None:34    jid = new_id("queue_job")35    row = await fetch_one(conn, """insert into jobs (id, kind, payload, priority, run_after, max_attempts, batch_id, dedupe_key)36                                   values (:id, :k, cast(:p as jsonb), :pr, :ra, :ma, :b, :d)37                                   on conflict (dedupe_key) where status in ('queued','running') do nothing returning id""",38                         id=jid, k=kind, p=jsonb(payload), pr=priority, ra=run_after or datetime.now(UTC), ma=max_attempts, b=batch_id, d=dedupe_key)39    return row["id"] if row else None404142async def claim_next(conn: AsyncConnection, worker: str, kinds: list[str] | None = None, *, sticky_kind: str | None = None) -> dict[str, Any] | None:43    """Claim the next job. `sticky_kind` keeps a worker on the kind it just ran while any remain: LLM jobs of one kind share a loaded44    model, and alternating kinds (extraction ↔ embeddings) makes the inference server swap models every job."""45    kind_filter = "and kind = any(cast(:kinds as text[]))" if kinds else ""46    row = await fetch_one(conn, f"""with next as (47            select id from jobs where status = 'queued' and run_after <= now() {kind_filter}48            order by case when kind = cast(:sticky as text) then 0 else 1 end, priority, run_after limit 1 for update skip locked)49        update jobs j set status = 'running', locked_by = :w, locked_at = now(), started_at = now(), attempts = attempts + 150        from next where j.id = next.id returning j.*""", w=worker, kinds=kinds or [], sticky=sticky_kind or "")51    return row525354async def complete(conn: AsyncConnection, job_id: str, result: dict[str, Any] | None = None) -> None:55    await execute(conn, "update jobs set status = 'done', finished_at = now(), error = null, payload = payload || cast(:r as jsonb) where id = :id",56                  r=jsonb({"_result": result} if result else {}), id=job_id)575859async def fail(conn: AsyncConnection, job: dict[str, Any], error: str) -> None:60    dead = job["attempts"] >= job["max_attempts"]61    backoff = timedelta(seconds=min(3600, 30 * (3 ** max(0, job["attempts"] - 1))))62    await execute(conn, """update jobs set status = :st, finished_at = case when :dead then now() else null end, error = :e, locked_by = null, locked_at = null,63                           run_after = case when :dead then run_after else now() + :backoff end where id = :id""",64                  st="dead" if dead else "queued", dead=dead, e=error[:4000], backoff=backoff, id=job["id"])656667async def queue_depth(conn: AsyncConnection) -> dict[str, Any]:68    rows = await fetch_all(conn, "select kind, status, count(*) as n from jobs where status in ('queued','running','dead') or finished_at > now() - interval '1 day' group by 1, 2")69    out: dict[str, dict[str, int]] = {}70    for r in rows:71        out.setdefault(r["kind"], {})[r["status"]] = int(r["n"])72    return out737475async def run_worker(*, concurrency: int | None = None, kinds: list[str] | None = None, stop: asyncio.Event | None = None, idle_sleep: float = 3.0) -> None:76    """Long-running worker loop. Handlers are registered with `@handler(kind)` in `aiatlas.services.handlers`."""77    import aiatlas.services.handlers  # noqa: F401  (registers handlers)7879    worker = f"{socket.gethostname()}:{new_id('queue_job')[-6:]}"80    stop = stop or asyncio.Event()81    async with transaction() as conn:82        reclaimed = await requeue_stale(conn, older_than_minutes=0, worker_prefix=socket.gethostname())83    if reclaimed:84        log.info("reclaimed jobs left running by a previous worker on this host", extra={"n": reclaimed})85    sem = asyncio.Semaphore(concurrency or settings.worker_concurrency)86    log.info("worker started", extra={"worker": worker, "kinds": kinds or "all"})8788    async def one(job: dict[str, Any]) -> None:89        try:90            fn = _HANDLERS.get(job["kind"])91            if fn is None:92                async with transaction() as conn:93                    await fail(conn, {**job, "attempts": job["max_attempts"]}, f"no handler for {job['kind']}")94                return95            try:96                result = await fn(job["payload"], job)97                async with transaction() as conn:98                    await complete(conn, job["id"], result)99            except Exception as exc:  # noqa: BLE001100                log.warning("job failed", extra={"job": job["id"], "kind": job["kind"], "error": str(exc)})101                async with transaction() as conn:102                    await fail(conn, job, f"{exc.__class__.__name__}: {exc}\n{traceback.format_exc()[-1500:]}")103        finally:104            sem.release()105106    tasks: set[asyncio.Task[None]] = set()107    last_kind: str | None = None108    while not stop.is_set():109        await sem.acquire()                      # hold a slot BEFORE claiming, so concurrency bounds claims too110        if stop.is_set():111            sem.release()112            break113        async with transaction() as conn:114            job = await claim_next(conn, worker, kinds, sticky_kind=last_kind)115        if job is not None:116            last_kind = job["kind"]117        if job is None:118            sem.release()119            try:120                await asyncio.wait_for(stop.wait(), timeout=idle_sleep)121            except TimeoutError:122                pass123            continue124        t = asyncio.create_task(one(job))125        tasks.add(t)126        t.add_done_callback(tasks.discard)127    if tasks:128        await asyncio.gather(*tasks, return_exceptions=True)129130131async def requeue_stale(conn: AsyncConnection, *, older_than_minutes: int = 120, worker_prefix: str | None = None) -> int:132    """Put back jobs left `running` by a dead worker. With `worker_prefix` (hostname), only this host's jobs are reclaimed — used at133    worker start-up so a restart (deploy, reboot) never strands jobs for two hours."""134    row = await fetch_one(conn, """with s as (update jobs set status = 'queued', locked_by = null, locked_at = null135                                   where status = 'running' and locked_at < now() - make_interval(mins => :m)136                                   and (cast(:p as text) = '' or locked_by like cast(:p as text) || ':%') returning 1) select count(*) as n from s""",137                          m=older_than_minutes, p=worker_prefix or "")138    return int(row["n"]) if row else 0139140141__all__ = ["claim_next", "complete", "enqueue", "fail", "handler", "queue_depth", "requeue_stale", "run_worker"]142