Filter sinfo by slurm cluster and partition dataclasses

This commit is contained in:
Benoît Seignovert 2024-02-19 18:09:41 +01:00
parent 02166be4d2
commit 21935f15f5
Signed by: Benoît Seignovert
GPG key ID: F5D8895227D18A0B
2 changed files with 42 additions and 30 deletions

View file

@ -156,16 +156,20 @@ def sinfo_filter(resources: list, with_states=('idle', 'mixed')) -> dict:
Grouped by cluster and partition names.
"""
resources = {
cluster: {
partition: available
for partition, nodes in groupby(partitions, key=attrgetter('partition'))
if (available := [node for node in nodes if node.state in with_states])
}
clusters = [
SlurmCluster(
cluster,
[
SlurmPartition(partition, nodes_with_states)
for partition, nodes in groupby(partitions, key=attrgetter('partition'))
if (nodes_with_states := [node for node in nodes if node.state in with_states])
],
)
for cluster, partitions in groupby(resources, key=attrgetter('cluster'))
}
]
return {key: values for key, values in resources.items() if values}
# Remove empty cluster
return {cluster.name: cluster for cluster in clusters if cluster}
def sinfo_from_file(fname, with_states=('idle', 'mixed')) -> dict: