Make readiness-check work when updatedReplicas is None
[oom.git] / kubernetes / readiness / docker / init / ready.py
1 #!/usr/bin/env python
2 import getopt
3 import logging
4 import os
5 import sys
6 import time
7 import random
8
9 from kubernetes import client
10
11 # extract env variables.
12 namespace = os.environ['NAMESPACE']
13 cert = os.environ['CERT']
14 host = os.environ['KUBERNETES_SERVICE_HOST']
15 token_path = os.environ['TOKEN']
16
17 with open(token_path, 'r') as token_file:
18     token = token_file.read().replace('\n', '')
19
20 # setup logging
21 log = logging.getLogger(__name__)
22 handler = logging.StreamHandler(sys.stdout)
23 formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
24 handler.setFormatter(formatter)
25 handler.setLevel(logging.INFO)
26 log.addHandler(handler)
27 log.setLevel(logging.INFO)
28
29 configuration = client.Configuration()
30 configuration.host = "https://" + host
31 configuration.ssl_ca_cert = cert
32 configuration.api_key['authorization'] = token
33 configuration.api_key_prefix['authorization'] = 'Bearer'
34 coreV1Api = client.CoreV1Api(client.ApiClient(configuration))
35 api_instance = client.ExtensionsV1beta1Api(client.ApiClient(configuration))
36 api = client.AppsV1beta1Api(client.ApiClient(configuration))
37 batchV1Api = client.BatchV1Api(client.ApiClient(configuration))
38
39
40 def is_job_complete(job_name):
41     complete = False
42     log.info("Checking if " + job_name + "  is complete")
43     try:
44         response = batchV1Api.read_namespaced_job_status(job_name, namespace)
45         if response.status.succeeded == 1:
46             job_status_type = response.status.conditions[0].type
47             if job_status_type == "Complete":
48                 complete = True
49                 log.info(job_name + " is complete")
50             else:
51                 log.info(job_name + " is not complete")
52         else:
53             log.info(job_name + " has not succeeded yet")
54         return complete
55     except Exception as e:
56         log.error("Exception when calling read_namespaced_job_status: %s\n" % e)
57
58
59 def wait_for_statefulset_complete(statefulset_name):
60     try:
61         response = api.read_namespaced_stateful_set(statefulset_name, namespace)
62         s = response.status
63         if (s.updated_replicas == response.spec.replicas and
64                 s.replicas == response.spec.replicas and
65                 s.ready_replicas == response.spec.replicas and
66                 s.current_replicas == response.spec.replicas and
67                 s.observed_generation == response.metadata.generation):
68             log.info("Statefulset " + statefulset_name + "  is ready")
69             return True
70         else:
71             log.info("Statefulset " + statefulset_name + "  is not ready")
72         return False
73     except Exception as e:
74         log.error("Exception when waiting for Statefulset status: %s\n" % e)
75
76
77 def wait_for_deployment_complete(deployment_name):
78     try:
79         response = api.read_namespaced_deployment(deployment_name, namespace)
80         s = response.status
81         if (s.unavailable_replicas is None and
82                 ( s.updated_replicas is None or s.updated_replicas == response.spec.replicas ) and
83                 s.replicas == response.spec.replicas and
84                 s.ready_replicas == response.spec.replicas and
85                 s.observed_generation == response.metadata.generation):
86             log.info("Deployment " + deployment_name + "  is ready")
87             return True
88         else:
89             log.info("Deployment " + deployment_name + "  is not ready")
90         return False
91     except Exception as e:
92         log.error("Exception when waiting for deployment status: %s\n" % e)
93
94
95 def is_ready(container_name):
96     ready = False
97     log.info("Checking if " + container_name + "  is ready")
98     try:
99         response = coreV1Api.list_namespaced_pod(namespace=namespace,
100                                                  watch=False)
101         for i in response.items:
102             # container_statuses can be None, which is non-iterable.
103             if i.status.container_statuses is None:
104                 continue
105             for s in i.status.container_statuses:
106                 if s.name == container_name:
107                     name = read_name(i)
108                     if i.metadata.owner_references[0].kind == "StatefulSet":
109                         ready = wait_for_statefulset_complete(name)
110                     elif i.metadata.owner_references[0].kind == "ReplicaSet":
111                         deployment_name = get_deployment_name(name)
112                         ready = wait_for_deployment_complete(deployment_name)
113                     elif i.metadata.owner_references[0].kind == "Job":
114                         ready = is_job_complete(name)
115
116                     return ready
117
118                 else:
119                     continue
120         return ready
121     except Exception as e:
122         log.error("Exception when calling list_namespaced_pod: %s\n" % e)
123
124
125 def read_name(item):
126     return item.metadata.owner_reference[0].name
127
128
129 def get_deployment_name(replicaset):
130     api_response = api_instance.read_namespaced_replica_set_status(replicaset,
131                                                                    namespace)
132     deployment_name = read_name(api_response)
133     return deployment_name
134
135
136 DEF_TIMEOUT = 10
137 DESCRIPTION = "Kubernetes container readiness check utility"
138 USAGE = "Usage: ready.py [-t <timeout>] -c <container_name> " \
139         "[-c <container_name> ...]\n" \
140         "where\n" \
141         "<timeout> - wait for container readiness timeout in min, " \
142         "default is " + str(DEF_TIMEOUT) + "\n" \
143         "<container_name> - name of the container to wait for\n"
144
145
146 def main(argv):
147     # args are a list of container names
148     container_names = []
149     timeout = DEF_TIMEOUT
150     try:
151         opts, args = getopt.getopt(argv, "hc:t:", ["container-name=",
152                                                    "timeout=",
153                                                    "help"])
154         for opt, arg in opts:
155             if opt in ("-h", "--help"):
156                 print("%s\n\n%s" % (DESCRIPTION, USAGE))
157                 sys.exit()
158             elif opt in ("-c", "--container-name"):
159                 container_names.append(arg)
160             elif opt in ("-t", "--timeout"):
161                 timeout = float(arg)
162     except (getopt.GetoptError, ValueError) as e:
163         print("Error parsing input parameters: %s\n" % e)
164         print(USAGE)
165         sys.exit(2)
166     if container_names.__len__() == 0:
167         print("Missing required input parameter(s)\n")
168         print(USAGE)
169         sys.exit(2)
170
171     for container_name in container_names:
172         timeout = time.time() + timeout * 60
173         while True:
174             ready = is_ready(container_name)
175             if ready is True:
176                 break
177             elif time.time() > timeout:
178                 log.warning("timed out waiting for '" + container_name +
179                             "' to be ready")
180                 exit(1)
181             else:
182                 # spread in time potentially parallel execution in multiple
183                 # containers
184                 time.sleep(random.randint(5, 11))
185
186
187 if __name__ == "__main__":
188     main(sys.argv[1:])
189