From 255bf6201eab4823bf10ace74b90169e0cb46c13 Mon Sep 17 00:00:00 2001 From: Benoit Seignovert Date: Tue, 21 May 2024 17:52:26 +0200 Subject: [PATCH] Fix support for GPU gres with a dot in their name --- src/glicid_spawner/slurm.py | 2 +- src/glicid_spawner/templates/views/resources.jinja | 2 +- tests/data/sinfo.txt | 4 ++-- tests/test_slurm.py | 14 +++++++++++--- 4 files changed, 15 insertions(+), 7 deletions(-) diff --git a/src/glicid_spawner/slurm.py b/src/glicid_spawner/slurm.py index 1aea52d..15da6fd 100644 --- a/src/glicid_spawner/slurm.py +++ b/src/glicid_spawner/slurm.py @@ -69,7 +69,7 @@ class SlurmNode: self.state = state.strip().lower() self.cpu = SlurmCpu(*re.findall(r'(\d+)/(\d+)/\d+/(\d+)', cpus_state)[0]) self.mem = int(memory_mb) // 1000 # in GB - self.gpu = SlurmGpu(*re.findall(r'gpu:(\w+):(\d+)', gres)[0] if 'gpu:' in gres else []) + self.gpu = SlurmGpu(*re.findall(r'gpu:([\w\.]+):(\d+)', gres)[0] if 'gpu:' in gres else []) def __str__(self): return self.hostname diff --git a/src/glicid_spawner/templates/views/resources.jinja b/src/glicid_spawner/templates/views/resources.jinja index efae401..d29de24 100644 --- a/src/glicid_spawner/templates/views/resources.jinja +++ b/src/glicid_spawner/templates/views/resources.jinja @@ -41,7 +41,7 @@ data-max-duration="{{max_duration}}" {%- if loop.first %} checked{% endif %}> {% endfor -%} diff --git a/tests/data/sinfo.txt b/tests/data/sinfo.txt index 37d0b25..8079789 100644 --- a/tests/data/sinfo.txt +++ b/tests/data/sinfo.txt @@ -9,13 +9,13 @@ nautilus standard cnode001 completing nautilus bigmem cnode002 planned 0/96/0/96 768000 (null) nautilus gpu gnode1 mixed 4/92/0/96 768000 gpu:A100:1(S:0-1) nautilus gpu gnode2 idle 0/96/0/96 256000 gpu:A100:2(S:0-1) -nautilus gpu gnode3 allocated 96/0/0/96 128000 gpu:A100:4(S:0-1) +nautilus gpu gnode3 allocated 96/0/0/96 128000 gpu:A100_2g.10gb:6(S nautilus visu visu1 idle 0/96/0/96 768000 (null) nautilus all cnode001 completing 0/96/0/96 384000 (null) nautilus all cnode002 planned 0/96/0/96 768000 (null) nautilus all gnode1 mixed 4/92/0/96 768000 gpu:A100:1(S:0-1) nautilus all gnode2 idle 0/96/0/96 256000 gpu:A100:2(S:0-1) -nautilus all gnode3 allocated 96/0/0/96 128000 gpu:A100:4(S:0-1) +nautilus all gnode3 allocated 96/0/0/96 128000 gpu:A100_2g.10gb:6(S nautilus all visu1 idle 0/96/0/96 768000 (null) waves standard cribbar001 mixed 30/10/0/40 16000 (null) waves gpu budbud006 allocated 64/0/0/64 256000 gpu:a100:2,mps:a100: diff --git a/tests/test_slurm.py b/tests/test_slurm.py index a4d8693..c38829d 100644 --- a/tests/test_slurm.py +++ b/tests/test_slurm.py @@ -161,9 +161,17 @@ def test_slurm_sinfo_reader(): 'P100', 'K80', 'P100', - ] + 8 * ['A100'] + 'A100', + 'A100', + 'A100_2g.10gb', + 'A100', + 'A100', + 'A100_2g.10gb', + 'A100', + 'A100', + ] - assert [node.gpu.nb for node in nodes if node.gpu] == [2, 2, 2, 4, 1, 1, 2, 4, 1, 2, 4, 2, 2] + assert [node.gpu.nb for node in nodes if node.gpu] == [2, 2, 2, 4, 1, 1, 2, 6, 1, 2, 6, 2, 2] def test_slurm_sinfo_filter(monkeypatch): @@ -253,4 +261,4 @@ def test_slurm_gres(): gpus = gres(resources) # Sorted and without duplicates - assert gpus == ['A100', 'A40', 'None', 'P100', 'T4'] + assert gpus == ['A100', 'A100_2g.10gb', 'A40', 'None', 'P100', 'T4']