def test_worker_set_list_data():
with WorkerWithSocket() as send_and_recv:
dataToSend = toCP(list(range(10)))
data_id = b"daterz-idz"
send_and_recv(Commands.SETDATA, data_id, dataToSend, expect=Commands.OK)
send_and_recv(Commands.LISTDATA, expect=toCP([data_id]))
def test_worker_stdout_stream():
with WorkerWithSocket(("send_and_recv", "stdsocket")) as (send_and_recv, stdsocket):
data = list(range(10))
data_pickle = toCP(data)
data_id = b"data1"
send_and_recv(Commands.SETDATA, data_id, data_pickle, expect=Commands.OK)
def fun(partition):
ret = [x + 1 for x in partition]
print(ret)
return ret
fun_pickle = toCP(fun)
fun_id = b"fun"
send_and_recv(Commands.SETDATA, fun_id, fun_pickle, expect=Commands.OK)
map_output_id = b"data2"
send_and_recv(Commands.CALL, fun_id, data_id, map_output_id, expect=Commands.OK)
# Poll the worker repeatedly until it's done.
isworking = Commands.FALSE
while isworking != Commands.TRUE:
isworking = send_and_recv(Commands.ISWORKING)
time.sleep(0.1)
received_stdout = stdsocket.recv()
# Ignore the outer worker's logging.
while received_stdout.startswith(b"WORKER"):
received_stdout = stdsocket.recv()
assert received_stdout == str(fun(data)).encode("ascii"), received_stdout
listing = send_and_recv(Commands.LISTDATA)
assert set(fromCP(listing)) == {data_id, fun_id, map_output_id}, listing
send_and_recv(Commands.GETDATA, map_output_id, expect=toCP(fun(data)))
print("") # Make sure this test doesn't end on half a line
def test_worker_set_get_data():
with WorkerWithSocket() as send_and_recv:
data = list(range(10))
data_pickle = toCP(data)
send_and_recv(Commands.SETDATA, b"daterz-idz", data_pickle, expect=Commands.OK)
send_and_recv(Commands.GETDATA, b"daterz-idz", expect=data_pickle)
def set_data(self, data_id, iterator):
get_args_for_element = lambda element: (data_id, toCP(element))
cmd_tuple_iterator = map(get_args_for_element, iterator)
responses = self.transact_to_all_workers(Commands.SETDATA, cmd_tuple_iterator)
responses_are_ok = [response == Commands.OK for response in responses]
if not all(responses_are_ok):
failed_worker_ids = [self.worker_ids[ix] for ix, ok in enumerate(responses_are_ok) if not ok]
raise ValueError("Workers {} failed 'setdata' command.".format(failed_worker_ids))
def test_worker_del_data():
with WorkerWithSocket() as send_and_recv:
data = list(range(10))
data_pickle = toCP(data)
data_id = b"data1"
send_and_recv(Commands.SETDATA, data_id, data_pickle, expect=Commands.OK)
send_and_recv(Commands.DELDATA, data_id, expect=Commands.TRUE)
send_and_recv(Commands.DELDATA, data_id, expect=Commands.FALSE)
send_and_recv(Commands.GETDATA, data_id, expect=lambda e: isinstance(fromCP(e), Exception))
def test_worker_set_call_get_data():
with WorkerWithSocket() as send_and_recv:
data = list(range(10))
data_pickle = toCP(data)
data_id = b"data1"
send_and_recv(Commands.SETDATA, data_id, data_pickle, expect=Commands.OK)
fun_id = b"fun1"
fun = lambda partition: [x + 1 for x in partition]
fun_pickle = toCP(fun)
send_and_recv(Commands.SETDATA, fun_id, fun_pickle, expect=Commands.OK)
map_output_id = b"data2"
send_and_recv(Commands.CALL, fun_id, data_id, map_output_id, expect=Commands.OK)
# Poll the worker repeatedly until it's done.
isworking = Commands.FALSE
while isworking != Commands.TRUE:
isworking = send_and_recv(Commands.ISWORKING)
time.sleep(0.1)
send_and_recv(Commands.LISTDATA, expect=lambda d: set(fromCP(d)) == {data_id, fun_id, map_output_id})
send_and_recv(Commands.GETDATA, map_output_id, expect=toCP(fun(data)))
def _handle_GETDATA(self, id):
try:
return self.data[id]
except KeyError as e:
return toCP(e)
def _handle_LISTDATA(self):
return toCP(list(self.data.keys()))
def test_worker_set_data():
with WorkerWithSocket() as send_and_recv:
dataToSend = toCP(list(range(10)))
send_and_recv(Commands.SETDATA, b"daterz-idz", dataToSend, expect=Commands.OK)
def send_function_async(self, workerID, functionID, functionOrCP):
if callable(functionOrCP):
functionOrCP = toCP(functionOrCP)
return self._send_setdata(workerID, functionOrCP)
