scripts/run_appearance_pipeline.py

#!/usr/bin/env python
"""Run the appearance-pair pipeline as N independent (depth, control) shards.

One shard is one MoGe process plus one render/encode process on the same GPU,
with its own scratch directory so two consumers can never race for the same
depth file. Per shard the two stages measure 19.4 and 18.8 fps -- close enough
that neither starves the other, and separate processes are what lets decode
(61-74 fps), the GPU forward pass and two ffmpeg encoders all overlap. Merged
into one process the same chain would run sequentially at ~14 fps.

**Both stages get pinned to the shard's GPU, by different mechanisms**, and
missing either one silently loses most of the parallelism: ``CUDA_VISIBLE_DEVICES``
for MoGe, ``EGL_DEVICE_ID`` for pyrender. Rendering is 52% of stage 2's budget
and it drops from 57 to 36 fps when it lands on a GPU somebody else is using --
so leaving the renderer on the default device puts every shard on GPU 0 at once.

Politeness: ``--gpus`` is explicit and has no default that grabs everything.
MoGe holds ~6 GB and the renderer ~1 GB, so a shard fits alongside most jobs,
but pick GPUs that are actually idle.

Usage:
    /home/quang/miniconda3/envs/fpgm/bin/python scripts/run_appearance_pipeline.py \\
        --gpus 1 2 3 --uuid-file configs/appearance_400.txt \\
        --out outputs/appearance_pairs
"""

from __future__ import annotations

import argparse
import os
import shutil
import signal
import subprocess
import sys
import time
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]
MOGE_PY = "/home/quang/miniconda3/envs/moge/bin/python"
FPGM_PY = "/home/quang/miniconda3/envs/fpgm/bin/python"


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--gpus", type=int, nargs="+", required=True,
                    help="one shard per GPU id; pick idle ones")
    ap.add_argument("--out", type=Path, default=REPO_ROOT / "outputs/appearance_pairs")
    ap.add_argument("--work-root", type=Path,
                    default=Path(os.environ.get("TMPDIR", "/tmp")) / "appearance_work")
    ap.add_argument("--log-dir", type=Path, default=REPO_ROOT / "outputs/appearance_logs")
    ap.add_argument("--episodes-root", type=Path, default=REPO_ROOT / "data/droid_raw")
    ap.add_argument("--uuid-file", type=Path, default=None)
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--crf", type=int, default=18)
    ap.add_argument("--depth-tol-m", type=float, default=0.15)
    ap.add_argument("--batch", type=int, default=4)
    ap.add_argument("--keep-work", action="store_true")
    args = ap.parse_args()

    n = len(args.gpus)
    args.log_dir.mkdir(parents=True, exist_ok=True)
    args.out.mkdir(parents=True, exist_ok=True)
    procs: list[tuple[str, subprocess.Popen]] = []
    works: list[Path] = []

    for i, gpu in enumerate(args.gpus):
        work = args.work_root / f"shard{i}"
        work.mkdir(parents=True, exist_ok=True)
        works.append(work)

        depth_cmd = [MOGE_PY, "-u", str(REPO_ROOT / "scripts/appearance_depth.py"),
                     "--work", str(work), "--shard", str(i), str(n),
                     "--episodes-root", str(args.episodes_root), "--batch", str(args.batch)]
        ctrl_cmd = [FPGM_PY, "-u", str(REPO_ROOT / "scripts/appearance_control.py"),
                    "--work", str(work), "--out", str(args.out), "--follow",
                    "--episodes-root", str(args.episodes_root),
                    "--crf", str(args.crf), "--depth-tol-m", str(args.depth_tol_m)]
        if args.uuid_file:
            depth_cmd += ["--uuid-file", str(args.uuid_file)]
        if args.limit:
            depth_cmd += ["--limit", str(args.limit)]

        depth_env = {**os.environ, "CUDA_VISIBLE_DEVICES": str(gpu)}
        ctrl_env = {**os.environ, "PYOPENGL_PLATFORM": "egl",
                    "EGL_DEVICE_ID": str(gpu),
                    "PYTHONPATH": f"{REPO_ROOT/'src'}:{os.environ.get('PYTHONPATH','')}"}
        for tag, cmd, env in (("depth", depth_cmd, depth_env),
                              ("control", ctrl_cmd, ctrl_env)):
            log = (args.log_dir / f"shard{i}_gpu{gpu}_{tag}.log").open("w")
            procs.append((f"shard{i}/{tag}",
                          subprocess.Popen(cmd, env=env, stdout=log, stderr=log)))
        print(f"shard {i} -> GPU {gpu}  work={work}", flush=True)

    print(f"{n} shard(s) running; logs in {args.log_dir}", flush=True)

    def stop(*_):
        for name, p in procs:
            p.terminate()
        sys.exit(130)

    signal.signal(signal.SIGINT, stop)
    signal.signal(signal.SIGTERM, stop)

    failed = []
    for name, p in procs:
        if p.wait() != 0:
            failed.append(f"{name} (exit {p.returncode})")
    if not args.keep_work:
        for w in works:
            shutil.rmtree(w, ignore_errors=True)
    if failed:
        print(f"FAILED: {', '.join(failed)} -- see {args.log_dir}", flush=True)
        return 1
    print("all shards done", flush=True)
    return 0


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