from __future__ import annotations import os import re import shutil import subprocess import sys import threading import time import math from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path from typing import Any, Literal from uuid import uuid4 from app.persistence import PersistenceStore Status = Literal["queued", "running", "completed", "failed", "stopped"] _PROGRESS_RE = re.compile(r"(\d[\d,]*)/(\d[\d,]*)") _METRIC_ROW_RE = re.compile(r"\|\s*([a-zA-Z0-9_ ]+?)\s*\|\s*(-?\d+(?:\.\d+)?)\s*\|") _EVAL_PROGRESS_RE = re.compile( r"Eval\s+num_timesteps=(\d+),\s*episode_reward=([-]?\d+(?:\.\d+)?)", re.IGNORECASE, ) _EVAL_ROW_RE = re.compile( r"^\[Eval\]\s+([a-z_]+)\s+score=([0-9.]+)\s+reward=([-0-9.]+)\s+completed=(\d+)\s+sla_breaches=(\d+)$" ) _AVG_RE = re.compile(r"^\[Eval\]\s+Average grader score:\s+([0-9.]+)$") _BEST_GRADER_RE = re.compile( r"\[Eval\]\s+New best(?: recurrent)? grader score:\s+([0-9.]+)", re.IGNORECASE, ) def _now() -> float: return time.time() def _tail_append(lines: list[str], line: str, max_size: int = 500) -> None: lines.append(line.rstrip("\n")) if len(lines) > max_size: del lines[: len(lines) - max_size] def _normalize_metric_key(raw: str) -> str: return raw.strip().lower().replace(" ", "_") def _parse_eval(stdout: str) -> tuple[list[dict[str, Any]], float | None]: rows: list[dict[str, Any]] = [] avg: float | None = None for line in stdout.splitlines(): line = line.strip() if not line: continue row = _EVAL_ROW_RE.match(line) if row: rows.append( { "task_id": row.group(1), "grader_score": float(row.group(2)), "total_reward": float(row.group(3)), "total_completed": int(row.group(4)), "total_sla_breaches": int(row.group(5)), } ) continue m = _AVG_RE.match(line) if m: avg = float(m.group(1)) return rows, avg @dataclass class TrainingJob: job_id: str phase: int timesteps: int n_envs: int seed: int config_path: str created_at: float = field(default_factory=_now) started_at: float | None = None updated_at: float = field(default_factory=_now) ended_at: float | None = None status: Status = "queued" progress: float = 0.0 process_id: int | None = None command: list[str] = field(default_factory=list) output_model_path: str | None = None output_model_name: str | None = None latest_metrics: dict[str, float] = field(default_factory=dict) metric_history: list[dict[str, Any]] = field(default_factory=list) evaluation_rows: list[dict[str, Any]] = field(default_factory=list) evaluation_avg_score: float | None = None logs_tail: list[str] = field(default_factory=list) error_message: str | None = None return_code: int | None = None process: subprocess.Popen[str] | None = field(default=None, repr=False) lock: threading.Lock = field(default_factory=threading.Lock, repr=False) last_persist_at: float = field(default_factory=lambda: 0.0, repr=False) def snapshot(self) -> dict[str, Any]: with self.lock: return { "job_id": self.job_id, "phase": self.phase, "timesteps": self.timesteps, "n_envs": self.n_envs, "seed": self.seed, "config_path": self.config_path, "created_at": self.created_at, "started_at": self.started_at, "updated_at": self.updated_at, "ended_at": self.ended_at, "status": self.status, "progress": self.progress, "process_id": self.process_id, "command": self.command, "output_model_path": self.output_model_path, "output_model_name": self.output_model_name, "latest_metrics": dict(self.latest_metrics), "metric_history": list(self.metric_history), "evaluation_rows": list(self.evaluation_rows), "evaluation_avg_score": self.evaluation_avg_score, "logs_tail": list(self.logs_tail), "error_message": self.error_message, "return_code": self.return_code, } class TrainingJobManager: def __init__(self, repo_root: Path, persistence: PersistenceStore | None = None) -> None: self._repo_root = repo_root self._persistence = persistence self._jobs: dict[str, TrainingJob] = {} self._lock = threading.Lock() self._training_runs_root = ( self._persistence.training_runs_dir if self._persistence is not None and self._persistence.enabled else self._repo_root / "results" / "training_runs" ) self._load_persisted_jobs() def _load_persisted_jobs(self) -> None: if self._persistence is None or not self._persistence.enabled: return persisted = self._persistence.list_training_jobs(limit=500) with self._lock: for snap in persisted: try: job = TrainingJob( job_id=str(snap["job_id"]), phase=int(snap["phase"]), timesteps=int(snap["timesteps"]), n_envs=int(snap["n_envs"]), seed=int(snap["seed"]), config_path=str(snap.get("config_path") or ""), created_at=float(snap.get("created_at") or _now()), started_at=float(snap["started_at"]) if snap.get("started_at") is not None else None, updated_at=float(snap.get("updated_at") or _now()), ended_at=float(snap["ended_at"]) if snap.get("ended_at") is not None else None, status=str(snap.get("status") or "failed"), progress=float(snap.get("progress") or 0.0), process_id=int(snap["process_id"]) if snap.get("process_id") is not None else None, command=list(snap.get("command") or []), output_model_path=snap.get("output_model_path"), output_model_name=snap.get("output_model_name"), latest_metrics=dict(snap.get("latest_metrics") or {}), metric_history=list(snap.get("metric_history") or []), evaluation_rows=list(snap.get("evaluation_rows") or []), evaluation_avg_score=( float(snap["evaluation_avg_score"]) if snap.get("evaluation_avg_score") is not None else None ), logs_tail=list(snap.get("logs_tail") or []), error_message=snap.get("error_message"), return_code=int(snap["return_code"]) if snap.get("return_code") is not None else None, ) except Exception: continue # Process handles cannot survive a server restart. Recover to terminal state. if job.status in ("queued", "running"): job.status = "failed" msg = "Recovered after restart: previous process state unavailable." job.error_message = f"{job.error_message} {msg}".strip() if job.error_message else msg if job.ended_at is None: job.ended_at = _now() job.process = None self._jobs[job.job_id] = job def clear_jobs(self, *, clear_artifacts: bool = False) -> int: to_stop: list[subprocess.Popen[str]] = [] with self._lock: removed = len(self._jobs) for job in self._jobs.values(): with job.lock: proc = job.process if proc is not None and job.status in ("queued", "running"): to_stop.append(proc) self._jobs.clear() for proc in to_stop: try: proc.terminate() except Exception: pass if self._persistence is not None and self._persistence.enabled: self._persistence.clear_training_jobs() if clear_artifacts: try: if self._training_runs_root.exists(): shutil.rmtree(self._training_runs_root, ignore_errors=True) self._training_runs_root.mkdir(parents=True, exist_ok=True) except Exception: pass return removed def _persist_job(self, job: TrainingJob) -> None: if self._persistence is None or not self._persistence.enabled: return snapshot = job.snapshot() self._persistence.upsert_training_job(snapshot) with job.lock: job.last_persist_at = _now() def list_jobs(self) -> list[dict[str, Any]]: with self._lock: jobs = list(self._jobs.values()) jobs.sort(key=lambda x: x.created_at, reverse=True) return [job.snapshot() for job in jobs] def get_job(self, job_id: str) -> dict[str, Any] | None: with self._lock: job = self._jobs.get(job_id) return None if job is None else job.snapshot() def start_job( self, *, phase: int, timesteps: int, n_envs: int, seed: int | None, config_path: str | None, ) -> dict[str, Any]: job_id = str(uuid4()) job_seed = int(seed if seed is not None else int(time.time()) % 1_000_000) cfg = config_path or ( "rl/configs/ppo_easy.yaml" if phase == 1 else "rl/configs/curriculum.yaml" ) job = TrainingJob( job_id=job_id, phase=phase, timesteps=timesteps, n_envs=n_envs, seed=job_seed, config_path=cfg, ) with self._lock: self._jobs[job_id] = job cmd = [ sys.executable, "-u", "-m", "rl.train_ppo", "--phase", str(phase), "--timesteps", str(timesteps), "--n-envs", str(n_envs), "--seed", str(job_seed), ] if phase == 1: # Keep Phase 1 UI responsive by emitting multiple eval checkpoints # across the requested run length instead of only near the end. phase1_eval_freq = max(128, int((timesteps / max(n_envs, 1)) / 15)) cmd.extend( [ "--phase1-config", cfg, "--phase1-eval-freq", str(phase1_eval_freq), ] ) else: cmd.extend(["--phase2-config", cfg]) env = os.environ.copy() env["PYTHONUNBUFFERED"] = "1" proc = subprocess.Popen( cmd, cwd=str(self._repo_root), env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, ) with job.lock: job.command = cmd job.status = "running" job.started_at = _now() job.updated_at = _now() job.process_id = proc.pid job.process = proc _tail_append(job.logs_tail, f"[training_jobs] started pid={proc.pid}") _tail_append(job.logs_tail, f"[training_jobs] command: {' '.join(cmd)}") self._persist_job(job) t = threading.Thread(target=self._watch_job, args=(job,), daemon=True) t.start() return job.snapshot() @staticmethod def _append_metric_point_locked( job: TrainingJob, *, timesteps: float | None, reward: float | None = None, score: float | None = None, source: str | None = None, max_points: int = 5000, ) -> None: """ Append (or merge) a structured metric point while holding job.lock. """ if timesteps is None or not math.isfinite(float(timesteps)): return payload: dict[str, Any] = {"t": float(timesteps)} if reward is not None and math.isfinite(float(reward)): payload["ep_rew_mean"] = float(reward) if score is not None and math.isfinite(float(score)): payload["grader_score"] = float(score) if source: payload["source"] = str(source) if "ep_rew_mean" not in payload and "grader_score" not in payload: return if job.metric_history and float(job.metric_history[-1].get("t", -1.0)) == float(payload["t"]): job.metric_history[-1].update(payload) else: job.metric_history.append(payload) if len(job.metric_history) > max_points: del job.metric_history[: len(job.metric_history) - max_points] def stop_job(self, job_id: str) -> dict[str, Any] | None: with self._lock: job = self._jobs.get(job_id) if job is None: return None with job.lock: proc = job.process if proc is None or job.status not in ("running", "queued"): return job.snapshot() job.status = "stopped" job.updated_at = _now() self._persist_job(job) try: proc.terminate() except Exception: pass return job.snapshot() def delete_job(self, job_id: str, *, clear_artifacts: bool = False) -> bool: with self._lock: job = self._jobs.pop(job_id, None) if job is None: return False with job.lock: proc = job.process status = job.status output_model_path = job.output_model_path if proc is not None and status in ("queued", "running"): try: proc.terminate() except Exception: pass if self._persistence is not None and self._persistence.enabled: self._persistence.delete_training_job(job_id) if clear_artifacts and output_model_path: try: out = Path(output_model_path) if out.exists() and out.is_file(): out.unlink(missing_ok=True) parent = out.parent if parent.exists() and parent.is_dir() and not any(parent.iterdir()): parent.rmdir() except Exception: pass return True def _watch_job(self, job: TrainingJob) -> None: proc = job.process if proc is None or proc.stdout is None: with job.lock: job.status = "failed" job.error_message = "Training process failed to start." job.updated_at = _now() job.ended_at = _now() self._persist_job(job) return for line in proc.stdout: self._update_from_line(job, line) return_code = proc.wait() with job.lock: job.return_code = int(return_code) if job.status == "stopped": job.ended_at = _now() job.updated_at = _now() job.process = None return if return_code == 0: job.status = "completed" job.progress = 1.0 else: job.status = "failed" base_error = f"Training exited with code {return_code}." if not job.logs_tail: _tail_append( job.logs_tail, "[training_jobs] Process ended before producing logs. " "Check RL dependencies/environment and training command arguments.", ) job.error_message = base_error job.ended_at = _now() job.updated_at = _now() job.process = None self._persist_job(job) if return_code == 0: self._finalize_artifacts(job) def _update_from_line(self, job: TrainingJob, line: str) -> None: line = line.rstrip("\n") should_persist = False with job.lock: _tail_append(job.logs_tail, line) job.updated_at = _now() p = _PROGRESS_RE.search(line) if p: num = int(p.group(1).replace(",", "")) den = int(p.group(2).replace(",", "")) if den > 0: job.progress = max(0.0, min(1.0, num / den)) ep = _EVAL_PROGRESS_RE.search(line) if ep: ts = int(ep.group(1)) rew = float(ep.group(2)) job.latest_metrics["total_timesteps"] = float(ts) job.latest_metrics["ep_rew_mean"] = rew self._append_metric_point_locked( job, timesteps=float(ts), reward=rew, source="eval_progress", ) if job.timesteps > 0: job.progress = max(0.0, min(1.0, ts / float(job.timesteps))) m = _METRIC_ROW_RE.search(line) if m: key = _normalize_metric_key(m.group(1)) val = float(m.group(2)) interesting = { "total_timesteps", "ep_rew_mean", "ep_len_mean", "grader_score", "mean_reward", "mean_ep_length", "episode_mean_sla_penalty", "episode_mean_fairness_penalty", "explained_variance", "approx_kl", } if key in interesting: job.latest_metrics[key] = val current_ts = job.latest_metrics.get("total_timesteps") if key == "total_timesteps": self._append_metric_point_locked( job, timesteps=val, reward=job.latest_metrics.get("ep_rew_mean"), score=job.latest_metrics.get("grader_score") or job.latest_metrics.get("avg_grader_score"), source="metrics_row_ts", ) elif key in {"ep_rew_mean", "mean_reward"}: self._append_metric_point_locked( job, timesteps=float(current_ts) if current_ts is not None else None, reward=val, source="metrics_row_reward", ) elif key in {"grader_score", "avg_grader_score"}: self._append_metric_point_locked( job, timesteps=float(current_ts) if current_ts is not None else None, score=val, source="metrics_row_score", ) best = _BEST_GRADER_RE.search(line) if best: score = float(best.group(1)) job.latest_metrics["grader_score"] = score fallback_ts = ( float(job.latest_metrics.get("total_timesteps")) if "total_timesteps" in job.latest_metrics else float(job.metric_history[-1]["t"]) if job.metric_history else 0.0 ) self._append_metric_point_locked( job, timesteps=fallback_ts if fallback_ts > 0 else float(len(job.metric_history) + 1), score=score, source="best_grader", ) avg_line = _AVG_RE.match(line.strip()) if avg_line: avg_score = float(avg_line.group(1)) job.latest_metrics["avg_grader_score"] = avg_score fallback_ts = ( float(job.latest_metrics.get("total_timesteps")) if "total_timesteps" in job.latest_metrics else float(job.metric_history[-1]["t"]) if job.metric_history else 0.0 ) self._append_metric_point_locked( job, timesteps=fallback_ts if fallback_ts > 0 else float(len(job.metric_history) + 1), score=avg_score, source="avg_grader", ) if job.updated_at - job.last_persist_at >= 1.5: should_persist = True if should_persist: self._persist_job(job) def _finalize_artifacts(self, job: TrainingJob) -> None: src_name = "phase1_final.zip" if job.phase == 1 else "phase2_final.zip" src = self._repo_root / "results" / "best_model" / src_name run_dir = self._training_runs_root / job.job_id run_dir.mkdir(parents=True, exist_ok=True) # Keep a mirror under repo/results for local developer convenience. mirror_dir = self._repo_root / "results" / "training_runs" / job.job_id if mirror_dir != run_dir: mirror_dir.mkdir(parents=True, exist_ok=True) if src.exists(): ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") unique_name = f"phase{job.phase}_seed{job.seed}_{ts}_{job.job_id[:8]}.zip" out = run_dir / unique_name shutil.copy2(src, out) if mirror_dir != run_dir: try: shutil.copy2(src, mirror_dir / unique_name) except Exception: pass with job.lock: job.output_model_path = str(out.resolve()) job.output_model_name = unique_name job.updated_at = _now() model_type = "maskable" eval_cmd = [ sys.executable, "-m", "rl.evaluate", "--model", str(out), "--episodes", "3", "--model-type", model_type, ] proc = subprocess.run( eval_cmd, cwd=str(self._repo_root), env=os.environ.copy(), capture_output=True, text=True, check=False, ) rows, avg = _parse_eval(proc.stdout or "") with job.lock: job.evaluation_rows = rows job.evaluation_avg_score = avg if avg is not None: job.latest_metrics["avg_grader_score"] = float(avg) fallback_ts = ( float(job.latest_metrics.get("total_timesteps")) if "total_timesteps" in job.latest_metrics else float(job.timesteps) ) self._append_metric_point_locked( job, timesteps=fallback_ts if fallback_ts > 0 else float(len(job.metric_history) + 1), score=float(avg), source="final_eval_avg", ) _tail_append(job.logs_tail, "----- EVALUATION -----") for ln in (proc.stdout or "").splitlines(): _tail_append(job.logs_tail, ln) if proc.returncode != 0 and not job.error_message: job.error_message = f"Evaluation exited with code {proc.returncode}." job.updated_at = _now() self._persist_job(job) else: self._persist_job(job)