toil 7.0.0__py3-none-any.whl → 8.1.0b1__py3-none-any.whl

toil/batchSystems/slurm.py

@@ -11,19 +11,32 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
+import errno
 import logging
 import math
 import os
-from argparse import ArgumentParser, _ArgumentGroup
-from shlex import quote
-from typing import Dict, List, Optional, Set, Tuple, TypeVar, Union
-
-from toil.batchSystems.abstractBatchSystem import BatchJobExitReason, EXIT_STATUS_UNAVAILABLE_VALUE, InsufficientSystemResources
-from toil.batchSystems.abstractGridEngineBatchSystem import \
-    AbstractGridEngineBatchSystem
+import sys
+from argparse import SUPPRESS, ArgumentParser, _ArgumentGroup
+import shlex
+from typing import Callable, NamedTuple, TypeVar
+
+from toil.batchSystems.abstractBatchSystem import (
+    EXIT_STATUS_UNAVAILABLE_VALUE,
+    BatchJobExitReason,
+    InsufficientSystemResources,
+)
+from toil.batchSystems.abstractGridEngineBatchSystem import (
+    AbstractGridEngineBatchSystem,
+)
 from toil.batchSystems.options import OptionSetter
-from toil.job import Requirer
+from toil.bus import get_job_kind
+from toil.common import Config
+from toil.job import JobDescription, Requirer
+from toil.lib.conversions import strtobool
 from toil.lib.misc import CalledProcessErrorStderr, call_command
+from toil.statsAndLogging import TRACE
 
 logger = logging.getLogger(__name__)
 
@@ -36,7 +49,7 @@ logger = logging.getLogger(__name__)
 # it, so Toil should wait for it.
 #
 # We map from each terminal state to the Toil-ontology exit reason.
-TERMINAL_STATES: Dict[str, BatchJobExitReason] = {
+TERMINAL_STATES: dict[str, BatchJobExitReason] = {
     "BOOT_FAIL": BatchJobExitReason.LOST,
     "CANCELLED": BatchJobExitReason.KILLED,
     "COMPLETED": BatchJobExitReason.FINISHED,
@@ -47,12 +60,12 @@ TERMINAL_STATES: Dict[str, BatchJobExitReason] = {
     "PREEMPTED": BatchJobExitReason.KILLED,
     "REVOKED": BatchJobExitReason.KILLED,
     "SPECIAL_EXIT": BatchJobExitReason.FAILED,
-    "TIMEOUT": BatchJobExitReason.KILLED
+    "TIMEOUT": BatchJobExitReason.KILLED,
 }
 
 # If a job is in one of these states, it might eventually move to a different
 # state.
-NONTERMINAL_STATES: Set[str] = {
+NONTERMINAL_STATES: set[str] = {
     "CONFIGURING",
     "COMPLETING",
     "PENDING",
@@ -65,51 +78,215 @@ NONTERMINAL_STATES: Set[str] = {
     "SIGNALING",
     "STAGE_OUT",
     "STOPPED",
-    "SUSPENDED"
-} 
+    "SUSPENDED",
+}
+
+
+def parse_slurm_time(slurm_time: str) -> int:
+    """
+    Parse a Slurm-style time duration like 7-00:00:00 to a number of seconds.
+
+    Raises ValueError if not parseable.
+    """
+    # slurm returns time in days-hours:minutes:seconds format
+    # Sometimes it will only return minutes:seconds, so days may be omitted
+    # For ease of calculating, we'll make sure all the delimeters are ':'
+    # Then reverse the list so that we're always counting up from seconds -> minutes -> hours -> days
+    total_seconds = 0
+    elapsed_split: list[str] = slurm_time.replace("-", ":").split(":")
+    elapsed_split.reverse()
+    seconds_per_unit = [1, 60, 3600, 86400]
+    for index, multiplier in enumerate(seconds_per_unit):
+        if index < len(elapsed_split):
+            total_seconds += multiplier * int(elapsed_split[index])
+    return total_seconds
+
 
 class SlurmBatchSystem(AbstractGridEngineBatchSystem):
+    class PartitionInfo(NamedTuple):
+        partition_name: str
+        gres: bool
+        time_limit: float
+        priority: int
+        cpus: str
+        memory: str
+
+    class PartitionSet:
+        """
+        Set of available partitions detected on the slurm batch system
+        """
+
+        default_gpu_partition: SlurmBatchSystem.PartitionInfo | None
+        all_partitions: list[SlurmBatchSystem.PartitionInfo]
+        gpu_partitions: set[str]
+
+        def __init__(self) -> None:
+            self._get_partition_info()
+            self._get_gpu_partitions()
+
+        def _get_gpu_partitions(self) -> None:
+            """
+            Get all available GPU partitions. Also get the default GPU partition.
+            :return: None
+            """
+            gpu_partitions = [
+                partition for partition in self.all_partitions if partition.gres
+            ]
+            self.gpu_partitions = {p.partition_name for p in gpu_partitions}
+            # Grab the lowest priority GPU partition
+            # If no GPU partitions are available, then set the default to None
+            self.default_gpu_partition = None
+            if len(gpu_partitions) > 0:
+                self.default_gpu_partition = sorted(
+                    gpu_partitions, key=lambda x: x.priority
+                )[0]
+
+        def _get_partition_info(self) -> None:
+            """
+            Call the Slurm batch system with sinfo to grab all available partitions.
+            Then parse the output and store all available Slurm partitions
+            :return: None
+            """
+            sinfo_command = ["sinfo", "-a", "-o", "%P %G %l %p %c %m"]
+
+            sinfo = call_command(sinfo_command)
+
+            parsed_partitions = []
+            for line in sinfo.split("\n")[1:]:
+                if line.strip():
+                    partition_name, gres, time, priority, cpus, memory = line.split(" ")
+                    try:
+                        # Parse time to a number so we can compute on it
+                        partition_time: float = parse_slurm_time(time)
+                    except ValueError:
+                        # Maybe time is unlimited?
+                        partition_time = float("inf")
+                    try:
+                        # Parse priority to an int so we can sort on it
+                        partition_priority = int(priority)
+                    except ValueError:
+                        logger.warning(
+                            "Could not parse priority %s for partition %s, assuming high priority",
+                            partition_name,
+                            priority,
+                        )
+                        partition_priority = sys.maxsize
+                    parsed_partitions.append(
+                        SlurmBatchSystem.PartitionInfo(
+                            partition_name.rstrip("*"),
+                            gres != "(null)",
+                            partition_time,
+                            partition_priority,
+                            cpus,
+                            memory,
+                        )
+                    )
+            self.all_partitions = parsed_partitions
+
+        def get_partition(self, time_limit: float | None) -> str | None:
+            """
+            Get the partition name to use for a job with the given time limit.
+
+            :param time_limit: Time limit in seconds.
+            """
+
+            if time_limit is None:
+                # Just use Slurm's default
+                return None
+
+            winning_partition = None
+            for partition in self.all_partitions:
+                if partition.time_limit < time_limit:
+                    # Can't use this
+                    continue
+                if winning_partition is None:
+                    # Anything beats None
+                    winning_partition = partition
+                    continue
+                if partition.gres and not winning_partition.gres:
+                    # Never use a partition witn GRES if you can avoid it
+                    continue
+                elif not partition.gres and winning_partition.gres:
+                    # Never keep a partition with GRES if we find one without
+                    winning_partition = partition
+                    continue
+                if partition.priority > winning_partition.priority:
+                    # After that, don't raise priority
+                    continue
+                elif partition.priority < winning_partition.priority:
+                    # And always lower it
+                    winning_partition = partition
+                    continue
+                if partition.time_limit < winning_partition.time_limit:
+                    # Finally, lower time limit
+                    winning_partition = partition
+
+            # TODO: Store partitions in a better indexed way
+            if winning_partition is None and len(self.all_partitions) > 0:
+                # We have partitions and none of them can fit this
+                raise RuntimeError(
+                    f"Could not find a Slurm partition that can fit a job that runs for {time_limit} seconds"
+                )
+
+            if winning_partition is None:
+                return None
+            else:
+                return winning_partition.partition_name
 
     class GridEngineThread(AbstractGridEngineBatchSystem.GridEngineThread):
+        # Our boss is always the enclosing class
+        boss: SlurmBatchSystem
 
-        def getRunningJobIDs(self):
+        def getRunningJobIDs(self) -> dict[int, int]:
             # Should return a dictionary of Job IDs and number of seconds
             times = {}
             with self.runningJobsLock:
-                currentjobs = {str(self.batchJobIDs[x][0]): x for x in self.runningJobs}
+                currentjobs: dict[str, int] = {
+                    str(self.batchJobIDs[x][0]): x for x in self.runningJobs
+                }
             # currentjobs is a dictionary that maps a slurm job id (string) to our own internal job id
             # squeue arguments:
             # -h for no header
             # --format to get jobid i, state %t and time days-hours:minutes:seconds
 
-            lines = call_command(['squeue', '-h', '--format', '%i %t %M'], quiet=True).split('\n')
+            lines = call_command(
+                ["squeue", "-h", "--format", "%i %t %M"], quiet=True
+            ).split("\n")
             for line in lines:
                 values = line.split()
                 if len(values) < 3:
                     continue
                 slurm_jobid, state, elapsed_time = values
-                if slurm_jobid in currentjobs and state == 'R':
-                    seconds_running = self.parse_elapsed(elapsed_time)
+                if slurm_jobid in currentjobs and state == "R":
+                    try:
+                        seconds_running = parse_slurm_time(elapsed_time)
+                    except ValueError:
+                        # slurm may return INVALID instead of a time
+                        seconds_running = 0
                     times[currentjobs[slurm_jobid]] = seconds_running
 
             return times
 
-        def killJob(self, jobID):
-            call_command(['scancel', self.getBatchSystemID(jobID)])
-
-        def prepareSubmission(self,
-                              cpu: int,
-                              memory: int,
-                              jobID: int,
-                              command: str,
-                              jobName: str,
-                              job_environment: Optional[Dict[str, str]] = None,
-                              gpus: Optional[int] = None) -> List[str]:
+        def killJob(self, jobID: int) -> None:
+            call_command(["scancel", self.getBatchSystemID(jobID)])
+
+        def prepareSubmission(
+            self,
+            cpu: int,
+            memory: int,
+            jobID: int,
+            command: str,
+            jobName: str,
+            job_environment: dict[str, str] | None = None,
+            gpus: int | None = None,
+        ) -> list[str]:
             # 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}']
+            return self.prepareSbatch(
+                cpu, memory, jobID, jobName, job_environment, gpus
+            ) + [f"--wrap=exec {command}"]
 
-        def submitJob(self, subLine):
+        def submitJob(self, subLine: list[str]) -> int:
             try:
                 # Slurm is not quite clever enough to follow the XDG spec on
                 # its own. If the submission command sees e.g. XDG_RUNTIME_DIR
@@ -125,7 +302,11 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
                 # This doesn't get us a trustworthy XDG session in Slurm, but
                 # it does let us see the one Slurm tries to give us.
                 no_session_environment = os.environ.copy()
-                session_names = [n for n in no_session_environment.keys() if n.startswith('XDG_') or n.startswith('DBUS_')]
+                session_names = [
+                    n
+                    for n in no_session_environment.keys()
+                    if n.startswith("XDG_") or n.startswith("DBUS_")
+                ]
                 for name in session_names:
                     del no_session_environment[name]
 
@@ -138,36 +319,44 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
                 logger.error(f"sbatch command failed with error: {e}")
                 raise e
 
-        def coalesce_job_exit_codes(self, batch_job_id_list: list) -> List[Union[int, Tuple[int, Optional[BatchJobExitReason]], None]]:
+        def coalesce_job_exit_codes(
+            self, batch_job_id_list: list[str]
+        ) -> list[int | tuple[int, BatchJobExitReason | None] | 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 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)
+            logger.log(
+                TRACE, "Getting exit codes for slurm jobs: %s", batch_job_id_list
+            )
             # Convert batch_job_id_list to list of integer job IDs.
-            job_id_list = [int(id.split('.')[0]) for id in batch_job_id_list]
+            job_id_list = [int(id.split(".")[0]) for id in batch_job_id_list]
             status_dict = self._get_job_details(job_id_list)
-            exit_codes = []
+            exit_codes: list[int | tuple[int, BatchJobExitReason | None] | None] = []
             for _, status in status_dict.items():
                 exit_codes.append(self._get_job_return_code(status))
             return exit_codes
 
-        def getJobExitCode(self, batchJobID: str) -> Union[int, Tuple[int, Optional[BatchJobExitReason]], None]:
+        def getJobExitCode(
+            self, batchJobID: str
+        ) -> int | tuple[int, BatchJobExitReason | None] | None:
             """
             Get job exit code for given batch job ID.
             :param batchJobID: string of the form "<job>[.<task>]".
             :return: integer job exit code.
             """
-            logger.debug("Getting exit code for slurm job: %s", batchJobID)
+            logger.log(TRACE, "Getting exit code for slurm job: %s", batchJobID)
             # Convert batchJobID to an integer job ID.
-            job_id = int(batchJobID.split('.')[0])
+            job_id = int(batchJobID.split(".")[0])
             status_dict = self._get_job_details([job_id])
             status = status_dict[job_id]
             return self._get_job_return_code(status)
 
-        def _get_job_details(self, job_id_list: list) -> dict:
+        def _get_job_details(
+            self, job_id_list: list[int]
+        ) -> dict[int, tuple[str | None, int | None]]:
             """
             Helper function for `getJobExitCode` and `coalesce_job_exit_codes`.
             Fetch job details from Slurm's accounting system or job control system.
@@ -177,11 +366,15 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
             """
             try:
                 status_dict = self._getJobDetailsFromSacct(job_id_list)
-            except CalledProcessErrorStderr:
+            except (CalledProcessErrorStderr, OSError) as e:
+                if isinstance(e, OSError):
+                    logger.warning("Could not run sacct: %s", e)
                 status_dict = self._getJobDetailsFromScontrol(job_id_list)
             return status_dict
 
-        def _get_job_return_code(self, status: tuple) -> Union[int, Tuple[int, Optional[BatchJobExitReason]], None]:
+        def _get_job_return_code(
+            self, status: tuple[str | None, int | None]
+        ) -> int | tuple[int, BatchJobExitReason | None] | None:
             """
             Given a Slurm return code, status pair, summarize them into a Toil return code, exit reason pair.
 
@@ -215,7 +408,7 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
                 # 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)
+                return (rc, exit_reason)  # type: ignore[return-value] # mypy doesn't understand enums well
 
             if rc == 0:
                 # The job claims to be in a state other than COMPLETED, but
@@ -224,7 +417,7 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
                 return (EXIT_STATUS_UNAVAILABLE_VALUE, exit_reason)
 
             # If the code is nonzero, pass it along.
-            return (rc, exit_reason)
+            return (rc, exit_reason)  # type: ignore[return-value] # mypy doesn't understand enums well
 
         def _canonicalize_state(self, state: str) -> str:
             """
@@ -233,18 +426,23 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
 
             # Slurm will sometimes send something like "CANCELED by 30065" in
             # the state column for some reason.
-            
+
             state_token = state
 
             if " " in state_token:
                 state_token = state.split(" ", 1)[0]
 
-            if state_token not in TERMINAL_STATES and state_token not in NONTERMINAL_STATES:
+            if (
+                state_token not in TERMINAL_STATES
+                and state_token not in NONTERMINAL_STATES
+            ):
                 raise RuntimeError("Toil job in unimplemented Slurm state " + state)
-            
+
             return state_token
 
-        def _getJobDetailsFromSacct(self, job_id_list: list) -> dict:
+        def _getJobDetailsFromSacct(
+            self, job_id_list: list[int]
+        ) -> dict[int, tuple[str | None, int | None]]:
             """
             Get SLURM job exit codes for the jobs in `job_id_list` by running `sacct`.
             :param job_id_list: list of integer batch job IDs.
@@ -252,52 +450,82 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
             containing the job's state and exit code.
             """
             job_ids = ",".join(str(id) for id in job_id_list)
-            args = ['sacct',
-                    '-n',  # no header
-                    '-j', job_ids,  # job
-                    '--format', 'JobIDRaw,State,ExitCode',  # specify output columns
-                    '-P',  # separate columns with pipes
-                    '-S', '1970-01-01']  # override start time limit
-            stdout = call_command(args, quiet=True)
+            args = [
+                "sacct",
+                "-n",  # no header
+                "-j",
+                job_ids,  # job
+                "--format",
+                "JobIDRaw,State,ExitCode",  # specify output columns
+                "-P",  # separate columns with pipes
+                "-S",
+                "1970-01-01",
+            ]  # override start time limit
 
             # Collect the job statuses in a dict; key is the job-id, value is a tuple containing
             # job state and exit status. Initialize dict before processing output of `sacct`.
-            job_statuses = {}
+            job_statuses: dict[int, tuple[str | None, int | None]] = {}
+
+            try:
+                stdout = call_command(args, quiet=True)
+            except OSError as e:
+                if e.errno == errno.E2BIG:
+                    # Argument list is too big, recurse on half the argument list
+                    if len(job_id_list) == 1:
+                        # 1 is too big, we can't recurse further, bail out
+                        raise
+                    job_statuses.update(self._getJobDetailsFromSacct(job_id_list[:len(job_id_list)//2]))
+                    job_statuses.update(self._getJobDetailsFromSacct(job_id_list[len(job_id_list)//2:]))
+                    return job_statuses
+                else:
+                    raise
+
             for job_id in job_id_list:
                 job_statuses[job_id] = (None, None)
 
             for line in stdout.splitlines():
-                values = line.strip().split('|')
+                values = line.strip().split("|")
                 if len(values) < 3:
                     continue
+                state: str
                 job_id_raw, state, exitcode = values
                 state = self._canonicalize_state(state)
-                logger.debug("%s state of job %s is %s", args[0], job_id_raw, state)
+                logger.log(
+                    TRACE, "%s state of job %s is %s", args[0], job_id_raw, state
+                )
                 # JobIDRaw is in the form JobID[.JobStep]; we're not interested in job steps.
                 job_id_parts = job_id_raw.split(".")
                 if len(job_id_parts) > 1:
                     continue
                 job_id = int(job_id_parts[0])
-                status, signal = (int(n) for n in exitcode.split(':'))
+                status: int
+                signal: int
+                status, signal = (int(n) for n in exitcode.split(":"))
                 if signal > 0:
                     # A non-zero signal may indicate e.g. an out-of-memory killed job
                     status = 128 + signal
-                logger.debug("%s exit code of job %d is %s, return status %d",
-                             args[0], job_id, exitcode, status)
+                logger.log(
+                    TRACE,
+                    "%s exit code of job %d is %s, return status %d",
+                    args[0],
+                    job_id,
+                    exitcode,
+                    status,
+                )
                 job_statuses[job_id] = state, status
-            logger.debug("%s returning job statuses: %s", args[0], job_statuses)
+            logger.log(TRACE, "%s returning job statuses: %s", args[0], job_statuses)
             return job_statuses
 
-        def _getJobDetailsFromScontrol(self, job_id_list: list) -> dict:
+        def _getJobDetailsFromScontrol(
+            self, job_id_list: list[int]
+        ) -> dict[int, tuple[str | None, int | None]]:
             """
             Get SLURM job exit codes for the jobs in `job_id_list` by running `scontrol`.
             :param job_id_list: list of integer batch job IDs.
             :return: dict of job statuses, where key is the job-id, and value is a tuple
             containing the job's state and exit code.
             """
-            args = ['scontrol',
-                    'show',
-                    'job']
+            args = ["scontrol", "show", "job"]
             # `scontrol` can only return information about a single job,
             # or all the jobs it knows about.
             if len(job_id_list) == 1:
@@ -306,14 +534,16 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
             stdout = call_command(args, quiet=True)
 
             # Job records are separated by a blank line.
+            job_records = None
             if isinstance(stdout, str):
-                job_records = stdout.strip().split('\n\n')
+                job_records = stdout.strip().split("\n\n")
             elif isinstance(stdout, bytes):
-                job_records = stdout.decode('utf-8').strip().split('\n\n')
+                job_records = stdout.decode("utf-8").strip().split("\n\n")
 
             # Collect the job statuses in a dict; key is the job-id, value is a tuple containing
             # job state and exit status. Initialize dict before processing output of `scontrol`.
-            job_statuses = {}
+            job_statuses: dict[int, tuple[str | None, int | None]] = {}
+            job_id: int | None
             for job_id in job_id_list:
                 job_statuses[job_id] = (None, None)
 
@@ -323,7 +553,8 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
                 return job_statuses
 
             for record in job_records:
-                job = {}
+                job: dict[str, str] = {}
+                job_id = None
                 for line in record.splitlines():
                     for item in line.split():
                         # Output is in the form of many key=value pairs, multiple pairs on each line
@@ -331,59 +562,68 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
                         # added to a dictionary.
                         # Note: In some cases, the value itself may contain white-space. So, if we find
                         # a key without a value, we consider that key part of the previous value.
-                        bits = item.split('=', 1)
+                        bits = item.split("=", 1)
                         if len(bits) == 1:
-                            job[key] += ' ' + bits[0]
+                            job[key] += " " + bits[0]  # type: ignore[has-type]  # we depend on the previous iteration to populate key
                         else:
                             key = bits[0]
                             job[key] = bits[1]
                     # The first line of the record contains the JobId. Stop processing the remainder
                     # of this record, if we're not interested in this job.
-                    job_id = int(job['JobId'])
+                    job_id = int(job["JobId"])
                     if job_id not in job_id_list:
-                        logger.debug("%s job %d is not in the list", args[0], job_id)
+                        logger.log(
+                            TRACE, "%s job %d is not in the list", args[0], job_id
+                        )
                         break
-                if job_id not in job_id_list:
+                if job_id is None or job_id not in job_id_list:
                     continue
-                state = job['JobState']
+                state = job["JobState"]
                 state = self._canonicalize_state(state)
-                logger.debug("%s state of job %s is %s", args[0], job_id, state)
+                logger.log(TRACE, "%s state of job %s is %s", args[0], job_id, state)
                 try:
-                    exitcode = job['ExitCode']
+                    exitcode = job["ExitCode"]
                     if exitcode is not None:
-                        status, signal = (int(n) for n in exitcode.split(':'))
+                        status, signal = (int(n) for n in exitcode.split(":"))
                         if signal > 0:
                             # A non-zero signal may indicate e.g. an out-of-memory killed job
                             status = 128 + signal
-                        logger.debug("%s exit code of job %d is %s, return status %d",
-                                     args[0], job_id, exitcode, status)
+                        logger.log(
+                            TRACE,
+                            "%s exit code of job %d is %s, return status %d",
+                            args[0],
+                            job_id,
+                            exitcode,
+                            status,
+                        )
                         rc = status
                     else:
                         rc = None
                 except KeyError:
                     rc = None
                 job_statuses[job_id] = (state, rc)
-            logger.debug("%s returning job statuses: %s", args[0], job_statuses)
+            logger.log(TRACE, "%s returning job statuses: %s", args[0], job_statuses)
             return job_statuses
 
         ###
         ### Implementation-specific helper methods
         ###
 
-        def prepareSbatch(self,
-                          cpu: int,
-                          mem: int,
-                          jobID: int,
-                          jobName: str,
-                          job_environment: Optional[Dict[str, str]],
-                          gpus: Optional[int]) -> List[str]:
-
+        def prepareSbatch(
+            self,
+            cpu: int,
+            mem: int,
+            jobID: int,
+            jobName: str,
+            job_environment: dict[str, str] | None,
+            gpus: int | None,
+        ) -> list[str]:
             """
             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}']
+            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
@@ -398,83 +638,271 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
             # 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 = {}
             environment.update(self.boss.environment)
             if job_environment:
                 environment.update(job_environment)
 
             # "Native extensions" for SLURM (see DRMAA or SAGA)
-            nativeConfig = os.getenv('TOIL_SLURM_ARGS')
+            # Also any extra arguments from --slurmArgs or TOIL_SLURM_ARGS
+            nativeConfig: str = self.boss.config.slurm_args  # type: ignore[attr-defined]
+
+            # For parsing user-provided option overrides (or self-generated
+            # options) we need a way to recognize long, long-with-equals, and
+            # short forms.
+            def option_detector(long: str, short: str | None = None) -> Callable[[str], bool]:
+                """
+                Get a function that returns true if it sees the long or short
+                option.
+                """
+                def is_match(option: str) -> bool:
+                    return option == f"--{long}" or option.startswith(f"--{long}=") or (short is not None and option == f"-{short}")
+                return is_match
+
+            def any_option_detector(options: list[str | tuple[str, str]]) -> Callable[[str], bool]:
+                """
+                Get a function that returns true if it sees any of the long
+                options or long or short option pairs.
+                """
+                detectors = [option_detector(o) if isinstance(o, str) else option_detector(*o) for o in options]
+                def is_match(option: str) -> bool:
+                    for detector in detectors:
+                        if detector(option):
+                            return True
+                    return False
+                return is_match
+
+            is_any_mem_option = any_option_detector(["mem", "mem-per-cpu", "mem-per-gpu"])
+            is_any_cpus_option = any_option_detector([("cpus-per-task", "c"), "cpus-per-gpu"])
+            is_export_option = option_detector("export")
+            is_export_file_option = option_detector("export-file")
+            is_time_option = option_detector("time", "t")
+            is_partition_option = option_detector("partition", "p")
+
+            # We will fill these in with stuff parsed from TOIL_SLURM_ARGS, or
+            # with our own determinations if they aren't there.
 
             # --export=[ALL,]<environment_toil_variables>
-            set_exports = "--export=ALL"
+            export_all = True
+            export_list = [] # Some items here may be multiple comma-separated values
+            time_limit: int | None = self.boss.config.slurm_time  # type: ignore[attr-defined]
+            partition: str | None = None
 
             if nativeConfig is not None:
-                logger.debug("Native SLURM options appended to sbatch from TOIL_SLURM_ARGS env. variable: %s", nativeConfig)
+                logger.debug(
+                    "Native SLURM options appended to sbatch: %s", nativeConfig
+                )
+
+                # Do a mini argument parse to pull out export and parse time if
+                # needed
+                args = shlex.split(nativeConfig)
+                i = 0
+                while i < len(args):
+                    arg = args[i]
+                    if is_any_mem_option(arg) or is_any_cpus_option(arg):
+                        # Prohibit arguments that set CPUs or memory
+                        raise ValueError(
+                            f"Cannot use Slurm argument {arg} which conflicts "
+                            f"with Toil's own arguments to Slurm"
+                        )
+                    elif is_export_option(arg):
+                        # Capture the export argument value so we can modify it
+                        export_all = False
+                        if "=" not in arg:
+                            if i + 1 >= len(args):
+                                raise ValueError(
+                                    f"No value supplied for Slurm {arg} argument"
+                                )
+                            i += 1
+                            export_list.append(args[i])
+                        else:
+                            export_list.append(arg.split("=", 1)[1])
+                    elif is_export_file_option(arg):
+                        # Keep --export-file but turn off --export=ALL in that
+                        # case.
+                        export_all = False
+                        sbatch_line.append(arg)
+                    elif is_time_option(arg):
+                        # Capture the time limit in seconds so we can use it for picking a partition
+                        if "=" not in arg:
+                            if i + 1 >= len(args):
+                                raise ValueError(
+                                    f"No value supplied for Slurm {arg} argument"
+                                )
+                            i += 1
+                            time_string = args[i]
+                        else:
+                            time_string = arg.split("=", 1)[1]
+                        time_limit = parse_slurm_time(time_string)
+                    elif is_partition_option(arg):
+                        # Capture the partition so we can run checks on it and know not to assign one
+                        if "=" not in arg:
+                            if i + 1 >= len(args):
+                                raise ValueError(
+                                    f"No value supplied for Slurm {arg} argument"
+                                )
+                            i += 1
+                            partition = args[i]
+                        else:
+                            partition = arg.split("=", 1)[1]
+                    else:
+                        # Other arguments pass through.
+                        sbatch_line.append(arg)
+                    i += 1
 
-                for arg in nativeConfig.split():
-                    if arg.startswith("--mem") or arg.startswith("--cpus-per-task"):
-                        raise ValueError(f"Some resource arguments are incompatible: {nativeConfig}")
-                    # repleace default behaviour by the one stated at TOIL_SLURM_ARGS
-                    if arg.startswith("--export"):
-                        set_exports = arg
-                sbatch_line.extend(nativeConfig.split())
+            if export_all:
+                # We don't have any export overrides so we ened to start with
+                # an ALL
+                export_list.append("ALL")
 
             if environment:
                 argList = []
 
                 for k, v in environment.items():
-                    quoted_value = quote(os.environ[k] if v is None else v)
-                    argList.append(f'{k}={quoted_value}')
-
-                set_exports += ',' + ','.join(argList)
-
-            # add --export to the sbatch
-            sbatch_line.append(set_exports)
-
-            parallel_env = os.getenv('TOIL_SLURM_PE')
-            if cpu and cpu > 1 and parallel_env:
-                sbatch_line.append(f'--partition={parallel_env}')
-
-            if mem is not None and self.boss.config.allocate_mem:
+                    # TODO: The sbatch man page doesn't say we can quote these;
+                    # if we need to send characters like , itself we need to
+                    # use --export-file and clean it up when the command has
+                    # been issued.
+                    quoted_value = shlex.quote(os.environ[k] if v is None else v)
+                    argList.append(f"{k}={quoted_value}")
+
+                export_list.extend(argList)
+
+            # If partition isn't set and we have a GPU partition override
+            # that applies, apply it
+            gpu_partition_override: str | None = self.boss.config.slurm_gpu_partition  # type: ignore[attr-defined]
+            if partition is None and gpus and gpu_partition_override:
+                partition = gpu_partition_override
+
+            # If partition isn't set and we have a parallel partition override
+            # that applies, apply it
+            parallel_env: str | None = self.boss.config.slurm_pe  # type: ignore[attr-defined]
+            if partition is None and cpu and cpu > 1 and parallel_env:
+                partition = parallel_env
+
+            # If partition isn't set and we have a general partition override
+            # that applies, apply it
+            partition_override: str | None = self.boss.config.slurm_partition  # type: ignore[attr-defined]
+            if partition is None and partition_override:
+                partition = partition_override
+
+            if partition is None and gpus:
+                # Send to a GPU partition
+                gpu_partition = self.boss.partitions.default_gpu_partition
+                if gpu_partition is None:
+                    # no gpu partitions are available, raise an error
+                    raise RuntimeError(
+                        f"The job {jobName} is requesting GPUs, but the Slurm cluster does not appear to have an accessible partition with GPUs"
+                    )
+                if (
+                    time_limit is not None
+                    and gpu_partition.time_limit < time_limit
+                ):
+                    # TODO: find the lowest-priority GPU partition that has at least each job's time limit!
+                    logger.warning(
+                        "Trying to submit a job that needs %s seconds to partition %s that has a limit of %s seconds",
+                        time_limit,
+                        gpu_partition.partition_name,
+                        gpu_partition.time_limit,
+                    )
+                partition = gpu_partition.partition_name
+
+            if partition is None:
+                # Pick a partition based on time limit
+                partition = self.boss.partitions.get_partition(time_limit)
+
+            # Now generate all the arguments
+            if len(export_list) > 0:
+                # add --export to the sbatch
+                sbatch_line.append("--export=" + ",".join(export_list))
+            if partition is not None:
+                sbatch_line.append(f"--partition={partition}")
+            if gpus:
+                # Generate GPU assignment argument
+                sbatch_line.append(f"--gres=gpu:{gpus}")
+                if partition is not None and partition not in self.boss.partitions.gpu_partitions:
+                    # the specified partition is not compatible, so warn the user that the job may not work
+                    logger.warning(
+                        f"Job {jobName} needs GPUs, but specified partition {partition} does not have them. This job may not work."
+                        f"Try specifying one of these partitions instead: {', '.join(self.boss.partitions.gpu_partitions)}."
+                    )
+            if mem is not None and self.boss.config.slurm_allocate_mem:  # type: ignore[attr-defined]
                 # memory passed in is in bytes, but slurm expects megabytes
-                sbatch_line.append(f'--mem={math.ceil(mem / 2 ** 20)}')
+                sbatch_line.append(f"--mem={math.ceil(mem / 2 ** 20)}")
             if cpu is not None:
-                sbatch_line.append(f'--cpus-per-task={math.ceil(cpu)}')
+                sbatch_line.append(f"--cpus-per-task={math.ceil(cpu)}")
+            if time_limit is not None:
+                # Put all the seconds in the seconds slot
+                sbatch_line.append(f"--time=0:{time_limit}")
 
-            stdoutfile: str = self.boss.format_std_out_err_path(jobID, '%j', 'out')
-            stderrfile: str = self.boss.format_std_out_err_path(jobID, '%j', 'err')
-            sbatch_line.extend(['-o', stdoutfile, '-e', stderrfile])
+            stdoutfile: str = self.boss.format_std_out_err_path(jobID, "%j", "out")
+            stderrfile: str = self.boss.format_std_out_err_path(jobID, "%j", "err")
+            sbatch_line.extend(["-o", stdoutfile, "-e", stderrfile])
 
             return sbatch_line
 
-        def parse_elapsed(self, elapsed):
-            # slurm returns elapsed time in days-hours:minutes:seconds format
-            # Sometimes it will only return minutes:seconds, so days may be omitted
-            # For ease of calculating, we'll make sure all the delimeters are ':'
-            # Then reverse the list so that we're always counting up from seconds -> minutes -> hours -> days
-            total_seconds = 0
-            try:
-                elapsed = elapsed.replace('-', ':').split(':')
-                elapsed.reverse()
-                seconds_per_unit = [1, 60, 3600, 86400]
-                for index, multiplier in enumerate(seconds_per_unit):
-                    if index < len(elapsed):
-                        total_seconds += multiplier * int(elapsed[index])
-            except ValueError:
-                pass  # slurm may return INVALID instead of a time
-            return total_seconds
+    def __init__(
+        self, config: Config, maxCores: float, maxMemory: int, maxDisk: int
+    ) -> None:
+        super().__init__(config, maxCores, maxMemory, maxDisk)
+        self.partitions = SlurmBatchSystem.PartitionSet()
+
+    # Override issuing jobs so we can check if we need to use Slurm's magic
+    # whole-node-memory feature.
+    def issueBatchJob(
+        self,
+        command: str,
+        job_desc: JobDescription,
+        job_environment: dict[str, str] | None = None,
+    ) -> int:
+        # Avoid submitting internal jobs to the batch queue, handle locally
+        local_id = self.handleLocalJob(command, job_desc)
+        if local_id is not None:
+            return local_id
+        else:
+            self.check_resource_request(job_desc)
+            gpus = self.count_needed_gpus(job_desc)
+            job_id = self.getNextJobID()
+            self.currentJobs.add(job_id)
+
+            if "memory" not in job_desc.requirements and self.config.slurm_default_all_mem:  # type: ignore[attr-defined]
+                # The job doesn't have its own memory requirement, and we are
+                # defaulting to whole node memory. Use Slurm's 0-memory sentinel.
+                memory = 0
+            else:
+                # Use the memory actually on the job, or the Toil default memory
+                memory = job_desc.memory
+
+            self.newJobsQueue.put(
+                (
+                    job_id,
+                    job_desc.cores,
+                    memory,
+                    command,
+                    get_job_kind(job_desc.get_names()),
+                    job_environment,
+                    gpus,
+                )
+            )
+            logger.debug(
+                "Issued the job command: %s with job id: %s and job name %s",
+                command,
+                str(job_id),
+                get_job_kind(job_desc.get_names()),
+            )
+        return job_id
 
     def _check_accelerator_request(self, requirer: Requirer) -> None:
         for accelerator in requirer.accelerators:
-            if accelerator['kind'] != 'gpu':
-                raise InsufficientSystemResources(requirer, 'accelerators', details=
-                                                  [
-                                                      f'The accelerator {accelerator} could not be provided'
-                                                      'The Toil Slurm batch system only supports gpu accelerators at the moment.'
-                                                  ])
+            if accelerator["kind"] != "gpu":
+                raise InsufficientSystemResources(
+                    requirer,
+                    "accelerators",
+                    details=[
+                        f"The accelerator {accelerator} could not be provided"
+                        "The Toil Slurm batch system only supports gpu accelerators at the moment."
+                    ],
+                )
 
     ###
     ### The interface for SLURM
@@ -488,17 +916,85 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
     # implement getWaitDuration().
 
     @classmethod
-    def add_options(cls, parser: Union[ArgumentParser, _ArgumentGroup]):
-        allocate_mem = parser.add_mutually_exclusive_group()
-        allocate_mem_help = ("A flag that can block allocating memory with '--mem' for job submissions "
-                             "on SLURM since some system servers may reject any job request that "
-                             "explicitly specifies the memory allocation.  The default is to always allocate memory.")
-        allocate_mem.add_argument("--dont_allocate_mem", action='store_false', dest="allocate_mem", help=allocate_mem_help)
-        allocate_mem.add_argument("--allocate_mem", action='store_true', dest="allocate_mem", help=allocate_mem_help)
-        allocate_mem.set_defaults(allocate_mem=True)
-
-    OptionType = TypeVar('OptionType')
+    def add_options(cls, parser: ArgumentParser | _ArgumentGroup) -> None:
+
+        parser.add_argument(
+            "--slurmAllocateMem",
+            dest="slurm_allocate_mem",
+            type=strtobool,
+            default=True,
+            env_var="TOIL_SLURM_ALLOCATE_MEM",
+            help="If False, do not use --mem. Used as a workaround for Slurm clusters that reject jobs "
+            "with memory allocations.",
+        )
+        # Keep these deprcated options for backward compatibility
+        parser.add_argument(
+            "--dont_allocate_mem",
+            action="store_false",
+            dest="slurm_allocate_mem",
+            help=SUPPRESS,
+        )
+        parser.add_argument(
+            "--allocate_mem",
+            action="store_true",
+            dest="slurm_allocate_mem",
+            help=SUPPRESS,
+        )
+
+        parser.add_argument(
+            "--slurmDefaultAllMem",
+            dest="slurm_default_all_mem",
+            type=strtobool,
+            default=False,
+            env_var="TOIL_SLURM_DEFAULT_ALL_MEM",
+            help="If True, assign Toil jobs without their own memory requirements all available "
+            "memory on a Slurm node (via Slurm --mem=0).",
+        )
+        parser.add_argument(
+            "--slurmTime",
+            dest="slurm_time",
+            type=parse_slurm_time,
+            default=None,
+            env_var="TOIL_SLURM_TIME",
+            help="Slurm job time limit, in [DD-]HH:MM:SS format.",
+        )
+        parser.add_argument(
+            "--slurmPartition",
+            dest="slurm_partition",
+            default=None,
+            env_var="TOIL_SLURM_PARTITION",
+            help="Partition to send Slurm jobs to.",
+        )
+        parser.add_argument(
+            "--slurmGPUPartition",
+            dest="slurm_gpu_partition",
+            default=None,
+            env_var="TOIL_SLURM_GPU_PARTITION",
+            help="Partition to send Slurm jobs to if they ask for GPUs.",
+        )
+        parser.add_argument(
+            "--slurmPE",
+            dest="slurm_pe",
+            default=None,
+            env_var="TOIL_SLURM_PE",
+            help="Special partition to send Slurm jobs to if they ask for more than 1 CPU.",
+        )
+        parser.add_argument(
+            "--slurmArgs",
+            dest="slurm_args",
+            default="",
+            env_var="TOIL_SLURM_ARGS",
+            help="Extra arguments to pass to Slurm.",
+        )
+
+    OptionType = TypeVar("OptionType")
+
     @classmethod
     def setOptions(cls, setOption: OptionSetter) -> None:
-        setOption("allocate_mem")
-
+        setOption("slurm_allocate_mem")
+        setOption("slurm_default_all_mem")
+        setOption("slurm_time")
+        setOption("slurm_partition")
+        setOption("slurm_gpu_partition")
+        setOption("slurm_pe")
+        setOption("slurm_args")