Update commandline and get local runs working.

PiperOrigin-RevId: 289937063

3 files changed, 27 insertions, 4 deletions

diff --git a/benchmarks/harness/__init__.py b/benchmarks/harness/__init__.py
index a7f34da9e..7b96d1666 100644
--- a/benchmarks/harness/__init__.py
+++ b/benchmarks/harness/__init__.py
@@ -18,7 +18,7 @@ import os
 # LOCAL_WORKLOADS_PATH defines the path to use for local workloads. This is a
 # format string that accepts a single string parameter.
 LOCAL_WORKLOADS_PATH = os.path.join(
-    os.path.dirname(__file__), "../workloads/{}")
+    os.path.dirname(__file__), "../workloads/{}/tar.tar")
 
 # REMOTE_WORKLOADS_PATH defines the path to use for storing the workloads on the
 # remote host. This is a format string that accepts a single string parameter.
diff --git a/benchmarks/harness/machine.py b/benchmarks/harness/machine.py
index 66b719b63..af037dbcc 100644
--- a/benchmarks/harness/machine.py
+++ b/benchmarks/harness/machine.py
@@ -160,15 +160,17 @@ class LocalMachine(Machine):
     stdout, stderr = process.communicate()
     return stdout.decode("utf-8"), stderr.decode("utf-8")
 
-  def read(self, path: str) -> str:
+  def read(self, path: str) -> bytes:
     # Read the exact path locally.
     return open(path, "r").read()
 
   def pull(self, workload: str) -> str:
     # Run the docker build command locally.
     logging.info("Building %s@%s locally...", workload, self._name)
-    self.run("docker build --tag={} {}".format(
-        workload, harness.LOCAL_WORKLOADS_PATH.format(workload)))
+    with open(harness.LOCAL_WORKLOADS_PATH.format(workload),
+              "rb") as dockerfile:
+      self._docker_client.images.build(
+          fileobj=dockerfile, tag=workload, custom_context=True)
     return workload  # Workload is the tag.
 
   def container(self, image: str, **kwargs) -> container.Container:
diff --git a/benchmarks/harness/machine_producers/machine_producer.py b/benchmarks/harness/machine_producers/machine_producer.py
index 124ee14cc..f5591c026 100644
--- a/benchmarks/harness/machine_producers/machine_producer.py
+++ b/benchmarks/harness/machine_producers/machine_producer.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Abstract types."""
 
+import threading
 from typing import List
 
 from benchmarks.harness import machine
@@ -28,3 +29,23 @@ class MachineProducer:
   def release_machines(self, machine_list: List[machine.Machine]):
     """Releases the given set of machines."""
     raise NotImplementedError
+
+
+class LocalMachineProducer(MachineProducer):
+  """Produces Local Machines."""
+
+  def __init__(self, limit: int):
+    self.limit_sem = threading.Semaphore(value=limit)
+
+  def get_machines(self, num_machines: int) -> List[machine.Machine]:
+    """Returns the request number of MockMachines."""
+
+    self.limit_sem.acquire()
+    return [machine.LocalMachine("local") for _ in range(num_machines)]
+
+  def release_machines(self, machine_list: List[machine.MockMachine]):
+    """No-op."""
+    if not machine_list:
+      raise ValueError("Cannot release an empty list!")
+    self.limit_sem.release()
+    machine_list.clear()