Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-evaluation/src/lexigram/ai/evaluation/tracking.py: 33%

89 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

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"]