Unify cache root override: loader.get/set_cache_root as single source, fixes --cache-dir not relocating .raw blobs
This commit is contained in:
parent
f8cd15fea1
commit
414a89216c
@ -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:
|
||||||
|
|||||||
@ -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:
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user