Use NCCL_DEBUG_FILE for transport logging

Replace os.dup2 stderr redirect with NCCL_DEBUG_FILE to avoid
interfering with NCCL's internal logging. NCCL writes TRACE output
to a temp file, then we filter and print only transport-related
lines (NET/IB, GPU Direct, Channel, via NET/Socket).

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

dags/test-train-pytorch.py

@@ -39,12 +39,13 @@ def run_training_node_func(rank, world_size):
   from datetime import datetime
   import time
   import sys
-
+  import tempfile
-  # Redirect NCCL stderr to stdout so Airflow captures it
-  os.dup2(sys.stdout.fileno(), sys.stderr.fileno())
 
   # CRITICAL: Set NCCL environment variables BEFORE importing torch
+  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'
   os.environ['NCCL_IB_DISABLE'] = '0'  # Enable InfiniBand
@@ -302,6 +303,19 @@ def run_training_node_func(rank, world_size):
   print(f"[{rank}] Training Complete!")
   print(f"[{rank}] Total training time: {total_time:.3f}s")
   print(f"[{rank}] {'='*60}")
+
+  # 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()}")