diff --git a/src/glicid_spawner/slurm.py b/src/glicid_spawner/slurm.py index 6f059c7..22b5739 100644 --- a/src/glicid_spawner/slurm.py +++ b/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 + } + ) diff --git a/tests/test_slurm.py b/tests/test_slurm.py index 39b6052..ff57c33 100644 --- a/tests/test_slurm.py +++ b/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']