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()
@@ -305,6 +305,18 @@ def gate(self, metric_value_condition: Callable[[Any], bool]):
         cmd.extend(get_namespace_args(self.namespace, self.cluster))
 
         pf_process = None
+
+        def _terminate_port_forward_process(pf_process: subprocess.Popen):
+            if pf_process and pf_process.poll() is None:
+                print_colored(f"Terminating kubectl port-forward process (PID: {pf_process.pid})")
+                pf_process.terminate()
+                try:
+                    pf_process.wait(timeout=5)
+                except subprocess.TimeoutExpired:
+                    print_colored("Force killing kubectl port-forward process")
+                    pf_process.kill()
+                    pf_process.wait()
+
         try:
             pf_process = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
             print("Waiting for forwarding to start")
@@ -315,36 +327,16 @@ def gate(self, metric_value_condition: Callable[[Any], bool]):
 
             # Set up signal handler to ensure forwarding subprocess is terminated on interruption
             def signal_handler(signum, frame):
-                if pf_process and pf_process.poll() is None:
-                    print_colored(
-                        f"Terminating kubectl port-forward process (PID: {pf_process.pid})",
-                        Colors.RED,
-                    )
-                    pf_process.terminate()
-                    try:
-                        pf_process.wait(timeout=5)
-                    except subprocess.TimeoutExpired:
-                        print_colored("Force killing kubectl port-forward process", Colors.RED)
-                        pf_process.kill()
-                        pf_process.wait()
+                _terminate_port_forward_process(pf_process)
                 sys.exit(0)
 
             signal.signal(signal.SIGINT, signal_handler)
             signal.signal(signal.SIGTERM, signal_handler)
 
-            self._poll_condition_met(metric_value_condition)
+            self._poll_until_condition_met(metric_value_condition)
 
         finally:
-            # Ensure subprocess is always terminated
-            if pf_process and pf_process.poll() is None:
-                print_colored(f"Terminating kubectl port-forward process (PID: {pf_process.pid})")
-                pf_process.terminate()
-                try:
-                    pf_process.wait(timeout=5)
-                except subprocess.TimeoutExpired:
-                    print_colored("Force killing kubectl port-forward process")
-                    pf_process.kill()
-                    pf_process.wait()
+            _terminate_port_forward_process(pf_process)
 
 
 class NamespaceAndInstructionArgs:
@@ -409,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."""
 
@@ -425,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}.")
@@ -522,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."""