Repository: DataBiosphere/toil

Issue:
Ask for and handle Slurm signals when a Toil job is about to time out
Slurm supports a `--signal` option with which you can ask for a signal when a job is within, say, 60 seconds of its time limit.

Toil should ask for these signals, and it should shut down the Python process in an orderly fashion when they are received.

I think Python signal handlers (including the default one for Ctrl+C) should work fine here; [they can raise exceptions in the main thread](https://discuss.python.org/t/how-to-create-a-custom-keyboardinterrupt/19864/4), and MiniWDL [Uses a `finally` block to clean up its containers](https://github.com/chanzuckerberg/miniwdl/blob/c5bb92bb5e2c76a90d26d8c70e41cdb85db68627/WDL/runtime/backend/docker_swarm.py#L292), so implementing this ought to make a good dent in #4800.

┆Issue is synchronized with this [Jira Story](https://ucsc-cgl.atlassian.net/browse/TOIL-1504)
┆Issue Number: TOIL-1504



Discussion:

Resolving patch:
diff --git a/src/toil/batchSystems/slurm.py b/src/toil/batchSystems/slurm.py
--- a/src/toil/batchSystems/slurm.py
+++ b/src/toil/batchSystems/slurm.py
@@ -64,7 +64,9 @@ def prepareSubmission(self,
                               jobName: str,
                               job_environment: Optional[Dict[str, str]] = None,
                               gpus: Optional[int] = None) -> List[str]:
-            return self.prepareSbatch(cpu, memory, jobID, jobName, job_environment, gpus) + [f'--wrap={command}']
+            # Make sure to use exec so we can get Slurm's signals in the Toil
+            # worker instead of having an intervening Bash
+            return self.prepareSbatch(cpu, memory, jobID, jobName, job_environment, gpus) + [f'--wrap=exec {command}']
 
         def submitJob(self, subLine):
             try:
@@ -283,8 +285,26 @@ def prepareSbatch(self,
                           job_environment: Optional[Dict[str, str]],
                           gpus: Optional[int]) -> List[str]:
 
-            #  Returns the sbatch command line before the script to run
+            """
+            Returns the sbatch command line to run to queue the job.
+            """
+
+            # Start by naming the job
             sbatch_line = ['sbatch', '-J', f'toil_job_{jobID}_{jobName}']
+
+            # Make sure the job gets a signal before it disappears so that e.g.
+            # container cleanup finally blocks can run. Ask for SIGINT so we
+            # can get the default Python KeyboardInterrupt which third-party
+            # code is likely to plan for. Make sure to send it to the batch
+            # shell process with "B:", not to all the srun steps it launches
+            # (because there shouldn't be any). We cunningly replaced the batch
+            # shell process with the Toil worker process, so Toil should be
+            # able to get the signal.
+            #
+            # TODO: Add a way to detect when the job failed because it
+            # responded to this signal and use the right exit reason for it.
+            sbatch_line.append("--signal=B:INT@30")
+
             if gpus:
                 sbatch_line = sbatch_line[:1] + [f'--gres=gpu:{gpus}'] + sbatch_line[1:]
             environment = {}
diff --git a/src/toil/worker.py b/src/toil/worker.py
--- a/src/toil/worker.py
+++ b/src/toil/worker.py
@@ -496,13 +496,16 @@ def blockFn() -> bool:
     ##########################################
     #Trapping where worker goes wrong
     ##########################################
-    except Exception as e: #Case that something goes wrong in worker
+    except BaseException as e: #Case that something goes wrong in worker, or we are asked to stop
         traceback.print_exc()
         logger.error("Exiting the worker because of a failed job on host %s", socket.gethostname())
         if isinstance(e, CWL_UNSUPPORTED_REQUIREMENT_EXCEPTION):
             # We need to inform the leader that this is a CWL workflow problem
             # and it needs to inform its caller.
             failure_exit_code = CWL_UNSUPPORTED_REQUIREMENT_EXIT_CODE
+        elif isinstance(e, SystemExit) and isinstance(e.code, int) and e.code != 0:
+            # We're meant to be exiting with a particular code.
+            failure_exit_code = e.code
         AbstractFileStore._terminateEvent.set()
     finally:
         # Get rid of our deferred function manager now so we can't mistake it


Test patch: