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

Update slurm sinfo filter to discard empty cluster

Browse source
This commit is contained in:
Benoît Seignovert 2024-02-14 16:31:52 +01:00
parent 867e217d44
commit 52b96548c2
Signed by: Benoît Seignovert
GPG key ID: F5D8895227D18A0B
2 changed files with 64 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

19
src/glicid_spawner/slurm.py
View file

@ -59,7 +59,7 @@ class SlurmNode:
self.gpu = SlurmGpu(*re.findall(r'gpu:(\w+):(\d+)', gres)[0] if 'gpu:' in gres else [])
def _sinfo_run(username: str = None) -> str:
def sinfo_run(username: str = None) -> str:
"""SLURM SINFO run command."""
flags = '--federation --noheader --responding'
fmt = 'Cluster,PartitionName,NodeHost,StateLong,CPUsState,Memory,Gres'
@ -71,20 +71,18 @@ def _sinfo_run(username: str = None) -> str:
return subprocess.check_output(shlex.split(cmd, posix=False)).decode('utf-8')
def _sinfo_reader(result) -> list:
def sinfo_reader(result: str) -> list:
"""SLURM SINFO reader."""
return [SlurmNode(*re.findall('.{20}', node)) for node in result.splitlines()]
def sinfo(username: str = None, with_states=('idle', 'mixed')) -> dict:
"""SLURM SINFO resources available with a given state(s).
def sinfo_filter(resources: list, with_states=('idle', 'mixed')) -> dict:
"""SLURM SINFO filtered resources available with a given state(s).
Grouped by cluster and partition names.
"""
resources = _sinfo_reader(_sinfo_run(username=username))
return {
resources = {
cluster: {
partition: available
for partition, nodes in groupby(partitions, key=attrgetter('partition'))
@ -92,3 +90,10 @@ def sinfo(username: str = None, with_states=('idle', 'mixed')) -> dict:
}
for cluster, partitions in groupby(resources, key=attrgetter('cluster'))
}
return {key: values for key, values in resources.items() if values}
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)