bvandeusen
2cb0427868
So work an agent orphaned gets picked back up quickly, three layers: - GpuJobService.release(): a graceful agent stop hands its still-leased jobs back to pending instantly (POST /api/gpu/jobs/release), no waiting out the lease. - GpuJobService.recover_orphaned() + recover_orphaned_gpu_jobs Celery task on a 60s beat: resets expired leases (a hard-crashed agent) to pending and keeps the queue counts honest even when nothing is leasing. - Lease TTL 300→180s: still well above any single job (a capped-frame video embed is tens of seconds, and a live worker heartbeats), but a hard crash recovers faster once the sweep fires. Tests: release returns-to-pending (token-scoped), recover_orphaned resets only expired leases, release API round-trip. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
118 lines
4.3 KiB
Python
118 lines
4.3 KiB
Python
"""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
|
|
|
|
|
|
@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
|