FabledCurator/tests/test_api_gpu.py

"""GPU-job HTTP API (#114): bearer auth + lease/submit round-trip + backfill."""
from datetime import UTC, datetime, timedelta

import pytest
from sqlalchemy import func, select

from backend.app.models import GpuJob, ImageRecord
from backend.app.services.ml.gpu_jobs import GpuJobService
from backend.app.services.ml.regions import RegionService

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_agent_endpoints_require_bearer(client, db):
    resp = await client.post("/api/gpu/jobs/lease", json={"agent_id": "a1"})
    assert resp.status_code == 401
    # A wrong token is also rejected.
    await (await client.post("/api/gpu/token/rotate")).get_json()
    bad = await client.post(
        "/api/gpu/jobs/lease", json={"agent_id": "a1"},
        headers={"Authorization": "Bearer nope"},
    )
    assert bad.status_code == 401


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

    token = (await (await client.post("/api/gpu/token/rotate")).get_json())["token"]
    hdr = {"Authorization": f"Bearer {token}"}

    leased = await client.post(
        "/api/gpu/jobs/lease", json={"agent_id": "a1", "batch_size": 5}, headers=hdr,
    )
    assert leased.status_code == 200
    jobs = (await leased.get_json())["jobs"]
    assert len(jobs) == 1
    j = jobs[0]
    assert j["image_id"] == img.id and j["task"] == "ccip"
    assert j["image_url"].startswith("/images/")

    submitted = await client.post("/api/gpu/jobs/submit", json={
        "agent_id": "a1", "job_id": j["job_id"],
        "regions": [{
            "kind": "figure", "bbox": [0.1, 0.1, 0.4, 0.4],
            "ccip_embedding": [0.1] * 768, "embedding_version": "ccip-test",
        }],
    }, headers=hdr)
    assert submitted.status_code == 200
    assert (await submitted.get_json())["stored"] == 1

    # Job closed (read on the app's own connection via the status endpoint).
    st = await (await client.get("/api/gpu/status")).get_json()
    assert st["done"] == 1 and st["pending"] == 0 and st["leased"] == 0

    # Region persisted with its CCIP vector.
    regs = await RegionService(db).get_regions(img.id, kinds=["figure"])
    assert len(regs) == 1 and len(list(regs[0].ccip_embedding)) == 768


@pytest.mark.asyncio
async def test_lease_announces_embed_model_then_submit_embedding(client, db):
    # Whole-image GPU re-embed (#1190): the lease announces the embedder model so
    # the agent loads the right one, and submit_embedding writes it back onto
    # image_record with its version stamp.
    img = await _img(db, "b" * 64)
    await GpuJobService(db).enqueue(img.id, "embed")
    await db.commit()

    token = (await (await client.post("/api/gpu/token/rotate")).get_json())["token"]
    hdr = {"Authorization": f"Bearer {token}"}

    leased = await client.post(
        "/api/gpu/jobs/lease", json={"agent_id": "a1", "batch_size": 5}, headers=hdr,
    )
    j = (await leased.get_json())["jobs"][0]
    assert j["task"] == "embed"
    assert j["embed_model_name"] and j["embed_version"]   # server-announced model

    submitted = await client.post("/api/gpu/jobs/submit_embedding", json={
        "agent_id": "a1", "job_id": j["job_id"],
        "embedding": [0.2] * 1152, "embedding_version": "siglip2-test-v9",
    }, headers=hdr)
    assert submitted.status_code == 200

    st = await (await client.get("/api/gpu/status")).get_json()
    assert st["done"] == 1 and st["leased"] == 0

    await db.refresh(img)
    assert img.siglip_model_version == "siglip2-test-v9"
    assert img.siglip_embedding is not None and len(list(img.siglip_embedding)) == 1152


@pytest.mark.asyncio
async def test_submit_with_stale_lease_is_409(client, db):
    img = await _img(db, "b" * 64)
    await GpuJobService(db).enqueue(img.id, "ccip")
    await db.commit()
    token = (await (await client.post("/api/gpu/token/rotate")).get_json())["token"]
    hdr = {"Authorization": f"Bearer {token}"}
    j = (await (await client.post(
        "/api/gpu/jobs/lease", json={"agent_id": "a1"}, headers=hdr,
    )).get_json())["jobs"][0]
    # A different agent can't submit someone else's lease.
    resp = await client.post("/api/gpu/jobs/submit", json={
        "agent_id": "other", "job_id": j["job_id"], "regions": [],
    }, headers=hdr)
    assert resp.status_code == 409


@pytest.mark.asyncio
async def test_backfill_enqueues_then_is_idempotent(db):
    await _img(db, "c" * 64)
    await _img(db, "d" * 64)
    await db.commit()
    from backend.app.tasks.gpu_queue import enqueue_gpu_backfill

    n = enqueue_gpu_backfill("ccip")     # sync task, own session
    assert n >= 2
    assert enqueue_gpu_backfill("ccip") == 0   # all already pending


@pytest.mark.asyncio
async def test_release_hands_job_back_to_pending(client, db):
    img = await _img(db, "e" * 64)
    await GpuJobService(db).enqueue(img.id, "ccip")
    await db.commit()
    token = (await (await client.post("/api/gpu/token/rotate")).get_json())["token"]
    hdr = {"Authorization": f"Bearer {token}"}
    j = (await (await client.post(
        "/api/gpu/jobs/lease", json={"agent_id": "a1"}, headers=hdr,
    )).get_json())["jobs"][0]

    resp = await client.post("/api/gpu/jobs/release", json={
        "agent_id": "a1", "job_ids": [j["job_id"]],
    }, headers=hdr)
    assert resp.status_code == 200 and (await resp.get_json())["released"] == 1
    st = await (await client.get("/api/gpu/status")).get_json()
    assert st["pending"] == 1 and st["leased"] == 0


@pytest.mark.asyncio
async def test_retry_errors_requeues_only_errored(client, db):
    """/retry_errors prunes stale tombstones first (older duplicates + rows a
    later success made moot), then resets the SURVIVING errored jobs to pending
    with a fresh retry budget — and leaves done work untouched (NOT /reprocess).
    The prune is what stops one failing file fanning out into duplicate pending
    jobs (the 2026-07-02 tombstone loop minted one error row per hour)."""
    img1 = await _img(db, "1" * 64)
    img2 = await _img(db, "2" * 64)
    svc = GpuJobService(db)
    j_err = await svc.enqueue(img1.id, "ccip")
    j_done = await svc.enqueue(img2.id, "siglip")
    err_id = j_err.id
    done_id = j_done.id
    j_err.status = "error"
    j_err.attempts = 3
    j_err.error = "no frames sampled from video (unprocessable)"
    j_err.updated_at = datetime.now(UTC)
    j_done.status = "done"
    # Loop-era leftovers: an OLDER duplicate error row for img1's ccip, and a
    # tombstone img2's done row makes moot — both pruned, never requeued.
    dup = GpuJob(
        image_record_id=img1.id, task="ccip", status="error",
        error="older duplicate", updated_at=datetime.now(UTC) - timedelta(hours=1),
    )
    moot = GpuJob(
        image_record_id=img2.id, task="siglip", status="error",
        error="superseded by the done row",
    )
    db.add(dup)
    db.add(moot)
    await db.flush()
    dup_id = dup.id
    moot_id = moot.id
    await db.commit()

    resp = await client.post("/api/gpu/retry_errors")
    assert resp.status_code == 200
    body = await resp.get_json()
    assert body["requeued"] == 1
    assert body["pruned"] == 2

    # Column selects, not ORM refresh — the route wrote via Core DML.
    row = (await db.execute(
        select(GpuJob.status, GpuJob.attempts, GpuJob.error)
        .where(GpuJob.id == err_id)
    )).one()
    assert tuple(row) == ("pending", 0, None)
    done_status = await db.scalar(
        select(GpuJob.status).where(GpuJob.id == done_id)
    )
    assert done_status == "done"
    survivors = (await db.execute(
        select(func.count()).select_from(GpuJob)
        .where(GpuJob.id.in_([dup_id, moot_id]))
    )).scalar_one()
    assert survivors == 0

    st = await (await client.get("/api/gpu/status")).get_json()
    assert st["pending"] == 1 and st["error"] == 0


@pytest.mark.asyncio
async def test_retry_errors_keeps_triaged_defects(client, db):
    """A probe-confirmed DEFECT is a bad FILE — requeueing it just burns agent
    time re-minting the tombstone, so /retry_errors leaves it for the recovery
    surface and reports it as defects_kept."""
    img1 = await _img(db, "4" * 64)
    img2 = await _img(db, "5" * 64)
    db.add(GpuJob(image_record_id=img1.id, task="ccip", status="error",
                  attempts=3, error="moov atom not found",
                  triage_status="defect"))
    db.add(GpuJob(image_record_id=img2.id, task="ccip", status="error",
                  attempts=3, error="ffmpeg timed out after 1200s"))
    await db.commit()

    body = await (await client.post("/api/gpu/retry_errors")).get_json()
    assert body["requeued"] == 1
    assert body["defects_kept"] == 1

    rows = dict((await db.execute(
        select(GpuJob.image_record_id, GpuJob.status)
    )).all())
    assert rows[img1.id] == "error"       # defect stays tombstoned
    assert rows[img2.id] == "pending"     # operational failure requeued


@pytest.mark.asyncio
async def test_errors_endpoint_reports_triage_view(client, db):
    img = await _img(db, "6" * 64)
    db.add(GpuJob(image_record_id=img.id, task="ccip", status="error",
                  attempts=3,
                  error="no frames sampled from video — moov atom not found"))
    await db.commit()

    resp = await client.get("/api/gpu/errors")
    assert resp.status_code == 200
    body = await resp.get_json()
    assert body["total"] == 1
    assert body["by_class"] == {"truncated_or_corrupt": 1}
    assert body["triage"]["unclassified"] == 1
    item = body["items"][0]
    assert item["image_id"] == img.id
    assert item["task"] == "ccip"
    assert item["reason_class"] == "truncated_or_corrupt"
    assert item["triage_status"] is None
    assert item["image_url"].startswith("/images/")


@pytest.mark.asyncio
async def test_cpu_embed_never_blocks_gpu_crop_backfills(db):
    """B3 invariant (operator 2026-07-02): ccip (detect + character) and
    siglip (concept crops) completion is judged per-pipeline — gpu_job rows and
    image_region state — never inferred from image_record.siglip_embedding. So
    an image the CPU fallback already embedded still gets both crop jobs; only
    the whole-image 'embed' job (the SAME artifact the CPU path produces) is
    satisfied by it."""
    from backend.app.models import MLSettings
    from backend.app.tasks.gpu_queue import enqueue_gpu_backfill

    img = await _img(db, "7" * 64)
    cur = (await db.execute(
        select(MLSettings.embedder_model_version).where(MLSettings.id == 1)
    )).scalar_one()
    # As if the CPU fallback already embedded it under the current model.
    img.siglip_embedding = [0.1] * 1152
    img.siglip_model_version = cur
    await db.commit()

    assert enqueue_gpu_backfill("ccip") == 1      # crops still open
    assert enqueue_gpu_backfill("siglip") == 1    # concept crops still open
    assert enqueue_gpu_backfill("embed") == 0     # same artifact — already done

    tasks = set((await db.execute(
        select(GpuJob.task).where(GpuJob.image_record_id == img.id)
    )).scalars().all())
    assert tasks == {"ccip", "siglip"}