fix: Account for maint and down nodes when calculating available tasks

zenotech · Jan 28, 2025 · c867ebc · c867ebc
1 parent e45ddfd
commit c867ebc
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/mycluster/schedulers/slurm.py b/mycluster/schedulers/slurm.py
@@ -37,7 +37,7 @@
 import subprocess
 
 from .base import Scheduler
-from mycluster.exceptions import SchedulerException, ConfigurationException
+from mycluster.exceptions import SchedulerException
 
 
 class Slurm(Scheduler):
@@ -92,12 +92,14 @@ def available_tasks(self, queue_id):
         max_tasks = 0
         queue_name = queue_id
         nc = self.node_config(queue_id)
-        with os.popen("sinfo -sh -p " + queue_name) as f:
+        with os.popen(f"sinfo -sh -p {queue_name} -o '%C'") as f:
             line = f.readline()
-            new_line = re.sub(" +", " ", line.strip())
-            line = new_line.split(" ")[3]
-            free_tasks = int(line.split("/")[1]) * nc["max task"]
-            max_tasks = int(line.split("/")[3]) * nc["max task"]
+            max_tasks = int(line.split("/")[3])
+
+        with os.popen(f"sinfo -sh -p {queue_name} -o '%C' | grep 'idle'") as f:
+            for line in f:
+                free_tasks += int(line.split("/")[1])
+
         return {"available": free_tasks, "max tasks": max_tasks}
 
     def accounts(self):