Honor executor srun_args for Ray command srun

Author: hemildesai

nemo_run/run/ray/slurm.py

@@ -278,6 +278,19 @@ def get_srun_flags(mounts: list[str], container_image: Optional[str]) -> str:
 
             return " ".join(_srun_flags)
 
+        def get_command_srun_args() -> str:
+            if (
+                self.executor.run_as_group
+                and self.executor.heterogeneous
+                and self.executor.resource_group
+                and self.executor.resource_group[0].srun_args is not None
+            ):
+                command_srun_args = self.executor.resource_group[0].srun_args
+            else:
+                command_srun_args = self.executor.srun_args or []
+
+            return " ".join(shlex.quote(arg) for arg in command_srun_args)
+
         ray_log_prefix = job_details.ray_log_prefix
         vars_to_fill = {
             "sbatch_flags": sbatch_flags,
@@ -296,6 +309,7 @@ def get_srun_flags(mounts: list[str], container_image: Optional[str]) -> str:
             "ray_log_prefix": ray_log_prefix,
             "heterogeneous": self.executor.heterogeneous,
             "resource_group": self.executor.resource_group if self.executor.heterogeneous else [],
+            "command_srun_args": get_command_srun_args(),
         }
 
         if self.command_groups:
@@ -1257,9 +1271,9 @@ def start(
                 if isinstance(self.executor.tunnel, SSHTunnel):
                     # Rsync workdir honouring .gitignore
                     self.executor.tunnel.connect()
-                    assert self.executor.tunnel.session is not None, (
-                        "Tunnel session is not connected"
-                    )
+                    assert (
+                        self.executor.tunnel.session is not None
+                    ), "Tunnel session is not connected"
                     rsync(
                         self.executor.tunnel.session,
                         workdir,
@@ -1314,9 +1328,9 @@ def start(
 
                 if isinstance(self.executor.tunnel, SSHTunnel):
                     self.executor.tunnel.connect()
-                    assert self.executor.tunnel.session is not None, (
-                        "Tunnel session is not connected"
-                    )
+                    assert (
+                        self.executor.tunnel.session is not None
+                    ), "Tunnel session is not connected"
                     rsync(
                         self.executor.tunnel.session,
                         os.path.join(local_code_extraction_path, ""),

nemo_run/run/ray/templates/ray.sub.j2

@@ -454,7 +454,7 @@ COMMAND="${COMMAND:-{{ command | default('', true) }}}"
 COMMAND_WORKDIR={{ command_workdir | default('$CONTAINER_CWD') }}
 
 if [[ -n "$COMMAND" ]]; then
-  srun {% if heterogeneous %}--het-group=0 {% endif %}--no-container-mount-home --gpus=0 --overlap --container-name=ray-head --container-workdir=$COMMAND_WORKDIR --nodes=1 --ntasks=1 -w "$head_node" -o $LOG_DIR/{{ ray_log_prefix }}job.log bash -c "$COMMAND"
+  srun {% if heterogeneous %}--het-group=0 {% endif %}--no-container-mount-home --gpus=0 --overlap {% if command_srun_args %}{{ command_srun_args }} {% endif %}--container-name=ray-head --container-workdir=$COMMAND_WORKDIR --nodes=1 --ntasks=1 -w "$head_node" -o $LOG_DIR/{{ ray_log_prefix }}job.log bash -c "$COMMAND"
 else
   echo "[INFO]: Ray Cluster is idled, run this on the slurm head node to get a shell to the head node:"
   cat <<EOF >$CLUSTER_DIR/scripts/${SLURM_JOB_ID}-attach.sh

test/run/ray/test_slurm_ray_request.py

@@ -627,6 +627,24 @@ def test_command_groups_without_resource_group(self):
         assert "--overlap" in script
         assert "cmd1" in script  # Second command in the list (index 1)
 
+    def test_command_srun_honors_executor_srun_args(self):
+        """Test that the COMMAND launch srun includes executor srun_args."""
+        executor = SlurmExecutor(account="test_account", srun_args=["--mpi=pmix"])
+        executor.tunnel = Mock(spec=SSHTunnel)
+        executor.tunnel.job_dir = "/tmp/test_jobs"
+
+        request = SlurmRayRequest(
+            name="test-ray-cluster",
+            cluster_dir="/tmp/test_jobs/test-ray-cluster",
+            template_name="ray.sub.j2",
+            executor=executor,
+            command="echo hello",
+            launch_cmd=["sbatch", "--parsable"],
+        )
+
+        script = request.materialize()
+        assert "--gpus=0 --overlap --mpi=pmix --container-name=ray-head" in script
+
     def test_env_vars_formatting(self):
         """Test that environment variables are properly formatted as export statements."""
         executor = SlurmExecutor(