FabledCurator/backend/app/services/ml/gpu_jobs.py

"""GPU-job queue engine (#114): enqueue / lease / heartbeat / complete / fail
/ release / recover_orphaned.

Backs the HTTP API the desktop agent pulls work from. The lease claims pending
OR expired-leased jobs with FOR UPDATE SKIP LOCKED, so concurrent agents/workers
never grab the same job. Orphan recovery is three-layered: a graceful agent stop
calls release() to hand its in-flight jobs back instantly; a hard crash is caught
by recover_orphaned() (a 60s beat sweep) which resets expired leases to pending;
and the lease itself reclaims expired leases as a final backstop. Result-writing
(regions) is done by the API handler via RegionService; complete() just closes.
"""

from datetime import UTC, datetime, timedelta

from sqlalchemy import and_, delete, exists, func, or_, select, update
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import aliased

from ...models import GpuJob

# Lease window. Kept comfortably above any single job (a capped-frame video embed
# is tens of seconds) so a live, heartbeating worker is never falsely expired,
# but short enough that a hard crash recovers fast once the sweep fires.
DEFAULT_LEASE_TTL = 180     # seconds an agent holds a job before it can be re-leased
DEFAULT_BATCH = 8
MAX_ATTEMPTS = 3

# Poison-loop backstops. `attempts` counts LEASES GRANTED (incremented in
# lease()), but fail()'s MAX_ATTEMPTS cap only fires when the agent reports a
# failure — a job that keeps coming back via release() (transient handback) or
# lease expiry (agent crash/wedge) never gets a verdict and would cycle forever.
# The orphan sweep converts those to 'error': an expired lease that has already
# been granted EXPIRED_POISON_CAP leases is presumed to kill/wedge the agent,
# and a pending job granted PENDING_POISON_CAP leases without ever completing is
# presumed poisoned (e.g. a transfer that stalls every time). Both stay
# resurrectable via /retry_errors, which resets attempts.
EXPIRED_POISON_CAP = MAX_ATTEMPTS + 2
PENDING_POISON_CAP = 10


def error_dedupe_statements():
    """DELETEs enforcing: at most ONE error row per (image, task), and none that
    a live or succeeded row makes moot. The 2026-07-02 tombstone loop (backfill
    skip-lists lacked 'error') minted a duplicate error row per bad file per
    hour; running these before every backfill and inside /retry_errors keeps the
    error count reading as "distinct failing files" and stops a retry fanning
    one file out into several duplicate pending jobs. Shared by the sync beat
    task and the async API route so both prune by the SAME predicate.
    Execution order matters: moot rows first, then older duplicates (the newest
    error — the freshest reason — survives)."""
    other = aliased(GpuJob)
    same_pair = and_(
        other.image_record_id == GpuJob.image_record_id,
        other.task == GpuJob.task,
    )
    moot = (
        delete(GpuJob)
        .where(
            GpuJob.status == "error",
            exists().where(
                same_pair, other.status.in_(["pending", "leased", "done"]),
            ),
        )
        .execution_options(synchronize_session=False)
    )
    older_dupe = (
        delete(GpuJob)
        .where(
            GpuJob.status == "error",
            exists().where(
                same_pair,
                other.status == "error",
                or_(
                    other.updated_at > GpuJob.updated_at,
                    and_(other.updated_at == GpuJob.updated_at,
                         other.id > GpuJob.id),
                ),
            ),
        )
        .execution_options(synchronize_session=False)
    )
    return [moot, older_dupe]


def recover_statements(now: datetime) -> dict:
    """UPDATEs for the orphan sweep, keyed by outcome; insertion order IS the
    required execution order ('recovered' must run after 'poison_expired', which
    claims the crash-loopers out of the same expired-lease pool)."""
    expired = and_(GpuJob.status == "leased", GpuJob.lease_expires_at < now)
    unlease = {"lease_token": None, "leased_at": None, "lease_expires_at": None,
               "updated_at": now}
    return {
        "poison_expired": (
            update(GpuJob)
            .where(expired, GpuJob.attempts >= EXPIRED_POISON_CAP)
            .values(
                status="error",
                # Keep the job's last stored failure reason — it's the triage
                # signal for WHY the loop happened.
                error=func.concat(
                    f"poisoned: lease expired after {EXPIRED_POISON_CAP}+ lease "
                    "attempts (job repeatedly crashes or wedges the agent?); "
                    "last error: ",
                    func.coalesce(GpuJob.error, "none"),
                ),
                **unlease,
            )
        ),
        "recovered": update(GpuJob).where(expired).values(
            status="pending", **unlease,
        ),
        "poison_pending": (
            update(GpuJob)
            .where(GpuJob.status == "pending",
                   GpuJob.attempts >= PENDING_POISON_CAP)
            .values(
                status="error",
                error=func.concat(
                    f"poisoned: {PENDING_POISON_CAP}+ lease attempts without "
                    "ever completing (transfer stalls every time?); "
                    "last error: ",
                    func.coalesce(GpuJob.error, "none"),
                ),
                updated_at=now,
            )
        ),
    }


class GpuJobService:
    def __init__(self, session: AsyncSession):
        self.session = session

    async def enqueue(self, image_id: int, task: str) -> GpuJob | None:
        """Queue a (image, task) job. Idempotent: returns None if one is already
        pending/leased for the same pair (no duplicate work)."""
        dup = (
            await self.session.execute(
                select(GpuJob.id).where(
                    GpuJob.image_record_id == image_id,
                    GpuJob.task == task,
                    GpuJob.status.in_(["pending", "leased"]),
                )
            )
        ).first()
        if dup:
            return None
        job = GpuJob(image_record_id=image_id, task=task, status="pending")
        self.session.add(job)
        await self.session.flush()
        return job

    async def lease(
        self, token: str, batch_size: int = DEFAULT_BATCH, ttl: int = DEFAULT_LEASE_TTL
    ) -> list[GpuJob]:
        """Claim up to batch_size pending (or expired-leased) jobs for `token`.

        Two phases so each hits a partial index (0070) and stays O(batch) no
        matter how many done/error rows have accumulated: the pending pool is the
        hot path; expired leases are reclaimed only when pending can't fill the
        batch (a crashed agent's work — rare). The old single OR-query walked the
        primary key past the whole done-prefix in id order → O(done), which is
        why leasing crawled — and the DB saturated — as the run progressed."""
        now = datetime.now(UTC)

        async def _claim(condition, limit: int) -> list[int]:
            return list(
                (
                    await self.session.execute(
                        select(GpuJob.id).where(condition)
                        .order_by(GpuJob.id).limit(limit)
                        .with_for_update(skip_locked=True)
                    )
                ).scalars().all()
            )

        picked = await _claim(GpuJob.status == "pending", batch_size)
        if len(picked) < batch_size:     # pending exhausted → reclaim expired leases
            picked += await _claim(
                and_(GpuJob.status == "leased", GpuJob.lease_expires_at < now),
                batch_size - len(picked),
            )
        if not picked:
            return []
        await self.session.execute(
            update(GpuJob)
            .where(GpuJob.id.in_(picked))
            .values(
                status="leased", lease_token=token, leased_at=now,
                lease_expires_at=now + timedelta(seconds=ttl),
                attempts=GpuJob.attempts + 1, updated_at=now,
            )
        )
        # populate_existing: overwrite identity-map copies with the post-UPDATE
        # values so the returned jobs reflect the new lease/attempts, not stale
        # pre-lease state.
        return list(
            (
                await self.session.execute(
                    select(GpuJob)
                    .where(GpuJob.id.in_(picked))
                    .order_by(GpuJob.id)
                    .execution_options(populate_existing=True)
                )
            ).scalars()
        )

    async def heartbeat(
        self, token: str, job_ids: list[int], ttl: int = DEFAULT_LEASE_TTL
    ) -> int:
        """Extend the lease on the agent's in-flight jobs. Returns rows touched."""
        now = datetime.now(UTC)
        res = await self.session.execute(
            update(GpuJob)
            .where(
                GpuJob.id.in_(job_ids),
                GpuJob.lease_token == token,
                GpuJob.status == "leased",
            )
            .values(lease_expires_at=now + timedelta(seconds=ttl), updated_at=now)
        )
        return res.rowcount or 0

    async def complete(self, token: str, job_id: int) -> bool:
        """Close a leased job (after its results were stored). False if the job
        isn't leased by this token (a stale/expired submit)."""
        job = await self.session.get(GpuJob, job_id)
        if job is None or job.status != "leased" or job.lease_token != token:
            return False
        job.status = "done"
        job.lease_token = None
        job.lease_expires_at = None
        job.error = None
        job.updated_at = datetime.now(UTC)
        return True

    async def fail(self, token: str, job_id: int, error: str) -> bool:
        """Report a failure: re-queue (pending) until MAX_ATTEMPTS, then 'error'."""
        job = await self.session.get(GpuJob, job_id)
        if job is None or job.lease_token != token:
            return False
        if job.attempts >= MAX_ATTEMPTS:
            job.status = "error"
        else:
            job.status = "pending"
            job.lease_token = None
            job.lease_expires_at = None
        job.error = (error or "")[:1000]
        job.updated_at = datetime.now(UTC)
        return True

    async def release(self, token: str, job_ids: list[int]) -> int:
        """Hand the agent's still-leased jobs back to pending NOW (graceful stop),
        so another worker picks them up immediately instead of waiting out the
        lease. Scoped to the token's own leases. Returns rows released."""
        if not job_ids:
            return 0
        now = datetime.now(UTC)
        res = await self.session.execute(
            update(GpuJob)
            .where(
                GpuJob.id.in_(job_ids),
                GpuJob.lease_token == token,
                GpuJob.status == "leased",
            )
            .values(
                status="pending", lease_token=None, leased_at=None,
                lease_expires_at=None, updated_at=now,
            )
        )
        return res.rowcount or 0

    async def recover_orphaned(self) -> int:
        """Reset every expired lease back to pending — catches agents that died
        mid-job (no graceful release) — and convert poison-loopers to 'error'
        (see the *_POISON_CAP rationale above). Run on a short beat so the queue
        recovers + reads honestly even when no worker is actively leasing.
        Returns rows recovered to pending (poison conversions are extra)."""
        counts = {}
        for name, stmt in recover_statements(datetime.now(UTC)).items():
            counts[name] = (await self.session.execute(stmt)).rowcount or 0
        return counts["recovered"]