1#!/usr/bin/env python3 2 3# Copyright 2023 gRPC authors. 4# 5# Licensed under the Apache License, Version 2.0 (the "License"); 6# you may not use this file except in compliance with the License. 7# You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, software 12# distributed under the License is distributed on an "AS IS" BASIS, 13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14# See the License for the specific language governing permissions and 15# limitations under the License. 16""" 17Local QPS benchmark runner for the OSS Benchmark loadtest configurations. 18 19This tool will run a scenario locally, either already extracted from 20scenario_config_exporter, or extracted from a benchmark loadtest config. The 21driver, client, and server all in the same process. You can run the process 22under a custom runner using the --runner_cmd="<COMMAND>" flag, and with custom 23environment variables if needed. 24 25This example will run an optimized build of the loadtest under gdb 26 27GRPC_VERBOSITY=debug \ 28 bazel run \ 29 --config=opt \ 30 --cxxopt="-gmlt" \ 31 test/cpp/qps:scenario_runner -- \ 32 --loadtest_file=/path/to/loadtest.config \ 33 --runner_cmd="gdb --args" 34 35This builds the binary and runs: 36 37 gdb --args bazel-bin/.../scenario_runner -- \ 38 --loadtest_config=/tmp/path/extracted_scenario_json.config 39 40 41If you have already extracted the JSON scenario using scenario_config_exporter, 42you can replace `--loadtest_file=loadtest.yaml` with 43`--scenario_file=scenario.json`. 44 45 46Other --runner_cmd examples: 47 --runner_cmd="perf record -F 777 -o $(pwd)/perf.data -g --event=cpu-cycles", 48 --runner_cmd="perf stat record -o $(pwd)/perf.stat.data", 49" 50""" 51 52 53import os 54import subprocess 55import sys 56import tempfile 57 58from absl import app 59from absl import flags 60import yaml 61 62_LOADTEST_YAML = flags.DEFINE_string( 63 "loadtest_file", default=None, help="Path to the benchmark loadtest file" 64) 65_SCENARIO_JSON = flags.DEFINE_string( 66 "scenario_file", default=None, help="Path to a scenario JSON file" 67) 68_RUNNER_CMD = flags.DEFINE_string( 69 "runner_cmd", 70 default="", 71 help="Run the scearnio runner under a custom command (example: bazel ... --cmd='perf lock record -o $(pwd)/out')", 72) 73_RUN_FIRST = flags.DEFINE_bool( 74 "run_first", 75 default=False, 76 help="Only run the first scenario in the loadtest", 77) 78_RUN_ALL = flags.DEFINE_bool( 79 "run_all", default=False, help="Run all scenarios in the loadtest" 80) 81 82 83def run_command(filename): 84 cmd = [ 85 os.path.join( 86 os.path.dirname(os.path.abspath(__file__)), 87 "scenario_runner_cc", 88 ), 89 "--loadtest_config", 90 filename, 91 ] 92 if _RUNNER_CMD.value: 93 cmd = _RUNNER_CMD.value.split(" ") + cmd 94 print(cmd) 95 subprocess.run(cmd, check=True) 96 if _RUN_FIRST.value: 97 print("Exiting due to --run_first") 98 sys.exit(0) 99 100 101def run_loadtests(): 102 loadtests = [] 103 with open( 104 os.path.join( 105 os.path.dirname(os.path.abspath(__file__)), _LOADTEST_YAML.value 106 ) 107 ) as f: 108 loadtests = list(yaml.safe_load_all(f)) 109 if len(loadtests) > 1 and not (_RUN_FIRST.value or _RUN_ALL.value): 110 print( 111 "The loadtest configuration file contains more than one loadtest. Please specify --run_first or --run_all.", 112 file=sys.stderr, 113 ) 114 sys.exit(1) 115 for loadtest in loadtests: 116 with tempfile.NamedTemporaryFile() as tmp_f: 117 tmp_f.write( 118 "".join(loadtest["spec"]["scenariosJSON"]).encode("utf-8") 119 ) 120 tmp_f.flush() 121 run_command(tmp_f.name) 122 123 124def run_scenario_file(): 125 run_command(_SCENARIO_JSON.value) 126 127 128def main(args): 129 if _LOADTEST_YAML.value: 130 run_loadtests() 131 elif _SCENARIO_JSON.value: 132 run_scenario_file() 133 else: 134 "You must provide either a scenario.json or loadtest.yaml" 135 136 137if __name__ == "__main__": 138 app.run(main) 139