Repository: DataBiosphere/toil

Issue:
Slurm runs can get so big that we can't list all the job IDs on the `sacct` command line and crash
This was reported in https://github.com/ComparativeGenomicsToolkit/cactus/issues/1465

We attempt to put too many job IDs into an `sacct` command, and we get an argument list too long error. That's not among the ones that prompt us to fall back to `scontrol` here:

https://github.com/DataBiosphere/toil/blob/5a8090e86e7dca2be45728da925bf68f0658ba1c/src/toil/batchSystems/slurm.py#L301

So the net result is that Toil crashes.

We should set up to run multiple `sacct` commands when we have too many jobs for one, and if weird stuff goes wrong with launching `sacct` in general (not just if we launch it and it fails) we should fall back to `scontrol`.

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

Discussion:

Solution patch:
diff --git a/src/toil/batchSystems/slurm.py b/src/toil/batchSystems/slurm.py
index 3ddddff6..df6ebc9b 100644
--- a/src/toil/batchSystems/slurm.py
+++ b/src/toil/batchSystems/slurm.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
+import errno
 import logging
 import math
 import os
@@ -185,6 +186,8 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
         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:
@@ -193,17 +196,36 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
 
             winning_partition = None
             for partition in self.all_partitions:
-                if partition.time_limit >= time_limit and (
-                    winning_partition is None
-                    or partition.time_limit < winning_partition.time_limit
-                ):
-                    # If this partition can fit the job and is faster than the current winner, take it
+                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(
-                    "Could not find a Slurm partition that can fit a job that runs for {time_limit} seconds"
+                    f"Could not find a Slurm partition that can fit a job that runs for {time_limit} seconds"
                 )
 
             if winning_partition is None:
@@ -344,7 +366,9 @@ 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
 
@@ -437,11 +461,25 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
                 "-S",
                 "1970-01-01",
             ]  # override start time limit
-            stdout = call_command(args, quiet=True)
 
             # 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: 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)
 
@@ -649,6 +687,9 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
             if cpu is not None:
                 sbatch_line.append(f"--cpus-per-task={math.ceil(cpu)}")
 
+            if any(option.startswith("--time=") or option == "--time" or option == "-t" for option in sbatch_line):
+                raise RuntimeError("Support for manual --time option has been replaced by Toil --slurmTime")
+
             time_limit: int = self.boss.config.slurm_time  # type: ignore[attr-defined]
             if time_limit is not None:
                 # Put all the seconds in the seconds slot


Test patch:
diff --git a/src/toil/test/batchSystems/test_slurm.py b/src/toil/test/batchSystems/test_slurm.py
index b2818359..2c6aaeb4 100644
--- a/src/toil/test/batchSystems/test_slurm.py
+++ b/src/toil/test/batchSystems/test_slurm.py
@@ -1,3 +1,4 @@
+import errno
 import textwrap
 from queue import Queue
 
@@ -29,6 +30,9 @@ def call_sacct(args, **_) -> str:
         1236|FAILED|0:2
         1236.extern|COMPLETED|0:0
     """
+    if sum(len(a) for a in args) > 1000:
+        # Simulate if the argument list is too long
+        raise OSError(errno.E2BIG, "Argument list is too long")
     # Fake output per fake job-id.
     sacct_info = {
         609663: "609663|FAILED|0:2\n609663.extern|COMPLETED|0:0\n",
@@ -173,6 +177,26 @@ def call_sacct_raises(*_):
         1, "sacct: error: Problem talking to the database: " "Connection timed out"
     )
 
+def call_sinfo(*_) -> str:
+    """
+    Simulate asking for partition info from Slurm
+    """
+    stdout = textwrap.dedent(
+        """\
+        PARTITION GRES TIMELIMIT PRIO_TIER CPUS MEMORY
+        short* (null) 1:00:00 500 256+ 1996800+
+        medium (null) 12:00:00 500 256+ 1996800+
+        long (null) 14-00:00:00 500 256+ 1996800+
+        gpu gpu:A100:8 7-00:00:00 5000 256 996800
+        gpu gpu:A5500:8 7-00:00:00 5000 256 1996800
+        high_priority gpu:A5500:8 7-00:00:00 65000 256 1996800
+        high_priority (null) 7-00:00:00 65000 256+ 1996800+
+        simple_nodelist gpu:A100:8 1:00 65000 256 996800
+        simple_nodelist gpu:A5500:8 1:00 65000 256 1996800
+        simple_nodelist (null) 1:00 65000 256+ 1996800+
+        """
+    )
+    return stdout
 
 class FakeBatchSystem:
     """
@@ -262,6 +286,13 @@ class SlurmTest(ToilTest):
         result = self.worker._getJobDetailsFromSacct(list(expected_result))
         assert result == expected_result, f"{result} != {expected_result}"
 
+    def test_getJobDetailsFromSacct_argument_list_too_big(self):
+        self.monkeypatch.setattr(toil.batchSystems.slurm, "call_command", call_sacct)
+        expected_result = {i: (None, None) for i in range(2000)}
+        result = self.worker._getJobDetailsFromSacct(list(expected_result))
+        assert result == expected_result, f"{result} != {expected_result}"
+
+
     ####
     #### tests for _getJobDetailsFromScontrol()
     ####
@@ -449,3 +480,34 @@ class SlurmTest(ToilTest):
             pass
         else:
             assert False, "Exception CalledProcessErrorStderr not raised"
+
+    ###
+    ### Tests for partition selection
+    ##
+
+    def test_PartitionSet_get_partition(self):
+        self.monkeypatch.setattr(toil.batchSystems.slurm, "call_command", call_sinfo)
+        ps = toil.batchSystems.slurm.SlurmBatchSystem.PartitionSet()
+
+        # At zero. short will win because simple_nodelist has higher priority.
+        self.assertEqual(ps.get_partition(0), "short")
+        # Easily within the partition
+        self.assertEqual(ps.get_partition(10 * 60), "short")
+        # Exactly on the boundary
+        self.assertEqual(ps.get_partition(60 * 60), "short")
+        # Well within the next partition
+        self.assertEqual(ps.get_partition(2 * 60 * 60), "medium")
+        # Can only fit in long
+        self.assertEqual(ps.get_partition(8 * 24 * 60 * 60), "long")
+        # Could fit in gpu or long
+        self.assertEqual(ps.get_partition(6 * 24 * 60 * 60), "long")
+        # Can't fit in anything
+        with self.assertRaises(Exception):
+            ps.get_partition(365 * 24 * 60 * 60)
+
+    def test_PartitionSet_default_gpu_partition(self):
+        self.monkeypatch.setattr(toil.batchSystems.slurm, "call_command", call_sinfo)
+        ps = toil.batchSystems.slurm.SlurmBatchSystem.PartitionSet()
+
+        # Make sure we picked the useful-length GPU partition and not the super short one.
+        self.assertEqual(ps.default_gpu_partition.partition_name, "gpu")