From c4c8fa33aabf5bc700e60bda48b7fa421a2e0eed Mon Sep 17 00:00:00 2001 From: Benoit Seignovert Date: Thu, 7 Mar 2024 15:27:22 +0100 Subject: [PATCH] Add completing to states available --- src/glicid_spawner/slurm.py | 12 +++++++++--- tests/test_slurm.py | 6 +++--- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/src/glicid_spawner/slurm.py b/src/glicid_spawner/slurm.py index 59e64af..011193e 100644 --- a/src/glicid_spawner/slurm.py +++ b/src/glicid_spawner/slurm.py @@ -8,6 +8,12 @@ from itertools import groupby from operator import attrgetter from pathlib import Path +STATES_AVAILABLE = ( + 'idle', + 'mixed', + 'completing', +) + @dataclass class SlurmCpu: @@ -150,7 +156,7 @@ def sinfo_reader(result: str) -> list: return [SlurmNode(*re.findall('.{20}', node)) for node in result.splitlines()] -def sinfo_filter(resources: list, with_states=('idle', 'mixed')) -> dict: +def sinfo_filter(resources: list, with_states=STATES_AVAILABLE) -> dict: """SLURM SINFO filtered resources available with a given state(s). Grouped by cluster and partition names. @@ -172,13 +178,13 @@ def sinfo_filter(resources: list, with_states=('idle', 'mixed')) -> dict: return {cluster.name: cluster for cluster in clusters if cluster} -def sinfo_from_file(fname, with_states=('idle', 'mixed')) -> dict: +def sinfo_from_file(fname, with_states=STATES_AVAILABLE) -> dict: """SLURM SINFO resources available from a given file.""" content = Path(fname).read_text() return sinfo_filter(sinfo_reader(content), with_states=with_states) -def sinfo(username: str = None, with_states=('idle', 'mixed')) -> dict: +def sinfo(username: str = None, with_states=STATES_AVAILABLE) -> dict: """SLURM SINFO resources available for a given user.""" return sinfo_filter(sinfo_reader(sinfo_run(username=username)), with_states=with_states) diff --git a/tests/test_slurm.py b/tests/test_slurm.py index bcb9429..28e23f1 100644 --- a/tests/test_slurm.py +++ b/tests/test_slurm.py @@ -174,13 +174,13 @@ def test_slurm_sinfo_filter(monkeypatch): assert len(clusters) == 3 assert list(clusters) == ['N/A', 'nautilus', 'waves'] # __eq__ on cluster.name - assert [len(partitions) for partitions in clusters.values()] == [2, 3, 2] + assert [len(partitions) for partitions in clusters.values()] == [2, 4, 2] nautilus = clusters['nautilus'] assert isinstance(nautilus, SlurmCluster) - assert len(nautilus) == 3 - assert nautilus.partitions == ['gpu', 'visu', 'all'] # __eq__ on partition.name + assert len(nautilus) == 4 + assert nautilus.partitions == ['standard', 'gpu', 'visu', 'all'] # __eq__ on partition.name gpu = nautilus['gpu']