Optimized the retrieval of nodes info from Slurm scheduler

demartinofra · tilne · commit 3abca9b333a6 · 2020-11-11T13:48:47.000-08:00
Retrieving all nodes with scontrol show nodes is much faster than having Slurm filter those from inactive partition out

Signed-off-by: Francesco De Martino &lt;fdm@amazon.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,7 @@ This file is used to list changes made in each version of the aws-parallelcluste
 -----
 
 **CHANGES**
+- Optimized the retrieval of nodes info from Slurm scheduler.
 - Increase default timeout for Slurm commands submitted by clustermgtd and computemgtd from 10 to 30 seconds.
 
 **BUG FIXES**
diff --git a/src/common/schedulers/slurm_commands.py b/src/common/schedulers/slurm_commands.py
@@ -73,13 +73,14 @@ class SlurmNode:
     SLURM_SCONTROL_POWERING_DOWN_STATE = "POWERING_DOWN"
     SLURM_SCONTROL_POWER_STATE = "IDLE+CLOUD+POWER"
 
-    def __init__(self, name, nodeaddr, nodehostname, state):
+    def __init__(self, name, nodeaddr, nodehostname, state, partitions=None):
         """Initialize slurm node with attributes."""
         self.name = name
         self.is_static = is_static_node(name)
         self.nodeaddr = nodeaddr
         self.nodehostname = nodehostname
         self.state = state
+        self.partitions = partitions.strip().split(",") if partitions else None
 
     def is_nodeaddr_set(self):
         """Check if nodeaddr(private ip) for the node is set."""
@@ -328,15 +329,18 @@ def set_nodes_down_and_power_save(node_list, reason):
     set_nodes_power_down(node_list, reason=reason)
 
 
-def get_nodes_info(nodes, command_timeout=DEFAULT_GET_INFO_COMMAND_TIMEOUT):
+def get_nodes_info(nodes="", command_timeout=DEFAULT_GET_INFO_COMMAND_TIMEOUT):
     """
     Retrieve SlurmNode list from slurm nodelist notation.
 
     Sample slurm nodelist notation: queue1-dy-c5_xlarge-[1-3],queue2-st-t2_micro-5.
     """
+    # awk is used to replace the \n\n record separator with '---\n'
+    # Note: In case the node does not belong to any partition the Partitions field is missing from Slurm output
     show_node_info_command = (
-        f'{SCONTROL} show nodes {nodes} | grep -oP "^NodeName=\\K(\\S+)| '
-        'NodeAddr=\\K(\\S+)| NodeHostName=\\K(\\S+)| State=\\K(\\S+)"'
+        f'{SCONTROL} show nodes {nodes} | awk \'BEGIN{{RS="\\n\\n" ; ORS="---\\n";}} {{print}}\' | '
+        'grep -oP "^NodeName=\\K(\\S+)| NodeAddr=\\K(\\S+)| NodeHostName=\\K(\\S+)| State=\\K(\\S+)|'
+        ' Partitions=\\K(\\S+)|(---)"'
     )
     nodeinfo_str = check_command_output(show_node_info_command, timeout=command_timeout, shell=True)
 
@@ -393,4 +397,36 @@ def _get_partition_nodes(partition_name, command_timeout=DEFAULT_GET_INFO_COMMAN
 
 def _parse_nodes_info(slurm_node_info):
     """Parse slurm node info into SlurmNode objects."""
-    return [SlurmNode(*node) for node in grouper(slurm_node_info.splitlines(), 4)]
+    # [ec2-user@ip-10-0-0-58 ~]$ /opt/slurm/bin/scontrol show nodes compute-dy-c5xlarge-[1-3],compute-dy-c5xlarge-50001\
+    # | awk 'BEGIN{RS="\n\n" ; ORS="---\n";} {print}' | grep -oP "^NodeName=\K(\S+)| NodeAddr=\K(\S+)|\
+    # NodeHostName=\K(\S+)| State=\K(\S+)| Partitions=\K(\S+)|(---)"
+    # compute-dy-c5xlarge-1
+    # 1.2.3.4
+    # compute-dy-c5xlarge-1
+    # IDLE+CLOUD+POWER
+    # compute,compute2
+    # ---
+    # compute-dy-c5xlarge-2
+    # 1.2.3.4
+    # compute-dy-c5xlarge-2
+    # IDLE+CLOUD+POWER
+    # compute,compute2
+    # ---
+    # compute-dy-c5xlarge-3
+    # 1.2.3.4
+    # compute-dy-c5xlarge-3
+    # IDLE+CLOUD+POWER
+    # compute,compute2
+    # ---
+    # compute-dy-c5xlarge-50001
+    # 1.2.3.4
+    # compute-dy-c5xlarge-50001
+    # IDLE+CLOUD+POWER
+    # ---
+    node_info = slurm_node_info.split("---")
+    slurm_nodes = []
+    for node in node_info:
+        lines = node.strip().splitlines()
+        if lines:
+            slurm_nodes.append(SlurmNode(*lines))
+    return slurm_nodes
diff --git a/src/slurm_plugin/clustermgtd.py b/src/slurm_plugin/clustermgtd.py
@@ -439,7 +439,7 @@ def _write_timestamp_to_file(self):
 
     @staticmethod
     @retry(stop_max_attempt_number=2, wait_fixed=1000)
-    def _get_node_info_with_retry(nodes):
+    def _get_node_info_with_retry(nodes=""):
         return get_nodes_info(nodes)
 
     @staticmethod
@@ -457,18 +457,22 @@ def _get_node_info_from_partition():
         try:
             inactive_nodes = []
             active_nodes = []
-            partitions = ClusterManager._get_partition_info_with_retry()
+            ignored_nodes = []
+            partitions = {partition.name: partition for partition in ClusterManager._get_partition_info_with_retry()}
             log.debug("Partitions: %s", partitions)
-            for part in partitions:
-                # Get node info if partition has relevant node
-                # See SlurmNode.SLURM_RELEVANT_SINFO_STATES
-                if part.nodes:
-                    nodes = ClusterManager._get_node_info_with_retry(part.nodes)
-                    if "INACTIVE" in part.state:
-                        inactive_nodes.extend(nodes)
-                    else:
-                        active_nodes.extend(nodes)
+            nodes = ClusterManager._get_node_info_with_retry()
+            log.debug("Nodes: %s", nodes)
+            for node in nodes:
+                if not node.partitions or any(p not in partitions for p in node.partitions):
+                    # ignore nodes not belonging to any partition
+                    ignored_nodes.append(node)
+                elif any("INACTIVE" not in partitions[p].state for p in node.partitions):
+                    active_nodes.append(node)
+                else:
+                    inactive_nodes.append(node)
 
+            if ignored_nodes:
+                log.warning("Ignoring following nodes because they do not belong to any partition: %s", ignored_nodes)
             return active_nodes, inactive_nodes
         except Exception as e:
             log.error("Failed when getting partition/node states from scheduler with exception %s", e)
diff --git a/tests/common/schedulers/test_slurm_commands.py b/tests/common/schedulers/test_slurm_commands.py
@@ -87,29 +87,62 @@ def test_is_static_node(nodename, expected_is_static):
                 "172.31.10.155\n"
                 "172-31-10-155\n"
                 "MIXED+CLOUD\n"
+                "multiple\n"
+                "---\n"
                 "multiple-dy-c5xlarge-2\n"
                 "172.31.7.218\n"
                 "172-31-7-218\n"
                 "IDLE+CLOUD+POWER\n"
+                "multiple\n"
+                "---\n"
                 "multiple-dy-c5xlarge-3\n"
                 "multiple-dy-c5xlarge-3\n"
                 "multiple-dy-c5xlarge-3\n"
-                "IDLE+CLOUD+POWER"
+                "IDLE+CLOUD+POWER\n"
+                "multiple\n"
+                "---\n"
+                "multiple-dy-c5xlarge-4\n"
+                "multiple-dy-c5xlarge-4\n"
+                "multiple-dy-c5xlarge-4\n"
+                "IDLE+CLOUD+POWER\n"
+                "multiple,multiple2\n"
+                "---\n"
+                "multiple-dy-c5xlarge-5\n"
+                "multiple-dy-c5xlarge-5\n"
+                "multiple-dy-c5xlarge-5\n"
+                "IDLE+CLOUD+POWER\n"
+                # missing partitions
+                "---"
             ),
             [
-                SlurmNode("multiple-dy-c5xlarge-1", "172.31.10.155", "172-31-10-155", "MIXED+CLOUD"),
-                SlurmNode("multiple-dy-c5xlarge-2", "172.31.7.218", "172-31-7-218", "IDLE+CLOUD+POWER"),
+                SlurmNode("multiple-dy-c5xlarge-1", "172.31.10.155", "172-31-10-155", "MIXED+CLOUD", "multiple"),
+                SlurmNode("multiple-dy-c5xlarge-2", "172.31.7.218", "172-31-7-218", "IDLE+CLOUD+POWER", "multiple"),
                 SlurmNode(
                     "multiple-dy-c5xlarge-3",
                     "multiple-dy-c5xlarge-3",
                     "multiple-dy-c5xlarge-3",
                     "IDLE+CLOUD+POWER",
+                    "multiple",
+                ),
+                SlurmNode(
+                    "multiple-dy-c5xlarge-4",
+                    "multiple-dy-c5xlarge-4",
+                    "multiple-dy-c5xlarge-4",
+                    "IDLE+CLOUD+POWER",
+                    "multiple,multiple2",
+                ),
+                SlurmNode(
+                    "multiple-dy-c5xlarge-5",
+                    "multiple-dy-c5xlarge-5",
+                    "multiple-dy-c5xlarge-5",
+                    "IDLE+CLOUD+POWER",
+                    None,
                 ),
             ],
         )
     ],
 )
-def test_parse_nodes_info(node_info, expected_parsed_nodes_output, mocker):
+def test_parse_nodes_info(node_info, expected_parsed_nodes_output):
     assert_that(_parse_nodes_info(node_info)).is_equal_to(expected_parsed_nodes_output)
 
 
@@ -426,10 +459,13 @@ def test_update_nodes(batch_node_info, state, reason, raise_on_error, run_comman
 @pytest.mark.parametrize(
     "node, expected_output",
     [
-        (SlurmNode("queue-name-st-t2micro-1", "nodeip", "nodehostname", "somestate"), True),
-        (SlurmNode("queue-name-st-dy-t2micro-1", "nodeip", "nodehostname", "somestate"), False),
-        (SlurmNode("queuename-dy-t2micro-1", "nodeip", "nodehostname", "somestate"), False),
-        (SlurmNode("queuename-dy-dy-dy-st-t2micro-1", "nodeip", "nodehostname", "somestate"), True),
+        (SlurmNode("queue-name-st-t2micro-1", "nodeip", "nodehostname", "somestate", "queue-name"), True),
+        (SlurmNode("queue-name-st-dy-t2micro-1", "nodeip", "nodehostname", "somestate", "queue-name-st"), False),
+        (SlurmNode("queuename-dy-t2micro-1", "nodeip", "nodehostname", "somestate", "queuename"), False),
+        (
+            SlurmNode("queuename-dy-dy-dy-st-t2micro-1", "nodeip", "nodehostname", "somestate", "queuename-dy-dy-dy"),
+            True,
+        ),
     ],
 )
 def test_slurm_node_is_static(node, expected_output):
@@ -439,8 +475,11 @@ def test_slurm_node_is_static(node, expected_output):
 @pytest.mark.parametrize(
     "node, expected_output",
     [
-        (SlurmNode("queue-name-st-t2micro-1", "nodeip", "nodehostname", "somestate"), True),
-        (SlurmNode("queuename-dy-t2micro-1", "queuename-dy-t2micro-1", "nodehostname", "somestate"), False),
+        (SlurmNode("queue-name-st-t2micro-1", "nodeip", "nodehostname", "somestate", "queue-name"), True),
+        (
+            SlurmNode("queuename-dy-t2micro-1", "queuename-dy-t2micro-1", "nodehostname", "somestate", "queuename"),
+            False,
+        ),
     ],
 )
 def test_slurm_node_is_nodeaddr_set(node, expected_output):
@@ -450,12 +489,12 @@ def test_slurm_node_is_nodeaddr_set(node, expected_output):
 @pytest.mark.parametrize(
     "node, expected_output",
     [
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "somestate"), False),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "MIXED#+CLOUD+DRAIN"), True),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "ALLOCATED*+CLOUD+DRAIN"), True),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "IDLE+CLOUD"), False),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "DOWN+CLOUD"), False),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "COMPLETING+DRAIN"), True),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "somestate", "queue1"), False),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "MIXED#+CLOUD+DRAIN", "queue1"), True),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "ALLOCATED*+CLOUD+DRAIN", "queue1"), True),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "IDLE+CLOUD", "queue1"), False),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "DOWN+CLOUD", "queue1"), False),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "COMPLETING+DRAIN", "queue1"), True),
     ],
 )
 def test_slurm_node_has_job(node, expected_output):
@@ -465,11 +504,11 @@ def test_slurm_node_has_job(node, expected_output):
 @pytest.mark.parametrize(
     "node, expected_output",
     [
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "somestate"), False),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "MIXED#+CLOUD+DRAIN"), False),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "ALLOCATED*+CLOUD+DRAIN"), False),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "IDLE*+CLOUD+DRAIN"), True),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "DOWN+CLOUD+DRAIN"), True),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "somestate", "queue1"), False),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "MIXED#+CLOUD+DRAIN", "queue1"), False),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "ALLOCATED*+CLOUD+DRAIN", "queue1"), False),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "IDLE*+CLOUD+DRAIN", "queue1"), True),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "DOWN+CLOUD+DRAIN", "queue1"), True),
     ],
 )
 def test_slurm_node_is_drained(node, expected_output):
@@ -479,12 +518,12 @@ def test_slurm_node_is_drained(node, expected_output):
 @pytest.mark.parametrize(
     "node, expected_output",
     [
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "somestate"), False),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "MIXED#+CLOUD+DOWN"), True),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "ALLOCATED*+CLOUD+DRAIN"), False),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "DOWN*+CLOUD"), True),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "DOWN+CLOUD+POWER"), True),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "IDLE~+CLOUD+POWERING_DOWN"), False),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "somestate", "queue1"), False),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "MIXED#+CLOUD+DOWN", "queue1"), True),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "ALLOCATED*+CLOUD+DRAIN", "queue1"), False),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "DOWN*+CLOUD", "queue1"), True),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "DOWN+CLOUD+POWER", "queue1"), True),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "IDLE~+CLOUD+POWERING_DOWN", "queue1"), False),
     ],
 )
 def test_slurm_node_is_down(node, expected_output):
@@ -494,11 +533,11 @@ def test_slurm_node_is_down(node, expected_output):
 @pytest.mark.parametrize(
     "node, expected_output",
     [
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "IDLE+CLOUD+POWER"), True),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "MIXED#+CLOUD+DRAIN"), False),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "ALLOCATED*+CLOUD+DOWN"), False),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "IDLE+CLOUD+POWERING_DOWN"), False),
-        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "IDLE#+CLOUD"), True),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "IDLE+CLOUD+POWER", "queue1"), True),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "MIXED#+CLOUD+DRAIN", "queue1"), False),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "ALLOCATED*+CLOUD+DOWN", "queue1"), False),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "IDLE+CLOUD+POWERING_DOWN", "queue1"), False),
+        (SlurmNode("queue1-st-c5xlarge-1", "nodeip", "nodehostname", "IDLE#+CLOUD", "queue1"), True),
     ],
 )
 def test_slurm_node_is_up(node, expected_output):
diff --git a/tests/slurm_plugin/test_clustermgtd.py b/tests/slurm_plugin/test_clustermgtd.py
diff --git a/tests/slurm_plugin/test_computemgtd.py b/tests/slurm_plugin/test_computemgtd.py
