"""Postgres-backed job queue (`jobs` table, `FOR UPDATE SKIP LOCKED`): priorities, retries with backoff, dead letters, batches. Local-first: no extra broker. Redis is used only for cross-process locks and caches.""" from __future__ import annotations import asyncio import logging import socket import traceback from collections.abc import Awaitable, Callable from datetime import UTC, datetime, timedelta from typing import Any from sqlalchemy.ext.asyncio import AsyncConnection from aiatlas.config import settings from aiatlas.db import execute, fetch_all, fetch_one, jsonb, transaction from aiatlas.ids import new_id log = logging.getLogger(__name__) Handler = Callable[[dict[str, Any], dict[str, Any]], Awaitable[dict[str, Any] | None]] _HANDLERS: dict[str, Handler] = {} def handler(kind: str) -> Callable[[Handler], Handler]: def deco(fn: Handler) -> Handler: _HANDLERS[kind] = fn return fn return deco async def enqueue(conn: AsyncConnection, kind: str, payload: dict[str, Any], *, priority: int = 5, run_after: datetime | None = None, max_attempts: int = 3, batch_id: str | None = None, dedupe_key: str | None = None) -> str | None: jid = new_id("queue_job") row = await fetch_one(conn, """insert into jobs (id, kind, payload, priority, run_after, max_attempts, batch_id, dedupe_key) values (:id, :k, cast(:p as jsonb), :pr, :ra, :ma, :b, :d) on conflict (dedupe_key) where status in ('queued','running') do nothing returning id""", 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) return row["id"] if row else None async def claim_next(conn: AsyncConnection, worker: str, kinds: list[str] | None = None, *, sticky_kind: str | None = None) -> dict[str, Any] | None: """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 loaded model, and alternating kinds (extraction ↔ embeddings) makes the inference server swap models every job.""" kind_filter = "and kind = any(cast(:kinds as text[]))" if kinds else "" row = await fetch_one(conn, f"""with next as ( select id from jobs where status = 'queued' and run_after <= now() {kind_filter} order by case when kind = cast(:sticky as text) then 0 else 1 end, priority, run_after limit 1 for update skip locked) update jobs j set status = 'running', locked_by = :w, locked_at = now(), started_at = now(), attempts = attempts + 1 from next where j.id = next.id returning j.*""", w=worker, kinds=kinds or [], sticky=sticky_kind or "") return row async def complete(conn: AsyncConnection, job_id: str, result: dict[str, Any] | None = None) -> None: await execute(conn, "update jobs set status = 'done', finished_at = now(), error = null, payload = payload || cast(:r as jsonb) where id = :id", r=jsonb({"_result": result} if result else {}), id=job_id) async def fail(conn: AsyncConnection, job: dict[str, Any], error: str) -> None: dead = job["attempts"] >= job["max_attempts"] backoff = timedelta(seconds=min(3600, 30 * (3 ** max(0, job["attempts"] - 1)))) 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, run_after = case when :dead then run_after else now() + :backoff end where id = :id""", st="dead" if dead else "queued", dead=dead, e=error[:4000], backoff=backoff, id=job["id"]) async def queue_depth(conn: AsyncConnection) -> dict[str, Any]: 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") out: dict[str, dict[str, int]] = {} for r in rows: out.setdefault(r["kind"], {})[r["status"]] = int(r["n"]) return out async def run_worker(*, concurrency: int | None = None, kinds: list[str] | None = None, stop: asyncio.Event | None = None, idle_sleep: float = 3.0) -> None: """Long-running worker loop. Handlers are registered with `@handler(kind)` in `aiatlas.services.handlers`.""" import aiatlas.services.handlers # noqa: F401 (registers handlers) worker = f"{socket.gethostname()}:{new_id('queue_job')[-6:]}" stop = stop or asyncio.Event() async with transaction() as conn: reclaimed = await requeue_stale(conn, older_than_minutes=0, worker_prefix=socket.gethostname()) if reclaimed: log.info("reclaimed jobs left running by a previous worker on this host", extra={"n": reclaimed}) sem = asyncio.Semaphore(concurrency or settings.worker_concurrency) log.info("worker started", extra={"worker": worker, "kinds": kinds or "all"}) async def one(job: dict[str, Any]) -> None: try: fn = _HANDLERS.get(job["kind"]) if fn is None: async with transaction() as conn: await fail(conn, {**job, "attempts": job["max_attempts"]}, f"no handler for {job['kind']}") return try: result = await fn(job["payload"], job) async with transaction() as conn: await complete(conn, job["id"], result) except Exception as exc: # noqa: BLE001 log.warning("job failed", extra={"job": job["id"], "kind": job["kind"], "error": str(exc)}) async with transaction() as conn: await fail(conn, job, f"{exc.__class__.__name__}: {exc}\n{traceback.format_exc()[-1500:]}") finally: sem.release() tasks: set[asyncio.Task[None]] = set() last_kind: str | None = None while not stop.is_set(): await sem.acquire() # hold a slot BEFORE claiming, so concurrency bounds claims too if stop.is_set(): sem.release() break async with transaction() as conn: job = await claim_next(conn, worker, kinds, sticky_kind=last_kind) if job is not None: last_kind = job["kind"] if job is None: sem.release() try: await asyncio.wait_for(stop.wait(), timeout=idle_sleep) except TimeoutError: pass continue t = asyncio.create_task(one(job)) tasks.add(t) t.add_done_callback(tasks.discard) if tasks: await asyncio.gather(*tasks, return_exceptions=True) async def requeue_stale(conn: AsyncConnection, *, older_than_minutes: int = 120, worker_prefix: str | None = None) -> int: """Put back jobs left `running` by a dead worker. With `worker_prefix` (hostname), only this host's jobs are reclaimed — used at worker start-up so a restart (deploy, reboot) never strands jobs for two hours.""" row = await fetch_one(conn, """with s as (update jobs set status = 'queued', locked_by = null, locked_at = null where status = 'running' and locked_at < now() - make_interval(mins => :m) and (cast(:p as text) = '' or locked_by like cast(:p as text) || ':%') returning 1) select count(*) as n from s""", m=older_than_minutes, p=worker_prefix or "") return int(row["n"]) if row else 0 __all__ = ["claim_next", "complete", "enqueue", "fail", "handler", "queue_depth", "requeue_stale", "run_worker"]