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 pathlib import Path
from typing import Any, Callable, Dict, Iterator, List, Optional, Union 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 .sample import Sample
from .spec import DatasetSpec 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: 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() 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: class Dataset:
"""Sequence-like lazy dataset. Nothing is downloaded until first use.""" """Sequence-like lazy dataset. Nothing is downloaded until first use."""
@ -71,7 +74,7 @@ class Dataset:
if self.spec.version: if self.spec.version:
parts += f'_{safe_filename(self.spec.version)}' parts += f'_{safe_filename(self.spec.version)}'
subdir = f'{parts}-{gen_hash(key)[:6]}' 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 @property
def is_materialized(self) -> bool: 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' _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]]: 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. """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]}' f'modelscope {spec.source!r}. Available (first 10): {files[:10]}'
) )
# shared blob store: download once per repo, hardlink into each cache entry # 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 = get_cache_root() / '.raw' / hashlib.md5(spec.source.encode()).hexdigest()[:10]
blob_dir = blob_dir / '.raw' / hashlib.md5(spec.source.encode()).hexdigest()[:10]
records: List[Dict[str, Any]] = [] records: List[Dict[str, Any]] = []
blobs: List[Path] = [] blobs: List[Path] = []
for path in selected: for path in selected: