jupyter/spawner
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix support for GPU gres with a dot in their name

Browse source
This commit is contained in:
Benoît Seignovert 2024-05-21 17:52:26 +02:00
parent d2c0631956
commit 255bf6201e
Signed by: Benoît Seignovert
GPG key ID: F5D8895227D18A0B
4 changed files with 15 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
src/glicid_spawner/slurm.py
View file

@ -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

2
src/glicid_spawner/templates/views/resources.jinja
View file

@ -41,7 +41,7 @@
data-max-duration="{{max_duration}}"
{%- if loop.first %} checked{% endif %}>
<label for="gpu_{{gpu}}" class="btn btn-default btn-block">
{{ gpu }}
{{ gpu | replace("_"," ") | replace("."," ") | replace("gb","GB") }}
</label>
</div>
{% endfor -%}

4
tests/data/sinfo.txt
View file

@ -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:

14
tests/test_slurm.py
View file

@ -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']