77d2c4569a
A finished 12-minute humaneval bench died at the save line: some scorer detail carried a non-JSON object (Ellipsis) and report.save's json.dumps had no default. Every dumps on the save path now stringifies exotic objects instead of killing the run; verified by replaying the exact crashed checkpoint end-to-end (pass 82.3%, 15s, no crash). Co-Authored-By: Claude <noreply@anthropic.com>
125 lines
4.9 KiB
Python
125 lines
4.9 KiB
Python
"""Resumable evaluation: checkpoint plugin.
|
|
|
|
Run-level checkpointing: every completed sample is appended to a jsonl
|
|
checkpoint file; on restart (crash/OOM/service down), completed samples are
|
|
restored and ONLY the missing ones are generated. The report is then
|
|
assembled from restored + fresh predictions.
|
|
|
|
from evalharness.eval.checkpoint import CheckpointStore
|
|
store = CheckpointStore(path='run.jsonl', model=model_spec, dataset=name)
|
|
done = store.load() # {sample_key: prediction-dict}
|
|
... generate only missing ...
|
|
store.append(sample_key, pred) # after EACH sample (crash-safe)
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional
|
|
|
|
|
|
class CheckpointStore:
|
|
"""Append-only jsonl checkpoint; keyed by stable sample identity."""
|
|
|
|
def __init__(self, path: str, model: str = '', dataset: str = ''):
|
|
self.path = Path(path)
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.model = model
|
|
self.dataset = dataset
|
|
self._entries: Dict[str, Dict[str, Any]] = {}
|
|
self._scores: Dict[str, Dict[str, Any]] = {}
|
|
self._fh = None
|
|
|
|
@staticmethod
|
|
def key_for(sample, idx: int) -> str:
|
|
"""Stable per-sample key: prefer dataset-native ids, fall back to
|
|
a hash of the input text (survives re-orderings)."""
|
|
native = (sample.metadata or {}).get('task_id') \
|
|
or (sample.metadata or {}).get('id') \
|
|
or (sample.metadata or {}).get('instance_id')
|
|
if native:
|
|
return str(native)
|
|
import hashlib
|
|
|
|
h = hashlib.md5((sample.input_text or '').encode('utf-8')).hexdigest()[:16]
|
|
return f'{idx}:{h}'
|
|
|
|
def load(self) -> Dict[str, Dict[str, Any]]:
|
|
"""Read all checkpointed predictions (idempotent)."""
|
|
self._entries = {}
|
|
self._scores = {}
|
|
if not self.path.exists():
|
|
return self._entries
|
|
with open(self.path, encoding='utf-8') as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
rec = json.loads(line)
|
|
self._entries[rec['key']] = rec.get('pred', {})
|
|
if rec.get('score'):
|
|
self._scores[rec['key']] = rec['score']
|
|
except (ValueError, KeyError):
|
|
continue # torn tail line from a crash -- safe to skip
|
|
return self._entries
|
|
|
|
def append(self, key: str, pred: Dict[str, Any]) -> None:
|
|
"""Persist one prediction immediately (fsync-free append is fine:
|
|
worst case loses the last in-flight sample on crash)."""
|
|
rec = {'key': key, 'ts': time.time(), 'pred': pred}
|
|
with open(self.path, 'a', encoding='utf-8') as f:
|
|
f.write(json.dumps(rec, ensure_ascii=False, default=str) + '\n')
|
|
self._entries[key] = pred
|
|
|
|
def scores(self) -> Dict[str, Dict[str, Any]]:
|
|
"""Cached per-sample score records ({key: {'fp', 'scores', ...}}).
|
|
|
|
Scores are bound to predictions and live in the SAME file -- one
|
|
--resume flag controls both layers."""
|
|
return self._scores
|
|
|
|
def put_scores(self, records: Dict[str, Dict[str, Any]]) -> None:
|
|
"""Attach score records to checkpoint entries (end-of-evaluation
|
|
writeback). Rewrites the file atomically; legacy lines without a
|
|
score field simply gain one."""
|
|
if not records:
|
|
return
|
|
lines: Dict[str, str] = {}
|
|
if self.path.exists():
|
|
with open(self.path, encoding='utf-8') as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
rec = json.loads(line)
|
|
except ValueError:
|
|
continue
|
|
if rec.get('key') in records:
|
|
rec['score'] = records[rec['key']]
|
|
lines[rec['key']] = json.dumps(rec, ensure_ascii=False,
|
|
default=str)
|
|
tmp = self.path.with_suffix('.tmp')
|
|
with open(tmp, 'w', encoding='utf-8') as f:
|
|
for v in lines.values():
|
|
f.write(v + '\n')
|
|
os.replace(tmp, self.path)
|
|
self._scores.update(records)
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._entries)
|
|
|
|
def summary(self) -> Dict[str, Any]:
|
|
return {'path': str(self.path), 'restored': len(self._entries),
|
|
'size_kb': self.path.stat().st_size // 1024 if self.path.exists() else 0}
|
|
|
|
|
|
def checkpoint_path(root: str, dataset: str, model: str, tag: str = '') -> str:
|
|
"""Deterministic path per (dataset, model[, tag]) so re-runs resume."""
|
|
import hashlib
|
|
|
|
h = hashlib.md5(f'{dataset}|{model}|{tag}'.encode()).hexdigest()[:10]
|
|
return os.path.join(root, 'ckpt', f'{dataset}-{h}.jsonl')
|