Import benchmark-tools to main repository.

This has adapted for use with bazel from the original commit
a26e93769ebefd82593a43e22fb13a09717cfa6d.

In particular, the style has been made consistent with internal python style
guidelines, and the packages (including the main entrypoint) have been
refactored in order to allow bazel testing targets.

PiperOrigin-RevId: 283484433

1 files changed, 181 insertions, 0 deletions

diff --git a/benchmarks/harness/container.py b/benchmarks/harness/container.py
new file mode 100644
index 000000000..585436e20
--- /dev/null
+++ b/benchmarks/harness/container.py
@@ -0,0 +1,181 @@
+# 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.
+"""Container definitions."""
+
+import contextlib
+import logging
+import pydoc
+import types
+from typing import Tuple
+
+import docker
+import docker.errors
+
+from benchmarks import workloads
+
+
+class Container:
+  """Abstract container.
+
+  Must be a context manager.
+
+  Usage:
+
+    with Container(client, image, ...):
+        ...
+  """
+
+  def run(self, **env) -> str:
+    """Run the container synchronously."""
+    raise NotImplementedError
+
+  def detach(self, **env):
+    """Run the container asynchronously."""
+    raise NotImplementedError
+
+  def address(self) -> Tuple[str, int]:
+    """Return the bound address for the container."""
+    raise NotImplementedError
+
+  def get_names(self) -> types.GeneratorType:
+    """Return names of all containers."""
+    raise NotImplementedError
+
+
+# pylint: disable=too-many-instance-attributes
+class DockerContainer(Container):
+  """Class that handles creating a docker container."""
+
+  # pylint: disable=too-many-arguments
+  def __init__(self,
+               client: docker.DockerClient,
+               host: str,
+               image: str,
+               count: int = 1,
+               runtime: str = "runc",
+               port: int = 0,
+               **kwargs):
+    """Trys to setup "count" containers.
+
+    Args:
+      client: A docker client from dockerpy.
+      host: The host address the image is running on.
+      image: The name of the image to run.
+      count: The number of containers to setup.
+      runtime: The container runtime to use.
+      port: The port to reserve.
+      **kwargs: Additional container options.
+    """
+    assert count >= 1
+    assert port == 0 or count == 1
+    self._client = client
+    self._host = host
+    self._containers = []
+    self._count = count
+    self._image = image
+    self._runtime = runtime
+    self._port = port
+    self._kwargs = kwargs
+    if port != 0:
+      self._ports = {"%d/tcp" % port: None}
+    else:
+      self._ports = {}
+
+  @contextlib.contextmanager
+  def detach(self, **env):
+    env = ["%s=%s" % (key, value) for (key, value) in env.items()]
+    # Start all containers.
+    for _ in range(self._count):
+      try:
+        # Start the container in a detached mode.
+        container = self._client.containers.run(
+            self._image,
+            detach=True,
+            remove=True,
+            runtime=self._runtime,
+            ports=self._ports,
+            environment=env,
+            **self._kwargs)
+        logging.info("Started detached container %s -> %s", self._image,
+                     container.attrs["Id"])
+        self._containers.append(container)
+      except Exception as exc:
+        self._clean_containers()
+        raise exc
+    try:
+      # Wait for all containers to be up.
+      for container in self._containers:
+        while not container.attrs["State"]["Running"]:
+          container = self._client.containers.get(container.attrs["Id"])
+      yield self
+    finally:
+      self._clean_containers()
+
+  def address(self) -> Tuple[str, int]:
+    assert self._count == 1
+    assert self._port != 0
+    container = self._client.containers.get(self._containers[0].attrs["Id"])
+    port = container.attrs["NetworkSettings"]["Ports"][
+        "%d/tcp" % self._port][0]["HostPort"]
+    return (self._host, port)
+
+  def get_names(self) -> types.GeneratorType:
+    for container in self._containers:
+      yield container.name
+
+  def run(self, **env) -> str:
+    env = ["%s=%s" % (key, value) for (key, value) in env.items()]
+    return self._client.containers.run(
+        self._image,
+        runtime=self._runtime,
+        ports=self._ports,
+        remove=True,
+        environment=env,
+        **self._kwargs).decode("utf-8")
+
+  def _clean_containers(self):
+    """Kills all containers."""
+    for container in self._containers:
+      try:
+        container.kill()
+      except docker.errors.NotFound:
+        pass
+
+
+class MockContainer(Container):
+  """Mock of Container."""
+
+  def __init__(self, workload: str):
+    self._workload = workload
+
+  def __enter__(self):
+    return self
+
+  def run(self, **env):
+    # Lookup sample data if any exists for the workload module. We use a
+    # well-defined test locate and a well-defined sample function.
+    mod = pydoc.locate(workloads.__name__ + "." + self._workload)
+    if hasattr(mod, "sample"):
+      return mod.sample(**env)
+    return ""  # No output.
+
+  def address(self) -> Tuple[str, int]:
+    return ("example.com", 80)
+
+  def get_names(self) -> types.GeneratorType:
+    yield "mock"
+
+  @contextlib.contextmanager
+  def detach(self, **env):
+    yield self