diff --git a/.gitignore b/.gitignore index a8c46f9..f151518 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,7 @@ dist/ wheels/ *.egg-info runs/ +/runs # Cython-generated sources src/coolrl_lost_cities/games/classic/game.c diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py b/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py index a585644..f4f5bd8 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/cli.py @@ -60,6 +60,8 @@ def _resolve_resume_path(config: DeepCFRConfig, resume: str | None) -> str | Non def _train_overrides_from_args(args: argparse.Namespace) -> dict[str, Any]: overrides: dict[str, Any] = {} + if args.no_save and (args.save_latest_only or args.save_iteration_interval is not None): + raise ValueError("--no-save cannot be combined with checkpoint save overrides") run_overrides = overrides.setdefault("run", {}) if args.iterations is not None: run_overrides["iterations"] = args.iterations @@ -85,7 +87,19 @@ def _train_overrides_from_args(args: argparse.Namespace) -> dict[str, Any]: if args.eval_games is not None: overrides.setdefault("evaluation", {})["games"] = args.eval_games if args.no_save: - overrides.setdefault("checkpoint", {})["save_every_iteration"] = False + checkpoint_overrides = overrides.setdefault("checkpoint", {}) + checkpoint_overrides["save_latest"] = False + checkpoint_overrides["save_every_iteration"] = False + checkpoint_overrides["save_iteration_interval"] = 0 + if args.save_latest_only: + checkpoint_overrides = overrides.setdefault("checkpoint", {}) + checkpoint_overrides["save_latest"] = True + checkpoint_overrides["save_latest_only"] = True + checkpoint_overrides["save_every_iteration"] = False + if args.save_iteration_interval is not None: + overrides.setdefault("checkpoint", {})["save_iteration_interval"] = ( + args.save_iteration_interval + ) if args.exact_resume: overrides.setdefault("checkpoint", {})["exact_resume"] = True return overrides @@ -206,6 +220,8 @@ def main(argv: list[str] | None = None) -> None: train.add_argument("--eval-games", type=int) train.add_argument("--seed", type=int) train.add_argument("--no-save", action="store_true") + train.add_argument("--save-latest-only", action="store_true") + train.add_argument("--save-iteration-interval", type=int) train.set_defaults(func=train_command) evaluate = subparsers.add_parser("eval") diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/config.py b/src/coolrl_lost_cities/games/classic/deep_cfr/config.py index e74f73e..04bce04 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/config.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/config.py @@ -214,6 +214,7 @@ class MemoryConfig(StrictModel): class CheckpointConfig(StrictModel): directory: str = "runs/deep_cfr/default" + save_latest: bool = True save_every_iteration: bool = True save_iteration_interval: int = 0 save_latest_only: bool = False diff --git a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py b/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py index 6402f2d..700c175 100644 --- a/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py +++ b/src/coolrl_lost_cities/games/classic/deep_cfr/trainer.py @@ -456,7 +456,8 @@ class DeepCFRTrainer: checkpoint_dir = self.run_dir if self._should_save_iteration(iteration) and not self.config.checkpoint.save_latest_only: self.save_checkpoint(checkpoint_dir / f"iteration_{iteration:05d}.pt", item) - self.save_checkpoint(checkpoint_dir / "latest.pt", item) + if self.config.checkpoint.save_latest: + self.save_checkpoint(checkpoint_dir / "latest.pt", item) def _start_run_logging(self) -> None: self.run_dir.mkdir(parents=True, exist_ok=True) diff --git a/tests/games/classic/test_deep_cfr_trainer.py b/tests/games/classic/test_deep_cfr_trainer.py index 2439592..1f16da1 100644 --- a/tests/games/classic/test_deep_cfr_trainer.py +++ b/tests/games/classic/test_deep_cfr_trainer.py @@ -80,6 +80,8 @@ def test_deep_cfr_train_cli_count_overrides_disable_duration_limits() -> None: "eval_every": None, "eval_games": None, "no_save": True, + "save_latest_only": False, + "save_iteration_interval": None, "exact_resume": False, }, )() @@ -94,6 +96,36 @@ def test_deep_cfr_train_cli_count_overrides_disable_duration_limits() -> None: assert overridden.traversal.resolved_traversals_per_player() == 1 assert overridden.traversal.resolved_num_workers() == 0 assert overridden.checkpoint.save_every_iteration is False + assert overridden.checkpoint.save_latest is False + + +def test_deep_cfr_train_cli_checkpoint_save_overrides() -> None: + args = type( + "Args", + (), + { + "iterations": None, + "max_hours": None, + "max_iterations": None, + "seed": None, + "traversals_per_iteration": None, + "num_workers": None, + "checkpoint_dir": None, + "eval_every": None, + "eval_games": None, + "no_save": False, + "save_latest_only": True, + "save_iteration_interval": 1, + "exact_resume": False, + }, + )() + + overridden = _with_overrides(DeepCFRConfig(), _train_overrides_from_args(args)) + + assert overridden.checkpoint.save_latest is True + assert overridden.checkpoint.save_latest_only is True + assert overridden.checkpoint.save_every_iteration is False + assert overridden.checkpoint.save_iteration_interval == 1 def test_deep_cfr_resume_latest_resolution_uses_config_checkpoint_dir(tmp_path) -> None: