#!/usr/bin/env python

import argparse
import os
import subprocess
import sys

parser = argparse.ArgumentParser(description="Nebula multi-node torchrun entry")
parser.add_argument("--training_script", required=True)
parser.add_argument("--nnodes", type=int, required=True)
parser.add_argument("--nproc_per_node", type=int, required=True)
launcher_args, training_args = parser.parse_known_args()

env = os.environ
print("===== distributed-related env =====", flush=True)
for key in sorted(env):
    if any(name in key for name in ("MASTER", "RANK", "WORLD", "NODE", "WORKER", "NPROC")):
        print(f"{key}={env[key]}", flush=True)
print("===================================", flush=True)

master_addr = env.get("MASTER_ADDR")
master_port = env.get("MASTER_PORT", "23456")
node_rank = env.get("RANK") or env.get("NODE_RANK")

if launcher_args.nnodes > 1 and not master_addr:
    parser.error("MASTER_ADDR is required for multi-node training")
if node_rank is None:
    parser.error("RANK or NODE_RANK is required for distributed training")

command = [
    "torchrun",
    f"--nnodes={launcher_args.nnodes}",
    f"--nproc_per_node={launcher_args.nproc_per_node}",
    f"--node_rank={node_rank}",
    f"--master_addr={master_addr or '127.0.0.1'}",
    f"--master_port={master_port}",
    launcher_args.training_script,
    *training_args,
]
print("command", " ".join(command), flush=True)
sys.exit(subprocess.run(command, check=False).returncode)