[Develop] Backing instance max count updates - `_is_node_in_replaceme…

…nt_valid()`, `is_backing_instance_valid()`, and unit tests (#622)

* Add logic to `is_backing_instance_valid()` to check the IP to make sure the instance matches what is being tracked (#618)

The missing instance map did not track what the IP address was that was associated with the slurm node.
Because of this if a new instance is launched before an instance becomes healthy, the increment is not reset
for the instance count map.  This change uses a class object to track the data and links the node name to the ip.

Also use the `is_backing_instance_valid()` function in `is_state_healthy()` instead of the plain `node.instance`
object check to allow for the delay in EC2 consistency.

* Refactor logic in `_is_node_in_replacement_valid()` to account for `node.instance` being `None` (#620)

* Add unit tests to cover max_count > 0 in _is_node_in_replacement_valid

src/slurm_plugin/clustermgtd.py

@@ -1200,17 +1200,31 @@ def _find_active_nodes(partitions_name_map):
  active_nodes += partition.slurm_nodes
  return list(dict.fromkeys(active_nodes))
 
- def _is_node_in_replacement_valid(self, node, check_node_is_valid):
+ def _is_node_in_replacement_valid(self, node: SlurmNode, check_node_is_valid):
  """
  Check node is replacement timeout or in replacement.
 
  If check_node_is_valid=True, check whether a node is in replacement,
  If check_node_is_valid=False, check whether a node is replacement timeout.
  """
- if node.instance and node.name in self._static_nodes_in_replacement:
- time_is_expired = time_is_up(
- node.instance.launch_time, self._current_time, grace_time=self._config.node_replacement_timeout
+ log.debug(f"Checking if node is in replacement {node}")
+ if (
+ node.is_backing_instance_valid(
+ self._config.ec2_instance_missing_max_count,
+ self._nodes_without_backing_instance_count_map,
+ log_warn_if_unhealthy=True,
+ )
+ and node.name in self._static_nodes_in_replacement
+ ):
+ # Set `time_is_expired` to `False` if `node.instance` is `None` since we don't have a launch time yet
+ time_is_expired = (
+ False
+ if not node.instance
+ else time_is_up(
+ node.instance.launch_time, self._current_time, grace_time=self._config.node_replacement_timeout
+ )
  )
+ log.debug(f"Node {node} is in replacement and timer expired? {time_is_expired}, instance? {node.instance}")
  return not time_is_expired if check_node_is_valid else time_is_expired
  return False
 

src/slurm_plugin/slurm_resources.py

@@ -180,6 +180,17 @@ class SlurmReservation:
  users: str
 
 
+class MissingInstance:
+ name: str
+ ip: str
+ count: int
+
+ def __init__(self, name, ip, count):
+ self.name = name
+ self.ip = ip
+ self.count = count
+
+
 class SlurmNode(metaclass=ABCMeta):
  SLURM_SCONTROL_COMPLETING_STATE = "COMPLETING"
  SLURM_SCONTROL_BUSY_STATES = {"MIXED", "ALLOCATED", SLURM_SCONTROL_COMPLETING_STATE}
@@ -427,7 +438,7 @@ def is_powering_down_with_nodeaddr(self):
  def is_backing_instance_valid(
  self,
  ec2_instance_missing_max_count,
- nodes_without_backing_instance_count_map: dict,
+ nodes_without_backing_instance_count_map: dict[str, MissingInstance],
  log_warn_if_unhealthy=True,
  ):
  """Check if a slurm node's addr is set, it points to a valid instance in EC2."""
@@ -445,7 +456,11 @@ def is_backing_instance_valid(
  )
  # Allow a few iterations for the eventual consistency of EC2 data
  logger.debug(f"Map of slurm nodes without backing instances {nodes_without_backing_instance_count_map}")
- missing_instance_loop_count = nodes_without_backing_instance_count_map.get(self.name, 0)
+ missing_instance = nodes_without_backing_instance_count_map.get(self.name, None)
+ missing_instance_loop_count = missing_instance.count if missing_instance else 0
+ if missing_instance and self.nodeaddr != missing_instance.ip:
+ # Reset the loop count since the nodeaddr has changed
+ missing_instance_loop_count = 0
  # If the loop count has been reached, the instance is unhealthy and will be terminated
  if missing_instance_loop_count >= ec2_instance_missing_max_count:
  if log_warn_if_unhealthy:
@@ -454,11 +469,12 @@ def is_backing_instance_valid(
  nodes_without_backing_instance_count_map.pop(self.name, None)
  self.ec2_backing_instance_valid = False
  else:
- nodes_without_backing_instance_count_map[self.name] = missing_instance_loop_count + 1
+ instance_to_add = MissingInstance(self.name, self.nodeaddr, missing_instance_loop_count + 1)
+ nodes_without_backing_instance_count_map[self.name] = instance_to_add
  if log_warn_if_unhealthy:
  logger.warning(
  f"Incrementing missing EC2 instance count for node {self.name} to "
- f"{nodes_without_backing_instance_count_map[self.name]}."
+ f"{nodes_without_backing_instance_count_map[self.name].count}."
  )
  else:
  # Remove the slurm node from the map since the instance is healthy

tests/slurm_plugin/slurm_resources/test_slurm_resources.py

@@ -18,6 +18,7 @@
  DynamicNode,
  EC2InstanceHealthState,
  InvalidNodenameError,
+ MissingInstance,
  SlurmPartition,
  SlurmResumeJob,
  StaticNode,
@@ -1185,26 +1186,34 @@ def test_slurm_node_is_powering_down_with_nodeaddr(node, expected_result):
  StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "IDLE+CLOUD+POWER", "queue1"),
  None,
  2,
- {"queue1-st-c5xlarge-1": 1},
- {"queue1-st-c5xlarge-1": 2},
+ {"queue1-st-c5xlarge-1": MissingInstance("queue1-st-c5xlarge-1", "ip-1", 1)},
+ {"queue1-st-c5xlarge-1": MissingInstance("queue1-st-c5xlarge-1", "ip-1", 2)},
  True,
  ),
  (
  StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "IDLE+CLOUD+POWER", "queue1"),
  None,
  2,
- {"queue1-st-c5xlarge-1": 2},
- {"queue1-st-c5xlarge-1": 2},
+ {"queue1-st-c5xlarge-1": MissingInstance("queue1-st-c5xlarge-1", "ip-1", 2)},
+ {"queue1-st-c5xlarge-1": MissingInstance("queue1-st-c5xlarge-1", "ip-1", 2)},
  False,
  ),
  (
  StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "IDLE+CLOUD+POWER", "queue1"),
  "Instance",
  2,
- {"queue1-st-c5xlarge-1": 3},
+ {"queue1-st-c5xlarge-1": MissingInstance("queue1-st-c5xlarge-1", "ip-1", 3)},
  {},
  True,
  ),
+ (
+ StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "IDLE+CLOUD+POWER", "queue1"),
+ "Instance",
+ 3,
+ {"queue1-st-c5xlarge-1": MissingInstance("queue1-st-c5xlarge-1", "ip-2", 2)},
+ {"queue1-st-c5xlarge-1": MissingInstance("queue1-st-c5xlarge-1", "ip-1", 1)},
+ True,
+ ),
  ],
  ids=[
  "static_no_backing_zero_max_count",
@@ -1214,6 +1223,7 @@ def test_slurm_node_is_powering_down_with_nodeaddr(node, expected_result):
  "static_no_backing_count_not_exceeded",
  "static_no_backing_with_count_exceeded",
  "static_backed_with_count_exceeded",
+ "static_no_backing_count_not_exceeded_with_wrong_ip",
  ],
 )
 def test_slurm_node_is_backing_instance_valid(node, instance, max_count, count_map, final_map, expected_result):
@@ -1225,7 +1235,9 @@ def test_slurm_node_is_backing_instance_valid(node, instance, max_count, count_m
  ).is_equal_to(expected_result)
  assert_that(node.ec2_backing_instance_valid).is_equal_to(expected_result)
  if count_map:
- assert_that(count_map[node.name]).is_equal_to(final_map.get(node.name, None))
+ assert_that(count_map[node.name].count).is_equal_to(final_map.get(node.name, None).count)
+ assert_that(count_map[node.name].ip).is_equal_to(final_map.get(node.name, None).ip)
+ assert_that(count_map[node.name].ip).is_equal_to(node.nodeaddr)
 
 
 @pytest.mark.parametrize(

tests/slurm_plugin/test_clustermgtd.py

@@ -767,6 +767,7 @@ def test_handle_health_check(
  region="region",
  boto3_config=None,
  fleet_config={},
+ ec2_instance_missing_max_count=0,
  )
 
  cluster_manager = ClusterManager(mock_sync_config)
@@ -831,6 +832,7 @@ def test_update_static_nodes_in_replacement(current_replacing_nodes, slurm_nodes
  region="region",
  boto3_config=None,
  fleet_config={},
+ ec2_instance_missing_max_count=0,
  )
  cluster_manager = ClusterManager(mock_sync_config)
  cluster_manager._static_nodes_in_replacement = current_replacing_nodes
@@ -2646,43 +2648,61 @@ def initialize_console_logger_mock(mocker):
 
 
 @pytest.mark.parametrize(
- "current_replacing_nodes, node, instance, current_time, expected_result",
+ "current_replacing_nodes, node, instance, current_time, max_count, expected_result",
  [
  (
  set(),
  StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "IDLE+CLOUD", "queue1"),
  EC2Instance("id-1", "ip-1", "hostname", {"ip-1"}, datetime(2020, 1, 1, 0, 0, 0)),
  datetime(2020, 1, 1, 0, 0, 29),
+ 0,
  False,
  ),
  (
  {"queue1-st-c5xlarge-1"},
  StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "IDLE+CLOUD", "queue1"),
  None,
  datetime(2020, 1, 1, 0, 0, 29),
+ 0,
  False,
  ),
  (
  {"queue1-st-c5xlarge-1"},
  StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "DOWN+CLOUD", "queue1"),
  EC2Instance("id-1", "ip-1", "hostname", {"ip-1"}, datetime(2020, 1, 1, 0, 0, 0)),
  datetime(2020, 1, 1, 0, 0, 29),
+ 0,
  True,
  ),
  (
  {"queue1-st-c5xlarge-1"},
  StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "IDLE+CLOUD", "queue1"),
  EC2Instance("id-1", "ip-1", "hostname", {"ip-1"}, datetime(2020, 1, 1, 0, 0, 0)),
  datetime(2020, 1, 1, 0, 0, 30),
+ 0,
  False,
  ),
+ (
+ {"queue1-st-c5xlarge-1"},
+ StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "IDLE+CLOUD", "queue1"),
+ None,
+ datetime(2020, 1, 1, 0, 0, 30),
+ 1,
+ True,
+ ),
+ ],
+ ids=[
+ "not_in_replacement",
+ "no-backing-instance",
+ "in_replacement",
+ "timeout",
+ "no-backing-instance-with-max-count",
  ],
- ids=["not_in_replacement", "no-backing-instance", "in_replacement", "timeout"],
 )
 @pytest.mark.usefixtures(
  "initialize_instance_manager_mock", "initialize_executor_mock", "initialize_console_logger_mock"
 )
-def test_is_node_being_replaced(current_replacing_nodes, node, instance, current_time, expected_result):
+def test_is_node_being_replaced(current_replacing_nodes, node, instance, current_time, max_count, expected_result):
  mock_sync_config = SimpleNamespace(
  node_replacement_timeout=30,
  insufficient_capacity_timeout=3,
@@ -2691,6 +2711,7 @@ def test_is_node_being_replaced(current_replacing_nodes, node, instance, current
  region="region",
  boto3_config=None,
  fleet_config={},
+ ec2_instance_missing_max_count=max_count,
  )
  cluster_manager = ClusterManager(mock_sync_config)
  cluster_manager._current_time = current_time
@@ -2700,24 +2721,34 @@ def test_is_node_being_replaced(current_replacing_nodes, node, instance, current
 
 
 @pytest.mark.parametrize(
- "node, instance, current_node_in_replacement, is_replacement_timeout",
+ "node, instance, current_node_in_replacement, max_count, is_replacement_timeout",
  [
  (
  StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "DOWN+CLOUD+NOT_RESPONDING", "queue1"),
  None,
  {"queue1-st-c5xlarge-1"},
+ 0,
+ False,
+ ),
+ (
+ StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "DOWN+CLOUD+NOT_RESPONDING", "queue1"),
+ None,
+ {"queue1-st-c5xlarge-1"},
+ 1,
  False,
  ),
  (
  StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "DOWN+CLOUD+NOT_RESPONDING", "queue1"),
  EC2Instance("id-1", "ip-1", "hostname", {"ip-1"}, datetime(2020, 1, 1, 0, 0, 0)),
  {"queue1-st-c5xlarge-1"},
+ 0,
  True,
  ),
  (
  DynamicNode("queue1-dy-c5xlarge-1", "ip-1", "hostname", "MIXED+CLOUD+NOT_RESPONDING+POWERING_UP", "queue1"),
  None,
  {"some_node_in_replacement"},
+ 0,
  False,
  ),
  (
@@ -2730,20 +2761,22 @@ def test_is_node_being_replaced(current_replacing_nodes, node, instance, current
  ),
  EC2Instance("id-1", "ip-1", "hostname", {"ip-1"}, datetime(2020, 1, 1, 0, 0, 0)),
  {"some_node_in_replacement"},
+ 0,
  False,
  ),
  (
  StaticNode("queue1-st-c5xlarge-1", "ip-1", "hostname", "DOWN+CLOUD+NOT_RESPONDING", "queue1"),
  EC2Instance("id-1", "ip-1", "hostname", {"ip-1"}, datetime(2020, 1, 1, 0, 0, 0)),
  {"some_node_in_replacement"},
+ 0,
  False,
  ),
  ],
 )
 @pytest.mark.usefixtures(
  "initialize_instance_manager_mock", "initialize_executor_mock", "initialize_console_logger_mock"
 )
-def test_is_node_replacement_timeout(node, current_node_in_replacement, is_replacement_timeout, instance):
+def test_is_node_replacement_timeout(node, current_node_in_replacement, max_count, is_replacement_timeout, instance):
  node.instance = instance
  mock_sync_config = SimpleNamespace(
  node_replacement_timeout=30,
@@ -2753,6 +2786,7 @@ def test_is_node_replacement_timeout(node, current_node_in_replacement, is_repla
  region="region",
  boto3_config=None,
  fleet_config={},
+ ec2_instance_missing_max_count=0,
  )
  cluster_manager = ClusterManager(mock_sync_config)
  cluster_manager._current_time = datetime(2020, 1, 2, 0, 0, 0)