dpython.py

#Change the path below to your Python file (D:\AIMOD\AIMOD\venv\Scripts\python.exe) insteafd of the one shown here.
PYTHON = r"D:\AIMOD\AIMOD\venv\Scripts\python.exe"

import sys
import os
import json
import subprocess
import shutil
import time

# ================= CONFIG =================
#Change the file location below to your cluster.json location
CONFIG = r"D:\distributed_env\cluster.json"
# ==========================================

# ---------- SAFETY CHECK ----------
if len(sys.argv) < 2:
    print("❌ ERROR: No training script provided")
    print("✅ Usage: dpython train.py [args]")
    sys.exit(1)

SCRIPT = os.path.abspath(sys.argv[1])
ARGS = sys.argv[2:]

if not os.path.exists(SCRIPT):
    print(f"❌ ERROR: Script not found: {SCRIPT}")
    sys.exit(1)

# ---------- LOAD CONFIG ----------
with open(CONFIG) as f:
    cfg = json.load(f)

MASTER_IP = cfg["master_ip"]
WORKER_IP = cfg["worker_ip"]
PORT = str(cfg["port"])
SSH_USER = cfg["ssh_user"]

# ---------- GPU HELPERS ----------
def get_gpu_stats(cmd):
    try:
        out = subprocess.check_output(cmd, stderr=subprocess.DEVNULL, timeout=5)
        stats = []
        for line in out.decode().strip().splitlines():
            name, util, mem_used, mem_total = line.split(", ")
            stats.append((name, util, mem_used, mem_total))
        return stats
    except Exception:
        return []

def get_local_gpu_stats():
    return get_gpu_stats([
        "nvidia-smi",
        "--query-gpu=name,utilization.gpu,memory.used,memory.total",
        "--format=csv,noheader,nounits"
    ])

def get_remote_gpu_stats():
    return get_gpu_stats([
        "ssh",
        f"{SSH_USER}@{WORKER_IP}",
        "nvidia-smi --query-gpu=name,utilization.gpu,memory.used,memory.total "
        "--format=csv,noheader,nounits"
    ])

def remote_available():
    try:
        subprocess.check_output(
            ["ssh", f"{SSH_USER}@{WORKER_IP}", "nvidia-smi"],
            stderr=subprocess.DEVNULL,
            timeout=5
        )
        return True
    except Exception:
        return False

# ---------- DISTRIBUTED CHILD PROCESS ----------
if "LOCAL_RANK" in os.environ:
    subprocess.run([PYTHON, SCRIPT] + ARGS)
    sys.exit(0)

# ---------- GPU STATUS ----------
print("\n================ GPU STATUS ================")

local_stats = get_local_gpu_stats()
if not local_stats:
    print("❌ No local NVIDIA GPU found")
    sys.exit(1)

print("🟢 LOCAL GPU:")
for i, (n, u, mu, mt) in enumerate(local_stats):
    print(f"   [{i}] {n} | Util {u}% | VRAM {mu}/{mt} MB")

use_remote = remote_available()

if use_remote:
    remote_stats = get_remote_gpu_stats()
    if remote_stats:
        print("\n🟣 REMOTE GPU:")
        for i, (n, u, mu, mt) in enumerate(remote_stats):
            print(f"   [{i}] {n} | Util {u}% | VRAM {mu}/{mt} MB")
    else:
        use_remote = False

if not use_remote:
    print("\n⚠️ REMOTE GPU NOT AVAILABLE")
    print("➡️ Falling back to LOCAL GPU ONLY")

print("============================================\n")

# ---------- SINGLE GPU FALLBACK ----------
if not use_remote:
    print("[DPYTHON] Running single-GPU training\n")
    subprocess.run([PYTHON, SCRIPT] + ARGS)
    sys.exit(0)

# ---------- MULTI-GPU LAUNCH ----------
print("[DPYTHON] Launching remote worker...")

worker_cmd = (
    f"{PYTHON} -m accelerate launch "
    f"--num_machines 2 "
    f"--machine_rank 1 "
    f"--main_process_ip {MASTER_IP} "
    f"--main_process_port {PORT} "
    f"{SCRIPT} {' '.join(ARGS)}"
)

subprocess.Popen([
    "ssh",
    f"{SSH_USER}@{WORKER_IP}",
    worker_cmd
])

time.sleep(2)

print("[DPYTHON] Launching local master...\n")

subprocess.run([
    PYTHON, "-m", "accelerate", "launch",
    "--num_machines", "2",
    "--machine_rank", "0",
    "--main_process_ip", MASTER_IP,
    "--main_process_port", PORT,
    SCRIPT,
    *ARGS
])