From 6968b18418b454cf914c5e8b516f31a20fb474cf Mon Sep 17 00:00:00 2001
From: George Stykalin <gstykalin@mail.ru>
Date: Wed, 15 Apr 2026 18:21:41 +0300
Subject: [PATCH] update dags

---
 dags/test-train-pytorch.py | 39 ++++++++++++++++++++------------------
 1 file changed, 21 insertions(+), 18 deletions(-)

diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py
index ade6574..612445d 100644
--- a/dags/test-train-pytorch.py
+++ b/dags/test-train-pytorch.py
@@ -36,12 +36,28 @@ def run_training_node_func(rank, world_size):
   """Execute distributed training with synchronization"""
   import os
   import socket
+  from datetime import datetime
+  import time
+
+  # CRITICAL: Set NCCL environment variables BEFORE importing torch
+  # This ensures NCCL reads these settings during initialization
+  os.environ['NCCL_DEBUG'] = 'INFO'
+  os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
+  os.environ['NCCL_BLOCKING_WAIT'] = '1'
+  os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'  # Force ethernet (reliable)
+  os.environ['NCCL_IB_DISABLE'] = '1'  # Disable IB - infrastructure not configured
+  os.environ['NCCL_P2P_DISABLE'] = '0'  # Enable P2P
+  os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1'  # Better compatibility
+  os.environ['NCCL_LL_THRESHOLD'] = '0'  # Disable LL for compatibility
+  os.environ['NCCL_ALGO'] = 'Ring'  # Use ring algorithm
+  os.environ['NCCL_PROTO'] = 'Simple'  # Use simple protocol
+  os.environ['NCCL_SOCKET_FAMILY'] = 'AF_INET'  # Force IPv4 sockets
+
+  # Now import torch after NCCL config is set
   import torch
   import torch.distributed as dist
   import torch.nn as nn
   import torch.optim as optim
-  from datetime import datetime
-  import time
 
   print(f"{'='*60}")
   print(f"Node Rank {rank}/{world_size} - Starting at {datetime.now()}")
@@ -115,23 +131,13 @@ def run_training_node_func(rank, world_size):
   os.environ['WORLD_SIZE'] = str(world_size)
   os.environ['RANK'] = str(rank)
 
-  # NCCL Configuration - Use ethernet (IB not accessible in containers)
-  os.environ['NCCL_DEBUG'] = 'INFO'
-  os.environ['NCCL_TIMEOUT'] = str(NCCL_TIMEOUT)
-  os.environ['NCCL_BLOCKING_WAIT'] = '1'
-  os.environ['NCCL_SOCKET_IFNAME'] = 'eth0'  # Use ethernet (reliable)
-  os.environ['NCCL_IB_DISABLE'] = '1'  # Disable IB - /proc/drivers/infiniband not accessible
-  os.environ['NCCL_P2P_DISABLE'] = '0'  # Enable P2P
-  os.environ['NCCL_IGNORE_CPU_AFFINITY'] = '1'  # Better compatibility
-  os.environ['NCCL_LL_THRESHOLD'] = '0'  # Disable LL for compatibility
-  os.environ['NCCL_ALGO'] = 'Ring'  # Use ring algorithm
-  os.environ['NCCL_PROTO'] = 'Simple'  # Use simple protocol
-
   print(f"[{rank}] Environment configured:")
-  print(f"  MASTER_ADDR: {MASTER_ADDR}")
+  print(f"  MASTER_ADDR: {master_addr}")
   print(f"  MASTER_PORT: {MASTER_PORT}")
   print(f"  RANK: {rank}")
   print(f"  WORLD_SIZE: {world_size}")
+  print(f"  NCCL_IB_DISABLE: {os.environ.get('NCCL_IB_DISABLE')}")
+  print(f"  NCCL_SOCKET_IFNAME: {os.environ.get('NCCL_SOCKET_IFNAME')}")
 
   # STEP 3.5: Pre-flight checks
   print(f"[{rank}] Running pre-flight checks...")
@@ -170,9 +176,6 @@ def run_training_node_func(rank, world_size):
   # STEP 5: Initialize process group
   print(f"[{rank}] Initializing process group (backend=nccl)...")
   print(f"[{rank}] This may take 30-60 seconds...")
-  print(f"[{rank}] NCCL_DEBUG={os.environ.get('NCCL_DEBUG', 'not set')}")
-  print(f"[{rank}] NCCL_IB_DISABLE={os.environ.get('NCCL_IB_DISABLE', 'not set')}")
-  print(f"[{rank}] NCCL_SOCKET_IFNAME={os.environ.get('NCCL_SOCKET_IFNAME', 'not set')}")
 
   # Flush output to ensure logs are visible
   import sys