rpx_benchmark.hub

  1"""HuggingFace Hub integration for RPX benchmark.
  2
  3Task-aware downloads: fetches only the modalities a given task needs,
  4reusing HF's content-addressed cache so switching tasks on the same
  5scenes only pulls the new label files.
  6
  7Repo layout (on HF)::
  8
  9    rpx-benchmark/
 10    ├── metadata/
 11    │   ├── scenes.parquet
 12    │   └── esd_scores.parquet
 13    ├── manifests/
 14    │   └── <task>/<difficulty>.json      # logical views, not duplicates
 15    └── scenes/scene_000/{0,1,2}/
 16        ├── rgb/*.png
 17        ├── depth/*.png                   # 16-bit mm
 18        ├── mask/*.png                    # integer instance IDs
 19        ├── pose/*.npz
 20        ├── tracklets.json
 21        ├── questionnaires.json
 22        ├── spatial_qa.json
 23        └── general_qa.json
 24"""
 25
 26from __future__ import annotations
 27
 28import fcntl
 29import hashlib
 30import json
 31import os
 32import tarfile
 33from pathlib import Path
 34from typing import Any, Dict, Iterable, List, Sequence, Set, Tuple
 35
 36from .api import Difficulty, TaskType
 37from .exceptions import ConfigError, DatasetError, DownloadError, ManifestError
 38from .loader import RPXDataset
 39from .logging_utils import get_logger
 40
 41
 42def _file_sha256(path: Path) -> str:
 43    """Stream the file through SHA-256, returning the hexdigest."""
 44    h = hashlib.sha256()
 45    with path.open("rb") as f:
 46        for chunk in iter(lambda: f.read(1 << 20), b""):
 47            h.update(chunk)
 48    return h.hexdigest()
 49
 50
 51def _load_tar_checksums(snapshot_root: Path) -> Dict[str, str]:
 52    """Read ``manifest/checksums.json`` if it exists.
 53
 54    Returns a mapping ``{repo_relative_tar_path: sha256_hex}``. An
 55    empty dict means no checksum manifest was found — verification is
 56    then skipped (legacy uploads), but the rest of the pipeline still
 57    works.
 58    """
 59    cp = snapshot_root / "manifest" / "checksums.json"
 60    if not cp.is_file():
 61        return {}
 62    try:
 63        payload = json.loads(cp.read_text(encoding="utf-8"))
 64    except json.JSONDecodeError:
 65        return {}
 66    sha = payload.get("sha256") or {}
 67    return {str(k): str(v) for k, v in sha.items() if v}
 68
 69
 70def _load_file_checksums(snapshot_root: Path) -> Dict[str, str]:
 71    """Read ``manifest/file_checksums.json`` if it exists.
 72
 73    Returns a mapping ``{extracted_path: sha256_hex}``. Empty dict
 74    means legacy upload — ``verify_dataset`` will report that there is
 75    nothing to compare against.
 76    """
 77    cp = snapshot_root / "manifest" / "file_checksums.json"
 78    if not cp.is_file():
 79        return {}
 80    try:
 81        payload = json.loads(cp.read_text(encoding="utf-8"))
 82    except json.JSONDecodeError:
 83        return {}
 84    sha = payload.get("sha256") or {}
 85    return {str(k): str(v) for k, v in sha.items() if v}
 86
 87
 88def verify_dataset(
 89    snapshot_root: Path | str,
 90    *,
 91    raise_on_error: bool = True,
 92) -> Dict[str, str]:
 93    """Walk the extracted tree under ``snapshot_root`` and verify each
 94    file's SHA-256 against ``manifest/file_checksums.json``.
 95
 96    Catches the class of failure that tar-level verification cannot:
 97    local disk corruption (cosmic rays, bad sectors), accidental
 98    overwrites of extracted files, or any state where the on-disk
 99    bytes diverge from what the operator originally packed.
100
101    Returns ``{extracted_path: status}`` where ``status`` is one of:
102
103    * ``"ok"``       — file's SHA-256 matches the manifest exactly.
104    * ``"missing"``  — file is in the manifest but not on disk.
105    * ``"mismatch"`` — file is on disk but its SHA-256 does not match.
106
107    Parameters
108    ----------
109    snapshot_root
110        Root of an HF snapshot (the path returned by ``snapshot_download``,
111        or any local staging dir that ``_extract_snapshot_tars`` has run on).
112    raise_on_error
113        If ``True`` (default), any ``mismatch`` or ``missing`` raises
114        :class:`DatasetError` listing the first ten offenders. Set
115        ``False`` to get the full report dict instead.
116
117    Raises
118    ------
119    DatasetError
120        If ``raise_on_error`` and any file fails verification, **or**
121        if ``manifest/file_checksums.json`` is missing entirely
122        (legacy dataset — verification cannot be done; the caller
123        must explicitly skip verification by examining the empty
124        return from :func:`_load_file_checksums` themselves).
125    """
126    snapshot_root = Path(snapshot_root)
127    expected = _load_file_checksums(snapshot_root)
128    if not expected:
129        if raise_on_error:
130            raise DatasetError(
131                f"no file_checksums.json under {snapshot_root}/manifest/",
132                hint=(
133                    "Per-file verification requires a manifest produced "
134                    "by build_frame_manifest at PR #60 or later. Legacy "
135                    "uploads do not have per-file SHA-256 entries."
136                ),
137            )
138        return {}
139
140    report: Dict[str, str] = {}
141    bad: List[str] = []
142    for rel_path, expected_sha in expected.items():
143        on_disk = snapshot_root / rel_path
144        if not on_disk.is_file():
145            report[rel_path] = "missing"
146            bad.append(f"{rel_path}: missing")
147            continue
148        actual_sha = _file_sha256(on_disk)
149        if actual_sha != expected_sha:
150            report[rel_path] = "mismatch"
151            bad.append(f"{rel_path}: expected {expected_sha[:16]}..., got {actual_sha[:16]}...")
152        else:
153            report[rel_path] = "ok"
154
155    if bad and raise_on_error:
156        summary = "\n  ".join(bad[:10])
157        more = f"\n  ... and {len(bad) - 10} more" if len(bad) > 10 else ""
158        raise DatasetError(
159            f"verify_dataset: {len(bad)} / {len(expected)} files failed "
160            f"verification:\n  {summary}{more}",
161            hint=(
162                "Re-extract the offending tars from the snapshot, or "
163                "re-download. If only a small number of files are bad, "
164                "the most likely cause is local disk corruption — run "
165                "the verifier again after deleting the extracted dir."
166            ),
167        )
168    return report
169
170
171log = get_logger(__name__)
172
173DEFAULT_REPO_ID = os.environ.get("RPX_HF_REPO", "IRVLUTD/RPX")
174# Consolidated, audited RPX dataset release.  Pinning the default keeps two
175# benchmark runs comparable even while the Hub's ``main`` branch moves.
176# Operators can select another immutable release through RPX_HF_REVISION or
177# the public ``revision=`` arguments.
178DEFAULT_REVISION = os.environ.get(
179    "RPX_HF_REVISION", "51432e5d48895bbafb490ae7c5e34e0f6b7da7d6"
180)
181# The VQA question lists live in a later commit than the images they point at:
182# a list pinned to DEFAULT_REVISION cannot also be stored at that revision, since
183# a commit can't contain its own hash. Lists are read from VQA_MANIFEST_REVISION;
184# every image inside them is read from DEFAULT_REVISION (IRVLUTD/RPX#96).
185VQA_MANIFEST_REVISION = os.environ.get(
186    "RPX_VQA_MANIFEST_REVISION", "4f5d13bf37ebf3e50a382e747e9ea63163f037b6"
187)
188VQA_BENCHMARK_LIST = "vqa/metadata/benchmark_subset_v1/benchmark_available_31500.jsonl"
189REPO_TYPE = "dataset"
190
191# ------------------------------------------------------------------ #
192# Task → modality mapping
193# ------------------------------------------------------------------ #
194
195# Modality globs, relative to a (scene, phase) directory on the hub.
196RGB = "rgb/*"
197DEPTH = "depth/*"
198MASK = "mask/*"
199SAM2_META = "sam2_meta/*"
200POSE = "pose/*"
201FISHEYE_L = "fisheye_left/*"
202FISHEYE_R = "fisheye_right/*"
203SPARSE_DEPTH_DIR = "sparse_depth/*"
204KEYPOINTS_DIR = "keypoints/*"
205TRACKLETS = "tracklets.json"
206QUESTIONNAIRES = "questionnaires.json"
207SPATIAL_QA = "spatial_qa.json"
208GENERAL_QA = "general_qa.json"
209
210TASK_MODALITIES: Dict[TaskType, List[str]] = {
211    TaskType.MONOCULAR_DEPTH: [RGB, DEPTH],
212    # VideoDepthDataset consumes pose_filenames whenever the official
213    # manifest provides them, even for stride sampling. Fetch the pose shard
214    # with RGB/depth so a clip cannot fail after download_split succeeds.
215    TaskType.VIDEO_DEPTH: [RGB, DEPTH, POSE],
216    TaskType.SPARSE_DEPTH: [RGB, DEPTH, SPARSE_DEPTH_DIR],
217    TaskType.OBJECT_SEGMENTATION: [RGB, MASK],
218    TaskType.OBJECT_DETECTION: [RGB, MASK, TRACKLETS],
219    TaskType.OPEN_VOCAB_DETECTION: [RGB, MASK, TRACKLETS, QUESTIONNAIRES],
220    # Tracking GT is the temporally consistent instance-mask sequence.
221    # ``sam2_meta`` carries mask_to_object.json for auditable ID/name
222    # provenance.  The pinned RPX release contains no loose tracklets.json;
223    # the historical TRACKLETS entry therefore could never be downloaded.
224    TaskType.OBJECT_TRACKING: [RGB, MASK, SAM2_META],
225    TaskType.RELATIVE_CAMERA_POSE: [RGB, POSE],
226    TaskType.VISUAL_GROUNDING: [RGB, QUESTIONNAIRES, SPATIAL_QA],
227    TaskType.KEYPOINT_MATCHING: [RGB, KEYPOINTS_DIR],
228}
229
230# QA-only tasks we may add later (not in TaskType enum yet); keep mapping
231# for forward compatibility with the CLI.
232EXTRA_TASK_ALIASES: Dict[str, List[str]] = {
233    "qa_spatial": [RGB, DEPTH, SPATIAL_QA],
234    "qa_general": [RGB, GENERAL_QA],
235}
236
237
238# ------------------------------------------------------------------ #
239# Hub helpers
240# ------------------------------------------------------------------ #
241
242
243def _hub():
244    """Lazy-import ``huggingface_hub`` and re-raise as :class:`DownloadError`.
245
246    Raises
247    ------
248    DownloadError
249        If ``huggingface_hub`` is not installed.
250    """
251    try:
252        import huggingface_hub as hf
253    except ImportError as e:
254        raise DownloadError(
255            "huggingface_hub is required for RPX hub operations.",
256            hint="Install with: pip install 'rpx-benchmark[hub]'",
257        ) from e
258    return hf
259
260
261def _rpx_cache_dir() -> Path:
262    """Location for resolved-manifest files (outside the HF blob cache)."""
263    base = os.environ.get("RPX_CACHE_DIR")
264    if base:
265        return Path(base)
266    return Path.home() / ".cache" / "rpx_benchmark"
267
268
269def _manifest_repo_path(
270    task: TaskType | str,
271    split: Difficulty | str,
272    manifest_name: str | None = None,
273) -> str:
274    task_name = manifest_name or (task.value if isinstance(task, TaskType) else str(task))
275    split_name = split.value if isinstance(split, Difficulty) else str(split)
276    return f"manifests/{task_name}/{split_name}.json"
277
278
279def _modalities_for(task: TaskType | str) -> List[str]:
280    if isinstance(task, TaskType):
281        return TASK_MODALITIES[task]
282    if task in EXTRA_TASK_ALIASES:
283        return EXTRA_TASK_ALIASES[task]
284    return TASK_MODALITIES[TaskType(task)]
285
286
287# ------------------------------------------------------------------ #
288# Manifest handling
289# ------------------------------------------------------------------ #
290
291
292def fetch_manifest(
293    task: TaskType | str,
294    split: Difficulty | str,
295    repo_id: str = DEFAULT_REPO_ID,
296    cache_dir: str | Path | None = None,
297    revision: str | None = DEFAULT_REVISION,
298    manifest_name: str | None = None,
299) -> Dict[str, Any]:
300    """Download and parse the task-level manifest for ``(task, split)``.
301
302    Manifests are small (hundreds of KB) and are fetched eagerly so the
303    caller can discover which (scene, phase) dirs the split references
304    before kicking off a bulk download.
305
306    Parameters
307    ----------
308    task : TaskType or str
309    split : Difficulty or str
310    repo_id : str
311        HuggingFace dataset repo id. Defaults to
312        :data:`DEFAULT_REPO_ID` (``"IRVLUTD/RPX"``).
313    cache_dir : str or Path, optional
314    revision : str, optional
315
316    Returns
317    -------
318    dict
319        Parsed manifest JSON.
320
321    Raises
322    ------
323    DownloadError
324        If the download fails (network, auth, bad repo id) or the
325        manifest file does not exist on the hub.
326    ManifestError
327        If the downloaded file is not valid JSON.
328    """
329    hf = _hub()
330    repo_path = _manifest_repo_path(task, split, manifest_name)
331    try:
332        local = hf.hf_hub_download(
333            repo_id=repo_id,
334            repo_type=REPO_TYPE,
335            filename=repo_path,
336            cache_dir=str(cache_dir) if cache_dir else None,
337            revision=revision,
338        )
339    except Exception as e:
340        raise DownloadError(
341            f"Failed to download manifest {repo_path!r} from {repo_id}: {e}",
342            hint=(
343                "Check your network connection and that the repo id is "
344                "spelled correctly. Private repos need HF_TOKEN set."
345            ),
346        ) from e
347    try:
348        with open(local, "r", encoding="utf-8") as f:
349            return json.load(f)
350    except json.JSONDecodeError as e:
351        raise ManifestError(
352            f"Manifest file at {local} is not valid JSON.",
353        ) from e
354
355
356def _extract_scene_phase_pairs(manifest: Dict[str, Any]) -> Set[Tuple[str, str]]:
357    """Unique (scene_dir, phase_dir) pairs referenced by a manifest."""
358    pairs: Set[Tuple[str, str]] = set()
359    if "scenes" in manifest:
360        for entry in manifest["scenes"]:
361            pairs.add((str(entry["scene"]), str(entry["phase"])))
362        return pairs
363    for sample in manifest.get("samples", []):
364        # Prefer the physical path over the logical ``phase`` field. Ego
365        # manifests intentionally use phase=0 for metric grouping while their
366        # files live under scenes/<scene>/ego/ on the Hub.
367        found_path = False
368        for key in ("rgb", "depth", "mask"):
369            p = sample.get(key)
370            if not p:
371                continue
372            parts = Path(p).parts
373            try:
374                scene_part = parts.index("scenes")
375            except ValueError:
376                continue
377            if len(parts) >= scene_part + 3:
378                pairs.add((parts[scene_part + 1], parts[scene_part + 2]))
379                found_path = True
380                break
381        if not found_path and sample.get("scene_id") is not None and sample.get("phase") is not None:
382            pairs.add((str(sample["scene_id"]), str(sample["phase"])))
383    return pairs
384
385
386def _build_allow_patterns(
387    modalities: Sequence[str],
388    scene_phase_pairs: Iterable[Tuple[str, str]],
389) -> List[str]:
390    # Keep legacy extracted-file globs while also requesting the current
391    # lossless tar shards. A snapshot may contain one layout or the other.
392    tar_for_modality = {
393        RGB: "rgb.tar",
394        DEPTH: "depth.tar",
395        MASK: "labels/masks/v1.tar",
396        SAM2_META: "labels/sam2_meta/v1.tar",
397        POSE: "labels/cam_pose/v1.tar",
398        FISHEYE_L: "fisheye.tar",
399        FISHEYE_R: "fisheye.tar",
400    }
401    patterns: List[str] = []
402    for scene, phase in sorted(scene_phase_pairs):
403        base = f"scenes/{scene}/{phase}"
404        for m in modalities:
405            patterns.append(f"{base}/{m}")
406            tar_path = tar_for_modality.get(m)
407            if tar_path is not None:
408                patterns.append(f"{base}/{tar_path}")
409    return list(dict.fromkeys(patterns))
410
411
412# ------------------------------------------------------------------ #
413# Public download API
414# ------------------------------------------------------------------ #
415
416
417def download_split(
418    task: TaskType | str,
419    split: Difficulty | str,
420    repo_id: str = DEFAULT_REPO_ID,
421    cache_dir: str | Path | None = None,
422    revision: str | None = DEFAULT_REVISION,
423    extra_modalities: Sequence[str] | None = None,
424    max_workers: int = 8,
425    max_samples: int | None = None,
426    manifest_name: str | None = None,
427) -> Path:
428    """Download only the files (task, split) needs, return resolved manifest path.
429
430    The resolved manifest is a JSON file whose ``root`` field points to
431    the local HF snapshot directory, so it can be fed directly to
432    :meth:`RPXDataset.from_manifest`.
433    """
434    hf = _hub()
435    task_enum = (
436        TaskType(task) if isinstance(task, str) and task in TaskType._value2member_map_ else task
437    )
438    split_enum = Difficulty(split) if isinstance(split, str) else split
439
440    manifest = fetch_manifest(
441        task_enum,
442        split_enum,
443        repo_id,
444        cache_dir,
445        revision,
446        manifest_name,
447    )
448    if max_samples is not None:
449        if max_samples < 1:
450            raise ConfigError(
451                f"max_samples must be >= 1, got {max_samples}",
452                hint="Use None for the complete split.",
453            )
454        manifest = dict(manifest)
455        manifest["samples"] = list(manifest.get("samples") or [])[:max_samples]
456
457    pairs = _extract_scene_phase_pairs(manifest)
458    if not pairs:
459        raise ManifestError(
460            f"Manifest {task}/{split} references no scenes; cannot derive download patterns.",
461            hint="This usually means the manifest was generated against "
462            "an empty scene list — re-run the upload script.",
463        )
464
465    modalities = list(_modalities_for(task_enum))
466    if extra_modalities:
467        modalities.extend(extra_modalities)
468
469    allow_patterns = _build_allow_patterns(modalities, pairs)
470    allow_patterns.append(_manifest_repo_path(task_enum, split_enum, manifest_name))
471    allow_patterns.append("manifest/checksums.json")
472
473    log.info(
474        "downloading %d file patterns for task=%s split=%s from %s",
475        len(allow_patterns),
476        task_enum.value if isinstance(task_enum, TaskType) else task_enum,
477        split_enum.value if isinstance(split_enum, Difficulty) else split_enum,
478        repo_id,
479    )
480    try:
481        snapshot_root = hf.snapshot_download(
482            repo_id=repo_id,
483            repo_type=REPO_TYPE,
484            allow_patterns=allow_patterns,
485            cache_dir=str(cache_dir) if cache_dir else None,
486            revision=revision,
487            max_workers=max_workers,
488        )
489    except Exception as e:
490        raise DownloadError(
491            f"snapshot_download failed for {repo_id}: {e}",
492            hint="Rerun with --cache-dir pointing at a writable directory "
493            "or set HF_HUB_OFFLINE=1 to use a prebuilt local cache.",
494        ) from e
495
496    # The HF tree ships tar shards; the per-task per-split JSON we just
497    # downloaded references *extracted* PNG/NPZ paths. Run the extraction
498    # step here so the user can hand the manifest to RPXDataset directly.
499    # Idempotent — files already on disk are skipped.
500    n_extracted, n_skipped = _extract_snapshot_tars(Path(snapshot_root))
501    log.info(
502        "extracted %d new files (%d already on disk) from tar shards",
503        n_extracted,
504        n_skipped,
505    )
506
507    resolved = dict(manifest)
508    resolved["root"] = str(snapshot_root)
509    resolved.setdefault(
510        "task",
511        task_enum.value if isinstance(task_enum, TaskType) else str(task_enum),
512    )
513
514    task_name = manifest_name or (
515        task_enum.value if isinstance(task_enum, TaskType) else str(task_enum)
516    )
517    out_dir = _rpx_cache_dir() / repo_id.replace("/", "__") / "manifests" / task_name
518    out_dir.mkdir(parents=True, exist_ok=True)
519    out_path = (
520        out_dir / f"{split_enum.value if isinstance(split_enum, Difficulty) else split_enum}.json"
521    )
522    with out_path.open("w", encoding="utf-8") as f:
523        json.dump(resolved, f)
524    return out_path
525
526
527def load(
528    task: TaskType | str,
529    split: Difficulty | str,
530    repo_id: str = DEFAULT_REPO_ID,
531    cache_dir: str | Path | None = None,
532    revision: str | None = DEFAULT_REVISION,
533    batch_size: int = 1,
534    max_samples: int | None = None,
535) -> RPXDataset:
536    """Download (task, split) and return an iterable :class:`RPXDataset`.
537
538    Incremental re-use::
539
540        # First run: fetches rgb + depth for 'hard' scenes.
541        depth_ds = rpx.load("monocular_depth", "hard")
542
543        # Second run: rgb/depth already cached, only spatial_qa.json fetched.
544        qa_ds    = rpx.load("visual_grounding", "hard")
545    """
546    manifest_path = download_split(
547        task=task,
548        split=split,
549        repo_id=repo_id,
550        cache_dir=cache_dir,
551        revision=revision,
552        max_samples=max_samples,
553    )
554    return RPXDataset.from_manifest(
555        manifest_path,
556        batch_size=batch_size,
557        max_samples=max_samples,
558    )
559
560
561# ------------------------------------------------------------------ #
562# fsspec mount (preview / debug only — do not use for training)
563# ------------------------------------------------------------------ #
564
565
566def mount(repo_id: str = DEFAULT_REPO_ID):
567    """Return an ``HfFileSystem`` rooted at the RPX repo for lazy browsing.
568
569    Each read goes over the network; prefer :func:`load` for real workloads.
570    """
571    hf = _hub()
572    fs = hf.HfFileSystem()
573    return fs, f"datasets/{repo_id}"
574
575
576# ------------------------------------------------------------------ #
577# Tar extraction (post-download)
578# ------------------------------------------------------------------ #
579
580
581def _extract_snapshot_tars_unlocked(snapshot_root: Path) -> Tuple[int, int]:
582    """Extract every tar shard under ``snapshot_root`` into ``snapshot_root/extracted/``.
583
584    The HF dataset tree ships tar shards (``scenes/<scene>/<phase>/rgb.tar``,
585    ``scenes/<scene>/<phase>/labels/masks/v1.tar``, ...). Per-task per-split
586    manifests reference *extracted* paths
587    (``extracted/scenes/<scene>/<phase>/rgb/<frame>.png``,
588     ``extracted/scenes/<scene>/<phase>/sam2/masks/<frame>.png``, ...). This
589    helper materialises that layout post-download.
590
591    Idempotency
592    -----------
593    For each tar member, we check if the destination file already exists
594    (correct size). Only missing or wrong-size files are extracted. The
595    write is atomic-ish: extract to ``<dest>.part`` then ``rename`` so a
596    half-written file never collides with a re-run.
597
598    Path mapping
599    ------------
600    Tar location ``scenes/<scene>/<phase>/<...>.tar`` → extraction root
601    ``extracted/scenes/<scene>/<phase>/`` (with the tar's *own* member
602    names appended). Examples:
603
604    * ``scenes/scene1/0/rgb.tar``               members ``rgb/00000.png``  → ``extracted/scenes/scene1/0/rgb/00000.png``
605    * ``scenes/scene1/0/labels/masks/v1.tar``   members ``sam2/masks/00000.png`` → ``extracted/scenes/scene1/0/sam2/masks/00000.png``
606    * ``scenes/scene1/0/labels/cam_pose/v1.tar`` members ``cam_pose/00000.npz`` → ``extracted/scenes/scene1/0/cam_pose/00000.npz``
607
608    Parameters
609    ----------
610    snapshot_root
611        Root of an HF snapshot (the path returned by ``snapshot_download``).
612
613    Returns
614    -------
615    (n_extracted, n_skipped)
616        Number of files newly extracted and number already present.
617    """
618    extracted_root = snapshot_root / "extracted"
619    # Load the manifest's checksums.json (written by build_frame_manifest)
620    # and use it to verify each tar's SHA-256 before extraction. If the
621    # file is missing (legacy datasets) we skip verification — the worst
622    # case is "no defense beyond HF's own SHA check" rather than "abort".
623    expected_checksums = _load_tar_checksums(snapshot_root)
624    n_new = 0
625    n_skip = 0
626    for tar_path in sorted(snapshot_root.rglob("*.tar")):
627        if extracted_root in tar_path.parents:
628            # Don't re-process artefacts already under extracted/ from a
629            # previous run that happened to leave tars behind.
630            continue
631        try:
632            rel = tar_path.relative_to(snapshot_root)
633        except ValueError:  # symlinks pointing outside; skip
634            continue
635        # Fault-proof gate: if this tar has a declared SHA-256, the
636        # bytes on disk must match before any byte is extracted. Any
637        # mismatch is a hard error — we refuse to extract bytes that
638        # could feed corrupted data into downstream benchmarks.
639        rel_posix = rel.as_posix()
640        if expected_checksums and rel_posix in expected_checksums:
641            expected = expected_checksums[rel_posix]
642            if expected is not None:
643                actual = _file_sha256(tar_path)
644                if actual != expected:
645                    raise DatasetError(
646                        (
647                            f"SHA-256 mismatch for {rel_posix}: "
648                            f"expected {expected[:16]}..., got {actual[:16]}..."
649                        ),
650                        hint=(
651                            "The downloaded tar does not match the operator's "
652                            "manifest. Re-download (delete the snapshot dir "
653                            "and re-run download_for_task) or contact the "
654                            "publisher."
655                        ),
656                    )
657        # Locate the physical capture prefix. MOS uses numeric phase names,
658        # while the ego-preview release stores each capture under the literal
659        # directory ``scenes/<scene>/ego``. Requiring ``str.isdigit()`` here
660        # silently skipped every downloaded ego tar even though the manifest
661        # and allow-pattern generation were correct.
662        if len(rel.parts) >= 4 and rel.parts[0] == "scenes":
663            prefix_parts = list(rel.parts[:3])
664        else:
665            # Retain the legacy numeric-prefix fallback for old snapshots that
666            # did not use the canonical scenes/<scene>/<capture>/ layout.
667            prefix_parts = []
668            for p in rel.parts[:-1]:  # stop before the .tar filename
669                prefix_parts.append(p)
670                if p.isdigit():
671                    break
672            if not prefix_parts or not prefix_parts[-1].isdigit():
673                # Unrelated tar at repo root; skip rather than guess.
674                continue
675        out_base = extracted_root.joinpath(*prefix_parts)
676        out_base.mkdir(parents=True, exist_ok=True)
677        try:
678            with tarfile.open(tar_path, "r") as tf:
679                for member in tf.getmembers():
680                    if not member.isfile():
681                        continue
682                    out = out_base / member.name
683                    if out.exists() and out.stat().st_size == member.size:
684                        n_skip += 1
685                        continue
686                    out.parent.mkdir(parents=True, exist_ok=True)
687                    tmp = out.with_suffix(out.suffix + ".part")
688                    f = tf.extractfile(member)
689                    if f is None:
690                        continue
691                    with tmp.open("wb") as g:
692                        # Stream-copy in chunks; member.size can be ~MB
693                        # for masks, no point loading whole thing in RAM.
694                        while True:
695                            chunk = f.read(1 << 20)
696                            if not chunk:
697                                break
698                            g.write(chunk)
699                    tmp.rename(out)
700                    n_new += 1
701        except tarfile.TarError as e:
702            log.warning("failed to extract %s: %s", tar_path, e)
703            continue
704    return n_new, n_skip
705
706
707def _extract_snapshot_tars(snapshot_root: Path) -> Tuple[int, int]:
708    """Serialize extraction into a shared HF snapshot across processes.
709
710    Tracking jobs deliberately share one content-addressed cache. Without a
711    snapshot-level lock, concurrent processes can overwrite and rename the
712    same ``.part`` file, causing one process to fail after another moves it.
713    ``flock`` keeps the existing idempotent extraction implementation while
714    ensuring only one process mutates ``extracted/`` at a time.
715    """
716
717    snapshot_root = Path(snapshot_root)
718    lock_path = snapshot_root / ".rpx-extraction.lock"
719    with lock_path.open("a+b") as lock:
720        fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
721        try:
722            return _extract_snapshot_tars_unlocked(snapshot_root)
723        finally:
724            fcntl.flock(lock.fileno(), fcntl.LOCK_UN)