feat: support container-image None in slurm

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

Author: hemildesai

nemo_run/core/execution/slurm.py

@@ -369,12 +369,12 @@ def merge(
         main_executor.run_as_group = True
 
         if main_executor.het_group_indices:
-            assert main_executor.heterogeneous, (
-                "heterogeneous must be True if het_group_indices is provided"
-            )
-            assert len(main_executor.het_group_indices) == num_tasks, (
-                "het_group_indices must be the same length as the number of tasks"
-            )
+            assert (
+                main_executor.heterogeneous
+            ), "heterogeneous must be True if het_group_indices is provided"
+            assert (
+                len(main_executor.het_group_indices) == num_tasks
+            ), "het_group_indices must be the same length as the number of tasks"
             assert all(
                 x <= y
                 for x, y in zip(
@@ -858,9 +858,9 @@ def materialize(self) -> str:
 
         sbatch_flags = []
         if self.executor.heterogeneous:
-            assert len(self.jobs) == len(self.executor.resource_group), (
-                f"Number of jobs {len(self.jobs)} must match number of resource group requests {len(self.executor.resource_group)}.\nIf you are just submitting a single job, make sure that heterogeneous=False in the executor."
-            )
+            assert (
+                len(self.jobs) == len(self.executor.resource_group)
+            ), f"Number of jobs {len(self.jobs)} must match number of resource group requests {len(self.executor.resource_group)}.\nIf you are just submitting a single job, make sure that heterogeneous=False in the executor."
             final_group_index = len(self.executor.resource_group) - 1
             if self.executor.het_group_indices:
                 final_group_index = self.executor.het_group_indices.index(
@@ -870,9 +870,9 @@ def materialize(self) -> str:
             for i in range(len(self.executor.resource_group)):
                 resource_req = self.executor.resource_group[i]
                 if resource_req.het_group_index is not None:
-                    assert self.executor.resource_group[i - 1].het_group_index is not None, (
-                        "het_group_index must be set for all requests in resource_group"
-                    )
+                    assert (
+                        self.executor.resource_group[i - 1].het_group_index is not None
+                    ), "het_group_index must be set for all requests in resource_group"
                     if (
                         i > 0
                         and resource_req.het_group_index
@@ -938,7 +938,18 @@ def get_container_flags(
             container_image: Optional[str],
             container_env: Optional[list[str]] = None,
         ) -> list[str]:
-            _container_flags = ["--container-image", container_image] if container_image else []
+            """Get srun flags for container or non-container mode.
+
+            For non-container mode, returns --chdir flag to set working directory.
+            For container mode, returns container-related flags (image, mounts, workdir, env).
+            """
+            if container_image is None:
+                # Non-container mode: use --chdir to set working directory
+                workdir = os.path.join(src_job_dir, "code")
+                return ["--chdir", workdir]
+
+            # Container mode: set up container mounts and workdir
+            _container_flags = ["--container-image", container_image]
 
             new_mounts = copy.deepcopy(base_mounts)
             for i, mount in enumerate(new_mounts):
@@ -1079,6 +1090,12 @@ def get_container_flags(
             vars_to_fill["fault_tol_job_results_file"] = self.launcher.job_results_file
 
         sbatch_script = fill_template("slurm.sh.j2", vars_to_fill)
+
+        # For non-container mode, substitute /{RUNDIR_NAME} paths with actual job directory
+        if self.executor.container_image is None:
+            actual_job_dir = os.path.join(slurm_job_dir, job_directory_name)
+            sbatch_script = sbatch_script.replace(f"/{RUNDIR_NAME}", actual_job_dir)
+
         return sbatch_script
 
     def __repr__(self) -> str:

test/core/execution/test_slurm.py

@@ -18,8 +18,10 @@
 
 import pytest
 
+from nemo_run.config import RUNDIR_NAME
 from nemo_run.core.execution.launcher import SlurmTemplate, Torchrun
 from nemo_run.core.execution.slurm import (
+    SlurmBatchRequest,
     SlurmExecutor,
     SlurmJobDetails,
     SlurmTunnelCallback,
@@ -403,3 +405,148 @@ def test_merge_mismatch(self):
                 [SlurmExecutor(account="account1"), SlurmExecutor(account="account2")],
                 num_tasks=3,
             )
+
+
+class TestSlurmBatchRequestNonContainerMode:
+    """Tests for non-container mode support (container_image=None)."""
+
+    @pytest.fixture
+    def executor_with_container(self):
+        """Create an executor with container image."""
+        executor = SlurmExecutor(
+            account="test_account",
+            partition="gpu",
+            nodes=2,
+            ntasks_per_node=8,
+            container_image="nvcr.io/nvidia/pytorch:24.01-py3",
+            container_mounts=["/data:/data"],
+        )
+        executor.job_name = "test-job"
+        executor.experiment_dir = "/local/experiments"
+        executor.job_dir = "/local/experiments/test-job"
+        executor.experiment_id = "exp-123"
+
+        # Mock tunnel
+        tunnel = MagicMock(spec=LocalTunnel)
+        tunnel.job_dir = "/remote/experiments/exp-123"
+        executor.tunnel = tunnel
+
+        return executor
+
+    @pytest.fixture
+    def executor_without_container(self):
+        """Create an executor without container image (non-container mode)."""
+        executor = SlurmExecutor(
+            account="test_account",
+            partition="gpu",
+            nodes=2,
+            ntasks_per_node=8,
+            container_image=None,  # Non-container mode
+        )
+        executor.job_name = "test-job"
+        executor.experiment_dir = "/local/experiments"
+        executor.job_dir = "/local/experiments/test-job"
+        executor.experiment_id = "exp-123"
+
+        # Mock tunnel
+        tunnel = MagicMock(spec=LocalTunnel)
+        tunnel.job_dir = "/remote/experiments/exp-123"
+        executor.tunnel = tunnel
+
+        return executor
+
+    def test_materialize_with_container_uses_container_flags(self, executor_with_container):
+        """Test that materialize uses container flags when container_image is set."""
+        request = SlurmBatchRequest(
+            launch_cmd=["sbatch", "--parsable"],
+            jobs=["test-job"],
+            command_groups=[["python train.py"]],
+            executor=executor_with_container,
+            max_retries=0,
+            extra_env={},
+        )
+
+        script = request.materialize()
+
+        # Should contain container flags
+        assert "--container-image" in script
+        assert "--container-mounts" in script
+        assert "--container-workdir" in script
+        # Should NOT contain --chdir (used for non-container mode)
+        assert "--chdir" not in script
+        # Should contain /nemo_run paths (not substituted)
+        assert f"/{RUNDIR_NAME}" in script
+
+    def test_materialize_without_container_uses_chdir(self, executor_without_container):
+        """Test that materialize uses --chdir when container_image is None."""
+        request = SlurmBatchRequest(
+            launch_cmd=["sbatch", "--parsable"],
+            jobs=["test-job"],
+            command_groups=[["python train.py"]],
+            executor=executor_without_container,
+            max_retries=0,
+            extra_env={},
+        )
+
+        script = request.materialize()
+
+        # Should contain --chdir flag for working directory
+        assert "--chdir" in script
+        # Should NOT contain container flags
+        assert "--container-image" not in script
+        assert "--container-mounts" not in script
+        assert "--container-workdir" not in script
+
+    def test_materialize_without_container_substitutes_rundir_paths(
+        self, executor_without_container
+    ):
+        """Test that /{RUNDIR_NAME} paths are substituted with actual paths in non-container mode."""
+        request = SlurmBatchRequest(
+            launch_cmd=["sbatch", "--parsable"],
+            jobs=["test-job"],
+            command_groups=[["python train.py"]],
+            executor=executor_without_container,
+            max_retries=0,
+            extra_env={},
+        )
+
+        script = request.materialize()
+
+        # Should NOT contain /nemo_run paths (should be substituted)
+        assert f"/{RUNDIR_NAME}/code" not in script
+        # Should contain the actual job directory path
+        actual_job_dir = "/remote/experiments/exp-123/test-job"
+        assert f"{actual_job_dir}/code" in script
+
+    def test_materialize_with_container_preserves_rundir_paths(self, executor_with_container):
+        """Test that /{RUNDIR_NAME} paths are NOT substituted when using container."""
+        request = SlurmBatchRequest(
+            launch_cmd=["sbatch", "--parsable"],
+            jobs=["test-job"],
+            command_groups=[["python train.py"]],
+            executor=executor_with_container,
+            max_retries=0,
+            extra_env={},
+        )
+
+        script = request.materialize()
+
+        # Should contain /nemo_run paths (not substituted for container mode)
+        assert f"/{RUNDIR_NAME}" in script
+
+    def test_non_container_mode_chdir_points_to_code_directory(self, executor_without_container):
+        """Test that --chdir in non-container mode points to the code directory."""
+        request = SlurmBatchRequest(
+            launch_cmd=["sbatch", "--parsable"],
+            jobs=["test-job"],
+            command_groups=[["python train.py"]],
+            executor=executor_without_container,
+            max_retries=0,
+            extra_env={},
+        )
+
+        script = request.materialize()
+
+        # The --chdir should point to {job_dir}/code
+        expected_chdir = "--chdir /remote/experiments/exp-123/test-job/code"
+        assert expected_chdir in script