Initialize new run w/replay buffer of old run

deep_quoridor/src/run_benchmarks_v2.py

@@ -0,0 +1,119 @@
+"""Run just the benchmark schedules from an existing run's config.yaml.
+
+Usage:
+    python deep_quoridor/src/run_benchmarks_v2.py <run_dir> [-o key=val ...]
+
+Spawns one process per `config.benchmarks` schedule and waits until Ctrl-C.
+Reuses `benchmarks.create_benchmark_processes` from the v2 package; does not
+train, run self-play, or generate AI reports.
+"""
+
+import argparse
+import multiprocessing as mp
+import os
+import time
+from pathlib import Path
+
+from v2 import benchmarks
+from v2.common import ShutdownSignal
+from v2.config import Config, load_user_config
+
+# Match train_v2.py: suppress wandb's "install weave" log spam.
+os.environ["WANDB_DISABLE_WEAVE"] = "true"
+
+
+def _derive_base_dir(run_dir: Path) -> str:
+    """Given a run dir laid out as `base_dir/runs/<run_id>/`, return `base_dir`.
+
+    The run-dir convention used by `train_v2.py`'s `load_config_and_setup_run`
+    places each run under `<base_dir>/runs/<run_id>/`, so the parent of `runs/`
+    is the base_dir the rest of the v2 machinery expects.
+    """
+    return str(run_dir.parent.parent)
+
+
+def _load_config(run_dir: Path, overrides: list[str] | None) -> Config:
+    """Load `<run_dir>/config.yaml` and build a Config without touching disk.
+
+    Uses `Config.from_user(..., create_dirs=False)` so the existing run directory
+    isn't disturbed and no `config.yaml` snapshot is rewritten. Raises
+    `FileNotFoundError` if the config file is missing.
+    """
+    config_yaml = run_dir / "config.yaml"
+    if not config_yaml.is_file():
+        raise FileNotFoundError(f"No config.yaml in {run_dir}")
+    user_config = load_user_config(str(config_yaml), overrides=overrides)
+    return Config.from_user(user_config, _derive_base_dir(run_dir), create_dirs=False)
+
+
+def _check_run_dir(run_dir: Path) -> None:
+    """Verify the run directory has the layout we need before spawning processes.
+
+    Aborts early on a missing `latest.yaml` so the benchmark processes don't enter
+    `LatestModel.wait_for_creation`'s blocking wait (no training is producing
+    models in this script).
+    """
+    if not run_dir.is_dir():
+        raise FileNotFoundError(f"Run directory not found: {run_dir}")
+    if not (run_dir / "config.yaml").is_file():
+        raise FileNotFoundError(f"No config.yaml in {run_dir}")
+    latest_yaml = run_dir / "models" / "latest.yaml"
+    if not latest_yaml.is_file():
+        raise FileNotFoundError(f"No models/latest.yaml in {run_dir}; the run has no trained model to benchmark.")
+
+
+def main(args) -> int:
+    """Entry point. Returns the exit code."""
+    run_dir = Path(args.run_dir).resolve()
+    _check_run_dir(run_dir)
+    config = _load_config(run_dir, args.overrides)
+
+    if not config.benchmarks:
+        print(f"No benchmarks configured in {run_dir}/config.yaml; nothing to run.")
+        return 0
+
+    mp.set_start_method("spawn", force=True)
+    ShutdownSignal.clear(config)
+
+    benchmark_processes = benchmarks.create_benchmark_processes(config)
+    for p in benchmark_processes:
+        p.start()
+    print(f"Started {len(benchmark_processes)} benchmark processes")
+
+    try:
+        b_count_prev = -1
+        while True:
+            b_count = sum(p.is_alive() for p in benchmark_processes)
+            if b_count != b_count_prev:
+                print(f"Waiting for {b_count} benchmark processes")
+                b_count_prev = b_count
+            if b_count == 0:
+                break
+            time.sleep(1)
+    except KeyboardInterrupt:
+        print("\nCaught Ctrl-C; signaling shutdown...")
+        ShutdownSignal.signal(config)
+        for p in benchmark_processes:
+            p.join()
+
+    ShutdownSignal.clear(config)
+    return 0
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Run just the benchmark schedules from an existing run's config.yaml.",
+    )
+    parser.add_argument(
+        "run_dir",
+        type=str,
+        help="Path to an existing run directory (e.g. /path/to/runs/<run_id>/).",
+    )
+    parser.add_argument(
+        "-o",
+        "--overrides",
+        nargs="*",
+        help="Configuration overrides (e.g., benchmarks.0.every=2 minutes).",
+    )
+    args = parser.parse_args()
+    raise SystemExit(main(args))

deep_quoridor/src/train_v2.py

@@ -10,6 +10,7 @@
     check_ai_available,
     load_config_and_setup_run,
     metrics_dir_for,
+    preload_symlinks,
     run_ai_reporter,
     run_selfplay_metrics,
     self_play,
@@ -85,7 +86,17 @@ def _selfplay_subprocess_env():
     # Make sure we don't have the shutdown signal from a previous run
     ShutdownSignal.clear(config)
 
-    train_process = mp.Process(target=train, args=[config])
+    games_already_trained_on = 0
+    if config.training.initial_replay_buffer is not None:
+        n_loaded = preload_symlinks(
+            source_run=Path(config.training.initial_replay_buffer.run),
+            dest_ready=config.paths.replay_buffers_ready,
+            buffer_size=config.training.replay_buffer_size,
+        )
+        print(f"Preloaded {n_loaded} games from {config.training.initial_replay_buffer.run}")
+        games_already_trained_on = n_loaded
+
+    train_process = mp.Process(target=train, args=[config, games_already_trained_on])
     train_process.start()
 
     benchmark_processes = benchmarks.create_benchmark_processes(config)
@@ -99,40 +110,41 @@ def _selfplay_subprocess_env():
     self_play_processes = []
     rust_subprocesses = []
 
-    if config.self_play.program == "rust":
-        # Spawn Rust self-play processes in continuous mode
-        selfplay_env = _selfplay_subprocess_env()
-        if selfplay_env is not None:
-            print(f"Self-play GPU env: ORT_DYLIB_PATH={selfplay_env['ORT_DYLIB_PATH']}")
-        metrics_dir = metrics_dir_for(config)
-        os.makedirs(metrics_dir, exist_ok=True)
-        config_file_path = str(config.paths.config_file)
-        for i in range(config.self_play.num_processes):
-            cmd = [
-                config.self_play.rust_selfplay_binary,
-                "--config",
-                config_file_path,
-                "--output-dir",
-                str(config.paths.replay_buffers_ready),
-                "--continuous",
-                "--latest-model-yaml",
-                str(config.paths.latest_model_yaml),
-                "--shutdown-file",
-                str(ShutdownSignal.file_path(config)),
-                "--metrics-dir",
-                metrics_dir,
-            ]
-            proc = subprocess.Popen(cmd, env=selfplay_env)
-            rust_subprocesses.append(proc)
-            print(f"Started Rust self-play process {proc.pid}")
-        selfplay_metrics_process = mp.Process(target=run_selfplay_metrics, args=[config])
-        selfplay_metrics_process.start()
-        self_play_processes.append(selfplay_metrics_process)
-    else:
-        for i in range(config.self_play.num_processes):
-            p = mp.Process(target=self_play, args=[config])
-            p.start()
-            self_play_processes.append(p)
+    if config.self_play.enabled:
+        if config.self_play.program == "rust":
+            # Spawn Rust self-play processes in continuous mode
+            selfplay_env = _selfplay_subprocess_env()
+            if selfplay_env is not None:
+                print(f"Self-play GPU env: ORT_DYLIB_PATH={selfplay_env['ORT_DYLIB_PATH']}")
+            metrics_dir = metrics_dir_for(config)
+            os.makedirs(metrics_dir, exist_ok=True)
+            config_file_path = str(config.paths.config_file)
+            for i in range(config.self_play.num_processes):
+                cmd = [
+                    config.self_play.rust_selfplay_binary,
+                    "--config",
+                    config_file_path,
+                    "--output-dir",
+                    str(config.paths.replay_buffers_ready),
+                    "--continuous",
+                    "--latest-model-yaml",
+                    str(config.paths.latest_model_yaml),
+                    "--shutdown-file",
+                    str(ShutdownSignal.file_path(config)),
+                    "--metrics-dir",
+                    metrics_dir,
+                ]
+                proc = subprocess.Popen(cmd, env=selfplay_env)
+                rust_subprocesses.append(proc)
+                print(f"Started Rust self-play process {proc.pid}")
+            selfplay_metrics_process = mp.Process(target=run_selfplay_metrics, args=[config])
+            selfplay_metrics_process.start()
+            self_play_processes.append(selfplay_metrics_process)
+        else:
+            for i in range(config.self_play.num_processes):
+                p = mp.Process(target=self_play, args=[config])
+                p.start()
+                self_play_processes.append(p)
 
     train_process.join()
     ShutdownSignal.signal(config)

deep_quoridor/src/v2/__init__.py

@@ -15,13 +15,15 @@
     "generate_on_demand_report",
     "metrics_dir_for",
     "run_selfplay_metrics",
+    "preload_symlinks",
 ]
 
 from v2.ai_report import check_ai_available, generate_on_demand_report, run_ai_reporter
 from v2.benchmarks import create_benchmark_processes
 from v2.selfplay_metrics import metrics_dir_for, run_selfplay_metrics
 from v2.common import JobTrigger, MockWandb, ShutdownSignal, create_alphazero, upload_model
 from v2.config import load_config_and_setup_run
+from v2.offline_preload import preload_symlinks
 from v2.self_play import self_play
 from v2.trainer import train
 from v2.yaml_models import GameInfo, LatestModel

deep_quoridor/src/v2/common.py

@@ -1,10 +1,12 @@
 import re
 import time
 from abc import abstractmethod
+from pathlib import Path
 from typing import Any, Callable, Optional
 
 import wandb
 from agents.alphazero import AlphaZeroAgent, AlphaZeroParams
+from pydantic_yaml import parse_yaml_file_as
 from v2.config import AlphaZeroPlayConfig, AlphaZeroSelfPlayConfig, Config
 from v2.yaml_models import LatestModel
 
@@ -135,11 +137,15 @@ def alphazero_params_dict_from_config(
         im = config.training.initial_model
         if im.file:
             params_dict["model_filename"] = im.file
-        if im.wandb_alias:
+        elif im.wandb_alias:
             params_dict["wandb_alias"] = im.wandb_alias
             params_dict["wandb_project"] = im.wandb_project or (
                 config.wandb.project if config.wandb else "deep_quoridor"
             )
+        elif im.run:
+            latest_yaml = Path(im.run) / "models" / "latest.yaml"
+            latest = parse_yaml_file_as(LatestModel, latest_yaml)
+            params_dict["model_filename"] = latest.filename
 
     # Add network config
     if config.alphazero.network.type == "mlp":

deep_quoridor/src/v2/config.py

@@ -3,7 +3,7 @@
 from typing import Annotated, Literal, Optional, Union
 
 import yaml
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 
 class StrictBaseModel(BaseModel):
@@ -64,6 +64,7 @@ class AlphaZeroSelfPlayConfig(StrictBaseModel):
 
 
 class SelfPlayConfig(StrictBaseModel):
+    enabled: bool = True
     num_processes: int
     games_per_process: int
     # Leaf-parallel MCTS knobs (Rust self-play only).
@@ -83,13 +84,29 @@ class InitialModel(StrictBaseModel):
     file: Optional[str] = None
     wandb_project: Optional[str] = None
     wandb_alias: Optional[str] = None
+    run: Optional[str] = None
+
+    @model_validator(mode="after")
+    def at_most_one_source(self) -> "InitialModel":
+        sources = [
+            ("file", self.file),
+            ("wandb_alias", self.wandb_alias),
+            ("run", self.run),
+        ]
+        set_sources = [name for name, val in sources if val is not None]
+        if len(set_sources) > 1:
+            raise ValueError(f"At most one of file, wandb_alias, run may be set in initial_model; got: {set_sources}")
+        return self
+
+
+class InitialReplayBuffer(StrictBaseModel):
+    """Configures preloading the replay buffer from a previous run.
+
+    `run` points at a run directory (parent of `replay_buffers/`), mirroring
+    `InitialModel.run`. At preload time the loader reads `<run>/replay_buffers/`.
+    """
 
-    @field_validator("wandb_alias")
-    @classmethod
-    def file_and_wandb_mutually_exclusive(cls, v, info):
-        if v is not None and info.data.get("file") is not None:
-            raise ValueError("Cannot specify both 'file' and 'wandb_alias' in initial_model")
-        return v
+    run: str
 
 
 class CosineWarmRestartsSchedulerConfig(StrictBaseModel):
@@ -113,6 +130,7 @@ class TrainingConfig(StrictBaseModel):
     save_onnx: bool = False
     finish_after: Optional[str] = None
     initial_model: Optional[InitialModel] = None
+    initial_replay_buffer: Optional[InitialReplayBuffer] = None
     lr_scheduler: Optional[LRSchedulerConfig] = None
 
 
@@ -196,6 +214,15 @@ def replace_datetime_placeholder(cls, v: str) -> str:
             return v.replace("$DATETIME", current_datetime)
         return v
 
+    @model_validator(mode="after")
+    def selfplay_off_requires_replay_buffer(self) -> "UserConfig":
+        if not self.self_play.enabled and self.training.initial_replay_buffer is None:
+            raise ValueError(
+                "When self_play.enabled is False, training.initial_replay_buffer must be set "
+                "(otherwise the trainer would hang forever waiting for games)."
+            )
+        return self
+
 
 class PathsConfig(StrictBaseModel):
     run_dir: Path
@@ -372,7 +399,7 @@ def load_config_and_setup_run(
     with config_filename.open(mode="w") as f:
         f.write(to_yaml_str_ordered(user_config))
 
-    use_rust = config.self_play.program == "rust"
+    use_rust = config.self_play.enabled and config.self_play.program == "rust"
     if use_rust:
         # Apply default Rust binary path if not specified in config
         if config.self_play.rust_selfplay_binary is None: