rebase on simplify_slurm_system

NVIDIA · Aug 7, 2024 · 8d34614 · 8d34614
1 parent 7a6809e
commit 8d34614
Showing 1 changed file with 2 additions and 11 deletions.
diff --git 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))