Update commandline and get local runs working.

PiperOrigin-RevId: 289937063

4 files changed, 113 insertions, 42 deletions

diff --git a/benchmarks/runner/BUILD b/benchmarks/runner/BUILD
index de24824cc..e1b2ea550 100644
--- a/benchmarks/runner/BUILD
+++ b/benchmarks/runner/BUILD
@@ -10,7 +10,9 @@ py_library(
     ],
     visibility = ["//benchmarks:__pkg__"],
     deps = [
+        ":commands",
         "//benchmarks/harness:benchmark_driver",
+        "//benchmarks/harness/machine_producers:machine_producer",
         "//benchmarks/harness/machine_producers:mock_producer",
         "//benchmarks/harness/machine_producers:yaml_producer",
         "//benchmarks/suites",
@@ -30,6 +32,14 @@ py_library(
     ],
 )
 
+py_library(
+    name = "commands",
+    srcs = ["commands.py"],
+    deps = [
+        requirement("click", True),
+    ],
+)
+
 py_test(
     name = "runner_test",
     srcs = ["runner_test.py"],
diff --git a/benchmarks/runner/__init__.py b/benchmarks/runner/__init__.py
index 9bf9cfd65..6f56704d8 100644
--- a/benchmarks/runner/__init__.py
+++ b/benchmarks/runner/__init__.py
@@ -28,8 +28,10 @@ import click
 
 from benchmarks import suites
 from benchmarks.harness import benchmark_driver
+from benchmarks.harness.machine_producers import machine_producer
 from benchmarks.harness.machine_producers import mock_producer
 from benchmarks.harness.machine_producers import yaml_producer
+from benchmarks.runner import commands
 
 
 @click.group()
@@ -100,30 +102,22 @@ def list_all(method):
     print("\n")
 
 
-# pylint: disable=too-many-arguments
-# pylint: disable=too-many-branches
-# pylint: disable=too-many-locals
-@runner.command(
-    context_settings=dict(ignore_unknown_options=True, allow_extra_args=True))
+@runner.command("run-local", commands.LocalCommand)
 @click.pass_context
-@click.argument("method")
-@click.option("--mock/--no-mock", default=False, help="Mock the machines.")
-@click.option("--env", default=None, help="Specify a yaml file with machines.")
-@click.option(
-    "--runtime", default=["runc"], help="The runtime to use.", multiple=True)
-@click.option("--metric", help="The metric to extract.", multiple=True)
-@click.option(
-    "--runs", default=1, help="The number of times to run each benchmark.")
-@click.option(
-    "--stat",
-    default="median",
-    help="How to aggregate the data from all runs."
-    "\nmedian - returns the median of all runs (default)"
-    "\nall - returns all results comma separated"
-    "\nmeanstd - returns result as mean,std")
-# pylint: disable=too-many-statements
-def run(ctx, method: str, runs: int, env: str, mock: bool, runtime: List[str],
-        metric: List[str], stat: str, **kwargs):
+def run_local(ctx, limit: float, **kwargs):
+  """Runs benchmarks locally."""
+  run(ctx, machine_producer.LocalMachineProducer(limit=limit), **kwargs)
+
+
+@runner.command("run-mock", commands.RunCommand)
+@click.pass_context
+def run_mock(ctx, **kwargs):
+  """Runs benchmarks on Mock machines. Used for testing."""
+  run(ctx, mock_producer.MockMachineProducer(), **kwargs)
+
+
+def run(ctx, producer: machine_producer.MachineProducer, method: str, runs: int,
+        runtime: List[str], metric: List[str], stat: str, **kwargs):
   """Runs arbitrary benchmarks.
 
   All unknown command line flags are passed through to the underlying benchmark
@@ -139,16 +133,13 @@ def run(ctx, method: str, runs: int, env: str, mock: bool, runtime: List[str],
   All benchmarks are run in parallel where possible, but have exclusive
   ownership over the individual machines.
 
-  Exactly one of the --mock and --env flag must be specified.
-
   Every benchmark method will be run the times indicated by --runs.
 
   Args:
     ctx: Click context.
+    producer: A Machine Producer from which to get Machines.
     method: A regular expression for methods to be run.
     runs: Number of runs.
-    env: Environment to use.
-    mock: If true, use mocked environment (supercedes env).
     runtime: A list of runtimes to test.
     metric: A list of metrics to extract.
     stat: The class of statistics to extract.
@@ -218,20 +209,6 @@ def run(ctx, method: str, runs: int, env: str, mock: bool, runtime: List[str],
     sys.exit(1)
   fold("method", list(methods.keys()), allow_flatten=True)
 
-  # Construct the environment.
-  if mock and env:
-    # You can't provide both.
-    logging.error("both --mock and --env are set: which one is it?")
-    sys.exit(1)
-  elif mock:
-    producer = mock_producer.MockMachineProducer()
-  elif env:
-    producer = yaml_producer.YamlMachineProducer(env)
-  else:
-    # You must provide one of mock or env.
-    logging.error("no enviroment provided: use --mock or --env.")
-    sys.exit(1)
-
   # Spin up the drivers.
   #
   # We ensure that metric is the last entry, because we have special behavior.
diff --git a/benchmarks/runner/commands.py b/benchmarks/runner/commands.py
new file mode 100644
index 000000000..4973843b9
--- /dev/null
+++ b/benchmarks/runner/commands.py
@@ -0,0 +1,84 @@
+# 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.
+"""Module with the guts of `click` commands.
+
+Overrides of the click.core.Command. This is done so flags are inherited between
+similar commands (the run command). The classes below are meant to be used in
+click templates like so.
+
+@runner.command("run-mock", RunCommand)
+def run_mock(**kwargs):
+  # mock implementation
+
+"""
+import click
+
+
+class RunCommand(click.core.Command):
+  """Base Run Command with flags.
+
+  Attributes:
+    method: regex of which suite to choose (e.g. sysbench would run
+      sysbench.cpu, sysbench.memory, and sysbench.mutex) See list command for
+      details.
+    metric: metric(s) to extract. See list command for details.
+    runtime: the runtime(s) on which to run.
+    runs: the number of runs to do of each method.
+    stat: how to compile results in the case of multiple run (e.g. median).
+  """
+
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    method = click.core.Argument(("method",))
+
+    metric = click.core.Option(("--metric",),
+                               help="The metric to extract.",
+                               multiple=True)
+
+    runtime = click.core.Option(("--runtime",),
+                                default=["runc"],
+                                help="The runtime to use.",
+                                multiple=True)
+    runs = click.core.Option(("--runs",),
+                             default=1,
+                             help="The number of times to run each benchmark.")
+    stat = click.core.Option(
+        ("--stat",),
+        default="median",
+        help="How to aggregate the data from all runs."
+        "\nmedian - returns the median of all runs (default)"
+        "\nall - returns all results comma separated"
+        "\nmeanstd - returns result as mean,std")
+    self.params.extend([method, runtime, runs, stat, metric])
+    self.ignore_unknown_options = True
+    self.allow_extra_args = True
+
+
+class LocalCommand(RunCommand):
+  """LocalCommand inherits all flags from RunCommand.
+
+  Attributes:
+    limit: limits the number of machines on which to run benchmarks. This limits
+      for local how many benchmarks may run at a time. e.g. "startup" requires
+      one machine -- passing two machines would limit two startup jobs at a
+      time. Default is infinity.
+  """
+
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+    self.params.append(
+        click.core.Option(
+            ("--limit",),
+            default=1,
+            help="Limit of number of benchmarks that can run at a given time."))
diff --git a/benchmarks/runner/runner_test.py b/benchmarks/runner/runner_test.py
index 5719c2838..7818d631a 100644
--- a/benchmarks/runner/runner_test.py
+++ b/benchmarks/runner/runner_test.py
@@ -49,7 +49,7 @@ def test_list():
 
 def test_run():
   cli_runner = testing.CliRunner()
-  result = cli_runner.invoke(runner.runner, ["run", "--mock", "."])
+  result = cli_runner.invoke(runner.runner, ["run-mock", "."])
   print(result.output)
   assert result.exit_code == 0