diff --git a/dags/test-train-pytorch.py b/dags/test-train-pytorch.py
index 497d4fd..50a4c22 100644
--- a/dags/test-train-pytorch.py
+++ b/dags/test-train-pytorch.py
@@ -78,7 +78,29 @@ def run_training_node_func(rank, world_size):
   # Small delay to ensure all workers see the ready state
   time.sleep(3)
 
-  # STEP 3: Configure distributed environment
+  # STEP 3: Configure distributed Environment
+
+  import socket
+
+  if rank == 0:
+      # Rank 0 becomes master
+      master_addr = socket.gethostbyname(socket.gethostname())
+      Variable.set("MASTER_ADDR_DYNAMIC", master_addr)
+      print(f"[{rank}] Acting as MASTER at {master_addr}")
+  else:
+      # Other ranks wait for master
+      print(f"[{rank}] Waiting for MASTER_ADDR...")
+      while True:
+          try:
+              master_addr = Variable.get("MASTER_ADDR_DYNAMIC", default_var=None)
+              if master_addr:
+                  break
+          except:
+              pass
+          time.sleep(1)
+
+      print(f"[{rank}] Found MASTER at {master_addr}")
+
   os.environ['MASTER_ADDR'] = MASTER_ADDR
   os.environ['MASTER_PORT'] = MASTER_PORT
   os.environ['WORLD_SIZE'] = str(world_size)