diff options
Diffstat (limited to 'benchmarks/harness/machine_producers/machine_producer.py')
-rw-r--r-- | benchmarks/harness/machine_producers/machine_producer.py | 21 |
1 files changed, 21 insertions, 0 deletions
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() |