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()