diff --git a/kaievolve/cli.py b/kaievolve/cli.py index 4b77eed..4c3af53 100644 --- a/kaievolve/cli.py +++ b/kaievolve/cli.py @@ -85,6 +85,17 @@ def parse_args() -> argparse.Namespace: default=None, ) + parser.add_argument( + "--init-from", + default=None, + help=( + "Seed a fresh run from a previous run's best program. Pass a run " + "directory (e.g. bench_results//run_0); its best program replaces " + "the initial program, so you can pass just the evaluator: " + "`kai run --init-from evaluator.py -c config.yaml`." + ), + ) + parser.add_argument("--api-base", help="Base URL for the LLM API", default=None) parser.add_argument("--primary-model", help="Primary LLM model name", default=None) @@ -94,6 +105,30 @@ def parse_args() -> argparse.Namespace: return parser.parse_args() +def _resolve_best_program(run_dir): + """For --init-from: locate a previous run's best evolved program (.py). + + Looks for ``/best/best_program.py`` first, then the latest + ``/checkpoints/checkpoint_*/best_program.py``. Returns the path as a + string, or None if no best program is found. + """ + import glob + import re + + candidates = [os.path.join(run_dir, "best", "best_program.py")] + ckpts = glob.glob(os.path.join(run_dir, "checkpoints", "checkpoint_*")) + if ckpts: + def _ckpt_num(p): + m = re.search(r"checkpoint_(\d+)$", p.rstrip("/")) + return int(m.group(1)) if m else -1 + + candidates.append(os.path.join(max(ckpts, key=_ckpt_num), "best_program.py")) + for c in candidates: + if os.path.isfile(c): + return c + return None + + async def main_async() -> int: """ Main asynchronous entry point @@ -157,6 +192,22 @@ async def main_async() -> int: print(f"Error: Evaluation file '{args.evaluation_file}' not found") return 1 else: + # --init-from: seed a fresh run from a previous run's best program (issue #28). + if getattr(args, "init_from", None): + best = _resolve_best_program(args.init_from) + if best is None: + print( + f"Error: no best program found under '{args.init_from}' " + "(looked for best/best_program.py and checkpoints/*/best_program.py)" + ) + return 1 + # Convenience: `--init-from evaluator.py` -- with only one + # positional given it is the evaluator, since the seed comes from --init-from. + if not args.evaluation_file and args.initial_program: + args.evaluation_file = args.initial_program + args.initial_program = best + print(f"Seeding initial program from {args.init_from}: {best}") + # Single-file mode validation if not args.initial_program: print("Error: Either provide initial_program or use --directory for multi-file mode") diff --git a/tests/test_cli_init_from.py b/tests/test_cli_init_from.py new file mode 100644 index 0000000..fdaf5fd --- /dev/null +++ b/tests/test_cli_init_from.py @@ -0,0 +1,59 @@ +"""Tests for `--init-from` (issue #28): seed a run from a previous run's best +program. Covers the resolver (best/ dir, checkpoint fallback, absent) and the +argument wiring.""" + +import os +import unittest +from tempfile import TemporaryDirectory + +from kaievolve.cli import _resolve_best_program, parse_args + + +def _write(path, text="x = 1\n"): + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w") as f: + f.write(text) + + +class TestResolveBestProgram(unittest.TestCase): + def test_prefers_best_dir(self): + with TemporaryDirectory() as d: + _write(os.path.join(d, "best", "best_program.py")) + _write(os.path.join(d, "checkpoints", "checkpoint_10", "best_program.py")) + self.assertEqual( + _resolve_best_program(d), os.path.join(d, "best", "best_program.py") + ) + + def test_falls_back_to_latest_checkpoint(self): + with TemporaryDirectory() as d: + _write(os.path.join(d, "checkpoints", "checkpoint_10", "best_program.py")) + _write(os.path.join(d, "checkpoints", "checkpoint_40", "best_program.py")) + _write(os.path.join(d, "checkpoints", "checkpoint_30", "best_program.py")) + # latest checkpoint by numeric suffix, not lexicographic + self.assertEqual( + _resolve_best_program(d), + os.path.join(d, "checkpoints", "checkpoint_40", "best_program.py"), + ) + + def test_returns_none_when_absent(self): + with TemporaryDirectory() as d: + self.assertIsNone(_resolve_best_program(d)) + + +class TestInitFromArg(unittest.TestCase): + def test_flag_parses(self): + import sys + + argv = sys.argv + try: + sys.argv = ["kai", "--init-from", "some/run_0", "evaluator.py", "-c", "cfg.yaml"] + args = parse_args() + self.assertEqual(args.init_from, "some/run_0") + # with one positional + --init-from, it lands in initial_program (shifted later) + self.assertEqual(args.initial_program, "evaluator.py") + finally: + sys.argv = argv + + +if __name__ == "__main__": + unittest.main()