From 7b7ce29af326ccd247ee5225e9b5b55a9d0330ce Mon Sep 17 00:00:00 2001 From: Zach Koopmans Date: Wed, 15 Jan 2020 14:24:55 -0800 Subject: Update commandline and get local runs working. PiperOrigin-RevId: 289937063 --- .../harness/machine_producers/machine_producer.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) (limited to 'benchmarks/harness/machine_producers/machine_producer.py') 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() -- cgit v1.2.3