vh
b617a8b674
comfy-dev's explicit-over-implicit call: arbo now sends train_id, so the
worker no longer derives the loras/trained/{train_id}/ namespace from
output_dir.parent (which coupled it to arbo's handoff layout). train_id is
optional + path-safe-validated; when present it wins, else the path
derivation remains as the fallback. Wired through TrainRequest ->
validate_request -> published_relative_path -> _publish_lora. 18 tests green.
87 lines
2.8 KiB
Python
87 lines
2.8 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)
|
|
train_id: str | None = None # Phase 2: arbo's id, the publish-path namespace (optional)
|
|
|
|
|
|
@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}
|