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