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"]),