AXONVERTEX-AI-RESEARCH's picture
Publish GraphShieldMistral interactive explorer (part 2)
7ad3625 verified
Raw
History Blame Contribute Delete
5.84 kB
# Copyright 2026-present, the HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""On-disk cache for repository tree listings.
A tree listing is the set of files (with their download metadata) contained in a repo at a given commit. Because a
commit hash is immutable, its tree listing never changes and can be cached forever without any invalidation logic.
The listing is stored as a human-readable JSON file under `<tree_cache_folder>/trees/<commit_hash>.json`. The
folder depends on the download target: the per-repo `storage_folder` for `cache_dir` downloads, or
`local_dir/.cache/huggingface/` for `local_dir` downloads (see `tree_cache_folder_for_local_dir`).
```json
{
"format_version": 1,
"files": {
"config.json": {"size": 519, "blob_id": "<git sha1>"},
"model.safetensors": {"size": 1234, "blob_id": "...", "lfs_sha256": "<sha256>", "lfs_size": 1234, "xet_hash": "..."}
}
}
```
"""
import json
import os
import tempfile
import threading
from dataclasses import dataclass
from .utils import logging
from .utils._xet import is_valid_xet_hash
logger = logging.get_logger(__name__)
TREE_CACHE_FORMAT_VERSION = 1
# In-memory cache of parsed tree listings, keyed by absolute file path.
_IN_MEMORY_TREE_CACHE: dict[str, "dict[str, TreeCacheEntry]"] = {}
_IN_MEMORY_TREE_CACHE_LOCK = threading.Lock()
@dataclass(frozen=True)
class TreeCacheEntry:
"""Raw metadata of a single file in a cached tree listing, mirroring the `/tree` endpoint fields."""
size: int
blob_id: str
lfs_sha256: str | None = None
lfs_size: int | None = None
xet_hash: str | None = None
def to_json(self) -> dict:
info: dict = {"size": self.size, "blob_id": self.blob_id}
if self.lfs_sha256 is not None:
info["lfs_sha256"] = self.lfs_sha256
info["lfs_size"] = self.lfs_size
if self.xet_hash is not None:
info["xet_hash"] = self.xet_hash
return info
@classmethod
def from_json(cls, info: dict) -> "TreeCacheEntry":
return cls(
size=info["size"],
blob_id=info["blob_id"],
lfs_sha256=info.get("lfs_sha256"),
lfs_size=info.get("lfs_size"),
xet_hash=info.get("xet_hash"),
)
def is_valid_tree_entries(entries: dict[str, TreeCacheEntry]) -> bool:
"""Return whether all Xet hashes in the tree listing are valid."""
return all(entry.xet_hash is None or is_valid_xet_hash(entry.xet_hash) for entry in entries.values())
def _tree_cache_path(tree_cache_folder: str, commit_hash: str) -> str:
return os.path.join(tree_cache_folder, "trees", f"{commit_hash}.json")
def tree_cache_folder_for_local_dir(local_dir: str) -> str:
"""Folder under which the `trees/` cache lives for a `local_dir` download."""
return os.path.join(local_dir, ".cache", "huggingface")
def read_tree_cache(tree_cache_folder: str, commit_hash: str) -> dict[str, TreeCacheEntry] | None:
"""Return the cached tree listing for a commit hash, or `None` if not cached, invalid, or unreadable."""
path = _tree_cache_path(tree_cache_folder, commit_hash)
with _IN_MEMORY_TREE_CACHE_LOCK:
if path in _IN_MEMORY_TREE_CACHE:
cached_entries = _IN_MEMORY_TREE_CACHE[path]
return cached_entries if is_valid_tree_entries(cached_entries) else None
entries = _read_tree_cache_from_disk(path)
if entries is None or not is_valid_tree_entries(entries):
return None
with _IN_MEMORY_TREE_CACHE_LOCK:
_IN_MEMORY_TREE_CACHE[path] = entries
return entries
def _read_tree_cache_from_disk(path: str) -> dict[str, TreeCacheEntry] | None:
try:
with open(path, encoding="utf-8") as f:
data = json.load(f)
if data.get("format_version") != TREE_CACHE_FORMAT_VERSION:
# Unknown format (e.g. written by a newer version) => ignore and re-fetch.
return None
return {file_path: TreeCacheEntry.from_json(info) for file_path, info in data["files"].items()}
except FileNotFoundError:
return None
except (OSError, ValueError, KeyError, TypeError) as e:
logger.warning(f"Ignoring corrupted tree cache file {path}: {e}")
return None
def write_tree_cache(tree_cache_folder: str, commit_hash: str, entries: dict[str, TreeCacheEntry]) -> None:
"""Write a valid tree listing to the cache (ignoring invalid entries and any failures)."""
if not is_valid_tree_entries(entries):
return
path = _tree_cache_path(tree_cache_folder, commit_hash)
data = {
"format_version": TREE_CACHE_FORMAT_VERSION,
"files": {file_path: entries[file_path].to_json() for file_path in sorted(entries)},
}
try:
os.makedirs(os.path.dirname(path), exist_ok=True)
tmp_fd, tmp_path = tempfile.mkstemp(dir=os.path.dirname(path), suffix=".tmp")
with os.fdopen(tmp_fd, "w", encoding="utf-8") as f:
json.dump(data, f, indent=1)
os.replace(tmp_path, path)
except OSError as e:
logger.warning(f"Ignored error while writing tree cache file {path}: {e}")
return
# Seed the in-memory cache so later readers of this commit skip re-reading and re-parsing the file.
with _IN_MEMORY_TREE_CACHE_LOCK:
_IN_MEMORY_TREE_CACHE[path] = dict(entries)