93 lines
3.5 KiB
Python
93 lines
3.5 KiB
Python
"""Background image prefetcher: overlap docker pulls with evaluation.
|
|
|
|
While the runner executes sample N (its image already local), worker
|
|
threads pull the images of upcoming samples N+1.. — so per-sample runs
|
|
never wait on a cold multi-GB pull unless the queue drains.
|
|
|
|
from evalharness.sandbox.prefetch import BackgroundPrefetcher
|
|
images = images_for_dataset(ds) # ordered like the dataset
|
|
with BackgroundPrefetcher(images, workers=4, lookahead=8) as bp:
|
|
for sample in ds:
|
|
bp.ensure(sample.sandbox.image) # blocks only if still pulling
|
|
... run in sandbox ...
|
|
"""
|
|
|
|
import threading
|
|
import time
|
|
from typing import Dict, List, Optional
|
|
|
|
from .prefetch import _pull_one, local_images
|
|
|
|
|
|
class BackgroundPrefetcher:
|
|
"""Pull upcoming images on worker threads; evaluation thread consumes."""
|
|
|
|
def __init__(self, images: List[str], workers: int = 4, lookahead: int = 8):
|
|
self.images = list(images)
|
|
self.workers = max(1, workers)
|
|
self.lookahead = max(1, lookahead)
|
|
self._cursor = 0
|
|
self._lock = threading.Lock()
|
|
self._ready: Dict[str, bool] = {}
|
|
self._failed: Dict[str, str] = {}
|
|
self._stop = threading.Event()
|
|
self._threads: List[threading.Thread] = []
|
|
for img in self.images:
|
|
self._ready[img] = True if img in local_images() else False
|
|
|
|
def __enter__(self) -> 'BackgroundPrefetcher':
|
|
for i in range(self.workers):
|
|
t = threading.Thread(target=self._work, name=f'eh-prefetch-{i}', daemon=True)
|
|
t.start()
|
|
self._threads.append(t)
|
|
return self
|
|
|
|
def __exit__(self, *exc) -> None:
|
|
self._stop.set()
|
|
for t in self._threads:
|
|
t.join(timeout=5)
|
|
|
|
def _work(self) -> None:
|
|
while not self._stop.is_set():
|
|
img = self._next_pending()
|
|
if img is None:
|
|
time.sleep(0.5)
|
|
continue
|
|
try:
|
|
_pull_one(img)
|
|
with self._lock:
|
|
self._ready[img] = True
|
|
except Exception as e:
|
|
with self._lock:
|
|
self._failed[img] = str(e)[:200]
|
|
|
|
def _next_pending(self) -> Optional[str]:
|
|
"""Claim the next not-ready image within the lookahead window."""
|
|
with self._lock:
|
|
hi = min(self._cursor + self.lookahead, len(self.images))
|
|
for i in range(self._cursor, hi):
|
|
img = self.images[i]
|
|
if not self._ready.get(img) and img not in self._failed:
|
|
return img # claimed (pull is idempotent; duplicates are cheap)
|
|
return None
|
|
|
|
def ensure(self, image: Optional[str]) -> bool:
|
|
"""Advance the cursor to `image`; wait (bounded) until pulled."""
|
|
if not image:
|
|
return True
|
|
with self._lock: # allow random access ordering too
|
|
if image in self.images and self.images.index(image) >= self._cursor:
|
|
self._cursor = self.images.index(image)
|
|
deadline = time.time() + 3600
|
|
while time.time() < deadline and not self._stop.is_set():
|
|
with self._lock:
|
|
if self._ready.get(image) or image in self._failed:
|
|
return self._ready.get(image, False)
|
|
time.sleep(1.0)
|
|
return self._ready.get(image, False)
|
|
|
|
def stats(self) -> Dict[str, int]:
|
|
with self._lock:
|
|
ready = sum(1 for v in self._ready.values() if v)
|
|
return {'total': len(self.images), 'ready': ready, 'failed': len(self._failed)}
|