"""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}