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_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 get_data(self, data_id): return [fromCP(r) for r in self.transact_to_all_workers(Commands.GETDATA, itertools.repeat((data_id,)))]