From 21935f15f5f41087ab5d88a3b9465964411c5b36 Mon Sep 17 00:00:00 2001
From: Benoit Seignovert <benoit.seignovert@univ-nantes.fr>
Date: Mon, 19 Feb 2024 18:09:41 +0100
Subject: [PATCH] Filter sinfo by slurm cluster and partition dataclasses

---
 src/glicid_spawner/slurm.py | 20 ++++++++------
 tests/test_slurm.py         | 52 +++++++++++++++++++++----------------
 2 files changed, 42 insertions(+), 30 deletions(-)

diff --git a/src/glicid_spawner/slurm.py b/src/glicid_spawner/slurm.py
index c059b0e..6f059c7 100644
--- a/src/glicid_spawner/slurm.py
+++ b/src/glicid_spawner/slurm.py
@@ -156,16 +156,20 @@ def sinfo_filter(resources: list, with_states=('idle', 'mixed')) -> dict:
     Grouped by cluster and partition names.
 
     """
-    resources = {
-        cluster: {
-            partition: available
-            for partition, nodes in groupby(partitions, key=attrgetter('partition'))
-            if (available := [node for node in nodes if node.state in with_states])
-        }
+    clusters = [
+        SlurmCluster(
+            cluster,
+            [
+                SlurmPartition(partition, nodes_with_states)
+                for partition, nodes in groupby(partitions, key=attrgetter('partition'))
+                if (nodes_with_states := [node for node in nodes if node.state in with_states])
+            ],
+        )
         for cluster, partitions in groupby(resources, key=attrgetter('cluster'))
-    }
+    ]
 
-    return {key: values for key, values in resources.items() if values}
+    # Remove empty cluster
+    return {cluster.name: cluster for cluster in clusters if cluster}
 
 
 def sinfo_from_file(fname, with_states=('idle', 'mixed')) -> dict:
diff --git a/tests/test_slurm.py b/tests/test_slurm.py
index 3f8185c..39b6052 100644
--- a/tests/test_slurm.py
+++ b/tests/test_slurm.py
@@ -172,25 +172,25 @@ def test_slurm_sinfo_filter(monkeypatch):
 
     assert isinstance(clusters, dict)
     assert len(clusters) == 3
-    assert list(clusters) == ['N/A', 'nautilus', 'waves']
+    assert list(clusters) == ['N/A', 'nautilus', 'waves']  # __eq__ on cluster.name
 
     assert [len(partitions) for partitions in clusters.values()] == [2, 3, 2]
 
-    partitions = clusters['nautilus']
+    nautilus = clusters['nautilus']
 
-    assert isinstance(partitions, dict)
-    assert len(partitions) == 3
-    assert list(partitions) == ['gpu', 'visu', 'all']
+    assert isinstance(nautilus, SlurmCluster)
+    assert len(nautilus) == 3
+    assert nautilus.partitions == ['gpu', 'visu', 'all']  # __eq__ on partition.name
 
-    gpu_nodes = partitions['gpu']
+    gpu = nautilus['gpu']
 
-    assert len(gpu_nodes) == 2
-    assert [node.hostname for node in gpu_nodes] == ['gnode1', 'gnode2']
-    assert [node.cpu.allocated for node in gpu_nodes] == [4, 0]
-    assert [node.cpu.idle for node in gpu_nodes] == [92, 96]
-    assert [node.mem for node in gpu_nodes] == [768, 256]
-    assert [node.gpu.name for node in gpu_nodes] == ['A100', 'A100']
-    assert [node.gpu.nb for node in gpu_nodes] == [1, 2]
+    assert len(gpu) == 2
+    assert [node.hostname for node in gpu] == ['gnode1', 'gnode2']
+    assert [node.cpu.allocated for node in gpu] == [4, 0]
+    assert [node.cpu.idle for node in gpu] == [92, 96]
+    assert [node.mem for node in gpu] == [768, 256]
+    assert [node.gpu.name for node in gpu] == ['A100', 'A100']
+    assert [node.gpu.nb for node in gpu] == [1, 2]
 
     # Get only `idle` nodes
     clusters = sinfo_filter(resources, with_states=('idle'))
@@ -211,9 +211,9 @@ def test_slurm_sinfo_from_file(monkeypatch):
 
     assert [
         node.hostname
-        for cluster, partitions in resources.items()
-        for nodes in partitions.values()
-        for node in nodes
+        for cluster in resources.values()
+        for partition in cluster
+        for node in partition
     ] == ['nazare001', 'gnode2', 'visu1', 'gnode2', 'visu1']
 
 
@@ -223,14 +223,22 @@ def test_slurm_sinfo_resources(monkeypatch):
 
     clusters = sinfo(username='john-doe', with_states=('completing'))
 
+    assert 'nautilus' in clusters
     assert list(clusters) == ['nautilus']
 
-    partitions = clusters['nautilus']
+    nautilus = clusters['nautilus']
 
-    assert list(partitions) == ['standard', 'all']
+    assert isinstance(nautilus, SlurmCluster)
+    assert 'standard' in nautilus
+    assert list(nautilus) == ['standard', 'all']
 
-    std_nodes = partitions['standard']
+    standard = nautilus['standard']
 
-    assert std_nodes == [
-        SlurmNode(*'nautilus standard cnode001 completing 0/96/0/96 384000 (null)'.split())
-    ]
+    assert isinstance(standard, SlurmPartition)
+    assert 'cnode001' in standard
+    assert list(standard) == ['cnode001']
+
+    cnode = standard['cnode001']
+
+    assert isinstance(cnode, SlurmNode)
+    assert cnode == 'cnode001'