Add run-gcp command.

Add command to run benchmarks on GCP backed machines
using the gcloud producer.

Run with:
`bazel run :benchmarks -- run-gcp [BENCHMARK_NAME]`

Tested with the startup benchmark.

PiperOrigin-RevId: 290126444

2 files changed, 111 insertions, 0 deletions

diff --git a/benchmarks/runner/__init__.py b/benchmarks/runner/__init__.py
index 6f56704d8..ba80d83d7 100644
--- a/benchmarks/runner/__init__.py
+++ b/benchmarks/runner/__init__.py
@@ -15,10 +15,13 @@
 
 import copy
 import csv
+import json
 import logging
+import os
 import pkgutil
 import pydoc
 import re
+import subprocess
 import sys
 import types
 from typing import List
@@ -26,8 +29,10 @@ from typing import Tuple
 
 import click
 
+from benchmarks import harness
 from benchmarks import suites
 from benchmarks.harness import benchmark_driver
+from benchmarks.harness.machine_producers import gcloud_producer
 from benchmarks.harness.machine_producers import machine_producer
 from benchmarks.harness.machine_producers import mock_producer
 from benchmarks.harness.machine_producers import yaml_producer
@@ -116,6 +121,61 @@ def run_mock(ctx, **kwargs):
   run(ctx, mock_producer.MockMachineProducer(), **kwargs)
 
 
+@runner.command("run-gcp", commands.GCPCommand)
+@click.pass_context
+def run_gcp(ctx, project: str, ssh_key_file: str, image: str,
+            image_project: str, machine_type: str, zone: str, ssh_user: str,
+            ssh_password: str, **kwargs):
+  """Runs all benchmarks on GCP instances."""
+
+  if not ssh_user:
+    ssh_user = harness.DEFAULT_USER
+
+  # Get the default project if one was not provided.
+  if not project:
+    sub = subprocess.run(
+        "gcloud config get-value project".split(" "), stdout=subprocess.PIPE)
+    if sub.returncode:
+      raise ValueError(
+          "Cannot get default project from gcloud. Is it configured>")
+    project = sub.stdout.decode("utf-8").strip("\n")
+
+  if not image_project:
+    image_project = project
+
+  # Check that the ssh-key exists and is readable.
+  if not os.access(ssh_key_file, os.R_OK):
+    raise ValueError(
+        "ssh key given `{ssh_key}` is does not exist or is not readable."
+        .format(ssh_key=ssh_key_file))
+
+  # Check that the image exists.
+  sub = subprocess.run(
+      "gcloud compute images describe {image} --project {image_project} --format=json"
+      .format(image=image, image_project=image_project).split(" "),
+      stdout=subprocess.PIPE)
+  if sub.returncode or "READY" not in json.loads(sub.stdout)["status"]:
+    raise ValueError(
+        "given image was not found or is not ready: {image} {image_project}."
+        .format(image=image, image_project=image_project))
+
+  # Check and set zone to default.
+  if not zone:
+    sub = subprocess.run(
+        "gcloud config get-value compute/zone".split(" "),
+        stdout=subprocess.PIPE)
+    if sub.returncode:
+      raise ValueError(
+          "Default zone is not set in gcloud. Set one or pass a zone with the --zone flag."
+      )
+    zone = sub.stdout.decode("utf-8").strip("\n")
+
+  producer = gcloud_producer.GCloudProducer(project, ssh_key_file, image,
+                                            image_project, machine_type, zone,
+                                            ssh_user, ssh_password)
+  run(ctx, producer, **kwargs)
+
+
 def run(ctx, producer: machine_producer.MachineProducer, method: str, runs: int,
         runtime: List[str], metric: List[str], stat: str, **kwargs):
   """Runs arbitrary benchmarks.
diff --git a/benchmarks/runner/commands.py b/benchmarks/runner/commands.py
index 4973843b9..7ab12fac6 100644
--- a/benchmarks/runner/commands.py
+++ b/benchmarks/runner/commands.py
@@ -24,6 +24,8 @@ def run_mock(**kwargs):
 """
 import click
 
+from benchmarks import harness
+
 
 class RunCommand(click.core.Command):
   """Base Run Command with flags.
@@ -82,3 +84,52 @@ class LocalCommand(RunCommand):
             ("--limit",),
             default=1,
             help="Limit of number of benchmarks that can run at a given time."))
+
+
+class GCPCommand(RunCommand):
+  """GCPCommand inherits all flags from RunCommand and adds flags for run_gcp method.
+
+  Attributes:
+    project: GCP project
+    ssh_key_path: path to the ssh-key to use for the run
+    image: name of the image to build machines from
+    image_project: GCP project under which to find image
+    zone: a GCP zone (e.g. us-west1-b)
+    ssh_user: username to use for the ssh-key
+    ssh_password: password to use for the ssh-key
+  """
+
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+
+    project = click.core.Option(
+        ("--project",),
+        help="Project to run on if not default value given by 'gcloud config get-value project'."
+    )
+    ssh_key_path = click.core.Option(
+        ("--ssh-key-file",),
+        help="Path to a valid ssh private key to use. See README on generating a valid ssh key. Set to ~/.ssh/benchmark-tools by default.",
+        default=harness.DEFAULT_USER_HOME + "/.ssh/benchmark-tools")
+    image = click.core.Option(("--image",),
+                              help="The image on which to build VMs.",
+                              default="bm-tools-testing")
+    image_project = click.core.Option(
+        ("--image_project",),
+        help="The project under which the image to be used is listed.",
+        default="")
+    machine_type = click.core.Option(("--machine_type",),
+                                     help="Type to make all machines.",
+                                     default="n1-standard-4")
+    zone = click.core.Option(("--zone",),
+                             help="The GCP zone to run on.",
+                             default="")
+    ssh_user = click.core.Option(("--ssh-user",),
+                                 help="User for the ssh key.",
+                                 default=harness.DEFAULT_USER)
+    ssh_password = click.core.Option(("--ssh-password",),
+                                     help="Password for the ssh key.",
+                                     default="")
+    self.params.extend([
+        project, ssh_key_path, image, image_project, machine_type, zone,
+        ssh_user, ssh_password
+    ])