scripts/run_datagen_batch.py

#!/usr/bin/env python
"""Thin CLI over :class:`fpgm.datagen.batch.runner.BatchRunner`.

All batch policy (pilot sizing, GPU admission, thread counts, checkpointing,
pruning) lives in :mod:`fpgm.datagen.batch` -- this script only parses
arguments, loads the profile, and reports the result. See
``fpgm/datagen/batch/runner.py``'s own module docstring for the
process-per-worker/multi-GPU design this drives.

Usage:
    PYTHONPATH=src python scripts/run_datagen_batch.py \\
        [--profile configs/datagen_droid.yaml] [--camera ext1] \\
        [--mode auto|pilot|full] [--limit N] [--uuid UUID [UUID ...]] [--force]
"""

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"))

_NVSHIM_DIR = REPO_ROOT / ".nvshim"
if _NVSHIM_DIR.is_dir():
    import os

    _existing = os.environ.get("LD_LIBRARY_PATH", "")
    os.environ["LD_LIBRARY_PATH"] = (
        f"{_NVSHIM_DIR}{os.pathsep}{_existing}" if _existing else str(_NVSHIM_DIR)
    )


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("--camera", default=None,
                    help="camera role; defaults to the profile's batch.camera_role")
    p.add_argument("--mode", choices=("auto", "pilot", "full"), default="auto")
    p.add_argument("--uuid", nargs="+", default=None,
                    help="explicit episode uuid(s); default: discover all")
    p.add_argument("--limit", type=int, default=None, help="cap the number of episodes (debug)")
    p.add_argument("--force", action="store_true",
                    help="re-run episodes even if already recorded ok")
    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.batch.runner import BatchRunner
    from fpgm.utils.logging import setup_logging

    setup_logging(args.log_level)

    profile = DatagenProfile.from_yaml(args.profile)
    runner = BatchRunner(profile, args.profile)

    uuids = args.uuid if args.uuid else runner.discover_episode_uuids()
    if args.limit:
        uuids = uuids[: args.limit]

    jobs = runner.build_jobs(uuids, camera_role=args.camera, force=args.force)
    summary = runner.run(jobs, mode=args.mode)

    n_ok = sum(1 for e in summary.episodes if e.get("status") == "ok")
    n_failed = sum(1 for e in summary.episodes if e.get("status") == "failed")
    n_skipped = sum(1 for e in summary.episodes if e.get("status") == "skipped")
    n_motion = sum(1 for e in summary.episodes if e.get("objects_from") == "motion")
    print(
        f"batch done: {n_ok} ok, {n_failed} failed, {n_skipped} skipped "
        f"(of {len(summary.episodes)} recorded). Summary: {summary.json_path}"
    )
    # No episode is refused for an unparseable DROID instruction any more -- see
    # fpgm.datagen.episode_spec's module docstring. `n_motion` is how many episodes
    # needed the objects_from="motion" fallback (recorded per-episode as
    # "objects_from" in the summary, and in the aggregate
    # "objects_from_motion_rate"/"objects_from_motion_episodes" fields), which is a
    # measurement, never a skip: any remaining "skipped" episodes above failed a
    # different, stage-level precondition (e.g. a missing mp4/URDF), not this one.
    print(
        f"objects_from='motion' fallback used for {n_motion}/{len(summary.episodes)} "
        "episode(s) (unparseable instruction -- not a refusal, see summary for details)"
    )


if __name__ == "__main__":
    main()