File size: 6,468 Bytes
2facf1f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
#!/usr/bin/env python3
"""
Fork an existing multi-problem experiment at a given generation.

Copies the first N generations from a baseline run into a new directory,
so the runner can resume from gen N+1 with a different configuration
(e.g., with eval agent enabled).

Usage:
    python scripts/fork_experiment.py \
        results/frontier_cs_algorithmic/vanilla_g50_20260327_055051 \
        results/frontier_cs_algorithmic/agent_fork_g5 \
        --fork-at 5

This will:
1. For each problem (p0, p1, ...) in the source:
   - Copy gen_0/ through gen_{fork_at-1}/ directories
   - Copy experiment_config.yaml
   - Create a new evolution_db.sqlite with only programs from gen 0..fork_at-1
   - Set last_iteration = fork_at-1 so the runner resumes from fork_at
"""

import argparse
import shutil
import sqlite3
import sys
from pathlib import Path


def fork_problem(src_dir: Path, dst_dir: Path, fork_at: int) -> bool:
    """Fork a single problem directory at the given generation."""
    problem_name = src_dir.name

    # Check source DB exists
    src_db = src_dir / "evolution_db.sqlite"
    if not src_db.exists():
        print(f"  SKIP {problem_name}: no evolution_db.sqlite")
        return False

    dst_dir.mkdir(parents=True, exist_ok=True)

    # 1. Copy experiment_config.yaml
    src_config = src_dir / "experiment_config.yaml"
    if src_config.exists():
        shutil.copy2(src_config, dst_dir / "experiment_config.yaml")

    # 2. Copy gen_0 through gen_{fork_at-1}
    for gen in range(fork_at):
        src_gen = src_dir / f"gen_{gen}"
        dst_gen = dst_dir / f"gen_{gen}"
        if src_gen.exists():
            shutil.copytree(src_gen, dst_gen, dirs_exist_ok=True)

    # 3. Create forked DB with only programs from gen 0..fork_at-1
    dst_db_path = dst_dir / "evolution_db.sqlite"
    src_conn = sqlite3.connect(str(src_db))
    dst_conn = sqlite3.connect(str(dst_db_path))

    # Copy schema
    src_conn.row_factory = sqlite3.Row
    schema_rows = src_conn.execute(
        "SELECT sql FROM sqlite_master WHERE type='table' AND sql IS NOT NULL"
    ).fetchall()
    for row in schema_rows:
        dst_conn.execute(row["sql"])

    # Copy indexes
    index_rows = src_conn.execute(
        "SELECT sql FROM sqlite_master WHERE type='index' AND sql IS NOT NULL"
    ).fetchall()
    for row in index_rows:
        try:
            dst_conn.execute(row["sql"])
        except sqlite3.OperationalError:
            pass  # Index may already exist from CREATE TABLE

    # Copy programs from gen 0..fork_at-1
    programs = src_conn.execute(
        "SELECT * FROM programs WHERE generation < ?", (fork_at,)
    ).fetchall()

    if programs:
        cols = [desc[0] for desc in src_conn.execute("SELECT * FROM programs LIMIT 0").description]
        placeholders = ", ".join(["?"] * len(cols))
        col_names = ", ".join(cols)
        for prog in programs:
            dst_conn.execute(
                f"INSERT INTO programs ({col_names}) VALUES ({placeholders})",
                tuple(prog),
            )

    # Copy archive entries that reference copied programs
    program_ids = set()
    for prog in programs:
        program_ids.add(prog[0])  # id is first column

    archive_rows = src_conn.execute("SELECT program_id FROM archive").fetchall()
    for row in archive_rows:
        if row[0] in program_ids:
            dst_conn.execute("INSERT INTO archive (program_id) VALUES (?)", (row[0],))

    # Set metadata
    dst_conn.execute(
        "INSERT OR REPLACE INTO metadata_store (key, value) VALUES (?, ?)",
        ("last_iteration", str(fork_at - 1)),
    )

    # Find best program in the forked range
    best_row = dst_conn.execute(
        "SELECT id FROM programs WHERE combined_score = (SELECT MAX(combined_score) FROM programs) LIMIT 1"
    ).fetchone()
    if best_row:
        dst_conn.execute(
            "INSERT OR REPLACE INTO metadata_store (key, value) VALUES (?, ?)",
            ("best_program_id", best_row[0]),
        )

    dst_conn.commit()

    n_programs = dst_conn.execute("SELECT COUNT(*) FROM programs").fetchone()[0]
    n_archive = dst_conn.execute("SELECT COUNT(*) FROM archive").fetchone()[0]

    src_conn.close()
    dst_conn.close()

    print(f"  OK   {problem_name}: {n_programs} programs, {n_archive} archived, last_iteration={fork_at-1}")
    return True


def main():
    parser = argparse.ArgumentParser(
        description="Fork an experiment at a given generation for controlled comparison"
    )
    parser.add_argument("source", help="Source experiment directory")
    parser.add_argument("destination", help="Destination directory for forked experiment")
    parser.add_argument("--fork-at", type=int, required=True,
                        help="Fork at this generation (copies gen 0..fork_at-1)")
    parser.add_argument("--problems", type=str, default=None,
                        help="Comma-separated problem IDs to fork (e.g., 'p0,p1,p5'). Default: all")

    args = parser.parse_args()

    src_root = Path(args.source)
    dst_root = Path(args.destination)

    if not src_root.exists():
        print(f"Error: source directory not found: {src_root}")
        sys.exit(1)

    if dst_root.exists():
        print(f"Error: destination already exists: {dst_root}")
        print("Remove it first or choose a different path.")
        sys.exit(1)

    # Discover problems
    all_problems = sorted(
        [d for d in src_root.iterdir() if d.is_dir() and d.name.startswith("p") and d.name[1:].isdigit()],
        key=lambda d: int(d.name[1:]),
    )

    if args.problems:
        selected = set(args.problems.split(","))
        all_problems = [d for d in all_problems if d.name in selected]

    print(f"Forking experiment at generation {args.fork_at}")
    print(f"  Source:      {src_root}")
    print(f"  Destination: {dst_root}")
    print(f"  Problems:    {len(all_problems)}")
    print()

    dst_root.mkdir(parents=True, exist_ok=True)

    success = 0
    for problem_dir in all_problems:
        dst_problem = dst_root / problem_dir.name
        if fork_problem(problem_dir, dst_problem, args.fork_at):
            success += 1

    print()
    print(f"Done: {success}/{len(all_problems)} problems forked successfully")
    print(f"Output: {dst_root}")
    print()
    print(f"To run with eval agent, point your experiment runner at {dst_root}/pN/")
    print("The runner will auto-resume from gen", args.fork_at)


if __name__ == "__main__":
    main()