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_, select, update
from sqlalchemy.ext.asyncio import AsyncSession

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


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). Run on a short beat so the queue recovers
        + reads honestly even when no worker is actively leasing. Returns rows
        recovered."""
        now = datetime.now(UTC)
        res = await self.session.execute(
            update(GpuJob)
            .where(GpuJob.status == "leased", GpuJob.lease_expires_at < now)
            .values(
                status="pending", lease_token=None, leased_at=None,
                lease_expires_at=None, updated_at=now,
            )
        )
        return res.rowcount or 0