Add getitem and eq to slurm objects

This commit is contained in:
Benoît Seignovert 2024-02-19 17:53:35 +01:00
parent e8ae32cc86
commit 02166be4d2
Signed by: Benoît Seignovert
GPG key ID: F5D8895227D18A0B
2 changed files with 31 additions and 20 deletions

View file

@ -40,6 +40,9 @@ class SlurmGpu:
def __str__(self):
return self.name
def __eq__(self, other):
return str(self) == str(other)
@dataclass
class SlurmNode:
@ -65,6 +68,9 @@ class SlurmNode:
def __str__(self):
return self.hostname
def __eq__(self, other):
return str(self) == str(other)
@dataclass
class SlurmPartition:
@ -76,14 +82,17 @@ class SlurmPartition:
def __str__(self):
return self.name
def __eq__(self, other):
return str(self) == str(other)
def __iter__(self):
return iter(self.nodes)
def __len__(self):
return len(self.nodes)
def __eq__(self, other):
return str(self) == str(other)
def __getitem__(self, other):
return self.nodes[self.nodes.index(other)]
@property
def gpus(self) -> str:
@ -111,14 +120,17 @@ class SlurmCluster:
def __str__(self):
return self.name
def __eq__(self, other):
return str(self) == str(other)
def __iter__(self):
return iter(self.partitions)
def __len__(self):
return len(self.partitions)
def __eq__(self, other):
return str(self) == str(other)
def __getitem__(self, other):
return self.partitions[self.partitions.index(other)]
def sinfo_run(username: str = None) -> str: