scripts/train_wan_vace_lora.py

#!/usr/bin/env python
"""LoRA-finetune Wan2.1-VACE-1.3B on the ``vace/window_*`` bundles.

Thin CLI over :class:`fpgm.training.launcher.TrainingLauncher`; all policy
(GPU admission, env setup, the VRAM sweep grid) lives there. Replaces the
former ``train_wan_vace_lora.sh`` -- see the launcher's module docstring for
what the shell version got wrong.

Usage::

    PYTHONPATH=src python scripts/train_wan_vace_lora.py train  [train.py args...]
    PYTHONPATH=src python scripts/train_wan_vace_lora.py smoke  [train.py args...]
    PYTHONPATH=src python scripts/train_wan_vace_lora.py measure-vram

Unrecognised arguments are forwarded verbatim to ``fpgm.training.train``.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

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

from fpgm.training.launcher import TrainingLaunchConfig, TrainingLauncher  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

logger = get_logger("train_wan_vace_lora")


def parse_args() -> tuple[argparse.Namespace, list[str]]:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("mode", choices=sorted(TrainingLauncher.MODES))
    p.add_argument(
        "--gpu", type=int, default=None,
        help="pin a device index and skip the headroom wait (default: wait for the freest)",
    )
    p.add_argument(
        "--min-free-gb", type=float, default=30.0,
        help="free VRAM required before launching. Measured peaks: 23.2 GB at 21 frames, "
             "27.7 GB at 41, 36.8 GB at 81.",
    )
    p.add_argument("--wait-poll-s", type=float, default=60.0)
    p.add_argument("--wait-max-s", type=float, default=3600.0)
    p.add_argument("--log-level", default="INFO")
    return p.parse_known_args()


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

    launcher = TrainingLauncher(TrainingLaunchConfig(
        min_free_mb=int(args.min_free_gb * 1024),
        poll_interval_s=args.wait_poll_s,
        acquire_timeout_s=args.wait_max_s,
        device_index=args.gpu,
    ))
    results = launcher.run(args.mode, extra)

    print(json.dumps([r.to_json() for r in results], indent=2))
    n_ok = sum(1 for r in results if r.ok)
    logger.info("%s: %d/%d trial(s) ok", args.mode, n_ok, len(results))
    # measure-vram expects failures at the largest frame counts -- a non-zero
    # trial there is a *measurement*, not a launcher error, so only the
    # single-run modes propagate a failing status.
    if args.mode == "measure-vram":
        return 0
    return 0 if n_ok == len(results) else 1


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