def run(self): contol_stub = beam_fn_api_pb2.BeamFnControlStub(self._control_channel) # TODO(robertwb): Wire up to new state api. state_stub = None self.worker = SdkWorker(state_stub, self._data_channel_factory) responses = queue.Queue() no_more_work = object() def get_responses(): while True: response = responses.get() if response is no_more_work: return yield response def process_requests(): for work_request in contol_stub.Control(get_responses()): logging.info('Got work %s', work_request.instruction_id) try: response = self.worker.do_instruction(work_request) except Exception: # pylint: disable=broad-except response = beam_fn_api_pb2.InstructionResponse( instruction_id=work_request.instruction_id, error=traceback.format_exc()) responses.put(response) t = threading.Thread(target=process_requests) t.start() t.join() # get_responses may be blocked on responses.get(), but we need to return # control to its caller. responses.put(no_more_work) self._data_channel_factory.close() logging.info('Done consuming work.')
def run(self): contol_stub = beam_fn_api_pb2.BeamFnControlStub(self._control_channel) # TODO(robertwb): Wire up to new state api. state_stub = None self.worker = SdkWorker(state_stub, self._data_channel_factory) responses = queue.Queue() no_more_work = object() def get_responses(): while True: response = responses.get() if response is no_more_work: return yield response for work_request in contol_stub.Control(get_responses()): logging.info('Got work %s', work_request.instruction_id) request_type = work_request.WhichOneof('request') if request_type == ['process_bundle_progress']: thread_pool = self._progress_thread_pool else: thread_pool = self._default_work_thread_pool # Need this wrapper to capture the original stack trace. def do_instruction(request): try: return self.worker.do_instruction(request) except Exception as e: # pylint: disable=broad-except traceback_str = traceback.format_exc(e) raise StandardError( "Error processing request. Original traceback " "is\n%s\n" % traceback_str) def handle_response(request, response_future): try: response = response_future.result() except Exception as e: # pylint: disable=broad-except logging.error('Error processing instruction %s', request.instruction_id, exc_info=True) response = beam_fn_api_pb2.InstructionResponse( instruction_id=request.instruction_id, error=str(e)) responses.put(response) thread_pool.submit(do_instruction, work_request).add_done_callback( functools.partial(handle_response, work_request)) logging.info("No more requests from control plane") logging.info("SDK Harness waiting for in-flight requests to complete") # Wait until existing requests are processed. self._progress_thread_pool.shutdown() self._default_work_thread_pool.shutdown() # get_responses may be blocked on responses.get(), but we need to return # control to its caller. responses.put(no_more_work) self._data_channel_factory.close() logging.info('Done consuming work.')