File size: 8,320 Bytes
a15e50d 38d12b7 a15e50d 979c011 a15e50d 38d12b7 a15e50d b60c4c8 7faa636 a15e50d 979c011 a15e50d 979c011 a15e50d 112a4cb a15e50d 38d12b7 a15e50d 38d12b7 a15e50d 38d12b7 a15e50d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 | 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")
|