scripts/run_datagen.py

#!/usr/bin/env python
"""DROID -> VACE training-data pipeline CLI: per-episode, per-camera stage runner.

A thin wrapper over :class:`fpgm.datagen.pipeline.EpisodePipeline`: parse args, load a
:class:`~fpgm.config_datagen.DatagenProfile` (deployment configuration -- paths, role
treatment, GPU pool -- see ``configs/datagen_droid.yaml``), resolve this episode's
:class:`~fpgm.datagen.episode_spec.EpisodeSpec` (object identity/roles/mesh paths,
derived from the DROID instruction text -- nothing here is hard-coded to any one
episode), run the pipeline, and print/write the resulting
:class:`~fpgm.datagen.pipeline.EpisodeReport`. All stage logic (S2 robot buffers, S3
background/dense depth + RGB plate, S4' prompt-driven segmentation, S6 per-object SE(3)
poses + events, S8 VACE export) lives in :mod:`fpgm.datagen.pipeline`; this file owns
only CLI parsing and episode/camera resolution.

Usage:
    # GPU host note: nvidia-smi/CUDA need the driver shim active for this whole
    # process -- see scripts/nvidia_lib_shim.sh. This script sources its effect
    # itself by prepending .nvshim to LD_LIBRARY_PATH before any GPU-touching import.

    PYTHONPATH=src python scripts/run_datagen.py \\
        --uuid AUTOLab+0d4edc83+2023-10-21-19h-07m-04s --camera ext1

    # Explicit stage list, a different output root, and ignoring the on-disk cache:
    PYTHONPATH=src python scripts/run_datagen.py \\
        --uuid AUTOLab+0d4edc83+2023-10-21-19h-07m-04s --camera 22008760 \\
        --stages s2,s3 --out-root /tmp/scratch --force
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

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

# This process's CUDA (torch, and therefore SAM3) needs the driver shim on
# LD_LIBRARY_PATH -- see scripts/nvidia_lib_shim.sh's docstring. Must happen before
# any torch/CUDA-touching import; nothing above this line imports one.
_NVSHIM_DIR = REPO_ROOT / ".nvshim"
if _NVSHIM_DIR.is_dir():
    _existing = os.environ.get("LD_LIBRARY_PATH", "")
    os.environ["LD_LIBRARY_PATH"] = (
        f"{_NVSHIM_DIR}{os.pathsep}{_existing}" if _existing else str(_NVSHIM_DIR)
    )

from fpgm.config_datagen import DatagenProfile  # noqa: E402
from fpgm.datagen.episode_spec import (  # noqa: E402
    EpisodeMetadata,
    EpisodeSpecResolver,
    TaskNotParseableError,
)
from fpgm.datagen.pipeline import EpisodePipeline  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

logger = get_logger("run_datagen")

_DEFAULT_CONFIG = REPO_ROOT / "configs" / "datagen_droid.yaml"


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--uuid", required=True, help="DROID episode uuid")
    p.add_argument(
        "--camera", required=True,
        help="camera serial (e.g. 22008760) or role (ext1/ext2)",
    )
    p.add_argument(
        "--stages", default=",".join(EpisodePipeline.STAGES),
        help=f"comma-separated stage list; known stages: {EpisodePipeline.STAGES}",
    )
    p.add_argument("--config", type=Path, default=_DEFAULT_CONFIG, help="DatagenProfile YAML")
    p.add_argument(
        "--out-root", type=Path, default=None,
        help="override the profile's paths.output_root (e.g. for a scratch run)",
    )
    p.add_argument("--force", action="store_true", help="ignore the on-disk stage cache")
    p.add_argument("--log-level", default="INFO")
    return p.parse_args()


def main() -> int:
    args = parse_args()
    setup_logging(args.log_level)

    profile = DatagenProfile.from_yaml(args.config)
    if args.out_root is not None:
        profile.paths.output_root = args.out_root.resolve()

    stages = [s.strip() for s in args.stages.split(",") if s.strip()]

    flows_path = profile.paths.flows_h5(args.uuid)
    metadata = EpisodeMetadata.from_flows_h5(flows_path)
    try:
        spec = EpisodeSpecResolver(profile).resolve(metadata, camera_role=args.camera)
    except TaskNotParseableError as exc:
        logger.error(
            "episode %s: task does not parse into a usable object spec: %s", args.uuid, exc
        )
        return 2

    logger.info(
        "episode=%s camera=%s (serial=%s) objects=%s stages=%s",
        spec.uuid, spec.camera_role, spec.camera_serial, spec.labels, stages,
    )

    pipeline = EpisodePipeline(profile)
    report = pipeline.run(spec, stages=stages, force=args.force)

    master_dir = profile.paths.master_dir(spec.uuid, spec.camera_serial)
    report_path = master_dir / "meta.json"
    report_path.write_text(json.dumps(report.to_json(), indent=2, default=str))
    logger.info("wrote %s", report_path)

    for stage in report.stages:
        print(f"  {stage.name}: {stage.status} ({stage.elapsed_s:.1f}s)"
              + (f" -- {stage.error}" if stage.error else ""))
    print(f"\nepisode status: {report.status}")
    print(f"outputs: {master_dir}")
    return 0 if report.status == "ok" else 1


if __name__ == "__main__":
    raise SystemExit(main())