vh
888ba6a714
Host service (runs as llmuser, owns /opt/fluxgym + GPU access) that runs
sd-scripts SDXL LoRA training on demand for arbo — the infra-ops half of the
in-arbo LoRA training Phase 1 ownership split (vh/arbo
docs/contracts/in-arbo-lora-training-phase1.contract.md §4.1/§2).
- Fixed-invocation only (INV-T7): bounded params -> one sd-scripts command
shape; every param range/allowlist/path-containment checked before spawn;
bad request = 422, never a silent downgrade. 14 unit tests green.
- Thin supervisor: never imports torch; subprocesses the fluxgym venv's
accelerate. 1-job-at-a-time (arbo lease is the serializer, 409 is backstop).
Durable job records + boot reconciliation (§4.3).
- API: POST /train, GET /train/{id}[/log], POST /train/{id}/cancel,
GET /gpu-status (per-device VRAM + tts_on_3090 co-OOM signal), GET /healthz.
- Wire-shape (§7 resolved with comfy-dev): shared /worktank/arbo/train handoff
(group arbotrain, setgid 2770); worker binds 0.0.0.0:8203, arbo reaches via
host.docker.internal:host-gateway (reachability proven on 172.20.0.1:8203);
device-aware TTS steering via /gpu-status.
Deployed to irv-ml1 via playbooks/deploy-lora-training-worker.yaml (elway,
idempotent); systemd unit active; /healthz + /gpu-status verified live.
86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
"""FastAPI surface for the LoRA training worker (§4.1 API — arbo is the client).
|
|
|
|
Endpoints:
|
|
POST /train dispatch a train (409 if busy, 422 on invalid params) → {worker_job_id}
|
|
GET /train/{id} status {status, step, total_steps, loss, eta_s, lora_path?, error?}
|
|
GET /train/{id}/log tail of the run log
|
|
POST /train/{id}/cancel best-effort kill → cancelled
|
|
GET /gpu-status per-device VRAM + tts_on_3090 (arbo's device-aware steering, §7)
|
|
GET /healthz liveness + whether a train is active
|
|
|
|
The worker holds ONE job at a time; arbo's lease is the real serializer (INV-T2), the 409
|
|
here is the backstop. Nothing here interprets free-form arguments — see invocation.py (INV-T7).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.responses import PlainTextResponse
|
|
from pydantic import BaseModel, Field
|
|
|
|
from . import config
|
|
from .gpu import gpu_status
|
|
from .invocation import InvalidTrainRequest
|
|
from .jobs import Busy, manager
|
|
|
|
app = FastAPI(title="LoRA Training Worker", version="0.1.0")
|
|
|
|
|
|
class TrainRequest(BaseModel):
|
|
# Bounded params only (INV-T7). Full validation (ranges, path containment, tier/device
|
|
# fit) happens in invocation.validate_request; pydantic just pins the shape + types.
|
|
dataset_dir: str
|
|
base_model_path: str
|
|
output_dir: str
|
|
output_name: str
|
|
trigger: str
|
|
subject_class: str
|
|
repeats: int = Field(ge=1, le=100)
|
|
tier: str
|
|
device_index: int = Field(ge=0, le=1)
|
|
seed: int = Field(default=42, ge=0)
|
|
|
|
|
|
@app.post("/train")
|
|
def post_train(req: TrainRequest):
|
|
try:
|
|
return manager.submit(req.model_dump())
|
|
except InvalidTrainRequest as exc:
|
|
raise HTTPException(status_code=422, detail=str(exc))
|
|
except Busy as exc:
|
|
raise HTTPException(status_code=409, detail=str(exc))
|
|
|
|
|
|
@app.get("/train/{job_id}")
|
|
def get_train(job_id: str):
|
|
status = manager.get(job_id)
|
|
if status is None:
|
|
raise HTTPException(status_code=404, detail=f"unknown worker_job_id: {job_id}")
|
|
return status
|
|
|
|
|
|
@app.get("/train/{job_id}/log", response_class=PlainTextResponse)
|
|
def get_train_log(job_id: str):
|
|
tail = manager.log_tail(job_id)
|
|
if tail is None:
|
|
raise HTTPException(status_code=404, detail=f"no log for worker_job_id: {job_id}")
|
|
return tail
|
|
|
|
|
|
@app.post("/train/{job_id}/cancel")
|
|
def post_cancel(job_id: str):
|
|
status = manager.cancel(job_id)
|
|
if status is None:
|
|
raise HTTPException(status_code=404, detail=f"unknown worker_job_id: {job_id}")
|
|
return status
|
|
|
|
|
|
@app.get("/gpu-status")
|
|
def get_gpu_status():
|
|
return gpu_status()
|
|
|
|
|
|
@app.get("/healthz")
|
|
def healthz():
|
|
return {"ok": True, "active_job": manager.active(), "port": config.PORT}
|