Spaces:
Running
Running
| from __future__ import annotations | |
| import os | |
| import pickle | |
| import time | |
| from typing import TYPE_CHECKING | |
| from fsspec.utils import atomic_write | |
| try: | |
| import ujson as json | |
| except ImportError: | |
| if not TYPE_CHECKING: | |
| import json | |
| if TYPE_CHECKING: | |
| from typing import Any, Dict, Iterator, Literal | |
| from typing_extensions import TypeAlias | |
| from .cached import CachingFileSystem | |
| Detail: TypeAlias = Dict[str, Any] | |
| class CacheMetadata: | |
| """Cache metadata. | |
| All reading and writing of cache metadata is performed by this class, | |
| accessing the cached files and blocks is not. | |
| Metadata is stored in a single file per storage directory in JSON format. | |
| For backward compatibility, also reads metadata stored in pickle format | |
| which is converted to JSON when next saved. | |
| """ | |
| def __init__(self, storage: list[str]): | |
| """ | |
| Parameters | |
| ---------- | |
| storage: list[str] | |
| Directories containing cached files, must be at least one. Metadata | |
| is stored in the last of these directories by convention. | |
| """ | |
| if not storage: | |
| raise ValueError("CacheMetadata expects at least one storage location") | |
| self._storage = storage | |
| self.cached_files: list[Detail] = [{}] | |
| # Private attribute to force saving of metadata in pickle format rather than | |
| # JSON for use in tests to confirm can read both pickle and JSON formats. | |
| self._force_save_pickle = False | |
| def _load(self, fn: str) -> Detail: | |
| """Low-level function to load metadata from specific file""" | |
| try: | |
| with open(fn, "r") as f: | |
| loaded = json.load(f) | |
| except ValueError: | |
| with open(fn, "rb") as f: | |
| loaded = pickle.load(f) | |
| for c in loaded.values(): | |
| if isinstance(c.get("blocks"), list): | |
| c["blocks"] = set(c["blocks"]) | |
| return loaded | |
| def _save(self, metadata_to_save: Detail, fn: str) -> None: | |
| """Low-level function to save metadata to specific file""" | |
| if self._force_save_pickle: | |
| with atomic_write(fn) as f: | |
| pickle.dump(metadata_to_save, f) | |
| else: | |
| with atomic_write(fn, mode="w") as f: | |
| json.dump(metadata_to_save, f) | |
| def _scan_locations( | |
| self, writable_only: bool = False | |
| ) -> Iterator[tuple[str, str, bool]]: | |
| """Yield locations (filenames) where metadata is stored, and whether | |
| writable or not. | |
| Parameters | |
| ---------- | |
| writable: bool | |
| Set to True to only yield writable locations. | |
| Returns | |
| ------- | |
| Yields (str, str, bool) | |
| """ | |
| n = len(self._storage) | |
| for i, storage in enumerate(self._storage): | |
| writable = i == n - 1 | |
| if writable_only and not writable: | |
| continue | |
| yield os.path.join(storage, "cache"), storage, writable | |
| def check_file( | |
| self, path: str, cfs: CachingFileSystem | None | |
| ) -> Literal[False] | tuple[Detail, str]: | |
| """If path is in cache return its details, otherwise return ``False``. | |
| If the optional CachingFileSystem is specified then it is used to | |
| perform extra checks to reject possible matches, such as if they are | |
| too old. | |
| """ | |
| for (fn, base, _), cache in zip(self._scan_locations(), self.cached_files): | |
| if path not in cache: | |
| continue | |
| detail = cache[path].copy() | |
| if cfs is not None: | |
| if cfs.check_files and detail["uid"] != cfs.fs.ukey(path): | |
| # Wrong file as determined by hash of file properties | |
| continue | |
| if cfs.expiry and time.time() - detail["time"] > cfs.expiry: | |
| # Cached file has expired | |
| continue | |
| fn = os.path.join(base, detail["fn"]) | |
| if os.path.exists(fn): | |
| return detail, fn | |
| return False | |
| def clear_expired(self, expiry_time: int) -> tuple[list[str], bool]: | |
| """Remove expired metadata from the cache. | |
| Returns names of files corresponding to expired metadata and a boolean | |
| flag indicating whether the writable cache is empty. Caller is | |
| responsible for deleting the expired files. | |
| """ | |
| expired_files = [] | |
| for path, detail in self.cached_files[-1].copy().items(): | |
| if time.time() - detail["time"] > expiry_time: | |
| fn = detail.get("fn", "") | |
| if not fn: | |
| raise RuntimeError( | |
| f"Cache metadata does not contain 'fn' for {path}" | |
| ) | |
| fn = os.path.join(self._storage[-1], fn) | |
| expired_files.append(fn) | |
| self.cached_files[-1].pop(path) | |
| if self.cached_files[-1]: | |
| cache_path = os.path.join(self._storage[-1], "cache") | |
| self._save(self.cached_files[-1], cache_path) | |
| writable_cache_empty = not self.cached_files[-1] | |
| return expired_files, writable_cache_empty | |
| def load(self) -> None: | |
| """Load all metadata from disk and store in ``self.cached_files``""" | |
| cached_files = [] | |
| for fn, _, _ in self._scan_locations(): | |
| if os.path.exists(fn): | |
| # TODO: consolidate blocks here | |
| cached_files.append(self._load(fn)) | |
| else: | |
| cached_files.append({}) | |
| self.cached_files = cached_files or [{}] | |
| def on_close_cached_file(self, f: Any, path: str) -> None: | |
| """Perform side-effect actions on closing a cached file. | |
| The actual closing of the file is the responsibility of the caller. | |
| """ | |
| # File must be writeble, so in self.cached_files[-1] | |
| c = self.cached_files[-1][path] | |
| if c["blocks"] is not True and len(c["blocks"]) * f.blocksize >= f.size: | |
| c["blocks"] = True | |
| def pop_file(self, path: str) -> str | None: | |
| """Remove metadata of cached file. | |
| If path is in the cache, return the filename of the cached file, | |
| otherwise return ``None``. Caller is responsible for deleting the | |
| cached file. | |
| """ | |
| details = self.check_file(path, None) | |
| if not details: | |
| return None | |
| _, fn = details | |
| if fn.startswith(self._storage[-1]): | |
| self.cached_files[-1].pop(path) | |
| self.save() | |
| else: | |
| raise PermissionError( | |
| "Can only delete cached file in last, writable cache location" | |
| ) | |
| return fn | |
| def save(self) -> None: | |
| """Save metadata to disk""" | |
| for (fn, _, writable), cache in zip(self._scan_locations(), self.cached_files): | |
| if not writable: | |
| continue | |
| if os.path.exists(fn): | |
| cached_files = self._load(fn) | |
| for k, c in cached_files.items(): | |
| if k in cache: | |
| if c["blocks"] is True or cache[k]["blocks"] is True: | |
| c["blocks"] = True | |
| else: | |
| # self.cached_files[*][*]["blocks"] must continue to | |
| # point to the same set object so that updates | |
| # performed by MMapCache are propagated back to | |
| # self.cached_files. | |
| blocks = cache[k]["blocks"] | |
| blocks.update(c["blocks"]) | |
| c["blocks"] = blocks | |
| c["time"] = max(c["time"], cache[k]["time"]) | |
| c["uid"] = cache[k]["uid"] | |
| # Files can be added to cache after it was written once | |
| for k, c in cache.items(): | |
| if k not in cached_files: | |
| cached_files[k] = c | |
| else: | |
| cached_files = cache | |
| cache = {k: v.copy() for k, v in cached_files.items()} | |
| for c in cache.values(): | |
| if isinstance(c["blocks"], set): | |
| c["blocks"] = list(c["blocks"]) | |
| self._save(cache, fn) | |
| self.cached_files[-1] = cached_files | |
| def update_file(self, path: str, detail: Detail) -> None: | |
| """Update metadata for specific file in memory, do not save""" | |
| self.cached_files[-1][path] = detail | |