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

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

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 (or a
retry after an agent died) never grab the same job and the queue self-heals
without a separate recovery sweep. Result-writing (regions) is done by the API
handler via RegionService; complete() just closes the job.
"""

from datetime import UTC, datetime, timedelta

from sqlalchemy import and_, or_, select, update
from sqlalchemy.ext.asyncio import AsyncSession

from ...models import GpuJob

DEFAULT_LEASE_TTL = 300     # 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`."""
        now = datetime.now(UTC)
        picked = (
            await self.session.execute(
                select(GpuJob.id)
                .where(
                    or_(
                        GpuJob.status == "pending",
                        and_(
                            GpuJob.status == "leased",
                            GpuJob.lease_expires_at < now,
                        ),
                    )
                )
                .order_by(GpuJob.id)
                .limit(batch_size)
                .with_for_update(skip_locked=True)
            )
        ).scalars().all()
        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