scripts/launch_ddp_no_libuv.py

"""Windows-friendly two-process DDP launcher for PyTorch builds without libuv."""

from __future__ import annotations

import argparse
import os
import subprocess
import sys


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--nproc_per_node", type=int, default=2)
    parser.add_argument("--master_addr", default="127.0.0.1")
    parser.add_argument("--master_port", default="29500")
    parser.add_argument("training_script")
    parser.add_argument("training_args", nargs=argparse.REMAINDER)
    args = parser.parse_args()

    processes = []
    for local_rank in range(args.nproc_per_node):
        env = os.environ.copy()
        env.update(
            {
                "MASTER_ADDR": args.master_addr,
                "MASTER_PORT": args.master_port,
                "WORLD_SIZE": str(args.nproc_per_node),
                "RANK": str(local_rank),
                "LOCAL_RANK": str(local_rank),
                "USE_LIBUV": "0",
            }
        )
        cmd = [sys.executable, args.training_script, *args.training_args]
        processes.append(subprocess.Popen(cmd, env=env))

    exit_code = 0
    for process in processes:
        exit_code = max(exit_code, process.wait())

    if exit_code:
        for process in processes:
            if process.poll() is None:
                process.terminate()
    raise SystemExit(exit_code)


if __name__ == "__main__":
    main()