FabledCurator/tests/test_api_gpu.py

"""GPU-job HTTP API (#114): bearer auth + lease/submit round-trip + backfill."""
import pytest

from backend.app.models import 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.ml 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