1"""Seed-stable local experiment tracking for the AI evaluation framework.
2
3The :class:`LocalTracker` persists runs as JSON documents plus JSONL
4metric/error streams under ``<root>/runs/<run_id>/``. Run ids are
5derived deterministically from the experiment name, seed, and
6canonicalized knob config, so rerunning the same seed and knobs resumes
7the same run and produces byte-identical artifacts.
8"""
9
10from __future__ import annotations
11
12from datetime import UTC, datetime
13import hashlib
14from pathlib import Path
15from typing import Any
16
17from lexigram.ai.evaluation.exceptions import TrackingError
18from lexigram.contracts.ai.experiment import (
19 ErrorRecord,
20 ExperimentConfig,
21 ExperimentRun,
22 ExperimentTrackerProtocol,
23 MetricRecord,
24 RunStatus,
25)
26from lexigram.logging import get_logger
27from lexigram.serialization import dumps_str, loads_str
28
29logger = get_logger(__name__)
30
31
32def canonical_json(data: Any) -> str:
33 """Serialize ``data`` to a stable, key-sorted JSON string."""
34 return dumps_str(data, sort_keys=True)
35
36
37def make_run_id(name: str, seed: int, config: dict[str, Any]) -> str:
38 """Derive a deterministic run id from name, seed, and canonical config.
39
40 Args:
41 name: Experiment name.
42 seed: Seed value.
43 config: Knob configuration dict.
44
45 Returns:
46 Run id of the form ``<name>-<seed>-<8-char digest>``.
47
48 Example:
49 ```python
50 run_id = make_run_id("probe", 42, {"model": "gpt-4o"})
51 assert run_id == make_run_id("probe", 42, {"model": "gpt-4o"})
52 ```
53 """
54 digest = hashlib.sha256(
55 f"{name}|{seed}|{canonical_json(config)}".encode()
56 ).hexdigest()[:8]
57 return f"{name}-{seed}-{digest}"
58
59
60class LocalTracker(ExperimentTrackerProtocol):
61 """JSON/JSONL-backed tracker persisting runs under ``<root>/runs/``.
62
63 Args:
64 root: Base directory for run artifacts. Defaults to ``runs``.
65 """
66
67 def __init__(self, root: str | Path = "runs") -> None:
68 self._root = Path(root)
69
70 def _run_dir(self, run_id: str) -> Path:
71 return self._root / "runs" / run_id
72
73 def _load_run(self, run_id: str) -> ExperimentRun | None:
74 run_file = self._run_dir(run_id) / "run.json"
75 if not run_file.exists():
76 return None
77 try:
78 data = loads_str(run_file.read_text())
79 except (OSError, ValueError) as exc:
80 raise TrackingError(f"cannot read run {run_id!r}: {exc}") from exc
81 return ExperimentRun(
82 **{key: value for key, value in data.items() if key != "status"},
83 status=RunStatus(data["status"]),
84 )
85
86 def _store_run(self, run: ExperimentRun) -> ExperimentRun:
87 run_file = self._run_dir(run.run_id) / "run.json"
88 try:
89 run_file.write_text(
90 dumps_str(
91 {
92 "run_id": run.run_id,
93 "experiment": run.experiment,
94 "seed": run.seed,
95 "config": run.config,
96 "config_hash": run.config_hash,
97 "status": run.status.value,
98 "started_at": run.started_at,
99 "finished_at": run.finished_at,
100 },
101 sort_keys=True,
102 )
103 )
104 except OSError as exc:
105 raise TrackingError(f"cannot persist run {run.run_id!r}: {exc}") from exc
106 return run
107
108 async def start(self, config: ExperimentConfig) -> ExperimentRun:
109 """Start (or resume) an experiment run for the given config.
110
111 Args:
112 config: Seed and knob configuration of the run.
113
114 Returns:
115 The started or resumed run.
116
117 Raises:
118 TrackingError: If the run manifest cannot be persisted.
119 """
120 run_id = make_run_id(config.name, config.seed, config.config)
121 existing = self._load_run(run_id)
122 if existing is not None:
123 return existing
124 self._run_dir(run_id).mkdir(parents=True, exist_ok=True)
125 run = ExperimentRun(
126 run_id=run_id,
127 experiment=config.name,
128 seed=config.seed,
129 config=config.config,
130 config_hash=hashlib.sha256(
131 canonical_json(config.config).encode()
132 ).hexdigest(),
133 status=RunStatus.RUNNING,
134 started_at=datetime.now(UTC).isoformat(),
135 )
136 self._store_run(run)
137 return run
138
139 async def log_metric(
140 self, run_id: str, name: str, value: float, step: int = 0
141 ) -> None:
142 """Record a scalar metric for a run.
143
144 Args:
145 run_id: Run identifier.
146 name: Metric name.
147 value: Metric value.
148 step: Step index the metric was recorded at. Defaults to 0.
149
150 Raises:
151 TrackingError: If the metric line cannot be appended.
152 """
153 self._append(
154 run_id, "metrics.jsonl", {"step": step, "name": name, "value": value}
155 )
156
157 async def log_error(
158 self, run_id: str, kind: str, message: str, step: int = 0
159 ) -> None:
160 """Record an error for a run.
161
162 Args:
163 run_id: Run identifier.
164 kind: Error kind or code.
165 message: Human-readable error message.
166 step: Step index the error occurred at. Defaults to 0.
167
168 Raises:
169 TrackingError: If the error line cannot be appended.
170 """
171 self._append(
172 run_id, "errors.jsonl", {"kind": kind, "message": message, "step": step}
173 )
174
175 def _append(self, run_id: str, filename: str, record: dict[str, Any]) -> None:
176 stream = self._run_dir(run_id) / filename
177 try:
178 stream.parent.mkdir(parents=True, exist_ok=True)
179 with stream.open("a") as handle:
180 handle.write(canonical_json(record) + "\n")
181 except OSError as exc:
182 raise TrackingError(
183 f"cannot append {filename} for run {run_id!r}: {exc}"
184 ) from exc
185
186 async def metrics(self, run_id: str) -> list[MetricRecord]:
187 """Return all metric records for a run.
188
189 Args:
190 run_id: Run identifier.
191
192 Returns:
193 Metric records in recording order.
194 """
195 return [MetricRecord(**line) for line in self._stream(run_id, "metrics.jsonl")]
196
197 async def errors(self, run_id: str) -> list[ErrorRecord]:
198 """Return all error records for a run.
199
200 Args:
201 run_id: Run identifier.
202
203 Returns:
204 Error records in recording order.
205 """
206 return [ErrorRecord(**line) for line in self._stream(run_id, "errors.jsonl")]
207
208 def _stream(self, run_id: str, filename: str) -> list[dict[str, Any]]:
209 stream = self._run_dir(run_id) / filename
210 if not stream.exists():
211 return []
212 try:
213 return [
214 loads_str(line)
215 for line in stream.read_text().splitlines()
216 if line.strip()
217 ]
218 except (OSError, ValueError) as exc:
219 raise TrackingError(
220 f"cannot read {filename} for run {run_id!r}: {exc}"
221 ) from exc
222
223 async def snapshot(self, run_id: str) -> dict[str, Any]:
224 """Return a compact summary of a run's current state.
225
226 Args:
227 run_id: Run identifier.
228
229 Returns:
230 Latest value per metric, error counts, and run metadata.
231
232 Raises:
233 TrackingError: If the run is unknown.
234 """
235 run = self._load_run(run_id)
236 if run is None:
237 raise TrackingError(f"unknown run {run_id!r}")
238 latest: dict[str, float] = {}
239 for metric in await self.metrics(run_id):
240 latest[metric.name] = metric.value
241 kinds: dict[str, int] = {}
242 for error in await self.errors(run_id):
243 kinds[error.kind] = kinds.get(error.kind, 0) + 1
244 return {
245 "run_id": run_id,
246 "experiment": run.experiment,
247 "seed": run.seed,
248 "config_hash": run.config_hash,
249 "status": run.status.value,
250 "metrics": latest,
251 "error_kinds": kinds,
252 }
253
254 async def resume(self, run_id: str) -> ExperimentRun | None:
255 """Return an already-started run, or ``None`` when unknown.
256
257 Args:
258 run_id: Run identifier.
259
260 Returns:
261 The existing run, or ``None``.
262 """
263 return self._load_run(run_id)
264
265 async def finish(
266 self, run_id: str, status: RunStatus = RunStatus.COMPLETED
267 ) -> None:
268 """Mark a run finished.
269
270 Args:
271 run_id: Run identifier.
272 status: Terminal status. Defaults to ``COMPLETED``.
273
274 Raises:
275 TrackingError: If the run is unknown.
276 """
277 run = self._load_run(run_id)
278 if run is None:
279 raise TrackingError(f"unknown run {run_id!r}")
280 finished = self._store_run(
281 ExperimentRun(
282 run_id=run.run_id,
283 experiment=run.experiment,
284 seed=run.seed,
285 config=run.config,
286 config_hash=run.config_hash,
287 status=status,
288 started_at=run.started_at,
289 finished_at=datetime.now(UTC).isoformat(),
290 )
291 )
292 logger.info(
293 "experiment_finished", run_id=run_id, status=status.value, finished=finished
294 )
295
296
297__all__ = ["LocalTracker", "canonical_json", "make_run_id"]