240 lines
9.5 KiB
Python
240 lines
9.5 KiB
Python
"""Dataset: a lazy handle over a registered dataset.
|
|
|
|
``get_dataset('gsm8k')`` costs nothing -- no download, no parsing. The first
|
|
actual use (iteration / len / indexing) triggers ``materialize()``:
|
|
|
|
cache hit -> read the cached samples.jsonl
|
|
cache miss -> download raw -> record_to_sample -> atomic cache write -> read
|
|
|
|
Cache directory name: ``{safe_name}-{md5(source+split+subset+version+params)}``
|
|
so any config change yields a different cache entry (zero invalidation logic).
|
|
|
|
Each cache entry is self-contained:
|
|
raw/ native source data, exactly as downloaded (never converted)
|
|
samples.jsonl the unified Sample stream converted from raw/
|
|
meta.json spec + provenance
|
|
"""
|
|
|
|
import fcntl
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
import string
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Dict, Iterator, List, Optional, Union
|
|
|
|
from .loader import load_raw_records
|
|
from .sample import Sample
|
|
from .spec import DatasetSpec
|
|
|
|
CACHE_ROOT = Path(os.environ.get('EVALHARNESS_CACHE', '~/.cache/evalharness')).expanduser() / 'datasets'
|
|
|
|
|
|
def safe_filename(s: str, max_length: int = 255) -> str:
|
|
safe_chars = string.ascii_letters + string.digits + '.-_'
|
|
s = ''.join(c if c in safe_chars else '_' for c in s)
|
|
s = re.sub(r'_+', '_', s).strip('._')
|
|
return (s or 'untitled')[:max_length]
|
|
|
|
|
|
def gen_hash(s: str) -> str:
|
|
return hashlib.md5(s.encode('utf-8')).hexdigest()
|
|
|
|
|
|
def set_cache_root(path) -> None:
|
|
"""Override the cache root at runtime (used by the CLI --cache-dir flag)."""
|
|
global CACHE_ROOT
|
|
CACHE_ROOT = Path(path).expanduser() / 'datasets'
|
|
|
|
|
|
class Dataset:
|
|
"""Sequence-like lazy dataset. Nothing is downloaded until first use."""
|
|
|
|
def __init__(self, spec: DatasetSpec, record_fn: Callable, samples: Optional[List[Sample]] = None):
|
|
self.spec = spec
|
|
self._record_fn = record_fn
|
|
self._samples = samples # None => not materialized yet
|
|
self.lineage: Dict[str, Any] = {}
|
|
|
|
# ---------------- materialization ----------------
|
|
|
|
@property
|
|
def cache_dir(self) -> Path:
|
|
key = f'{self.spec.source}{self.spec.split}{self.spec.subset}{self.spec.version}{self.spec.params}'
|
|
# Layout: datasets/<benchmark_name>/<subset>_<split>[-<version>]-<hash6>/
|
|
# Readable benchmark folder + readable subset/split; the 6-char hash
|
|
# suffix disambiguates different sources/params that would otherwise
|
|
# collide on the same subset_split name (correctness requirement).
|
|
parts = f'{safe_filename(self.spec.subset)}_{safe_filename(self.spec.split)}'
|
|
if self.spec.version:
|
|
parts += f'_{safe_filename(self.spec.version)}'
|
|
subdir = f'{parts}-{gen_hash(key)[:6]}'
|
|
return CACHE_ROOT / safe_filename(self.spec.name) / subdir
|
|
|
|
@property
|
|
def is_materialized(self) -> bool:
|
|
return self._samples is not None
|
|
|
|
def materialize(self, force: bool = False) -> 'Dataset':
|
|
if self._samples is not None and not force:
|
|
return self
|
|
cache_dir = self.cache_dir
|
|
cache_file = cache_dir / 'samples.jsonl'
|
|
if cache_file.exists() and not force:
|
|
self._samples = self._read_cache(cache_file)
|
|
self.lineage = {'from': 'cache', 'cache_dir': str(cache_dir)}
|
|
return self
|
|
|
|
cache_dir.parent.mkdir(parents=True, exist_ok=True)
|
|
lock_path = cache_dir.with_suffix('.lock')
|
|
with open(lock_path, 'w') as lock_f: # noqa: PTH123
|
|
fcntl.flock(lock_f, fcntl.LOCK_EX)
|
|
try:
|
|
if cache_file.exists() and not force: # double-check under lock
|
|
self._samples = self._read_cache(cache_file)
|
|
self.lineage = {'from': 'cache', 'cache_dir': str(cache_dir)}
|
|
return self
|
|
cache_dir.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp_dir = cache_dir.with_name(cache_dir.name + f'.tmp-{os.getpid()}')
|
|
if tmp_dir.exists():
|
|
import shutil
|
|
|
|
shutil.rmtree(tmp_dir)
|
|
tmp_dir.mkdir(parents=True)
|
|
records = load_raw_records(self.spec, raw_dir=tmp_dir / 'raw')
|
|
samples = [self._to_sample(r) for r in records]
|
|
self._assign_ids(samples)
|
|
self._write_cache_atomic(samples, cache_dir, tmp_dir)
|
|
self._samples = samples
|
|
self.lineage = {'from': 'source', 'cache_dir': str(cache_dir)}
|
|
finally:
|
|
fcntl.flock(lock_f, fcntl.LOCK_UN)
|
|
return self
|
|
|
|
def _to_sample(self, record: Dict[str, Any]) -> Sample:
|
|
sample = self._record_fn(record)
|
|
if not sample.task_type:
|
|
sample.task_type = self.spec.task_type
|
|
return sample
|
|
|
|
@staticmethod
|
|
def _assign_ids(samples: List[Sample]) -> None:
|
|
"""Assign sequential ids to samples that don't carry one."""
|
|
for i, s in enumerate(samples):
|
|
if s.id is None:
|
|
s.id = i
|
|
|
|
def _read_cache(self, cache_file: Path) -> List[Sample]:
|
|
with open(cache_file, encoding='utf-8') as f:
|
|
samples = [Sample.model_validate(json.loads(line)) for line in f if line.strip()]
|
|
return samples
|
|
|
|
def _write_cache_atomic(self, samples: List[Sample], cache_dir: Path, tmp_dir: Path) -> None:
|
|
"""Finalize the tmp dir (samples + meta next to raw/) and swap it in.
|
|
|
|
``raw/`` was already populated by ``load_raw_records`` inside tmp_dir;
|
|
the tmp+rename swap makes the whole entry (raw + samples + meta)
|
|
appear atomically.
|
|
"""
|
|
with open(tmp_dir / 'samples.jsonl', 'w', encoding='utf-8') as f:
|
|
for s in samples:
|
|
f.write(json.dumps(s.model_dump(), ensure_ascii=False) + '\n')
|
|
raw_dir = tmp_dir / 'raw'
|
|
meta = {
|
|
'spec': {k: v for k, v in vars(self.spec).items()},
|
|
'num_samples': len(samples),
|
|
'raw_files': sorted(p.name for p in raw_dir.iterdir()) if raw_dir.exists() else [],
|
|
'note': 'raw/ holds the native source data exactly as downloaded; samples.jsonl is the converted view',
|
|
'created_at': time.strftime('%Y-%m-%d %H:%M:%S'),
|
|
}
|
|
with open(tmp_dir / 'meta.json', 'w', encoding='utf-8') as f:
|
|
json.dump(meta, f, ensure_ascii=False, indent=2)
|
|
final = cache_dir
|
|
if final.exists():
|
|
import shutil
|
|
|
|
# retire the old entry, then swap the new one in
|
|
old = cache_dir.with_name(cache_dir.name + f'.old-{os.getpid()}')
|
|
os.rename(final, old)
|
|
try:
|
|
os.rename(tmp_dir, final)
|
|
except OSError:
|
|
os.rename(old, final)
|
|
raise
|
|
shutil.rmtree(old, ignore_errors=True)
|
|
else:
|
|
os.rename(tmp_dir, final)
|
|
|
|
def unload(self) -> bool:
|
|
"""Drop this dataset's cache entry (raw/ + samples.jsonl + meta.json).
|
|
|
|
Pure cache management: in-memory samples (if any) stay usable; the
|
|
next materialize rebuilds from source. Execution-environment
|
|
resources (docker images declared by ``Sample.sandbox``) are NOT
|
|
touched -- those belong to the sandbox layer's lifecycle.
|
|
"""
|
|
import shutil
|
|
|
|
removed = False
|
|
if self.cache_dir.exists():
|
|
shutil.rmtree(self.cache_dir)
|
|
removed = True
|
|
lock_path = self.cache_dir.with_suffix('.lock')
|
|
if lock_path.exists():
|
|
lock_path.unlink()
|
|
if self._samples is not None or self.lineage:
|
|
self._samples = None
|
|
self.lineage = {}
|
|
return removed
|
|
|
|
# ---------------- sequence protocol (triggers materialize) ----------------
|
|
|
|
def _require(self) -> List[Sample]:
|
|
self.materialize()
|
|
return self._samples
|
|
|
|
def __iter__(self) -> Iterator[Sample]:
|
|
return iter(self._require())
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._require())
|
|
|
|
def __getitem__(self, i: Union[int, slice]) -> Union[Sample, List[Sample]]:
|
|
return self._require()[i]
|
|
|
|
# ---------------- views / derived data ----------------
|
|
|
|
def view(self, samples: List[Sample], lineage: Optional[Dict[str, Any]] = None) -> 'Dataset':
|
|
"""An in-memory derived dataset (filter/sample/dedup results).
|
|
|
|
Same class, same interface; ``lineage`` records how it was produced.
|
|
"""
|
|
derived = Dataset(self.spec, self._record_fn, samples=samples)
|
|
derived.lineage = {'from': 'derived', 'parent': self.spec.name, **(lineage or {})}
|
|
return derived
|
|
|
|
# ---------------- introspection ----------------
|
|
|
|
def stats(self) -> Dict[str, Any]:
|
|
samples = self._require()
|
|
lengths = [len(s.input_text) for s in samples]
|
|
targets = [s.target if isinstance(s.target, str) else ','.join(s.target) for s in samples]
|
|
return {
|
|
'name': self.spec.name,
|
|
'task_type': self.spec.task_type,
|
|
'num_samples': len(samples),
|
|
'input_len': {
|
|
'min': min(lengths) if lengths else 0,
|
|
'max': max(lengths) if lengths else 0,
|
|
'mean': round(sum(lengths) / len(lengths), 1) if lengths else 0,
|
|
},
|
|
'target_top': sorted({t: targets.count(t) for t in set(targets)}.items(), key=lambda kv: -kv[1])[:10],
|
|
'cache_dir': str(self.cache_dir),
|
|
}
|
|
|
|
def __repr__(self) -> str:
|
|
state = 'materialized' if self.is_materialized else 'lazy'
|
|
return f"Dataset(name={self.spec.name!r}, type={self.spec.task_type!r}, {state})"
|