Add SLURM gres extraction from resources

src/glicid_spawner/slurm.py

@@ -181,3 +181,15 @@ def sinfo_from_file(fname, with_states=('idle', 'mixed')) -> dict:
 def sinfo(username: str = None, with_states=('idle', 'mixed')) -> dict:
     """SLURM SINFO resources available for a given user."""
     return sinfo_filter(sinfo_reader(sinfo_run(username=username)), with_states=with_states)
+
+
+def gres(resources: dict) -> list:
+    """List SLURM GPU resources."""
+    return sorted(
+        {
+            node.gpu.name
+            for cluster in resources.values()
+            for partition in cluster
+            for node in partition
+        }
+    )

tests/test_slurm.py

@@ -8,6 +8,7 @@ from glicid_spawner.slurm import (
     SlurmGpu,
     SlurmNode,
     SlurmPartition,
+    gres,
     sinfo,
     sinfo_filter,
     sinfo_from_file,
@@ -242,3 +243,13 @@ def test_slurm_sinfo_resources(monkeypatch):
 
     assert isinstance(cnode, SlurmNode)
     assert cnode == 'cnode001'
+
+
+def test_slurm_gres():
+    """Test SLURM GPU resources extraction."""
+    resources = sinfo_from_file(SINFO_FILE, with_states=('idle', 'idle~', 'mixed', 'allocated'))
+
+    gpus = gres(resources)
+
+    # Sorted and without duplicates
+    assert gpus == ['A100', 'A40', 'None', 'P100', 'T4']