Internal change.

PiperOrigin-RevId: 303773475

3 files changed, 23 insertions, 11 deletions

diff --git a/benchmarks/harness/machine_producers/gcloud_producer.py b/benchmarks/harness/machine_producers/gcloud_producer.py
index 1a624df2e..44d72f575 100644
--- a/benchmarks/harness/machine_producers/gcloud_producer.py
+++ b/benchmarks/harness/machine_producers/gcloud_producer.py
@@ -53,6 +53,8 @@ class GCloudProducer(machine_producer.MachineProducer):
     ssh_key_file: path to a valid ssh private key. See README on vaild ssh keys.
     ssh_user: string of user name for ssh_key
     ssh_password: string of password for ssh key
+    internal: if true, use internal IPs of instances. Used if bm-tools is
+    running on a GCP vm when a firewall is set for external IPs.
     mock: a mock printer which will print mock data if required. Mock data is
       recorded output from subprocess calls (returncode, stdout, args).
     condition: mutex for this class around machine creation and deleteion.
@@ -66,6 +68,7 @@ class GCloudProducer(machine_producer.MachineProducer):
                ssh_key_file: str,
                ssh_user: str,
                ssh_password: str,
+               internal: bool,
                mock: gcloud_mock_recorder.MockPrinter = None):
     self.image = image
     self.zone = zone
@@ -74,6 +77,7 @@ class GCloudProducer(machine_producer.MachineProducer):
     self.ssh_key_file = ssh_key_file
     self.ssh_user = ssh_user
     self.ssh_password = ssh_password
+    self.internal = internal
     self.mock = mock
     self.condition = threading.Condition()
 
@@ -129,15 +133,13 @@ class GCloudProducer(machine_producer.MachineProducer):
     machines = []
     for instance in instances:
       name = instance["name"]
+      external = instance["networkInterfaces"][0]["accessConfigs"][0]["natIP"]
+      internal = instance["networkInterfaces"][0]["networkIP"]
       kwargs = {
-          "hostname":
-              instance["networkInterfaces"][0]["accessConfigs"][0]["natIP"],
-          "key_path":
-              self.ssh_key_file,
-          "username":
-              self.ssh_user,
-          "key_password":
-              self.ssh_password
+          "hostname": internal if self.internal else external,
+          "key_path": self.ssh_key_file,
+          "username": self.ssh_user,
+          "key_password": self.ssh_password
       }
       machines.append(machine.RemoteMachine(name=name, **kwargs))
     return machines
@@ -190,6 +192,8 @@ class GCloudProducer(machine_producer.MachineProducer):
     for name in names:
       cmd = "gcloud compute ssh {user}@{name}".format(
           user=self.ssh_user, name=name).split(" ")
+      if self.internal:
+        cmd.append("--internal-ip")
       cmd.append("--ssh-key-file={key}".format(key=self.ssh_key_file))
       cmd.append("--zone={zone}".format(zone=self.zone))
       cmd.append("--command=uname")
diff --git a/benchmarks/runner/__init__.py b/benchmarks/runner/__init__.py
index ba27dc69f..ca785a148 100644
--- a/benchmarks/runner/__init__.py
+++ b/benchmarks/runner/__init__.py
@@ -120,8 +120,8 @@ def run_mock(ctx, **kwargs):
 
 @runner.command("run-gcp", commands.GCPCommand)
 @click.pass_context
-def run_gcp(ctx, image_file: str, zone_file: str, machine_type: str,
-            installers: List[str], **kwargs):
+def run_gcp(ctx, image_file: str, zone_file: str, internal: bool,
+            machine_type: str, installers: List[str], **kwargs):
   """Runs all benchmarks on GCP instances."""
 
   # Resolve all files.
@@ -137,7 +137,8 @@ def run_gcp(ctx, image_file: str, zone_file: str, machine_type: str,
       installers,
       ssh_key_file=key_file,
       ssh_user=harness.DEFAULT_USER,
-      ssh_password="")
+      ssh_password="",
+      internal=internal)
 
   try:
     run(ctx, producer, **kwargs)
diff --git a/benchmarks/runner/commands.py b/benchmarks/runner/commands.py
index 0fccb2fad..194804527 100644
--- a/benchmarks/runner/commands.py
+++ b/benchmarks/runner/commands.py
@@ -111,6 +111,12 @@ class GCPCommand(RunCommand):
         default=os.path.join(
             os.path.dirname(__file__), "../../tools/images/zone.txt"),
     )
+    internal = click.core.Option(
+        ("--internal/--no-internal",),
+        help="""Use instance internal IPs. Used if bm-tools runner is running on
+        GCP instance with firewall rules blocking external IPs.""",
+        default=False,
+    )
     installers = click.core.Option(
         ("--installers",),
         help="The set of installers to use.",
@@ -124,6 +130,7 @@ class GCPCommand(RunCommand):
     self.params.extend([
         image_file,
         zone_file,
+        internal,
         machine_type,
         installers,
     ])