Add slurm sinfo parser from file

This commit is contained in:
Benoît Seignovert 2024-02-14 17:09:12 +01:00
parent 52b96548c2
commit b0f8e336bc
Signed by: Benoît Seignovert
GPG key ID: F5D8895227D18A0B
2 changed files with 25 additions and 4 deletions

View file

@ -6,6 +6,7 @@ import subprocess
from dataclasses import dataclass, field
from itertools import groupby
from operator import attrgetter
from pathlib import Path
@dataclass
@ -94,6 +95,12 @@ def sinfo_filter(resources: list, with_states=('idle', 'mixed')) -> dict:
return {key: values for key, values in resources.items() if values}
def sinfo_from_file(fname, with_states=('idle', 'mixed')) -> 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:
"""SLURM SINFO resources available for a given user."""
return sinfo_filter(sinfo_reader(sinfo_run(username=username)), with_states=with_states)

View file

@ -8,13 +8,15 @@ from glicid_spawner.slurm import (
SlurmNode,
sinfo,
sinfo_filter,
sinfo_from_file,
sinfo_reader,
sinfo_run,
subprocess,
)
DATA = Path(__file__).parent / 'data'
SINFO = (DATA / 'sinfo.txt').read_text()
SINFO_FILE = DATA / 'sinfo.txt'
SINFO_CONTENT = SINFO_FILE.read_text()
def test_slurm_dataclasses():
@ -71,7 +73,7 @@ def test_slurm_sinfo_run(monkeypatch):
def test_slurm_sinfo_reader():
"""Test SLURM SINFO reader."""
nodes = sinfo_reader(SINFO)
nodes = sinfo_reader(SINFO_CONTENT)
for node in nodes:
assert isinstance(node, SlurmNode)
@ -117,7 +119,7 @@ def test_slurm_sinfo_reader():
def test_slurm_sinfo_filter(monkeypatch):
"""Test SLURM SINFO filtered resources."""
resources = sinfo_reader(SINFO)
resources = sinfo_reader(SINFO_CONTENT)
clusters = sinfo_filter(resources)
@ -156,9 +158,21 @@ def test_slurm_sinfo_filter(monkeypatch):
assert [len(partitions) for partitions in clusters.values()] == [1]
def test_slurm_sinfo_from_file(monkeypatch):
"""Test SLURM SINFO resources from file."""
resources = sinfo_from_file(SINFO_FILE, with_states=('idle'))
assert [
node.hostname
for cluster, partitions in resources.items()
for nodes in partitions.values()
for node in nodes
] == ['nazare001', 'gnode2', 'visu1', 'cribbar001']
def test_slurm_sinfo_resources(monkeypatch):
"""Test SLURM SINFO resources."""
monkeypatch.setattr(subprocess, 'check_output', lambda _: SINFO.encode())
monkeypatch.setattr(subprocess, 'check_output', lambda _: SINFO_CONTENT.encode())
clusters = sinfo(username='john-doe', with_states=('completing'))