kubernetes_rolling_restart.py

import sys
import time

from kubernetes import config, client


class KubeControl:
    def __init__(self):
        # config.load_kube_config(context='staging')
        config.load_kube_config()
        self.v1 = client.CoreV1Api()
        self.api_instance = client.AppsV1beta1Api()

    def pods(self, namespace):
        ret = self.v1.list_namespaced_pod(namespace)
        return ret.items

    def deployments(self, namespace):
        api_response = self.api_instance.list_namespaced_deployment(namespace)
        return api_response.items

    def display(self, namespace):
        for pod in self.pods(namespace):
            print("%s\t\t%s\t\t%s\t\t%s" % (pod.status.pod_ip, pod.metadata.namespace, pod.metadata.name, pod.status.phase))

    def restart_namespace_pods(self, namespace='default'):
        for pod in self.pods(namespace):
            if pod.metadata.namespace == namespace:
                print('Deleting pod: {}'.format(pod.metadata.name))
                self.v1.delete_namespaced_pod(pod.metadata.name, namespace, body={})
            else:
                print('Ignoring pod: {}, out of namespace'.format(pod.metadata.name))

    def _below_threshold(self, deployment, threshold):
        ready_replicas = deployment.status.ready_replicas or 0
        if ((ready_replicas * 100.0)/deployment.status.replicas) <= threshold:
            return True
        return False

    def rolling_restart(self, namespace='default', kill_percent=50):
        alive_percent = 100 - kill_percent

        for deployment in self.deployments(namespace):
            print('Killing deployment: {}'.format(deployment.metadata.name))

            deployment_pods = []
            for pod in self.pods(namespace):
                if pod.metadata.name.startswith(deployment.metadata.name):
                    deployment_pods.append(pod)

            for dep_pod in deployment_pods:
                self.v1.delete_namespaced_pod(dep_pod.metadata.name, namespace, body={})

                counter = 0
                while self._below_threshold(deployment, alive_percent) is True and counter < 5:
                    print('Alive pod count is below threshold. Sleeping for {}'.format(counter))
                    time.sleep(counter)
                    counter += 1

                if counter >= 5:
                    print(('Alive pod count lower than the threshold even after waiting. '
                           'Ignoring pod deletion for deployment: {}').format(deployment.metadata.name))
                    break


if __name__ == '__main__':
    namespace = sys.argv[1]
    kc = KubeControl()
    kc.rolling_restart(namespace=namespace)