run_final_window_queue.py
末期逐 episode actor 评估
原文件 · 387 行 · 14607 字节
scripts/experiments/run_final_window_queue.py
SHA-256: 6a4bd044765a4b6d49cf3085cae6c89cc7a2c7298a1b25a04d8ed5f105531a76
1"""Evaluate every available actor checkpoint in the final training window.23Each training job is one subprocess containing all of its final-window models.4This avoids starting thousands of Python processes while preserving one result5record per actor checkpoint. Planned checkpoints that were not frozen by the6rollout-granular trainer remain explicit missing work in ``status.json``.7"""89import argparse10import json11import os12import re13import subprocess14import sys15import threading16import time17from concurrent.futures import ThreadPoolExecutor, as_completed18from datetime import datetime19from pathlib import Path2021import yaml222324REPO_ROOT = Path(__file__).resolve().parents[2]252627def _now():28 return datetime.now().astimezone().isoformat(timespec="seconds")293031def _load_json(path):32 try:33 return json.loads(Path(path).read_text(encoding="utf-8"))34 except (OSError, ValueError, json.JSONDecodeError):35 return None363738def _atomic_json(path, payload):39 path = Path(path)40 path.parent.mkdir(parents=True, exist_ok=True)41 temporary = path.with_name(f"{path.name}.{os.getpid()}.tmp")42 temporary.write_text(43 json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"44 )45 temporary.replace(path)464748def _atomic_yaml(path, payload):49 path = Path(path)50 path.parent.mkdir(parents=True, exist_ok=True)51 temporary = path.with_name(f"{path.name}.{os.getpid()}.tmp")52 temporary.write_text(53 yaml.safe_dump(payload, sort_keys=False, allow_unicode=True),54 encoding="utf-8",55 )56 temporary.replace(path)575859def _parse_named_paths(values, multiple=False):60 parsed = {}61 for value in values:62 if "=" not in value:63 raise ValueError(f"expected NAME=PATH, got {value!r}")64 name, path = value.split("=", 1)65 if multiple:66 parsed.setdefault(name, []).append(Path(path).resolve())67 else:68 parsed[name] = Path(path).resolve()69 return parsed707172class FinalWindowQueue:73 def __init__(74 self,75 output_dir,76 stages,77 recovery_stages,78 cuda,79 max_workers,80 fraction,81 eval_episodes,82 eval_seed,83 job_pattern=None,84 ):85 self.root = Path(output_dir).resolve()86 self.status_path = self.root / "status.json"87 self.stages = stages88 self.recovery_stages = recovery_stages89 self.cuda = int(cuda)90 self.max_workers = int(max_workers)91 self.fraction = float(fraction)92 self.eval_episodes = int(eval_episodes)93 self.eval_seed = int(eval_seed)94 self.job_pattern = re.compile(job_pattern) if job_pattern else None95 self.lock = threading.Lock()96 self.state = _load_json(self.status_path) or {97 "name": "final-window-v1",98 "created_at": _now(),99 "state": "initializing",100 "phase": "final_window",101 "cuda": self.cuda,102 "max_workers": self.max_workers,103 "fraction": self.fraction,104 "eval_episodes_per_checkpoint": self.eval_episodes,105 "eval_seed": self.eval_seed,106 "tasks": {},107 }108 # Requeue interrupted/failed per-job suites. A completed suite is109 # immutable and will be reused on process restart.110 for task in self.state.get("tasks", {}).values():111 if task.get("state") in ("pending", "running", "failed"):112 task.update(state="pending")113 for key in ("pid", "return_code", "finished_at", "started_at"):114 task.pop(key, None)115 self._write_state()116117 def _write_state(self):118 tasks = list(self.state.get("tasks", {}).values())119 planned = sum(int(task.get("planned_units", 0)) for task in tasks)120 submitted = sum(int(task.get("submitted_units", 0)) for task in tasks)121 completed = sum(int(task.get("completed_units", 0)) for task in tasks)122 self.state.update(123 updated_at=_now(),124 counts={125 key: sum(task.get("state") == key for task in tasks)126 for key in ("pending", "running", "completed", "failed")127 },128 planned_units=planned,129 submitted_units=submitted,130 completed_units=completed,131 missing_units=max(0, planned - submitted),132 )133 _atomic_json(self.status_path, self.state)134135 def _stage_jobs(self, stage):136 status = _load_json(self.stages[stage] / "status.json") or {}137 jobs = dict(status.get("jobs", {}))138 for recovery_dir in self.recovery_stages.get(stage, []):139 recovery = _load_json(recovery_dir / "status.json") or {}140 for name, job in recovery.get("jobs", {}).items():141 if job.get("state") == "completed":142 jobs[name] = job143 return jobs144145 def _discover(self):146 discovered = []147 for stage in self.stages:148 for job_name, job in self._stage_jobs(stage).items():149 if job.get("state") != "completed":150 continue151 if self.job_pattern and not self.job_pattern.search(job_name):152 continue153 config_path = Path(job["config"])154 config = yaml.safe_load(config_path.read_text(encoding="utf-8")) or {}155 max_episode = int(config["max_train_episodes"])156 first_episode = int(max_episode * (1.0 - self.fraction)) + 1157 planned_episodes = list(range(first_episode, max_episode + 1))158 models_dir = Path(job["run_dir"]) / "models"159 models = []160 missing = []161 for episode in planned_episodes:162 checkpoint = (163 models_dir / "actor.pt"164 if episode == max_episode165 else models_dir / f"actor_ep{episode:06d}.pt"166 )167 if checkpoint.is_file():168 models.append(169 {170 "name": f"ep{episode:06d}",171 "checkpoint": str(checkpoint),172 "overrides": {},173 }174 )175 else:176 missing.append(episode)177 task_id = f"{stage}:{job_name}"178 task = self.state["tasks"].setdefault(task_id, {})179 task.update(180 task_id=task_id,181 stage=stage,182 job_name=job_name,183 config_path=str(config_path),184 run_dir=str(job["run_dir"]),185 first_episode=first_episode,186 last_episode=max_episode,187 planned_checkpoints=len(planned_episodes),188 available_checkpoints=len(models),189 missing_checkpoint_episodes=missing,190 planned_units=len(planned_episodes) * self.eval_episodes,191 submitted_units=len(models) * self.eval_episodes,192 models=models,193 )194 task.setdefault("state", "pending")195 task.setdefault("completed_units", 0)196 discovered.append(task_id)197 self.state["state"] = "running"198 self._write_state()199 return discovered200201 def _prepare(self, task):202 task_dir = self.root / task["stage"] / task["job_name"]203 config = yaml.safe_load(204 Path(task["config_path"]).read_text(encoding="utf-8")205 ) or {}206 progress_path = task_dir / "progress.json"207 config.update(208 cuda=self.cuda,209 use_eval=True,210 model_dir=None,211 n_eval_rollout_threads=1,212 eval_episodes=self.eval_episodes,213 eval_seed=self.eval_seed,214 eval_reset_mode="full_random",215 eval_deterministic=True,216 progress_file=str(progress_path),217 )218 config_path = task_dir / "config.yaml"219 manifest_path = task_dir / "manifest.yaml"220 _atomic_yaml(config_path, config)221 _atomic_yaml(222 manifest_path,223 {224 "name": f"final-window-{task['stage']}-{task['job_name']}",225 "output_root": str(task_dir / "results"),226 "models": task["models"],227 },228 )229 return task_dir, config_path, manifest_path, progress_path230231 def _update_progress(self, task_id, log_path, progress_path):232 try:233 started_models = log_path.read_text(234 encoding="utf-8", errors="replace"235 ).count("Evaluating ")236 except OSError:237 started_models = 0238 current_episodes = 0239 progress = _load_json(progress_path)240 if progress:241 current_episodes = min(242 self.eval_episodes,243 max(0, int(progress.get("completed_episodes", 0))),244 )245 completed_units = (246 max(0, started_models - 1) * self.eval_episodes + current_episodes247 )248 with self.lock:249 task = self.state["tasks"][task_id]250 task["completed_units"] = min(251 int(task["submitted_units"]), completed_units252 )253 task["heartbeat_at"] = _now()254 self._write_state()255256 def _run_task(self, task_id):257 with self.lock:258 task = self.state["tasks"][task_id]259 task.update(state="running", started_at=_now())260 self._write_state()261 task = dict(task)262 task_dir, config_path, manifest_path, progress_path = self._prepare(task)263 log_path = task_dir / "eval.log"264 command = [265 sys.executable,266 str(REPO_ROOT / "scripts" / "gpu_thermal_guard.py"),267 "--gpu", str(self.cuda),268 "--pause-at", "93",269 "--resume-at", "85",270 "--poll-seconds", "2",271 "--safe-readings", "3",272 "--query-failures", "2",273 "--status-file", str(task_dir / "thermal_status.json"),274 "--",275 sys.executable,276 str(REPO_ROOT / "train" / "eval.py"),277 "--config_file", str(config_path),278 "--models_file", str(manifest_path),279 "--output_root", str(task_dir / "results"),280 "--progress_file", str(progress_path),281 ]282 with log_path.open("w", encoding="utf-8") as log:283 process = subprocess.Popen(284 command,285 cwd=REPO_ROOT,286 stdin=subprocess.DEVNULL,287 stdout=log,288 stderr=subprocess.STDOUT,289 env={**os.environ, "PYTHONUNBUFFERED": "1"},290 )291 with self.lock:292 self.state["tasks"][task_id]["pid"] = process.pid293 self._write_state()294 while process.poll() is None:295 self._update_progress(task_id, log_path, progress_path)296 time.sleep(10)297 return_code = process.returncode298299 summaries = sorted(300 (task_dir / "results").glob("*/summary.json"),301 key=lambda path: path.stat().st_mtime,302 )303 records = []304 summary_path = None305 if return_code == 0 and summaries:306 summary_path = summaries[-1]307 payload = _load_json(summary_path) or {}308 records = payload.get("models") or []309 ok = return_code == 0 and len(records) == int(task["available_checkpoints"])310 metrics = {}311 if records:312 for key in (313 "success", "task_completion_ratio", "sent_percentage",314 "episode_length", "episode_reward",315 ):316 values = [float(record[key]) for record in records if key in record]317 if values:318 metrics[f"{key}_mean"] = sum(values) / len(values)319 metrics[f"{key}_min"] = min(values)320 metrics[f"{key}_max"] = max(values)321 with self.lock:322 target = self.state["tasks"][task_id]323 target.update(324 state="completed" if ok else "failed",325 return_code=return_code,326 finished_at=_now(),327 completed_checkpoints=len(records),328 completed_units=len(records) * self.eval_episodes,329 summary_path=str(summary_path) if summary_path else None,330 metrics=metrics,331 )332 self._write_state()333 return ok334335 def run(self):336 task_ids = self._discover()337 pending = [338 task_id for task_id in task_ids339 if self.state["tasks"][task_id].get("state") != "completed"340 ]341 results = []342 with ThreadPoolExecutor(max_workers=self.max_workers) as executor:343 futures = [executor.submit(self._run_task, task_id) for task_id in pending]344 for future in as_completed(futures):345 results.append(future.result())346 with self.lock:347 failed = any(348 task.get("state") != "completed"349 for task in self.state["tasks"].values()350 )351 self.state.update(352 state="partial_failure" if failed else "completed",353 finished_at=_now(),354 )355 self._write_state()356 return 1 if failed else 0357358359def main(argv=None):360 parser = argparse.ArgumentParser()361 parser.add_argument("--output-dir", required=True)362 parser.add_argument("--stage", action="append", required=True, metavar="NAME=PATH")363 parser.add_argument("--recovery-stage", action="append", default=[], metavar="NAME=PATH")364 parser.add_argument("--cuda", type=int, default=1)365 parser.add_argument("--max-workers", type=int, default=12)366 parser.add_argument("--fraction", type=float, default=0.1)367 parser.add_argument("--eval-episodes", type=int, default=8)368 parser.add_argument("--eval-seed", type=int, default=30000)369 parser.add_argument("--job-pattern")370 args = parser.parse_args(argv)371 stages = _parse_named_paths(args.stage)372 recovery = _parse_named_paths(args.recovery_stage, multiple=True)373 return FinalWindowQueue(374 args.output_dir,375 stages,376 recovery,377 args.cuda,378 args.max_workers,379 args.fraction,380 args.eval_episodes,381 args.eval_seed,382 args.job_pattern,383 ).run()384385386if __name__ == "__main__":387 raise SystemExit(main())