Tresor26/Coffee-beans-classification

1"""FastAPI service. All model work happens here; the UI is a pure client."""2import json3import threading4import time5import uuid6from contextlib import asynccontextmanager7 8from fastapi import BackgroundTasks, FastAPI, File, HTTPException, UploadFile9from fastapi.responses import JSONResponse10from pydantic import BaseModel11 12from src import config, database, model as model_module, prediction, preprocessing13 14START_TIME = time.time()15 16_JOBS: dict = {}17_retrain_lock = threading.Lock()18 19 20@asynccontextmanager21async def lifespan(app: FastAPI):22    database.init_db()23    try:24        prediction.load_champion()25        prediction.warmup()26    except Exception as exc:  # noqa: BLE001 — service must start to report this27        print(f"WARNING: could not load champion model: {exc}")28    yield29 30 31app = FastAPI(32    title="Coffee Bean Grading API",33    description="Grades green Arabica coffee beans into four classes.",34    version="1.0.0",35    lifespan=lifespan,36)37 38 39class RetrainRequest(BaseModel):40    epochs: int | None = None41    replay_n: int | None = None42    force: bool = False43 44 45@app.get("/api/health")46def health():47    """Liveness only. Deliberately does not touch the model."""48    return {"status": "ok"}49 50 51@app.get("/api/status")52def status():53    counts = preprocessing.pending_counts()54    pending_total = sum(counts.values())55    stats = database.prediction_stats()56    champion = database.get_champion()57    return {58        "status": "ok",59        "uptime_seconds": round(time.time() - START_TIME, 1),60        "model_loaded": prediction.is_ready(),61        "model_version": prediction.get_version(),62        "model_accuracy": champion["accuracy"] if champion else None,63        "classes": config.CLASS_NAMES,64        "pending_counts": counts,65        "pending_total": pending_total,66        "retrain_threshold": config.RETRAIN_THRESHOLD,67        "retrain_ready": pending_total >= config.RETRAIN_THRESHOLD,68        "predictions_served": stats["total"],69        "mean_latency_ms": stats["mean_latency_ms"],70        "p95_latency_ms": stats["p95_latency_ms"],71        "class_counts": stats["class_counts"],72    }73 74 75@app.post("/api/predict")76async def predict(file: UploadFile = File(...)):77    if not prediction.is_ready():78        raise HTTPException(503, "Model is not loaded")79    payload = await file.read()80    try:81        return prediction.predict_image(payload)82    except RuntimeError as exc:83        raise HTTPException(503, str(exc)) from exc84    except Exception as exc:  # noqa: BLE00185        raise HTTPException(415, "Could not decode image") from exc86 87 88@app.post("/api/upload")89async def upload(file: UploadFile = File(...)):90    payload = await file.read()91    if len(payload) > config.MAX_UPLOAD_BYTES:92        raise HTTPException(413, "Upload exceeds the maximum size")93    try:94        result = preprocessing.stage_upload(payload, str(uuid.uuid4())[:8])95    except ValueError as exc:96        raise HTTPException(400, str(exc)) from exc97    result["retrain_threshold"] = config.RETRAIN_THRESHOLD98    result["retrain_ready"] = result["pending_total"] >= config.RETRAIN_THRESHOLD99    return result100 101 102def _run_retrain_job(run_id: int, epochs, replay_n) -> None:103    """Synchronous by design.104 105    Starlette dispatches sync background tasks to a threadpool. An async def106    here would block the event loop and freeze the status polling that draws107    the progress log.108    """109    job = _JOBS[run_id]110    try:111        result = model_module.retrain(112            progress_cb=job["log"].append, epochs=epochs, replay_n=replay_n)113        job.update(status="completed", promoted=result["promoted"],114                   metrics=result["metrics"],115                   candidate_accuracy=result["candidate_accuracy"],116                   champion_accuracy=result["champion_accuracy"])117        database.finish_retrain_run(118            run_id, "completed", result["candidate_accuracy"],119            result["champion_accuracy"], result["promoted"],120            result["model_path"], "\n".join(job["log"]))121    except Exception as exc:  # noqa: BLE001122        job.update(status="failed", error=str(exc))123        job["log"].append(f"FAILED: {exc}")124        database.finish_retrain_run(125            run_id, "failed", None, None, False, None, "\n".join(job["log"]))126    finally:127        if _retrain_lock.locked():128            _retrain_lock.release()129 130 131@app.post("/api/retrain", status_code=202)132def trigger_retrain(request: RetrainRequest, background: BackgroundTasks):133    counts = preprocessing.pending_counts()134    pending_total = sum(counts.values())135    if pending_total == 0:136        raise HTTPException(422, "No pending images to retrain on")137    if pending_total < config.RETRAIN_THRESHOLD and not request.force:138        raise HTTPException(139            422,140            f"Only {pending_total} pending images; threshold is "141            f"{config.RETRAIN_THRESHOLD}. Upload more or use force.")142    if not _retrain_lock.acquire(blocking=False):143        raise HTTPException(409, "A retraining run is already in progress")144 145    epochs = request.epochs or config.RETRAIN_EPOCHS146    replay_n = request.replay_n or config.REPLAY_SAMPLES147    run_id = database.create_retrain_run(pending_total, replay_n, epochs)148    _JOBS[run_id] = {"status": "running", "log": [], "promoted": None,149                     "metrics": None, "error": None}150    background.add_task(_run_retrain_job, run_id, epochs, replay_n)151    return {"job_id": run_id, "status": "running",152            "n_pending": pending_total, "n_replay": replay_n}153 154 155# Declared before /api/retrain/{job_id} so "history" is not parsed as an int.156@app.get("/api/retrain/history")157def retrain_history():158    return {"runs": database.list_retrain_runs()}159 160 161@app.get("/api/retrain/{job_id}")162def retrain_status(job_id: int):163    job = _JOBS.get(job_id)164    if job is None:165        record = database.get_retrain_run(job_id)166        if record is None:167            raise HTTPException(404, f"No retraining job with id {job_id}")168        return {"job_id": job_id, "status": record["status"],169                "log": (record["log"] or "").splitlines(),170                "promoted": bool(record["promoted"]),171                "candidate_accuracy": record["candidate_accuracy"],172                "champion_accuracy": record["champion_accuracy"]}173    return {"job_id": job_id, **job}174 175 176@app.get("/api/insights")177def insights():178    if not config.INSIGHTS_PATH.exists():179        raise HTTPException(180            404, "Insights not generated. Run scripts/build_insights.py")181    return JSONResponse(json.loads(config.INSIGHTS_PATH.read_text()))182 183 184@app.get("/api/metrics")185def metrics():186    """Evaluate the deployed champion against the full test set.187 188    This is the production-evaluation surface: its numbers should line up189    with the notebook's.190    """191    if not prediction.is_ready():192        raise HTTPException(503, "Model is not loaded")193    ds = preprocessing.load_dataset(config.test_dir(), shuffle=False)194    result = model_module.evaluate(prediction.get_model(), ds)195    result["model_version"] = prediction.get_version()196    return result197