scripts/publish_dataset.py

#!/usr/bin/env python
"""Thin CLI over :class:`fpgm.datagen.publish.DatasetPublisher`.

Stages the shippable tier of ``outputs/datagen`` into a local directory,
writes ``dataset_manifest.json`` + ``README.md``, and (unless ``--dry-run``,
which is the default) syncs it to the configured HF bucket.

``--dry-run`` is on by default deliberately: an accidental invocation of this
script must never publish anything. Pass ``--no-dry-run`` to actually upload.

Usage:
    PYTHONPATH=src python scripts/publish_dataset.py \\
        [--profile configs/datagen_droid.yaml] [--staging-dir DIR] \\
        [--uuid UUID [UUID ...]] [--dry-run | --no-dry-run]
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--profile", default=str(REPO_ROOT / "configs" / "datagen_droid.yaml"))
    p.add_argument("--staging-dir", default=None,
                    help="default: <output_root>/_publish_staging")
    p.add_argument("--uuid", nargs="+", default=None,
                    help="explicit episode uuid(s) to publish; default: every episode with a "
                         "master/ dir under output_root")
    p.add_argument("--dry-run", action=argparse.BooleanOptionalAction, default=True,
                    help="stage + manifest only, no upload (default: on)")
    p.add_argument("--delete", action="store_true",
                    help="let sync_bucket delete remote files with no local counterpart")
    p.add_argument("--log-level", default="INFO")
    return p.parse_args()


def main() -> None:
    args = parse_args()

    from fpgm.config_datagen import DatagenProfile
    from fpgm.datagen.publish import DatasetPublisher
    from fpgm.utils.logging import get_logger, setup_logging

    setup_logging(args.log_level)
    logger = get_logger("publish_dataset")

    profile = DatagenProfile.from_yaml(args.profile)
    staging_dir = (
        Path(args.staging_dir) if args.staging_dir
        else profile.paths.output_root / "_publish_staging"
    )

    publisher = DatasetPublisher(profile)
    report = publisher.stage(staging_dir, episode_uuids=args.uuid)
    print(
        f"staged {report.n_episodes} episode(s), {report.n_files} file(s), "
        f"{report.total_bytes / 1e9:.3f} GB under {report.staging_dir}"
    )
    print(f"manifest: {report.manifest_path}")
    print(f"README:   {report.readme_path}")

    if report.n_episodes == 0:
        logger.warning("publish: nothing staged -- skipping sync_bucket entirely")
        return

    plan = publisher.publish(staging_dir, dry_run=args.dry_run, delete=args.delete)
    plan_summary = plan.summary()
    print(f"sync plan ({'DRY RUN' if args.dry_run else 'LIVE'}): {plan_summary}")
    if args.dry_run:
        print("(pass --no-dry-run to actually upload)")


if __name__ == "__main__":
    main()