From 4eef9d7016d7dc38ea98dec655b730977bf1406a Mon Sep 17 00:00:00 2001 From: Benoit Seignovert Date: Mon, 19 Feb 2024 17:32:23 +0100 Subject: [PATCH] Add slurm cluster and partition len property --- src/glicid_spawner/slurm.py | 6 ++++++ tests/test_slurm.py | 4 ++++ 2 files changed, 10 insertions(+) diff --git a/src/glicid_spawner/slurm.py b/src/glicid_spawner/slurm.py index 6a29a94..bdc1029 100644 --- a/src/glicid_spawner/slurm.py +++ b/src/glicid_spawner/slurm.py @@ -79,6 +79,9 @@ class SlurmPartition: def __iter__(self): return iter(self.nodes) + def __len__(self): + return len(self.nodes) + @property def gpus(self) -> str: """List of GPUs available.""" @@ -108,6 +111,9 @@ class SlurmCluster: def __iter__(self): return iter(self.partitions) + def __len__(self): + return len(self.partitions) + def __eq__(self, other): return str(self) == str(other) diff --git a/tests/test_slurm.py b/tests/test_slurm.py index 7901cc3..e93070f 100644 --- a/tests/test_slurm.py +++ b/tests/test_slurm.py @@ -73,6 +73,8 @@ def test_slurm_dataclasses(): assert str(partition) == 'standard' # = name assert partition.name == 'standard' + assert len(partition) == 1 + for _node in partition: assert str(_node) == 'cnode001' @@ -87,6 +89,8 @@ def test_slurm_dataclasses(): assert cluster.name == 'nautilus' assert cluster == 'nautilus' # __eq__ + assert len(cluster) == 1 + for _partition in cluster: assert str(_partition) == 'standard'