Add --workers parallel push (ThreadPoolExecutor)

scripts/push_metadata_stubs.py

@@ -33,11 +33,13 @@ from __future__ import annotations
 
 import argparse
 import ast
+import concurrent.futures
 import json
 import logging
 import os
 import sys
 import tempfile
+import threading
 import time
 import urllib.error
 import urllib.request
@@ -754,6 +756,12 @@ def main(argv: list[str] | None = None) -> int:
     parser.add_argument("--dry-run-out", type=Path, default=Path("/tmp/stub_preview"))
     parser.add_argument("--private", action="store_true")
     parser.add_argument("--token", default=os.environ.get("HF_TOKEN"))
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=1,
+        help="Parallel pushes (IO-bound — 8-16 is safe; higher risks rate-limits).",
+    )
     parser.add_argument("-v", "--verbose", action="count", default=0)
     args = parser.parse_args(argv)
 
@@ -790,28 +798,58 @@ def main(argv: list[str] | None = None) -> int:
         logger.info("Dry-run output: %s", args.dry_run_out)
         return 0
 
-    failed: list[tuple[str, str]] = []
+    pending = [r for r in rows if str(r["dataset"]).lower() not in existing]
     for r in rows:
         slug = str(r["dataset"]).lower()
         if slug in existing:
             logger.info("skipping %s (exists)", slug)
-            continue
+
+    failed: list[tuple[str, str]] = []
+    done = 0
+    done_lock = threading.Lock()
+
+    def _one(r: pd.Series) -> tuple[str, Exception | None]:
+        slug = str(r["dataset"]).lower()
         try:
             ctx = _build_context(r)
-            repo_id = _push_one(ctx, args)
-            logger.info("pushed %s", repo_id)
+            _push_one(ctx, args)
+            return slug, None
         except Exception as exc:  # noqa: BLE001
-            logger.exception("failed %s", slug)
-            failed.append((slug, str(exc)))
-        # Be polite to the API and HF.
-        time.sleep(0.25)
+            return slug, exc
+
+    if args.workers and args.workers > 1:
+        logger.info(
+            "parallel push: %d workers, %d pending", args.workers, len(pending)
+        )
+        with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as pool:
+            futures = {pool.submit(_one, r): r for r in pending}
+            for fut in concurrent.futures.as_completed(futures):
+                slug, err = fut.result()
+                if err is None:
+                    with done_lock:
+                        done += 1
+                    logger.info("pushed EEGDash/%s (%d/%d)", slug, done, len(pending))
+                else:
+                    logger.exception("failed %s", slug, exc_info=err)
+                    failed.append((slug, str(err)))
+    else:
+        for r in pending:
+            slug, err = _one(r)
+            if err is None:
+                done += 1
+                logger.info("pushed EEGDash/%s (%d/%d)", slug, done, len(pending))
+            else:
+                logger.exception("failed %s", slug, exc_info=err)
+                failed.append((slug, str(err)))
+            # Serial mode only — parallel mode doesn't need the spacer.
+            time.sleep(0.15)
 
     if failed:
         logger.error("%d failures:", len(failed))
         for slug, err in failed:
             logger.error("  %s — %s", slug, err)
         return 1
-    logger.info("done — %d stubs processed", len(rows) - len(existing))
+    logger.info("done — %d stubs processed (%d skipped)", done, len(existing))
     return 0