blob: b190ebde7a4f44d697b7eab4b23515755a6553d6 [file] [log] [blame]
Siddharth Shukla8e64d902017-03-12 19:50:18 +01001#!/usr/bin/env python
Sree Kuchibhotlaaecac882016-04-04 22:48:15 -07002# Copyright 2015-2016, Google Inc.
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -07003# 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.
Siddharth Shuklad194f592017-03-11 19:12:43 +010030
31from __future__ import print_function
32
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -070033import argparse
34import datetime
35import json
36import os
37import subprocess
38import sys
39import time
40
41stress_test_utils_dir = os.path.abspath(os.path.join(
42 os.path.dirname(__file__), '../../gcp/stress_test'))
43sys.path.append(stress_test_utils_dir)
44from stress_test_utils import BigQueryHelper
45
46kubernetes_api_dir = os.path.abspath(os.path.join(
47 os.path.dirname(__file__), '../../gcp/utils'))
48sys.path.append(kubernetes_api_dir)
49
50import kubernetes_api
51
52
53class GlobalSettings:
54
55 def __init__(self, gcp_project_id, build_docker_images,
56 test_poll_interval_secs, test_duration_secs,
57 kubernetes_proxy_port, dataset_id_prefix, summary_table_id,
58 qps_table_id, pod_warmup_secs):
59 self.gcp_project_id = gcp_project_id
60 self.build_docker_images = build_docker_images
61 self.test_poll_interval_secs = test_poll_interval_secs
62 self.test_duration_secs = test_duration_secs
63 self.kubernetes_proxy_port = kubernetes_proxy_port
64 self.dataset_id_prefix = dataset_id_prefix
65 self.summary_table_id = summary_table_id
66 self.qps_table_id = qps_table_id
67 self.pod_warmup_secs = pod_warmup_secs
68
69
70class ClientTemplate:
Sree Kuchibhotla9c9644b2016-03-28 09:30:51 -070071 """ Contains all the common settings that are used by a stress client """
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -070072
Sree Kuchibhotla83a9a652016-04-04 21:49:15 -070073 def __init__(self, name, stress_client_cmd, metrics_client_cmd, metrics_port,
74 wrapper_script_path, poll_interval_secs, client_args_dict,
Sree Kuchibhotla1bdd5312016-05-02 19:57:04 -070075 metrics_args_dict, will_run_forever, env_dict):
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -070076 self.name = name
Sree Kuchibhotla7a054362016-04-04 14:08:02 -070077 self.stress_client_cmd = stress_client_cmd
78 self.metrics_client_cmd = metrics_client_cmd
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -070079 self.metrics_port = metrics_port
80 self.wrapper_script_path = wrapper_script_path
81 self.poll_interval_secs = poll_interval_secs
82 self.client_args_dict = client_args_dict
83 self.metrics_args_dict = metrics_args_dict
Sree Kuchibhotla5bc112c2016-04-18 15:34:04 -070084 self.will_run_forever = will_run_forever
Sree Kuchibhotla1bdd5312016-05-02 19:57:04 -070085 self.env_dict = env_dict
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -070086
87
88class ServerTemplate:
Sree Kuchibhotla9c9644b2016-03-28 09:30:51 -070089 """ Contains all the common settings used by a stress server """
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -070090
Sree Kuchibhotla7a054362016-04-04 14:08:02 -070091 def __init__(self, name, server_cmd, wrapper_script_path, server_port,
Sree Kuchibhotla1bdd5312016-05-02 19:57:04 -070092 server_args_dict, will_run_forever, env_dict):
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -070093 self.name = name
Sree Kuchibhotla7a054362016-04-04 14:08:02 -070094 self.server_cmd = server_cmd
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -070095 self.wrapper_script_path = wrapper_script_path
96 self.server_port = server_port
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -070097 self.server_args_dict = server_args_dict
Sree Kuchibhotla5bc112c2016-04-18 15:34:04 -070098 self.will_run_forever = will_run_forever
Sree Kuchibhotla1bdd5312016-05-02 19:57:04 -070099 self.env_dict = env_dict
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700100
101
102class DockerImage:
Sree Kuchibhotlae68ec432016-03-28 13:55:01 -0700103 """ Represents properties of a Docker image. Provides methods to build the
104 image and push it to GKE registry
Sree Kuchibhotla9c9644b2016-03-28 09:30:51 -0700105 """
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700106
107 def __init__(self, gcp_project_id, image_name, build_script_path,
Sree Kuchibhotla8d41d512016-03-25 14:50:31 -0700108 dockerfile_dir, build_type):
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700109 """Args:
110
111 image_name: The docker image name
112 tag_name: The additional tag name. This is the name used when pushing the
113 docker image to GKE registry
114 build_script_path: The path to the build script that builds this docker
115 image
116 dockerfile_dir: The name of the directory under
117 '<grpc_root>/tools/dockerfile' that contains the dockerfile
118 """
119 self.image_name = image_name
120 self.gcp_project_id = gcp_project_id
121 self.build_script_path = build_script_path
122 self.dockerfile_dir = dockerfile_dir
Sree Kuchibhotla8d41d512016-03-25 14:50:31 -0700123 self.build_type = build_type
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700124 self.tag_name = self._make_tag_name(gcp_project_id, image_name)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700125
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700126 def _make_tag_name(self, project_id, image_name):
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700127 return 'gcr.io/%s/%s' % (project_id, image_name)
128
129 def build_image(self):
Siddharth Shuklad194f592017-03-11 19:12:43 +0100130 print('Building docker image: %s (tag: %s)' % (self.image_name,
131 self.tag_name))
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700132 os.environ['INTEROP_IMAGE'] = self.image_name
133 os.environ['INTEROP_IMAGE_REPOSITORY_TAG'] = self.tag_name
134 os.environ['BASE_NAME'] = self.dockerfile_dir
Sree Kuchibhotla8d41d512016-03-25 14:50:31 -0700135 os.environ['BUILD_TYPE'] = self.build_type
Siddharth Shuklad194f592017-03-11 19:12:43 +0100136 print('DEBUG: path: ', self.build_script_path)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700137 if subprocess.call(args=[self.build_script_path]) != 0:
Siddharth Shuklad194f592017-03-11 19:12:43 +0100138 print('Error in building the Docker image')
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700139 return False
140 return True
141
142 def push_to_gke_registry(self):
143 cmd = ['gcloud', 'docker', 'push', self.tag_name]
Siddharth Shuklad194f592017-03-11 19:12:43 +0100144 print('Pushing %s to the GKE registry..' % self.tag_name)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700145 if subprocess.call(args=cmd) != 0:
Siddharth Shuklad194f592017-03-11 19:12:43 +0100146 print('Error in pushing the image %s to the GKE registry' %
147 self.tag_name)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700148 return False
149 return True
150
151
152class ServerPodSpec:
Sree Kuchibhotla9c9644b2016-03-28 09:30:51 -0700153 """ Contains the information required to launch server pods. """
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700154
155 def __init__(self, name, server_template, docker_image, num_instances):
156 self.name = name
157 self.template = server_template
158 self.docker_image = docker_image
159 self.num_instances = num_instances
160
161 def pod_names(self):
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700162 """ Return a list of names of server pods to create. """
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700163 return ['%s-%d' % (self.name, i) for i in range(1, self.num_instances + 1)]
164
165 def server_addresses(self):
166 """ Return string of server addresses in the following format:
167 '<server_pod_name_1>:<server_port>,<server_pod_name_2>:<server_port>...'
168 """
169 return ','.join(['%s:%d' % (pod_name, self.template.server_port)
170 for pod_name in self.pod_names()])
171
172
173class ClientPodSpec:
Sree Kuchibhotla9c9644b2016-03-28 09:30:51 -0700174 """ Contains the information required to launch client pods """
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700175
176 def __init__(self, name, client_template, docker_image, num_instances,
177 server_addresses):
178 self.name = name
179 self.template = client_template
180 self.docker_image = docker_image
181 self.num_instances = num_instances
182 self.server_addresses = server_addresses
183
184 def pod_names(self):
185 """ Return a list of names of client pods to create """
186 return ['%s-%d' % (self.name, i) for i in range(1, self.num_instances + 1)]
187
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700188 # The client args in the template do not have server addresses. This function
189 # adds the server addresses and returns the updated client args
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700190 def get_client_args_dict(self):
191 args_dict = self.template.client_args_dict.copy()
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700192 args_dict['server_addresses'] = self.server_addresses
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700193 return args_dict
194
195
196class Gke:
Sree Kuchibhotla9c9644b2016-03-28 09:30:51 -0700197 """ Class that has helper methods to interact with GKE """
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700198
199 class KubernetesProxy:
200 """Class to start a proxy on localhost to talk to the Kubernetes API server"""
201
202 def __init__(self, port):
203 cmd = ['kubectl', 'proxy', '--port=%d' % port]
204 self.p = subprocess.Popen(args=cmd)
205 time.sleep(2)
Siddharth Shuklad194f592017-03-11 19:12:43 +0100206 print('\nStarted kubernetes proxy on port: %d' % port)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700207
208 def __del__(self):
209 if self.p is not None:
Siddharth Shuklad194f592017-03-11 19:12:43 +0100210 print('Shutting down Kubernetes proxy..')
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700211 self.p.kill()
212
213 def __init__(self, project_id, run_id, dataset_id, summary_table_id,
214 qps_table_id, kubernetes_port):
215 self.project_id = project_id
216 self.run_id = run_id
217 self.dataset_id = dataset_id
218 self.summary_table_id = summary_table_id
219 self.qps_table_id = qps_table_id
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700220
221 # The environment variables we would like to pass to every pod (both client
222 # and server) launched in GKE
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700223 self.gke_env = {
224 'RUN_ID': self.run_id,
225 'GCP_PROJECT_ID': self.project_id,
226 'DATASET_ID': self.dataset_id,
227 'SUMMARY_TABLE_ID': self.summary_table_id,
228 'QPS_TABLE_ID': self.qps_table_id
229 }
230
231 self.kubernetes_port = kubernetes_port
232 # Start kubernetes proxy
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700233 self.kubernetes_proxy = Gke.KubernetesProxy(kubernetes_port)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700234
235 def _args_dict_to_str(self, args_dict):
236 return ' '.join('--%s=%s' % (k, args_dict[k]) for k in args_dict.keys())
237
238 def launch_servers(self, server_pod_spec):
239 is_success = True
240
241 # The command to run inside the container is the wrapper script (which then
242 # launches the actual server)
243 container_cmd = server_pod_spec.template.wrapper_script_path
244
245 # The parameters to the wrapper script (defined in
246 # server_pod_spec.template.wrapper_script_path) are are injected into the
247 # container via environment variables
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700248 server_env = self.gke_env.copy()
Sree Kuchibhotla1bdd5312016-05-02 19:57:04 -0700249 server_env.update(server_pod_spec.template.env_dict)
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700250 server_env.update({
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700251 'STRESS_TEST_IMAGE_TYPE': 'SERVER',
Sree Kuchibhotla7a054362016-04-04 14:08:02 -0700252 'STRESS_TEST_CMD': server_pod_spec.template.server_cmd,
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700253 'STRESS_TEST_ARGS_STR': self._args_dict_to_str(
Sree Kuchibhotla5bc112c2016-04-18 15:34:04 -0700254 server_pod_spec.template.server_args_dict),
255 'WILL_RUN_FOREVER': str(server_pod_spec.template.will_run_forever)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700256 })
257
258 for pod_name in server_pod_spec.pod_names():
259 server_env['POD_NAME'] = pod_name
Siddharth Shuklad194f592017-03-11 19:12:43 +0100260 print('Creating server: %s' % pod_name)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700261 is_success = kubernetes_api.create_pod_and_service(
262 'localhost',
263 self.kubernetes_port,
264 'default', # Use 'default' namespace
265 pod_name,
266 server_pod_spec.docker_image.tag_name,
267 [server_pod_spec.template.server_port], # Ports to expose on the pod
268 [container_cmd],
269 [], # Args list is empty since we are passing all args via env variables
270 server_env,
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700271 True # Headless = True for server to that GKE creates a DNS record for pod_name
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700272 )
273 if not is_success:
Siddharth Shuklad194f592017-03-11 19:12:43 +0100274 print('Error in launching server: %s' % pod_name)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700275 break
276
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700277 if is_success:
Siddharth Shuklad194f592017-03-11 19:12:43 +0100278 print('Successfully created server(s)')
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700279
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700280 return is_success
281
282 def launch_clients(self, client_pod_spec):
283 is_success = True
284
285 # The command to run inside the container is the wrapper script (which then
286 # launches the actual stress client)
287 container_cmd = client_pod_spec.template.wrapper_script_path
288
289 # The parameters to the wrapper script (defined in
290 # client_pod_spec.template.wrapper_script_path) are are injected into the
291 # container via environment variables
292 client_env = self.gke_env.copy()
Sree Kuchibhotla1bdd5312016-05-02 19:57:04 -0700293 client_env.update(client_pod_spec.template.env_dict)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700294 client_env.update({
295 'STRESS_TEST_IMAGE_TYPE': 'CLIENT',
Sree Kuchibhotla7a054362016-04-04 14:08:02 -0700296 'STRESS_TEST_CMD': client_pod_spec.template.stress_client_cmd,
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700297 'STRESS_TEST_ARGS_STR': self._args_dict_to_str(
298 client_pod_spec.get_client_args_dict()),
Sree Kuchibhotla83a9a652016-04-04 21:49:15 -0700299 'METRICS_CLIENT_CMD': client_pod_spec.template.metrics_client_cmd,
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700300 'METRICS_CLIENT_ARGS_STR': self._args_dict_to_str(
301 client_pod_spec.template.metrics_args_dict),
Sree Kuchibhotla5bc112c2016-04-18 15:34:04 -0700302 'POLL_INTERVAL_SECS': str(client_pod_spec.template.poll_interval_secs),
303 'WILL_RUN_FOREVER': str(client_pod_spec.template.will_run_forever)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700304 })
305
306 for pod_name in client_pod_spec.pod_names():
307 client_env['POD_NAME'] = pod_name
Siddharth Shuklad194f592017-03-11 19:12:43 +0100308 print('Creating client: %s' % pod_name)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700309 is_success = kubernetes_api.create_pod_and_service(
310 'localhost',
311 self.kubernetes_port,
312 'default', # default namespace,
313 pod_name,
314 client_pod_spec.docker_image.tag_name,
315 [client_pod_spec.template.metrics_port], # Ports to expose on the pod
316 [container_cmd],
317 [], # Empty args list since all args are passed via env variables
318 client_env,
Sree Kuchibhotla081b6032017-02-02 16:52:30 -0800319 True # Client is a headless service (no need for an external ip)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700320 )
321
322 if not is_success:
Siddharth Shuklad194f592017-03-11 19:12:43 +0100323 print('Error in launching client %s' % pod_name)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700324 break
325
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700326 if is_success:
Siddharth Shuklad194f592017-03-11 19:12:43 +0100327 print('Successfully created all client(s)')
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700328
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700329 return is_success
330
331 def _delete_pods(self, pod_name_list):
332 is_success = True
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700333 for pod_name in pod_name_list:
Siddharth Shuklad194f592017-03-11 19:12:43 +0100334 print('Deleting %s' % pod_name)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700335 is_success = kubernetes_api.delete_pod_and_service(
336 'localhost',
337 self.kubernetes_port,
338 'default', # default namespace
339 pod_name)
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700340
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700341 if not is_success:
Siddharth Shuklad194f592017-03-11 19:12:43 +0100342 print('Error in deleting pod %s' % pod_name)
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700343 break
344
345 if is_success:
Siddharth Shuklad194f592017-03-11 19:12:43 +0100346 print('Successfully deleted all pods')
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700347
348 return is_success
349
350 def delete_servers(self, server_pod_spec):
351 return self._delete_pods(server_pod_spec.pod_names())
352
353 def delete_clients(self, client_pod_spec):
354 return self._delete_pods(client_pod_spec.pod_names())
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700355
356
357class Config:
358
Sree Kuchibhotla815c5892016-03-25 15:03:50 -0700359 def __init__(self, config_filename, gcp_project_id):
Siddharth Shuklad194f592017-03-11 19:12:43 +0100360 print('Loading configuration...')
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700361 config_dict = self._load_config(config_filename)
362
363 self.global_settings = self._parse_global_settings(config_dict,
Sree Kuchibhotla9c9644b2016-03-28 09:30:51 -0700364 gcp_project_id)
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700365 self.docker_images_dict = self._parse_docker_images(
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700366 config_dict, self.global_settings.gcp_project_id)
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700367 self.client_templates_dict = self._parse_client_templates(config_dict)
368 self.server_templates_dict = self._parse_server_templates(config_dict)
369 self.server_pod_specs_dict = self._parse_server_pod_specs(
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700370 config_dict, self.docker_images_dict, self.server_templates_dict)
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700371 self.client_pod_specs_dict = self._parse_client_pod_specs(
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700372 config_dict, self.docker_images_dict, self.client_templates_dict,
373 self.server_pod_specs_dict)
Siddharth Shuklad194f592017-03-11 19:12:43 +0100374 print('Loaded Configuaration.')
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700375
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700376 def _parse_global_settings(self, config_dict, gcp_project_id):
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700377 global_settings_dict = config_dict['globalSettings']
Sree Kuchibhotla815c5892016-03-25 15:03:50 -0700378 return GlobalSettings(gcp_project_id,
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700379 global_settings_dict['buildDockerImages'],
380 global_settings_dict['pollIntervalSecs'],
381 global_settings_dict['testDurationSecs'],
382 global_settings_dict['kubernetesProxyPort'],
383 global_settings_dict['datasetIdNamePrefix'],
384 global_settings_dict['summaryTableId'],
385 global_settings_dict['qpsTableId'],
386 global_settings_dict['podWarmupSecs'])
387
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700388 def _parse_docker_images(self, config_dict, gcp_project_id):
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700389 """Parses the 'dockerImages' section of the config file and returns a
390 Dictionary of 'DockerImage' objects keyed by docker image names"""
391 docker_images_dict = {}
Sree Kuchibhotla8d41d512016-03-25 14:50:31 -0700392
393 docker_config_dict = config_dict['dockerImages']
394 for image_name in docker_config_dict.keys():
395 build_script_path = docker_config_dict[image_name]['buildScript']
396 dockerfile_dir = docker_config_dict[image_name]['dockerFileDir']
Sree Kuchibhotla7a054362016-04-04 14:08:02 -0700397 build_type = docker_config_dict[image_name].get('buildType', 'opt')
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700398 docker_images_dict[image_name] = DockerImage(gcp_project_id, image_name,
399 build_script_path,
Sree Kuchibhotla8d41d512016-03-25 14:50:31 -0700400 dockerfile_dir, build_type)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700401 return docker_images_dict
402
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700403 def _parse_client_templates(self, config_dict):
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700404 """Parses the 'clientTemplates' section of the config file and returns a
405 Dictionary of 'ClientTemplate' objects keyed by client template names.
406
407 Note: The 'baseTemplates' sub section of the config file contains templates
408 with default values and the 'templates' sub section contains the actual
409 client templates (which refer to the base template name to use for default
410 values).
411 """
412 client_templates_dict = {}
413
414 templates_dict = config_dict['clientTemplates']['templates']
415 base_templates_dict = config_dict['clientTemplates'].get('baseTemplates',
416 {})
417 for template_name in templates_dict.keys():
418 # temp_dict is a temporary dictionary that merges base template dictionary
419 # and client template dictionary (with client template dictionary values
420 # overriding base template values)
421 temp_dict = {}
422
423 base_template_name = templates_dict[template_name].get('baseTemplate')
424 if not base_template_name is None:
425 temp_dict = base_templates_dict[base_template_name].copy()
426
427 temp_dict.update(templates_dict[template_name])
428
429 # Create and add ClientTemplate object to the final client_templates_dict
Sree Kuchibhotla83a9a652016-04-04 21:49:15 -0700430 stress_client_cmd = ' '.join(temp_dict['stressClientCmd'])
431 metrics_client_cmd = ' '.join(temp_dict['metricsClientCmd'])
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700432 client_templates_dict[template_name] = ClientTemplate(
Sree Kuchibhotla83a9a652016-04-04 21:49:15 -0700433 template_name, stress_client_cmd, metrics_client_cmd,
434 temp_dict['metricsPort'], temp_dict['wrapperScriptPath'],
435 temp_dict['pollIntervalSecs'], temp_dict['clientArgs'].copy(),
Sree Kuchibhotla1bdd5312016-05-02 19:57:04 -0700436 temp_dict['metricsArgs'].copy(), temp_dict.get('willRunForever', 1),
437 temp_dict.get('env', {}).copy())
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700438
439 return client_templates_dict
440
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700441 def _parse_server_templates(self, config_dict):
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700442 """Parses the 'serverTemplates' section of the config file and returns a
443 Dictionary of 'serverTemplate' objects keyed by server template names.
444
445 Note: The 'baseTemplates' sub section of the config file contains templates
446 with default values and the 'templates' sub section contains the actual
447 server templates (which refer to the base template name to use for default
448 values).
449 """
450 server_templates_dict = {}
451
452 templates_dict = config_dict['serverTemplates']['templates']
453 base_templates_dict = config_dict['serverTemplates'].get('baseTemplates',
454 {})
455
456 for template_name in templates_dict.keys():
457 # temp_dict is a temporary dictionary that merges base template dictionary
458 # and server template dictionary (with server template dictionary values
459 # overriding base template values)
460 temp_dict = {}
461
462 base_template_name = templates_dict[template_name].get('baseTemplate')
463 if not base_template_name is None:
464 temp_dict = base_templates_dict[base_template_name].copy()
465
466 temp_dict.update(templates_dict[template_name])
467
468 # Create and add ServerTemplate object to the final server_templates_dict
Sree Kuchibhotla83a9a652016-04-04 21:49:15 -0700469 stress_server_cmd = ' '.join(temp_dict['stressServerCmd'])
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700470 server_templates_dict[template_name] = ServerTemplate(
Sree Kuchibhotla83a9a652016-04-04 21:49:15 -0700471 template_name, stress_server_cmd, temp_dict['wrapperScriptPath'],
Sree Kuchibhotla5bc112c2016-04-18 15:34:04 -0700472 temp_dict['serverPort'], temp_dict['serverArgs'].copy(),
Sree Kuchibhotla1bdd5312016-05-02 19:57:04 -0700473 temp_dict.get('willRunForever', 1), temp_dict.get('env', {}).copy())
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700474
475 return server_templates_dict
476
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700477 def _parse_server_pod_specs(self, config_dict, docker_images_dict,
Sree Kuchibhotla9c9644b2016-03-28 09:30:51 -0700478 server_templates_dict):
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700479 """Parses the 'serverPodSpecs' sub-section (under 'testMatrix' section) of
480 the config file and returns a Dictionary of 'ServerPodSpec' objects keyed
481 by server pod spec names"""
482 server_pod_specs_dict = {}
483
484 pod_specs_dict = config_dict['testMatrix'].get('serverPodSpecs', {})
485
486 for pod_name in pod_specs_dict.keys():
487 server_template_name = pod_specs_dict[pod_name]['serverTemplate']
488 docker_image_name = pod_specs_dict[pod_name]['dockerImage']
489 num_instances = pod_specs_dict[pod_name].get('numInstances', 1)
490
491 # Create and add the ServerPodSpec object to the final
492 # server_pod_specs_dict
493 server_pod_specs_dict[pod_name] = ServerPodSpec(
494 pod_name, server_templates_dict[server_template_name],
495 docker_images_dict[docker_image_name], num_instances)
496
497 return server_pod_specs_dict
498
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700499 def _parse_client_pod_specs(self, config_dict, docker_images_dict,
Sree Kuchibhotla9c9644b2016-03-28 09:30:51 -0700500 client_templates_dict, server_pod_specs_dict):
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700501 """Parses the 'clientPodSpecs' sub-section (under 'testMatrix' section) of
502 the config file and returns a Dictionary of 'ClientPodSpec' objects keyed
503 by client pod spec names"""
504 client_pod_specs_dict = {}
505
506 pod_specs_dict = config_dict['testMatrix'].get('clientPodSpecs', {})
507 for pod_name in pod_specs_dict.keys():
508 client_template_name = pod_specs_dict[pod_name]['clientTemplate']
509 docker_image_name = pod_specs_dict[pod_name]['dockerImage']
510 num_instances = pod_specs_dict[pod_name]['numInstances']
511
512 # Get the server addresses from the server pod spec object
513 server_pod_spec_name = pod_specs_dict[pod_name]['serverPodSpec']
514 server_addresses = server_pod_specs_dict[
515 server_pod_spec_name].server_addresses()
516
517 client_pod_specs_dict[pod_name] = ClientPodSpec(
518 pod_name, client_templates_dict[client_template_name],
519 docker_images_dict[docker_image_name], num_instances,
520 server_addresses)
521
522 return client_pod_specs_dict
523
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700524 def _load_config(self, config_filename):
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700525 """Opens the config file and converts the Json text to Dictionary"""
526 if not os.path.isabs(config_filename):
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700527 raise Exception('Config objects expects an absolute file path. '
528 'config file name passed: %s' % config_filename)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700529 with open(config_filename) as config_file:
530 return json.load(config_file)
531
532
533def run_tests(config):
Sree Kuchibhotla9c9644b2016-03-28 09:30:51 -0700534 """ The main function that launches the stress tests """
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700535 # Build docker images and push to GKE registry
536 if config.global_settings.build_docker_images:
537 for name, docker_image in config.docker_images_dict.iteritems():
538 if not (docker_image.build_image() and
539 docker_image.push_to_gke_registry()):
540 return False
541
542 # Create a unique id for this run (Note: Using timestamp instead of UUID to
543 # make it easier to deduce the date/time of the run just by looking at the run
544 # run id. This is useful in debugging when looking at records in Biq query)
545 run_id = datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
546 dataset_id = '%s_%s' % (config.global_settings.dataset_id_prefix, run_id)
Siddharth Shuklad194f592017-03-11 19:12:43 +0100547 print('Run id:', run_id)
548 print('Dataset id:', dataset_id)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700549
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700550 bq_helper = BigQueryHelper(run_id, '', '',
551 config.global_settings.gcp_project_id, dataset_id,
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700552 config.global_settings.summary_table_id,
553 config.global_settings.qps_table_id)
554 bq_helper.initialize()
555
556 gke = Gke(config.global_settings.gcp_project_id, run_id, dataset_id,
557 config.global_settings.summary_table_id,
558 config.global_settings.qps_table_id,
559 config.global_settings.kubernetes_proxy_port)
560
561 is_success = True
562
563 try:
Siddharth Shuklad194f592017-03-11 19:12:43 +0100564 print('Launching servers..')
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700565 for name, server_pod_spec in config.server_pod_specs_dict.iteritems():
566 if not gke.launch_servers(server_pod_spec):
567 is_success = False # is_success is checked in the 'finally' block
568 return False
569
570 print('Launched servers. Waiting for %d seconds for the server pods to be '
571 'fully online') % config.global_settings.pod_warmup_secs
572 time.sleep(config.global_settings.pod_warmup_secs)
573
574 for name, client_pod_spec in config.client_pod_specs_dict.iteritems():
575 if not gke.launch_clients(client_pod_spec):
576 is_success = False # is_success is checked in the 'finally' block
577 return False
578
579 print('Launched all clients. Waiting for %d seconds for the client pods to '
580 'be fully online') % config.global_settings.pod_warmup_secs
581 time.sleep(config.global_settings.pod_warmup_secs)
582
583 start_time = datetime.datetime.now()
584 end_time = start_time + datetime.timedelta(
585 seconds=config.global_settings.test_duration_secs)
Siddharth Shuklad194f592017-03-11 19:12:43 +0100586 print('Running the test until %s' % end_time.isoformat())
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700587
588 while True:
589 if datetime.datetime.now() > end_time:
Siddharth Shuklad194f592017-03-11 19:12:43 +0100590 print('Test was run for %d seconds' %
591 config.global_settings.test_duration_secs)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700592 break
593
594 # Check if either stress server or clients have failed (btw, the bq_helper
595 # monitors all the rows in the summary table and checks if any of them
596 # have a failure status)
597 if bq_helper.check_if_any_tests_failed():
598 is_success = False
Siddharth Shuklad194f592017-03-11 19:12:43 +0100599 print('Some tests failed.')
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700600 break # Don't 'return' here. We still want to call bq_helper to print qps/summary tables
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700601
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700602 # Tests running fine. Wait until next poll time to check the status
Siddharth Shuklad194f592017-03-11 19:12:43 +0100603 print('Sleeping for %d seconds..' %
604 config.global_settings.test_poll_interval_secs)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700605 time.sleep(config.global_settings.test_poll_interval_secs)
606
607 # Print BiqQuery tables
608 bq_helper.print_qps_records()
609 bq_helper.print_summary_records()
610
611 finally:
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700612 # If there was a test failure, we should not delete the pods since they
613 # would contain useful debug information (logs, core dumps etc)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700614 if is_success:
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700615 for name, server_pod_spec in config.server_pod_specs_dict.iteritems():
616 gke.delete_servers(server_pod_spec)
617 for name, client_pod_spec in config.client_pod_specs_dict.iteritems():
618 gke.delete_clients(client_pod_spec)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700619
620 return is_success
621
622
Sree Kuchibhotla478bd442016-04-12 17:40:24 -0700623def tear_down(config):
624 gke = Gke(config.global_settings.gcp_project_id, '', '',
625 config.global_settings.summary_table_id,
626 config.global_settings.qps_table_id,
627 config.global_settings.kubernetes_proxy_port)
628 for name, server_pod_spec in config.server_pod_specs_dict.iteritems():
629 gke.delete_servers(server_pod_spec)
630 for name, client_pod_spec in config.client_pod_specs_dict.iteritems():
631 gke.delete_clients(client_pod_spec)
632
633
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700634argp = argparse.ArgumentParser(
635 description='Launch stress tests in GKE',
636 formatter_class=argparse.ArgumentDefaultsHelpFormatter)
Sree Kuchibhotla815c5892016-03-25 15:03:50 -0700637argp.add_argument('--gcp_project_id',
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700638 required=True,
639 help='The Google Cloud Platform Project Id')
640argp.add_argument('--config_file',
641 required=True,
642 type=str,
643 help='The test config file')
Sree Kuchibhotla478bd442016-04-12 17:40:24 -0700644argp.add_argument('--tear_down', action='store_true', default=False)
Sree Kuchibhotla575f0fa2016-03-25 14:27:07 -0700645
646if __name__ == '__main__':
647 args = argp.parse_args()
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700648
649 config_filename = args.config_file
650
Sree Kuchibhotla5cadf512016-03-28 08:55:11 -0700651 # Since we will be changing the current working directory to grpc root in the
652 # next step, we should check if the config filename path is a relative path
653 # (i.e a path relative to the current working directory) and if so, convert it
654 # to abosulte path
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700655 if not os.path.isabs(config_filename):
Sree Kuchibhotla5cadf512016-03-28 08:55:11 -0700656 config_filename = os.path.abspath(config_filename)
Sree Kuchibhotlacdf77342016-03-25 15:37:34 -0700657
658 config = Config(config_filename, args.gcp_project_id)
659
660 # Change current working directory to grpc root
661 # (This is important because all relative file paths in the config file are
662 # supposed to interpreted as relative to the GRPC root)
663 grpc_root = os.path.abspath(os.path.join(
664 os.path.dirname(sys.argv[0]), '../../..'))
665 os.chdir(grpc_root)
666
Sree Kuchibhotla478bd442016-04-12 17:40:24 -0700667 # Note that tear_down is only in cases where we want to manually tear down a
668 # test that for some reason run_tests() could not cleanup
669 if args.tear_down:
670 tear_down(config)
671 sys.exit(1)
672
Sree Kuchibhotla247b34f2016-03-29 16:56:14 -0700673 if not run_tests(config):
674 sys.exit(1)