train_api.py

import json
import os
import subprocess
import sys
import threading
import time
from pathlib import Path

from fastapi import File, Request, UploadFile
from fastapi.responses import FileResponse, JSONResponse, PlainTextResponse


ROOT = Path(__file__).resolve().parent
RUNS = ROOT / "training_space" / "runs"
RUNS.mkdir(parents=True, exist_ok=True)
IMAGE_FOLDERS = {"train": ROOT / "hpt_data", "validation": ROOT / "hpt_data_val"}
MODEL_FOLDERS = (ROOT / "pbc3_students", RUNS)
LOCK = threading.Lock()
PROCESS = None
CURRENT = {}


def _safe_model_path(value):
    path = (ROOT / value).resolve() if not os.path.isabs(value) else Path(value).resolve()
    if not any(path == folder.resolve() or folder.resolve() in path.parents for folder in MODEL_FOLDERS):
        return None
    return path if path.is_file() else None


def _output_path(value):
    name = Path(value or "rl_training.npz").name
    if not name.endswith(".npz"):
        name += ".npz"
    return (RUNS / name).resolve()


def _json_lines(log_path):
    rows = []
    if not log_path.exists():
        return rows
    for line in log_path.read_text(encoding="utf-8", errors="replace").splitlines():
        candidate = line.strip()
        if candidate.startswith("initial "):
            candidate = candidate[8:]
        try:
            row = json.loads(candidate)
        except Exception:
            continue
        if "epoch" in row or "reward" in row:
            rows.append(row)
    return rows


def _refresh_process():
    global PROCESS
    if PROCESS is not None and PROCESS.poll() is not None:
        CURRENT["return_code"] = PROCESS.returncode
        PROCESS = None


def status():
    with LOCK:
        _refresh_process()
        output = Path(CURRENT["output"]) if CURRENT.get("output") else None
        log = Path(CURRENT["log"]) if CURRENT.get("log") else None
        history_path = Path(f"{output}.json") if output else None
        manifest_path = Path(f"{output}.checkpoints.json") if output else None
        history = []
        checkpoint_groups = {"mse": [], "bpp": [], "work": [], "overall": []}
        if history_path and history_path.exists():
            try:
                saved = json.loads(history_path.read_text(encoding="utf-8"))
                initial = saved.get("initial")
                if initial:
                    history.append(initial if "validation" in initial else {"validation": initial})
                history.extend(saved.get("history", []))
            except Exception:
                pass
        if not history:
            history = _json_lines(log) if log else []
        if manifest_path and manifest_path.exists():
            try:
                checkpoint_groups = json.loads(manifest_path.read_text(encoding="utf-8")).get("groups", checkpoint_groups)
            except Exception:
                pass
        return {
            "running": PROCESS is not None,
            "return_code": CURRENT.get("return_code"),
            "output": str(output.relative_to(ROOT)) if output else None,
            "log": str(log.relative_to(ROOT)) if log else None,
            "started": CURRENT.get("started"),
            "spec": CURRENT.get("spec", {}),
            "history": history,
            "checkpoint_groups": checkpoint_groups,
            "log_tail": log.read_text(encoding="utf-8", errors="replace")[-10000:] if log and log.exists() else "",
        }


def start(spec):
    global PROCESS
    with LOCK:
        _refresh_process()
        if PROCESS is not None:
            return {"error": "A training run is already active."}
        presets = str(spec.get("presets", "high_quality"))
        selected = [p.strip() for p in presets.split(",") if p.strip()]
        allowed = {"compression", "balanced", "quality", "high_quality"}
        if not selected or any(p not in allowed for p in selected):
            return {"error": "Presets must be compression, balanced, quality, or high_quality."}
        init = _safe_model_path(str(spec.get("init", "pbc3_students/patch_policy_f26_a20_h512_l2_e1200.npz")))
        if init is None:
            return {"error": "Initial model was not found in pbc3_students or training_space/runs."}
        default_output = f"rl_{'_'.join(selected)}_top2.npz"
        output = _output_path(spec.get("output") or default_output)
        output.parent.mkdir(parents=True, exist_ok=True)
        log = Path(f"{output}.log")
        resume = Path(f"{output}.resume.pt")
        command = [
            sys.executable, "pbc3_rl_a20.py",
            "--presets", ",".join(selected),
            "--epochs", str(int(spec.get("epochs", 100))),
            "--batch", str(int(spec.get("batch", 4))),
            "--init", os.fspath(init),
            "--out", os.fspath(output),
            "--rate-weight", str(float(spec.get("rate_weight", 1.5))),
            "--speed-weight", str(float(spec.get("speed_weight", 0.15))),
            "--temperature", str(float(spec.get("temperature", 0.9))),
            "--entropy-weight", str(float(spec.get("entropy_weight", 0.003))),
            "--kl-weight", str(float(spec.get("kl_weight", 0.015))),
            "--quality-weight", str(float(spec.get("quality_weight", 2.0))),
            "--worse-quality-weight", str(float(spec.get("worse_quality_weight", 7.0))),
        ]
        if bool(spec.get("resume", True)) and resume.exists():
            command.extend(["--resume", os.fspath(resume)])
        log.write_text("", encoding="utf-8")
        handle = log.open("a", encoding="utf-8")
        PROCESS = subprocess.Popen(command, cwd=ROOT, stdout=handle, stderr=subprocess.STDOUT)
        handle.close()
        CURRENT.clear()
        CURRENT.update({"output": str(output), "log": str(log), "started": time.time(), "spec": spec, "return_code": None})
        return {"ok": True, "output": str(output.relative_to(ROOT)), "resuming": resume.exists()}


def stop():
    with LOCK:
        _refresh_process()
        if PROCESS is None:
            return {"ok": False, "error": "No training run is active."}
        PROCESS.terminate()
        return {"ok": True}


def checkpoints():
    rows = []
    for folder in MODEL_FOLDERS:
        for path in sorted(folder.glob("*.npz"), key=lambda p: p.stat().st_mtime, reverse=True):
            rows.append({"name": path.name, "path": str(path.relative_to(ROOT)), "bytes": path.stat().st_size, "modified": path.stat().st_mtime})
    return {"checkpoints": rows}


def download_checkpoint(path):
    safe = _safe_model_path(path)
    if safe is None or safe.suffix != ".npz":
        return JSONResponse({"error": "Checkpoint not found or not allowed."}, status_code=404)
    return FileResponse(safe, filename=safe.name, media_type="application/octet-stream")


async def upload_checkpoint(file: UploadFile):
    name = Path(file.filename or "checkpoint.npz").name
    if Path(name).suffix.lower() != ".npz":
        return JSONResponse({"error": "Only .npz checkpoints are supported."}, status_code=400)
    stem = Path(name).stem.replace(" ", "_")
    path = (RUNS / f"uploaded_{int(time.time())}_{stem}.npz").resolve()
    path.write_bytes(await file.read())
    return {"path": str(path.relative_to(ROOT)), "name": path.name}


def register(app):
    @app.get("/api/train/status")
    def train_status():
        return status()

    @app.post("/api/train/start")
    async def train_start(request: Request):
        result = start(await request.json())
        return JSONResponse(result, status_code=400 if result.get("error") else 200)

    @app.post("/api/train/stop")
    def train_stop():
        return stop()

    @app.get("/api/train/checkpoints")
    def train_checkpoints():
        return checkpoints()

    @app.get("/api/train/download_checkpoint")
    def train_download_checkpoint(path: str):
        return download_checkpoint(path)

    @app.post("/api/train/upload_checkpoint")
    async def train_upload_checkpoint(file: UploadFile = File(...)):
        return await upload_checkpoint(file)

    @app.get("/api/train/log")
    def train_log():
        data = status()
        path = ROOT / data["log"] if data.get("log") else None
        if not path or not path.exists():
            return PlainTextResponse("No training log is available.")
        return FileResponse(path, filename=path.name, media_type="text/plain")