Merge pull request #9 from vub-hpc/fix

Fix
vub-hpc · Nov 13, 2024 · 7860462 · 7860462
2 parents d3cdca6 + 46a7207
commit 7860462
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 9 deletions.
diff --git a/jupyterhub_moss/spawner.py b/jupyterhub_moss/spawner.py
@@ -84,7 +84,7 @@ def _validate_partitions(self, proposal: dict) -> dict[str, dict]:
 
     slurm_info_cmd = traitlets.Unicode(
         # Get number of nodes/state, cores/node, cores/state, gpus, total memory for all partitions
-        r"sinfo -N -p {partition} {clusters} --noheader -O 'Cluster,PartitionName,StateCompact,CPUsState,Gres:64,GresUsed:64,Memory,Time,OverSubscribe'",
+        r"sinfo -N -p {partition} {clusters} --noheader -O 'Cluster,PartitionName,StateCompact,CPUsState,Gres:128,GresUsed:128,Memory,Time,OverSubscribe'",
         help="Command to query cluster information from Slurm. Formatted using req_xyz traits as {xyz}."
         "Output will be parsed by ``slurm_info_resources``.",
     ).tag(config=True)
@@ -138,7 +138,7 @@ def _slurm_info_resources(
                 continue
 
             # unique reference from cluster and partition names
-            if cluster != 'N/A':
+            if cluster != "N/A":
                 partition = f"{cluster}.{partition}"
 
             # core count - allocated/idle/other/total
@@ -167,7 +167,7 @@ def _slurm_info_resources(
                 # Slurm does not report usage of oversubscribed resources
                 # sinfo shows partition as full once each core has a single job running on them
                 # even though more jobs will be allowed according to the oversubscription factor
-                _, oversub_factor= oversubscribe.split(":", 1)
+                _, oversub_factor = oversubscribe.split(":", 1)
                 ncores_total *= int(oversub_factor)
                 ncores_idle = ncores_total
                 shared = True
@@ -213,8 +213,12 @@ def _slurm_info_resources(
 
             if partition in partitions_info:
                 # update display counters of existing partition
-                slots_counters = zip(partitions_info[partition]["job_slots"], resources["job_slots"])
-                partitions_info[partition]["job_slots"] = [old + new for old, new in slots_counters]
+                slots_counters = zip(
+                    partitions_info[partition]["job_slots"], resources["job_slots"]
+                )
+                partitions_info[partition]["job_slots"] = [
+                    old + new for old, new in slots_counters
+                ]
             else:
                 # add new partition
                 partitions_info[partition] = resources
@@ -278,7 +282,9 @@ async def _get_partitions_info(self) -> dict[str, PartitionInfo]:
                     format_template(self.slurm_info_cmd, **subvars),
                 )
             )
-            self.log.debug("Slurm info command for partition ID '%s': %s", partition_id, sinfo_cmd)
+            self.log.debug(
+                "Slurm info command for partition ID '%s': %s", partition_id, sinfo_cmd
+            )
             partition_sinfo_out = await self.run_command(sinfo_cmd)
             # self.log.debug("Slurm info command output: %s", partition_sinfo_out)
 

diff --git a/jupyterhub_moss/utils.py b/jupyterhub_moss/utils.py
@@ -89,13 +89,12 @@ def create_prologue(
 
 def parse_partition_id(partition_id: str) -> tuple[str, str]:
     """Parse partition ID string from configuration file"""
-    id_sections = partition_id.split('.', 1)
+    id_sections = partition_id.split(".", 1)
     id_sections.reverse()
     partition_name = id_sections.pop(0)
 
-    cluster_name = ''
+    cluster_name = ""
     if id_sections:
         cluster_name = id_sections.pop(0)
 
     return partition_name, cluster_name
-