fix [slurm]: Put srun in the right place

martius-lab · Apr 24, 2024 · 97211c2 · 97211c2
1 parent 6719fad
commit 97211c2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/cluster_utils/slurm_cluster_system.py b/cluster_utils/slurm_cluster_system.py
@@ -24,7 +24,7 @@
 echo "Job id: {id}, cluster id: ${{SLURM_JOB_ID}}, hostname: $(hostname), time: $(date)"
 echo
 
-srun {cmd}
+{cmd}
 rc=$?
 if [[ $rc == %(RETURN_CODE_FOR_RESUME)d ]]; then
     echo "exit with code %(RETURN_CODE_FOR_RESUME)d for resume"
@@ -231,7 +231,8 @@ def _generate_run_script(self, job: Job):
         runs_script_name = "job_{}_{}.sh".format(job.iteration, job.id)
         submission_dir = pathlib.Path(self.submission_dir)
         run_script_file_path = submission_dir / runs_script_name
-        cmd = job.generate_execution_cmd(self.paths)
+        # need to prefix the actual job command with `srun` so that --signal works.
+        cmd = job.generate_execution_cmd(self.paths, cmd_prefix="srun")
 
         stdout_file = run_script_file_path.with_suffix(".out")
         stderr_file = run_script_file_path.with_suffix(".err")

diff --git a/tests/test_slurm_cluster_system.py b/tests/test_slurm_cluster_system.py
@@ -76,7 +76,7 @@ def test_generate_run_script(job_data):
 
     assert job_data.job.run_script_path == str(expected_path)
 
-    job_cmd = job_data.job.generate_execution_cmd(job_data.paths)
+    job_cmd = job_data.job.generate_execution_cmd(job_data.paths, cmd_prefix="srun")
     run_script = run_script_path.read_text()
     assert (
         run_script