Unify cache root override: loader.get/set_cache_root as single source, fixes --cache-dir not relocating .raw blobs

This commit is contained in:
sora 2026-08-24 03:37:15 +00:00
parent f8cd15fea1
commit 414a89216c
2 changed files with 23 additions and 11 deletions

View File

@ -25,11 +25,20 @@ import time
from pathlib import Path
from typing import Any, Callable, Dict, Iterator, List, Optional, Union
from .loader import load_raw_records
from .loader import get_cache_root, load_raw_records
from .sample import Sample
from .spec import DatasetSpec
CACHE_ROOT = Path(os.environ.get('EVALHARNESS_CACHE', '~/.cache/evalharness')).expanduser() / 'datasets'
# Backward-compatible alias: the canonical root lives in loader.get_cache_root()
# (a function, so runtime overrides via set_cache_root are always respected).
CACHE_ROOT = get_cache_root()
def set_cache_root(path) -> None:
"""Override the cache root at runtime (used by the CLI --cache-dir flag)."""
from . import loader
loader.set_cache_root(path)
def safe_filename(s: str, max_length: int = 255) -> str:
@ -43,12 +52,6 @@ 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."""
@ -71,7 +74,7 @@ class Dataset:
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
return get_cache_root() / 'datasets' / safe_filename(self.spec.name) / subdir
@property
def is_materialized(self) -> bool:

View File

@ -28,6 +28,16 @@ _RESERVED_PARAMS = {'hub', 'ms_files', 'filter_column'}
_MS_API = 'https://www.modelscope.cn/api/v1/datasets'
def get_cache_root() -> Path:
"""Canonical evalharness cache root ($EVALHARNESS_CACHE or ~/.cache/evalharness)."""
return Path(os.environ.get('EVALHARNESS_CACHE', '~/.cache/evalharness')).expanduser()
def set_cache_root(path) -> None:
"""Override the cache root at runtime (also moves the .raw blob store)."""
os.environ['EVALHARNESS_CACHE'] = str(path)
def load_raw_records(spec: DatasetSpec, raw_dir: Optional[Path] = None) -> List[Dict[str, Any]]:
"""Load raw records (list of dicts) from the spec's source.
@ -220,8 +230,7 @@ def _load_from_modelscope(spec: DatasetSpec, raw_dir: Optional[Path] = None) ->
f'modelscope {spec.source!r}. Available (first 10): {files[:10]}'
)
# shared blob store: download once per repo, hardlink into each cache entry
blob_dir = Path(os.environ.get('EVALHARNESS_CACHE', '~/.cache/evalharness')).expanduser()
blob_dir = blob_dir / '.raw' / hashlib.md5(spec.source.encode()).hexdigest()[:10]
blob_dir = get_cache_root() / '.raw' / hashlib.md5(spec.source.encode()).hexdigest()[:10]
records: List[Dict[str, Any]] = []
blobs: List[Path] = []
for path in selected: