pyxecm 2.0.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

pyxecm/customizer/k8s.py

@@ -202,7 +202,7 @@ class K8s:
             )
         except ApiException as e:
             if e.status == 404:
-                self.logger.info("Pod -> '%s' not found (may be deleted).", pod_name)
+                self.logger.debug("Pod -> '%s' not found (may be deleted).", pod_name)
                 return None
             else:
                 self.logger.error("Failed to get Pod -> '%s'!", pod_name)
@@ -246,7 +246,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to list Pods with field_selector -> '%s' and label_selector -> '%s'",
+                "Failed to list pods with field selector -> '%s' and label selector -> '%s'",
                 field_selector,
                 label_selector,
             )
@@ -273,7 +273,7 @@ class K8s:
                 The number of seconds to wait between repetitive status checks.
 
         Returns:
-            True once the pod reaches the condition - otherwise wait forever.
+            None
 
         """
 
@@ -319,6 +319,7 @@ class K8s:
         max_retry: int = 3,
         time_retry: int = 10,
         container: str | None = None,
+        timeout: int = 60,
     ) -> str:
         """Execute a command inside a Kubernetes Pod (similar to kubectl exec on command line).
 
@@ -336,6 +337,9 @@ class K8s:
                 Wait time in seconds between retries.
             container (str):
                 The container name if the pod runs multiple containers inside.
+            timeout (int):
+                Timeout duration that is waited for any response in seconds.
+                Each time a response is found in stdout or stderr we wait another timeout duration [60]
 
         Returns:
             str:
@@ -363,6 +367,7 @@ class K8s:
                     stdin=False,
                     stdout=True,
                     tty=False,
+                    _request_timeout=timeout,
                 )
             except ApiException as exc:
                 self.logger.warning(
@@ -382,7 +387,7 @@ class K8s:
                 time.sleep(time_retry)
                 continue
             else:
-                self.logger.debug(response)
+                self.logger.debug("Command execution response -> %s", response if response else "<empty>")
                 return response
 
         self.logger.error(
@@ -437,7 +442,7 @@ class K8s:
             self.logger.error("Pod -> '%s' does not exist", pod_name)
 
         if not commands:
-            self.logger.error("No commands to execute on Pod ->'%s'!", pod_name)
+            self.logger.error("No commands to execute on pod ->'%s'!", pod_name)
             return None
 
         # Get first command - this should be the shell:
@@ -454,6 +459,7 @@ class K8s:
                 stdout=True,
                 tty=False,
                 _preload_content=False,  # This is important!
+                _request_timeout=timeout,
             )
         except ApiException:
             self.logger.error(
@@ -526,7 +532,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to delete Pod -> '%s'",
+                "Failed to delete pod -> '%s'",
                 pod_name,
             )
             return None
@@ -571,7 +577,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to get Config Map -> '%s'",
+                "Failed to get config map -> '%s'",
                 config_map_name,
             )
             return None
@@ -582,8 +588,8 @@ class K8s:
 
     def list_config_maps(
         self,
-        field_selector: str = "",
-        label_selector: str = "",
+        field_selector: str | None = None,
+        label_selector: str | None = None,
     ) -> V1ConfigMapList:
         """List all Kubernetes Config Maps in the current namespace.
 
@@ -616,7 +622,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to list Config Maps with field_selector -> '%s' and label_selector -> '%s'",
+                "Failed to list config maps with field selector -> '%s' and label selector -> '%s'",
                 field_selector,
                 label_selector,
             )
@@ -648,7 +654,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to find Config Map -> '%s'",
+                "Failed to find config map -> '%s'",
                 config_map_name,
             )
             return None
@@ -692,7 +698,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to replace Config Map -> '%s'",
+                "Failed to replace config map -> '%s'",
                 config_map_name,
             )
             return None
@@ -724,7 +730,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to get Stateful Set -> '%s'",
+                "Failed to get stateful set -> '%s'",
                 sts_name,
             )
             return None
@@ -756,7 +762,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to get scaling (replicas) of Stateful Set -> '%s'",
+                "Failed to get scaling (replicas) of stateful set -> '%s'",
                 sts_name,
             )
             return None
@@ -791,9 +797,9 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to patch Stateful Set -> '%s' with -> %s",
+                "Failed to patch stateful set -> '%s' with -> %s",
                 sts_name,
-                sts_body,
+                str(sts_body),
             )
             return None
 
@@ -826,7 +832,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to scale Stateful Set -> '%s' to -> %s replicas",
+                "Failed to scale stateful set -> '%s' to -> %s replicas",
                 sts_name,
                 scale,
             )
@@ -858,7 +864,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to get Service -> '%s'",
+                "Failed to get service -> '%s'",
                 service_name,
             )
             return None
@@ -901,7 +907,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to list Services with field_selector -> '%s' and label_selector -> '%s'",
+                "Failed to list services with field selector -> '%s' and label selector -> '%s'",
                 field_selector,
                 label_selector,
             )
@@ -938,7 +944,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to patch Service -> '%s' with -> %s",
+                "Failed to patch service -> '%s' with -> %s",
                 service_name,
                 service_body,
             )
@@ -970,7 +976,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to get Ingress -> '%s'!",
+                "Failed to get ingress -> '%s'!",
                 ingress_name,
             )
             return None
@@ -1008,7 +1014,7 @@ class K8s:
             )
         except ApiException:
             self.logger.error(
-                "Failed to patch Ingress -> '%s' with -> %s",
+                "Failed to patch ingress -> '%s' with -> %s",
                 ingress_name,
                 ingress_body,
             )
@@ -1275,9 +1281,21 @@ class K8s:
                 self.logger.info("Pod -> '%s' has been deleted successfully.", pod_name)
                 return True
 
-            self.logger.debug(
-                "Pod -> '%s' still exists. Waiting %s seconds before next check.",
+            pod_status = self.get_core_v1_api().read_namespaced_pod_status(
                 pod_name,
+                self.get_namespace(),
+            )
+
+            self.logger.info(
+                "Pod -> '%s' still exists with conditions -> %s. Waiting %s seconds before next check.",
+                pod_name,
+                str(
+                    [
+                        pod_condition.type
+                        for pod_condition in pod_status.status.conditions
+                        if pod_condition.status == "True"
+                    ]
+                ),
                 retry_interval,
             )
             time.sleep(retry_interval)
@@ -1289,12 +1307,20 @@ class K8s:
 
     # end method definition
 
-    def restart_deployment(self, deployment_name: str) -> bool:
+    def restart_deployment(
+        self, deployment_name: str, force: bool = False, wait: bool = False, wait_timeout: int = 1800
+    ) -> bool:
         """Restart a Kubernetes deployment using rolling restart.
 
         Args:
             deployment_name (str):
                 Name of the Kubernetes deployment.
+            force (bool):
+                If True, all pod instances will be forcefully deleted. [False]
+            wait (bool):
+                If True, wait for the stateful set to be ready again. [False]
+            wait_timeout (int):
+                Maximum time to wait for the stateful set to be ready again (in seconds). [1800]
 
         Returns:
             bool:
@@ -1302,45 +1328,137 @@ class K8s:
 
         """
 
-        now = datetime.now(timezone.utc).isoformat(timespec="seconds") + "Z"
+        success = True
 
-        body = {
-            "spec": {
-                "template": {
-                    "metadata": {
-                        "annotations": {
-                            "kubectl.kubernetes.io/restartedAt": now,
+        if not force:
+            now = datetime.now(timezone.utc).isoformat(timespec="seconds") + "Z"
+
+            body = {
+                "spec": {
+                    "template": {
+                        "metadata": {
+                            "annotations": {
+                                "kubectl.kubernetes.io/restartedAt": now,
+                            },
                         },
                     },
                 },
-            },
-        }
-        try:
-            self.get_apps_v1_api().patch_namespaced_deployment(
-                deployment_name,
-                self.get_namespace(),
-                body,
-                pretty="true",
-            )
+            }
+            try:
+                self.get_apps_v1_api().patch_namespaced_deployment(
+                    deployment_name,
+                    self.get_namespace(),
+                    body,
+                    pretty="true",
+                )
+                self.logger.info("Triggered restart of deployment -> '%s'.", deployment_name)
 
-        except ApiException:
-            self.logger.exception(
-                "Failed to restart deployment -> '%s'!",
-                deployment_name,
+            except ApiException as api_exception:
+                self.logger.error(
+                    "Failed to restart deployment -> '%s'; error -> %s!", deployment_name, str(api_exception)
+                )
+                return False
+
+        # If force is set, all pod instances will be forcefully deleted.
+        elif force:
+            self.logger.info("Force deleting all pods of deployment -> '%s'.", deployment_name)
+
+            try:
+                # Get the Deployment to retrieve its pod labels
+                deployment = self.get_apps_v1_api().read_namespaced_deployment(
+                    name=deployment_name, namespace=self.get_namespace()
+                )
+
+                # Get the label selector for the Deployment
+                label_selector = deployment.spec.selector.match_labels
+
+                # List pods matching the label selector
+                pods = (
+                    self.get_core_v1_api()
+                    .list_namespaced_pod(
+                        namespace=self.get_namespace(),
+                        label_selector=",".join([f"{k}={v}" for k, v in label_selector.items()]),
+                    )
+                    .items
+                )
+
+                # Loop through the pods and delete each one
+                for pod in pods:
+                    pod_name = pod.metadata.name
+                    try:
+                        # Define the delete options with force and grace period set to 0
+                        body = client.V1DeleteOptions(grace_period_seconds=0, propagation_policy="Foreground")
+
+                        # Call the delete_namespaced_pod method
+                        self.get_core_v1_api().delete_namespaced_pod(
+                            name=pod_name, namespace=self.get_namespace(), body=body
+                        )
+                        self.logger.info(
+                            "Pod '%s' in namespace '%s' has been deleted forcefully.", pod_name, self.get_namespace()
+                        )
+                    except Exception as e:
+                        self.logger.error("Error occurred while deleting pod '%s': %s", pod_name, e)
+                        success = False
+
+            except Exception as e:
+                self.logger.error("Error occurred while getting Deployment '%s': %s", deployment_name, e)
+                success = False
+
+        start_time = time.time()
+        while wait:
+            self.logger.info("Waiting for restart of deployment -> '%s' to complete.", deployment_name)
+            # Get the deployment
+            deployment = self.get_apps_v1_api().read_namespaced_deployment_status(deployment_name, self.get_namespace())
+
+            # Check the availability status
+            ready_replicas = deployment.status.ready_replicas or 0
+            updated_replicas = deployment.status.updated_replicas or 0
+            unavailable_replicas = deployment.status.unavailable_replicas or 0
+            total_replicas = deployment.status.replicas or 0
+            desired_replicas = deployment.spec.replicas or 0
+
+            self.logger.debug(
+                "Deployment status -> updated pods: %s/%s -> ready replicas: %s/%s",
+                updated_replicas,
+                desired_replicas,
+                ready_replicas,
+                total_replicas,
             )
-            return False
 
-        else:
-            return True
+            if (
+                updated_replicas == desired_replicas
+                and unavailable_replicas == 0
+                and total_replicas == desired_replicas
+            ):
+                self.logger.info("Restart of deployment -> '%s' completed successfully", deployment_name)
+                break
+
+            if (time.time() - start_time) > wait_timeout:
+                self.logger.error("Timed out waiting for restart of deployment -> '%s' to complete.", deployment_name)
+                success = False
+                break
+
+            # Sleep for a while before checking again
+            time.sleep(20)
+
+        return success
 
     # end method definition
 
-    def restart_stateful_set(self, sts_name: str) -> bool:
+    def restart_stateful_set(
+        self, sts_name: str, force: bool = False, wait: bool = False, wait_timeout: int = 1800
+    ) -> bool:
         """Restart a Kubernetes stateful set using rolling restart.
 
         Args:
             sts_name (str):
                 Name of the Kubernetes statefulset.
+            force (bool, optional):
+                If True, all pod instances will be forcefully deleted. [False]
+            wait (bool, optional):
+                If True, wait for the stateful set to be ready again. [False]
+            wait_timeout (int, optional):
+                Maximum time to wait for the stateful set to be ready again (in seconds). [1800]
 
         Returns:
             bool:
@@ -1348,6 +1466,8 @@ class K8s:
 
         """
 
+        success = True
+
         now = datetime.now(timezone.utc).isoformat(timespec="seconds") + "Z"
 
         body = {
@@ -1361,17 +1481,84 @@ class K8s:
                 },
             },
         }
+
         try:
             self.get_apps_v1_api().patch_namespaced_stateful_set(sts_name, self.get_namespace(), body, pretty="true")
+            self.logger.info("Triggered restart of stateful set -> '%s'.", sts_name)
 
-        except ApiException:
-            self.logger.exception(
-                "Failed to restart stateful set -> '%s'!",
-                sts_name,
-            )
+        except ApiException as api_exception:
+            self.logger.error("Failed to restart stateful set -> '%s'; error -> %s!", sts_name, str(api_exception))
             return False
 
-        else:
-            return True
+        # If force is set, all pod instances will be forcefully deleted.
+        if force:
+            self.logger.info("Force deleting all pods of stateful set -> '%s'.", sts_name)
+
+            try:
+                # Get the StatefulSet
+                statefulset = self.get_apps_v1_api().read_namespaced_stateful_set(
+                    name=sts_name, namespace=self.get_namespace()
+                )
+
+                # Loop through the replicas of the StatefulSet
+                for i in range(statefulset.spec.replicas):
+                    pod_name = f"{statefulset.metadata.name}-{i}"
+                    try:
+                        # Define the delete options with force and grace period set to 0
+                        body = client.V1DeleteOptions(grace_period_seconds=0, propagation_policy="Foreground")
+
+                        # Call the delete_namespaced_pod method
+                        self.get_core_v1_api().delete_namespaced_pod(
+                            name=pod_name, namespace=self.get_namespace(), body=body
+                        )
+                        self.logger.info(
+                            "Pod -> '%s' in namespace -> '%s' has been deleted forcefully.",
+                            pod_name,
+                            self.get_namespace(),
+                        )
+
+                    except Exception as e:
+                        self.logger.error("Error occurred while deleting pod -> '%s': %s", pod_name, str(e))
+                        success = False
+
+            except Exception as e:
+                self.logger.error("Error occurred while getting stateful set -> '%s': %s", sts_name, str(e))
+                success = False
+
+        start_time = time.time()
+
+        while wait:
+            time.sleep(10)  # Add delay before checking that the stateful set is ready again.
+            self.logger.info("Waiting for restart of stateful set -> '%s' to complete...", sts_name)
+            # Get the deployment
+            statefulset = self.get_apps_v1_api().read_namespaced_stateful_set_status(sts_name, self.get_namespace())
+
+            # Check the availability status
+            available_replicas = statefulset.status.available_replicas or 0
+            desired_replicas = statefulset.spec.replicas or 0
+
+            current_revision = statefulset.status.current_revision or ""
+            update_revision = statefulset.status.update_revision or ""
+
+            self.logger.debug(
+                "Stateful set status -> available pods: %s/%s, revision updated: %s",
+                available_replicas,
+                desired_replicas,
+                current_revision == update_revision,
+            )
+
+            if available_replicas == desired_replicas and update_revision == current_revision:
+                self.logger.info("Stateful set -> '%s' completed restart successfully", sts_name)
+                break
+
+            if (time.time() - start_time) > wait_timeout:
+                self.logger.error("Timed out waiting for restart of stateful set -> '%s' to complete.", sts_name)
+                success = False
+                break
+
+            # Sleep for a while before checking again
+            time.sleep(10)
+
+        return success
 
     # end method definition