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.')