Repository: NVIDIA-NeMo/Run

Issue:
Add support for parallel srun commands in Slurm based Ray APIs

Discussion:

Solution patch:
diff --git a/.github/workflows/build-test-publish-wheel.yml b/.github/workflows/build-test-publish-wheel.yml
index 61aa075..68b7087 100644
--- a/.github/workflows/build-test-publish-wheel.yml
+++ b/.github/workflows/build-test-publish-wheel.yml
@@ -28,7 +28,7 @@ defaults:
 
 jobs:
   build-test-publish-wheel:
-    uses: NVIDIA/NeMo-FW-CI-templates/.github/workflows/_build_test_publish_wheel.yml@v0.23.0
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_build_test_publish_wheel.yml@v0.23.0
     with:
       dry-run: true
       python-package: nemo_run
diff --git a/.github/workflows/cherry-pick-release-commit.yml b/.github/workflows/cherry-pick-release-commit.yml
index 2430e82..8b21bd3 100644
--- a/.github/workflows/cherry-pick-release-commit.yml
+++ b/.github/workflows/cherry-pick-release-commit.yml
@@ -7,7 +7,7 @@ on:
 
 jobs:
   cherry-pick:
-    uses: NVIDIA/NeMo-FW-CI-templates/.github/workflows/_cherry_pick.yml@v0.22.7
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cherry_pick.yml@v0.22.7
     secrets:
       PAT: ${{ secrets.PAT }}
       SLACK_WEBHOOK_ADMIN: ${{ secrets.SLACK_WEBHOOK_ADMIN }}
diff --git a/.github/workflows/release-freeze.yml b/.github/workflows/release-freeze.yml
index f37b549..af02b66 100644
--- a/.github/workflows/release-freeze.yml
+++ b/.github/workflows/release-freeze.yml
@@ -21,7 +21,7 @@ on:
         default: true
 jobs:
   code-freeze:
-    uses: NVIDIA/NeMo-FW-CI-templates/.github/workflows/_code_freeze.yml@v0.25.2
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_code_freeze.yml@v0.25.2
     with:
       library-name: NeMo Run
       python-package: nemo_run
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 327215a..d782dd5 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -31,7 +31,7 @@ on:
         description: Branch to target for version bump
 jobs:
   release:
-    uses: NVIDIA/NeMo-FW-CI-templates/.github/workflows/_release_library.yml@v0.22.6
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_release_library.yml@v0.22.6
     with:
       release-ref: ${{ inputs.release-ref }}
       python-package: nemo_run
diff --git a/nemo_run/core/execution/base.py b/nemo_run/core/execution/base.py
index 66c2ec8..870ea8b 100644
--- a/nemo_run/core/execution/base.py
+++ b/nemo_run/core/execution/base.py
@@ -165,6 +165,9 @@ class Executor(ConfigurableMixin):
             os.makedirs(os.path.join(self.job_dir, launcher.nsys_folder), exist_ok=True)
             return launcher.get_nsys_prefix(profile_dir=self.job_dir)
 
+    def get_nsys_entrypoint(self) -> str:
+        return ("nsys", "")
+
     def supports_launcher_transform(self) -> bool:
         return False
 
diff --git a/nemo_run/core/execution/launcher.py b/nemo_run/core/execution/launcher.py
index f24a1d2..7e6c14a 100644
--- a/nemo_run/core/execution/launcher.py
+++ b/nemo_run/core/execution/launcher.py
@@ -24,6 +24,7 @@ class Launcher(ConfigurableMixin):
             "--cuda-event-trace=false",
         ]
     )
+    nsys_gpu_metrics: bool = False
 
     def get_nsys_prefix(self, profile_dir: str) -> Optional[list[str]]:
         """Make a command prefix for nsys profiling"""
diff --git a/nemo_run/core/execution/lepton.py b/nemo_run/core/execution/lepton.py
index 254cc92..d5e8511 100644
--- a/nemo_run/core/execution/lepton.py
+++ b/nemo_run/core/execution/lepton.py
@@ -13,15 +13,12 @@ from typing import Any, List, Optional, Set, Type
 from invoke.context import Context
 from leptonai.api.v2.client import APIClient
 from leptonai.api.v1.types.affinity import LeptonResourceAffinity
-from leptonai.api.v1.types.common import LeptonVisibility, Metadata
+from leptonai.api.v1.types.common import Metadata, LeptonVisibility
 from leptonai.api.v1.types.dedicated_node_group import DedicatedNodeGroup
 from leptonai.api.v1.types.deployment import (
     EnvVar,
     LeptonContainer,
-    LeptonDeployment,
-    LeptonDeploymentUserSpec,
     Mount,
-    ResourceRequirement,
 )
 from leptonai.api.v1.types.job import LeptonJob, LeptonJobState, LeptonJobUserSpec
 from leptonai.api.v1.types.replica import Replica
@@ -85,7 +82,7 @@ class LeptonExecutor(Executor):
             full_command = ["sh", "-c", cmd]
             return full_command
 
-    def move_data(self, sleep: float = 10) -> None:
+    def move_data(self, sleep: float = 10, timeout: int = 600, poll_interval: int = 5) -> None:
         """
         Moves job directory into remote storage and deletes the workload after completion.
         """
@@ -94,34 +91,60 @@ class LeptonExecutor(Executor):
         node_group_id = self._node_group_id(client)
         valid_node_ids = self._valid_node_ids(node_group_id, client)
 
-        spec = LeptonDeploymentUserSpec(
-            container=LeptonContainer(
-                image="busybox:1.37.0",  # Use a very low resource container
-                command=cmd,
-            ),
-            mounts=[Mount(**mount) for mount in self.mounts],
-        )
-        spec.resource_requirement = ResourceRequirement(
+        job_spec = LeptonJobUserSpec(
             resource_shape="cpu.small",
             affinity=LeptonResourceAffinity(
                 allowed_dedicated_node_groups=[node_group_id.metadata.id_],
                 allowed_nodes_in_node_group=valid_node_ids,
             ),
-            min_replicas=1,
-            max_replicas=1,
+            container=LeptonContainer(
+                image="busybox:1.37.0",
+                command=cmd,
+            ),
+            completions=1,
+            parallelism=1,
+            mounts=[Mount(**mount) for mount in self.mounts],
         )
+
         custom_name = f"data-mover-{int(datetime.now().timestamp())}"
 
-        deployment = LeptonDeployment(
+        job = LeptonJob(
             metadata=Metadata(
                 id=custom_name,
                 name=custom_name,
                 visibility=LeptonVisibility("private"),
             ),
-            spec=spec,
+            spec=job_spec,
         )
 
-        client.deployment.create(deployment)
+        response = client.job.create(job)
+        job_id = response.metadata.id_
+
+        start_time = time.time()
+        count = 0
+
+        while True:
+            if time.time() - start_time > timeout:
+                raise TimeoutError(f"Job {job_id} did not complete within {timeout} seconds.")
+            current_job = client.job.get(job_id)
+            current_job_status = current_job.status.state
+            if count > 0 and current_job_status in [
+                LeptonJobState.Completed,
+                LeptonJobState.Failed,
+                LeptonJobState.Unknown,
+            ]:
+                break
+            count += 1
+            time.sleep(poll_interval)
+
+        if current_job_status != LeptonJobState.Completed:
+            raise RuntimeError(f"Job {job_id} failed with status: {current_job_status}")
+
+        delete_success = client.job.delete(job_id)
+        if delete_success:
+            logging.info(f"Successfully deleted job {job_id}")
+        else:
+            logging.error(f"Failed to delete job {job_id}")
 
     def _node_group_id(self, client: APIClient) -> DedicatedNodeGroup:
         """
diff --git a/nemo_run/core/execution/slurm.py b/nemo_run/core/execution/slurm.py
index 3f66272..2432385 100644
--- a/nemo_run/core/execution/slurm.py
+++ b/nemo_run/core/execution/slurm.py
@@ -414,7 +414,6 @@ class SlurmExecutor(Executor):
                 )
             )
 
-        main_executor.env_vars = {}
         return main_executor
 
     def __post_init__(self):
@@ -547,7 +546,18 @@ class SlurmExecutor(Executor):
     def get_launcher_prefix(self) -> Optional[list[str]]:
         launcher = self.get_launcher()
         if launcher.nsys_profile:
-            return launcher.get_nsys_prefix(profile_dir=f"/{RUNDIR_NAME}")
+            nsys_prefix = launcher.get_nsys_prefix(profile_dir=f"/{RUNDIR_NAME}")
+            if launcher.nsys_gpu_metrics:
+                nsys_prefix += ["$GPU_METRICS_FLAG"]
+        return nsys_prefix
+
+    def get_nsys_entrypoint(self) -> str:
+        launcher = self.get_launcher()
+        entrypoint, postfix = "nsys", ""
+        if launcher.nsys_gpu_metrics:
+            entrypoint = 'bash -c \'GPU_METRICS_FLAG=""; if [ "$SLURM_PROCID" -eq 0 ]; then GPU_METRICS_FLAG="--gpu-metrics-devices=all"; fi; nsys'
+            postfix = "'"
+        return (entrypoint, postfix)
 
     def supports_launcher_transform(self) -> bool:
         return True if isinstance(self.get_launcher(), SlurmTemplate) else False
diff --git a/nemo_run/run/ray/slurm.py b/nemo_run/run/ray/slurm.py
index abb63fa..639247a 100644
--- a/nemo_run/run/ray/slurm.py
+++ b/nemo_run/run/ray/slurm.py
@@ -117,6 +117,7 @@ class SlurmRayRequest:
     command: Optional[str] = None
     workdir: Optional[str] = None
     nemo_run_dir: Optional[str] = None
+    command_groups: Optional[list[list[str]]] = None
     launch_cmd: list[str]
 
     @staticmethod
@@ -234,6 +235,60 @@ class SlurmRayRequest:
             "gres_specification": get_gres_specification(),
         }
 
+        if self.command_groups:
+            srun_commands: list[str] = []
+            group_env_vars: list[list[str]] = []
+
+            for idx, group in enumerate(self.command_groups):
+                if idx == 0:
+                    continue
+
+                if self.executor.run_as_group and len(self.executor.resource_group) == len(
+                    self.command_groups
+                ):
+                    req = self.executor.resource_group[idx]
+                    env_list = [f"export {k.upper()}={v}" for k, v in req.env_vars.items()]
+                    group_env_vars.append(env_list)
+                    container_flags = get_srun_flags(req.container_mounts, req.container_image)
+                    srun_args = ["--wait=60", "--kill-on-bad-exit=1", "--overlap"]
+                    srun_args.extend(req.srun_args or [])
+                else:
+                    container_flags = get_srun_flags(
+                        self.executor.container_mounts, self.executor.container_image
+                    )
+                    srun_args = ["--wait=60", "--kill-on-bad-exit=1", "--overlap"]
+                    srun_args.extend(self.executor.srun_args or [])
+                    group_env_vars.append([])
+
+                stdout_path = os.path.join(self.cluster_dir, "logs", f"ray-overlap-{idx}.out")
+                stderr_flags = []
+                if not self.executor.stderr_to_stdout:
+                    stderr_flags = [
+                        "--error",
+                        os.path.join(self.cluster_dir, "logs", f"ray-overlap-{idx}.err"),
+                    ]
+
+                srun_cmd = " ".join(
+                    list(
+                        map(
+                            lambda arg: arg if isinstance(arg, noquote) else shlex.quote(arg),
+                            [
+                                "srun",
+                                "--output",
+                                noquote(stdout_path),
+                                *stderr_flags,
+                                container_flags,
+                                *srun_args,
+                            ],
+                        )
+                    )
+                )
+                command = " ".join(group)
+                srun_commands.append(f"{srun_cmd} {command} &")
+
+            vars_to_fill["srun_commands"] = srun_commands
+            vars_to_fill["group_env_vars"] = group_env_vars
+
         if self.pre_ray_start_commands:
             vars_to_fill["pre_ray_start_commands"] = "\n".join(self.pre_ray_start_commands)
 
@@ -398,6 +453,7 @@ Useful Commands
         dryrun: bool = False,
         command: Optional[str] = None,
         workdir: Optional[str] = None,
+        command_groups: Optional[list[list[str]]] = None,
     ) -> Any:
         """Create (or reuse) a Slurm-backed Ray cluster and return its job-id.
 
@@ -416,6 +472,9 @@ Useful Commands
             Optional command executed after the Ray head node is ready (e.g. ``ray job submit``).
         workdir : str | None
             Remote working directory that becomes the CWD inside the container.
+        command_groups : list[list[str]] | None
+            Additional commands (one per group) executed via ``srun`` with ``--overlap``
+            after the cluster is started.
 
         Returns
         -------
@@ -433,6 +492,7 @@ Useful Commands
             pre_ray_start_commands=pre_ray_start_commands,
             command=command,
             workdir=workdir,
+            command_groups=command_groups,
             launch_cmd=["sbatch", "--requeue", "--parsable", "--dependency=singleton"],
         ).materialize()
 
@@ -1094,6 +1154,7 @@ Useful Commands (to be run on the login node of the Slurm cluster)
         runtime_env_yaml: Optional[str] | None = None,
         pre_ray_start_commands: Optional[list[str]] = None,
         dryrun: bool = False,
+        command_groups: Optional[list[list[str]]] = None,
     ):
         """Submit a Ray job via Slurm and return a *live* SlurmRayJob helper.
 
@@ -1106,6 +1167,7 @@ Useful Commands (to be run on the login node of the Slurm cluster)
                 executor=my_slurm_executor,
                 command="python train.py",
                 workdir="./src",
+                command_groups=[["echo", "hello"]],
             )
         """
         # ------------------------------------------------------------------
@@ -1212,6 +1274,7 @@ Useful Commands (to be run on the login node of the Slurm cluster)
             dryrun=dryrun,
             command=command,
             workdir=remote_workdir,
+            command_groups=command_groups,
         )
 
         self.job_id = job_id
diff --git a/nemo_run/run/ray/templates/ray.sub.j2 b/nemo_run/run/ray/templates/ray.sub.j2
index 4375e15..6d0eed1 100644
--- a/nemo_run/run/ray/templates/ray.sub.j2
+++ b/nemo_run/run/ray/templates/ray.sub.j2
@@ -295,6 +295,20 @@ echo "[INFO] Ray cluster information saved to $CLUSTER_DIR/ray_cluster_info.json
 
 ########################################################
 
+{% if srun_commands %}
+# Run extra commands
+{% for srun_command in srun_commands %}
+{%- if loop.index <= group_env_vars|length %}
+{%- for env_var in group_env_vars[loop.index - 1] %}
+{{env_var}}
+{%- endfor %}
+{%- endif %}
+
+{{srun_command}}
+{% endfor %}
+########################################################
+{% endif -%}
+
 # We can now launch a job on this cluster
 # We do so by launching a driver process on the physical node that the head node is on
 # This driver process is responsible for launching a job on the Ray cluster
diff --git a/nemo_run/run/torchx_backend/packaging.py b/nemo_run/run/torchx_backend/packaging.py
index e915e9b..093d190 100644
--- a/nemo_run/run/torchx_backend/packaging.py
+++ b/nemo_run/run/torchx_backend/packaging.py
@@ -225,15 +225,14 @@ def package(
         nsys_prefix = executor.get_launcher_prefix()
         if nsys_prefix:
             role.args = [role.entrypoint] + role.args
-            role.entrypoint = "nsys"
-            role.args = nsys_prefix + role.args
+            role.entrypoint, nsys_postfix = executor.get_nsys_entrypoint()
+            role.args = nsys_prefix + role.args + [nsys_postfix]
 
     if metadata:
         if USE_WITH_RAY_CLUSTER_KEY in metadata:
             assert isinstance(executor, SlurmExecutor), (
                 f"{USE_WITH_RAY_CLUSTER_KEY} is only supported for SlurmExecutor"
             )
-            assert len(app_def.roles) == 1, "Only one command is supported for Ray jobs."
 
     app_def.metadata = metadata
     return app_def
@@ -241,7 +240,13 @@ def package(
 
 def merge_executables(app_defs: Iterator[specs.AppDef], name: str) -> specs.AppDef:
     result = specs.AppDef(name=name, roles=[])
-    for app_def in app_defs:
+    result.metadata = {}
+    for idx, app_def in enumerate(app_defs):
+        metadata = app_def.metadata or {}
+        if USE_WITH_RAY_CLUSTER_KEY in metadata:
+            assert idx == 0, f"{USE_WITH_RAY_CLUSTER_KEY} is only supported for the first command"
+
+        result.metadata.update(metadata)
         result.roles.extend(app_def.roles)
     return result
 
diff --git a/nemo_run/run/torchx_backend/schedulers/slurm.py b/nemo_run/run/torchx_backend/schedulers/slurm.py
index b65e050..48c6808 100644
--- a/nemo_run/run/torchx_backend/schedulers/slurm.py
+++ b/nemo_run/run/torchx_backend/schedulers/slurm.py
@@ -102,8 +102,17 @@ class SlurmTunnelScheduler(SchedulerMixin, SlurmScheduler):  # type: ignore
 
         executor.package(packager=executor.packager, job_name=Path(job_dir).name)
 
+        values = executor.macro_values()
+
         if app.metadata and app.metadata.get(USE_WITH_RAY_CLUSTER_KEY, False):
-            assert len(app.roles) == 1, "Only one command is supported for Ray jobs."
+            srun_cmds: list[list[str]] = []
+
+            for role in app.roles:
+                if values:
+                    role = values.apply(role)
+                srun_cmd = [role.entrypoint] + role.args
+                srun_cmds.append([" ".join(srun_cmd)])
+
             command = [app.roles[0].entrypoint] + app.roles[0].args
             req = SlurmRayRequest(
                 name=app.roles[0].name,
@@ -114,12 +123,12 @@ class SlurmTunnelScheduler(SchedulerMixin, SlurmScheduler):  # type: ignore
                 executor=executor,
                 workdir=f"/{RUNDIR_NAME}/code",
                 nemo_run_dir=os.path.join(executor.tunnel.job_dir, Path(job_dir).name),
+                command_groups=srun_cmds,
             )
         else:
             srun_cmds: list[list[str]] = []
             jobs = []
             envs = {}
-            values = executor.macro_values()
 
             if values:
                 executor.env_vars = {


Test patch:
diff --git a/test/core/execution/artifacts/group_resource_req_slurm.sh b/test/core/execution/artifacts/group_resource_req_slurm.sh
index 3bbde92..a68b77d 100644
--- a/test/core/execution/artifacts/group_resource_req_slurm.sh
+++ b/test/core/execution/artifacts/group_resource_req_slurm.sh
@@ -30,6 +30,7 @@ nodes_array=($nodes)
 head_node=${nodes_array[0]}
 head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)
 
+export CUSTOM_ENV_1=some_value_1
 export ENV_VAR=value
 
 
diff --git a/test/core/execution/test_base.py b/test/core/execution/test_base.py
index 8b6d69d..c122b51 100644
--- a/test/core/execution/test_base.py
+++ b/test/core/execution/test_base.py
@@ -126,6 +126,11 @@ class TestExecutor:
         executor = Executor(launcher="torchrun")
         assert isinstance(executor.get_launcher(), Torchrun)
 
+    def test_get_nsys_entrypoint(self):
+        mock_launcher = Launcher()
+        executor = Executor(launcher=mock_launcher)
+        assert executor.get_nsys_entrypoint() == ("nsys", "")
+
     def test_cleanup(self):
         executor = Executor()
         assert executor.cleanup("handle") is None
diff --git a/test/core/execution/test_lepton.py b/test/core/execution/test_lepton.py
index cbe93a9..821c8d9 100644
--- a/test/core/execution/test_lepton.py
+++ b/test/core/execution/test_lepton.py
@@ -22,13 +22,10 @@ import pytest
 from leptonai.api.v1.types.common import LeptonVisibility, Metadata
 from leptonai.api.v1.types.deployment import (
     LeptonContainer,
-    LeptonDeployment,
-    LeptonDeploymentUserSpec,
     LeptonResourceAffinity,
     Mount,
-    ResourceRequirement,
 )
-from leptonai.api.v1.types.job import LeptonJob
+from leptonai.api.v1.types.job import LeptonJob, LeptonJobUserSpec
 
 from nemo_run.core.execution.lepton import LeptonExecutor, LeptonJobState
 from nemo_run.core.packaging.git import GitArchivePackager
@@ -173,19 +170,15 @@ class TestLeptonExecutor:
     @patch("nemo_run.core.execution.lepton.APIClient")
     def test_move_data_success(self, mock_APIClient, mock_datetime, mock_copy):
         mock_instance = MagicMock()
-        mock_deployment_api = MagicMock()
-        mock_instance.deployment = mock_deployment_api
-        mock_deployment_api.create = MagicMock()
+        mock_job_api = MagicMock()
+        mock_instance.job = mock_job_api
         mock_copy.return_value = ["sh", "-c", "echo 'hello world'"]
         mock_APIClient.return_value = mock_instance
         mock_client = mock_APIClient.return_value
         mock_nodegroup = MagicMock()
-        mock_datetime = MagicMock()
         mock_datetime_now = MagicMock()
-        mock_datetime.now = mock_datetime_now
-        mock_timestamp = MagicMock()
-        mock_datetime_now.timestamp = MagicMock()
-        mock_timestamp.return_value = "1"
+        mock_datetime.now.return_value = mock_datetime_now
+        mock_datetime_now.timestamp.return_value = 1
         mock_client.nodegroup = mock_nodegroup
         mock_nodegroup.list_all.return_value = [
             SimpleNamespace(metadata=SimpleNamespace(name="123456", id_="my-node-id"))
@@ -193,6 +186,7 @@ class TestLeptonExecutor:
         mock_nodegroup.list_nodes.return_value = [
             SimpleNamespace(metadata=SimpleNamespace(id_="10-10-10-10"))
         ]
+        mock_job_api.get.return_value = SimpleNamespace(status=SimpleNamespace(state="Completed"))
 
         executor = LeptonExecutor(
             container_image="nvcr.io/nvidia/test:latest",
@@ -203,32 +197,35 @@ class TestLeptonExecutor:
 
         executor.move_data()
 
-        spec = LeptonDeploymentUserSpec(
-            container=LeptonContainer(
-                image="busybox:1.37.0",
-                command=["sh", "-c", "echo 'hello world'"],
-            ),
-            mounts=[Mount(path="/workspace", mount_path="/workspace")],
-        )
-        spec.resource_requirement = ResourceRequirement(
+        expected_cmd = ["sh", "-c", "echo 'hello world'"]
+        expected_spec = LeptonJobUserSpec(
             resource_shape="cpu.small",
             affinity=LeptonResourceAffinity(
                 allowed_dedicated_node_groups=["my-node-id"],
                 allowed_nodes_in_node_group=["10-10-10-10"],
             ),
-            min_replicas=1,
-            max_replicas=1,
+            container=LeptonContainer(
+                image="busybox:1.37.0",
+                command=expected_cmd,
+            ),
+            completions=1,
+            parallelism=1,
+            mounts=[Mount(path="/workspace", mount_path="/workspace")],
         )
+
         custom_name = "data-mover-1"
-        deployment = LeptonDeployment(
+        expected_job = LeptonJob(
             metadata=Metadata(
-                id=custom_name, name=custom_name, visibility=LeptonVisibility("private")
+                id=custom_name,
+                name=custom_name,
+                visibility=LeptonVisibility("private"),
             ),
-            spec=spec,
+            spec=expected_spec,
         )
 
-        mock_copy.assert_called_once_with("", "")
-        mock_deployment_api.create.assert_called_once_with(deployment)
+        mock_copy.assert_called_once_with(executor.job_dir, executor.lepton_job_dir)
+        mock_job_api.create.assert_called_once_with(expected_job)
+        mock_job_api.delete.assert_called_once_with(mock_job_api.create.return_value.metadata.id_)
 
     def test_node_group_id(self):
         mock_client = MagicMock(
diff --git a/test/core/execution/test_slurm.py b/test/core/execution/test_slurm.py
index 0a2063b..4c68111 100644
--- a/test/core/execution/test_slurm.py
+++ b/test/core/execution/test_slurm.py
@@ -170,10 +170,38 @@ class TestSlurmExecutorExtended:
         launcher_mock = MagicMock()
         launcher_mock.nsys_profile = True
         launcher_mock.get_nsys_prefix.return_value = ["nsys", "profile"]
+        launcher_mock.nsys_gpu_metrics = False
 
         with patch.object(executor, "get_launcher", return_value=launcher_mock):
             assert executor.get_launcher_prefix() == ["nsys", "profile"]
 
+    def test_get_launcher_prefix_with_gpu_metrics(self):
+        """Test the get_launcher_prefix method with nsys_profile when gpu metrics is enabled."""
+        executor = SlurmExecutor(account="test")
+
+        # Test with launcher that has nsys_profile
+        launcher_mock = MagicMock()
+        launcher_mock.nsys_profile = True
+        launcher_mock.get_nsys_prefix.return_value = ["nsys", "profile"]
+        launcher_mock.nsys_gpu_metrics = True
+
+        with patch.object(executor, "get_launcher", return_value=launcher_mock):
+            assert executor.get_launcher_prefix() == ["nsys", "profile", "$GPU_METRICS_FLAG"]
+
+    def test_get_nsys_entrypoint(self):
+        """Test the get_nsys_entrypoint method with nsys_profile."""
+        executor = SlurmExecutor(account="test")
+
+        # Test with launcher that has nsys_profile
+        launcher_mock = MagicMock()
+        launcher_mock.nsys_gpu_metrics = True
+
+        with patch.object(executor, "get_launcher", return_value=launcher_mock):
+            assert executor.get_nsys_entrypoint() == (
+                'bash -c \'GPU_METRICS_FLAG=""; if [ "$SLURM_PROCID" -eq 0 ]; then GPU_METRICS_FLAG="--gpu-metrics-devices=all"; fi; nsys',
+                "'",
+            )
+
     def test_supports_launcher_transform(self):
         """Test the supports_launcher_transform method."""
         executor = SlurmExecutor(account="test")
diff --git a/test/run/ray/test_slurm_ray_request.py b/test/run/ray/test_slurm_ray_request.py
index 32a9673..6609872 100644
--- a/test/run/ray/test_slurm_ray_request.py
+++ b/test/run/ray/test_slurm_ray_request.py
@@ -405,3 +405,179 @@ class TestSlurmRayRequest:
 
         with pytest.raises(AssertionError, match="array is not supported"):
             request.materialize()
+
+    def test_command_groups_env_vars(self):
+        """Test environment variables are properly set for each command group."""
+        # Create executor with environment variables
+        executor = SlurmExecutor(
+            account="test_account",
+            env_vars={"GLOBAL_ENV": "global_value"},
+        )
+        executor.run_as_group = True
+
+        # Create resource groups with different env vars
+        resource_group = [
+            SlurmExecutor.ResourceRequest(
+                packager=Mock(),
+                nodes=1,
+                ntasks_per_node=1,
+                container_image="image1",
+                env_vars={"GROUP1_ENV": "group1_value"},
+                container_mounts=["/mount1"],
+            ),
+            SlurmExecutor.ResourceRequest(
+                packager=Mock(),
+                nodes=1,
+                ntasks_per_node=1,
+                container_image="image2",
+                env_vars={"GROUP2_ENV": "group2_value"},
+                container_mounts=["/mount2"],
+            ),
+        ]
+        executor.resource_group = resource_group
+        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_groups=[["cmd0"], ["cmd1"], ["cmd2"]],
+            launch_cmd=["sbatch", "--parsable"],
+        )
+
+        script = request.materialize()
+
+        # Check global env vars are set in setup section
+        assert "export GLOBAL_ENV=global_value" in script
+
+        # Check that command groups generate srun commands (excluding the first one)
+        # The template should have a section for srun_commands
+        assert "# Run extra commands" in script
+        assert "srun" in script
+        assert "cmd1" in script  # First command group after skipping index 0
+        assert "cmd2" in script  # Second command group
+
+    def test_command_groups_without_resource_group(self):
+        """Test command groups work without resource groups."""
+        executor = SlurmExecutor(
+            account="test_account",
+            env_vars={"GLOBAL_ENV": "global_value"},
+        )
+        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_groups=[["cmd0"], ["cmd1"]],
+            launch_cmd=["sbatch", "--parsable"],
+        )
+
+        script = request.materialize()
+
+        # Should have global env vars
+        assert "export GLOBAL_ENV=global_value" in script
+
+        # Should have srun commands for overlapping groups (skipping first)
+        assert "srun" in script
+        assert "--overlap" in script
+        assert "cmd1" in script  # Second command in the list (index 1)
+
+    def test_env_vars_formatting(self):
+        """Test that environment variables are properly formatted as export statements."""
+        executor = SlurmExecutor(
+            account="test_account",
+            env_vars={
+                "VAR_WITH_SPACES": "value with spaces",
+                "PATH_VAR": "/usr/bin:/usr/local/bin",
+                "EMPTY_VAR": "",
+                "NUMBER_VAR": "123",
+            },
+        )
+        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,
+            launch_cmd=["sbatch", "--parsable"],
+        )
+
+        script = request.materialize()
+
+        # Check all environment variables are properly exported
+        assert "export VAR_WITH_SPACES=value with spaces" in script
+        assert "export PATH_VAR=/usr/bin:/usr/local/bin" in script
+        assert "export EMPTY_VAR=" in script
+        assert "export NUMBER_VAR=123" in script
+
+    def test_group_env_vars_integration(self):
+        """Test full integration of group environment variables matching the artifact pattern."""
+        # This test verifies the behavior seen in group_resource_req_slurm.sh
+        executor = SlurmExecutor(
+            account="your_account",
+            partition="your_partition",
+            time="00:30:00",
+            nodes=1,
+            ntasks_per_node=8,
+            gpus_per_node=8,
+            container_image="some-image",
+            container_mounts=["/some/job/dir/sample_job:/nemo_run"],
+            env_vars={"ENV_VAR": "value"},
+        )
+        executor.run_as_group = True
+
+        # Set up resource groups with specific env vars
+        resource_group = [
+            # First group (index 0) - for the head/main command
+            SlurmExecutor.ResourceRequest(
+                packager=Mock(),
+                nodes=1,
+                ntasks_per_node=8,
+                container_image="some-image",
+                env_vars={"CUSTOM_ENV_1": "some_value_1"},
+                container_mounts=["/some/job/dir/sample_job:/nemo_run"],
+            ),
+            # Second group (index 1)
+            SlurmExecutor.ResourceRequest(
+                packager=Mock(),
+                nodes=1,
+                ntasks_per_node=8,
+                container_image="different_container_image",
+                env_vars={"CUSTOM_ENV_1": "some_value_1"},
+                container_mounts=["/some/job/dir/sample_job:/nemo_run"],
+            ),
+        ]
+        executor.resource_group = resource_group
+
+        # Mock tunnel
+        tunnel_mock = Mock(spec=SSHTunnel)
+        tunnel_mock.job_dir = "/some/job/dir"
+        executor.tunnel = tunnel_mock
+
+        request = SlurmRayRequest(
+            name="sample_job",
+            cluster_dir="/some/job/dir/sample_job",
+            template_name="ray.sub.j2",
+            executor=executor,
+            command_groups=[
+                ["bash ./scripts/start_server.sh"],
+                ["bash ./scripts/echo.sh server_host=$het_group_host_0"],
+            ],
+            launch_cmd=["sbatch", "--parsable"],
+        )
+
+        script = request.materialize()
+
+        # Verify the pattern matches the artifact:
+        # 1. Global env vars should be exported in setup
+        assert "export ENV_VAR=value" in script
+
+        # The template should include group_env_vars for proper env var handling per command
+        # (The actual env var exports per command happen in the template rendering)