rpx_benchmark.dataset_hub.staging

  1"""File-staging helpers — non-tar artefacts that ship to the HF repo root.
  2
  3The packer/manifest/uploader chain handles per-(scene, phase, modality)
  4tar shards and the per-frame Parquet manifest. This module covers the
  5small companion files HF Datasets needs alongside them:
  6
  7* ``splits/{scene_splits.json, easy.txt, medium.txt, hard.txt}`` —
  8  copied from ``benchmark/data/splits/`` so downloaders can resolve a
  9  split to a scene list before fetching the manifest.
 10* ``rpx_croissant.json`` — copied from
 11  ``paper-submission/croissant/`` with path reconciliation
 12  against the new tar layout (TODO 3 — handled in a sibling helper).
 13* ``README.md`` (HF dataset card) — generated from the scan totals
 14  (TODO 2 — handled in :mod:`.dataset_card`).
 15* ``preview/`` (HF Dataset Viewer thumbnails) — handled in
 16  :mod:`.preview` (TODO 4).
 17
 18Each helper returns a list of :class:`StagedFile` records so the upload
 19step can verify what landed.
 20"""
 21
 22from __future__ import annotations
 23
 24import json
 25import shutil
 26from dataclasses import dataclass
 27from pathlib import Path
 28from typing import List, Optional, Sequence
 29
 30from ..exceptions import ConfigError, DatasetError
 31from ..logging_utils import get_logger
 32
 33log = get_logger(__name__)
 34
 35
 36# Default location of the splits files relative to the repo root. Resolved
 37# lazily so the helper still works when called from a non-repo cwd.
 38def _default_splits_src() -> Path:
 39    """``<repo_root>/benchmark/data/splits/``."""
 40    here = Path(__file__).resolve()
 41    # __file__ -> .../benchmark/rpx_benchmark/dataset_hub/staging.py
 42    benchmark_root = here.parents[2]
 43    return benchmark_root / "data" / "splits"
 44
 45
 46SPLIT_TIER_FILES = ("easy.txt", "medium.txt", "hard.txt")
 47SCENE_SPLITS_JSON = "scene_splits.json"
 48
 49
 50@dataclass(frozen=True)
 51class StagedFile:
 52    """One file the staging step copied / wrote."""
 53
 54    repo_path: str  # path relative to staging root, forward-slashes
 55    src: Optional[Path]
 56    bytes_: int
 57
 58
 59def stage_splits(
 60    staging_root: Path,
 61    splits_src: Optional[Path] = None,
 62    overwrite: bool = False,
 63    require_all: bool = True,
 64) -> List[StagedFile]:
 65    """Copy ``splits/{scene_splits.json, easy.txt, medium.txt, hard.txt}``
 66    into ``<staging_root>/splits/``.
 67
 68    Parameters
 69    ----------
 70    staging_root : Path
 71        Where the packed/upload-ready tree lives.
 72    splits_src : Path, optional
 73        Source dir holding the splits files. Defaults to
 74        ``benchmark/data/splits/`` next to this package.
 75    overwrite : bool
 76        If False (default), raises when a destination file already exists.
 77    require_all : bool
 78        If True (default), raises when any of the four expected files is
 79        missing from ``splits_src``. Set False for partial-stage scenarios
 80        (e.g., shipping with only the JSON manifest).
 81
 82    Returns
 83    -------
 84    List[StagedFile]
 85        One entry per copied file.
 86    """
 87    src_root = Path(splits_src) if splits_src else _default_splits_src()
 88    if not src_root.is_dir():
 89        raise ConfigError(
 90            f"splits source dir does not exist: {src_root}",
 91            hint="Splits files are bundled on the HuggingFace dataset "
 92            "(splits/{scene_splits.json,easy.txt,medium.txt,hard.txt}); "
 93            "make sure you've downloaded them, or pass splits_src= "
 94            "pointing at a local copy.",
 95        )
 96
 97    staging_root = Path(staging_root)
 98    out_dir = staging_root / "splits"
 99    out_dir.mkdir(parents=True, exist_ok=True)
100
101    expected = (SCENE_SPLITS_JSON, *SPLIT_TIER_FILES)
102    missing = [name for name in expected if not (src_root / name).is_file()]
103    if missing and require_all:
104        raise DatasetError(
105            f"missing splits files in {src_root}: {missing}",
106            hint=(
107                "Re-download the HuggingFace dataset to recover the "
108                "splits/ tree, or pass require_all=False to ship a "
109                "partial splits dir."
110            ),
111        )
112
113    staged: List[StagedFile] = []
114    for name in expected:
115        src = src_root / name
116        if not src.is_file():
117            log.warning("splits source missing: %s (skipping)", src)
118            continue
119        dst = out_dir / name
120        if dst.exists() and not overwrite:
121            raise DatasetError(
122                f"refusing to overwrite existing staged file: {dst}",
123                hint="Pass overwrite=True or remove the file first.",
124            )
125        shutil.copy2(src, dst)
126        staged.append(
127            StagedFile(
128                repo_path=f"splits/{name}",
129                src=src,
130                bytes_=dst.stat().st_size,
131            )
132        )
133        log.info("staged splits/%s (%d bytes)", name, dst.stat().st_size)
134    return staged
135
136
137def stage_paths(
138    paths: Sequence[Path],
139    staging_root: Path,
140    repo_subdir: str = "",
141    overwrite: bool = False,
142) -> List[StagedFile]:
143    """Generic helper: copy a list of files into ``<staging>/<repo_subdir>/``.
144
145    Useful for ad-hoc staging of small files (Croissant JSON, custom
146    docs) without writing a dedicated module. ``repo_subdir`` may be
147    empty to land files at the staging root.
148    """
149    staging_root = Path(staging_root)
150    target = staging_root / repo_subdir if repo_subdir else staging_root
151    target.mkdir(parents=True, exist_ok=True)
152
153    staged: List[StagedFile] = []
154    for src in paths:
155        src = Path(src)
156        if not src.is_file():
157            raise ConfigError(
158                f"source file does not exist: {src}",
159                hint="Check the path; staging does not create files.",
160            )
161        dst = target / src.name
162        if dst.exists() and not overwrite:
163            raise DatasetError(
164                f"refusing to overwrite existing staged file: {dst}",
165                hint="Pass overwrite=True or remove the file first.",
166            )
167        shutil.copy2(src, dst)
168        rel = (Path(repo_subdir) / src.name).as_posix() if repo_subdir else src.name
169        staged.append(
170            StagedFile(
171                repo_path=rel,
172                src=src,
173                bytes_=dst.stat().st_size,
174            )
175        )
176        log.info("staged %s (%d bytes)", rel, dst.stat().st_size)
177    return staged
178
179
180def load_scene_splits(staging_root: Path) -> dict[str, str]:
181    """Read ``<staging_root>/splits/scene_splits.json`` and return a
182    flat ``{scene_id: tier}`` mapping. Convenience for downstream steps
183    (manifest builder, dataset card) that need to look up a scene's tier.
184
185    Accepts both the canonical schema with a top-level ``"splits"`` key
186    holding ``{tier: [scene_id, ...]}`` and a flat ``{scene_id: tier}``
187    shape.
188    """
189    path = Path(staging_root) / "splits" / SCENE_SPLITS_JSON
190    if not path.is_file():
191        raise ConfigError(
192            f"scene_splits.json not found at {path}",
193            hint="Run stage_splits() before load_scene_splits().",
194        )
195    payload = json.loads(path.read_text(encoding="utf-8"))
196    if "splits" in payload and isinstance(payload["splits"], dict):
197        return {sid: tier for tier, ids in payload["splits"].items() for sid in ids}
198    return dict(payload)