update dags

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",