migration: add prune_destination.py to remove errored / mislabeled-SO datasets

Reads the run manifest(s), selects datasets that errored during migration and/or
were labelled SO but aren't a real 6-DOF SO arm, intersects with what's actually
present in the destination repo, and deletes their folders. Dry-run by default.

community_v3_migration/prune_destination.py

@@ -0,0 +1,101 @@
+"""Prune datasets from the migrated destination repo based on the run manifest(s).
+
+Selects, from the manifest CSV(s) written by run_migration.py / slurm_migrate.py, the datasets to
+remove and deletes their folders from the destination repo. Two independent filters:
+
+  --errored        rows whose migration action starts with "ERROR:" (failed to convert; usually
+                   absent from the repo, but any partial upload left behind is cleaned up).
+  --mislabeled-so  rows whose robot_type claims SO but the dataset isn't a real 6-DOF SO arm
+                   (action_dim not a multiple of 6, or a classification note saying so).
+
+Only folders actually present in the destination repo are touched. Dry-run by default; pass --yes
+to perform the deletions.
+
+  python prune_destination.py --manifest /fsx/$USER/cdv3_manifests/manifest_*.csv \
+      --dst-repo lerobot/community_dataset_v3 --errored --mislabeled-so        # dry-run
+  python prune_destination.py --manifest manifest_*.csv --errored --mislabeled-so --yes
+"""
+import argparse
+import sys
+
+import pandas as pd
+from huggingface_hub import HfApi
+
+from classify import is_so_robot_type
+
+DST_REPO = "lerobot/community_dataset_v3"
+
+
+def _is_errored(row) -> bool:
+    return str(row.get("action") or "").strip().upper().startswith("ERROR")
+
+
+def _is_mislabeled_so(row) -> bool:
+    if "claims SO but" in str(row.get("note") or ""):
+        return True
+    if not is_so_robot_type(str(row.get("robot_type") or "")):
+        return False
+    try:
+        return int(float(row["action_dim"])) % 6 != 0
+    except (TypeError, ValueError, KeyError):
+        return False
+
+
+def select(df: pd.DataFrame, present: set[str], errored: bool, mislabeled: bool) -> dict[str, str]:
+    """Map each dataset root that is present in the repo AND matches an enabled filter to a reason.
+    'errored' takes precedence over 'mislabeled-so' when a row matches both."""
+    out: dict[str, str] = {}
+    for _, row in df.iterrows():
+        root = str(row.get("root") or "").strip()
+        if not root or root not in present or root in out:
+            continue
+        if errored and _is_errored(row):
+            out[root] = "errored"
+        elif mislabeled and _is_mislabeled_so(row):
+            out[root] = "mislabeled-so"
+    return out
+
+
+def main():
+    ap = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--manifest", nargs="+", required=True, metavar="CSV",
+                    help="One or more manifest CSVs from the migration run (per-rank files ok).")
+    ap.add_argument("--dst-repo", default=DST_REPO, metavar="ORG/NAME",
+                    help="Destination HF dataset repo to prune.")
+    ap.add_argument("--errored", action="store_true",
+                    help="Delete datasets that errored during migration.")
+    ap.add_argument("--mislabeled-so", action="store_true",
+                    help="Delete datasets labelled SO but not a real 6-DOF SO arm.")
+    ap.add_argument("--yes", action="store_true",
+                    help="Actually delete. Without it, only prints what would be deleted (dry-run).")
+    args = ap.parse_args()
+    if not (args.errored or args.mislabeled_so):
+        ap.error("enable at least one filter: --errored and/or --mislabeled-so")
+
+    df = pd.concat([pd.read_csv(p) for p in args.manifest], ignore_index=True)
+    if "root" not in df.columns:
+        ap.error("manifest has no 'root' column; is this a run_migration.py manifest?")
+
+    api = HfApi()
+    present = {p[: -len("/meta/info.json")]
+               for p in api.list_repo_files(args.dst_repo, repo_type="dataset")
+               if p.endswith("/meta/info.json")}
+
+    to_delete = select(df, present, args.errored, args.mislabeled_so)
+    for root, reason in sorted(to_delete.items()):
+        print(f"{reason:14s} {root}")
+    print(f"\n{len(to_delete)} dataset(s) present in {args.dst_repo} match "
+          f"({df['root'].nunique()} rows in manifest, {len(present)} datasets in repo).", file=sys.stderr)
+
+    if not args.yes:
+        print("dry-run: nothing deleted. re-run with --yes to delete.", file=sys.stderr)
+        return
+    for root, reason in sorted(to_delete.items()):
+        api.delete_folder(path_in_repo=root, repo_id=args.dst_repo, repo_type="dataset",
+                          commit_message=f"Prune {root} ({reason})")
+        print(f"deleted {root} ({reason})", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()