From dae4643f2c27482b5b2ecfb4c4e6396f9f0b73d6 Mon Sep 17 00:00:00 2001
From: Terry Kong <terryk@nvidia.com>
Date: Sun, 23 Mar 2025 14:28:58 -0700
Subject: [PATCH] fix: ray.sub race condition when overlapping srun commands on
 same node

trying a different approach

Signed-off-by: Terry Kong <terryk@nvidia.com>
---
 ray.sub | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/ray.sub b/ray.sub
index a600d96063..7258aec045 100644
--- a/ray.sub
+++ b/ray.sub
@@ -59,7 +59,11 @@ ip_head=$head_node_ip:$port
 
 # First we start the head of the ray cluster on one of the physical nodes
 # Set GPU/CPU resources to 0 to avoid scheduling on the head node
+
 head_cmd=$(cat <<EOF
+# Touch a file to indicate that the head node has started
+# Overlapping srun commands will check this file to determine if we can overlap a container command
+touch $LOG_DIR/STARTED_RAY_HEAD
 env
 cat <<EOFINNER | tee /launch-head.sh
 ray start --head \
@@ -121,11 +125,12 @@ EOF
   sleep 3
 done
 
-echo "[INFO] Querying head node too early can cause pyxis to fail"
-for i in {20..1}; do
-  echo "[INFO] Waiting for $i seconds before querying head node..."
-  sleep 1
+# Then we wait here for the file to be created by the head node container
+while ! srun --overlap --nodes=1 --ntasks=1 -w $head_node test -f $LOG_DIR/STARTED_RAY_HEAD; do 
+  echo "[INFO][$(date)] Waiting for head node container to start..."
+  sleep 2
 done
+
 # At this stage the Ray cluster bringup has started on the physical nodes in the allocation
 # Before we launch a job on this cluster we need to make sure that the bringup is complete
 # We do so by querying the number of worker_units in the ray cluster and asserting = NUM_ACTORS