Update commandline and get local runs working.

PiperOrigin-RevId: 289937063

1 files changed, 18 insertions, 41 deletions

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.