Spaces:
Running
Running
| """Byte-bounded disk cache and HTTP range reader. | |
| Shared verbatim with the OCR half (06-multilingual/ocr): the caching problem is identical once the unit | |
| is "a parquet row group", and a divergent copy would be a second thing to get right. | |
| """ | |
| import fcntl | |
| import hashlib | |
| import io | |
| import os | |
| import threading | |
| import time | |
| import uuid | |
| from collections import Counter | |
| from contextlib import contextmanager | |
| from pathlib import Path | |
| import requests | |
| class DiskCache: | |
| def __init__(self, directory, max_bytes, max_entry_bytes=None): | |
| self.directory = Path(directory) | |
| self.directory.mkdir(parents=True, exist_ok=True) | |
| self.max_bytes = int(max_bytes) | |
| self.max_entry_bytes = min(int(max_entry_bytes or max_bytes), self.max_bytes) | |
| if self.max_entry_bytes < 1: | |
| raise ValueError("Cache budgets must be positive") | |
| self.counters = Counter() | |
| self._stats_lock = threading.Lock() | |
| # Recover staging space after a process was killed before its finally block. | |
| # A live writer owns this stripe, including writers in another process. | |
| for path in self.directory.glob("*.part"): | |
| stripe = path.name[:3] | |
| with (self.directory / f"stripe-{stripe}.lock").open("a+b") as lock: | |
| try: | |
| fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) | |
| except BlockingIOError: | |
| continue | |
| path.unlink(missing_ok=True) | |
| def increment(self, key, value=1): | |
| with self._stats_lock: | |
| self.counters[key] += value | |
| def stats(self): | |
| with self._stats_lock: | |
| result = dict(self.counters) | |
| files = list(self.directory.glob("*.bin")) | |
| sizes = [] | |
| for path in files: | |
| try: | |
| sizes.append(path.stat().st_size) | |
| except FileNotFoundError: | |
| pass | |
| return { | |
| **result, | |
| "bytes": sum(sizes), | |
| "entries": len(sizes), | |
| "max_bytes": self.max_bytes, | |
| "max_entry_bytes": self.max_entry_bytes, | |
| } | |
| def lease(self, key, loader): | |
| digest = hashlib.sha256(key.encode()).hexdigest() | |
| path = self.directory / f"{digest}.bin" | |
| # Fixed lock stripes avoid millions of permanent lock files over full epochs. | |
| # A stripe covers both loading and reading; collisions only serialize work. | |
| stripe = digest[:3] | |
| with (self.directory / f"stripe-{stripe}.lock").open("a+b") as lock: | |
| fcntl.flock(lock, fcntl.LOCK_EX) | |
| if path.exists(): | |
| self.increment("hits") | |
| else: | |
| self.increment("misses") | |
| temporary = self.directory / f"{digest}.{uuid.uuid4().hex}.part" | |
| try: | |
| loader(temporary) | |
| size = temporary.stat().st_size | |
| if size > self.max_entry_bytes: | |
| raise ValueError( | |
| f"Cache entry {size} exceeds {self.max_entry_bytes} byte limit" | |
| ) | |
| with (self.directory / ".eviction.lock").open("a+b") as eviction: | |
| fcntl.flock(eviction, fcntl.LOCK_EX) | |
| self._make_room(size, stripe) | |
| os.replace(temporary, path) | |
| self.increment("loads") | |
| self.increment("loaded_bytes", size) | |
| except BaseException: | |
| self.increment("load_failures") | |
| raise | |
| finally: | |
| temporary.unlink(missing_ok=True) | |
| os.utime(path, None) | |
| yield path | |
| def _make_room(self, needed, owned_stripe): | |
| entries = [ | |
| (p.stat().st_mtime_ns, p.stat().st_size, p) | |
| for p in self.directory.glob("*.bin") | |
| ] | |
| used = sum(size for _, size, _ in entries) | |
| for _, size, path in sorted(entries): | |
| if used + needed <= self.max_bytes: | |
| return | |
| stripe = path.stem[:3] | |
| if stripe == owned_stripe: | |
| # Our exclusive stripe already proves no other reader holds this entry. | |
| path.unlink(missing_ok=True) | |
| used -= size | |
| self.increment("evictions") | |
| continue | |
| with (self.directory / f"stripe-{stripe}.lock").open("a+b") as lock: | |
| try: | |
| fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) | |
| except BlockingIOError: | |
| continue | |
| path.unlink(missing_ok=True) | |
| used -= size | |
| self.increment("evictions") | |
| if used + needed > self.max_bytes: | |
| raise RuntimeError( | |
| "Cache budget is occupied by active readers; reduce prefetch or increase cache bytes" | |
| ) | |
| class RangeReader(io.RawIOBase): | |
| """Seekable HTTP object: require exact 206 responses and meter actual payload bytes.""" | |
| def __init__(self, url, size, *, headers=None, on_read=None, timeout=120): | |
| super().__init__() | |
| self.url, self.size, self.position = url, size, 0 | |
| self.headers = headers or {} | |
| self.on_read = on_read or (lambda amount: None) | |
| self.timeout = timeout | |
| self.session = requests.Session() | |
| self._lock = threading.RLock() | |
| def readable(self): | |
| return True | |
| def seekable(self): | |
| return True | |
| def tell(self): | |
| return self.position | |
| def seek(self, offset, whence=0): | |
| with self._lock: | |
| position = ( | |
| offset | |
| if whence == 0 | |
| else self.position + offset | |
| if whence == 1 | |
| else self.size + offset | |
| ) | |
| if position < 0 or whence not in (0, 1, 2): | |
| raise ValueError("Invalid seek") | |
| self.position = position | |
| return position | |
| def read(self, size=-1): | |
| with self._lock: | |
| start = self.position | |
| stop = self.size if size < 0 else min(start + size, self.size) | |
| if start >= stop: | |
| return b"" | |
| for attempt in range(3): | |
| try: | |
| with self.session.get( | |
| self.url, | |
| headers={**self.headers, "Range": f"bytes={start}-{stop - 1}"}, | |
| timeout=(15, self.timeout), | |
| stream=True, | |
| ) as response: | |
| response.raise_for_status() | |
| if response.status_code != 206: | |
| raise ValueError( | |
| "Storage ignored byte range; refusing a whole-shard download" | |
| ) | |
| if ( | |
| response.headers.get("Content-Range") | |
| != f"bytes {start}-{stop - 1}/{self.size}" | |
| ): | |
| raise ValueError( | |
| "Storage returned an unexpected byte range or object size" | |
| ) | |
| data = bytearray() | |
| for chunk in response.iter_content(1024 * 1024): | |
| self.on_read(len(chunk)) | |
| data.extend(chunk) | |
| if len(data) > stop - start: | |
| raise ValueError( | |
| "Storage returned more bytes than requested" | |
| ) | |
| if len(data) != stop - start: | |
| raise requests.ConnectionError( | |
| "Incomplete byte-range response" | |
| ) | |
| self.position = stop | |
| return bytes(data) | |
| except (requests.ConnectionError, requests.Timeout): | |
| if attempt == 2: | |
| raise | |
| time.sleep(0.25 * (attempt + 1)) | |
| def readinto(self, buffer): | |
| value = self.read(len(buffer)) | |
| buffer[: len(value)] = value | |
| return len(value) | |
| def close(self): | |
| self.session.close() | |
| super().close() | |