Add SLURM partition and cluster data classes

src/glicid_spawner/slurm.py

@@ -11,7 +11,7 @@ from pathlib import Path
 
 @dataclass
 class SlurmCpu:
-    """SLURM CPU resource."""
+    """SLURM CPU."""
 
     allocated: int
     idle: int
@@ -25,7 +25,7 @@ class SlurmCpu:
 
 @dataclass
 class SlurmGpu:
-    """SLURM GPU resource."""
+    """SLURM GPU."""
 
     name: str = field(default='None')
     nb: int = field(default=0)
@@ -37,10 +37,13 @@ class SlurmGpu:
     def __bool__(self):
         return self.nb > 0
 
+    def __str__(self):
+        return self.name
+
 
 @dataclass
 class SlurmNode:
-    """SLURM node resource."""
+    """SLURM node."""
 
     cluster: str
     partition: str
@@ -59,6 +62,55 @@ class SlurmNode:
         self.mem = int(memory_mb) // 1000  # in GB
         self.gpu = SlurmGpu(*re.findall(r'gpu:(\w+):(\d+)', gres)[0] if 'gpu:' in gres else [])
 
+    def __str__(self):
+        return self.hostname
+
+
+@dataclass
+class SlurmPartition:
+    """SLURM partition."""
+
+    name: str
+    nodes: list
+
+    def __str__(self):
+        return self.name
+
+    def __iter__(self):
+        return iter(self.nodes)
+
+    @property
+    def gpus(self) -> str:
+        """List of GPUs available."""
+        return ':'.join({node.gpu.name for node in self.nodes})
+
+    @property
+    def max_idle_cpu(self) -> int:
+        """Maximum of idle CPU available."""
+        return max(node.cpu.idle for node in self.nodes)
+
+    @property
+    def max_mem(self) -> int:
+        """Maximum of memory available."""
+        return max(node.mem for node in self.nodes)
+
+
+@dataclass
+class SlurmCluster:
+    """SLURM cluster."""
+
+    name: str
+    partitions: list
+
+    def __str__(self):
+        return self.name
+
+    def __iter__(self):
+        return iter(self.partitions)
+
+    def __eq__(self, other):
+        return str(self) == str(other)
+
 
 def sinfo_run(username: str = None) -> str:
     """SLURM SINFO run command."""