toil 5.12.0__py3-none-any.whl → 6.1.0__py3-none-any.whl

toil/batchSystems/slurm.py

@@ -16,13 +16,14 @@ import math
 import os
 from argparse import ArgumentParser, _ArgumentGroup
 from shlex import quote
-from typing import Callable, Dict, List, Optional, TypeVar, Union
+from typing import Dict, List, Optional, Tuple, TypeVar, Union
 
+from toil.batchSystems.abstractBatchSystem import BatchJobExitReason, EXIT_STATUS_UNAVAILABLE_VALUE
 from toil.batchSystems.abstractGridEngineBatchSystem import \
     AbstractGridEngineBatchSystem
 from toil.batchSystems.options import OptionSetter
-from toil.lib.misc import CalledProcessErrorStderr, call_command
 from toil.job import Requirer
+from toil.lib.misc import CalledProcessErrorStderr, call_command
 
 logger = logging.getLogger(__name__)
 
@@ -64,7 +65,9 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
                               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:
@@ -95,12 +98,12 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
                 logger.error("sbatch command failed")
                 raise e
 
-        def coalesce_job_exit_codes(self, batch_job_id_list: list) -> list:
+        def coalesce_job_exit_codes(self, batch_job_id_list: list) -> List[Union[int, Tuple[int, Optional[BatchJobExitReason]], None]]:
             """
             Collect all job exit codes in a single call.
             :param batch_job_id_list: list of Job ID strings, where each string has the form
             "<job>[.<task>]".
-            :return: list of job exit codes, associated with the list of job IDs.
+            :return: list of job exit codes or exit code, exit reason pairs associated with the list of job IDs.
             """
             logger.debug("Getting exit codes for slurm jobs: %s", batch_job_id_list)
             # Convert batch_job_id_list to list of integer job IDs.
@@ -111,7 +114,7 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
                 exit_codes.append(self._get_job_return_code(status))
             return exit_codes
 
-        def getJobExitCode(self, batchJobID: str) -> int:
+        def getJobExitCode(self, batchJobID: str) -> Union[int, Tuple[int, Optional[BatchJobExitReason]], None]:
             """
             Get job exit code for given batch job ID.
             :param batchJobID: string of the form "<job>[.<task>]".
@@ -138,18 +141,68 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
                 status_dict = self._getJobDetailsFromScontrol(job_id_list)
             return status_dict
 
-        def _get_job_return_code(self, status: tuple) -> list:
+        def _get_job_return_code(self, status: tuple) -> Union[int, Tuple[int, Optional[BatchJobExitReason]], None]:
             """
+            Given a Slurm return code, status pair, summarize them into a Toil return code, exit reason pair.
+
+            The return code may have already been OR'd with the 128-offset
+            Slurm-reported signal.
+
+            Slurm will report return codes of 0 even if jobs time out instead
+            of succeeding:
+
+                2093597|TIMEOUT|0:0
+                2093597.batch|CANCELLED|0:15
+
+            So we guarantee here that, if the Slurm status string is not a
+            successful one as defined in
+            <https://slurm.schedmd.com/squeue.html#SECTION_JOB-STATE-CODES>, we
+            will not return a successful return code.
+
             Helper function for `getJobExitCode` and `coalesce_job_exit_codes`.
-            :param status: tuple containing the job's state and it's return code.
-            :return: the job's return code if it's completed, otherwise None.
+            :param status: tuple containing the job's state and it's return code from Slurm.
+            :return: the job's return code for Toil if it's completed, otherwise None.
             """
             state, rc = status
-            # If job is in a running state, set return code to None to indicate we don't have
-            # an update.
-            if state in ('PENDING', 'RUNNING', 'CONFIGURING', 'COMPLETING', 'RESIZING', 'SUSPENDED'):
-                rc = None
-            return rc
+
+            # If a job is in one of these states, Slurm can't run it anymore.
+            # We don't include states where the job is held or paused here;
+            # those mean it could run and needs to wait for someone to un-hold
+            # it, so Toil should wait for it.
+            #
+            # We map from each terminal state to the Toil-ontology exit reason.
+            TERMINAL_STATES: Dict[str, BatchJobExitReason] = {
+                "BOOT_FAIL": BatchJobExitReason.LOST,
+                "CANCELLED": BatchJobExitReason.KILLED,
+                "COMPLETED": BatchJobExitReason.FINISHED,
+                "DEADLINE": BatchJobExitReason.KILLED,
+                "FAILED": BatchJobExitReason.FAILED,
+                "NODE_FAIL": BatchJobExitReason.LOST,
+                "OUT_OF_MEMORY": BatchJobExitReason.MEMLIMIT,
+                "PREEMPTED": BatchJobExitReason.KILLED,
+                "TIMEOUT": BatchJobExitReason.KILLED
+            }
+
+            if state not in TERMINAL_STATES:
+                # Don't treat the job as exited yet
+                return None
+
+            exit_reason = TERMINAL_STATES[state]
+
+            if exit_reason == BatchJobExitReason.FINISHED:
+                # The only state that should produce a 0 ever is COMPLETED. So
+                # if the job is COMPLETED and the exit reason is thus FINISHED,
+                # pass along the code it has.
+                return (rc, exit_reason)
+
+            if rc == 0:
+                # The job claims to be in a state other than COMPLETED, but
+                # also to have not encountered a problem. Say the exit status
+                # is unavailable.
+                return (EXIT_STATUS_UNAVAILABLE_VALUE, exit_reason)
+
+            # If the code is nonzero, pass it along.
+            return (rc, exit_reason)
 
         def _getJobDetailsFromSacct(self, job_id_list: list) -> dict:
             """
@@ -283,8 +336,26 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
                           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 = {}
@@ -387,5 +458,5 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
     OptionType = TypeVar('OptionType')
     @classmethod
     def setOptions(cls, setOption: OptionSetter) -> None:
-        setOption("allocate_mem", bool, default=False)
+        setOption("allocate_mem")
 

toil/batchSystems/torque.py

@@ -16,7 +16,6 @@ import math
 import os
 import shlex
 import tempfile
-import time
 from queue import Empty
 from shlex import quote
 from typing import Dict, List, Optional

toil/bus.py

@@ -61,15 +61,13 @@ MessageBus.connect_output_file() and MessageBus.scan_bus_messages().
 """
 
 import collections
-from dataclasses import dataclass
-import inspect
+import json
 import logging
-
 import os
 import queue
-import json
 import tempfile
 import threading
+from dataclasses import dataclass
 from typing import (IO,
                     Any,
                     Callable,
@@ -80,7 +78,6 @@ from typing import (IO,
                     Optional,
                     Type,
                     TypeVar,
-                    Union,
                     cast)
 
 from pubsub.core import Publisher
@@ -90,6 +87,43 @@ from pubsub.core.topicutils import ALL_TOPICS
 
 logger = logging.getLogger( __name__ )
 
+# We define some ways to talk about jobs.
+
+class Names(NamedTuple):
+    """
+    Stores all the kinds of name a job can have.
+    """
+    # Name of the kind of job this is
+    job_name: str
+    # Name of this particular work unit
+    unit_name: str
+    # Human-readable name for the job
+    display_name: str
+    # What the job prints as, used for stats-and-logging log management
+    stats_name: str
+    # Job store ID of the job for the work unit
+    job_store_id: str
+
+def get_job_kind(names: Names) -> str:
+    """
+    Return an identifying string for the job.
+
+    The result may contain spaces.
+
+    Returns: Either the unit name, job name, or display name, which identifies
+             the kind of job it is to toil.
+             Otherwise "Unknown Job" in case no identifier is available
+    """
+    if names.unit_name:
+        return names.unit_name
+    elif names.job_name:
+        return names.job_name
+    elif names.display_name:
+        return names.display_name
+    else:
+        return "Unknown Job"
+
+
 # We define a bunch of named tuple message types.
 # These all need to be plain data: only hold ints, strings, etc.
 
@@ -396,8 +430,8 @@ class MessageBus:
             given topic.
             """
             # There should always be a "message"
-            assert len(message_data) == 1
-            assert 'message' in message_data
+            if len(message_data) != 1 or 'message' not in message_data:
+                raise RuntimeError("Cannot log the bus message. The message is either empty/malformed or there are too many messages provided.")
             message = message_data['message']
             topic = topic_object.getName()
             stream.write(topic.encode('utf-8'))
@@ -572,7 +606,8 @@ class MessageInbox(MessageBusClient):
                 handled = False
                 try:
                     # Emit the message
-                    assert isinstance(message, message_type), f"Unacceptable message type {type(message)} in list for type {message_type}"
+                    if not isinstance(message, message_type):
+                        raise RuntimeError(f"Unacceptable message type {type(message)} in list for type {message_type}")
                     yield message
                     # If we get here it was handled without error.
                     handled = True
@@ -650,6 +685,7 @@ class JobStatus:
 
     def __repr__(self) -> str:
         return json.dumps(self, default= lambda o: o.__dict__, indent=4)
+
 def replay_message_bus(path: str) -> Dict[str, JobStatus]:
     """
     Replay all the messages and work out what they mean for jobs.