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

Fix master address scope bug

Browse Source 
The sync_state variable was out of scope for non-zero ranks,
causing them to use an old sync_state value from an earlier
loop iteration. Now properly initialize and use master_addr
variable in both branches.

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
George Stykalin
2026-04-15 12:29:35 +03:00
parent d7aecf3cd8
commit 4fff9658d7
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -81,11 +81,13 @@ def run_training_node_func(rank, world_size):
my_fqdn = socket.getfqdn() my_fqdn = socket.getfqdn()
print(f"[{rank}] My FQDN: {my_fqdn}") print(f"[{rank}] My FQDN: {my_fqdn}")
master_addr = None # Initialize outside if/else
if rank == 0: if rank == 0:
# Rank 0 stores its address as the master # Rank 0 stores its address as the master
sync_state = json.loads(Variable.get('ddp_sync_state', default_var='{}')) sync_state = json.loads(Variable.get('ddp_sync_state', default_var='{}'))
sync_state['master_addr'] = my_fqdn sync_state['master_addr'] = my_fqdn
Variable.set('ddp_sync_state', json.dumps(sync_state)) Variable.set('ddp_sync_state', json.dumps(sync_state))
master_addr = my_fqdn # For rank 0, master is itself
print(f"[{rank}] I am the master. Stored my address: {my_fqdn}") print(f"[{rank}] I am the master. Stored my address: {my_fqdn}")
else: else:
# Other ranks wait for rank 0 to store its address # Other ranks wait for rank 0 to store its address
@@ -106,9 +108,9 @@ def run_training_node_func(rank, world_size):
time.sleep(3) time.sleep(3)
# STEP 3: Configure distributed environment # STEP 3: Configure distributed environment
# Use dynamic master address from rank 0 # Use the master_addr that was determined above
actual_master_addr = my_fqdn if rank == 0 else sync_state.get('master_addr', MASTER_ADDR) print(f"[{rank}] Using master address: {master_addr}")
os.environ['MASTER_ADDR'] = actual_master_addr os.environ['MASTER_ADDR'] = master_addr
os.environ['MASTER_PORT'] = MASTER_PORT os.environ['MASTER_PORT'] = MASTER_PORT
os.environ['WORLD_SIZE'] = str(world_size) os.environ['WORLD_SIZE'] = str(world_size)
os.environ['RANK'] = str(rank) os.environ['RANK'] = str(rank)