File size: 2,543 Bytes
fbd9366 | 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 | #!/usr/bin/env python3
"""
Download DreamZero DROID dataset from Hugging Face with rate-limit handling.
Use this script when `huggingface-cli download` hits 429 (Too Many Requests).
It uses a single worker and retries with backoff to stay within the 3000 req/5min limit.
Usage:
python scripts/data/download_droid_hf.py [--local-dir ./data/droid_lerobot] [--max-workers 1]
# Or with env (same as CLI default):
DROID_DATA_ROOT=./data/droid_lerobot python scripts/data/download_droid_hf.py
"""
from __future__ import annotations
import argparse
import os
import sys
import time
try:
from huggingface_hub import snapshot_download
except ImportError:
print("Install huggingface_hub: pip install huggingface_hub", file=sys.stderr)
sys.exit(1)
REPO_ID = "GEAR-Dreams/DreamZero-DROID-Data"
REPO_TYPE = "dataset"
def main() -> None:
p = argparse.ArgumentParser(description="Download DreamZero DROID dataset with rate-limit handling.")
p.add_argument(
"--local-dir",
default=os.environ.get("DROID_DATA_ROOT", "./data/droid_lerobot"),
help="Local directory to download into (default: DROID_DATA_ROOT or ./data/droid_lerobot)",
)
p.add_argument(
"--max-workers",
type=int,
default=1,
help="Concurrent download threads (default 1 to reduce API requests and avoid 429)",
)
p.add_argument(
"--retry-wait",
type=int,
default=320,
help="Seconds to wait on 429 before retry (default 320 ≈ 5min)",
)
args = p.parse_args()
local_dir = os.path.abspath(args.local_dir)
os.makedirs(local_dir, exist_ok=True)
attempt = 0
while True:
attempt += 1
try:
print(f"Download attempt {attempt} (max_workers={args.max_workers})...")
snapshot_download(
repo_id=REPO_ID,
repo_type=REPO_TYPE,
local_dir=local_dir,
local_dir_use_symlinks=False,
max_workers=args.max_workers,
resume_download=True,
)
print(f"Done. Dataset at: {local_dir}")
return
except Exception as e:
err_str = str(e).lower()
if "429" in err_str or "too many requests" in err_str or "rate limit" in err_str:
print(f"Rate limited (429). Waiting {args.retry_wait}s before retry...", file=sys.stderr)
time.sleep(args.retry_wait)
continue
raise
if __name__ == "__main__":
main()
|