FabledCurator/tests/test_gpu_jobs.py

"""GPU-job queue engine (#114): enqueue dedupe + lease/heartbeat/complete/fail."""
from datetime import UTC, datetime, timedelta

import pytest
from sqlalchemy import func, select

from backend.app.models import GpuJob, ImageRecord, ImageRegion
from backend.app.services.ml.gpu_jobs import (
    EXPIRED_POISON_CAP,
    PENDING_POISON_CAP,
    GpuJobService,
)

pytestmark = pytest.mark.integration


async def _img(db, sha) -> ImageRecord:
    img = ImageRecord(
        path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg",
        width=1, height=1, origin="imported_filesystem", integrity_status="unknown",
    )
    db.add(img)
    await db.flush()
    return img


@pytest.mark.asyncio
async def test_enqueue_siglip_backfill_gates_on_concept_region(db):
    # 'siglip' backfill enqueues images that lack a concept region (the
    # back-catalogue) and skips ones that already have one — and never double-
    # enqueues an image that already has a pending siglip job.
    from backend.app.models import MLSettings
    from backend.app.tasks.gpu_queue import enqueue_gpu_backfill

    cur = (await db.execute(
        select(MLSettings.embedder_model_version).where(MLSettings.id == 1)
    )).scalar_one()
    need = await _img(db, "e1" * 32)          # no concept region → wants one
    have = await _img(db, "e2" * 32)          # concept @ current version → skip
    db.add(ImageRegion(
        image_record_id=have.id, kind="concept", rx=0.0, ry=0.0, rw=1.0, rh=1.0,
        siglip_embedding=[0.0] * 1152, embedding_version=cur,
    ))
    await db.commit()

    assert enqueue_gpu_backfill("siglip") >= 1

    queued = {
        j.image_record_id for j in (
            await db.execute(select(GpuJob).where(GpuJob.task == "siglip"))
        ).scalars()
    }
    assert need.id in queued
    assert have.id not in queued

    # Idempotent: the now-pending job means a second run doesn't re-enqueue it.
    enqueue_gpu_backfill("siglip")
    n_for_need = (
        await db.execute(
            select(func.count()).select_from(GpuJob).where(
                GpuJob.task == "siglip", GpuJob.image_record_id == need.id
            )
        )
    ).scalar_one()
    assert n_for_need == 1


@pytest.mark.asyncio
async def test_enqueue_embed_backfill_selects_stale_and_unembedded(db):
    # Whole-image GPU re-embed (#1190): enqueue images with no embedding or one
    # stamped under a DIFFERENT model version (an operator swap); skip ones
    # already at the current version.
    from backend.app.models import MLSettings
    from backend.app.tasks.gpu_queue import enqueue_gpu_backfill

    cur = (await db.execute(
        select(MLSettings.embedder_model_version).where(MLSettings.id == 1)
    )).scalar_one()
    current = await _img(db, "f1" * 32)
    current.siglip_embedding = [0.0] * 1152
    current.siglip_model_version = cur                 # up to date → skip
    stale = await _img(db, "f2" * 32)
    stale.siglip_embedding = [0.0] * 1152
    stale.siglip_model_version = "old-embedder-v0"      # wrong space → re-embed
    unembedded = await _img(db, "f3" * 32)             # never embedded → embed
    await db.commit()

    assert enqueue_gpu_backfill("embed") >= 2
    queued = {
        j.image_record_id for j in (
            await db.execute(select(GpuJob).where(GpuJob.task == "embed"))
        ).scalars()
    }
    assert stale.id in queued and unembedded.id in queued
    assert current.id not in queued


@pytest.mark.asyncio
async def test_reprocess_resets_done_jobs_to_pending(db):
    # Re-process (#1202): done/error jobs of a task go back to pending so the
    # agent re-runs the whole library under the current pipeline.
    from backend.app.tasks.gpu_queue import reprocess_gpu_jobs

    img = await _img(db, "r1" * 32)
    job = await GpuJobService(db).enqueue(img.id, "ccip")
    job.status = "done"
    job.attempts = 2
    await db.commit()

    assert reprocess_gpu_jobs("ccip") >= 1
    await db.refresh(job)
    assert job.status == "pending"
    assert job.attempts == 0


@pytest.mark.asyncio
async def test_enqueue_dedupes_same_pair(db):
    img = await _img(db, "a" * 64)
    svc = GpuJobService(db)
    first = await svc.enqueue(img.id, "ccip")
    dup = await svc.enqueue(img.id, "ccip")
    other = await svc.enqueue(img.id, "siglip_region")
    await db.commit()
    assert first is not None
    assert dup is None              # same (image, task) already queued
    assert other is not None        # different task is fine


@pytest.mark.asyncio
async def test_lease_claims_then_skips_when_held(db):
    img = await _img(db, "b" * 64)
    svc = GpuJobService(db)
    await svc.enqueue(img.id, "ccip")
    await db.commit()

    leased = await svc.lease("agent-1", batch_size=8)
    await db.commit()
    assert len(leased) == 1
    assert leased[0].status == "leased" and leased[0].lease_token == "agent-1"
    assert leased[0].attempts == 1

    # Already leased + not expired → a second agent gets nothing.
    again = await svc.lease("agent-2", batch_size=8)
    await db.commit()
    assert again == []


@pytest.mark.asyncio
async def test_lease_skips_done_prefix(db):
    # A long prefix of already-done jobs at the LOW ids must not stop the lease
    # from reaching the pending ones (the O(done) scan bug this indexing fixes).
    svc = GpuJobService(db)
    for i in range(5):
        j = await svc.enqueue((await _img(db, f"d{i}" * 32)).id, "ccip")
        j.status = "done"
    target = await svc.enqueue((await _img(db, "dz" * 32)).id, "ccip")  # highest id
    await db.commit()

    leased = await svc.lease("agent-1", batch_size=8)
    await db.commit()
    assert [j.id for j in leased] == [target.id]


@pytest.mark.asyncio
async def test_expired_lease_is_reclaimed(db):
    img = await _img(db, "c" * 64)
    svc = GpuJobService(db)
    job = await svc.enqueue(img.id, "ccip")
    await db.commit()
    # Force the lease into the past.
    job.status = "leased"
    job.lease_token = "dead-agent"
    job.lease_expires_at = datetime.now(UTC) - timedelta(minutes=10)
    await db.commit()

    leased = await svc.lease("agent-2", batch_size=8)
    await db.commit()
    assert len(leased) == 1
    assert leased[0].lease_token == "agent-2"
    assert leased[0].attempts == 1   # re-lease incremented from 0 (was set directly)


@pytest.mark.asyncio
async def test_heartbeat_extends_only_own_lease(db):
    img = await _img(db, "d" * 64)
    svc = GpuJobService(db)
    await svc.enqueue(img.id, "ccip")
    await db.commit()
    job = (await svc.lease("agent-1"))[0]
    await db.commit()

    assert await svc.heartbeat("agent-1", [job.id]) == 1
    assert await svc.heartbeat("someone-else", [job.id]) == 0


@pytest.mark.asyncio
async def test_complete_closes_job(db):
    img = await _img(db, "e" * 64)
    svc = GpuJobService(db)
    await svc.enqueue(img.id, "ccip")
    await db.commit()
    job = (await svc.lease("agent-1"))[0]
    await db.commit()

    assert await svc.complete("wrong-token", job.id) is False
    assert await svc.complete("agent-1", job.id) is True
    await db.commit()
    fresh = await db.get(GpuJob, job.id)
    assert fresh.status == "done" and fresh.lease_token is None


@pytest.mark.asyncio
async def test_fail_requeues_until_cap(db):
    img = await _img(db, "f" * 64)
    svc = GpuJobService(db)
    await svc.enqueue(img.id, "ccip")
    await db.commit()
    job = (await svc.lease("agent-1"))[0]   # attempts -> 1
    await db.commit()
    # Under the cap → back to pending for a retry.
    assert await svc.fail("agent-1", job.id, "boom") is True
    await db.commit()
    assert (await db.get(GpuJob, job.id)).status == "pending"

    # At the attempt cap → terminal 'error'.
    j = await db.get(GpuJob, job.id)
    j.attempts = 3
    j.status = "leased"
    j.lease_token = "agent-1"
    j.lease_expires_at = datetime.now(UTC) + timedelta(minutes=5)
    await db.commit()
    assert await svc.fail("agent-1", job.id, "boom again") is True
    await db.commit()
    assert (await db.get(GpuJob, job.id)).status == "error"


@pytest.mark.asyncio
async def test_release_returns_to_pending(db):
    img = await _img(db, "01" + "a" * 62)
    svc = GpuJobService(db)
    await svc.enqueue(img.id, "ccip")
    await db.commit()
    job = (await svc.lease("agent-1"))[0]
    await db.commit()

    assert await svc.release("other", [job.id]) == 0      # not this token's lease
    assert await svc.release("agent-1", [job.id]) == 1     # graceful hand-back
    await db.commit()
    fresh = await db.get(GpuJob, job.id)
    assert fresh.status == "pending" and fresh.lease_token is None


@pytest.mark.asyncio
async def test_recover_orphaned_resets_only_expired(db):
    img1 = await _img(db, "02" + "a" * 62)
    img2 = await _img(db, "03" + "a" * 62)
    svc = GpuJobService(db)
    await svc.enqueue(img1.id, "ccip")
    await svc.enqueue(img2.id, "ccip")
    await db.commit()
    expired, fresh = await svc.lease("dead", batch_size=2)
    # One lease is in the past (orphaned), the other still valid.
    expired.lease_expires_at = datetime.now(UTC) - timedelta(minutes=10)
    await db.commit()

    assert await svc.recover_orphaned() == 1
    await db.commit()
    assert (await db.get(GpuJob, expired.id)).status == "pending"
    assert (await db.get(GpuJob, fresh.id)).status == "leased"   # untouched


@pytest.mark.asyncio
async def test_backfill_skips_errored_images(db):
    # An errored job is a TOMBSTONE for its (image, task): no backfill variant
    # re-enqueues it — retry is deliberate-only (/retry_errors). Pre-fix, the
    # hourly ccip run minted a fresh doomed job per bad file forever.
    from backend.app.tasks.gpu_queue import enqueue_gpu_backfill

    img = await _img(db, "f1" * 32)
    svc = GpuJobService(db)
    for task in ("ccip", "siglip", "embed"):
        job = await svc.enqueue(img.id, task)
        job.status = "error"
        job.error = "no frames sampled from video"
    await db.commit()

    assert enqueue_gpu_backfill("ccip") == 0
    assert enqueue_gpu_backfill("siglip") == 0
    assert enqueue_gpu_backfill("embed") == 0


@pytest.mark.asyncio
async def test_backfill_prunes_moot_error_tombstones(db):
    # Loop-era duplicates: several error rows for one (image, task), all made
    # moot by a later done row. The backfill's dedupe pass removes them, and
    # the done row still blocks re-enqueue.
    from backend.app.tasks.gpu_queue import enqueue_gpu_backfill

    img = await _img(db, "f2" * 32)
    for i in range(3):
        db.add(GpuJob(
            image_record_id=img.id, task="ccip", status="error",
            error=f"boom {i}",
        ))
    db.add(GpuJob(image_record_id=img.id, task="ccip", status="done"))
    await db.commit()

    assert enqueue_gpu_backfill("ccip") == 0

    statuses = (await db.execute(
        select(GpuJob.status).where(
            GpuJob.image_record_id == img.id, GpuJob.task == "ccip"
        )
    )).scalars().all()
    assert statuses == ["done"]


@pytest.mark.asyncio
async def test_recover_poisons_runaway_jobs(db):
    # release/expiry loops never reach fail()'s attempt cap — the sweep converts
    # them to 'error': an expired lease after EXPIRED_POISON_CAP grants (job
    # crashes/wedges the agent every time) and a pending job after
    # PENDING_POISON_CAP grants that never completed (transfer stalls forever).
    img1 = await _img(db, "06" + "a" * 62)
    img2 = await _img(db, "07" + "a" * 62)
    svc = GpuJobService(db)
    j1 = await svc.enqueue(img1.id, "ccip")
    j2 = await svc.enqueue(img2.id, "ccip")
    j1.status = "leased"
    j1.lease_token = "dead"
    j1.lease_expires_at = datetime.now(UTC) - timedelta(minutes=10)
    j1.attempts = EXPIRED_POISON_CAP
    j2.attempts = PENDING_POISON_CAP
    j1_id = j1.id
    j2_id = j2.id
    await db.commit()

    assert await svc.recover_orphaned() == 0   # nothing recovered — both poisoned
    await db.commit()

    # Column selects, not ORM refresh — the sweep wrote via Core DML.
    rows = (await db.execute(
        select(GpuJob.id, GpuJob.status, GpuJob.error)
        .where(GpuJob.id.in_([j1_id, j2_id]))
    )).all()
    by_id = {r.id: r for r in rows}
    assert by_id[j1_id].status == "error"
    assert "poisoned" in by_id[j1_id].error
    assert by_id[j2_id].status == "error"
    assert "poisoned" in by_id[j2_id].error