"""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"]