Spaces:
Running on RTX PRO 6000
Running on RTX PRO 6000
File size: 5,502 Bytes
1079f5b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | # 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()
@dataclass
class CacheStats:
hits: int = 0
misses: int = 0
loads: int = 0
negative_hits: int = 0
evictions: int = 0
@property
def hit_rate(self) -> float:
total = self.hits + self.misses
return self.hits / total if total else 0.0
@dataclass
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)
@dataclass
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)
|