Add completing to states available

This commit is contained in:
Benoît Seignovert 2024-03-07 15:27:22 +01:00
parent ae9ac12918
commit c4c8fa33aa
Signed by: Benoît Seignovert
GPG key ID: F5D8895227D18A0B
2 changed files with 12 additions and 6 deletions

View file

@ -8,6 +8,12 @@ from itertools import groupby
from operator import attrgetter from operator import attrgetter
from pathlib import Path from pathlib import Path
STATES_AVAILABLE = (
'idle',
'mixed',
'completing',
)
@dataclass @dataclass
class SlurmCpu: class SlurmCpu:
@ -150,7 +156,7 @@ def sinfo_reader(result: str) -> list:
return [SlurmNode(*re.findall('.{20}', node)) for node in result.splitlines()] 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). """SLURM SINFO filtered resources available with a given state(s).
Grouped by cluster and partition names. 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} 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.""" """SLURM SINFO resources available from a given file."""
content = Path(fname).read_text() content = Path(fname).read_text()
return sinfo_filter(sinfo_reader(content), with_states=with_states) 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.""" """SLURM SINFO resources available for a given user."""
return sinfo_filter(sinfo_reader(sinfo_run(username=username)), with_states=with_states) return sinfo_filter(sinfo_reader(sinfo_run(username=username)), with_states=with_states)

View file

@ -174,13 +174,13 @@ def test_slurm_sinfo_filter(monkeypatch):
assert len(clusters) == 3 assert len(clusters) == 3
assert list(clusters) == ['N/A', 'nautilus', 'waves'] # __eq__ on cluster.name 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'] nautilus = clusters['nautilus']
assert isinstance(nautilus, SlurmCluster) assert isinstance(nautilus, SlurmCluster)
assert len(nautilus) == 3 assert len(nautilus) == 4
assert nautilus.partitions == ['gpu', 'visu', 'all'] # __eq__ on partition.name assert nautilus.partitions == ['standard', 'gpu', 'visu', 'all'] # __eq__ on partition.name
gpu = nautilus['gpu'] gpu = nautilus['gpu']