adding test for slurm rescheduling

Author: djarecka

pydra/engine/tests/test_submitter.py

@@ -8,6 +8,7 @@
 
 from .utils import gen_basic_wf
 from ..core import Workflow
+from ..task import ShellCommandTask
 from ..submitter import Submitter
 from ... import mark
 
@@ -261,3 +262,58 @@ def test_slurm_args_2(tmpdir):
     with pytest.raises(RuntimeError, match="Error returned from sbatch:"):
         with Submitter("slurm", sbatch_args="-N1 --invalid") as sub:
             sub(task)
+
+
+@pytest.mark.skipif(not slurm_available, reason="slurm not installed")
+def test_slurm_cancel_rerun(tmpdir):
+    """ testing that tasks run with slurm is re-queue
+        Running wf with 2 tasks, one sleeps and the other trying to get
+        job_id of the first task and cancel it.
+        The first job should be re-queue and finish without problem.
+        (possibly has to be improved, in theory cancel job might finish before cancel)
+    """
+
+    @mark.task
+    def sleep(x):
+        time.sleep(x)
+        return x
+
+    @mark.task
+    def cancel(job_name_part):
+        import subprocess as sp
+
+        # getting the job_id of the first job that sleeps
+        job_id = ""
+        while job_id == "":
+            time.sleep(1)
+            id_p1 = sp.Popen(["squeue"], stdout=sp.PIPE)
+            id_p2 = sp.Popen(
+                ["grep", job_name_part], stdin=id_p1.stdout, stdout=sp.PIPE
+            )
+            id_p3 = sp.Popen(["awk", "{print $1}"], stdin=id_p2.stdout, stdout=sp.PIPE)
+            job_id = id_p3.communicate()[0].decode("utf-8").strip()
+
+        # # canceling the job
+        proc1 = sp.run(["scancel", job_id])
+        # checking the status of jobs with the name; returning the last item
+        proc2 = sp.run(["sacct", "-j", job_id], stdout=sp.PIPE, stderr=sp.PIPE)
+        return proc2.stdout.decode("utf-8").strip()  # .split("\n")[-1]
+
+    wf = Workflow(name="wf", input_spec=["x", "job_name"], cache_dir=tmpdir)
+    wf.add(sleep(name="sleep", x=wf.lzin.x))
+    wf.add(cancel(nane="cancel", job_name_part=wf.lzin.job_name))
+    # this is a job name for x=10, if x is different checksum and jobname would have to be updated
+    wf.inputs.x = 20
+    wf.inputs.job_name = "sleep"
+
+    wf.set_output([("out", wf.sleep.lzout.out), ("canc_out", wf.cancel.lzout.out)])
+    with Submitter("slurm") as sub:
+        sub(wf)
+
+    res = wf.result()
+    assert res.output.out == 20
+    # checking if indeed the sleep-task job was cancelled by cancel-task
+    assert "CANCELLED" in res.output.canc_out
+    breakpoint()
+    script_dir = tmpdir / "SlurmWorker_scripts"
+    assert script_dir.exists()

pydra/engine/workers.py

@@ -300,8 +300,10 @@ async def _submit_job(self, batchscript, name, checksum, cache_dir):
             # Exception: Polling / job failure
             done = await self._poll_job(jobid)
             if done:
-                if done == "CANCELLED":
-                    (cache_dir / f"{checksum}.lock").unlink(missing_ok=True)
+                if done in ["CANCELLED", "TIMEOUT", "PREEMPTED"]:
+                    if (cache_dir / f"{checksum}.lock").exists():
+                        # for pyt3.8 we could you missing_ok=True
+                        (cache_dir / f"{checksum}.lock").unlink()
                     cmd_re = ("scontrol", "requeue", jobid)
                     await read_and_display_async(*cmd_re, hide_display=True)
                 else:
@@ -326,8 +328,8 @@ async def _verify_exit_code(self, jobid):
         m = self._sacct_re.search(stdout)
         error_file = self.error[jobid]
         if int(m.group("exit_code")) != 0 or m.group("status") != "COMPLETED":
-            if m.group("status") == "CANCELLED":
-                return "CANCELLED"
+            if m.group("status") in ["CANCELLED", "TIMEOUT", "PREEMPTED"]:
+                return m.group("status")
             elif m.group("status") in ["RUNNING", "PENDING"]:
                 return False
             # TODO: potential for requeuing