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-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} DONE | Loss: {avg_loss:.6f} | Time: {epoch_time:.3f}s")
-            print(f"[Worker-{task_index}] Epoch {epoch+1}/{num_epochs} | Loss: {avg_loss:.6f}")
+
+    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]