Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 1 | #!/usr/bin/env python2.7 |
| 2 | # Copyright 2015-2016, Google Inc. |
| 3 | # All rights reserved. |
| 4 | # |
| 5 | # Redistribution and use in source and binary forms, with or without |
| 6 | # modification, are permitted provided that the following conditions are |
| 7 | # met: |
| 8 | # |
| 9 | # * Redistributions of source code must retain the above copyright |
| 10 | # notice, this list of conditions and the following disclaimer. |
| 11 | # * Redistributions in binary form must reproduce the above |
| 12 | # copyright notice, this list of conditions and the following disclaimer |
| 13 | # in the documentation and/or other materials provided with the |
| 14 | # distribution. |
| 15 | # * Neither the name of Google Inc. nor the names of its |
| 16 | # contributors may be used to endorse or promote products derived from |
| 17 | # this software without specific prior written permission. |
| 18 | # |
| 19 | # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS |
| 20 | # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT |
| 21 | # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR |
| 22 | # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT |
| 23 | # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, |
| 24 | # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT |
| 25 | # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, |
| 26 | # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY |
| 27 | # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT |
| 28 | # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE |
| 29 | # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
| 30 | import argparse |
| 31 | import datetime |
| 32 | import json |
| 33 | import os |
| 34 | import subprocess |
| 35 | import sys |
| 36 | import time |
| 37 | |
| 38 | stress_test_utils_dir = os.path.abspath(os.path.join( |
| 39 | os.path.dirname(__file__), '../../gcp/stress_test')) |
| 40 | sys.path.append(stress_test_utils_dir) |
| 41 | from stress_test_utils import BigQueryHelper |
| 42 | |
| 43 | kubernetes_api_dir = os.path.abspath(os.path.join( |
| 44 | os.path.dirname(__file__), '../../gcp/utils')) |
| 45 | sys.path.append(kubernetes_api_dir) |
| 46 | |
| 47 | import kubernetes_api |
| 48 | |
| 49 | |
| 50 | class GlobalSettings: |
| 51 | |
| 52 | def __init__(self, gcp_project_id, build_docker_images, |
| 53 | test_poll_interval_secs, test_duration_secs, |
| 54 | kubernetes_proxy_port, dataset_id_prefix, summary_table_id, |
| 55 | qps_table_id, pod_warmup_secs): |
| 56 | self.gcp_project_id = gcp_project_id |
| 57 | self.build_docker_images = build_docker_images |
| 58 | self.test_poll_interval_secs = test_poll_interval_secs |
| 59 | self.test_duration_secs = test_duration_secs |
| 60 | self.kubernetes_proxy_port = kubernetes_proxy_port |
| 61 | self.dataset_id_prefix = dataset_id_prefix |
| 62 | self.summary_table_id = summary_table_id |
| 63 | self.qps_table_id = qps_table_id |
| 64 | self.pod_warmup_secs = pod_warmup_secs |
| 65 | |
| 66 | |
| 67 | class ClientTemplate: |
Sree Kuchibhotla | 9c9644b | 2016-03-28 09:30:51 -0700 | [diff] [blame] | 68 | """ Contains all the common settings that are used by a stress client """ |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 69 | |
| 70 | def __init__(self, name, client_image_path, metrics_client_image_path, |
| 71 | metrics_port, wrapper_script_path, poll_interval_secs, |
| 72 | client_args_dict, metrics_args_dict): |
| 73 | self.name = name |
| 74 | self.client_image_path = client_image_path |
| 75 | self.metrics_client_image_path = metrics_client_image_path |
| 76 | self.metrics_port = metrics_port |
| 77 | self.wrapper_script_path = wrapper_script_path |
| 78 | self.poll_interval_secs = poll_interval_secs |
| 79 | self.client_args_dict = client_args_dict |
| 80 | self.metrics_args_dict = metrics_args_dict |
| 81 | |
| 82 | |
| 83 | class ServerTemplate: |
Sree Kuchibhotla | 9c9644b | 2016-03-28 09:30:51 -0700 | [diff] [blame] | 84 | """ Contains all the common settings used by a stress server """ |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 85 | |
| 86 | def __init__(self, name, server_image_path, wrapper_script_path, server_port, |
| 87 | server_args_dict): |
| 88 | self.name = name |
| 89 | self.server_image_path = server_image_path |
| 90 | self.wrapper_script_path = wrapper_script_path |
| 91 | self.server_port = server_port |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 92 | self.server_args_dict = server_args_dict |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 93 | |
| 94 | |
| 95 | class DockerImage: |
Sree Kuchibhotla | e68ec43 | 2016-03-28 13:55:01 -0700 | [diff] [blame^] | 96 | """ Represents properties of a Docker image. Provides methods to build the |
| 97 | image and push it to GKE registry |
Sree Kuchibhotla | 9c9644b | 2016-03-28 09:30:51 -0700 | [diff] [blame] | 98 | """ |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 99 | |
| 100 | def __init__(self, gcp_project_id, image_name, build_script_path, |
Sree Kuchibhotla | 8d41d51 | 2016-03-25 14:50:31 -0700 | [diff] [blame] | 101 | dockerfile_dir, build_type): |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 102 | """Args: |
| 103 | |
| 104 | image_name: The docker image name |
| 105 | tag_name: The additional tag name. This is the name used when pushing the |
| 106 | docker image to GKE registry |
| 107 | build_script_path: The path to the build script that builds this docker |
| 108 | image |
| 109 | dockerfile_dir: The name of the directory under |
| 110 | '<grpc_root>/tools/dockerfile' that contains the dockerfile |
| 111 | """ |
| 112 | self.image_name = image_name |
| 113 | self.gcp_project_id = gcp_project_id |
| 114 | self.build_script_path = build_script_path |
| 115 | self.dockerfile_dir = dockerfile_dir |
Sree Kuchibhotla | 8d41d51 | 2016-03-25 14:50:31 -0700 | [diff] [blame] | 116 | self.build_type = build_type |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 117 | self.tag_name = self._make_tag_name(gcp_project_id, image_name) |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 118 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 119 | def _make_tag_name(self, project_id, image_name): |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 120 | return 'gcr.io/%s/%s' % (project_id, image_name) |
| 121 | |
| 122 | def build_image(self): |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 123 | print 'Building docker image: %s (tag: %s)' % (self.image_name, |
| 124 | self.tag_name) |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 125 | os.environ['INTEROP_IMAGE'] = self.image_name |
| 126 | os.environ['INTEROP_IMAGE_REPOSITORY_TAG'] = self.tag_name |
| 127 | os.environ['BASE_NAME'] = self.dockerfile_dir |
Sree Kuchibhotla | 8d41d51 | 2016-03-25 14:50:31 -0700 | [diff] [blame] | 128 | os.environ['BUILD_TYPE'] = self.build_type |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 129 | print 'DEBUG: path: ', self.build_script_path |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 130 | if subprocess.call(args=[self.build_script_path]) != 0: |
| 131 | print 'Error in building the Docker image' |
| 132 | return False |
| 133 | return True |
| 134 | |
| 135 | def push_to_gke_registry(self): |
| 136 | cmd = ['gcloud', 'docker', 'push', self.tag_name] |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 137 | print 'Pushing %s to the GKE registry..' % self.tag_name |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 138 | if subprocess.call(args=cmd) != 0: |
| 139 | print 'Error in pushing the image %s to the GKE registry' % self.tag_name |
| 140 | return False |
| 141 | return True |
| 142 | |
| 143 | |
| 144 | class ServerPodSpec: |
Sree Kuchibhotla | 9c9644b | 2016-03-28 09:30:51 -0700 | [diff] [blame] | 145 | """ Contains the information required to launch server pods. """ |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 146 | |
| 147 | def __init__(self, name, server_template, docker_image, num_instances): |
| 148 | self.name = name |
| 149 | self.template = server_template |
| 150 | self.docker_image = docker_image |
| 151 | self.num_instances = num_instances |
| 152 | |
| 153 | def pod_names(self): |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 154 | """ Return a list of names of server pods to create. """ |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 155 | return ['%s-%d' % (self.name, i) for i in range(1, self.num_instances + 1)] |
| 156 | |
| 157 | def server_addresses(self): |
| 158 | """ Return string of server addresses in the following format: |
| 159 | '<server_pod_name_1>:<server_port>,<server_pod_name_2>:<server_port>...' |
| 160 | """ |
| 161 | return ','.join(['%s:%d' % (pod_name, self.template.server_port) |
| 162 | for pod_name in self.pod_names()]) |
| 163 | |
| 164 | |
| 165 | class ClientPodSpec: |
Sree Kuchibhotla | 9c9644b | 2016-03-28 09:30:51 -0700 | [diff] [blame] | 166 | """ Contains the information required to launch client pods """ |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 167 | |
| 168 | def __init__(self, name, client_template, docker_image, num_instances, |
| 169 | server_addresses): |
| 170 | self.name = name |
| 171 | self.template = client_template |
| 172 | self.docker_image = docker_image |
| 173 | self.num_instances = num_instances |
| 174 | self.server_addresses = server_addresses |
| 175 | |
| 176 | def pod_names(self): |
| 177 | """ Return a list of names of client pods to create """ |
| 178 | return ['%s-%d' % (self.name, i) for i in range(1, self.num_instances + 1)] |
| 179 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 180 | # The client args in the template do not have server addresses. This function |
| 181 | # adds the server addresses and returns the updated client args |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 182 | def get_client_args_dict(self): |
| 183 | args_dict = self.template.client_args_dict.copy() |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 184 | args_dict['server_addresses'] = self.server_addresses |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 185 | return args_dict |
| 186 | |
| 187 | |
| 188 | class Gke: |
Sree Kuchibhotla | 9c9644b | 2016-03-28 09:30:51 -0700 | [diff] [blame] | 189 | """ Class that has helper methods to interact with GKE """ |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 190 | |
| 191 | class KubernetesProxy: |
| 192 | """Class to start a proxy on localhost to talk to the Kubernetes API server""" |
| 193 | |
| 194 | def __init__(self, port): |
| 195 | cmd = ['kubectl', 'proxy', '--port=%d' % port] |
| 196 | self.p = subprocess.Popen(args=cmd) |
| 197 | time.sleep(2) |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 198 | print '\nStarted kubernetes proxy on port: %d' % port |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 199 | |
| 200 | def __del__(self): |
| 201 | if self.p is not None: |
| 202 | print 'Shutting down Kubernetes proxy..' |
| 203 | self.p.kill() |
| 204 | |
| 205 | def __init__(self, project_id, run_id, dataset_id, summary_table_id, |
| 206 | qps_table_id, kubernetes_port): |
| 207 | self.project_id = project_id |
| 208 | self.run_id = run_id |
| 209 | self.dataset_id = dataset_id |
| 210 | self.summary_table_id = summary_table_id |
| 211 | self.qps_table_id = qps_table_id |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 212 | |
| 213 | # The environment variables we would like to pass to every pod (both client |
| 214 | # and server) launched in GKE |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 215 | self.gke_env = { |
| 216 | 'RUN_ID': self.run_id, |
| 217 | 'GCP_PROJECT_ID': self.project_id, |
| 218 | 'DATASET_ID': self.dataset_id, |
| 219 | 'SUMMARY_TABLE_ID': self.summary_table_id, |
| 220 | 'QPS_TABLE_ID': self.qps_table_id |
| 221 | } |
| 222 | |
| 223 | self.kubernetes_port = kubernetes_port |
| 224 | # Start kubernetes proxy |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 225 | self.kubernetes_proxy = Gke.KubernetesProxy(kubernetes_port) |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 226 | |
| 227 | def _args_dict_to_str(self, args_dict): |
| 228 | return ' '.join('--%s=%s' % (k, args_dict[k]) for k in args_dict.keys()) |
| 229 | |
| 230 | def launch_servers(self, server_pod_spec): |
| 231 | is_success = True |
| 232 | |
| 233 | # The command to run inside the container is the wrapper script (which then |
| 234 | # launches the actual server) |
| 235 | container_cmd = server_pod_spec.template.wrapper_script_path |
| 236 | |
| 237 | # The parameters to the wrapper script (defined in |
| 238 | # server_pod_spec.template.wrapper_script_path) are are injected into the |
| 239 | # container via environment variables |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 240 | server_env = self.gke_env.copy() |
| 241 | server_env.update({ |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 242 | 'STRESS_TEST_IMAGE_TYPE': 'SERVER', |
| 243 | 'STRESS_TEST_IMAGE': server_pod_spec.template.server_image_path, |
| 244 | 'STRESS_TEST_ARGS_STR': self._args_dict_to_str( |
| 245 | server_pod_spec.template.server_args_dict) |
| 246 | }) |
| 247 | |
| 248 | for pod_name in server_pod_spec.pod_names(): |
| 249 | server_env['POD_NAME'] = pod_name |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 250 | print 'Creating server: %s' % pod_name |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 251 | is_success = kubernetes_api.create_pod_and_service( |
| 252 | 'localhost', |
| 253 | self.kubernetes_port, |
| 254 | 'default', # Use 'default' namespace |
| 255 | pod_name, |
| 256 | server_pod_spec.docker_image.tag_name, |
| 257 | [server_pod_spec.template.server_port], # Ports to expose on the pod |
| 258 | [container_cmd], |
| 259 | [], # Args list is empty since we are passing all args via env variables |
| 260 | server_env, |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 261 | True # Headless = True for server to that GKE creates a DNS record for pod_name |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 262 | ) |
| 263 | if not is_success: |
| 264 | print 'Error in launching server: %s' % pod_name |
| 265 | break |
| 266 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 267 | if is_success: |
| 268 | print 'Successfully created server(s)' |
| 269 | |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 270 | return is_success |
| 271 | |
| 272 | def launch_clients(self, client_pod_spec): |
| 273 | is_success = True |
| 274 | |
| 275 | # The command to run inside the container is the wrapper script (which then |
| 276 | # launches the actual stress client) |
| 277 | container_cmd = client_pod_spec.template.wrapper_script_path |
| 278 | |
| 279 | # The parameters to the wrapper script (defined in |
| 280 | # client_pod_spec.template.wrapper_script_path) are are injected into the |
| 281 | # container via environment variables |
| 282 | client_env = self.gke_env.copy() |
| 283 | client_env.update({ |
| 284 | 'STRESS_TEST_IMAGE_TYPE': 'CLIENT', |
| 285 | 'STRESS_TEST_IMAGE': client_pod_spec.template.client_image_path, |
| 286 | 'STRESS_TEST_ARGS_STR': self._args_dict_to_str( |
| 287 | client_pod_spec.get_client_args_dict()), |
| 288 | 'METRICS_CLIENT_IMAGE': |
| 289 | client_pod_spec.template.metrics_client_image_path, |
| 290 | 'METRICS_CLIENT_ARGS_STR': self._args_dict_to_str( |
| 291 | client_pod_spec.template.metrics_args_dict), |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 292 | 'POLL_INTERVAL_SECS': str(client_pod_spec.template.poll_interval_secs) |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 293 | }) |
| 294 | |
| 295 | for pod_name in client_pod_spec.pod_names(): |
| 296 | client_env['POD_NAME'] = pod_name |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 297 | print 'Creating client: %s' % pod_name |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 298 | is_success = kubernetes_api.create_pod_and_service( |
| 299 | 'localhost', |
| 300 | self.kubernetes_port, |
| 301 | 'default', # default namespace, |
| 302 | pod_name, |
| 303 | client_pod_spec.docker_image.tag_name, |
| 304 | [client_pod_spec.template.metrics_port], # Ports to expose on the pod |
| 305 | [container_cmd], |
| 306 | [], # Empty args list since all args are passed via env variables |
| 307 | client_env, |
| 308 | False # Client is not a headless service. |
| 309 | ) |
| 310 | |
| 311 | if not is_success: |
| 312 | print 'Error in launching client %s' % pod_name |
| 313 | break |
| 314 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 315 | if is_success: |
| 316 | print 'Successfully created all client(s)' |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 317 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 318 | return is_success |
| 319 | |
| 320 | def _delete_pods(self, pod_name_list): |
| 321 | is_success = True |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 322 | for pod_name in pod_name_list: |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 323 | print 'Deleting %s' % pod_name |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 324 | is_success = kubernetes_api.delete_pod_and_service( |
| 325 | 'localhost', |
| 326 | self.kubernetes_port, |
| 327 | 'default', # default namespace |
| 328 | pod_name) |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 329 | |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 330 | if not is_success: |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 331 | print 'Error in deleting pod %s' % pod_name |
| 332 | break |
| 333 | |
| 334 | if is_success: |
| 335 | print 'Successfully deleted all pods' |
| 336 | |
| 337 | return is_success |
| 338 | |
| 339 | def delete_servers(self, server_pod_spec): |
| 340 | return self._delete_pods(server_pod_spec.pod_names()) |
| 341 | |
| 342 | def delete_clients(self, client_pod_spec): |
| 343 | return self._delete_pods(client_pod_spec.pod_names()) |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 344 | |
| 345 | |
| 346 | class Config: |
| 347 | |
Sree Kuchibhotla | 815c589 | 2016-03-25 15:03:50 -0700 | [diff] [blame] | 348 | def __init__(self, config_filename, gcp_project_id): |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 349 | print 'Loading configuration...' |
| 350 | config_dict = self._load_config(config_filename) |
| 351 | |
| 352 | self.global_settings = self._parse_global_settings(config_dict, |
Sree Kuchibhotla | 9c9644b | 2016-03-28 09:30:51 -0700 | [diff] [blame] | 353 | gcp_project_id) |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 354 | self.docker_images_dict = self._parse_docker_images( |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 355 | config_dict, self.global_settings.gcp_project_id) |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 356 | self.client_templates_dict = self._parse_client_templates(config_dict) |
| 357 | self.server_templates_dict = self._parse_server_templates(config_dict) |
| 358 | self.server_pod_specs_dict = self._parse_server_pod_specs( |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 359 | config_dict, self.docker_images_dict, self.server_templates_dict) |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 360 | self.client_pod_specs_dict = self._parse_client_pod_specs( |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 361 | config_dict, self.docker_images_dict, self.client_templates_dict, |
| 362 | self.server_pod_specs_dict) |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 363 | print 'Loaded Configuaration.' |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 364 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 365 | def _parse_global_settings(self, config_dict, gcp_project_id): |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 366 | global_settings_dict = config_dict['globalSettings'] |
Sree Kuchibhotla | 815c589 | 2016-03-25 15:03:50 -0700 | [diff] [blame] | 367 | return GlobalSettings(gcp_project_id, |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 368 | global_settings_dict['buildDockerImages'], |
| 369 | global_settings_dict['pollIntervalSecs'], |
| 370 | global_settings_dict['testDurationSecs'], |
| 371 | global_settings_dict['kubernetesProxyPort'], |
| 372 | global_settings_dict['datasetIdNamePrefix'], |
| 373 | global_settings_dict['summaryTableId'], |
| 374 | global_settings_dict['qpsTableId'], |
| 375 | global_settings_dict['podWarmupSecs']) |
| 376 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 377 | def _parse_docker_images(self, config_dict, gcp_project_id): |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 378 | """Parses the 'dockerImages' section of the config file and returns a |
| 379 | Dictionary of 'DockerImage' objects keyed by docker image names""" |
| 380 | docker_images_dict = {} |
Sree Kuchibhotla | 8d41d51 | 2016-03-25 14:50:31 -0700 | [diff] [blame] | 381 | |
| 382 | docker_config_dict = config_dict['dockerImages'] |
| 383 | for image_name in docker_config_dict.keys(): |
| 384 | build_script_path = docker_config_dict[image_name]['buildScript'] |
| 385 | dockerfile_dir = docker_config_dict[image_name]['dockerFileDir'] |
| 386 | build_type = docker_config_dict[image_name]['buildType'] |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 387 | docker_images_dict[image_name] = DockerImage(gcp_project_id, image_name, |
| 388 | build_script_path, |
Sree Kuchibhotla | 8d41d51 | 2016-03-25 14:50:31 -0700 | [diff] [blame] | 389 | dockerfile_dir, build_type) |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 390 | return docker_images_dict |
| 391 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 392 | def _parse_client_templates(self, config_dict): |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 393 | """Parses the 'clientTemplates' section of the config file and returns a |
| 394 | Dictionary of 'ClientTemplate' objects keyed by client template names. |
| 395 | |
| 396 | Note: The 'baseTemplates' sub section of the config file contains templates |
| 397 | with default values and the 'templates' sub section contains the actual |
| 398 | client templates (which refer to the base template name to use for default |
| 399 | values). |
| 400 | """ |
| 401 | client_templates_dict = {} |
| 402 | |
| 403 | templates_dict = config_dict['clientTemplates']['templates'] |
| 404 | base_templates_dict = config_dict['clientTemplates'].get('baseTemplates', |
| 405 | {}) |
| 406 | for template_name in templates_dict.keys(): |
| 407 | # temp_dict is a temporary dictionary that merges base template dictionary |
| 408 | # and client template dictionary (with client template dictionary values |
| 409 | # overriding base template values) |
| 410 | temp_dict = {} |
| 411 | |
| 412 | base_template_name = templates_dict[template_name].get('baseTemplate') |
| 413 | if not base_template_name is None: |
| 414 | temp_dict = base_templates_dict[base_template_name].copy() |
| 415 | |
| 416 | temp_dict.update(templates_dict[template_name]) |
| 417 | |
| 418 | # Create and add ClientTemplate object to the final client_templates_dict |
| 419 | client_templates_dict[template_name] = ClientTemplate( |
| 420 | template_name, temp_dict['clientImagePath'], |
| 421 | temp_dict['metricsClientImagePath'], temp_dict['metricsPort'], |
| 422 | temp_dict['wrapperScriptPath'], temp_dict['pollIntervalSecs'], |
| 423 | temp_dict['clientArgs'].copy(), temp_dict['metricsArgs'].copy()) |
| 424 | |
| 425 | return client_templates_dict |
| 426 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 427 | def _parse_server_templates(self, config_dict): |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 428 | """Parses the 'serverTemplates' section of the config file and returns a |
| 429 | Dictionary of 'serverTemplate' objects keyed by server template names. |
| 430 | |
| 431 | Note: The 'baseTemplates' sub section of the config file contains templates |
| 432 | with default values and the 'templates' sub section contains the actual |
| 433 | server templates (which refer to the base template name to use for default |
| 434 | values). |
| 435 | """ |
| 436 | server_templates_dict = {} |
| 437 | |
| 438 | templates_dict = config_dict['serverTemplates']['templates'] |
| 439 | base_templates_dict = config_dict['serverTemplates'].get('baseTemplates', |
| 440 | {}) |
| 441 | |
| 442 | for template_name in templates_dict.keys(): |
| 443 | # temp_dict is a temporary dictionary that merges base template dictionary |
| 444 | # and server template dictionary (with server template dictionary values |
| 445 | # overriding base template values) |
| 446 | temp_dict = {} |
| 447 | |
| 448 | base_template_name = templates_dict[template_name].get('baseTemplate') |
| 449 | if not base_template_name is None: |
| 450 | temp_dict = base_templates_dict[base_template_name].copy() |
| 451 | |
| 452 | temp_dict.update(templates_dict[template_name]) |
| 453 | |
| 454 | # Create and add ServerTemplate object to the final server_templates_dict |
| 455 | server_templates_dict[template_name] = ServerTemplate( |
| 456 | template_name, temp_dict['serverImagePath'], |
| 457 | temp_dict['wrapperScriptPath'], temp_dict['serverPort'], |
| 458 | temp_dict['serverArgs'].copy()) |
| 459 | |
| 460 | return server_templates_dict |
| 461 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 462 | def _parse_server_pod_specs(self, config_dict, docker_images_dict, |
Sree Kuchibhotla | 9c9644b | 2016-03-28 09:30:51 -0700 | [diff] [blame] | 463 | server_templates_dict): |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 464 | """Parses the 'serverPodSpecs' sub-section (under 'testMatrix' section) of |
| 465 | the config file and returns a Dictionary of 'ServerPodSpec' objects keyed |
| 466 | by server pod spec names""" |
| 467 | server_pod_specs_dict = {} |
| 468 | |
| 469 | pod_specs_dict = config_dict['testMatrix'].get('serverPodSpecs', {}) |
| 470 | |
| 471 | for pod_name in pod_specs_dict.keys(): |
| 472 | server_template_name = pod_specs_dict[pod_name]['serverTemplate'] |
| 473 | docker_image_name = pod_specs_dict[pod_name]['dockerImage'] |
| 474 | num_instances = pod_specs_dict[pod_name].get('numInstances', 1) |
| 475 | |
| 476 | # Create and add the ServerPodSpec object to the final |
| 477 | # server_pod_specs_dict |
| 478 | server_pod_specs_dict[pod_name] = ServerPodSpec( |
| 479 | pod_name, server_templates_dict[server_template_name], |
| 480 | docker_images_dict[docker_image_name], num_instances) |
| 481 | |
| 482 | return server_pod_specs_dict |
| 483 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 484 | def _parse_client_pod_specs(self, config_dict, docker_images_dict, |
Sree Kuchibhotla | 9c9644b | 2016-03-28 09:30:51 -0700 | [diff] [blame] | 485 | client_templates_dict, server_pod_specs_dict): |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 486 | """Parses the 'clientPodSpecs' sub-section (under 'testMatrix' section) of |
| 487 | the config file and returns a Dictionary of 'ClientPodSpec' objects keyed |
| 488 | by client pod spec names""" |
| 489 | client_pod_specs_dict = {} |
| 490 | |
| 491 | pod_specs_dict = config_dict['testMatrix'].get('clientPodSpecs', {}) |
| 492 | for pod_name in pod_specs_dict.keys(): |
| 493 | client_template_name = pod_specs_dict[pod_name]['clientTemplate'] |
| 494 | docker_image_name = pod_specs_dict[pod_name]['dockerImage'] |
| 495 | num_instances = pod_specs_dict[pod_name]['numInstances'] |
| 496 | |
| 497 | # Get the server addresses from the server pod spec object |
| 498 | server_pod_spec_name = pod_specs_dict[pod_name]['serverPodSpec'] |
| 499 | server_addresses = server_pod_specs_dict[ |
| 500 | server_pod_spec_name].server_addresses() |
| 501 | |
| 502 | client_pod_specs_dict[pod_name] = ClientPodSpec( |
| 503 | pod_name, client_templates_dict[client_template_name], |
| 504 | docker_images_dict[docker_image_name], num_instances, |
| 505 | server_addresses) |
| 506 | |
| 507 | return client_pod_specs_dict |
| 508 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 509 | def _load_config(self, config_filename): |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 510 | """Opens the config file and converts the Json text to Dictionary""" |
| 511 | if not os.path.isabs(config_filename): |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 512 | raise Exception('Config objects expects an absolute file path. ' |
| 513 | 'config file name passed: %s' % config_filename) |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 514 | with open(config_filename) as config_file: |
| 515 | return json.load(config_file) |
| 516 | |
| 517 | |
| 518 | def run_tests(config): |
Sree Kuchibhotla | 9c9644b | 2016-03-28 09:30:51 -0700 | [diff] [blame] | 519 | """ The main function that launches the stress tests """ |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 520 | # Build docker images and push to GKE registry |
| 521 | if config.global_settings.build_docker_images: |
| 522 | for name, docker_image in config.docker_images_dict.iteritems(): |
| 523 | if not (docker_image.build_image() and |
| 524 | docker_image.push_to_gke_registry()): |
| 525 | return False |
| 526 | |
| 527 | # Create a unique id for this run (Note: Using timestamp instead of UUID to |
| 528 | # make it easier to deduce the date/time of the run just by looking at the run |
| 529 | # run id. This is useful in debugging when looking at records in Biq query) |
| 530 | run_id = datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S') |
| 531 | dataset_id = '%s_%s' % (config.global_settings.dataset_id_prefix, run_id) |
| 532 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 533 | bq_helper = BigQueryHelper(run_id, '', '', |
| 534 | config.global_settings.gcp_project_id, dataset_id, |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 535 | config.global_settings.summary_table_id, |
| 536 | config.global_settings.qps_table_id) |
| 537 | bq_helper.initialize() |
| 538 | |
| 539 | gke = Gke(config.global_settings.gcp_project_id, run_id, dataset_id, |
| 540 | config.global_settings.summary_table_id, |
| 541 | config.global_settings.qps_table_id, |
| 542 | config.global_settings.kubernetes_proxy_port) |
| 543 | |
| 544 | is_success = True |
| 545 | |
| 546 | try: |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 547 | print 'Launching servers..' |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 548 | for name, server_pod_spec in config.server_pod_specs_dict.iteritems(): |
| 549 | if not gke.launch_servers(server_pod_spec): |
| 550 | is_success = False # is_success is checked in the 'finally' block |
| 551 | return False |
| 552 | |
| 553 | print('Launched servers. Waiting for %d seconds for the server pods to be ' |
| 554 | 'fully online') % config.global_settings.pod_warmup_secs |
| 555 | time.sleep(config.global_settings.pod_warmup_secs) |
| 556 | |
| 557 | for name, client_pod_spec in config.client_pod_specs_dict.iteritems(): |
| 558 | if not gke.launch_clients(client_pod_spec): |
| 559 | is_success = False # is_success is checked in the 'finally' block |
| 560 | return False |
| 561 | |
| 562 | print('Launched all clients. Waiting for %d seconds for the client pods to ' |
| 563 | 'be fully online') % config.global_settings.pod_warmup_secs |
| 564 | time.sleep(config.global_settings.pod_warmup_secs) |
| 565 | |
| 566 | start_time = datetime.datetime.now() |
| 567 | end_time = start_time + datetime.timedelta( |
| 568 | seconds=config.global_settings.test_duration_secs) |
| 569 | print 'Running the test until %s' % end_time.isoformat() |
| 570 | |
| 571 | while True: |
| 572 | if datetime.datetime.now() > end_time: |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 573 | print 'Test was run for %d seconds' % config.global_settings.test_duration_secs |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 574 | break |
| 575 | |
| 576 | # Check if either stress server or clients have failed (btw, the bq_helper |
| 577 | # monitors all the rows in the summary table and checks if any of them |
| 578 | # have a failure status) |
| 579 | if bq_helper.check_if_any_tests_failed(): |
| 580 | is_success = False |
| 581 | print 'Some tests failed.' |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 582 | break # Don't 'return' here. We still want to call bq_helper to print qps/summary tables |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 583 | |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 584 | # Tests running fine. Wait until next poll time to check the status |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 585 | print 'Sleeping for %d seconds..' % config.global_settings.test_poll_interval_secs |
| 586 | time.sleep(config.global_settings.test_poll_interval_secs) |
| 587 | |
| 588 | # Print BiqQuery tables |
| 589 | bq_helper.print_qps_records() |
| 590 | bq_helper.print_summary_records() |
| 591 | |
| 592 | finally: |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 593 | # If there was a test failure, we should not delete the pods since they |
| 594 | # would contain useful debug information (logs, core dumps etc) |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 595 | if is_success: |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 596 | for name, server_pod_spec in config.server_pod_specs_dict.iteritems(): |
| 597 | gke.delete_servers(server_pod_spec) |
| 598 | for name, client_pod_spec in config.client_pod_specs_dict.iteritems(): |
| 599 | gke.delete_clients(client_pod_spec) |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 600 | |
| 601 | return is_success |
| 602 | |
| 603 | |
| 604 | argp = argparse.ArgumentParser( |
| 605 | description='Launch stress tests in GKE', |
| 606 | formatter_class=argparse.ArgumentDefaultsHelpFormatter) |
Sree Kuchibhotla | 815c589 | 2016-03-25 15:03:50 -0700 | [diff] [blame] | 607 | argp.add_argument('--gcp_project_id', |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 608 | required=True, |
| 609 | help='The Google Cloud Platform Project Id') |
| 610 | argp.add_argument('--config_file', |
| 611 | required=True, |
| 612 | type=str, |
| 613 | help='The test config file') |
| 614 | |
| 615 | if __name__ == '__main__': |
| 616 | args = argp.parse_args() |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 617 | |
| 618 | config_filename = args.config_file |
| 619 | |
Sree Kuchibhotla | 5cadf51 | 2016-03-28 08:55:11 -0700 | [diff] [blame] | 620 | # Since we will be changing the current working directory to grpc root in the |
| 621 | # next step, we should check if the config filename path is a relative path |
| 622 | # (i.e a path relative to the current working directory) and if so, convert it |
| 623 | # to abosulte path |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 624 | if not os.path.isabs(config_filename): |
Sree Kuchibhotla | 5cadf51 | 2016-03-28 08:55:11 -0700 | [diff] [blame] | 625 | config_filename = os.path.abspath(config_filename) |
Sree Kuchibhotla | cdf7734 | 2016-03-25 15:37:34 -0700 | [diff] [blame] | 626 | |
| 627 | config = Config(config_filename, args.gcp_project_id) |
| 628 | |
| 629 | # Change current working directory to grpc root |
| 630 | # (This is important because all relative file paths in the config file are |
| 631 | # supposed to interpreted as relative to the GRPC root) |
| 632 | grpc_root = os.path.abspath(os.path.join( |
| 633 | os.path.dirname(sys.argv[0]), '../../..')) |
| 634 | os.chdir(grpc_root) |
| 635 | |
Sree Kuchibhotla | 575f0fa | 2016-03-25 14:27:07 -0700 | [diff] [blame] | 636 | run_tests(config) |