bvandeusen
a4df279343
The GPU lease now carries the crop-proposer config from MLSettings in a per-job 'detectors' block (same pattern as embed_model_name). The agent's worker builds its Proposers from the announced config via _effective_cfg (lease block overlaid on env) + _proposers_for (rebuilds only when a config signature changes) — so an operator's UI edit takes effect on the next lease with NO restart, and env is now just the bootstrap fallback until the server announces. enabled-off maps to empty weights (proposer skipped); dedupe_iou + max_regions also come from the effective cfg. Test: lease announces the detectors block with the seeded default weights. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CDgx8bQS5YrGRK76v8HUnM
423 lines
17 KiB
Python
423 lines
17 KiB
Python
"""GPU-job API (#114): the HTTP surface the desktop agent pulls work from.
|
|
|
|
The agent stays HTTP-only — it leases jobs, fetches image pixels via the normal
|
|
FC image URLs, and submits embeddings/regions back, all over this API. Redis and
|
|
Postgres are never exposed. The agent endpoints are gated by a bearer token
|
|
(Authorization: Bearer <token>) stored in AppSetting; the admin endpoints
|
|
(token / backfill / status) ride the browser session like the rest of FC's
|
|
homelab admin.
|
|
"""
|
|
|
|
import secrets
|
|
from pathlib import Path
|
|
|
|
from quart import Blueprint, jsonify, request
|
|
from sqlalchemy import func, or_, select, update
|
|
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
|
|
|
from ..extensions import get_session
|
|
from ..models import AppSetting, GpuJob, ImageRecord, MLSettings
|
|
from ..services.gallery_service import image_url
|
|
from ..services.ml.gpu_jobs import GpuJobService, error_dedupe_statements
|
|
from ..services.ml.gpu_triage import classify_reason, recover_defective_image
|
|
from ..services.ml.regions import RegionService
|
|
|
|
gpu_bp = Blueprint("gpu", __name__, url_prefix="/api/gpu")
|
|
|
|
# Same container mount the maintenance tasks use (tasks/admin.py) — recovery
|
|
# deletes the defective original + thumbnail under it.
|
|
_IMAGES_ROOT = Path("/images")
|
|
|
|
_TOKEN_KEY = "gpu_agent_token"
|
|
|
|
|
|
def _bearer() -> str | None:
|
|
h = request.headers.get("Authorization", "")
|
|
return h[7:].strip() if h.startswith("Bearer ") else None
|
|
|
|
|
|
async def _agent_authed(session) -> bool:
|
|
supplied = _bearer()
|
|
if not supplied:
|
|
return False
|
|
stored = (
|
|
await session.execute(
|
|
select(AppSetting.value).where(AppSetting.key == _TOKEN_KEY)
|
|
)
|
|
).scalar_one_or_none()
|
|
return stored is not None and secrets.compare_digest(supplied, stored)
|
|
|
|
|
|
# --- Admin (browser): token + backfill + status -------------------------
|
|
|
|
@gpu_bp.route("/token", methods=["GET"])
|
|
async def get_token():
|
|
async with get_session() as session:
|
|
tok = (
|
|
await session.execute(
|
|
select(AppSetting.value).where(AppSetting.key == _TOKEN_KEY)
|
|
)
|
|
).scalar_one_or_none()
|
|
return jsonify({"token": tok, "configured": tok is not None})
|
|
|
|
|
|
@gpu_bp.route("/token/rotate", methods=["POST"])
|
|
async def rotate_token():
|
|
token = secrets.token_urlsafe(32)
|
|
async with get_session() as session:
|
|
await session.execute(
|
|
pg_insert(AppSetting)
|
|
.values(key=_TOKEN_KEY, value=token)
|
|
.on_conflict_do_update(index_elements=["key"], set_={"value": token})
|
|
)
|
|
await session.commit()
|
|
return jsonify({"token": token})
|
|
|
|
|
|
@gpu_bp.route("/status", methods=["GET"])
|
|
async def status():
|
|
async with get_session() as session:
|
|
rows = (
|
|
await session.execute(
|
|
select(GpuJob.status, func.count()).group_by(GpuJob.status)
|
|
)
|
|
).all()
|
|
counts = dict(rows)
|
|
return jsonify({
|
|
"pending": counts.get("pending", 0),
|
|
"leased": counts.get("leased", 0),
|
|
"done": counts.get("done", 0),
|
|
"error": counts.get("error", 0),
|
|
})
|
|
|
|
|
|
@gpu_bp.route("/backfill", methods=["POST"])
|
|
async def backfill():
|
|
"""Enqueue a job for every image that doesn't already have one for `task`."""
|
|
body = await request.get_json(silent=True) or {}
|
|
task = str(body.get("task") or "ccip")
|
|
from ..tasks.gpu_queue import enqueue_gpu_backfill
|
|
|
|
r = enqueue_gpu_backfill.delay(task)
|
|
return jsonify({"celery_task_id": r.id, "task": task}), 202
|
|
|
|
|
|
@gpu_bp.route("/reprocess", methods=["POST"])
|
|
async def reprocess():
|
|
"""Reset every done/error job of `task` back to pending so the agent re-runs
|
|
the WHOLE library under the current pipeline (e.g. after adding crop
|
|
detectors). Heavy — the back-catalogue is otherwise skipped by the backfills."""
|
|
body = await request.get_json(silent=True) or {}
|
|
task = str(body.get("task") or "ccip")
|
|
from ..tasks.gpu_queue import reprocess_gpu_jobs
|
|
|
|
r = reprocess_gpu_jobs.delay(task)
|
|
return jsonify({"celery_task_id": r.id, "task": task}), 202
|
|
|
|
|
|
@gpu_bp.route("/retry_errors", methods=["POST"])
|
|
async def retry_errors():
|
|
"""Requeue every ERRORED job (all task types) back to pending — the scoped
|
|
recovery after an agent-side fix (e.g. the short-video sampler), where
|
|
/reprocess would needlessly re-run the whole done library too. Attempts and
|
|
the stored error reset so each job gets its full retry budget under the
|
|
fixed pipeline. Stale tombstones are pruned FIRST (loop-era duplicates and
|
|
rows a later success made moot — the same statements the backfills run), so
|
|
one failing file requeues as ONE job, never a fan-out of duplicates. Small
|
|
row count (errors only) → inline statements; the response carries the
|
|
counts for the UI toast. Triage-confirmed defects are NOT requeued (see
|
|
the WHERE below) — they stay on the recovery surface."""
|
|
async with get_session() as session:
|
|
pruned = 0
|
|
for stmt in error_dedupe_statements():
|
|
pruned += (await session.execute(stmt)).rowcount or 0
|
|
res = await session.execute(
|
|
update(GpuJob)
|
|
.where(
|
|
GpuJob.status == "error",
|
|
# Triage-confirmed DEFECTS stay errored: the integrity probe
|
|
# already proved the FILE is bad, so re-running the job just
|
|
# burns agent time re-minting the same tombstone — those go
|
|
# through /errors/<id>/recover instead.
|
|
or_(GpuJob.triage_status.is_(None),
|
|
GpuJob.triage_status != "defect"),
|
|
)
|
|
.values(
|
|
status="pending", attempts=0, error=None, lease_token=None,
|
|
leased_at=None, lease_expires_at=None, triage_status=None,
|
|
updated_at=func.now(),
|
|
)
|
|
)
|
|
kept = (
|
|
await session.execute(
|
|
select(func.count()).select_from(GpuJob)
|
|
.where(GpuJob.status == "error")
|
|
)
|
|
).scalar_one()
|
|
await session.commit()
|
|
return jsonify({
|
|
"requeued": res.rowcount or 0, "pruned": pruned, "defects_kept": kept,
|
|
})
|
|
|
|
|
|
# --- Failure triage + recovery (#125) ------------------------------------
|
|
|
|
@gpu_bp.route("/errors", methods=["GET"])
|
|
async def errors():
|
|
"""The triage view of the error tombstones: every errored job joined with
|
|
its image's integrity verdict, bucketed by reason for the overview. The
|
|
probe sweep (triage_gpu_errors, 15-min beat) fills triage_status; 'defect'
|
|
rows are the recovery surface's list."""
|
|
async with get_session() as session:
|
|
rows = (
|
|
await session.execute(
|
|
select(
|
|
GpuJob.id, GpuJob.image_record_id, GpuJob.task,
|
|
GpuJob.error, GpuJob.triage_status, GpuJob.updated_at,
|
|
ImageRecord.integrity_status, ImageRecord.mime,
|
|
ImageRecord.path, ImageRecord.thumbnail_path,
|
|
)
|
|
.join(ImageRecord, ImageRecord.id == GpuJob.image_record_id)
|
|
.where(GpuJob.status == "error")
|
|
.order_by(GpuJob.updated_at.desc())
|
|
.limit(500)
|
|
)
|
|
).all()
|
|
total = (
|
|
await session.execute(
|
|
select(func.count()).select_from(GpuJob)
|
|
.where(GpuJob.status == "error")
|
|
)
|
|
).scalar_one()
|
|
by_class: dict[str, int] = {}
|
|
triage = {"defect": 0, "file_ok": 0, "unclassified": 0}
|
|
items = []
|
|
for r in rows:
|
|
cls = classify_reason(r.error)
|
|
by_class[cls] = by_class.get(cls, 0) + 1
|
|
bucket = r.triage_status or "unclassified"
|
|
triage[bucket] = triage.get(bucket, 0) + 1
|
|
items.append({
|
|
"job_id": r.id,
|
|
"image_id": r.image_record_id,
|
|
"task": r.task,
|
|
"error": r.error,
|
|
"reason_class": cls,
|
|
"triage_status": r.triage_status,
|
|
"integrity_status": r.integrity_status,
|
|
"mime": r.mime,
|
|
"image_url": image_url(r.path),
|
|
"thumbnail_url": (
|
|
image_url(r.thumbnail_path) if r.thumbnail_path else None
|
|
),
|
|
"updated_at": r.updated_at.isoformat() if r.updated_at else None,
|
|
})
|
|
return jsonify({
|
|
"total": total, "by_class": by_class, "triage": triage, "items": items,
|
|
})
|
|
|
|
|
|
@gpu_bp.route("/errors/triage", methods=["POST"])
|
|
async def errors_triage():
|
|
"""Run the probe sweep NOW (the card's button) instead of waiting out the
|
|
15-minute beat cadence."""
|
|
from ..tasks.maintenance import triage_gpu_errors
|
|
|
|
r = triage_gpu_errors.delay()
|
|
return jsonify({"celery_task_id": r.id}), 202
|
|
|
|
|
|
@gpu_bp.route("/errors/<int:image_id>/recover", methods=["POST"])
|
|
async def errors_recover(image_id: int):
|
|
"""Recover a defect-triaged original: delete the bad copy + record and
|
|
re-poll its subscription Source (a fresh fetch re-imports the file, which
|
|
re-enters the GPU pipeline). Returns status 'no_source' when nothing
|
|
pollable resolves — the file needs manual replacement there."""
|
|
async with get_session() as session:
|
|
result = await session.run_sync(
|
|
lambda s: recover_defective_image(
|
|
s, image_id, images_root=_IMAGES_ROOT,
|
|
)
|
|
)
|
|
return jsonify(result)
|
|
|
|
|
|
# --- Agent (bearer token): lease / submit / heartbeat / fail ------------
|
|
|
|
@gpu_bp.route("/jobs/lease", methods=["POST"])
|
|
async def lease():
|
|
body = await request.get_json(silent=True) or {}
|
|
agent_id = str(body.get("agent_id") or "agent")
|
|
try:
|
|
batch = min(max(int(body.get("batch_size", 8)), 1), 64)
|
|
except (TypeError, ValueError):
|
|
batch = 8
|
|
async with get_session() as session:
|
|
if not await _agent_authed(session):
|
|
return jsonify({"error": "unauthorized"}), 401
|
|
jobs = await GpuJobService(session).lease(agent_id, batch_size=batch)
|
|
ml = (
|
|
await session.execute(select(MLSettings).where(MLSettings.id == 1))
|
|
).scalar_one()
|
|
# image rows for url/mime in one shot
|
|
ids = [j.image_record_id for j in jobs]
|
|
imgs = {
|
|
i.id: i for i in (
|
|
await session.execute(
|
|
select(ImageRecord).where(ImageRecord.id.in_(ids))
|
|
)
|
|
).scalars()
|
|
} if ids else {}
|
|
await session.commit()
|
|
# Crop-proposer config, announced FROM THE SETTING like embed_model_name
|
|
# (#134): the agent builds its detectors from this, rebuilding live when
|
|
# it changes — so tuning is a DB/UI edit, never an agent restart. Same
|
|
# block for every job in the batch (it's global), built once. An enabled
|
|
# toggle off is carried through so the agent skips that proposer.
|
|
detectors = {
|
|
"person": {
|
|
"enabled": ml.detector_person_enabled,
|
|
"weights": ml.detector_person_weights,
|
|
"conf": ml.detector_person_conf,
|
|
},
|
|
"anatomy": {
|
|
"enabled": ml.detector_anatomy_enabled,
|
|
"weights": ml.detector_anatomy_weights,
|
|
"conf": ml.detector_anatomy_conf,
|
|
},
|
|
"panel": {
|
|
"enabled": ml.detector_panel_enabled,
|
|
"weights": ml.detector_panel_weights,
|
|
"conf": ml.detector_panel_conf,
|
|
},
|
|
"max_figures": ml.detector_max_figures,
|
|
"max_components": ml.detector_max_components,
|
|
"max_panels": ml.detector_max_panels,
|
|
"max_regions": ml.detector_max_regions,
|
|
"dedupe_iou": ml.detector_dedupe_iou,
|
|
}
|
|
out = []
|
|
for j in jobs:
|
|
img = imgs.get(j.image_record_id)
|
|
if img is None:
|
|
continue
|
|
out.append({
|
|
"job_id": j.id,
|
|
"image_id": j.image_record_id,
|
|
"task": j.task,
|
|
"mime": img.mime,
|
|
"image_url": image_url(img.path),
|
|
# For video/animated: the agent samples at this cadence.
|
|
"frame_interval_seconds": ml.video_frame_interval_seconds,
|
|
"max_frames": ml.video_max_frames,
|
|
# The embedding model the agent must use for concept crops + the
|
|
# whole-image 'embed' task, so its vectors land in the SAME space
|
|
# the heads trained in. Server-announced FROM THE SETTING → the
|
|
# agent stays model-agnostic; an operator swap is a setting + a
|
|
# re-embed, never an agent change.
|
|
"embed_model_name": ml.embedder_model_name,
|
|
"embed_version": ml.embedder_model_version,
|
|
"detectors": detectors,
|
|
})
|
|
return jsonify({"jobs": out})
|
|
|
|
|
|
@gpu_bp.route("/jobs/heartbeat", methods=["POST"])
|
|
async def heartbeat():
|
|
body = await request.get_json(silent=True) or {}
|
|
agent_id = str(body.get("agent_id") or "agent")
|
|
job_ids = [int(x) for x in (body.get("job_ids") or [])]
|
|
async with get_session() as session:
|
|
if not await _agent_authed(session):
|
|
return jsonify({"error": "unauthorized"}), 401
|
|
n = await GpuJobService(session).heartbeat(agent_id, job_ids)
|
|
await session.commit()
|
|
return jsonify({"extended": n})
|
|
|
|
|
|
@gpu_bp.route("/jobs/submit", methods=["POST"])
|
|
async def submit():
|
|
"""Store a job's regions + close it. regions: [{kind, bbox:[x,y,w,h],
|
|
frame_time?, score?, *_version?, ccip_embedding?, siglip_embedding?}].
|
|
replace_kinds defaults to the kinds present in the submitted regions."""
|
|
body = await request.get_json(silent=True) or {}
|
|
agent_id = str(body.get("agent_id") or "agent")
|
|
job_id = body.get("job_id")
|
|
regions = body.get("regions") or []
|
|
if job_id is None:
|
|
return jsonify({"error": "job_id required"}), 400
|
|
kinds = body.get("replace_kinds") or sorted({r["kind"] for r in regions})
|
|
async with get_session() as session:
|
|
if not await _agent_authed(session):
|
|
return jsonify({"error": "unauthorized"}), 401
|
|
job = await session.get(GpuJob, int(job_id))
|
|
if job is None or job.status != "leased" or job.lease_token != agent_id:
|
|
return jsonify({"error": "lease_invalid"}), 409
|
|
if kinds:
|
|
await RegionService(session).replace_regions(
|
|
job.image_record_id, kinds, regions
|
|
)
|
|
await GpuJobService(session).complete(agent_id, int(job_id))
|
|
await session.commit()
|
|
return jsonify({"ok": True, "stored": len(regions)})
|
|
|
|
|
|
@gpu_bp.route("/jobs/submit_embedding", methods=["POST"])
|
|
async def submit_embedding():
|
|
"""Store a whole-image SigLIP embedding (the 'embed' task) on image_record +
|
|
close the job. Body: {agent_id, job_id, embedding:[...], embedding_version}.
|
|
This is how the GPU agent re-embeds the library under a new model (#1190) —
|
|
much faster than the CPU ml-worker at higher resolutions."""
|
|
body = await request.get_json(silent=True) or {}
|
|
agent_id = str(body.get("agent_id") or "agent")
|
|
job_id = body.get("job_id")
|
|
embedding = body.get("embedding")
|
|
version = body.get("embedding_version")
|
|
if job_id is None or not embedding or not version:
|
|
return jsonify({"error": "job_id, embedding, embedding_version required"}), 400
|
|
async with get_session() as session:
|
|
if not await _agent_authed(session):
|
|
return jsonify({"error": "unauthorized"}), 401
|
|
job = await session.get(GpuJob, int(job_id))
|
|
if job is None or job.status != "leased" or job.lease_token != agent_id:
|
|
return jsonify({"error": "lease_invalid"}), 409
|
|
img = await session.get(ImageRecord, job.image_record_id)
|
|
if img is not None:
|
|
img.siglip_embedding = embedding
|
|
img.siglip_model_version = version
|
|
await GpuJobService(session).complete(agent_id, int(job_id))
|
|
await session.commit()
|
|
return jsonify({"ok": True})
|
|
|
|
|
|
@gpu_bp.route("/jobs/fail", methods=["POST"])
|
|
async def fail():
|
|
body = await request.get_json(silent=True) or {}
|
|
agent_id = str(body.get("agent_id") or "agent")
|
|
job_id = body.get("job_id")
|
|
if job_id is None:
|
|
return jsonify({"error": "job_id required"}), 400
|
|
async with get_session() as session:
|
|
if not await _agent_authed(session):
|
|
return jsonify({"error": "unauthorized"}), 401
|
|
ok = await GpuJobService(session).fail(
|
|
agent_id, int(job_id), str(body.get("error") or "")
|
|
)
|
|
await session.commit()
|
|
return jsonify({"ok": ok})
|
|
|
|
|
|
@gpu_bp.route("/jobs/release", methods=["POST"])
|
|
async def release():
|
|
"""Graceful stop: the agent hands its still-leased jobs back to pending so
|
|
they're picked up immediately instead of waiting out the lease."""
|
|
body = await request.get_json(silent=True) or {}
|
|
agent_id = str(body.get("agent_id") or "agent")
|
|
job_ids = [int(x) for x in (body.get("job_ids") or [])]
|
|
async with get_session() as session:
|
|
if not await _agent_authed(session):
|
|
return jsonify({"error": "unauthorized"}), 401
|
|
n = await GpuJobService(session).release(agent_id, job_ids)
|
|
await session.commit()
|
|
return jsonify({"released": n})
|