FastAPI and the training worker

sudo apt install -y redis-server && sudo systemctl enable --now redis-server
cd ~/sutraai && source .venv/bin/activate
pip install "fastapi[standard]" celery redis sse-starlette python-multipart
mkdir -p apps/api projects

import json, subprocess, sys
from pathlib import Path
import redis
from celery import Celery

ROOT = Path.home() / "sutraai"
app = Celery("sutraai", broker="redis://localhost:6379/0", backend="redis://localhost:6379/1")
r = redis.Redis()

@app.task(bind=True)
def train_job(self, project_id):
    proj = ROOT / "projects" / project_id
    channel, history = f"job:{self.request.id}", f"job:{self.request.id}:log"
    cmd = [sys.executable, "-m", "ml_core.run", str(proj / "pipeline.yaml"), "--data", str(proj / "data"), "--out", str(proj)]
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    for line in proc.stdout:
        line = line.strip()
        if line.startswith("{"):
            r.rpush(history, line); r.publish(channel, line)          # history lets late listeners catch up
    proc.wait()
    end = json.dumps({"event": "finished", "returncode": proc.returncode})
    r.rpush(history, end); r.publish(channel, end)
    return proc.returncode

import json, shutil, uuid, zipfile
from pathlib import Path
import redis.asyncio as aioredis
from fastapi import FastAPI, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from sse_starlette.sse import EventSourceResponse
from worker import train_job

ROOT = Path.home() / "sutraai"; PROJECTS = ROOT / "projects"
TEMPLATES = {"image", "sensor", "audio", "detect"}
app = FastAPI(title="SutraAI API", version="0.1")
app.add_middleware(CORSMiddleware, allow_origins=["http://localhost:3000"], allow_methods=["*"], allow_headers=["*"])

def project_dir(pid):
    d = PROJECTS / pid
    if not d.is_dir(): raise HTTPException(404, "project not found")
    return d

@app.post("/projects")
def create_project(template: str):
    if template not in TEMPLATES: raise HTTPException(400, f"template must be one of {sorted(TEMPLATES)}")
    pid = uuid.uuid4().hex[:8]; d = PROJECTS / pid; d.mkdir(parents=True)
    shutil.copy(ROOT / "configs" / f"{template}.yaml", d / "pipeline.yaml")
    return {"project_id": pid, "template": template}

@app.post("/projects/{pid}/dataset")
async def upload_dataset(pid: str, file: UploadFile):
    d = project_dir(pid); data = d / "data"; zpath = d / "upload.zip"
    zpath.write_bytes(await file.read())
    with zipfile.ZipFile(zpath) as z:
        for name in z.namelist():                               # refuse paths that escape the folder
            if not (data / name).resolve().is_relative_to(data.resolve()):
                raise HTTPException(400, "zip contains unsafe paths")
        z.extractall(data)
    return {"files": sum(1 for p in data.rglob("*") if p.is_file())}

@app.post("/projects/{pid}/train")
def start_training(pid: str):
    project_dir(pid)
    return {"job_id": train_job.delay(pid).id}

@app.get("/jobs/{job_id}/events")
async def job_events(job_id: str):
    r = aioredis.from_url("redis://localhost")
    async def stream():
        for old in await r.lrange(f"job:{job_id}:log", 0, -1):        # replay what already happened
            yield {"data": old.decode()}
            if b'"finished"' in old: return
        pubsub = r.pubsub(); await pubsub.subscribe(f"job:{job_id}")
        async for msg in pubsub.listen():
            if msg["type"] != "message": continue
            yield {"data": msg["data"].decode()}
            if b'"finished"' in msg["data"]: break
    return EventSourceResponse(stream())

@app.get("/projects/{pid}/metrics")
def metrics(pid: str):
    files = sorted(project_dir(pid).glob("reports/*/metrics.json"))
    if not files: raise HTTPException(404, "no metrics yet")
    return json.loads(files[-1].read_text())

@app.get("/projects/{pid}/download")
def download(pid: str, fmt: str = "onnx"):
    patterns = {"onnx": "*.onnx", "tflite_int8": "*full_integer_quant.tflite", "c_array": "*_model_data.cc",
                "pytorch": "best.pt", "engine": "*.engine"}
    if fmt not in patterns: raise HTTPException(400, f"fmt must be one of {sorted(patterns)}")
    hit = next(project_dir(pid).rglob(patterns[fmt]), None)
    if not hit: raise HTTPException(404, f"no {fmt} file yet")
    return FileResponse(hit, filename=hit.name)

cd ~/sutraai/apps/api && source ~/sutraai/.venv/bin/activate
tmux new -s worker -d "celery -A worker worker --concurrency 1 -l info"     # one training job at a time on the GPU
tmux new -s api -d "fastapi run app.py --host 0.0.0.0 --port 8000"
tmux ls