fix

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

Author: hemildesai

nemo_run/core/execution/slurm.py

@@ -1090,7 +1090,13 @@ def get_container_flags(
         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:
+        # Check both top-level container_image and resource_group container images
+        has_container = self.executor.container_image is not None
+        if self.executor.run_as_group and self.executor.resource_group:
+            has_container = has_container or any(
+                rg.container_image is not None for rg in self.executor.resource_group
+            )
+        if not has_container:
             actual_job_dir = os.path.join(slurm_job_dir, job_directory_name)
             sbatch_script = sbatch_script.replace(f"/{RUNDIR_NAME}", actual_job_dir)
 

test/core/execution/artifacts/dummy_slurm.sh

@@ -33,7 +33,7 @@ export ENV_VAR=value
 
 # Command 1
 
-srun --output /root/sample_job/log-account-account.sample_job_%j_${SLURM_RESTART_COUNT:-0}.out --container-mounts /root/sample_job:/nemo_run --container-workdir /nemo_run/code --wait=60 --kill-on-bad-exit=1 cmd3 cmd4
+srun --output /root/sample_job/log-account-account.sample_job_%j_${SLURM_RESTART_COUNT:-0}.out --container-image test_image --container-mounts /root/sample_job:/nemo_run --container-workdir /nemo_run/code --wait=60 --kill-on-bad-exit=1 cmd3 cmd4
 
 exitcode=$?
 

test/core/execution/artifacts/ft_slurm.sh

@@ -62,7 +62,7 @@ echo "$SLURM_JOB_ID ${SLURM_RESTART_COUNT:-0} X" >> "$JOB_RESULTS_FILE"
 
 # Command 1
 
-srun --output /root/sample_job/log-account-account.sample_job_%j_${SLURM_RESTART_COUNT:-0}.out --container-mounts /root/sample_job:/nemo_run --container-workdir /nemo_run/code --wait=60 --kill-on-bad-exit=1 ft_launcher --ft-param-workload_check_interval 10 --ft-param-rank_heartbeat_timeout 10 --rdzv-backend c10d --rdzv-endpoint localhost:0 --rdzv-id 7680 --nnodes 1 --nproc-per-node 1 --node-rank 0 --tee 3 --no-python test_ft.sh
+srun --output /root/sample_job/log-account-account.sample_job_%j_${SLURM_RESTART_COUNT:-0}.out --container-image test_image --container-mounts /root/sample_job:/nemo_run --container-workdir /nemo_run/code --wait=60 --kill-on-bad-exit=1 ft_launcher --ft-param-workload_check_interval 10 --ft-param-rank_heartbeat_timeout 10 --rdzv-backend c10d --rdzv-endpoint localhost:0 --rdzv-id 7680 --nnodes 1 --nproc-per-node 1 --node-rank 0 --tee 3 --no-python test_ft.sh
 
 exitcode=$?
 

test/core/execution/test_slurm_templates.py

@@ -54,6 +54,7 @@ def dummy_slurm_request_with_artifact(
             account="account",
             job_dir="/root/sample_job",
             tunnel=LocalTunnel(job_dir="/root"),
+            container_image="test_image",
         )
         slurm_config.job_name = "sample_job"
         max_retries = 3
@@ -79,6 +80,7 @@ def ft_slurm_request_with_artifact(
             account="account",
             job_dir="/root/sample_job",
             tunnel=LocalTunnel(job_dir="/root/"),
+            container_image="test_image",
         )
         slurm_config.job_name = "sample_job"
         slurm_config.launcher = FaultTolerance(