Merge pull request #49 from TaekyungHeo/parsing-bugfix

Fix bug in parse_node_list to handle multiple node ranges correctly
NVIDIA · May 29, 2024 · 02ac07d · 02ac07d
2 parents f2261f1 + b1f6cc5
commit 02ac07d
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 13 deletions.
diff --git a/src/cloudai/schema/system/slurm/slurm_system.py b/src/cloudai/schema/system/slurm/slurm_system.py
@@ -74,20 +74,27 @@ def parse_node_list(cls, node_list: str) -> List[str]:
         nodes = []
         if not node_list:
             return []
-        if "[" not in node_list:
-            return [node_list]
-        header, node_number = node_list.split("[")
-        node_number = node_number.replace("]", "")
-        ranges = node_number.split(",")
-        for r in ranges:
-            if "-" in r:
-                start_node, end_node = r.split("-")
-                number_of_digits = len(end_node)
-                nodes.extend(
-                    [f"{header}{str(i).zfill(number_of_digits)}" for i in range(int(start_node), int(end_node) + 1)]
-                )
+
+        components = re.split(r",\s*(?![^[]*\])", node_list)
+        for component in components:
+            if "[" not in component:
+                nodes.append(component)
             else:
-                nodes.append(f"{header}{r}")
+                header, node_number = component.split("[")
+                node_number = node_number.replace("]", "")
+                ranges = node_number.split(",")
+                for r in ranges:
+                    if "-" in r:
+                        start_node, end_node = r.split("-")
+                        number_of_digits = len(end_node)
+                        nodes.extend(
+                            [
+                                f"{header}{str(i).zfill(number_of_digits)}"
+                                for i in range(int(start_node), int(end_node) + 1)
+                            ]
+                        )
+                    else:
+                        nodes.append(f"{header}{r}")
 
         return nodes
 

diff --git a/tests/test_slurm_system.py b/tests/test_slurm_system.py
@@ -115,6 +115,7 @@ def test_update_node_states_with_mocked_outputs(mock_get_sinfo, mock_get_squeue,
     [
         ("node-[048-051]", ["node-048", "node-049", "node-050", "node-051"]),
         ("node-[055,114]", ["node-055", "node-114"]),
+        ("node-[055,114],node-[056,115]", ["node-055", "node-114", "node-056", "node-115"]),
         ("", []),
         ("node-001", ["node-001"]),
         ("node[1-4]", ["node1", "node2", "node3", "node4"]),