Enable InfiniBand for TensorFlow distributed training

- Set NCCL IB env vars before TF import (IB_DISABLE=0, GPU Direct RDMA)
- Use NCCL_DEBUG_FILE approach for transport logging (same as PyTorch DAG)
- Add IB pre-flight diagnostics (dev/infiniband, sys/class/infiniband, ibv_devinfo)
- Add per-epoch logging with timing for all workers
- Add NCCL transport summary after training (NET/IB, GPU Direct, Channel info)

Co-Authored-By: Claude <noreply@anthropic.com>

dags/test-train-pytorch.py

@@ -38,19 +38,18 @@ def run_training_node_func(rank, world_size):
   import socket
   from datetime import datetime
   import time
+  import sys
+  import tempfile
 
   # CRITICAL: Set NCCL environment variables BEFORE importing torch
-  # This ensures NCCL reads these settings during initialization
-  os.environ['NCCL_DEBUG'] = 'INFO'
+  nccl_log_file = tempfile.mktemp(suffix='.nccl.log')
+  os.environ['NCCL_DEBUG'] = 'TRACE'
+  os.environ['NCCL_DEBUG_FILE'] = nccl_log_file
+  os.environ['NCCL_DEBUG_SUBSYS'] = 'NET'
   os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
-  os.environ['TORCH_NCCL_BLOCKING_WAIT'] = '1'  # Updated naming
-  os.environ['NCCL_IB_DISABLE'] = '0'  # Enable IB
-  os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'  # Fallback to eth0
-  os.environ['NCCL_IB_GID_INDEX'] = '3'  # Use RoCE v2 (GID index 3)
-  os.environ['NCCL_IB_TC'] = '106'  # Traffic class for RoCE
-  os.environ['NCCL_IB_TIMEOUT'] = '22'  # IB timeout (higher = more tolerant)
-  os.environ['NCCL_IB_RETRY_CNT'] = '7'  # IB retry count
-  os.environ['NCCL_NET_GDR_LEVEL'] = '5'  # Enable GPUDirect RDMA
+  os.environ['TORCH_NCCL_BLOCKING_WAIT'] = '1'
+  os.environ['NCCL_IB_DISABLE'] = '0'  # Enable InfiniBand
+  os.environ['NCCL_IB_GPU_DIRECT_RDMA_MAX_DISTANCE'] = '6'  # GPU at bus 0, IB at bus 6
   os.environ['NCCL_P2P_DISABLE'] = '0'  # Enable P2P
   os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1'  # Better compatibility
 
@@ -138,8 +137,6 @@ def run_training_node_func(rank, world_size):
   print(f"  RANK: {rank}")
   print(f"  WORLD_SIZE: {world_size}")
   print(f"  NCCL_IB_DISABLE: {os.environ.get('NCCL_IB_DISABLE')} (IB enabled)")
-  print(f"  NCCL_IB_GID_INDEX: {os.environ.get('NCCL_IB_GID_INDEX')}")
-  print(f"  NCCL_IB_TIMEOUT: {os.environ.get('NCCL_IB_TIMEOUT')}")
 
   # STEP 3.5: Pre-flight checks
   print(f"[{rank}] Running pre-flight checks...")
@@ -163,9 +160,48 @@ def run_training_node_func(rank, world_size):
           print(f"[{rank}] Found IB interfaces:")
           for line in ib_lines[:5]:
               print(f"  {line.strip()}")
+      else:
+          print(f"[{rank}] No IB interfaces found via 'ip addr'")
   except Exception as e:
       print(f"[{rank}] Warning: Could not list interfaces: {e}")
 
+  # Check /dev/infiniband devices
+  print(f"[{rank}] Checking IB device access...")
+  ib_dev_path = '/dev/infiniband'
+  if os.path.exists(ib_dev_path):
+      ib_devices = os.listdir(ib_dev_path)
+      print(f"[{rank}] /dev/infiniband contents: {ib_devices if ib_devices else 'EMPTY'}")
+  else:
+      print(f"[{rank}] /dev/infiniband does NOT exist - IB devices not mounted in container")
+
+  # Check /sys/class/infiniband
+  ib_sys_path = '/sys/class/infiniband'
+  if os.path.exists(ib_sys_path):
+      ib_sys_devices = os.listdir(ib_sys_path)
+      print(f"[{rank}] /sys/class/infiniband contents: {ib_sys_devices if ib_sys_devices else 'EMPTY'}")
+  else:
+      print(f"[{rank}] /sys/class/infiniband does NOT exist")
+
+  # Try ibv_devinfo
+  try:
+      result = subprocess.run(['ibv_devinfo'], capture_output=True, text=True, timeout=10)
+      print(f"[{rank}] ibv_devinfo output:")
+      for line in result.stdout.split('\n')[:20]:
+          print(f"  {line}")
+      if result.returncode != 0:
+          print(f"[{rank}] ibv_devinfo stderr: {result.stderr.strip()}")
+  except FileNotFoundError:
+      print(f"[{rank}] ibv_devinfo not found - rdma-core tools not installed")
+  except Exception as e:
+      print(f"[{rank}] ibv_devinfo failed: {e}")
+
+  # Check RDMA devices via sysfs
+  rdma_path = '/sys/class/infiniband_verbs'
+  if os.path.exists(rdma_path):
+      print(f"[{rank}] RDMA verbs devices: {os.listdir(rdma_path)}")
+  else:
+      print(f"[{rank}] /sys/class/infiniband_verbs does NOT exist")
+
   # STEP 4: Setup GPU FIRST (before process group init)
   device = torch.device("cuda:0")
   torch.cuda.set_device(device)
@@ -231,9 +267,12 @@ def run_training_node_func(rank, world_size):
   batch_size = 32
   num_batches = 5
 
+  training_start = time.time()
+
   for epoch in range(num_epochs):
       ddp_model.train()
       epoch_loss = 0.0
+      epoch_start = time.time()
 
       for batch_idx in range(num_batches):
           torch.manual_seed(epoch * num_batches + batch_idx)
@@ -249,20 +288,34 @@ def run_training_node_func(rank, world_size):
           epoch_loss += loss.item()
 
       avg_loss = epoch_loss / num_batches
+      epoch_time = time.time() - epoch_start
 
       # Synchronize loss across ranks
       loss_tensor = torch.tensor([avg_loss]).to(device)
       dist.all_reduce(loss_tensor, op=dist.ReduceOp.AVG)
       global_avg_loss = loss_tensor.item()
 
-      if rank == 0:
-          print(f"[{rank}] Epoch {epoch+1}/{num_epochs} | Global Avg Loss: {global_avg_loss:.6f}")
+      print(f"[{rank}] Epoch {epoch+1}/{num_epochs} DONE | Local Avg: {avg_loss:.6f} | Global Avg: {global_avg_loss:.6f} | Time: {epoch_time:.3f}s")
+
+  total_time = time.time() - training_start
 
   print(f"\n[{rank}] {'='*60}")
   print(f"[{rank}] Training Complete!")
+  print(f"[{rank}] Total training time: {total_time:.3f}s")
   print(f"[{rank}] {'='*60}")
 
-  # STEP 8: Cleanup
+  # Print NCCL transport info from log file
+  print(f"\n[{rank}] === NCCL Transport Summary ===")
+  try:
+      with open(nccl_log_file, 'r') as f:
+          nccl_lines = f.readlines()
+      for line in nccl_lines:
+          line = line.strip()
+          if any(kw in line for kw in ['NET/IB', 'GPU Direct', 'Channel', 'via NET', 'via Socket', 'transport']):
+              print(f"[{rank}] NCCL: {line}")
+  except Exception as e:
+      print(f"[{rank}] Could not read NCCL log: {e}")
+
   dist.destroy_process_group()
   print(f"[{rank}] Process group destroyed. Finished at {datetime.now()}")
 

dags/test-train-tensorflow.py

@@ -38,6 +38,18 @@ def run_training_node_func(task_index):
     import socket
     from datetime import datetime
     import time
+    import sys
+    import tempfile
+
+    # Set NCCL environment variables BEFORE importing TensorFlow
+    nccl_log_file = tempfile.mktemp(suffix='.nccl.log')
+    os.environ['NCCL_DEBUG'] = 'TRACE'
+    os.environ['NCCL_DEBUG_FILE'] = nccl_log_file
+    os.environ['NCCL_DEBUG_SUBSYS'] = 'NET'
+    os.environ['NCCL_IB_DISABLE'] = '0'  # Enable InfiniBand
+    os.environ['NCCL_IB_GPU_DIRECT_RDMA_MAX_DISTANCE'] = '6'
+    os.environ['NCCL_P2P_DISABLE'] = '0'
+    os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1'
 
     hostname = socket.gethostname()
     my_address = f"{hostname}.airflow-worker-gpu:{WORKER_PORT}"
@@ -110,6 +122,30 @@ def run_training_node_func(task_index):
     print(f"  Cluster: {cluster_spec}")
     print(f"  Task type: worker")
     print(f"  Task index: {task_index}")
+    print(f"  NCCL_IB_DISABLE: {os.environ.get('NCCL_IB_DISABLE')} (IB enabled)")
+
+    # Pre-flight IB diagnostics
+    print(f"[Worker-{task_index}] Running IB pre-flight checks...")
+    import subprocess
+    ib_dev_path = '/dev/infiniband'
+    if os.path.exists(ib_dev_path):
+        print(f"[Worker-{task_index}] /dev/infiniband: {os.listdir(ib_dev_path)}")
+    else:
+        print(f"[Worker-{task_index}] /dev/infiniband does NOT exist")
+    ib_sys_path = '/sys/class/infiniband'
+    if os.path.exists(ib_sys_path):
+        ib_devs = os.listdir(ib_sys_path)
+        print(f"[Worker-{task_index}] /sys/class/infiniband: {ib_devs}")
+        for dev in ib_devs:
+            try:
+                result = subprocess.run(['ibv_devinfo', '-d', dev], capture_output=True, text=True, timeout=10)
+                for line in result.stdout.split('\n'):
+                    if any(k in line for k in ['hca_id', 'state', 'link_layer', 'port_lid']):
+                        print(f"  {line.strip()}")
+            except Exception:
+                pass
+    else:
+        print(f"[Worker-{task_index}] /sys/class/infiniband does NOT exist")
 
     # NOW import TensorFlow after TF_CONFIG is set
     import tensorflow as tf
@@ -214,9 +250,11 @@ def run_training_node_func(task_index):
         return loss
 
     history_loss = []
+    training_start = time.time()
 
     for epoch in range(num_epochs):
         epoch_losses = []
+        epoch_start = time.time()
 
         # Get batches from dataset
         batch_iter = iter(dataset)
@@ -232,14 +270,29 @@ def run_training_node_func(task_index):
 
         avg_loss = sum(epoch_losses) / len(epoch_losses)
         history_loss.append(avg_loss)
+        epoch_time = time.time() - epoch_start
 
-        if task_index == 0:
-            print(f"[Worker-{task_index}] Epoch {epoch+1}/{num_epochs} | Loss: {avg_loss:.6f}")
+        print(f"[Worker-{task_index}] Epoch {epoch+1}/{num_epochs} DONE | Loss: {avg_loss:.6f} | Time: {epoch_time:.3f}s")
+
+    total_time = time.time() - training_start
 
     print(f"\n[Worker-{task_index}] {'='*60}")
     print(f"[Worker-{task_index}] Training Complete!")
+    print(f"[Worker-{task_index}] Total training time: {total_time:.3f}s")
     print(f"[Worker-{task_index}] {'='*60}")
 
+    # Print NCCL transport info from log file
+    print(f"\n[Worker-{task_index}] === NCCL Transport Summary ===")
+    try:
+        with open(nccl_log_file, 'r') as f:
+            nccl_lines = f.readlines()
+        for line in nccl_lines:
+            line = line.strip()
+            if any(kw in line for kw in ['NET/IB', 'GPU Direct', 'Channel', 'via NET', 'via Socket', 'transport']):
+                print(f"[Worker-{task_index}] NCCL: {line}")
+    except Exception as e:
+        print(f"[Worker-{task_index}] Could not read NCCL log: {e}")
+
     # Get final metrics
     final_loss = history_loss[-1]