Files
FabledCurator/backend/app/models/gpu_job.py
T
bvandeusen 181f1c6a27
CI / lint (push) Successful in 3s
CI / frontend-build (push) Successful in 20s
CI / backend-lint-and-test (push) Successful in 26s
CI / integration (push) Successful in 3m25s
perf(gpu-queue): partial indexes + two-phase lease so leasing stays O(batch)
The throughput bottleneck was curator-side, not the network. lease() claimed the
lowest-id pending/expired jobs with `... ORDER BY id LIMIT n`, but with only a
plain `status` index Postgres walked the primary key from id=1, skipping the
entire prefix of already done/error rows before reaching pending ones. As `done`
grew (69k+), every lease became an O(done) scan — leasing crawled, the DB
saturated, and even /status (the queue GROUP BY count) stalled the agent.

- Migration 0070 adds two partial indexes over just the live slice: pending rows
  indexed by id (hot path), and leased rows by lease_expires_at (crash-recovery
  + orphan sweep). They stay tiny no matter how large the done/error history.
- lease() split into two phases so each uses a partial index: claim pending
  first (id-ordered, O(batch)); reclaim expired leases only when pending can't
  fill the batch. Same semantics (SKIP LOCKED, attempts++, expired reclaim).
- Model __table_args__ declares the indexes so ORM and schema agree.
- Test: a done-prefix at low ids must not stop the lease reaching pending.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
2026-06-30 21:12:12 -04:00

71 lines
2.7 KiB
Python

"""GpuJob — a unit of GPU work the desktop agent pulls over HTTP (#114).
The durable work list that lets the agent stay HTTP-only: the server enqueues a
job per (image, task) — e.g. detect figures + CCIP-embed — and the agent LEASES a
batch, computes on its GPU, then SUBMITS results, all over the already-exposed web
API. Redis/Postgres stay private. A lease has an expiry; the lease query itself
re-claims expired leases (agent died / stopped mid-batch), so the queue is
self-healing without a separate sweep. One job is per ITEM; the agent fans a
VIDEO out into per-frame instances internally (see image_region.frame_time).
State: pending → leased → done | error (a failure under the attempt cap returns to
pending for another agent).
"""
from datetime import datetime
from sqlalchemy import (
DateTime,
ForeignKey,
Index,
Integer,
String,
Text,
func,
text,
)
from sqlalchemy.orm import Mapped, mapped_column
from .base import Base
class GpuJob(Base):
__tablename__ = "gpu_job"
# Partial indexes over just the live slice (see migration 0070): the lease
# reads the lowest-id pending jobs on the hot path, and reclaims expired
# leases as a backstop — both stay O(batch) as done/error history grows.
__table_args__ = (
Index("ix_gpu_job_pending", "id", postgresql_where=text("status = 'pending'")),
Index(
"ix_gpu_job_leased_expires", "lease_expires_at",
postgresql_where=text("status = 'leased'"),
),
)
id: Mapped[int] = mapped_column(Integer, primary_key=True)
image_record_id: Mapped[int] = mapped_column(
ForeignKey("image_record.id", ondelete="CASCADE"), index=True
)
# What to compute, e.g. 'ccip' (detect figures + CCIP-embed) or 'siglip_region'.
task: Mapped[str] = mapped_column(String(32), nullable=False)
status: Mapped[str] = mapped_column(
String(16), nullable=False, default="pending", index=True
)
# pending | leased | done | error
lease_token: Mapped[str | None] = mapped_column(String(64), nullable=True)
leased_at: Mapped[datetime | None] = mapped_column(
DateTime(timezone=True), nullable=True
)
lease_expires_at: Mapped[datetime | None] = mapped_column(
DateTime(timezone=True), nullable=True
)
attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
error: Mapped[str | None] = mapped_column(Text, nullable=True)
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), nullable=False, server_default=func.now()
)
updated_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), nullable=False, server_default=func.now()
)