diff --git a/src/cloudai/systems/slurm/slurm_system.py b/src/cloudai/systems/slurm/slurm_system.py index 18ec3ff62..0c39a3bee 100644 --- a/src/cloudai/systems/slurm/slurm_system.py +++ b/src/cloudai/systems/slurm/slurm_system.py @@ -363,24 +363,15 @@ def get_available_nodes_from_partition(self, partition_name: str, number_of_node grouped_nodes = { SlurmNodeState.IDLE: [], SlurmNodeState.COMPLETING: [], - SlurmNodeState.ALLOCATED: [], } for node in self.partitions[partition_name]: if node.state in grouped_nodes: - # Exclude nodes allocated to the current user - if node.state == SlurmNodeState.ALLOCATED and node.user == current_user: - continue - if node.state in grouped_nodes: - grouped_nodes[node.state].append(node) + grouped_nodes[node.state].append(node) # Allocate nodes based on priority: idle, then completing, then allocated allocated_nodes = [] - for state in [ - SlurmNodeState.IDLE, - SlurmNodeState.COMPLETING, - SlurmNodeState.ALLOCATED, - ]: + for state in grouped_nodes: while grouped_nodes[state] and len(allocated_nodes) < number_of_nodes: allocated_nodes.append(grouped_nodes[state].pop(0))