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)