From a33888dbd1231661ce53a6c3e3bc08bba3757ebc Mon Sep 17 00:00:00 2001 From: George Stykalin Date: Thu, 16 Apr 2026 16:38:58 +0300 Subject: [PATCH] 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 --- dags/test-train-pytorch.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py index 7429a20..40cb5d8 100644 --- a/dags/test-train-pytorch.py +++ b/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 - - # Redirect NCCL stderr to stdout so Airflow captures it - os.dup2(sys.stdout.fileno(), sys.stderr.fileno()) + import tempfile # 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()}")