50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
"""Fetch preview cache on another machine if the stand's public IP is rate limited."""
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import json
|
|
from pathlib import Path
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
|
|
cache = Path(__file__).parent / 'preview-cache'
|
|
cache.mkdir(exist_ok=True)
|
|
|
|
|
|
def fetch(url):
|
|
for attempt in range(6):
|
|
try:
|
|
with urllib.request.urlopen(url, timeout=30) as response:
|
|
return response.read()
|
|
except Exception:
|
|
if attempt == 5:
|
|
raise
|
|
time.sleep(min(30, 3 * (attempt + 1)))
|
|
|
|
|
|
def metadata(args):
|
|
split, offset = args
|
|
path = cache / f'rows-{split}-{offset}.json'
|
|
if not path.exists():
|
|
query = urllib.parse.urlencode(dict(dataset='uoft-cs/cifar10', config='plain_text',
|
|
split=split, offset=offset, length=100))
|
|
path.write_bytes(fetch('https://datasets-server.huggingface.co/rows?' + query))
|
|
time.sleep(0.5)
|
|
return [(split, r) for r in json.loads(path.read_text())['rows'] if r['row']['label'] in (3, 5)]
|
|
|
|
|
|
def image(args):
|
|
split, row = args
|
|
path = cache / f'{split}-{row["row_idx"]}.jpg'
|
|
if not path.exists():
|
|
path.write_bytes(fetch(row['row']['img']['src']))
|
|
|
|
|
|
tasks = [('train', i) for i in range(0, 6000, 100)] + [('test', i) for i in range(0, 3000, 100)]
|
|
with ThreadPoolExecutor(max_workers=2) as pool:
|
|
images = [item for batch in pool.map(metadata, tasks) for item in batch]
|
|
print('metadata ready', len(images), flush=True)
|
|
with ThreadPoolExecutor(max_workers=4) as pool:
|
|
for i, _ in enumerate(pool.map(image, images)):
|
|
if i % 200 == 0:
|
|
print('images', i, flush=True)
|
|
print('Cache ready', flush=True)
|