Fix

Signed-off-by: Hemil Desai <[email protected]>

Author: hemildesai

src/nemo_run/core/execution/slurm.py

@@ -816,7 +816,7 @@ def materialize(self) -> str:
                     )
                 het_parameters.update(
                     {
-                        "job_name": f"{job_name_prefix}{self.jobs[i]}",
+                        "job_name": f"{job_details.job_name[:-2] if job_details.job_name.endswith('-0') else job_details.job_name}-{i}",
                         "nodes": resource_req.nodes,
                         "ntasks_per_node": resource_req.ntasks_per_node,
                         "gpus_per_node": resource_req.gpus_per_node,
@@ -995,7 +995,7 @@ def get_container_flags(
         return sbatch_script
 
     def __repr__(self) -> str:
-        return f"""{' '.join(self.cmd + ['$SBATCH_SCRIPT'])}
+        return f"""{" ".join(self.cmd + ["$SBATCH_SCRIPT"])}
 
 #----------------
 # SBATCH_SCRIPT

test/core/execution/test_slurm.py

@@ -691,3 +691,28 @@ def test_het_job_name_prefix(self, het_slurm_request_with_artifact):
         for job in het_request.jobs:
             expected = f"prefix_{job}"
             assert expected in sbatch_script, f"Expected job name '{expected}' not found in script"
+
+    def test_het_job_custom_details_job_name(self, het_slurm_request_with_artifact):
+        # Test that the job name from CustomJobDetails is used for heterogeneous slurm requests
+        from nemo_run.core.execution.slurm import SlurmJobDetails
+
+        het_request, _ = het_slurm_request_with_artifact
+
+        class CustomJobDetails(SlurmJobDetails):
+            @property
+            def stdout(self):
+                assert self.folder
+                return Path(self.folder) / "sbatch_job.out"
+
+            @property
+            def srun_stdout(self):
+                assert self.folder
+                return Path(self.folder) / "log_job.out"
+
+        custom_name = "custom_het_job"
+        het_request.slurm_config.job_details = CustomJobDetails(
+            job_name=custom_name, folder="/custom_folder"
+        )
+        sbatch_script = het_request.materialize()
+        for i in range(len(het_request.jobs)):
+            assert f"#SBATCH --job-name={custom_name}-{i}" in sbatch_script