Fix DDP master address for dynamic task scheduling

The fundamental issue was that MASTER_ADDR was hardcoded to
airflow-worker-gpu-0, but rank 0 might run on any worker due to
Celery's dynamic task scheduling.

Now rank 0 dynamically stores its FQDN as the master address,
and all other ranks connect to it.

Co-Authored-By: Claude <noreply@anthropic.com>

dags/test-train-pytorch.py

@@ -76,11 +76,39 @@ def run_training_node_func(rank, world_size):
   else:
       raise RuntimeError(f"[{rank}] Timeout waiting for all workers to be ready!")
 
+  # STEP 2.5: Rank 0 stores its address, others wait for it
+  print(f"[{rank}] Determining master address...")
+  my_fqdn = socket.getfqdn()
+  print(f"[{rank}] My FQDN: {my_fqdn}")
+
+  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))
+      print(f"[{rank}] I am the master. Stored my address: {my_fqdn}")
+  else:
+      # Other ranks wait for rank 0 to store its address
+      max_wait = 60
+      start_wait = time.time()
+      while time.time() - start_wait < max_wait:
+          sync_state = json.loads(Variable.get('ddp_sync_state', default_var='{}'))
+          if 'master_addr' in sync_state:
+              master_addr = sync_state['master_addr']
+              print(f"[{rank}] Master address retrieved: {master_addr}")
+              break
+          print(f"[{rank}] Waiting for master address...")
+          time.sleep(1)
+      else:
+          raise RuntimeError(f"[{rank}] Timeout waiting for master address!")
+
   # Small delay to ensure all workers see the ready state
   time.sleep(3)
 
   # STEP 3: Configure distributed environment
-  os.environ['MASTER_ADDR'] = MASTER_ADDR
+  # 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
   os.environ['MASTER_PORT'] = MASTER_PORT
   os.environ['WORLD_SIZE'] = str(world_size)
   os.environ['RANK'] = str(rank)