diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py index 720a87f..8a0c5f2 100644 --- a/dags/test-train-pytorch.py +++ b/dags/test-train-pytorch.py @@ -7,9 +7,10 @@ import json # --- CONFIGURATION --- WORLD_SIZE = 2 -MASTER_ADDR = "airflow-worker-gpu-0.airflow-worker-gpu" +MASTER_ADDR = "airflow-worker-gpu-0.airflow-worker-gpu.george-tests.svc.cluster.local" MASTER_PORT = "29500" NCCL_TIMEOUT = 1800 +NAMESPACE = "george-tests" default_args = { 'owner': 'airflow', @@ -83,10 +84,18 @@ def run_training_node_func(rank, world_size): os.environ['MASTER_PORT'] = MASTER_PORT os.environ['WORLD_SIZE'] = str(world_size) os.environ['RANK'] = str(rank) - os.environ['NCCL_SOCKET_IFNAME'] = 'eth0' + + # NCCL Configuration for InfiniBand os.environ['NCCL_DEBUG'] = 'INFO' os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT) os.environ['NCCL_BLOCKING_WAIT'] = '1' + os.environ['NCCL_IB_DISABLE'] = '0' # Enable IB + # Auto-detect IB interface - don't hardcode since names vary by node + os.environ['NCCL_SOCKET_IFNAME'] = '^ibp[0-9]+s[0-9]+$' # Regex pattern for IB interfaces + os.environ['NCCL_IB_HCA'] = '^mlx5_[0-9]+$' # Regex pattern for HCA devices + os.environ['NCCL_P2P_DISABLE'] = '0' # Enable P2P + os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1' # Better compatibility + os.environ['NCCL_NET_GDR_LEVEL'] = '5' # Enable GPUDirect for IB print(f"[{rank}] Environment configured:") print(f" MASTER_ADDR: {MASTER_ADDR}") @@ -94,8 +103,34 @@ def run_training_node_func(rank, world_size): print(f" RANK: {rank}") print(f" WORLD_SIZE: {world_size}") + # STEP 3.5: Pre-flight checks + print(f"[{rank}] Running pre-flight checks...") + print(f"[{rank}] Hostname: {socket.gethostname()}") + print(f"[{rank}] FQDN: {socket.getfqdn()}") + + # Test DNS resolution + try: + import subprocess + result = subprocess.run(['getent', 'hosts', MASTER_ADDR], capture_output=True, text=True) + print(f"[{rank}] DNS resolution for {MASTER_ADDR}:") + print(f" {result.stdout.strip()}") + except Exception as e: + print(f"[{rank}] Warning: Could not verify DNS: {e}") + + # Check network interfaces + try: + result = subprocess.run(['ip', 'addr'], capture_output=True, text=True) + ib_lines = [line for line in result.stdout.split('\n') if 'ibp' in line or 'ib0' in line] + if ib_lines: + print(f"[{rank}] Found IB interfaces:") + for line in ib_lines[:5]: + print(f" {line.strip()}") + except Exception as e: + print(f"[{rank}] Warning: Could not list interfaces: {e}") + # STEP 4: Initialize process group print(f"[{rank}] Initializing process group (backend=nccl)...") + print(f"[{rank}] This may take 30-60 seconds if using IB...") try: dist.init_process_group( backend="nccl",