"""Run declarative training sweeps as concurrent isolated processes.""" import argparse import copy import csv import json import os import re import subprocess import sys import threading from collections import defaultdict from contextlib import contextmanager from concurrent.futures import ( FIRST_COMPLETED, ProcessPoolExecutor, ThreadPoolExecutor, as_completed, wait, ) from datetime import datetime from pathlib import Path import yaml REPO_ROOT = Path(__file__).resolve().parents[1] def _gpu_memory_total_mb(gpu): """Best-effort physical-memory query used only to cap new submissions.""" try: result = subprocess.run( [ "nvidia-smi", f"--id={int(gpu)}", "--query-gpu=memory.total", "--format=csv,noheader,nounits", ], capture_output=True, text=True, timeout=10, ) if result.returncode == 0: return int(result.stdout.strip().splitlines()[0]) except (OSError, ValueError, subprocess.SubprocessError): pass return None def _load_yaml(path): with Path(path).open("r", encoding="utf-8") as file: return yaml.safe_load(file) or {} def _dump_yaml(path, value): with Path(path).open("w", encoding="utf-8") as file: yaml.safe_dump(value, file, sort_keys=False, allow_unicode=True) def _write_json_atomic(path, payload): path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name(f"{path.name}.{os.getpid()}.tmp") with temporary.open("w", encoding="utf-8") as file: json.dump(payload, file, ensure_ascii=False, indent=2) temporary.replace(path) def _safe_name(value): value = str(value) if not value or any(char not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-" for char in value): raise ValueError(f"Unsafe name: {value!r}") return value def _next_run_dir(config, group_dir): run_name = ( f"{_safe_name(config.get('algorithm_name', 'mappo'))}_" f"{_safe_name(config['experiment_name'])}" ) return Path(group_dir) / run_name def _job_key(job): overrides = job.get("overrides", {}) or {} algorithm_name = job.get("algorithm_name", overrides.get("algorithm_name")) experiment_name = job["experiment_name"] if algorithm_name: return _safe_name(f"{algorithm_name}_{experiment_name}") return _safe_name(experiment_name) class SweepRunner: def __init__(self, manifest_path, run_dir, no_eval=False): self.manifest_path = Path(manifest_path).resolve() self.manifest = _load_yaml(self.manifest_path) self.run_dir = Path(run_dir).resolve() self.no_eval = no_eval self.lock = threading.Lock() self.status_path = self.run_dir / "status.json" self._eval_task_metadata = {} self._eval_task_states = {} self._stream_eval_executor = None self._stream_eval_results = [] self._stream_eval_output_root = None self._stream_eval_config = None self._stream_eval_selectors = [] base_path = Path(self.manifest.get("base_config", "config.yaml")) if not base_path.is_absolute(): base_path = self.manifest_path.parent / base_path self.base_config_path = base_path.resolve() self.base_config = _load_yaml(self.base_config_path) self.jobs = self.manifest.get("jobs", []) if not self.jobs: raise ValueError("Sweep manifest requires a non-empty jobs list.") names = [_job_key(job) for job in self.jobs] if len(names) != len(set(names)): raise ValueError("algorithm_name/experiment_name job keys must be unique within a sweep.") self.status = { "name": self.manifest.get("name", self.manifest_path.stem), "manifest": str(self.manifest_path), "base_config": str(self.base_config_path), "created_at": datetime.now().isoformat(timespec="seconds"), "state": "initializing", "jobs": {}, "eval": {"state": "pending"}, } def _write_status(self): with self.lock: self._write_status_unlocked() def _write_status_unlocked(self): temporary = self.status_path.with_suffix(".json.tmp") with temporary.open("w", encoding="utf-8") as file: json.dump(self.status, file, ensure_ascii=False, indent=2) temporary.replace(self.status_path) def _update_job_status(self, name, values): with self.lock: self.status["jobs"][name].update(values) self._write_status_unlocked() def _touch_eval_heartbeat(self): """Publish liveness while an eval subprocess/task is still running.""" with self.lock: eval_status = self.status.setdefault("eval", {}) self._refresh_eval_episode_progress_unlocked() eval_status["heartbeat_at"] = datetime.now().isoformat(timespec="seconds") eval_status["heartbeat_seq"] = int(eval_status.get("heartbeat_seq", 0)) + 1 self._write_status_unlocked() def _refresh_eval_episode_progress_unlocked(self): if not self._eval_task_metadata: return completed_episodes = 0 total_episodes = 0 active_tasks = [] for task_id, metadata in self._eval_task_metadata.items(): total = int(metadata["total_episodes"]) total_episodes += total state = self._eval_task_states.get(task_id, "pending") completed = total if state == "completed" else 0 progress_path = Path(metadata["progress_path"]) if progress_path.is_file(): try: with progress_path.open("r", encoding="utf-8") as file: payload = json.load(file) completed = max( completed, min(total, max(0, int(payload.get("completed_episodes", 0)))), ) if state == "pending": state = str(payload.get("state", "running")) except (OSError, ValueError, json.JSONDecodeError): pass completed_episodes += completed if state == "running" and completed < total: active_tasks.append( { "job_name": metadata["job_name"], "checkpoint": metadata["checkpoint"], "completed_episodes": completed, "total_episodes": total, } ) eval_status = self.status.setdefault("eval", {}) eval_status["completed_episodes"] = completed_episodes eval_status["total_episodes"] = total_episodes eval_status["active_tasks"] = active_tasks @contextmanager def _eval_heartbeat(self, interval_seconds): interval_seconds = max(1.0, float(interval_seconds)) stopped = threading.Event() def publish(): while not stopped.wait(interval_seconds): self._touch_eval_heartbeat() self._touch_eval_heartbeat() thread = threading.Thread( target=publish, name="sweep-eval-heartbeat", daemon=True, ) thread.start() try: yield finally: stopped.set() thread.join(timeout=interval_seconds + 1.0) self._touch_eval_heartbeat() def prepare(self): (self.run_dir / "configs").mkdir(parents=True, exist_ok=True) (self.run_dir / "train_logs").mkdir(parents=True, exist_ok=True) prepared = [] for raw_job in self.jobs: job = copy.deepcopy(raw_job) job_name = _job_key(job) experiment_name = _safe_name(job["experiment_name"]) config = copy.deepcopy(self.base_config) # 应用 sweep 级 training_eval 默认值（可被 job overrides 覆盖） training_eval = self.manifest.get("training_eval", {}) or {} for key in ( "use_eval", "n_eval_rollout_threads", "eval_interval", "eval_episodes", "eval_seed", "eval_reset_mode", "eval_deterministic", "training_eval_cuda", ): if key in training_eval: config[key] = training_eval[key] overrides = job.get("overrides", {}) or {} config.update(overrides) if "algorithm_name" in job: config["algorithm_name"] = job["algorithm_name"] config["experiment_name"] = experiment_name group_name = self.manifest.get( "group_name", self.manifest.get("name", self.manifest_path.stem) ) config.setdefault("group_name", group_name) config.setdefault("experiment_group", group_name) config["cuda"] = int(job["cuda"]) config_path = self.run_dir / "configs" / f"{job_name}.yaml" expected_run_dir = _next_run_dir(config, self.run_dir) config["run_dir"] = str(expected_run_dir) _dump_yaml(config_path, config) prepared_job = { **job, "job_name": job_name, "experiment_name": experiment_name, "cuda": int(job["cuda"]), "config": config, "config_path": config_path, "run_dir": expected_run_dir, "log_path": self.run_dir / "train_logs" / f"{job_name}.log", } prepared.append(prepared_job) self.status["jobs"][job_name] = { "state": "pending", "algorithm_name": config.get("algorithm_name"), "experiment_name": experiment_name, "cuda": int(job["cuda"]), "config": str(config_path), "run_dir": str(expected_run_dir), "log": str(prepared_job["log_path"]), } self.status["state"] = "training" self._write_status() return prepared def _run_job(self, job): name = job["job_name"] command = [ sys.executable, str(REPO_ROOT / "train" / "train.py"), "--config_file", str(job["config_path"]), ] environment = os.environ.copy() environment["PYTHONUNBUFFERED"] = "1" with job["log_path"].open("w", encoding="utf-8") as log_file: process = subprocess.Popen( command, cwd=REPO_ROOT, env=environment, stdin=subprocess.DEVNULL, stdout=log_file, stderr=subprocess.STDOUT, ) self._update_job_status( name, { "state": "running", "pid": process.pid, "started_at": datetime.now().isoformat(timespec="seconds"), }, ) return_code = process.wait() actor_path = job["run_dir"] / "models" / "actor.pt" if job["config"].get("use_pcfa", False): state = "completed" if return_code == 0 else "failed" else: state = "completed" if return_code == 0 and actor_path.is_file() else "failed" self._update_job_status( name, { "state": state, "return_code": return_code, "actor_checkpoint": str(actor_path), "finished_at": datetime.now().isoformat(timespec="seconds"), }, ) return state == "completed" def _assign_cuda(self, job, cuda): """Rewrite a not-yet-started job config when the scheduler moves it to another GPU.""" cuda = int(cuda) if int(job["config"].get("cuda", job["cuda"])) == cuda: return job["cuda"] = cuda job["config"]["cuda"] = cuda _dump_yaml(job["config_path"], job["config"]) with self.lock: self.status["jobs"][job["job_name"]]["cuda"] = cuda self._write_status_unlocked() def _run_jobs_scheduled(self, jobs, on_job_finished=None): """Run jobs with optional per-GPU limits and GPU0 spillover. Manifest schema used here: scheduling: gpu_limits: {0: 2, 1: 3} spillover: from: 1 to: 0 after_jobs: [mappo_capcoal, mappo_delivcoal] to_limit: 1 """ scheduling = self.manifest.get("scheduling", {}) or {} gpu_limits = { int(gpu): int(limit) for gpu, limit in (scheduling.get("gpu_limits", {}) or {}).items() } memory_settings = { int(gpu): values or {} for gpu, values in (scheduling.get("gpu_memory", {}) or {}).items() } memory_budgets = {} for gpu, values in memory_settings.items(): total = _gpu_memory_total_mb(gpu) if total is None: total = int(values.get("assumed_total_mb", 0)) reserve = int(values.get("reserve_mb", 3072)) memory_budgets[gpu] = max(0, total - reserve) spillover = scheduling.get("spillover", {}) or {} spill_from = int(spillover.get("from", -1)) spill_to = int(spillover.get("to", -1)) spill_after = set(spillover.get("after_jobs", []) or []) spill_to_limit = int(spillover.get("to_limit", 0)) pending = list(jobs) completed = set() results = [] running = {} active_by_gpu = defaultdict(int) active_memory_by_gpu = defaultdict(int) def estimated_job_memory(job, gpu): settings = memory_settings.get(int(gpu), {}) return int( job.get( "gpu_memory_mb", settings.get("default_job_mb", 1200), ) ) def memory_available(job, gpu): if int(gpu) not in memory_budgets: return True estimate = estimated_job_memory(job, gpu) return active_memory_by_gpu[int(gpu)] + estimate <= memory_budgets[int(gpu)] def base_limit(gpu): return gpu_limits.get(int(gpu), len(jobs)) def spillover_ready(): return bool(spill_after) and spill_after.issubset(completed) def choose_assignment(job): original_cuda = int(job["cuda"]) if ( active_by_gpu[original_cuda] < base_limit(original_cuda) and memory_available(job, original_cuda) ): return original_cuda if ( spillover_ready() and original_cuda == spill_from and spill_to_limit > 0 and active_by_gpu[spill_to] < spill_to_limit and memory_available(job, spill_to) ): return spill_to return None with ThreadPoolExecutor(max_workers=len(jobs)) as executor: while pending or running: launched = False for job in list(pending): assigned_cuda = choose_assignment(job) if assigned_cuda is None: continue self._assign_cuda(job, assigned_cuda) active_by_gpu[assigned_cuda] += 1 active_memory_by_gpu[assigned_cuda] += estimated_job_memory( job, assigned_cuda ) future = executor.submit(self._run_job, job) running[future] = (job, assigned_cuda) pending.remove(job) launched = True if not running: raise RuntimeError("Scheduler has pending jobs but no runnable slot.") if launched and pending: continue done, _ = wait(running, return_when=FIRST_COMPLETED) for future in done: job, assigned_cuda = running.pop(future) active_by_gpu[assigned_cuda] -= 1 active_memory_by_gpu[assigned_cuda] -= estimated_job_memory( job, assigned_cuda ) ok = future.result() results.append(ok) completed.add(job["job_name"]) if on_job_finished is not None: on_job_finished(job, ok) return results @staticmethod def _eval_overrides(config): keys = ( "algorithm_name", "trainer_algorithm_name", "use_recurrent_policy", "use_naive_recurrent_policy", "recurrent_N", # Actor architecture must be reconstructed exactly before loading # a checkpoint. Omitting these silently rebuilt capacity-control # MAPPO checkpoints with the default 64-wide MLP at evaluation. "hidden_size", "layer_N", "use_feature_normalization", "use_orthogonal", "use_ReLU", "use_gat", "graph_model", "gat_mode", "gat_hidden_dim", "gat_out_dim", "hatt_hidden_dim", "hatt_num_heads", "hatt_dropout", "hatt_capability_scale", "hatt_shared_task_obs", "hatt_use_aa_attention", "hatt_use_coalition_context", "hatt_candidate_aa_mode", "hatt_aa_task_context", "qmix_agent_hidden_dim", "qmix_use_recurrent_agent", "qmix_transition_mode", "qmix_mixer_hidden_dim", "qmix_hypernet_hidden_dim", "qmix_double_q", "qmix_optimizer", "qmix_rmsprop_alpha", "env_name", "scenario_name", "reward_mode", "coalition_delivery_reward_mode", "explore_action_mode", "num_explore_actions", "explore_step_size", "explore_random_jitter", "sector_redecision_on_waypoint", "explore_event_redecision", "sector_min_leg_mobility_steps", "delivery_tasks_initially_discovered", "mask_explore_actions", "idle_time_cost_mode", "explore_observer_relative_geometry", "explore_coverage_features", "explore_coverage_grid_resolution", "reconfiguration_mode", "reconfiguration_interval", "reconfiguration_urgency_threshold", "reconfiguration_urgency_aggregation", "reconfiguration_switch_cost", "reconfiguration_gate_init", "reconfiguration_gate_floor", "reconfiguration_gate_parameterization", "reconfiguration_gate_override", "reconfiguration_decision_coef", "reconfiguration_context_features", "reconfiguration_context_dim", ) return {key: config[key] for key in keys if key in config} @staticmethod def _eval_suite_overrides(config): """Return environment settings that define an evaluation suite. These values cannot live in a model manifest because ``train/eval.py`` deliberately forbids a model from changing the shared scenario. They still need to be restored from training, however, or a checkpoint may be evaluated against the dimensions from the sweep's base config. """ keys = ( "map_size", "num_uav", "num_ugv", "uav_perception_rg", "ugv_perception_rg", "target_num", "max_steps", "episode_time_limit", "episode_length", "cap_dim", "rewards_config_file", "gymrt2a_map_size", "gymrt2a_num_agents", "gymrt2a_objects", "gymrt2a_view_range", "gymrt2a_comm_range", "gymrt2a_share_intention", ) return {key: config[key] for key in keys if key in config} @classmethod def _common_eval_suite_overrides(cls, jobs): """Require all jobs in a shared evaluation suite to use one scenario.""" common = None common_job = None for job in jobs: current = cls._eval_suite_overrides(job["config"]) if common is None: common = current common_job = job["job_name"] elif current != common: differing = sorted( key for key in set(common) | set(current) if common.get(key) != current.get(key) ) raise ValueError( "Shared evaluation requires identical suite settings; " f"{common_job} and {job['job_name']} differ in {differing}" ) return common or {} def _run_eval(self, jobs): eval_settings = self.manifest.get("eval", {}) or {} if self.no_eval or not eval_settings.get("enabled", True): self.status["eval"] = {"state": "skipped"} self._write_status() return True if eval_settings.get("all_checkpoints", False): return self._run_all_checkpoints_eval(jobs) return self._run_final_actor_eval(jobs) def _streaming_final_eval_enabled(self): """Whether selected final checkpoints should start as each job finishes.""" eval_settings = self.manifest.get("eval", {}) or {} return bool( not self.no_eval and eval_settings.get("enabled", True) and eval_settings.get("all_checkpoints", False) and eval_settings.get("checkpoint_selectors") and eval_settings.get("start_on_job_completion", True) ) @staticmethod def _selected_checkpoint_files(job, selectors): """Resolve a small semantic checkpoint set such as last/best/p90.""" models_dir = job["run_dir"] / "models" result = [] for raw_selector in selectors: selector = str(raw_selector).lower() checkpoint = None if selector == "last": checkpoint = models_dir / "actor.pt" elif selector == "latest": checkpoint = models_dir / "actor_eval_latest.pt" elif selector == "best": checkpoint = models_dir / "actor_eval_best.pt" elif re.fullmatch(r"p\d{1,3}", selector): percentage = int(selector[1:]) if not 0 <= percentage <= 100: raise ValueError(f"Invalid checkpoint selector: {raw_selector!r}") episode_files = [] for path in models_dir.glob("actor_ep*.pt"): match = re.fullmatch(r"actor_ep(\d+)\.pt", path.name) if match: episode_files.append((int(match.group(1)), path)) if episode_files: final_episode = int(job["config"].get("max_train_episodes", 0)) if final_episode <= 0: final_episode = max(item[0] for item in episode_files) target = final_episode * percentage / 100.0 checkpoint = min( episode_files, key=lambda item: (abs(item[0] - target), item[0]) )[1] else: raise ValueError(f"Unknown checkpoint selector: {raw_selector!r}") if checkpoint is not None and checkpoint.is_file(): result.append((selector, checkpoint)) return result @staticmethod def _write_csv(records, path): """Write a list of dict records to a CSV file.""" if not records: Path(path).write_text("", encoding="utf-8") return fields = [] for record in records: for key in record: if key not in fields: fields.append(key) with Path(path).open("w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=fields) writer.writeheader() writer.writerows(records) @staticmethod def _load_json_if_exists(path): path = Path(path) if not path.exists(): return None with path.open("r", encoding="utf-8") as file: return json.load(file) @staticmethod def _parse_training_eval_log(log_path): pattern = re.compile( r"eval:\s+episodes=(?P\d+),\s+" r"success=(?P[-+0-9.]+)%,\s+" r"reward=(?P[-+0-9.]+)[^,]*,\s+" r"length=(?P[-+0-9.]+)[^,]*,\s+" r"successful_length=(?P[-+0-9.]+|nan|NA)" ) records = [] for line in Path(log_path).read_text(encoding="utf-8", errors="replace").splitlines(): match = pattern.search(line) if not match: continue successful_length = match.group("successful_length") records.append( { "eval_episodes": int(match.group("episodes")), "success": float(match.group("success")) / 100.0, "episode_reward": float(match.group("reward")), "episode_length": float(match.group("length")), "successful_length": ( None if successful_length in ("nan", "NA") else float(successful_length) ), } ) return records @staticmethod def _training_eval_summary(payload): if not payload: return {} summary = payload.get("summary") if isinstance(summary, dict): flattened = dict(summary) else: flattened = dict(payload) if "total_num_steps" in payload: flattened["total_num_steps"] = payload["total_num_steps"] if "env_episode" in payload: flattened["env_episode"] = payload["env_episode"] return flattened def _write_post_summary(self, jobs): records = [] for job in jobs: eval_records = self._parse_training_eval_log(job["log_path"]) last_eval = eval_records[-1] if eval_records else {} best_eval = {} if eval_records: best_eval = max( eval_records, key=lambda item: ( item.get("success", 0.0), -item.get("episode_length", float("inf")), item.get("episode_reward", float("-inf")), ), ) eval_best = self._training_eval_summary( self._load_json_if_exists(job["run_dir"] / "models" / "eval_best.json") ) eval_latest = self._training_eval_summary( self._load_json_if_exists(job["run_dir"] / "models" / "eval_latest.json") ) record = { "job_name": job["job_name"], "algorithm_name": job["config"].get("algorithm_name"), "experiment_name": job["experiment_name"], "reward_mode": job["config"].get("reward_mode"), "graph_model": job["config"].get("graph_model"), "use_gat": job["config"].get("use_gat"), "cuda": job["config"].get("cuda"), "run_dir": str(job["run_dir"]), "train_log": str(job["log_path"]), } for prefix, source in ( ("last_eval", last_eval), ("best_log_eval", best_eval), ("saved_best_eval", eval_best), ("saved_latest_eval", eval_latest), ): for key in ( "success", "episode_reward", "episode_length", "successful_length", "successful_episode_length", "task_completion_ratio", "sent_percentage", "tasks_completed", "step", "episode", "total_num_steps", "env_episode", ): if key in source: record[f"{prefix}_{key}"] = source[key] records.append(record) csv_path = self.run_dir / "post_summary.csv" json_path = self.run_dir / "post_summary.json" self._write_csv(records, csv_path) with json_path.open("w", encoding="utf-8") as file: json.dump(records, file, ensure_ascii=False, indent=2) self.status["post_summary"] = { "csv": str(csv_path), "json": str(json_path), "written_at": datetime.now().isoformat(timespec="seconds"), } self._write_status() return records @staticmethod def _run_one_checkpoint_eval(task): ( task_id, job_name, ckpt_name, ckpt_path, manifest_path, task_output_root, eval_config_path, progress_path, ) = task _write_json_atomic( progress_path, { "state": "running", "completed_episodes": 0, "total_episodes": int(_load_yaml(eval_config_path).get("eval_episodes", 0)), "updated_at": datetime.now().astimezone().isoformat(timespec="seconds"), }, ) command = [ sys.executable, str(REPO_ROOT / "train" / "eval.py"), "--config_file", str(eval_config_path), "--models_file", str(manifest_path), "--output_root", str(task_output_root), "--progress_file", str(progress_path), ] eval_config = _load_yaml(eval_config_path) guard = eval_config.get("gpu_thermal_guard", {}) or {} if guard.get("enabled", False): gpu = int(eval_config.get("cuda", 0)) command = [ sys.executable, str(REPO_ROOT / "scripts" / "gpu_thermal_guard.py"), "--gpu", str(gpu), "--pause-at", str(int(guard.get("pause_at_c", 93))), "--resume-at", str(int(guard.get("resume_at_c", 85))), "--poll-seconds", str(float(guard.get("poll_seconds", 2))), "--safe-readings", str(int(guard.get("safe_readings", 3))), "--query-failures", str(int(guard.get("query_failures", 2))), "--status-file", str(Path(progress_path).with_name("thermal_status.json")), "--", *command, ] environment = os.environ.copy() environment["PYTHONUNBUFFERED"] = "1" proc = subprocess.run( command, cwd=REPO_ROOT, env=environment, capture_output=True, text=True, ) output_dir = None for line in proc.stdout.splitlines(): if "Evaluation suite written to" in line: output_dir = line.split("Evaluation suite written to")[-1].strip() return { "task_id": task_id, "job_name": job_name, "ckpt_name": ckpt_name, "ckpt_path": ckpt_path, "returncode": proc.returncode, "output_dir": output_dir, "stdout": proc.stdout, "stderr": proc.stderr, } def _start_streaming_final_eval(self, jobs): """Plan final eval up front, then accept checkpoints job by job.""" eval_settings = self.manifest.get("eval", {}) or {} selectors = [ str(selector).lower() for selector in eval_settings.get("checkpoint_selectors", []) ] if len(selectors) != len(set(selectors)): raise ValueError("eval.checkpoint_selectors must not contain duplicates") eval_config = copy.deepcopy(self.base_config) eval_config.update(eval_settings.get("overrides", {}) or {}) eval_config["cuda"] = int(eval_settings.get("cuda", 0)) output_root = ( self.run_dir / "eval_results" / f"all-checkpoints-{datetime.now().strftime('%Y%m%d-%H%M%S')}" ) output_root.mkdir(parents=True, exist_ok=True) self._stream_eval_results = [] self._stream_eval_output_root = output_root self._stream_eval_config = eval_config self._stream_eval_selectors = selectors eval_episodes = int(eval_config.get("eval_episodes", 0)) self._eval_task_metadata = {} self._eval_task_states = {} for job in jobs: for selector in selectors: task_id = f"{job['job_name']}:{selector}" task_output_root = output_root / job["job_name"] / selector self._eval_task_metadata[task_id] = { "job_name": job["job_name"], "checkpoint": selector, "progress_path": str(task_output_root / "progress.json"), "total_episodes": eval_episodes, } self._eval_task_states[task_id] = "pending" heartbeat_interval = float( eval_settings.get("heartbeat_interval_seconds", 15) ) self.status["eval"] = { "state": "running_all_checkpoints", "pid": None, "start_policy": "on_job_completion", "cuda": int(eval_settings.get("cuda", 0)), "output_root": str(output_root), "log": str(self.run_dir / "eval_all_checkpoints.log"), "started_at": datetime.now().isoformat(timespec="seconds"), "total_tasks": len(self._eval_task_metadata), "planned_tasks": len(self._eval_task_metadata), "submitted_tasks": 0, "completed_tasks": 0, "failed_tasks": 0, "heartbeat_interval_seconds": heartbeat_interval, } self._refresh_eval_episode_progress_unlocked() self._write_status() self._stream_eval_executor = ThreadPoolExecutor( max_workers=max(1, int(eval_settings.get("max_workers", 4))) ) return heartbeat_interval def _build_streaming_eval_tasks(self, job): """Resolve and materialize one completed job's selected eval tasks.""" eval_config = copy.deepcopy(self._stream_eval_config) eval_config.update(self._eval_suite_overrides(job["config"])) eval_config_path = self.run_dir / "configs" / f"eval_{job['job_name']}.yaml" _dump_yaml(eval_config_path, eval_config) selected = dict( self._selected_checkpoint_files(job, self._stream_eval_selectors) ) tasks = [] for selector in self._stream_eval_selectors: ckpt_file = selected.get(selector) if ckpt_file is None: continue task_output_root = ( self._stream_eval_output_root / job["job_name"] / selector ) task_output_root.mkdir(parents=True, exist_ok=True) manifest_name = f"{job['job_name']}_{selector}" manifest = { "name": manifest_name, "output_root": str(task_output_root), "models": [ { "name": manifest_name, "checkpoint": str(ckpt_file), "overrides": self._eval_overrides(job["config"]), } ], } manifest_path = task_output_root / "manifest.yaml" _dump_yaml(manifest_path, manifest) tasks.append( ( f"{job['job_name']}:{selector}", job["job_name"], selector, str(ckpt_file), str(manifest_path), str(task_output_root), str(eval_config_path), str(task_output_root / "progress.json"), ) ) return tasks def _record_streaming_eval_result(self, result): with self.lock: self._stream_eval_results.append(result) self._eval_task_states[result["task_id"]] = ( "completed" if result["returncode"] == 0 else "failed" ) eval_status = self.status["eval"] eval_status["completed_tasks"] = len(self._stream_eval_results) eval_status["failed_tasks"] = sum( item["returncode"] != 0 for item in self._stream_eval_results ) self._refresh_eval_episode_progress_unlocked() self._write_status_unlocked() status = "OK" if result["returncode"] == 0 else "FAILED" print(f"[{status}] {result['job_name']} {result['ckpt_name']}") if result["returncode"] != 0 and result.get("stderr"): print(result["stderr"][-500:]) def _streaming_eval_done(self, future, task): try: result = future.result() except Exception as error: # keep the planned unit visible as failed result = { "task_id": task[0], "job_name": task[1], "ckpt_name": task[2], "ckpt_path": task[3], "returncode": 1, "output_dir": None, "stdout": "", "stderr": repr(error), } self._record_streaming_eval_result(result) def _submit_streaming_final_eval(self, job, training_ok): """Submit this job immediately; failed/missing units remain in denominator.""" tasks = self._build_streaming_eval_tasks(job) if training_ok else [] tasks_by_id = {task[0]: task for task in tasks} for selector in self._stream_eval_selectors: task_id = f"{job['job_name']}:{selector}" task = tasks_by_id.get(task_id) if task is None: reason = ( "training failed before final evaluation" if not training_ok else f"selected checkpoint {selector!r} was not found" ) self._record_streaming_eval_result( { "task_id": task_id, "job_name": job["job_name"], "ckpt_name": selector, "ckpt_path": None, "returncode": 1, "output_dir": None, "stdout": "", "stderr": reason, } ) continue with self.lock: self.status["eval"]["submitted_tasks"] += 1 self._write_status_unlocked() future = self._stream_eval_executor.submit( self._run_one_checkpoint_eval, task ) future.add_done_callback( lambda completed, submitted_task=task: self._streaming_eval_done( completed, submitted_task ) ) def _finalize_checkpoint_eval_results(self, results, output_root): """Aggregate completed checkpoint tasks and publish their terminal state.""" per_job = defaultdict(list) for result in results: if result["returncode"] != 0 or not result["output_dir"]: continue summary_path = Path(result["output_dir"]) / "summary.json" if not summary_path.exists(): continue with summary_path.open("r", encoding="utf-8") as file: summary = json.load(file) for record in summary.get("models", []): per_job[result["job_name"]].append( { "checkpoint": result["ckpt_name"], "checkpoint_path": result["ckpt_path"], **record, } ) best_records = [] for job_name, records in per_job.items(): records.sort( key=lambda record: ( record["success"], record["episode_reward"], ), reverse=True, ) best = records[0] best_records.append(best) print( f"Best checkpoint for {job_name}: {best['checkpoint']} " f"(success={best['success']:.1%}, reward={best['episode_reward']:.2f})" ) aggregate_best_path = output_root / "aggregate_best.json" with aggregate_best_path.open("w", encoding="utf-8") as file: json.dump(best_records, file, ensure_ascii=False, indent=2) aggregate_all_path = output_root / "aggregate_all.json" with aggregate_all_path.open("w", encoding="utf-8") as file: json.dump(dict(per_job), file, ensure_ascii=False, indent=2) self._write_csv(best_records, self.run_dir / "eval_summary_best.csv") all_ok = bool(results) and all( result["returncode"] == 0 for result in results ) self.status["eval"].update( { "state": "completed" if all_ok else "partial_failure", "return_code": 0 if all_ok else 1, "aggregate_best": str(aggregate_best_path), "aggregate_all": str(aggregate_all_path), "finished_at": datetime.now().isoformat(timespec="seconds"), } ) self._write_status() return all_ok def _finish_streaming_final_eval(self): self._stream_eval_executor.shutdown(wait=True) self._stream_eval_executor = None return self._finalize_checkpoint_eval_results( list(self._stream_eval_results), self._stream_eval_output_root ) def _run_all_checkpoints_eval(self, jobs): eval_settings = self.manifest.get("eval", {}) or {} eval_config = copy.deepcopy(self.base_config) eval_config.update(eval_settings.get("overrides", {}) or {}) eval_config["cuda"] = int(eval_settings.get("cuda", 0)) max_workers = int(eval_settings.get("max_workers", 4)) output_root = ( self.run_dir / "eval_results" / f"all-checkpoints-{datetime.now().strftime('%Y%m%d-%H%M%S')}" ) output_root.mkdir(parents=True, exist_ok=True) tasks = [] selectors = eval_settings.get("checkpoint_selectors") for job in jobs: # Each all-checkpoint task has its own process, so preserve that # job's training scenario instead of assuming base_config matches. job_eval_config = copy.deepcopy(eval_config) job_eval_config.update(self._eval_suite_overrides(job["config"])) eval_config_path = ( self.run_dir / "configs" / f"eval_{job['job_name']}.yaml" ) _dump_yaml(eval_config_path, job_eval_config) models_dir = job["run_dir"] / "models" if selectors: checkpoint_files = self._selected_checkpoint_files(job, selectors) else: checkpoint_files = [ (path.stem, path) for path in sorted(models_dir.glob("actor*.pt")) ] for ckpt_label, ckpt_file in checkpoint_files: task_output_root = output_root / job["job_name"] / ckpt_file.stem task_output_root.mkdir(parents=True, exist_ok=True) manifest_name = f"{job['job_name']}_{ckpt_label}" manifest = { "name": manifest_name, "output_root": str(task_output_root), "models": [ { "name": manifest_name, "checkpoint": str(ckpt_file), "overrides": self._eval_overrides(job["config"]), } ], } manifest_path = task_output_root / "manifest.yaml" _dump_yaml(manifest_path, manifest) tasks.append( ( f"{job['job_name']}:{ckpt_label}", job["job_name"], ckpt_label, str(ckpt_file), str(manifest_path), str(task_output_root), str(eval_config_path), str(task_output_root / "progress.json"), ) ) eval_log_path = self.run_dir / "eval_all_checkpoints.log" self.status["eval"] = { "state": "running_all_checkpoints", "pid": None, "config": str(eval_config_path), "output_root": str(output_root), "log": str(eval_log_path), "started_at": datetime.now().isoformat(timespec="seconds"), "total_tasks": len(tasks), "completed_tasks": 0, "heartbeat_interval_seconds": float( eval_settings.get("heartbeat_interval_seconds", 15) ), } eval_episodes = int(eval_config.get("eval_episodes", 0)) self._eval_task_metadata = { task[0]: { "job_name": task[1], "checkpoint": task[2], "progress_path": task[7], "total_episodes": eval_episodes, } for task in tasks } self._eval_task_states = {task[0]: "pending" for task in tasks} self._refresh_eval_episode_progress_unlocked() self._write_status() results = [] heartbeat_interval = self.status["eval"]["heartbeat_interval_seconds"] with self._eval_heartbeat(heartbeat_interval): with ProcessPoolExecutor(max_workers=max_workers) as executor: futures = { executor.submit(self._run_one_checkpoint_eval, task): task for task in tasks } for future in as_completed(futures): result = future.result() results.append(result) with self.lock: self._eval_task_states[result["task_id"]] = ( "completed" if result["returncode"] == 0 else "failed" ) self.status["eval"]["completed_tasks"] = len(results) self._refresh_eval_episode_progress_unlocked() self._write_status_unlocked() status = "OK" if result["returncode"] == 0 else "FAILED" print(f"[{status}] {result['job_name']} {result['ckpt_name']}") if result["returncode"] != 0: print(result["stderr"][-500:]) return self._finalize_checkpoint_eval_results(results, output_root) def _run_final_actor_eval(self, jobs): eval_settings = self.manifest.get("eval", {}) or {} eval_config = copy.deepcopy(self.base_config) eval_config.update(eval_settings.get("overrides", {}) or {}) eval_config.update(self._common_eval_suite_overrides(jobs)) eval_config["cuda"] = int(eval_settings.get("cuda", 0)) eval_config_path = self.run_dir / "configs" / "eval.yaml" _dump_yaml(eval_config_path, eval_config) models = [] for job in jobs: models.append( { "name": job["job_name"], "checkpoint": str(job["run_dir"] / "models"), "overrides": self._eval_overrides(job["config"]), } ) eval_manifest = { "name": f"{self.manifest.get('name', 'sweep')}-eval", "output_root": str(self.run_dir / "eval_results"), "models": models, } eval_manifest_path = self.run_dir / "configs" / "eval_models.yaml" _dump_yaml(eval_manifest_path, eval_manifest) eval_log_path = self.run_dir / "eval.log" command = [ sys.executable, str(REPO_ROOT / "train" / "eval.py"), "--config_file", str(eval_config_path), "--models_file", str(eval_manifest_path), ] environment = os.environ.copy() environment["PYTHONUNBUFFERED"] = "1" with eval_log_path.open("w", encoding="utf-8") as log_file: process = subprocess.Popen( command, cwd=REPO_ROOT, env=environment, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, ) output_lines = [] self.status["eval"] = { "state": "running", "pid": process.pid, "config": str(eval_config_path), "models": str(eval_manifest_path), "log": str(eval_log_path), "started_at": datetime.now().isoformat(timespec="seconds"), "heartbeat_interval_seconds": float( eval_settings.get("heartbeat_interval_seconds", 15) ), } self._write_status() with self._eval_heartbeat( self.status["eval"]["heartbeat_interval_seconds"] ): for line in process.stdout: output_lines.append(line) log_file.write(line) log_file.flush() return_code = process.wait() output_dir = None for line in output_lines: if "Evaluation suite written to" in line: output_dir = line.split("Evaluation suite written to")[-1].strip() if output_dir: src_csv = Path(output_dir) / "summary.csv" dst_csv = self.run_dir / "eval_summary_final.csv" if src_csv.exists(): dst_csv.write_bytes(src_csv.read_bytes()) self.status["eval"].update( { "state": "completed" if return_code == 0 else "failed", "return_code": return_code, "finished_at": datetime.now().isoformat(timespec="seconds"), } ) self._write_status() return return_code == 0 def run(self): jobs = self.prepare() streaming_eval = self._streaming_final_eval_enabled() callback = None heartbeat_interval = None if streaming_eval: heartbeat_interval = self._start_streaming_final_eval(jobs) callback = self._submit_streaming_final_eval def run_training_jobs(): if self.manifest.get("scheduling"): return self._run_jobs_scheduled(jobs, callback) run_results = [] with ThreadPoolExecutor(max_workers=len(jobs)) as executor: futures = { executor.submit(self._run_job, job): job for job in jobs } for future in as_completed(futures): job = futures[future] ok = future.result() run_results.append(ok) if callback is not None: callback(job, ok) return run_results if streaming_eval: with self._eval_heartbeat(heartbeat_interval): results = run_training_jobs() if all(results): self.status["state"] = "evaluating" self._write_status() eval_ok = self._finish_streaming_final_eval() else: results = run_training_jobs() eval_ok = None if not all(results): self.status["state"] = "training_failed" if not streaming_eval: self.status["eval"] = {"state": "blocked_by_training_failure"} self.status["finished_at"] = datetime.now().isoformat(timespec="seconds") self._write_status() return 1 if not streaming_eval: self.status["state"] = "evaluating" self._write_status() eval_ok = self._run_eval(jobs) self._write_post_summary(jobs) self.status["state"] = "completed" if eval_ok else "eval_failed" self.status["finished_at"] = datetime.now().isoformat(timespec="seconds") self._write_status() return 0 if eval_ok else 1 def _default_run_dir(manifest_path, manifest): root = Path(manifest.get("output_root", "runs")) if not root.is_absolute(): root = REPO_ROOT / root base_path = Path(manifest.get("base_config", "config.yaml")) if not base_path.is_absolute(): base_path = Path(manifest_path).resolve().parent / base_path base_config = _load_yaml(base_path) if base_path.exists() else {} name = _safe_name( manifest.get("group_name", manifest.get("name", Path(manifest_path).stem)) ) env_name = _safe_name(manifest.get("env_name", base_config.get("env_name", "MyEnv"))) timestamp = datetime.now().strftime("%Y%m%d-%H%M") candidate = root / f"{timestamp}_{env_name}_{name}" if not candidate.exists(): return candidate for index in range(2, 1000): deduped = root / f"{timestamp}_{env_name}_{name}-{index}" if not deduped.exists(): return deduped raise RuntimeError(f"Could not allocate a unique sweep directory for {candidate}") def main(args=None): parser = argparse.ArgumentParser() parser.add_argument("--sweep_file", required=True) parser.add_argument("--run_dir") parser.add_argument("--detach", action="store_true") parser.add_argument("--dry_run", action="store_true") parser.add_argument("--no_eval", action="store_true") options = parser.parse_args(args) manifest_path = Path(options.sweep_file).resolve() manifest = _load_yaml(manifest_path) run_dir = ( Path(options.run_dir).resolve() if options.run_dir else _default_run_dir(manifest_path, manifest).resolve() ) run_dir.mkdir(parents=True, exist_ok=bool(options.run_dir)) if options.detach: log_path = run_dir / "supervisor.log" command = [ sys.executable, str(Path(__file__).resolve()), "--sweep_file", str(manifest_path), "--run_dir", str(run_dir), ] if options.no_eval: command.append("--no_eval") with log_path.open("w", encoding="utf-8") as log_file: process = subprocess.Popen( command, cwd=REPO_ROOT, stdin=subprocess.DEVNULL, stdout=log_file, stderr=subprocess.STDOUT, start_new_session=True, ) (run_dir / "supervisor.pid").write_text( f"{process.pid}\n", encoding="utf-8" ) print(f"Detached sweep pid={process.pid} run_dir={run_dir}") return 0 runner = SweepRunner(manifest_path, run_dir, no_eval=options.no_eval) if options.dry_run: jobs = runner.prepare() for job in jobs: print( f"cuda={job['cuda']} job={job['job_name']} " f"algorithm={job['config'].get('algorithm_name')} " f"experiment={job['experiment_name']} " f"run_dir={job['run_dir']}" ) runner.status["state"] = "dry_run" runner._write_status() return 0 return runner.run() if __name__ == "__main__": raise SystemExit(main())