File size: 7,044 Bytes
c27193b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
190
191
192
193
194
195
196
197
198
199
200
"""Run one measured autoresearch attempt (discipline 1).

Wraps the existing run path: enforce the git fence + dedup, append a `running` row,
shell the contract's `command` under a bounded wall-clock (identical per attempt so
scores are comparable), read the declared metric, capture the editable's diff for the
UI, and append the terminal row. Bounded teardown is a full process-group SIGTERM→
SIGKILL so a timed-out training subprocess never orphans (exit 124 on timeout).

The metric `source` is read from a conventional artifact the command writes into the
run dir — `metrics_jsonl` is the universal path; the eval-backed sources read the JSON
their (extra-gated) eval already produced, so this module stays base-install importable.
"""

from __future__ import annotations

import json
import os
import signal
import subprocess
import time
from pathlib import Path

from pydantic import BaseModel, ConfigDict, Field

from mindxtrain.research.contract import AttemptContract
from mindxtrain.research.guard import assert_clean_and_committed, assert_new_commit, git, repo_root
from mindxtrain.research.ledger import AttemptMetric, Ledger, LedgerStatus

# source → (artifact filename, fixed metric key or None to use contract.metric)
_SOURCE_ARTIFACT: dict[str, tuple[str, str | None]] = {
    "classroom_delta": ("classroom.json", "imprint_delta"),
    "mei": ("mei.json", "composite"),
    "bfcl": ("bfcl.json", None),
    "tau_bench": ("tau_bench.json", None),
    "held_out_loss": ("held_out_loss.json", None),
}
_TIMEOUT_EXIT = 124


class AttemptResult(BaseModel):
    """The outcome of one measured attempt."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    commit: str
    value: float | None = None
    status: LedgerStatus
    run_dir: str
    kind: str = Field(default="attempt")

    def summary(self) -> str:
        value = "n/a" if self.value is None else f"{self.value:.4f}"
        return f"champion commit={self.commit} value={value} status={self.status}"


def _terminate_group(pid: int, *, grace_s: float = 3.0) -> None:
    """SIGTERM the process group, wait briefly, then SIGKILL any survivors."""
    try:
        pgid = os.getpgid(pid)
    except OSError:
        return
    try:
        os.killpg(pgid, signal.SIGTERM)
    except OSError:
        return
    deadline = time.monotonic() + grace_s
    while time.monotonic() < deadline:
        try:
            os.killpg(pgid, 0)  # probe: raises once the group is gone
        except OSError:
            return
        time.sleep(0.1)
    try:
        os.killpg(pgid, signal.SIGKILL)
    except OSError:
        pass


def run_logged(command: str, log_path: Path, timeout_s: float, *, cwd: Path | None = None) -> int:
    """Run `command` (shell), tee output to `log_path`, return exit code (124 on timeout)."""
    log_path.parent.mkdir(parents=True, exist_ok=True)
    with log_path.open("w", encoding="utf-8") as log:
        proc = subprocess.Popen(  # contract commands are operator-authored
            command,
            shell=True,
            cwd=str(cwd) if cwd is not None else None,
            stdout=log,
            stderr=subprocess.STDOUT,
            start_new_session=True,  # own process group → clean teardown
            text=True,
        )
        try:
            return proc.wait(timeout=timeout_s)
        except subprocess.TimeoutExpired:
            _terminate_group(proc.pid)
            try:
                proc.wait(timeout=5.0)
            except subprocess.TimeoutExpired:
                pass
            return _TIMEOUT_EXIT


def _format_command(contract: AttemptContract, run_dir: Path, run_name: str) -> str:
    return (
        contract.command.replace("{run_dir}", str(run_dir))
        .replace("{run_name}", run_name)
        .replace("{attempt_timeout_minutes}", str(contract.attempt_timeout_minutes))
    )


def _read_last_jsonl(path: Path, key: str) -> float | None:
    if not path.exists():
        return None
    last: object = None
    for line in path.read_text(encoding="utf-8").splitlines():
        stripped = line.strip()
        if not stripped:
            continue
        try:
            last = json.loads(stripped)
        except ValueError:
            continue
    if not isinstance(last, dict):
        return None
    val = last.get(key)
    return float(val) if isinstance(val, (int, float)) else None


def read_metric(contract: AttemptContract, run_dir: Path) -> float | None:
    """Resolve the contract's scalar from the artifact the command wrote; None if absent."""
    if contract.source == "metrics_jsonl":
        return _read_last_jsonl(run_dir / "metrics.jsonl", contract.metric)
    filename, fixed_key = _SOURCE_ARTIFACT[contract.source]
    path = run_dir / filename
    if not path.exists():
        return None
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except ValueError:
        return None
    if not isinstance(data, dict):
        return None
    val = data.get(fixed_key or contract.metric)
    return float(val) if isinstance(val, (int, float)) else None


def _capture_diff(editable: list[Path], run_dir: Path, root: Path) -> None:
    """Best-effort: write the just-committed edit to run_dir/code.diff for the UI."""
    try:
        diff = git("diff", "HEAD~1", "HEAD", "--", *[str(p) for p in editable], cwd=root)
    except Exception:  # baseline (no parent) or detached state — diff is optional
        diff = ""
    (run_dir / "code.diff").write_text(diff, encoding="utf-8")


def run_attempt(
    contract: AttemptContract,
    researcher: str,
    run_dir: Path,
    ledger: Ledger,
    *,
    root: Path | None = None,
    kind: str = "attempt",
) -> AttemptResult:
    """Fence → dedup → run → measure → record. Returns the structured result."""
    root = root or repo_root()
    run_dir = Path(run_dir)
    run_dir.mkdir(parents=True, exist_ok=True)

    commit = assert_clean_and_committed(contract.editable, root)
    assert_new_commit(commit, researcher, ledger)
    ledger.append(
        researcher=researcher, commit=commit, status="running",
        kind=kind, run_dir=str(run_dir),  # type: ignore[arg-type]
    )

    command = _format_command(contract, run_dir, run_dir.name)
    code = run_logged(
        command, run_dir / "logs.log", contract.attempt_timeout_minutes * 60.0, cwd=root
    )
    value = read_metric(contract, run_dir)

    if code == _TIMEOUT_EXIT:
        status: LedgerStatus = "timed_out"
    elif code == 0 and value is not None:
        status = "completed"
    else:
        status = "failed"

    _capture_diff(contract.editable, run_dir, root)
    ledger.append(
        researcher=researcher, commit=commit, status=status,
        kind=kind, run_dir=str(run_dir),  # type: ignore[arg-type]
        metric=AttemptMetric(name=contract.metric, value=value, mode=contract.mode),
    )
    return AttemptResult(commit=commit, value=value, status=status, run_dir=str(run_dir), kind=kind)


__all__ = ["AttemptResult", "read_metric", "run_attempt", "run_logged"]