# Copyright 2024 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """File for checking cluster status during CI/CD. Intended to be run after the health runner has started, it repeatedly pings the cluster to check if labels look as expected. Will error if any nodes do not reach expected state before timeout. See parse_args() for information about expected calling format. """ from __future__ import annotations import abc import argparse from collections.abc import Iterable import re import sys import time from typing import Any import kubernetes from kubernetes.client.models import V1Node TIMEOUT_SECS: int = 60 * 15 POLL_INTERVAL_SECS: int = 60 V1API: kubernetes.client.CoreV1Api | None = None def list_nodes() -> list[V1Node]: global V1API if V1API is None: V1API = kubernetes.client.CoreV1Api() return V1API.list_node().items class NodePredicate(abc.ABC): """Checks that can be applied to a node.""" def __init__(self, title: str): self.title: str = title @abc.abstractmethod def __call__(self, node: V1Node) -> bool: ... class LabelRegexMatchPredicate(NodePredicate): """Checks if a label is a) present and b) matches the specified regex. returns False if either is validated """ title_template: str = "{label} = {regex}" def __init__(self, label: str, regex: str) -> None: super().__init__(self.title_template.format(label=label, regex=regex)) self.label: str = label self.regex: str = regex def __call__(self, node: V1Node) -> bool: label_value: str | None = node.metadata.labels.get(self.label) if label_value is None: return False return re.match(self.regex, label_value) is not None class ClusterRenderer(abc.ABC): """Classes for pretty-printing the status of nodes for inspection in logs.""" @abc.abstractmethod def __call__(self, nodes: Iterable[V1Node]) -> str: ... class ReprRenderer(ClusterRenderer): """Very raw dumping of each node.""" def __call__(self, nodes: Iterable[V1Node]) -> str: return "\n\n".join(map(repr, nodes)) class PredicateTableRenderer(ClusterRenderer): """Pretty prints the status of the checks in a table. A cell is True if the node associated with its row passes the check associated with its column. """ def __init__(self, predicates: Iterable[NodePredicate]) -> None: super().__init__() self.predicates: list[NodePredicate] = list(predicates) self.header_row = ["Name"] + [pred.title for pred in self.predicates] def __call__(self, nodes: Iterable[V1Node]) -> str: cells: list[list[bool]] = self._predicate_value_table(nodes) return self._render_table([self.header_row] + cells) def _predicate_value_table(self, nodes: Iterable[V1Node]) -> list[list[bool]]: return [ [node.metadata.name] + [pred(node) for pred in self.predicates] for node in nodes ] def _render_table(self, cells: list[list[Any]]) -> str: # Could just use pandas here, but choosing to avoid unnecessary deps # Convert all elements to strings str_data: list[list[str]] = [[str(item) for item in row] for row in cells] # Find the maximum width for each column num_cols: int = len(str_data[0]) col_widths: list[int] = [ max(len(row[i]) for row in str_data) for i in range(num_cols) ] result: str = "" # Pad each element in each column for row in str_data: for i, cell in enumerate(row): result += cell.ljust(col_widths[i]) + (" " * 5) result += "\n" return result class ClusterStatusChecker: """Main workhorse class, applies checks to a filtered set of nodes & prints their status.""" def __init__( self, check_predicates: Iterable[NodePredicate], filter_predicates: Iterable[NodePredicate] | None, renderers: Iterable[ClusterRenderer] | None, ): self.predicates_to_check: list[NodePredicate] = list(check_predicates) self.predicates_to_filter_for: list[NodePredicate] = ( list(filter_predicates) if filter_predicates else [] ) self.renderers: list[ClusterRenderer] = list(renderers) if renderers else [] def check_cluster_status(self, node_names: set[str]) -> bool: nodes: list[V1Node] = self._get_nodes(node_names) self._log_nodes(nodes) for node in nodes: for pred in self.predicates_to_check: if not pred(node): return False return True def check_cluster_status_with_retry( self, timeout_secs: int = TIMEOUT_SECS, poll_interval_secs: int = POLL_INTERVAL_SECS, ) -> bool: start_time: float = time.monotonic() nodes_to_check: set[str] = self._node_names_passing_filters() while time.monotonic() - start_time < timeout_secs: if self.check_cluster_status(nodes_to_check): return True time.sleep(poll_interval_secs) return False def _get_nodes(self, node_names: set[str]) -> list[V1Node]: nodes: Iterable[V1Node] = list_nodes() result: list[V1Node] = [ node for node in nodes if node.metadata.name in node_names ] if len(result) != len(node_names): missingno = set(node_names) - set([node.metadata.name for node in result]) print(f"Information missing for {missingno} number of nodes") return result def _log_nodes(self, nodes: Iterable[V1Node]) -> None: nodes = list(nodes) for renderer in self.renderers: print(renderer(nodes)) def _node_names_passing_filters(self) -> set[str]: nodes: Iterable[V1Node] = list_nodes() nodes_passing_filter: list[V1Node] = [ node for node in nodes if all([pred(node) for pred in self.predicates_to_filter_for]) ] if not nodes_passing_filter: print("No nodes passing filters found") sys.exit(1) return set([node.metadata.name for node in nodes_passing_filter]) def parse_dict_args_to_predicates( arg_strs: Iterable[str], ) -> list[LabelRegexMatchPredicate]: result: list[LabelRegexMatchPredicate] = [] for arg_str in arg_strs: label, regex = arg_str.split("=") if not label or not regex: print(f"Unable to parse filter argument: {arg_str}") sys.exit(1) result.append(LabelRegexMatchPredicate(label=label, regex=regex)) return result def parse_args(): parser = argparse.ArgumentParser() parser.add_argument( "--check", required=True, type=str, nargs="+", help=f"""The test will check if these labels match their expected format before {TIMEOUT_SECS} seconds. Each check should be formatted as label_to_check=regex_to_match "e.g. aiinfra/nccl-healthcheck-result=^pass$""", ) parser.add_argument( "--filter", required=False, type=str, nargs="*", default=[], help="""Each filter argument should be of the form 'label_name=acceptable_value_regex'. Nodes will be ignored from the test if they either a) are missing one of the filter labels or b) have values that do not match the corresponding regex. All filter labels are checked at the BEGINNING of the test. This allows the test to run on a subset of the nodes""", ) # Filter is added b/c gpu_healthchecks only run on a3-megagpu nodes args: argparse.Namespace = parser.parse_args() filter_predicates: list[LabelRegexMatchPredicate] = ( parse_dict_args_to_predicates(args.filter) ) check_predicates: list[LabelRegexMatchPredicate] = ( parse_dict_args_to_predicates(args.check) ) return ClusterStatusChecker( check_predicates=check_predicates, filter_predicates=filter_predicates, renderers=[ReprRenderer(), PredicateTableRenderer(check_predicates)], ) if __name__ == "__main__": kubernetes.config.load_kube_config() checker: ClusterStatusChecker = parse_args() if not checker.check_cluster_status_with_retry(): print("ERROR: All nodes did not pass before time out.") sys.exit(1)