feat(gpu): HTTP job API + token auth + backfill — the agent's server side (#114 slice 3b)
The thin HTTP surface over the queue so the desktop agent stays HTTP-only: - Agent endpoints (Authorization: Bearer <token>): POST /api/gpu/jobs/lease (returns jobs + image_url + mime + video frame cadence), /submit (stores regions via RegionService + closes the job; 409 on a stale lease), /heartbeat, /fail. Token validated against AppSetting (mirrors the extension-key pattern, constant-time compare). - Admin (browser): GET/POST /api/gpu/token[/rotate] (generate + show the agent token), GET /api/gpu/status (queue counts), POST /api/gpu/backfill → dispatches enqueue_gpu_backfill. - enqueue_gpu_backfill(task): one INSERT…SELECT enqueues a job per image lacking one for the task (scales to the full library; idempotent). Agent flow: lease over HTTP → fetch pixels via the normal FC image URL → compute on the GPU → submit. Redis/Postgres never exposed. Tests: bearer required (+ wrong-token 401), lease→submit round-trip (region+CCIP vector stored, job done via /status), stale-lease 409, backfill enqueue + idempotency. NEXT: the agent container + control UI, then the CCIP detector/embedder + matcher. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
This commit is contained in:
@@ -0,0 +1,98 @@
|
||||
"""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_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
|
||||
Reference in New Issue
Block a user