github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/testing/load_tests/microbenchmarks_test.py (about) 1 # 2 # Licensed to the Apache Software Foundation (ASF) under one or more 3 # contributor license agreements. See the NOTICE file distributed with 4 # this work for additional information regarding copyright ownership. 5 # The ASF licenses this file to You under the Apache License, Version 2.0 6 # (the "License"); you may not use this file except in compliance with 7 # the License. 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 # 17 18 """ 19 This is a load test that runs a set of basic microbenchmarks for the Python SDK 20 and the DirectRunner. 21 22 This test does not need any additional options passed to run, besides the 23 dataset information. 24 25 Example test run: 26 27 python -m apache_beam.testing.load_tests.microbenchmarks_test \ 28 --test-pipeline-options=" 29 --project=big-query-project 30 --input_options='{}' 31 --region=... 32 --publish_to_big_query=true 33 --metrics_dataset=python_load_tests 34 --metrics_table=microbenchmarks" 35 36 or: 37 38 ./gradlew -PloadTest.args=" 39 --publish_to_big_query=true 40 --project=... 41 --region=... 42 --input_options='{}' 43 --metrics_dataset=python_load_tests 44 --metrics_table=microbenchmarks 45 --runner=DirectRunner" \ 46 -PloadTest.mainClass=apache_beam.testing.load_tests.microbenchmarks_test \ 47 -Prunner=DirectRunner :sdks:python:apache_beam:testing:load_tests:run 48 """ 49 50 # pytype: skip-file 51 52 import logging 53 import time 54 55 from apache_beam.testing.load_tests.load_test import LoadTest 56 from apache_beam.tools import fn_api_runner_microbenchmark 57 from apache_beam.tools import teststream_microbenchmark 58 from apache_beam.transforms.util import _BatchSizeEstimator 59 60 61 class MicroBenchmarksLoadTest(LoadTest): 62 def __init__(self): 63 super().__init__() 64 65 def test(self): 66 self.extra_metrics.update(self._run_fn_api_runner_microbenchmark()) 67 self.extra_metrics.update(self._run_teststream_microbenchmark()) 68 69 def _run_teststream_microbenchmark(self): 70 start = time.perf_counter() 71 result = teststream_microbenchmark.run_benchmark(verbose=False) 72 sizes = list(result[0].values())[0] 73 costs = list(result[1].values())[0] 74 a, b = _BatchSizeEstimator.linear_regression_no_numpy(sizes, costs) 75 76 return { 77 'teststream_microbenchmark_runtime_sec': time.perf_counter() - start, 78 'teststream_microbenchmark_fixed_cost_ms': a * 1000, 79 'teststream_microbenchmark_per_element_cost_ms': b * 1000, 80 } 81 82 def _run_fn_api_runner_microbenchmark(self): 83 start = time.perf_counter() 84 result = fn_api_runner_microbenchmark.run_benchmark(verbose=False) 85 sizes = list(result[0].values())[0] 86 costs = list(result[1].values())[0] 87 a, b = _BatchSizeEstimator.linear_regression_no_numpy(sizes, costs) 88 89 return { 90 'fn_api_runner_microbenchmark_runtime_sec': time.perf_counter() - start, 91 'fn_api_runner_microbenchmark_fixed_cost_ms': a * 1000, 92 'fn_api_runner_microbenchmark_per_element_cost_ms': b * 1000, 93 } 94 95 96 if __name__ == '__main__': 97 logging.basicConfig(level=logging.INFO) 98 MicroBenchmarksLoadTest().run()