class MasterWithWorkers:
def __init__(self, num_workers=4):
self.num_workers = num_workers
# Increment the ports each time, so we don't hit the "address already in use" error.
self.workport = next(workport)
self.logport = next(logport)
self.master = Master(workport=self.workport, logport=self.logport)
self.workers = []
for i in range(num_workers):
worker = Worker("tcp://localhost", workport=self.workport, logport=self.logport, print_local=False)
self.workers.append(worker)
def __enter__(self):
for worker in self.workers:
worker.start()
self.master.wait_for_worker_connections(self.num_workers)
return self.master, self.workers
def __exit__(self, *args):
if any(arg is not None for arg in args):
print("MasterWithWorkers.__exit__() caught an error.", file=sys.stderr)
traceback.print_exc()
dead_workers = [worker for worker in self.workers if not worker.is_alive()]
if len(dead_workers) > 0:
raise RuntimeError("Workers {} died during testing.".format(dead_workers))
for worker in self.workers:
if worker.is_alive():
worker.terminate()
worker.join()
print("Killing workers.")
self.master.kill_workers()
self.master.release_ports()
print("Closed master sockets.", flush=True)
class ProcessWorkerContext:
def __init__(self, workport=23456, logport=34567, num_workers=4):
self.workport = workport
self.logport = logport
self.num_workers = num_workers
def __enter__(self):
self.master = Master(workport=self.workport, logport=self.logport)
self.workers = []
for i in range(self.num_workers):
worker = Worker("tcp://localhost", workport=self.workport, logport=self.logport)
self.workers.append(worker)
worker.start()
self.context = SplarkContext(master=self.master, num_workers=self.num_workers)
print("Created context.")
return self.context, self.master, self.workers
def __exit__(self, *args):
if any(arg is not None for arg in args):
print("ContextWithWorkers.__exit__() caught an error.", file=sys.stderr)
traceback.print_exc()
dead_workers = [worker for worker in self.workers if not worker.is_alive()]
if len(dead_workers) > 0:
raise RuntimeError("Workers {} died during testing.".format(dead_workers))
for worker in self.workers:
if worker.is_alive():
worker.terminate()
worker.join()
print("Killing workers.")
self.master.kill_workers()
self.master.release_ports()
print("Closed master sockets.", flush=True)
