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