From b1f6cc5a50cc0bedf20fca56c271681c7b2ea709 Mon Sep 17 00:00:00 2001
From: Taekyung Heo <7621438+TaekyungHeo@users.noreply.github.com>
Date: Wed, 29 May 2024 10:48:21 -0700
Subject: [PATCH] Fix bug in parse_node_list to handle multiple node ranges
 correctly

---
 .../schema/system/slurm/slurm_system.py       | 33 +++++++++++--------
 tests/test_slurm_system.py                    |  1 +
 2 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/src/cloudai/schema/system/slurm/slurm_system.py b/src/cloudai/schema/system/slurm/slurm_system.py
index 2c3290148..d8803b27d 100644
--- 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
index cdc8fda24..750bdff52 100644
--- 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"]),