src/fpgm/training/launcher.py

"""Launch LoRA training runs from Python, not shell.

Replaces ``scripts/train_wan_vace_lora.sh``. The shell version worked, but it
carried three liabilities that only a rewrite removes:

* it reimplemented GPU wait-for-headroom in ``awk``/``sort``, a *second*
  independent copy of the admission policy that
  :class:`~fpgm.datagen.batch.gpu_pool.GpuWorkerPool` already owns -- two copies
  of a policy drift, and this one had already drifted (it polled the single
  freest card, with no notion of a claim table);
* its ``measure-vram`` mode swallowed failures into a hand-written JSON string
  built by ``echo``, so a malformed trial record was indistinguishable from a
  real OOM measurement;
* it exported the **deprecated** ``PYTORCH_CUDA_ALLOC_CONF`` spelling, which
  newer PyTorch ignores in favour of ``PYTORCH_ALLOC_CONF`` -- silently losing
  the ``expandable_segments`` setting the training run depends on to fit.

Modes are strategies rather than a ``case`` statement, so adding one is a class
rather than another shell branch.
"""

from __future__ import annotations

import json
import os
import subprocess
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from pathlib import Path

from fpgm.config import REPO_ROOT
from fpgm.utils.logging import get_logger

logger = get_logger(__name__)

__all__ = [
    "LaunchResult",
    "MeasureVramMode",
    "SmokeMode",
    "TrainMode",
    "TrainingLaunchConfig",
    "TrainingLauncher",
]


@dataclass
class TrainingLaunchConfig:
    """Where the training env lives and how much headroom it needs."""

    #: Interpreter for the ``wan-train`` conda env (its own env: DiffSynth
    #: pins torch 2.6.0+cu124 and python 3.10, neither shared with ``fpgm``).
    python_executable: Path = Path.home() / "miniconda3" / "envs" / "wan-train" / "bin" / "python"
    output_root: Path = REPO_ROOT / "outputs" / "wan_lora"
    #: Minimum free VRAM before launching. Measured: a single fwd+bwd step
    #: peaks at 23.2 GB reserved at 21 frames, 27.7 GB at 41, 36.8 GB at 81.
    min_free_mb: int = 30720
    poll_interval_s: float = 60.0
    acquire_timeout_s: float = 3600.0
    #: Explicit device index, or ``None`` to wait for the freest one.
    device_index: int | None = None
    #: Directory holding the CUDA driver shim ``nvidia-smi`` needs here.
    nvshim_dir: Path = REPO_ROOT / ".nvshim"
    #: Frame counts and gradient-checkpointing settings the VRAM grid sweeps.
    #: 4n+1 values at or below Wan-VACE's 81-frame cap.
    vram_frame_counts: list[int] = field(default_factory=lambda: [81, 41, 21])


@dataclass(frozen=True)
class LaunchResult:
    """Outcome of one training subprocess."""

    name: str
    returncode: int
    device_index: int
    log_path: Path | None
    report_path: Path | None

    @property
    def ok(self) -> bool:
        return self.returncode == 0

    def to_json(self) -> dict:
        return {
            "name": self.name,
            "returncode": self.returncode,
            "ok": self.ok,
            "device_index": self.device_index,
            "log_path": str(self.log_path) if self.log_path else None,
            "report_path": str(self.report_path) if self.report_path else None,
        }


class LaunchMode(ABC):
    """One way of invoking ``fpgm.training.train``."""

    name: str

    @abstractmethod
    def run(self, launcher: TrainingLauncher, extra_args: list[str]) -> list[LaunchResult]:
        ...


class TrainMode(LaunchMode):
    """A normal LoRA run under the strict gate filter."""

    name = "train"

    def run(self, launcher: TrainingLauncher, extra_args: list[str]) -> list[LaunchResult]:
        out = launcher.config.output_root / "run"
        out.mkdir(parents=True, exist_ok=True)
        return [launcher.invoke("train", ["--output-path", str(out), *extra_args], log_dir=out)]


class SmokeMode(LaunchMode):
    """Prove the loop closes end to end on whatever windows exist today.

    Deliberately permissive (``--allow-pose-gap-overlap``): this measures
    *mechanics*, and a decreasing loss is explicitly **not** expected from a
    handful of samples. Recorded here rather than left to a reader's judgement,
    because a smoke run's loss curve is the single most tempting number in this
    project to over-read.
    """

    name = "smoke"

    def __init__(self, max_steps: int = 30) -> None:
        self._max_steps = max_steps

    def run(self, launcher: TrainingLauncher, extra_args: list[str]) -> list[LaunchResult]:
        out = launcher.config.output_root / "smoke"
        out.mkdir(parents=True, exist_ok=True)
        args = [
            "--output-path", str(out),
            "--allow-pose-gap-overlap",
            "--max-steps", str(self._max_steps),
            "--num-epochs", "100",
            "--save-steps", str(self._max_steps),
            "--report-json", str(out / "smoke_report.json"),
            *extra_args,
        ]
        return [launcher.invoke("smoke", args, log_dir=out, report=out / "smoke_report.json")]


class MeasureVramMode(LaunchMode):
    """Peak VRAM for one fwd+bwd step across a frames x checkpointing grid.

    A **fresh process per trial**, never one process looping: a peak-memory
    statistic collected after a previous trial in the same process is inflated
    by that trial's residual allocations, which would make the grid monotonic
    by construction rather than by measurement.

    A trial that dies (the expected outcome at the largest frame counts on a
    contended box) writes a structured record with the real return code. The
    shell version wrote this record with ``echo`` into a hand-built JSON
    string, so a malformed record and a genuine OOM looked identical.
    """

    name = "measure-vram"

    def run(self, launcher: TrainingLauncher, extra_args: list[str]) -> list[LaunchResult]:
        out_root = launcher.config.output_root / "vram_measurements"
        out_root.mkdir(parents=True, exist_ok=True)
        results: list[LaunchResult] = []
        for frames in launcher.config.vram_frame_counts:
            for checkpointing in (True, False):
                tag = f"frames{frames}_gc{'on' if checkpointing else 'off'}"
                report = out_root / f"{tag}.json"
                args = [
                    "--output-path", str(out_root / tag),
                    "--allow-pose-gap-overlap",
                    "--no-augment",
                    "--max-num-frames", str(frames),
                    (
                        "--use-gradient-checkpointing" if checkpointing
                        else "--no-gradient-checkpointing"
                    ),
                    "--measure-vram-only",
                    "--report-json", str(report),
                    *extra_args,
                ]
                result = launcher.invoke(tag, args, log_dir=out_root, report=report)
                if not result.ok and not report.exists():
                    report.write_text(json.dumps({
                        "trial": tag,
                        "max_num_frames": frames,
                        "requested_gradient_checkpointing": checkpointing,
                        "status": "failed",
                        "returncode": result.returncode,
                        "note": (
                            "the trial process exited non-zero without writing its own "
                            "report; typically CUDA OOM at this frame count"
                        ),
                    }, indent=2))
                results.append(result)
        return results


class TrainingLauncher:
    """Waits for GPU headroom, then runs ``fpgm.training.train`` in its own env."""

    MODES: dict[str, type[LaunchMode]] = {
        TrainMode.name: TrainMode,
        SmokeMode.name: SmokeMode,
        MeasureVramMode.name: MeasureVramMode,
    }

    def __init__(self, config: TrainingLaunchConfig | None = None) -> None:
        self.config = config or TrainingLaunchConfig()
        if not self.config.python_executable.exists():
            raise FileNotFoundError(
                f"no wan-train interpreter at {self.config.python_executable}. "
                "Create the env first (see scripts/setup_wan_train_env.sh)."
            )

    # -- GPU admission, delegated ------------------------------------------ #

    def acquire_device(self) -> int:
        """Block until a device has ``min_free_mb`` free, and return its index.

        Delegates to :class:`~fpgm.datagen.batch.gpu_pool.GpuWorkerPool` rather
        than reimplementing admission: one policy, one place. The shell version
        had its own ``nvidia-smi | sort | head -1`` copy that knew nothing
        about claims held by a concurrently-running datagen batch.
        """
        if self.config.device_index is not None:
            logger.info(
                "device_index set explicitly to %d -- skipping the headroom check",
                self.config.device_index,
            )
            return self.config.device_index

        from fpgm.config_datagen import GpuPoolConfig
        from fpgm.datagen.batch.gpu_pool import GpuWorkerPool

        pool = GpuWorkerPool(GpuPoolConfig(
            min_free_mb=self.config.min_free_mb,
            workers_per_device=1,
            poll_interval_s=self.config.poll_interval_s,
            acquire_timeout_s=self.config.acquire_timeout_s,
        ))
        # Released immediately: the claim table lives in this process, but the
        # work happens in a subprocess that outlives it, so holding the claim
        # would protect nothing. The value taken is the *admission decision*.
        with pool.acquire() as claim:
            return claim.device.index

    # -- subprocess -------------------------------------------------------- #

    def _env(self, device_index: int) -> dict[str, str]:
        env = dict(os.environ)
        env["CUDA_VISIBLE_DEVICES"] = str(device_index)
        env["PYTHONPATH"] = os.pathsep.join(
            [str(REPO_ROOT / "src"), env.get("PYTHONPATH", "")]
        ).rstrip(os.pathsep)
        if self.config.nvshim_dir.is_dir():
            env["LD_LIBRARY_PATH"] = os.pathsep.join(
                [str(self.config.nvshim_dir), env.get("LD_LIBRARY_PATH", "")]
            ).rstrip(os.pathsep)
        # Current spelling. The shell version set only PYTORCH_CUDA_ALLOC_CONF,
        # which newer PyTorch ignores -- silently dropping the setting that
        # makes the largest frame counts fit at all. Both are set so the
        # launcher works across the versions in the two envs here.
        env.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")
        env.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
        env["TOKENIZERS_PARALLELISM"] = "false"
        return env

    def invoke(
        self,
        name: str,
        args: list[str],
        *,
        log_dir: Path,
        report: Path | None = None,
    ) -> LaunchResult:
        device_index = self.acquire_device()
        log_path = log_dir / f"{name}.log"
        log_dir.mkdir(parents=True, exist_ok=True)
        cmd = [str(self.config.python_executable), "-m", "fpgm.training.train", *args]
        logger.info("launching %s on GPU %d: %s", name, device_index, " ".join(cmd))
        with open(log_path, "w") as fh:
            proc = subprocess.run(
                cmd, cwd=str(REPO_ROOT), env=self._env(device_index),
                stdout=fh, stderr=subprocess.STDOUT,
            )
        level = logger.info if proc.returncode == 0 else logger.error
        level("%s finished with returncode %d (log: %s)", name, proc.returncode, log_path)
        return LaunchResult(
            name=name, returncode=proc.returncode, device_index=device_index,
            log_path=log_path, report_path=report if report and report.exists() else None,
        )

    # -- entry point -------------------------------------------------------- #

    def run(self, mode: str, extra_args: list[str] | None = None) -> list[LaunchResult]:
        if mode not in self.MODES:
            raise ValueError(f"unknown mode {mode!r}; known: {sorted(self.MODES)}")
        return self.MODES[mode]().run(self, list(extra_args or []))