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_lease_announces_detector_config(client, db):
    # #134: the lease carries the crop-proposer config from MLSettings, so the
    # agent builds its detectors from the DB/UI (no restart). Defaults are all-on
    # with the pinned working weights.
    img = await _img(db, "c" * 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,
    )
    det = (await leased.get_json())["jobs"][0]["detectors"]
    assert det["person"]["enabled"]
    assert det["anatomy"]["enabled"]
    assert det["panel"]["enabled"]
    assert "yolov11m_aa22" in det["anatomy"]["weights"]      # booru_yolo default
    assert det["panel"]["weights"].endswith("::best.pt")     # mosesb default
    assert det["max_regions"] == 128
    assert det["dedupe_iou"] == 0.85


@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"}


# --- an idle agent still checks in -------------------------------------------


@pytest.mark.asyncio
async def test_a_heartbeat_with_no_jobs_still_records_the_check_in(client, db):
    """Operator, 2026-09-23: *"I'm running the gpu agent on my device and it
    currently reads as 'offline' but it's running and has checked in
    recently."*

    It had — twelve minutes ago. The roster takes its agent check-in from the
    `lease` and `heartbeat` calls, and calls an agent stopped after 300s of
    silence. The agent's heartbeat loop was gated on holding leases, so an
    IDLE agent sent none; the only check-in left was the lease poll, which
    sleep mode backs off to a 900s ceiling. 900 against 300 — an idle agent
    was structurally guaranteed to read as stopped.

    So the empty heartbeat has to be a real check-in on the server side, not
    merely tolerated. Asserted on `last_seen_at` moving, because "it returned
    200" would pass against an endpoint that recorded nothing.
    """
    from backend.app.models import ServiceSeen

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

    resp = await client.post(
        "/api/gpu/jobs/heartbeat",
        json={"agent_id": "desktop-agent", "job_ids": []}, headers=hdr,
    )

    assert resp.status_code == 200
    assert (await resp.get_json())["extended"] == 0, "it extends no lease"

    row = (await db.execute(
        select(ServiceSeen).where(ServiceSeen.key == "agent:desktop-agent")
    )).scalar_one()
    assert row.kind == "agent"
    assert row.display_name == "GPU agent (desktop-agent)"
    assert row.last_seen_at is not None


def test_the_agent_heartbeats_whether_or_not_it_holds_a_lease():
    """The other half, in the agent itself — the half that was actually wrong.

    Read from the source rather than by running the loop: it is a `while True`
    with a sleep, so exercising it means threads and timing, and the property
    is simply that the call is not behind a `if ids:`.
    """
    from pathlib import Path

    src = (
        Path(__file__).resolve().parents[1] / "agent" / "fc_agent" / "worker.py"
    ).read_text()
    loop = src[src.index("def _heartbeat_loop"):]
    loop = loop[:loop.index("\n    def ")]

    assert "self.client.heartbeat(ids)" in loop
    assert "if ids:" not in loop, (
        "the heartbeat is gated on holding leases again; an idle agent then "
        "reads as stopped after 300s while sleep mode backs its lease poll "
        "off to 900s"
    )