Import benchmark-tools to main repository.

This has adapted for use with bazel from the original commit
a26e93769ebefd82593a43e22fb13a09717cfa6d.

In particular, the style has been made consistent with internal python style
guidelines, and the packages (including the main entrypoint) have been
refactored in order to allow bazel testing targets.

PiperOrigin-RevId: 283484433

1 files changed, 119 insertions, 0 deletions

diff --git a/benchmarks/suites/__init__.py b/benchmarks/suites/__init__.py
new file mode 100644
index 000000000..360736cc3
--- /dev/null
+++ b/benchmarks/suites/__init__.py
@@ -0,0 +1,119 @@
+# python3
+# Copyright 2019 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Core benchmark annotations."""
+
+import functools
+import inspect
+import types
+from typing import List
+from typing import Tuple
+
+BENCHMARK_METRICS = '__benchmark_metrics__'
+BENCHMARK_MACHINES = '__benchmark_machines__'
+
+
+def is_benchmark(func: types.FunctionType) -> bool:
+  """Returns true if the given function is a benchmark."""
+  return isinstance(func, types.FunctionType) and \
+      hasattr(func, BENCHMARK_METRICS) and \
+      hasattr(func, BENCHMARK_MACHINES)
+
+
+def benchmark_metrics(func: types.FunctionType) -> List[Tuple[str, str]]:
+  """Returns the list of available metrics."""
+  return [(metric.__name__, metric.__doc__)
+          for metric in getattr(func, BENCHMARK_METRICS)]
+
+
+def benchmark_machines(func: types.FunctionType) -> int:
+  """Returns the number of machines required."""
+  return getattr(func, BENCHMARK_MACHINES)
+
+
+# pylint: disable=unused-argument
+def default(value, **kwargs):
+  """Returns the passed value."""
+  return value
+
+
+def benchmark(metrics: List[types.FunctionType] = None,
+              machines: int = 1) -> types.FunctionType:
+  """Define a benchmark function with metrics.
+
+  Args:
+    metrics: A list of metric functions.
+    machines: The number of machines required.
+
+  Returns:
+    A function that accepts the given number of machines, and iteratively
+    returns a set of (metric_name, metric_value) pairs when called repeatedly.
+  """
+  if not metrics:
+    # The default passes through.
+    metrics = [default]
+
+  def decorator(func: types.FunctionType) -> types.FunctionType:
+    """Decorator function."""
+    # Every benchmark should accept at least two parameters:
+    #   runtime: The runtime to use for the benchmark (str, required).
+    #   metrics: The metrics to use, if not the default (str, optional).
+    @functools.wraps(func)
+    def wrapper(*args, runtime: str, metric: list = None, **kwargs):
+      """Wrapper function."""
+      # First -- ensure that we marshall all types appropriately. In
+      # general, we will call this with only strings. These strings will
+      # need to be converted to their underlying types/classes.
+      sig = inspect.signature(func)
+      for param in sig.parameters.values():
+        if param.annotation != inspect.Parameter.empty and \
+           param.name in kwargs and not isinstance(kwargs[param.name], param.annotation):
+          try:
+            # Marshall to the appropriate type.
+            kwargs[param.name] = param.annotation(kwargs[param.name])
+          except Exception as exc:
+            raise ValueError(
+                'illegal type for %s(%s=%s): %s' %
+                (func.__name__, param.name, kwargs[param.name], exc))
+        elif param.default != inspect.Parameter.empty and \
+             param.name not in kwargs:
+          # Ensure that we have the value set, because it will
+          # be passed to the metric function for evaluation.
+          kwargs[param.name] = param.default
+
+      # Next, figure out how to apply a metric. We do this prior to
+      # running the underlying function to prevent having to wait a few
+      # minutes for a result just to see some error.
+      if not metric:
+        # Return all metrics in the iterator.
+        result = func(*args, runtime=runtime, **kwargs)
+        for metric_func in metrics:
+          yield (metric_func.__name__, metric_func(result, **kwargs))
+      else:
+        result = None
+        for single_metric in metric:
+          for metric_func in metrics:
+            # Is this a function that matches the name?
+            # Apply this function to the result.
+            if metric_func.__name__ == single_metric:
+              if not result:
+                # Lazy evaluation: only if metric matches.
+                result = func(*args, runtime=runtime, **kwargs)
+              yield single_metric, metric_func(result, **kwargs)
+
+    # Set metadata on the benchmark (used above).
+    setattr(wrapper, BENCHMARK_METRICS, metrics)
+    setattr(wrapper, BENCHMARK_MACHINES, machines)
+    return wrapper
+
+  return decorator