def run(self): contol_stub = beam_fn_api_pb2_grpc.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.')
def run(self): control_stub = beam_fn_api_pb2_grpc.BeamFnControlStub( self._control_channel) no_more_work = object() # Create workers bundle_processor_cache = BundleProcessorCache( state_handler_factory=self._state_handler_factory, data_channel_factory=self._data_channel_factory, fns=self._fns) for _ in range(self._worker_count): # SdkHarness manage function registration and share self._fns with all # the workers. This is needed because function registration (register) # and exceution(process_bundle) are send over different request and we # do not really know which woker is going to process bundle # for a function till we get process_bundle request. Moreover # same function is reused by different process bundle calls and # potentially get executed by different worker. Hence we need a # centralized function list shared among all the workers. self.workers.put( SdkWorker(bundle_processor_cache, profiler_factory=self._profiler_factory)) def get_responses(): while True: response = self._responses.get() if response is no_more_work: return yield response self._alive = True monitoring_thread = threading.Thread( target=self._monitor_process_bundle) monitoring_thread.daemon = True monitoring_thread.start() try: for work_request in control_stub.Control(get_responses()): logging.debug('Got work %s', work_request.instruction_id) request_type = work_request.WhichOneof('request') # Name spacing the request method with 'request_'. The called method # will be like self.request_register(request) getattr(self, SdkHarness.REQUEST_METHOD_PREFIX + request_type)(work_request) finally: self._alive = False 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._process_thread_pool.shutdown() # get_responses may be blocked on responses.get(), but we need to return # control to its caller. self._responses.put(no_more_work) # Stop all the workers and clean all the associated resources self._data_channel_factory.close() self._state_handler_factory.close() logging.info('Done consuming work.')
def run(self): # type: () -> None self._control_stub = beam_fn_api_pb2_grpc.BeamFnControlStub( self._control_channel) no_more_work = Sentinel.sentinel def get_responses(): # type: () -> Iterator[beam_fn_api_pb2.InstructionResponse] while True: response = self._responses.get() if response is no_more_work: return yield response self._alive = True try: for work_request in self._control_stub.Control(get_responses()): _LOGGER.debug('Got work %s', work_request.instruction_id) request_type = work_request.WhichOneof('request') # Name spacing the request method with 'request_'. The called method # will be like self.request_register(request) getattr(self, SdkHarness.REQUEST_METHOD_PREFIX + request_type)( work_request) finally: self._alive = False _LOGGER.info('No more requests from control plane') _LOGGER.info('SDK Harness waiting for in-flight requests to complete') # Wait until existing requests are processed. self._worker_thread_pool.shutdown() # get_responses may be blocked on responses.get(), but we need to return # control to its caller. self._responses.put(no_more_work) # Stop all the workers and clean all the associated resources self._data_channel_factory.close() self._state_handler_factory.close() self._bundle_processor_cache.shutdown() if self._status_handler: self._status_handler.close() _LOGGER.info('Done consuming work.')