Spaces:
Running on RTX PRO 6000
Running on RTX PRO 6000
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """TTL + LRU cache with single-flight loading (plan §5.2). | |
| Single-flight is not optional here. A crawl ``Enumerate``s and then ``Stat``s | |
| every item, and a thumbnail probe fires up to four ``Stat``s per asset for | |
| candidate ``.thumbs`` paths — all against the same folder, concurrently. Without | |
| collapsing concurrent loads of one key onto a single upstream request, laziness | |
| turns one preload into a burst of duplicates. | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import time | |
| from collections import OrderedDict | |
| from collections.abc import Awaitable, Callable, Hashable | |
| from dataclasses import dataclass, field | |
| from typing import Generic, TypeVar | |
| K = TypeVar("K", bound=Hashable) | |
| V = TypeVar("V") | |
| class Missing: | |
| """A cached negative result. | |
| Distinct from ``None`` so a miss can be cached and served without re-hitting | |
| HuggingFace: the thumbnail probe generates a steady stream of legitimate | |
| misses that would otherwise be upstream traffic every time. | |
| """ | |
| __slots__ = () | |
| def __bool__(self) -> bool: | |
| return False | |
| MISSING = Missing() | |
| class CacheStats: | |
| hits: int = 0 | |
| misses: int = 0 | |
| loads: int = 0 | |
| negative_hits: int = 0 | |
| evictions: int = 0 | |
| def hit_rate(self) -> float: | |
| total = self.hits + self.misses | |
| return self.hits / total if total else 0.0 | |
| class _Entry(Generic[V]): | |
| value: V | Missing | |
| expires_at: float | |
| class SingleFlightCache(Generic[K, V]): | |
| """Async cache: at most one in-flight load per key, TTL + LRU eviction.""" | |
| def __init__( | |
| self, | |
| *, | |
| ttl: float, | |
| negative_ttl: float, | |
| max_size: int, | |
| clock: Callable[[], float] = time.monotonic, | |
| ) -> None: | |
| self._ttl = ttl | |
| self._negative_ttl = negative_ttl | |
| self._max_size = max_size | |
| self._clock = clock | |
| self._entries: OrderedDict[K, _Entry[V]] = OrderedDict() | |
| self._inflight: dict[K, asyncio.Future[V | Missing]] = {} | |
| self.stats = CacheStats() | |
| def _get_live(self, key: K) -> _Entry[V] | None: | |
| entry = self._entries.get(key) | |
| if entry is None: | |
| return None | |
| if entry.expires_at <= self._clock(): | |
| del self._entries[key] | |
| return None | |
| self._entries.move_to_end(key) | |
| return entry | |
| def peek(self, key: K) -> V | Missing | None: | |
| """Return a cached value without loading. ``None`` means "not cached".""" | |
| entry = self._get_live(key) | |
| return entry.value if entry else None | |
| def put(self, key: K, value: V | Missing) -> None: | |
| ttl = self._negative_ttl if isinstance(value, Missing) else self._ttl | |
| self._entries[key] = _Entry(value=value, expires_at=self._clock() + ttl) | |
| self._entries.move_to_end(key) | |
| while len(self._entries) > self._max_size: | |
| self._entries.popitem(last=False) | |
| self.stats.evictions += 1 | |
| async def get(self, key: K, loader: Callable[[], Awaitable[V | Missing]]) -> V | Missing: | |
| """Return the cached value, or load it — once, however many callers race.""" | |
| entry = self._get_live(key) | |
| if entry is not None: | |
| self.stats.hits += 1 | |
| if isinstance(entry.value, Missing): | |
| self.stats.negative_hits += 1 | |
| return entry.value | |
| self.stats.misses += 1 | |
| inflight = self._inflight.get(key) | |
| if inflight is not None: | |
| # Someone else is already loading this key: wait for their result | |
| # rather than issuing a duplicate upstream request. | |
| return await asyncio.shield(inflight) | |
| future: asyncio.Future[V | Missing] = asyncio.get_running_loop().create_future() | |
| self._inflight[key] = future | |
| try: | |
| self.stats.loads += 1 | |
| value = await loader() | |
| except BaseException as exc: | |
| # Failures are not cached — a transient 503 must not pin a miss for | |
| # the whole TTL. | |
| if not future.done(): | |
| future.set_exception(exc) | |
| raise | |
| else: | |
| self.put(key, value) | |
| if not future.done(): | |
| future.set_result(value) | |
| return value | |
| finally: | |
| self._inflight.pop(key, None) | |
| # Nobody may be awaiting a failed future; keep the loop quiet about it. | |
| if future.done() and future.exception() is not None: | |
| future.exception() | |
| def clear(self) -> None: | |
| self._entries.clear() | |
| def __len__(self) -> int: | |
| return len(self._entries) | |
| class TokenBucket: | |
| """Simple rate limiter for upstream calls (plan §8.5).""" | |
| rate: float | |
| capacity: float | |
| _tokens: float = field(default=0.0, init=False) | |
| _updated: float = field(default=0.0, init=False) | |
| def __post_init__(self) -> None: | |
| self._tokens = self.capacity | |
| self._updated = time.monotonic() | |
| async def acquire(self) -> None: | |
| while True: | |
| now = time.monotonic() | |
| self._tokens = min(self.capacity, self._tokens + (now - self._updated) * self.rate) | |
| self._updated = now | |
| if self._tokens >= 1: | |
| self._tokens -= 1 | |
| return | |
| await asyncio.sleep((1 - self._tokens) / self.rate) | |