Repository: DataBiosphere/toil

Issue:
Slurm partition detection sends non-GPU jobs to GPU partitions
@glennhickey reports that in Toil 8, the Slurm batch system is sending non-GPU jobs to GPU partitions.

This is probably because https://github.com/DataBiosphere/toil/pull/5032 did not add comprehensive unit tests for partition selection.

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

Discussion:

Solution patch:
diff --git a/src/toil/batchSystems/slurm.py b/src/toil/batchSystems/slurm.py
index 3ddddff6..bfebad49 100644
--- a/src/toil/batchSystems/slurm.py
+++ b/src/toil/batchSystems/slurm.py
@@ -185,6 +185,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 +195,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:
@@ -649,6 +670,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..a2e026f8 100644
--- a/src/toil/test/batchSystems/test_slurm.py
+++ b/src/toil/test/batchSystems/test_slurm.py
@@ -173,6 +173,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:
     """
@@ -449,3 +469,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")