File size: 2,782 Bytes
1601a2f | 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 | #!/usr/bin/env python
"""Download Stage-2 dataset subsets under a hard disk budget.
Usage:
python scripts/download_subsets.py --root ~/tinyvla_data [--dry-run]
Downloads metadata first, measures actual on-disk size of a small episode
sample, then extrapolates before committing to the full subset download.
"""
from __future__ import annotations
import argparse
import shutil
from pathlib import Path
DISK_BUDGET_GB = 400
# (repo_id, num_episodes or None for all)
SUBSETS = [
("HuggingFaceVLA/community_dataset_v1", None),
("nvidia/BridgeData2_LeRobot_v3", 10_000),
("IPEC-COMMUNITY/fractal20220817_data_lerobot", 8_000),
]
def free_gb(path: Path) -> float:
return shutil.disk_usage(path).free / 1e9
def used_gb(path: Path) -> float:
return sum(f.stat().st_size for f in path.rglob("*") if f.is_file()) / 1e9
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--root", type=Path, default=Path.home() / "tinyvla_data")
parser.add_argument("--dry-run", action="store_true")
parser.add_argument("--only", type=str, default=None, help="substring filter on repo_id")
args = parser.parse_args()
args.root.mkdir(parents=True, exist_ok=True)
from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
for repo_id, n_eps in SUBSETS:
if args.only and args.only not in repo_id:
continue
print(f"\n=== {repo_id} (episodes: {n_eps or 'all'}) ===")
meta = LeRobotDatasetMetadata(repo_id)
total_eps = meta.total_episodes
print(f"total episodes: {total_eps}, fps: {meta.fps}")
episodes = list(range(min(n_eps, total_eps))) if n_eps else None
# probe with 1% to estimate size
probe_n = max(10, (len(episodes) if episodes else total_eps) // 100)
probe_dir = args.root / "_probe" / repo_id.replace("/", "__")
if not args.dry_run and not probe_dir.exists():
LeRobotDataset(repo_id, root=probe_dir, episodes=list(range(probe_n)))
probe_gb = used_gb(probe_dir)
est_gb = probe_gb / probe_n * (len(episodes) if episodes else total_eps)
print(f"probe: {probe_n} eps = {probe_gb:.2f}GB -> estimated full subset {est_gb:.0f}GB")
if est_gb > free_gb(args.root) - 100 or used_gb(args.root) + est_gb > DISK_BUDGET_GB:
print(f"SKIP {repo_id}: would exceed budget ({DISK_BUDGET_GB}GB) or disk")
continue
if args.dry_run:
continue
target = args.root / repo_id.replace("/", "__")
LeRobotDataset(repo_id, root=target, episodes=episodes)
print(f"downloaded {repo_id}: {used_gb(target):.1f}GB, total used {used_gb(args.root):.1f}GB")
if __name__ == "__main__":
main()
|