686808d3f3
After an agent-side fix (e.g. the short-video sampler), the errored jobs
(~2.8k) have exhausted their 3 attempts and stay parked: backfill skips
images that already have a job, and /reprocess is the nuclear option (it
resets the 179k DONE jobs too). There was no way to re-run just the errors.
POST /api/gpu/retry_errors resets every status='error' job (all task types)
to pending with attempts=0 and the stored error cleared — a small inline
UPDATE that returns {requeued: n} so the UI toast can show the count.
UI: a "Retry errored jobs" button on the GPU-agent card, right under the
queue tiles; disabled when errored==0. With the agent now logging ffmpeg's
stderr on failure, retrying also reveals which errors were real vs victims
of the fps-filter bug.
Test: retry_errors requeues the errored job (fresh attempts, error cleared)
and leaves done work untouched; asserts via column selects (Core-DML
gotcha), not ORM refresh.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
188 lines
6.9 KiB
Python
188 lines
6.9 KiB
Python
"""GPU-job HTTP API (#114): bearer auth + lease/submit round-trip + backfill."""
|
|
import pytest
|
|
from sqlalchemy import 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.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
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retry_errors_requeues_only_errored(client, db):
|
|
"""/retry_errors resets ERRORED jobs (any task) to pending with a fresh
|
|
retry budget — and leaves done work untouched (it is NOT /reprocess)."""
|
|
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_done.status = "done"
|
|
await db.commit()
|
|
|
|
resp = await client.post("/api/gpu/retry_errors")
|
|
assert resp.status_code == 200
|
|
assert (await resp.get_json())["requeued"] == 1
|
|
|
|
# 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"
|
|
|
|
st = await (await client.get("/api/gpu/status")).get_json()
|
|
assert st["pending"] == 1 and st["error"] == 0
|