sweep.py
批量训练、checkpoint 与状态调度
原文件 · 1446 行 · 56658 字节
train/sweep.py
SHA-256: 694c8588f955f82dc91251e8b94d64abb4bd97bef25b9a9d0974b3211e11d254
1"""Run declarative training sweeps as concurrent isolated processes."""23import argparse4import copy5import csv6import json7import os8import re9import subprocess10import sys11import threading12from collections import defaultdict13from contextlib import contextmanager14from concurrent.futures import (15 FIRST_COMPLETED,16 ProcessPoolExecutor,17 ThreadPoolExecutor,18 as_completed,19 wait,20)21from datetime import datetime22from pathlib import Path2324import yaml252627REPO_ROOT = Path(__file__).resolve().parents[1]282930def _gpu_memory_total_mb(gpu):31 """Best-effort physical-memory query used only to cap new submissions."""32 try:33 result = subprocess.run(34 [35 "nvidia-smi",36 f"--id={int(gpu)}",37 "--query-gpu=memory.total",38 "--format=csv,noheader,nounits",39 ],40 capture_output=True,41 text=True,42 timeout=10,43 )44 if result.returncode == 0:45 return int(result.stdout.strip().splitlines()[0])46 except (OSError, ValueError, subprocess.SubprocessError):47 pass48 return None495051def _load_yaml(path):52 with Path(path).open("r", encoding="utf-8") as file:53 return yaml.safe_load(file) or {}545556def _dump_yaml(path, value):57 with Path(path).open("w", encoding="utf-8") as file:58 yaml.safe_dump(value, file, sort_keys=False, allow_unicode=True)596061def _write_json_atomic(path, payload):62 path = Path(path)63 path.parent.mkdir(parents=True, exist_ok=True)64 temporary = path.with_name(f"{path.name}.{os.getpid()}.tmp")65 with temporary.open("w", encoding="utf-8") as file:66 json.dump(payload, file, ensure_ascii=False, indent=2)67 temporary.replace(path)686970def _safe_name(value):71 value = str(value)72 if not value or any(char not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-" for char in value):73 raise ValueError(f"Unsafe name: {value!r}")74 return value757677def _next_run_dir(config, group_dir):78 run_name = (79 f"{_safe_name(config.get('algorithm_name', 'mappo'))}_"80 f"{_safe_name(config['experiment_name'])}"81 )82 return Path(group_dir) / run_name838485def _job_key(job):86 overrides = job.get("overrides", {}) or {}87 algorithm_name = job.get("algorithm_name", overrides.get("algorithm_name"))88 experiment_name = job["experiment_name"]89 if algorithm_name:90 return _safe_name(f"{algorithm_name}_{experiment_name}")91 return _safe_name(experiment_name)929394class SweepRunner:95 def __init__(self, manifest_path, run_dir, no_eval=False):96 self.manifest_path = Path(manifest_path).resolve()97 self.manifest = _load_yaml(self.manifest_path)98 self.run_dir = Path(run_dir).resolve()99 self.no_eval = no_eval100 self.lock = threading.Lock()101 self.status_path = self.run_dir / "status.json"102 self._eval_task_metadata = {}103 self._eval_task_states = {}104 self._stream_eval_executor = None105 self._stream_eval_results = []106 self._stream_eval_output_root = None107 self._stream_eval_config = None108 self._stream_eval_selectors = []109110 base_path = Path(self.manifest.get("base_config", "config.yaml"))111 if not base_path.is_absolute():112 base_path = self.manifest_path.parent / base_path113 self.base_config_path = base_path.resolve()114 self.base_config = _load_yaml(self.base_config_path)115 self.jobs = self.manifest.get("jobs", [])116 if not self.jobs:117 raise ValueError("Sweep manifest requires a non-empty jobs list.")118119 names = [_job_key(job) for job in self.jobs]120 if len(names) != len(set(names)):121 raise ValueError("algorithm_name/experiment_name job keys must be unique within a sweep.")122123 self.status = {124 "name": self.manifest.get("name", self.manifest_path.stem),125 "manifest": str(self.manifest_path),126 "base_config": str(self.base_config_path),127 "created_at": datetime.now().isoformat(timespec="seconds"),128 "state": "initializing",129 "jobs": {},130 "eval": {"state": "pending"},131 }132133 def _write_status(self):134 with self.lock:135 self._write_status_unlocked()136137 def _write_status_unlocked(self):138 temporary = self.status_path.with_suffix(".json.tmp")139 with temporary.open("w", encoding="utf-8") as file:140 json.dump(self.status, file, ensure_ascii=False, indent=2)141 temporary.replace(self.status_path)142143 def _update_job_status(self, name, values):144 with self.lock:145 self.status["jobs"][name].update(values)146 self._write_status_unlocked()147148 def _touch_eval_heartbeat(self):149 """Publish liveness while an eval subprocess/task is still running."""150 with self.lock:151 eval_status = self.status.setdefault("eval", {})152 self._refresh_eval_episode_progress_unlocked()153 eval_status["heartbeat_at"] = datetime.now().isoformat(timespec="seconds")154 eval_status["heartbeat_seq"] = int(eval_status.get("heartbeat_seq", 0)) + 1155 self._write_status_unlocked()156157 def _refresh_eval_episode_progress_unlocked(self):158 if not self._eval_task_metadata:159 return160 completed_episodes = 0161 total_episodes = 0162 active_tasks = []163 for task_id, metadata in self._eval_task_metadata.items():164 total = int(metadata["total_episodes"])165 total_episodes += total166 state = self._eval_task_states.get(task_id, "pending")167 completed = total if state == "completed" else 0168 progress_path = Path(metadata["progress_path"])169 if progress_path.is_file():170 try:171 with progress_path.open("r", encoding="utf-8") as file:172 payload = json.load(file)173 completed = max(174 completed,175 min(total, max(0, int(payload.get("completed_episodes", 0)))),176 )177 if state == "pending":178 state = str(payload.get("state", "running"))179 except (OSError, ValueError, json.JSONDecodeError):180 pass181 completed_episodes += completed182 if state == "running" and completed < total:183 active_tasks.append(184 {185 "job_name": metadata["job_name"],186 "checkpoint": metadata["checkpoint"],187 "completed_episodes": completed,188 "total_episodes": total,189 }190 )191 eval_status = self.status.setdefault("eval", {})192 eval_status["completed_episodes"] = completed_episodes193 eval_status["total_episodes"] = total_episodes194 eval_status["active_tasks"] = active_tasks195196 @contextmanager197 def _eval_heartbeat(self, interval_seconds):198 interval_seconds = max(1.0, float(interval_seconds))199 stopped = threading.Event()200201 def publish():202 while not stopped.wait(interval_seconds):203 self._touch_eval_heartbeat()204205 self._touch_eval_heartbeat()206 thread = threading.Thread(207 target=publish,208 name="sweep-eval-heartbeat",209 daemon=True,210 )211 thread.start()212 try:213 yield214 finally:215 stopped.set()216 thread.join(timeout=interval_seconds + 1.0)217 self._touch_eval_heartbeat()218219 def prepare(self):220 (self.run_dir / "configs").mkdir(parents=True, exist_ok=True)221 (self.run_dir / "train_logs").mkdir(parents=True, exist_ok=True)222 prepared = []223 for raw_job in self.jobs:224 job = copy.deepcopy(raw_job)225 job_name = _job_key(job)226 experiment_name = _safe_name(job["experiment_name"])227 config = copy.deepcopy(self.base_config)228229 # 应用 sweep 级 training_eval 默认值（可被 job overrides 覆盖）230 training_eval = self.manifest.get("training_eval", {}) or {}231 for key in (232 "use_eval",233 "n_eval_rollout_threads",234 "eval_interval",235 "eval_episodes",236 "eval_seed",237 "eval_reset_mode",238 "eval_deterministic",239 "training_eval_cuda",240 ):241 if key in training_eval:242 config[key] = training_eval[key]243244 overrides = job.get("overrides", {}) or {}245 config.update(overrides)246 if "algorithm_name" in job:247 config["algorithm_name"] = job["algorithm_name"]248 config["experiment_name"] = experiment_name249 group_name = self.manifest.get(250 "group_name", self.manifest.get("name", self.manifest_path.stem)251 )252 config.setdefault("group_name", group_name)253 config.setdefault("experiment_group", group_name)254 config["cuda"] = int(job["cuda"])255256 config_path = self.run_dir / "configs" / f"{job_name}.yaml"257 expected_run_dir = _next_run_dir(config, self.run_dir)258 config["run_dir"] = str(expected_run_dir)259 _dump_yaml(config_path, config)260 prepared_job = {261 **job,262 "job_name": job_name,263 "experiment_name": experiment_name,264 "cuda": int(job["cuda"]),265 "config": config,266 "config_path": config_path,267 "run_dir": expected_run_dir,268 "log_path": self.run_dir / "train_logs" / f"{job_name}.log",269 }270 prepared.append(prepared_job)271 self.status["jobs"][job_name] = {272 "state": "pending",273 "algorithm_name": config.get("algorithm_name"),274 "experiment_name": experiment_name,275 "cuda": int(job["cuda"]),276 "config": str(config_path),277 "run_dir": str(expected_run_dir),278 "log": str(prepared_job["log_path"]),279 }280 self.status["state"] = "training"281 self._write_status()282 return prepared283284 def _run_job(self, job):285 name = job["job_name"]286 command = [287 sys.executable,288 str(REPO_ROOT / "train" / "train.py"),289 "--config_file",290 str(job["config_path"]),291 ]292 environment = os.environ.copy()293 environment["PYTHONUNBUFFERED"] = "1"294 with job["log_path"].open("w", encoding="utf-8") as log_file:295 process = subprocess.Popen(296 command,297 cwd=REPO_ROOT,298 env=environment,299 stdin=subprocess.DEVNULL,300 stdout=log_file,301 stderr=subprocess.STDOUT,302 )303 self._update_job_status(304 name,305 {306 "state": "running",307 "pid": process.pid,308 "started_at": datetime.now().isoformat(timespec="seconds"),309 },310 )311 return_code = process.wait()312313 actor_path = job["run_dir"] / "models" / "actor.pt"314 if job["config"].get("use_pcfa", False):315 state = "completed" if return_code == 0 else "failed"316 else:317 state = "completed" if return_code == 0 and actor_path.is_file() else "failed"318 self._update_job_status(319 name,320 {321 "state": state,322 "return_code": return_code,323 "actor_checkpoint": str(actor_path),324 "finished_at": datetime.now().isoformat(timespec="seconds"),325 },326 )327 return state == "completed"328329 def _assign_cuda(self, job, cuda):330 """Rewrite a not-yet-started job config when the scheduler moves it to another GPU."""331332 cuda = int(cuda)333 if int(job["config"].get("cuda", job["cuda"])) == cuda:334 return335 job["cuda"] = cuda336 job["config"]["cuda"] = cuda337 _dump_yaml(job["config_path"], job["config"])338 with self.lock:339 self.status["jobs"][job["job_name"]]["cuda"] = cuda340 self._write_status_unlocked()341342 def _run_jobs_scheduled(self, jobs, on_job_finished=None):343 """Run jobs with optional per-GPU limits and GPU0 spillover.344345 Manifest schema used here:346347 scheduling:348 gpu_limits: {0: 2, 1: 3}349 spillover:350 from: 1351 to: 0352 after_jobs: [mappo_capcoal, mappo_delivcoal]353 to_limit: 1354 """355356 scheduling = self.manifest.get("scheduling", {}) or {}357 gpu_limits = {358 int(gpu): int(limit)359 for gpu, limit in (scheduling.get("gpu_limits", {}) or {}).items()360 }361 memory_settings = {362 int(gpu): values or {}363 for gpu, values in (scheduling.get("gpu_memory", {}) or {}).items()364 }365 memory_budgets = {}366 for gpu, values in memory_settings.items():367 total = _gpu_memory_total_mb(gpu)368 if total is None:369 total = int(values.get("assumed_total_mb", 0))370 reserve = int(values.get("reserve_mb", 3072))371 memory_budgets[gpu] = max(0, total - reserve)372 spillover = scheduling.get("spillover", {}) or {}373 spill_from = int(spillover.get("from", -1))374 spill_to = int(spillover.get("to", -1))375 spill_after = set(spillover.get("after_jobs", []) or [])376 spill_to_limit = int(spillover.get("to_limit", 0))377378 pending = list(jobs)379 completed = set()380 results = []381 running = {}382 active_by_gpu = defaultdict(int)383 active_memory_by_gpu = defaultdict(int)384385 def estimated_job_memory(job, gpu):386 settings = memory_settings.get(int(gpu), {})387 return int(388 job.get(389 "gpu_memory_mb",390 settings.get("default_job_mb", 1200),391 )392 )393394 def memory_available(job, gpu):395 if int(gpu) not in memory_budgets:396 return True397 estimate = estimated_job_memory(job, gpu)398 return active_memory_by_gpu[int(gpu)] + estimate <= memory_budgets[int(gpu)]399400 def base_limit(gpu):401 return gpu_limits.get(int(gpu), len(jobs))402403 def spillover_ready():404 return bool(spill_after) and spill_after.issubset(completed)405406 def choose_assignment(job):407 original_cuda = int(job["cuda"])408 if (409 active_by_gpu[original_cuda] < base_limit(original_cuda)410 and memory_available(job, original_cuda)411 ):412 return original_cuda413 if (414 spillover_ready()415 and original_cuda == spill_from416 and spill_to_limit > 0417 and active_by_gpu[spill_to] < spill_to_limit418 and memory_available(job, spill_to)419 ):420 return spill_to421 return None422423 with ThreadPoolExecutor(max_workers=len(jobs)) as executor:424 while pending or running:425 launched = False426 for job in list(pending):427 assigned_cuda = choose_assignment(job)428 if assigned_cuda is None:429 continue430 self._assign_cuda(job, assigned_cuda)431 active_by_gpu[assigned_cuda] += 1432 active_memory_by_gpu[assigned_cuda] += estimated_job_memory(433 job, assigned_cuda434 )435 future = executor.submit(self._run_job, job)436 running[future] = (job, assigned_cuda)437 pending.remove(job)438 launched = True439440 if not running:441 raise RuntimeError("Scheduler has pending jobs but no runnable slot.")442 if launched and pending:443 continue444445 done, _ = wait(running, return_when=FIRST_COMPLETED)446 for future in done:447 job, assigned_cuda = running.pop(future)448 active_by_gpu[assigned_cuda] -= 1449 active_memory_by_gpu[assigned_cuda] -= estimated_job_memory(450 job, assigned_cuda451 )452 ok = future.result()453 results.append(ok)454 completed.add(job["job_name"])455 if on_job_finished is not None:456 on_job_finished(job, ok)457458 return results459460 @staticmethod461 def _eval_overrides(config):462 keys = (463 "algorithm_name",464 "trainer_algorithm_name",465 "use_recurrent_policy",466 "use_naive_recurrent_policy",467 "recurrent_N",468 # Actor architecture must be reconstructed exactly before loading469 # a checkpoint. Omitting these silently rebuilt capacity-control470 # MAPPO checkpoints with the default 64-wide MLP at evaluation.471 "hidden_size",472 "layer_N",473 "use_feature_normalization",474 "use_orthogonal",475 "use_ReLU",476 "use_gat",477 "graph_model",478 "gat_mode",479 "gat_hidden_dim",480 "gat_out_dim",481 "hatt_hidden_dim",482 "hatt_num_heads",483 "hatt_dropout",484 "hatt_capability_scale",485 "hatt_shared_task_obs",486 "hatt_use_aa_attention",487 "hatt_use_coalition_context",488 "hatt_candidate_aa_mode",489 "hatt_aa_task_context",490 "qmix_agent_hidden_dim",491 "qmix_use_recurrent_agent",492 "qmix_transition_mode",493 "qmix_mixer_hidden_dim",494 "qmix_hypernet_hidden_dim",495 "qmix_double_q",496 "qmix_optimizer",497 "qmix_rmsprop_alpha",498 "env_name",499 "scenario_name",500 "reward_mode",501 "coalition_delivery_reward_mode",502 "explore_action_mode",503 "num_explore_actions",504 "explore_step_size",505 "explore_random_jitter",506 "sector_redecision_on_waypoint",507 "explore_event_redecision",508 "sector_min_leg_mobility_steps",509 "delivery_tasks_initially_discovered",510 "mask_explore_actions",511 "idle_time_cost_mode",512 "explore_observer_relative_geometry",513 "explore_coverage_features",514 "explore_coverage_grid_resolution",515 "reconfiguration_mode",516 "reconfiguration_interval",517 "reconfiguration_urgency_threshold",518 "reconfiguration_urgency_aggregation",519 "reconfiguration_switch_cost",520 "reconfiguration_gate_init",521 "reconfiguration_gate_floor",522 "reconfiguration_gate_parameterization",523 "reconfiguration_gate_override",524 "reconfiguration_decision_coef",525 "reconfiguration_context_features",526 "reconfiguration_context_dim",527 )528 return {key: config[key] for key in keys if key in config}529530 @staticmethod531 def _eval_suite_overrides(config):532 """Return environment settings that define an evaluation suite.533534 These values cannot live in a model manifest because ``train/eval.py``535 deliberately forbids a model from changing the shared scenario. They536 still need to be restored from training, however, or a checkpoint may537 be evaluated against the dimensions from the sweep's base config.538 """539540 keys = (541 "map_size",542 "num_uav",543 "num_ugv",544 "uav_perception_rg",545 "ugv_perception_rg",546 "target_num",547 "max_steps",548 "episode_time_limit",549 "episode_length",550 "cap_dim",551 "rewards_config_file",552 "gymrt2a_map_size",553 "gymrt2a_num_agents",554 "gymrt2a_objects",555 "gymrt2a_view_range",556 "gymrt2a_comm_range",557 "gymrt2a_share_intention",558 )559 return {key: config[key] for key in keys if key in config}560561 @classmethod562 def _common_eval_suite_overrides(cls, jobs):563 """Require all jobs in a shared evaluation suite to use one scenario."""564565 common = None566 common_job = None567 for job in jobs:568 current = cls._eval_suite_overrides(job["config"])569 if common is None:570 common = current571 common_job = job["job_name"]572 elif current != common:573 differing = sorted(574 key575 for key in set(common) | set(current)576 if common.get(key) != current.get(key)577 )578 raise ValueError(579 "Shared evaluation requires identical suite settings; "580 f"{common_job} and {job['job_name']} differ in {differing}"581 )582 return common or {}583584 def _run_eval(self, jobs):585 eval_settings = self.manifest.get("eval", {}) or {}586 if self.no_eval or not eval_settings.get("enabled", True):587 self.status["eval"] = {"state": "skipped"}588 self._write_status()589 return True590591 if eval_settings.get("all_checkpoints", False):592 return self._run_all_checkpoints_eval(jobs)593594 return self._run_final_actor_eval(jobs)595596 def _streaming_final_eval_enabled(self):597 """Whether selected final checkpoints should start as each job finishes."""598599 eval_settings = self.manifest.get("eval", {}) or {}600 return bool(601 not self.no_eval602 and eval_settings.get("enabled", True)603 and eval_settings.get("all_checkpoints", False)604 and eval_settings.get("checkpoint_selectors")605 and eval_settings.get("start_on_job_completion", True)606 )607608 @staticmethod609 def _selected_checkpoint_files(job, selectors):610 """Resolve a small semantic checkpoint set such as last/best/p90."""611612 models_dir = job["run_dir"] / "models"613 result = []614 for raw_selector in selectors:615 selector = str(raw_selector).lower()616 checkpoint = None617 if selector == "last":618 checkpoint = models_dir / "actor.pt"619 elif selector == "latest":620 checkpoint = models_dir / "actor_eval_latest.pt"621 elif selector == "best":622 checkpoint = models_dir / "actor_eval_best.pt"623 elif re.fullmatch(r"p\d{1,3}", selector):624 percentage = int(selector[1:])625 if not 0 <= percentage <= 100:626 raise ValueError(f"Invalid checkpoint selector: {raw_selector!r}")627 episode_files = []628 for path in models_dir.glob("actor_ep*.pt"):629 match = re.fullmatch(r"actor_ep(\d+)\.pt", path.name)630 if match:631 episode_files.append((int(match.group(1)), path))632 if episode_files:633 final_episode = int(job["config"].get("max_train_episodes", 0))634 if final_episode <= 0:635 final_episode = max(item[0] for item in episode_files)636 target = final_episode * percentage / 100.0637 checkpoint = min(638 episode_files, key=lambda item: (abs(item[0] - target), item[0])639 )[1]640 else:641 raise ValueError(f"Unknown checkpoint selector: {raw_selector!r}")642 if checkpoint is not None and checkpoint.is_file():643 result.append((selector, checkpoint))644 return result645646 @staticmethod647 def _write_csv(records, path):648 """Write a list of dict records to a CSV file."""649 if not records:650 Path(path).write_text("", encoding="utf-8")651 return652 fields = []653 for record in records:654 for key in record:655 if key not in fields:656 fields.append(key)657 with Path(path).open("w", newline="", encoding="utf-8") as f:658 writer = csv.DictWriter(f, fieldnames=fields)659 writer.writeheader()660 writer.writerows(records)661662 @staticmethod663 def _load_json_if_exists(path):664 path = Path(path)665 if not path.exists():666 return None667 with path.open("r", encoding="utf-8") as file:668 return json.load(file)669670 @staticmethod671 def _parse_training_eval_log(log_path):672 pattern = re.compile(673 r"eval:\s+episodes=(?P<episodes>\d+),\s+"674 r"success=(?P<success>[-+0-9.]+)%,\s+"675 r"reward=(?P<reward>[-+0-9.]+)[^,]*,\s+"676 r"length=(?P<length>[-+0-9.]+)[^,]*,\s+"677 r"successful_length=(?P<successful_length>[-+0-9.]+|nan|NA)"678 )679 records = []680 for line in Path(log_path).read_text(encoding="utf-8", errors="replace").splitlines():681 match = pattern.search(line)682 if not match:683 continue684 successful_length = match.group("successful_length")685 records.append(686 {687 "eval_episodes": int(match.group("episodes")),688 "success": float(match.group("success")) / 100.0,689 "episode_reward": float(match.group("reward")),690 "episode_length": float(match.group("length")),691 "successful_length": (692 None if successful_length in ("nan", "NA") else float(successful_length)693 ),694 }695 )696 return records697698 @staticmethod699 def _training_eval_summary(payload):700 if not payload:701 return {}702 summary = payload.get("summary")703 if isinstance(summary, dict):704 flattened = dict(summary)705 else:706 flattened = dict(payload)707 if "total_num_steps" in payload:708 flattened["total_num_steps"] = payload["total_num_steps"]709 if "env_episode" in payload:710 flattened["env_episode"] = payload["env_episode"]711 return flattened712713 def _write_post_summary(self, jobs):714 records = []715 for job in jobs:716 eval_records = self._parse_training_eval_log(job["log_path"])717 last_eval = eval_records[-1] if eval_records else {}718 best_eval = {}719 if eval_records:720 best_eval = max(721 eval_records,722 key=lambda item: (723 item.get("success", 0.0),724 -item.get("episode_length", float("inf")),725 item.get("episode_reward", float("-inf")),726 ),727 )728729 eval_best = self._training_eval_summary(730 self._load_json_if_exists(job["run_dir"] / "models" / "eval_best.json")731 )732 eval_latest = self._training_eval_summary(733 self._load_json_if_exists(job["run_dir"] / "models" / "eval_latest.json")734 )735736 record = {737 "job_name": job["job_name"],738 "algorithm_name": job["config"].get("algorithm_name"),739 "experiment_name": job["experiment_name"],740 "reward_mode": job["config"].get("reward_mode"),741 "graph_model": job["config"].get("graph_model"),742 "use_gat": job["config"].get("use_gat"),743 "cuda": job["config"].get("cuda"),744 "run_dir": str(job["run_dir"]),745 "train_log": str(job["log_path"]),746 }747 for prefix, source in (748 ("last_eval", last_eval),749 ("best_log_eval", best_eval),750 ("saved_best_eval", eval_best),751 ("saved_latest_eval", eval_latest),752 ):753 for key in (754 "success",755 "episode_reward",756 "episode_length",757 "successful_length",758 "successful_episode_length",759 "task_completion_ratio",760 "sent_percentage",761 "tasks_completed",762 "step",763 "episode",764 "total_num_steps",765 "env_episode",766 ):767 if key in source:768 record[f"{prefix}_{key}"] = source[key]769 records.append(record)770771 csv_path = self.run_dir / "post_summary.csv"772 json_path = self.run_dir / "post_summary.json"773 self._write_csv(records, csv_path)774 with json_path.open("w", encoding="utf-8") as file:775 json.dump(records, file, ensure_ascii=False, indent=2)776 self.status["post_summary"] = {777 "csv": str(csv_path),778 "json": str(json_path),779 "written_at": datetime.now().isoformat(timespec="seconds"),780 }781 self._write_status()782 return records783784 @staticmethod785 def _run_one_checkpoint_eval(task):786 (787 task_id, job_name, ckpt_name, ckpt_path, manifest_path,788 task_output_root, eval_config_path, progress_path,789 ) = task790 _write_json_atomic(791 progress_path,792 {793 "state": "running",794 "completed_episodes": 0,795 "total_episodes": int(_load_yaml(eval_config_path).get("eval_episodes", 0)),796 "updated_at": datetime.now().astimezone().isoformat(timespec="seconds"),797 },798 )799 command = [800 sys.executable,801 str(REPO_ROOT / "train" / "eval.py"),802 "--config_file",803 str(eval_config_path),804 "--models_file",805 str(manifest_path),806 "--output_root",807 str(task_output_root),808 "--progress_file",809 str(progress_path),810 ]811 eval_config = _load_yaml(eval_config_path)812 guard = eval_config.get("gpu_thermal_guard", {}) or {}813 if guard.get("enabled", False):814 gpu = int(eval_config.get("cuda", 0))815 command = [816 sys.executable,817 str(REPO_ROOT / "scripts" / "gpu_thermal_guard.py"),818 "--gpu", str(gpu),819 "--pause-at", str(int(guard.get("pause_at_c", 93))),820 "--resume-at", str(int(guard.get("resume_at_c", 85))),821 "--poll-seconds", str(float(guard.get("poll_seconds", 2))),822 "--safe-readings", str(int(guard.get("safe_readings", 3))),823 "--query-failures", str(int(guard.get("query_failures", 2))),824 "--status-file", str(Path(progress_path).with_name("thermal_status.json")),825 "--",826 *command,827 ]828 environment = os.environ.copy()829 environment["PYTHONUNBUFFERED"] = "1"830 proc = subprocess.run(831 command,832 cwd=REPO_ROOT,833 env=environment,834 capture_output=True,835 text=True,836 )837 output_dir = None838 for line in proc.stdout.splitlines():839 if "Evaluation suite written to" in line:840 output_dir = line.split("Evaluation suite written to")[-1].strip()841 return {842 "task_id": task_id,843 "job_name": job_name,844 "ckpt_name": ckpt_name,845 "ckpt_path": ckpt_path,846 "returncode": proc.returncode,847 "output_dir": output_dir,848 "stdout": proc.stdout,849 "stderr": proc.stderr,850 }851852 def _start_streaming_final_eval(self, jobs):853 """Plan final eval up front, then accept checkpoints job by job."""854855 eval_settings = self.manifest.get("eval", {}) or {}856 selectors = [857 str(selector).lower()858 for selector in eval_settings.get("checkpoint_selectors", [])859 ]860 if len(selectors) != len(set(selectors)):861 raise ValueError("eval.checkpoint_selectors must not contain duplicates")862863 eval_config = copy.deepcopy(self.base_config)864 eval_config.update(eval_settings.get("overrides", {}) or {})865 eval_config["cuda"] = int(eval_settings.get("cuda", 0))866 output_root = (867 self.run_dir868 / "eval_results"869 / f"all-checkpoints-{datetime.now().strftime('%Y%m%d-%H%M%S')}"870 )871 output_root.mkdir(parents=True, exist_ok=True)872873 self._stream_eval_results = []874 self._stream_eval_output_root = output_root875 self._stream_eval_config = eval_config876 self._stream_eval_selectors = selectors877 eval_episodes = int(eval_config.get("eval_episodes", 0))878 self._eval_task_metadata = {}879 self._eval_task_states = {}880 for job in jobs:881 for selector in selectors:882 task_id = f"{job['job_name']}:{selector}"883 task_output_root = output_root / job["job_name"] / selector884 self._eval_task_metadata[task_id] = {885 "job_name": job["job_name"],886 "checkpoint": selector,887 "progress_path": str(task_output_root / "progress.json"),888 "total_episodes": eval_episodes,889 }890 self._eval_task_states[task_id] = "pending"891892 heartbeat_interval = float(893 eval_settings.get("heartbeat_interval_seconds", 15)894 )895 self.status["eval"] = {896 "state": "running_all_checkpoints",897 "pid": None,898 "start_policy": "on_job_completion",899 "cuda": int(eval_settings.get("cuda", 0)),900 "output_root": str(output_root),901 "log": str(self.run_dir / "eval_all_checkpoints.log"),902 "started_at": datetime.now().isoformat(timespec="seconds"),903 "total_tasks": len(self._eval_task_metadata),904 "planned_tasks": len(self._eval_task_metadata),905 "submitted_tasks": 0,906 "completed_tasks": 0,907 "failed_tasks": 0,908 "heartbeat_interval_seconds": heartbeat_interval,909 }910 self._refresh_eval_episode_progress_unlocked()911 self._write_status()912 self._stream_eval_executor = ThreadPoolExecutor(913 max_workers=max(1, int(eval_settings.get("max_workers", 4)))914 )915 return heartbeat_interval916917 def _build_streaming_eval_tasks(self, job):918 """Resolve and materialize one completed job's selected eval tasks."""919920 eval_config = copy.deepcopy(self._stream_eval_config)921 eval_config.update(self._eval_suite_overrides(job["config"]))922 eval_config_path = self.run_dir / "configs" / f"eval_{job['job_name']}.yaml"923 _dump_yaml(eval_config_path, eval_config)924 selected = dict(925 self._selected_checkpoint_files(job, self._stream_eval_selectors)926 )927 tasks = []928 for selector in self._stream_eval_selectors:929 ckpt_file = selected.get(selector)930 if ckpt_file is None:931 continue932 task_output_root = (933 self._stream_eval_output_root / job["job_name"] / selector934 )935 task_output_root.mkdir(parents=True, exist_ok=True)936 manifest_name = f"{job['job_name']}_{selector}"937 manifest = {938 "name": manifest_name,939 "output_root": str(task_output_root),940 "models": [941 {942 "name": manifest_name,943 "checkpoint": str(ckpt_file),944 "overrides": self._eval_overrides(job["config"]),945 }946 ],947 }948 manifest_path = task_output_root / "manifest.yaml"949 _dump_yaml(manifest_path, manifest)950 tasks.append(951 (952 f"{job['job_name']}:{selector}",953 job["job_name"],954 selector,955 str(ckpt_file),956 str(manifest_path),957 str(task_output_root),958 str(eval_config_path),959 str(task_output_root / "progress.json"),960 )961 )962 return tasks963964 def _record_streaming_eval_result(self, result):965 with self.lock:966 self._stream_eval_results.append(result)967 self._eval_task_states[result["task_id"]] = (968 "completed" if result["returncode"] == 0 else "failed"969 )970 eval_status = self.status["eval"]971 eval_status["completed_tasks"] = len(self._stream_eval_results)972 eval_status["failed_tasks"] = sum(973 item["returncode"] != 0 for item in self._stream_eval_results974 )975 self._refresh_eval_episode_progress_unlocked()976 self._write_status_unlocked()977 status = "OK" if result["returncode"] == 0 else "FAILED"978 print(f"[{status}] {result['job_name']} {result['ckpt_name']}")979 if result["returncode"] != 0 and result.get("stderr"):980 print(result["stderr"][-500:])981982 def _streaming_eval_done(self, future, task):983 try:984 result = future.result()985 except Exception as error: # keep the planned unit visible as failed986 result = {987 "task_id": task[0],988 "job_name": task[1],989 "ckpt_name": task[2],990 "ckpt_path": task[3],991 "returncode": 1,992 "output_dir": None,993 "stdout": "",994 "stderr": repr(error),995 }996 self._record_streaming_eval_result(result)997998 def _submit_streaming_final_eval(self, job, training_ok):999 """Submit this job immediately; failed/missing units remain in denominator."""10001001 tasks = self._build_streaming_eval_tasks(job) if training_ok else []1002 tasks_by_id = {task[0]: task for task in tasks}1003 for selector in self._stream_eval_selectors:1004 task_id = f"{job['job_name']}:{selector}"1005 task = tasks_by_id.get(task_id)1006 if task is None:1007 reason = (1008 "training failed before final evaluation"1009 if not training_ok1010 else f"selected checkpoint {selector!r} was not found"1011 )1012 self._record_streaming_eval_result(1013 {1014 "task_id": task_id,1015 "job_name": job["job_name"],1016 "ckpt_name": selector,1017 "ckpt_path": None,1018 "returncode": 1,1019 "output_dir": None,1020 "stdout": "",1021 "stderr": reason,1022 }1023 )1024 continue1025 with self.lock:1026 self.status["eval"]["submitted_tasks"] += 11027 self._write_status_unlocked()1028 future = self._stream_eval_executor.submit(1029 self._run_one_checkpoint_eval, task1030 )1031 future.add_done_callback(1032 lambda completed, submitted_task=task: self._streaming_eval_done(1033 completed, submitted_task1034 )1035 )10361037 def _finalize_checkpoint_eval_results(self, results, output_root):1038 """Aggregate completed checkpoint tasks and publish their terminal state."""10391040 per_job = defaultdict(list)1041 for result in results:1042 if result["returncode"] != 0 or not result["output_dir"]:1043 continue1044 summary_path = Path(result["output_dir"]) / "summary.json"1045 if not summary_path.exists():1046 continue1047 with summary_path.open("r", encoding="utf-8") as file:1048 summary = json.load(file)1049 for record in summary.get("models", []):1050 per_job[result["job_name"]].append(1051 {1052 "checkpoint": result["ckpt_name"],1053 "checkpoint_path": result["ckpt_path"],1054 **record,1055 }1056 )10571058 best_records = []1059 for job_name, records in per_job.items():1060 records.sort(1061 key=lambda record: (1062 record["success"],1063 record["episode_reward"],1064 ),1065 reverse=True,1066 )1067 best = records[0]1068 best_records.append(best)1069 print(1070 f"Best checkpoint for {job_name}: {best['checkpoint']} "1071 f"(success={best['success']:.1%}, reward={best['episode_reward']:.2f})"1072 )10731074 aggregate_best_path = output_root / "aggregate_best.json"1075 with aggregate_best_path.open("w", encoding="utf-8") as file:1076 json.dump(best_records, file, ensure_ascii=False, indent=2)1077 aggregate_all_path = output_root / "aggregate_all.json"1078 with aggregate_all_path.open("w", encoding="utf-8") as file:1079 json.dump(dict(per_job), file, ensure_ascii=False, indent=2)1080 self._write_csv(best_records, self.run_dir / "eval_summary_best.csv")10811082 all_ok = bool(results) and all(1083 result["returncode"] == 0 for result in results1084 )1085 self.status["eval"].update(1086 {1087 "state": "completed" if all_ok else "partial_failure",1088 "return_code": 0 if all_ok else 1,1089 "aggregate_best": str(aggregate_best_path),1090 "aggregate_all": str(aggregate_all_path),1091 "finished_at": datetime.now().isoformat(timespec="seconds"),1092 }1093 )1094 self._write_status()1095 return all_ok10961097 def _finish_streaming_final_eval(self):1098 self._stream_eval_executor.shutdown(wait=True)1099 self._stream_eval_executor = None1100 return self._finalize_checkpoint_eval_results(1101 list(self._stream_eval_results), self._stream_eval_output_root1102 )11031104 def _run_all_checkpoints_eval(self, jobs):1105 eval_settings = self.manifest.get("eval", {}) or {}1106 eval_config = copy.deepcopy(self.base_config)1107 eval_config.update(eval_settings.get("overrides", {}) or {})1108 eval_config["cuda"] = int(eval_settings.get("cuda", 0))1109 max_workers = int(eval_settings.get("max_workers", 4))1110 output_root = (1111 self.run_dir1112 / "eval_results"1113 / f"all-checkpoints-{datetime.now().strftime('%Y%m%d-%H%M%S')}"1114 )1115 output_root.mkdir(parents=True, exist_ok=True)11161117 tasks = []1118 selectors = eval_settings.get("checkpoint_selectors")1119 for job in jobs:1120 # Each all-checkpoint task has its own process, so preserve that1121 # job's training scenario instead of assuming base_config matches.1122 job_eval_config = copy.deepcopy(eval_config)1123 job_eval_config.update(self._eval_suite_overrides(job["config"]))1124 eval_config_path = (1125 self.run_dir / "configs" / f"eval_{job['job_name']}.yaml"1126 )1127 _dump_yaml(eval_config_path, job_eval_config)1128 models_dir = job["run_dir"] / "models"1129 if selectors:1130 checkpoint_files = self._selected_checkpoint_files(job, selectors)1131 else:1132 checkpoint_files = [1133 (path.stem, path) for path in sorted(models_dir.glob("actor*.pt"))1134 ]1135 for ckpt_label, ckpt_file in checkpoint_files:1136 task_output_root = output_root / job["job_name"] / ckpt_file.stem1137 task_output_root.mkdir(parents=True, exist_ok=True)1138 manifest_name = f"{job['job_name']}_{ckpt_label}"1139 manifest = {1140 "name": manifest_name,1141 "output_root": str(task_output_root),1142 "models": [1143 {1144 "name": manifest_name,1145 "checkpoint": str(ckpt_file),1146 "overrides": self._eval_overrides(job["config"]),1147 }1148 ],1149 }1150 manifest_path = task_output_root / "manifest.yaml"1151 _dump_yaml(manifest_path, manifest)1152 tasks.append(1153 (1154 f"{job['job_name']}:{ckpt_label}",1155 job["job_name"],1156 ckpt_label,1157 str(ckpt_file),1158 str(manifest_path),1159 str(task_output_root),1160 str(eval_config_path),1161 str(task_output_root / "progress.json"),1162 )1163 )11641165 eval_log_path = self.run_dir / "eval_all_checkpoints.log"1166 self.status["eval"] = {1167 "state": "running_all_checkpoints",1168 "pid": None,1169 "config": str(eval_config_path),1170 "output_root": str(output_root),1171 "log": str(eval_log_path),1172 "started_at": datetime.now().isoformat(timespec="seconds"),1173 "total_tasks": len(tasks),1174 "completed_tasks": 0,1175 "heartbeat_interval_seconds": float(1176 eval_settings.get("heartbeat_interval_seconds", 15)1177 ),1178 }1179 eval_episodes = int(eval_config.get("eval_episodes", 0))1180 self._eval_task_metadata = {1181 task[0]: {1182 "job_name": task[1],1183 "checkpoint": task[2],1184 "progress_path": task[7],1185 "total_episodes": eval_episodes,1186 }1187 for task in tasks1188 }1189 self._eval_task_states = {task[0]: "pending" for task in tasks}1190 self._refresh_eval_episode_progress_unlocked()1191 self._write_status()11921193 results = []1194 heartbeat_interval = self.status["eval"]["heartbeat_interval_seconds"]1195 with self._eval_heartbeat(heartbeat_interval):1196 with ProcessPoolExecutor(max_workers=max_workers) as executor:1197 futures = {1198 executor.submit(self._run_one_checkpoint_eval, task): task1199 for task in tasks1200 }1201 for future in as_completed(futures):1202 result = future.result()1203 results.append(result)1204 with self.lock:1205 self._eval_task_states[result["task_id"]] = (1206 "completed" if result["returncode"] == 0 else "failed"1207 )1208 self.status["eval"]["completed_tasks"] = len(results)1209 self._refresh_eval_episode_progress_unlocked()1210 self._write_status_unlocked()1211 status = "OK" if result["returncode"] == 0 else "FAILED"1212 print(f"[{status}] {result['job_name']} {result['ckpt_name']}")1213 if result["returncode"] != 0:1214 print(result["stderr"][-500:])12151216 return self._finalize_checkpoint_eval_results(results, output_root)12171218 def _run_final_actor_eval(self, jobs):1219 eval_settings = self.manifest.get("eval", {}) or {}1220 eval_config = copy.deepcopy(self.base_config)1221 eval_config.update(eval_settings.get("overrides", {}) or {})1222 eval_config.update(self._common_eval_suite_overrides(jobs))1223 eval_config["cuda"] = int(eval_settings.get("cuda", 0))1224 eval_config_path = self.run_dir / "configs" / "eval.yaml"1225 _dump_yaml(eval_config_path, eval_config)12261227 models = []1228 for job in jobs:1229 models.append(1230 {1231 "name": job["job_name"],1232 "checkpoint": str(job["run_dir"] / "models"),1233 "overrides": self._eval_overrides(job["config"]),1234 }1235 )1236 eval_manifest = {1237 "name": f"{self.manifest.get('name', 'sweep')}-eval",1238 "output_root": str(self.run_dir / "eval_results"),1239 "models": models,1240 }1241 eval_manifest_path = self.run_dir / "configs" / "eval_models.yaml"1242 _dump_yaml(eval_manifest_path, eval_manifest)1243 eval_log_path = self.run_dir / "eval.log"1244 command = [1245 sys.executable,1246 str(REPO_ROOT / "train" / "eval.py"),1247 "--config_file",1248 str(eval_config_path),1249 "--models_file",1250 str(eval_manifest_path),1251 ]1252 environment = os.environ.copy()1253 environment["PYTHONUNBUFFERED"] = "1"1254 with eval_log_path.open("w", encoding="utf-8") as log_file:1255 process = subprocess.Popen(1256 command,1257 cwd=REPO_ROOT,1258 env=environment,1259 stdin=subprocess.DEVNULL,1260 stdout=subprocess.PIPE,1261 stderr=subprocess.STDOUT,1262 text=True,1263 )1264 output_lines = []1265 self.status["eval"] = {1266 "state": "running",1267 "pid": process.pid,1268 "config": str(eval_config_path),1269 "models": str(eval_manifest_path),1270 "log": str(eval_log_path),1271 "started_at": datetime.now().isoformat(timespec="seconds"),1272 "heartbeat_interval_seconds": float(1273 eval_settings.get("heartbeat_interval_seconds", 15)1274 ),1275 }1276 self._write_status()1277 with self._eval_heartbeat(1278 self.status["eval"]["heartbeat_interval_seconds"]1279 ):1280 for line in process.stdout:1281 output_lines.append(line)1282 log_file.write(line)1283 log_file.flush()1284 return_code = process.wait()12851286 output_dir = None1287 for line in output_lines:1288 if "Evaluation suite written to" in line:1289 output_dir = line.split("Evaluation suite written to")[-1].strip()1290 if output_dir:1291 src_csv = Path(output_dir) / "summary.csv"1292 dst_csv = self.run_dir / "eval_summary_final.csv"1293 if src_csv.exists():1294 dst_csv.write_bytes(src_csv.read_bytes())12951296 self.status["eval"].update(1297 {1298 "state": "completed" if return_code == 0 else "failed",1299 "return_code": return_code,1300 "finished_at": datetime.now().isoformat(timespec="seconds"),1301 }1302 )1303 self._write_status()1304 return return_code == 013051306 def run(self):1307 jobs = self.prepare()1308 streaming_eval = self._streaming_final_eval_enabled()1309 callback = None1310 heartbeat_interval = None1311 if streaming_eval:1312 heartbeat_interval = self._start_streaming_final_eval(jobs)1313 callback = self._submit_streaming_final_eval13141315 def run_training_jobs():1316 if self.manifest.get("scheduling"):1317 return self._run_jobs_scheduled(jobs, callback)1318 run_results = []1319 with ThreadPoolExecutor(max_workers=len(jobs)) as executor:1320 futures = {1321 executor.submit(self._run_job, job): job for job in jobs1322 }1323 for future in as_completed(futures):1324 job = futures[future]1325 ok = future.result()1326 run_results.append(ok)1327 if callback is not None:1328 callback(job, ok)1329 return run_results13301331 if streaming_eval:1332 with self._eval_heartbeat(heartbeat_interval):1333 results = run_training_jobs()1334 if all(results):1335 self.status["state"] = "evaluating"1336 self._write_status()1337 eval_ok = self._finish_streaming_final_eval()1338 else:1339 results = run_training_jobs()1340 eval_ok = None1341 if not all(results):1342 self.status["state"] = "training_failed"1343 if not streaming_eval:1344 self.status["eval"] = {"state": "blocked_by_training_failure"}1345 self.status["finished_at"] = datetime.now().isoformat(timespec="seconds")1346 self._write_status()1347 return 11348 if not streaming_eval:1349 self.status["state"] = "evaluating"1350 self._write_status()1351 eval_ok = self._run_eval(jobs)1352 self._write_post_summary(jobs)1353 self.status["state"] = "completed" if eval_ok else "eval_failed"1354 self.status["finished_at"] = datetime.now().isoformat(timespec="seconds")1355 self._write_status()1356 return 0 if eval_ok else 1135713581359def _default_run_dir(manifest_path, manifest):1360 root = Path(manifest.get("output_root", "runs"))1361 if not root.is_absolute():1362 root = REPO_ROOT / root13631364 base_path = Path(manifest.get("base_config", "config.yaml"))1365 if not base_path.is_absolute():1366 base_path = Path(manifest_path).resolve().parent / base_path1367 base_config = _load_yaml(base_path) if base_path.exists() else {}13681369 name = _safe_name(1370 manifest.get("group_name", manifest.get("name", Path(manifest_path).stem))1371 )1372 env_name = _safe_name(manifest.get("env_name", base_config.get("env_name", "MyEnv")))1373 timestamp = datetime.now().strftime("%Y%m%d-%H%M")1374 candidate = root / f"{timestamp}_{env_name}_{name}"1375 if not candidate.exists():1376 return candidate1377 for index in range(2, 1000):1378 deduped = root / f"{timestamp}_{env_name}_{name}-{index}"1379 if not deduped.exists():1380 return deduped1381 raise RuntimeError(f"Could not allocate a unique sweep directory for {candidate}")138213831384def main(args=None):1385 parser = argparse.ArgumentParser()1386 parser.add_argument("--sweep_file", required=True)1387 parser.add_argument("--run_dir")1388 parser.add_argument("--detach", action="store_true")1389 parser.add_argument("--dry_run", action="store_true")1390 parser.add_argument("--no_eval", action="store_true")1391 options = parser.parse_args(args)13921393 manifest_path = Path(options.sweep_file).resolve()1394 manifest = _load_yaml(manifest_path)1395 run_dir = (1396 Path(options.run_dir).resolve()1397 if options.run_dir1398 else _default_run_dir(manifest_path, manifest).resolve()1399 )1400 run_dir.mkdir(parents=True, exist_ok=bool(options.run_dir))14011402 if options.detach:1403 log_path = run_dir / "supervisor.log"1404 command = [1405 sys.executable,1406 str(Path(__file__).resolve()),1407 "--sweep_file",1408 str(manifest_path),1409 "--run_dir",1410 str(run_dir),1411 ]1412 if options.no_eval:1413 command.append("--no_eval")1414 with log_path.open("w", encoding="utf-8") as log_file:1415 process = subprocess.Popen(1416 command,1417 cwd=REPO_ROOT,1418 stdin=subprocess.DEVNULL,1419 stdout=log_file,1420 stderr=subprocess.STDOUT,1421 start_new_session=True,1422 )1423 (run_dir / "supervisor.pid").write_text(1424 f"{process.pid}\n", encoding="utf-8"1425 )1426 print(f"Detached sweep pid={process.pid} run_dir={run_dir}")1427 return 014281429 runner = SweepRunner(manifest_path, run_dir, no_eval=options.no_eval)1430 if options.dry_run:1431 jobs = runner.prepare()1432 for job in jobs:1433 print(1434 f"cuda={job['cuda']} job={job['job_name']} "1435 f"algorithm={job['config'].get('algorithm_name')} "1436 f"experiment={job['experiment_name']} "1437 f"run_dir={job['run_dir']}"1438 )1439 runner.status["state"] = "dry_run"1440 runner._write_status()1441 return 01442 return runner.run()144314441445if __name__ == "__main__":1446 raise SystemExit(main())