scripts: Add code for waiting until the restarted job is ready and the metric condition is met

Author: guyf-starkware

scripts/prod/update_config_and_restart_nodes_lib.py

@@ -222,10 +222,10 @@ class MetricConditionGater:
     def __init__(
         self,
         metric_name: str,
+        namespace: str,
+        cluster: Optional[str],
         pod: str,
         metrics_port: int,
-        namespace: str,
-        cluster: Optional[str] = None,
         refresh_interval_seconds: int = 3,
     ):
         self.metric_name = metric_name
@@ -263,7 +263,7 @@ def _get_metrics(self) -> str:
             )
             sleep(self.refresh_interval_seconds)
 
-    def _poll_condition_met(self, metric_value_condition: Callable[[Any], bool]):
+    def _poll_until_condition_met(self, metric_value_condition: Callable[[Any], bool]):
         """Poll metrics until the condition is met for the metric."""
         while True:
             metrics = self._get_metrics()
@@ -401,6 +401,20 @@ def get_context_list_from_args(
         ]
 
 
+def _get_pod_names(
+    namespace: str, service: Service, index: int, cluster: Optional[str] = None
+) -> list[str]:
+    kubectl_args = [
+        "get",
+        "pods",
+        "-o",
+        "name",
+    ]
+    kubectl_args.extend(get_namespace_args(namespace, cluster))
+    pods = run_kubectl_command(kubectl_args, capture_output=True).stdout.splitlines()
+    return [pod.split("/")[1] for pod in pods if pod.startswith(f"pod/{service.pod_name}")]
+
+
 class ServiceRestarter(ABC):
     """Abstract class for restarting service instances."""
 
@@ -417,18 +431,7 @@ def _restart_pod(
         namespace: str, service: Service, index: int, cluster: Optional[str] = None
     ) -> None:
         """Restart pod by deleting it"""
-        # Get the list of pods (one string per line).
-        kubectl_args = [
-            "get",
-            "pods",
-            "-o",
-            "name",
-        ]
-        kubectl_args.extend(get_namespace_args(namespace, cluster))
-        pods = run_kubectl_command(kubectl_args, capture_output=True).stdout.splitlines()
-
-        # Filter the list of pods to only include the ones that match the service and extract the pod name.
-        pods = [pod.split("/")[1] for pod in pods if pod.startswith(f"pod/{service.pod_name}")]
+        pods = _get_pod_names(namespace, service, index, cluster)
 
         if not pods:
             print_error(f"Could not find pods for service {service.pod_name}.")
@@ -514,6 +517,66 @@ def restart_service(self, instance_index: int) -> bool:
         return self.check_between_restarts(instance_index)
 
 
+class WaitOnMetrticOneByOneRestarter(ChecksBetweenRestarts):
+    def __init__(
+        self,
+        namespace_and_instruction_args: NamespaceAndInstructionArgs,
+        service: Service,
+        metric_name: str,
+        metrics_port: int,
+        metric_value_condition: Callable[[Any], bool],
+    ):
+        def _check_between_restarts(instance_index: int) -> bool:
+            if not WaitOnMetrticOneByOneRestarter._wait_for_pods_to_be_ready(service):
+                return False
+            metric_condition_gater = MetricConditionGater(
+                metric_name,
+                namespace_and_instruction_args.get_namespace(instance_index),
+                namespace_and_instruction_args.get_cluster(instance_index),
+                service.pod_name,
+                metrics_port,
+            )
+            metric_condition_gater.gate(metric_value_condition)
+            if instance_index == namespace_and_instruction_args.size() - 1:
+                return True
+            return wait_until_y_or_n(f"Do you want to restart the next pod?")
+
+        super().__init__(namespace_and_instruction_args, service, _check_between_restarts)
+
+    @staticmethod
+    def _wait_for_pods_to_be_ready(
+        service: Service, wait_timeout: int = 180, num_retry: int = 3, refresh_delay_sec: int = 3
+    ) -> bool:
+        for i in range(num_retry):
+            pods = _get_pod_names(service)
+            if pods:
+                for pod in pods:
+                    print_colored(
+                        f"Waiting for pod {pod} to be ready... (timeout: {wait_timeout}s)"
+                    )
+                    result = run_kubectl_command(
+                        ["wait", "--for=condition=ready", "pod/", pod, "--timeout", "{timeout}s"],
+                        capture_output=False,
+                    )
+
+                    if result.returncode != 0:
+                        print_colored(
+                            f"Timed out waiting for pod {pod} to be ready: {result.stderr}, retrying... (attempt {i + 1}/{num_retry})",
+                            Colors.YELLOW,
+                        )
+                        break
+                    return True
+            else:
+                print_colored(
+                    f"Could not get pod names for service {service.pod_name}, retrying... (attempt {i + 1}/{num_retry})",
+                    Colors.YELLOW,
+                )
+            sleep(refresh_delay_sec)
+
+        print_error(f"Pods for service {service.pod_name} are not ready after {num_retry} attempts")
+        return False
+
+
 class NoOpServiceRestarter(ServiceRestarter):
     """No-op service restarter."""