From 4fff9658d75cc9671f44e954bc2c7fc2f2e9778e Mon Sep 17 00:00:00 2001 From: George Stykalin Date: Wed, 15 Apr 2026 12:29:35 +0300 Subject: [PATCH] Fix master address scope bug 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 --- dags/test-train-pytorch.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py index e98f455..9a34d56 100644 --- a/dags/test-train-pytorch.py +++ b/dags/test-train-pytorch.py @@ -81,11 +81,13 @@ def run_training_node_func(rank, world_size): my_fqdn = socket.getfqdn() print(f"[{rank}] My FQDN: {my_fqdn}") + master_addr = None # Initialize outside if/else if rank == 0: # Rank 0 stores its address as the master sync_state = json.loads(Variable.get('ddp_sync_state', default_var='{}')) sync_state['master_addr'] = my_fqdn 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}") else: # 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) # STEP 3: Configure distributed environment - # Use dynamic master address from rank 0 - actual_master_addr = my_fqdn if rank == 0 else sync_state.get('master_addr', MASTER_ADDR) - os.environ['MASTER_ADDR'] = actual_master_addr + # Use the master_addr that was determined above + print(f"[{rank}] Using master address: {master_addr}") + os.environ['MASTER_ADDR'] = master_addr os.environ['MASTER_PORT'] = MASTER_PORT os.environ['WORLD_SIZE'] = str(world_size) os.environ['RANK'] = str(rank)