gstykalin/test-dags
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update dags

Browse Source
This commit is contained in:
George Stykalin
2026-04-15 18:21:41 +03:00
parent 971bff62f7
commit 6968b18418
1 changed files with 21 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
dags/test-train-pytorch.py
View File

@@ -36,12 +36,28 @@ def run_training_node_func(rank, world_size):
"""Execute distributed training with synchronization""" """Execute distributed training with synchronization"""
import os import os
import socket 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
import torch.distributed as dist import torch.distributed as dist
import torch.nn as nn import torch.nn as nn
import torch.optim as optim import torch.optim as optim
from datetime import datetime
import time
print(f"{'='*60}") print(f"{'='*60}")
print(f"Node Rank {rank}/{world_size} - Starting at {datetime.now()}") 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['WORLD_SIZE'] = str(world_size)
os.environ['RANK'] = str(rank) 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"[{rank}] Environment configured:")
print(f" MASTER_ADDR: {MASTER_ADDR}") print(f" MASTER_ADDR: {master_addr}")
print(f" MASTER_PORT: {MASTER_PORT}") print(f" MASTER_PORT: {MASTER_PORT}")
print(f" RANK: {rank}") print(f" RANK: {rank}")
print(f" WORLD_SIZE: {world_size}") 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 # STEP 3.5: Pre-flight checks
print(f"[{rank}] Running 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 # STEP 5: Initialize process group
print(f"[{rank}] Initializing process group (backend=nccl)...") print(f"[{rank}] Initializing process group (backend=nccl)...")
print(f"[{rank}] This may take 30-60 seconds...") 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 # Flush output to ensure logs are visible
import sys import sys