Fix ray srun formatting and cover heterogeneous command args

Signed-off-by: Hemil Desai <hemild@nvidia.com>

Author: hemildesai

nemo_run/run/ray/slurm.py

@@ -1271,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,
@@ -1328,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, ""),

test/run/ray/test_slurm_ray_request.py

@@ -645,6 +645,52 @@ def test_command_srun_honors_executor_srun_args(self):
         script = request.materialize()
         assert "--gpus=0 --overlap --mpi=pmix --container-name=ray-head" in script
 
+    def test_command_srun_honors_head_resource_group_srun_args(self):
+        """Test that heterogeneous grouped runs use head resource-group srun_args for COMMAND."""
+        executor = SlurmExecutor(
+            account="test_account",
+            heterogeneous=True,
+            srun_args=["--mpi=none"],
+        )
+        executor.run_as_group = True
+        executor.resource_group = [
+            SlurmExecutor.ResourceRequest(
+                packager=Mock(),
+                nodes=1,
+                ntasks_per_node=1,
+                container_image="image1",
+                container_mounts=["/data:/data"],
+                srun_args=["--mpi=pmix"],
+                het_group_index=0,
+            ),
+            SlurmExecutor.ResourceRequest(
+                packager=Mock(),
+                nodes=1,
+                ntasks_per_node=1,
+                container_image="image2",
+                container_mounts=["/data:/data"],
+                het_group_index=1,
+            ),
+        ]
+        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 (
+            "--het-group=0 --no-container-mount-home --gpus=0 --overlap --mpi=pmix "
+            "--container-name=ray-head" in script
+        )
+        assert "--gpus=0 --overlap --mpi=none --container-name=ray-head" not in script
+
     def test_env_vars_formatting(self):
         """Test that environment variables are properly formatted as export statements."""
         executor = SlurmExecutor(