rpx_benchmark.dataset_hub.uploader

  1"""Upload a packed staging directory to a HuggingFace dataset repo.
  2
  3Wraps :func:`huggingface_hub.upload_large_folder` (the workhorse for
  4multi-hundred-GB pushes — chunked, multi-commit, resumable). For small
  5mock uploads / dry runs we fall back to :func:`upload_folder`.
  6
  7The uploader does *no* packing. Run ``pack_capture_tree`` first, then
  8``build_frame_manifest``, then point this at the resulting staging dir.
  9
 10Resumability: ``upload_large_folder`` writes its checkpoint state in
 11``staging_dir/.huggingface/`` so a re-invocation after a crash picks up
 12where it left off.
 13"""
 14
 15from __future__ import annotations
 16
 17from dataclasses import dataclass, field
 18from pathlib import Path
 19from typing import List, Optional, Sequence
 20
 21from ..exceptions import ConfigError, DownloadError
 22from ..logging_utils import get_logger
 23from .recipes import DEFAULT_REPO_ID
 24
 25log = get_logger(__name__)
 26
 27
 28# Patterns we always exclude from upload, regardless of caller settings.
 29# Build artefacts and editor leftovers should never end up in the repo.
 30DEFAULT_IGNORE_PATTERNS: tuple[str, ...] = (
 31    ".DS_Store",
 32    "__pycache__",
 33    "*.pyc",
 34    "*.tmp",
 35    "*.part",
 36    ".huggingface/*",  # HF's own checkpoint state
 37)
 38
 39
 40@dataclass(frozen=True)
 41class UploadPlan:
 42    """Inputs needed to push a staging directory to a HF dataset repo."""
 43
 44    staging_root: Path
 45    repo_id: str = DEFAULT_REPO_ID
 46    revision: str = "main"
 47    private: bool = True
 48    commit_message: str = "RPX dataset hub upload"
 49    create_if_missing: bool = True
 50    use_large_folder: bool = True
 51    allow_patterns: Optional[Sequence[str]] = None
 52    ignore_patterns: Sequence[str] = DEFAULT_IGNORE_PATTERNS
 53    dry_run: bool = False
 54
 55
 56@dataclass(frozen=True)
 57class UploadResult:
 58    """Reported back so callers can log / verify."""
 59
 60    repo_id: str
 61    revision: str
 62    files_planned: int
 63    bytes_planned: int
 64    commit_url: Optional[str] = None
 65    skipped: List[Path] = field(default_factory=list)
 66
 67
 68def _hub():
 69    try:
 70        import huggingface_hub as hf  # noqa: WPS433
 71    except ImportError as e:
 72        raise DownloadError(
 73            "huggingface_hub is required for RPX uploads.",
 74            hint="Install with: pip install 'rpx-benchmark[hub]'",
 75        ) from e
 76    return hf
 77
 78
 79def _enumerate_uploadable(
 80    root: Path,
 81    ignore: Sequence[str],
 82    allow: Optional[Sequence[str]] = None,
 83) -> tuple[List[Path], int]:
 84    """Walk ``root``, return (files, total_bytes), respecting ignore + allow.
 85
 86    ``allow`` mirrors ``huggingface_hub`` semantics: when set, *only*
 87    paths matching at least one allow glob are considered. ``ignore``
 88    is always applied. Both lists match against the POSIX path relative
 89    to ``root`` *and* against the bare file name.
 90    """
 91    import fnmatch
 92
 93    keep: List[Path] = []
 94    total = 0
 95    for p in root.rglob("*"):
 96        if not p.is_file():
 97            continue
 98        rel = p.relative_to(root).as_posix()
 99        if any(fnmatch.fnmatch(rel, pat) or fnmatch.fnmatch(p.name, pat) for pat in ignore):
100            continue
101        if allow is not None and not any(
102            fnmatch.fnmatch(rel, pat) or fnmatch.fnmatch(p.name, pat) for pat in allow
103        ):
104            continue
105        keep.append(p)
106        total += p.stat().st_size
107    return keep, total
108
109
110def upload_staging(plan: UploadPlan) -> UploadResult:
111    """Push everything under ``plan.staging_root`` to the HF dataset repo.
112
113    Behaviour:
114
115    * If the repo doesn't exist and ``create_if_missing`` is True, it is
116      created (private by default, dataset type).
117    * If ``dry_run`` is True, no network call is made — we walk the tree
118      and report what *would* be uploaded.
119    * Otherwise, ``upload_large_folder`` is preferred for true resumable
120      multi-commit pushes; ``upload_folder`` is the fallback for tiny
121      mock uploads where the multi-commit overhead isn't worth it.
122    """
123    if not plan.staging_root.is_dir():
124        raise ConfigError(
125            f"staging dir does not exist: {plan.staging_root}",
126            hint="Run pack_capture_tree() and build_frame_manifest() first.",
127        )
128
129    files, total = _enumerate_uploadable(
130        plan.staging_root,
131        plan.ignore_patterns,
132        plan.allow_patterns,
133    )
134    log.info(
135        "upload plan: %d files, %d bytes → %s (revision=%s, dry_run=%s)",
136        len(files),
137        total,
138        plan.repo_id,
139        plan.revision,
140        plan.dry_run,
141    )
142
143    if plan.dry_run:
144        return UploadResult(
145            repo_id=plan.repo_id,
146            revision=plan.revision,
147            files_planned=len(files),
148            bytes_planned=total,
149        )
150
151    hf = _hub()
152    api = hf.HfApi()
153
154    if plan.create_if_missing:
155        try:
156            api.create_repo(
157                repo_id=plan.repo_id,
158                repo_type="dataset",
159                private=plan.private,
160                exist_ok=True,
161            )
162        except Exception as e:
163            raise DownloadError(
164                f"failed to create or verify dataset repo {plan.repo_id}: {e}",
165                hint="Run `hf auth login` and confirm you can push to the org.",
166            ) from e
167
168    common_kwargs = dict(
169        repo_id=plan.repo_id,
170        repo_type="dataset",
171        folder_path=str(plan.staging_root),
172        revision=plan.revision,
173        ignore_patterns=list(plan.ignore_patterns),
174    )
175    if plan.allow_patterns:
176        common_kwargs["allow_patterns"] = list(plan.allow_patterns)
177
178    try:
179        if plan.use_large_folder and hasattr(api, "upload_large_folder"):
180            # `upload_large_folder` does its own multi-commit chunking
181            # and writes checkpoint state under .huggingface/ so re-runs
182            # resume cleanly.
183            api.upload_large_folder(**common_kwargs)
184            commit_url = None
185        else:
186            commit = api.upload_folder(
187                commit_message=plan.commit_message,
188                **common_kwargs,
189            )
190            commit_url = getattr(commit, "commit_url", None)
191    except Exception as e:
192        raise DownloadError(
193            f"upload to {plan.repo_id} failed: {e}",
194            hint=(
195                "Re-run to resume; HF stores checkpoint state under "
196                "<staging>/.huggingface/. If the failure is auth, run "
197                "`hf auth login` and confirm push access to the org."
198            ),
199        ) from e
200
201    return UploadResult(
202        repo_id=plan.repo_id,
203        revision=plan.revision,
204        files_planned=len(files),
205        bytes_planned=total,
206        commit_url=commit_url,
207    )
208
209
210# Allow uploading just one or a few patterns (e.g. only the manifest, or
211# only one modality) without materialising a separate staging dir.
212def upload_paths(
213    paths: Sequence[Path],
214    staging_root: Path,
215    repo_id: str = DEFAULT_REPO_ID,
216    revision: str = "main",
217    private: bool = True,
218    commit_message: str = "RPX dataset hub: partial upload",
219    dry_run: bool = False,
220) -> UploadResult:
221    """Upload a specific list of files (relative to ``staging_root``).
222
223    Useful for shipping a single modality or just the manifest after a
224    partial re-pack. Internally builds an ``allow_patterns`` list and
225    delegates to :func:`upload_staging` (which still calls
226    ``upload_large_folder``/``upload_folder`` under the hood).
227    """
228    rels = []
229    total = 0
230    for p in paths:
231        rel = p.relative_to(staging_root).as_posix()
232        rels.append(rel)
233        total += p.stat().st_size
234
235    plan = UploadPlan(
236        staging_root=staging_root,
237        repo_id=repo_id,
238        revision=revision,
239        private=private,
240        commit_message=commit_message,
241        allow_patterns=tuple(rels),
242        use_large_folder=False,
243        dry_run=dry_run,
244    )
245    return upload_staging(plan)