Set failed nodes down with reason to Errorcode

* When setting a dynamic node to down state, set the reason to Error code if the node fail with ClientError.
* If the error type is general Exception, set error code to Expcetion, if the exception is Client error, set errorcode to be equal to ClientError error code.
* Insufficient capacity error code: {"InsufficientInstanceCapacity", "MaxSpotInstanceCountExceeded", "InsufficientHostCapacity", "InsufficientReservedInstanceCapacity"}

When setting a dynamic node to down state, set the reason to Error code if the node fail with insufficient capacity.
Signed-off-by: chenwany <[email protected]>

Author: chenwany

src/slurm_plugin/clustermgtd.py

@@ -676,7 +676,8 @@ def _handle_unhealthy_static_nodes(self, unhealthy_static_nodes):
             node_list, self._config.launch_max_batch_size, self._config.update_node_address
         )
         # Add launched nodes to list of nodes being replaced, excluding any nodes that failed to launch
-        launched_nodes = set(node_list) - set(self._instance_manager.failed_nodes)
+        failed_nodes = set().union(*self._instance_manager.failed_nodes.values())
+        launched_nodes = set(node_list) - failed_nodes
         self._static_nodes_in_replacement |= launched_nodes
         log.info(
             "After node maintenance, following nodes are currently in replacement: %s",

src/slurm_plugin/instance_manager.py

@@ -61,7 +61,7 @@ def __init__(
         self._region = region
         self._cluster_name = cluster_name
         self._boto3_config = boto3_config
-        self.failed_nodes = []
+        self.failed_nodes = {}
         self._ddb_resource = boto3.resource("dynamodb", region_name=region, config=boto3_config)
         self._table = self._ddb_resource.Table(table_name) if table_name else None
         self._hosted_zone = hosted_zone
@@ -74,7 +74,7 @@ def __init__(
 
     def _clear_failed_nodes(self):
         """Clear and reset failed nodes list."""
-        self.failed_nodes = []
+        self.failed_nodes = {}
 
     def add_instances_for_nodes(
         self, node_list, launch_batch_size, update_node_address=True, all_or_nothing_batch=False
@@ -102,14 +102,22 @@ def add_instances_for_nodes(
                                 self._store_assigned_hostnames(assigned_nodes)
                                 self._update_dns_hostnames(assigned_nodes)
                             except Exception:
-                                self.failed_nodes.extend(list(assigned_nodes.keys()))
+                                self._update_failed_nodes(set(assigned_nodes.keys()))
+                    except ClientError as e:
+                        logger.error(
+                            "Encountered exception when launching instances for nodes %s: %s",
+                            print_with_count(batch_nodes),
+                            e,
+                        )
+                        error_code = e.response.get("Error", {}).get("Code")
+                        self._update_failed_nodes(set(batch_nodes), error_code)
                     except Exception as e:
                         logger.error(
                             "Encountered exception when launching instances for nodes %s: %s",
                             print_with_count(batch_nodes),
                             e,
                         )
-                        self.failed_nodes.extend(batch_nodes)
+                        self._update_failed_nodes(set(batch_nodes))
 
     def _update_slurm_node_addrs(self, slurm_nodes, launched_instances):
         """Update node information in slurm with info from launched EC2 instance."""
@@ -139,7 +147,7 @@ def _update_slurm_node_addrs(self, slurm_nodes, launched_instances):
                 )
             if fail_launch_nodes:
                 logger.info("Failed to launch instances for following nodes: %s", print_with_count(fail_launch_nodes))
-                self.failed_nodes.extend(fail_launch_nodes)
+                self._update_failed_nodes(set(fail_launch_nodes), "LimitedInstanceCapacity")
 
             return dict(zip(launched_nodes, launched_instances))
 
@@ -149,7 +157,7 @@ def _update_slurm_node_addrs(self, slurm_nodes, launched_instances):
                 print_with_count(slurm_nodes),
                 print_with_count(launched_instances),
             )
-            self.failed_nodes.extend(slurm_nodes)
+            self._update_failed_nodes(set(slurm_nodes))
 
     @log_exception(logger, "saving assigned hostnames in DynamoDB", raise_on_error=True)
     def _store_assigned_hostnames(self, nodes):
@@ -228,7 +236,7 @@ def _parse_requested_instances(self, node_list):
                 instances_to_launch[queue_name][compute_resource_name].append(node)
             except (InvalidNodenameError, KeyError):
                 logger.warning("Discarding NodeName with invalid format: %s", node)
-                self.failed_nodes.append(node)
+                self._update_failed_nodes({node}, "InvalidNodenameError")
         logger.debug("Launch configuration requested by nodes = %s", instances_to_launch)
 
         return instances_to_launch
@@ -362,3 +370,7 @@ def terminate_all_compute_nodes(self, terminate_batch_size):
         except Exception as e:
             logger.error("Failed when terminating compute fleet with error %s", e)
             return False
+
+    def _update_failed_nodes(self, nodeset, error_code="Exception"):
+        """Update failed nodes dict with error code as key and nodeset value."""
+        self.failed_nodes[error_code] = self.failed_nodes.get(error_code, set()).union(nodeset)

src/slurm_plugin/resume.py

@@ -116,7 +116,7 @@ def _get_config(self, config_file_path):
         log.info(self.__repr__())
 
 
-def _handle_failed_nodes(node_list):
+def _handle_failed_nodes(node_list, reason="Failure when resuming nodes"):
     """
     Fall back mechanism to handle failure when launching instances.
 
@@ -132,7 +132,7 @@ def _handle_failed_nodes(node_list):
     """
     try:
         log.info("Setting following failed nodes into DOWN state: %s", print_with_count(node_list))
-        set_nodes_down(node_list, reason="Failure when resuming nodes")
+        set_nodes_down(node_list, reason=reason)
     except Exception as e:
         log.error("Failed to place nodes %s into down with exception: %s", print_with_count(node_list), e)
 
@@ -175,14 +175,17 @@ def _resume(arg_nodes, resume_config):
         update_node_address=resume_config.update_node_address,
         all_or_nothing_batch=resume_config.all_or_nothing_batch,
     )
-    success_nodes = [node for node in node_list if node not in instance_manager.failed_nodes]
+    failed_nodes = set().union(*instance_manager.failed_nodes.values())
+    success_nodes = [node for node in node_list if node not in failed_nodes]
     log.info("Successfully launched nodes %s", print_with_count(success_nodes))
-    if instance_manager.failed_nodes:
+
+    if failed_nodes:
         log.error(
             "Failed to launch following nodes, setting nodes to down: %s",
-            print_with_count(instance_manager.failed_nodes),
+            print_with_count(failed_nodes),
         )
-        _handle_failed_nodes(instance_manager.failed_nodes)
+        for error_code, node_list in instance_manager.failed_nodes.items():
+            _handle_failed_nodes(node_list, reason=f"(Code:{error_code})Failure when resuming nodes")
 
 
 def main():

tests/slurm_plugin/slurm_resources/test_instance_manager.py

@@ -165,7 +165,7 @@ def instance_manager(self, mocker):
                         },
                     ),
                 ],
-                None,
+                {},
                 [
                     call(
                         ["queue1-st-c5xlarge-2"],
@@ -229,6 +229,7 @@ def instance_manager(self, mocker):
                             "LaunchTemplate": {"LaunchTemplateName": "hit-queue1-c5.2xlarge", "Version": "$Latest"},
                         },
                         generate_error=True,
+                        error_code="some_error_code",
                     ),
                     MockedBoto3Request(
                         method="run_instances",
@@ -257,7 +258,7 @@ def instance_manager(self, mocker):
                         },
                     ),
                 ],
-                ["queue1-st-c52xlarge-1"],
+                {"some_error_code": {"queue1-st-c52xlarge-1"}},
                 [
                     call(
                         ["queue1-st-c5xlarge-2"],
@@ -303,7 +304,7 @@ def instance_manager(self, mocker):
                         },
                     ),
                 ],
-                None,
+                {},
                 None,
             ),
             # batch_size1
@@ -353,6 +354,7 @@ def instance_manager(self, mocker):
                             "LaunchTemplate": {"LaunchTemplateName": "hit-queue1-c5.2xlarge", "Version": "$Latest"},
                         },
                         generate_error=True,
+                        error_code="InsufficientHostCapacity",
                     ),
                     MockedBoto3Request(
                         method="run_instances",
@@ -363,14 +365,13 @@ def instance_manager(self, mocker):
                             "LaunchTemplate": {"LaunchTemplateName": "hit-queue2-c5.xlarge", "Version": "$Latest"},
                         },
                         generate_error=True,
+                        error_code="ServiceUnavailable",
                     ),
                 ],
-                [
-                    "queue1-st-c52xlarge-1",
-                    "queue2-st-c5xlarge-1",
-                    "queue2-st-c5xlarge-2",
-                    "queue2-dy-c5xlarge-1",
-                ],
+                {
+                    "InsufficientHostCapacity": {"queue1-st-c52xlarge-1"},
+                    "ServiceUnavailable": {"queue2-st-c5xlarge-1", "queue2-dy-c5xlarge-1", "queue2-st-c5xlarge-2"},
+                },
                 [
                     call(
                         ["queue1-st-c5xlarge-2"],
@@ -425,6 +426,7 @@ def instance_manager(self, mocker):
                             "LaunchTemplate": {"LaunchTemplateName": "hit-queue1-c5.2xlarge", "Version": "$Latest"},
                         },
                         generate_error=True,
+                        error_code="InsufficientVolumeCapacity",
                     ),
                     MockedBoto3Request(
                         method="run_instances",
@@ -454,6 +456,7 @@ def instance_manager(self, mocker):
                             "LaunchTemplate": {"LaunchTemplateName": "hit-queue2-c5.xlarge", "Version": "$Latest"},
                         },
                         generate_error=True,
+                        error_code="InternalError",
                     ),
                     MockedBoto3Request(
                         method="run_instances",
@@ -475,7 +478,7 @@ def instance_manager(self, mocker):
                         },
                     ),
                 ],
-                ["queue1-st-c52xlarge-1", "queue2-st-c5xlarge-2"],
+                {"InsufficientVolumeCapacity": {"queue1-st-c52xlarge-1"}, "InternalError": {"queue2-st-c5xlarge-2"}},
                 [
                     call(
                         ["queue1-st-c5xlarge-2"],
@@ -525,7 +528,7 @@ def instance_manager(self, mocker):
                         "LaunchTemplate": {"LaunchTemplateName": "hit-queue2-c5.xlarge", "Version": "$Latest"},
                     },
                 ),
-                ["queue2-st-c5xlarge-2", "queue2-dy-c5xlarge-1"],
+                {"LimitedInstanceCapacity": {"queue2-st-c5xlarge-2", "queue2-dy-c5xlarge-1"}},
                 [
                     call(
                         ["queue2-st-c5xlarge-1", "queue2-st-c5xlarge-2", "queue2-dy-c5xlarge-1"],
@@ -592,9 +595,10 @@ def instance_manager(self, mocker):
                             "LaunchTemplate": {"LaunchTemplateName": "hit-queue2-c5.xlarge", "Version": "$Latest"},
                         },
                         generate_error=True,
+                        error_code="InsufficientInstanceCapacity",
                     ),
                 ],
-                ["queue2-dy-c5xlarge-2", "queue2-dy-c5xlarge-3"],
+                {"InsufficientInstanceCapacity": {"queue2-dy-c5xlarge-2", "queue2-dy-c5xlarge-3"}},
                 [
                     call(
                         ["queue2-st-c5xlarge-1", "queue2-st-c5xlarge-2", "queue2-dy-c5xlarge-1"],
@@ -716,7 +720,7 @@ def instance_manager(self, mocker):
                         },
                     ),
                 ],
-                None,
+                {},
                 [
                     call(
                         ["queue3-st-c5xlarge-2"],
@@ -941,35 +945,42 @@ def test_launch_ec2_instances(
                 ["queue1-st-c5xlarge-1"],
                 [EC2Instance("id-1", "ip-1", "hostname-1", "some_launch_time")],
                 call(["queue1-st-c5xlarge-1"], nodeaddrs=["ip-1"], nodehostnames=None),
-                [],
+                {},
+                False,
+                "dns.domain",
+            ),
+            (
+                ["queue1-st-c5xlarge-1"],
+                {},
+                None,
+                {"LimitedInstanceCapacity": {"queue1-st-c5xlarge-1"}},
                 False,
                 "dns.domain",
             ),
-            (["queue1-st-c5xlarge-1"], [], None, ["queue1-st-c5xlarge-1"], False, "dns.domain"),
             (
                 ["queue1-st-c5xlarge-1", "queue1-st-c5xlarge-2", "queue1-st-c5xlarge-3", "queue1-st-c5xlarge-4"],
                 [
                     EC2Instance("id-1", "ip-1", "hostname-1", "some_launch_time"),
                     EC2Instance("id-2", "ip-2", "hostname-2", "some_launch_time"),
                 ],
                 call(["queue1-st-c5xlarge-1", "queue1-st-c5xlarge-2"], nodeaddrs=["ip-1", "ip-2"], nodehostnames=None),
-                ["queue1-st-c5xlarge-3", "queue1-st-c5xlarge-4"],
+                {"LimitedInstanceCapacity": {"queue1-st-c5xlarge-4", "queue1-st-c5xlarge-3"}},
                 False,
                 "dns.domain",
             ),
             (
                 ["queue1-st-c5xlarge-1"],
                 [EC2Instance("id-1", "ip-1", "hostname-1", "some_launch_time")],
                 call(["queue1-st-c5xlarge-1"], nodeaddrs=["ip-1"], nodehostnames=["hostname-1"]),
-                [],
+                {},
                 True,
                 "dns.domain",
             ),
             (
                 ["queue1-st-c5xlarge-1"],
                 [EC2Instance("id-1", "ip-1", "hostname-1", "some_launch_time")],
                 call(["queue1-st-c5xlarge-1"], nodeaddrs=["ip-1"], nodehostnames=None),
-                [],
+                {},
                 False,
                 "",
             ),
@@ -1306,12 +1317,14 @@ def test_update_dns_hostnames(
                         "u6tb1metal": ["queue2-st-u6tb1metal-1"],
                     },
                 },
-                [
-                    "in-valid/queue.name-st-c5xlarge-2",
-                    "noBrackets-st-c5xlarge-[1-2]",
-                    "queue2-invalidnodetype-c5xlarge-12",
-                    "queuename-with-dash-and_underscore-st-i3enmetal2tb-1",
-                ],
+                {
+                    "InvalidNodenameError": {
+                        "queue2-invalidnodetype-c5xlarge-12",
+                        "noBrackets-st-c5xlarge-[1-2]",
+                        "queuename-with-dash-and_underscore-st-i3enmetal2tb-1",
+                        "in-valid/queue.name-st-c5xlarge-2",
+                    }
+                },
             ),
         ],
     )