def main(unused_argv):
"""Main entry point for SDK Fn Harness."""
logging_service_descriptor = beam_fn_api_pb2.ApiServiceDescriptor()
text_format.Merge(os.environ['LOGGING_API_SERVICE_DESCRIPTOR'],
logging_service_descriptor)
# Send all logs to the runner.
fn_log_handler = FnApiLogRecordHandler(logging_service_descriptor)
# TODO(vikasrk): This should be picked up from pipeline options.
logging.getLogger().setLevel(logging.INFO)
logging.getLogger().addHandler(fn_log_handler)
try:
logging.info('Python sdk harness started.')
service_descriptor = beam_fn_api_pb2.ApiServiceDescriptor()
text_format.Merge(os.environ['CONTROL_API_SERVICE_DESCRIPTOR'],
service_descriptor)
# TODO(robertwb): Support credentials.
assert not service_descriptor.oauth2_client_credentials_grant.url
channel = grpc.insecure_channel(service_descriptor.url)
SdkHarness(channel).run()
logging.info('Python sdk harness exiting.')
except: # pylint: disable=broad-except
logging.exception('Python sdk harness failed: ')
raise
finally:
fn_log_handler.close()
def setUp(self):
self.test_logging_service = BeamFnLoggingServicer()
self.server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
beam_fn_api_pb2.add_BeamFnLoggingServicer_to_server(
self.test_logging_service, self.server)
self.test_port = self.server.add_insecure_port('[::]:0')
self.server.start()
self.logging_service_descriptor = beam_fn_api_pb2.ApiServiceDescriptor(
)
self.logging_service_descriptor.url = 'localhost:%s' % self.test_port
self.fn_log_handler = log_handler.FnApiLogRecordHandler(
self.logging_service_descriptor)
logging.getLogger().setLevel(logging.INFO)
logging.getLogger().addHandler(self.fn_log_handler)
def run(self):
control_descriptor = text_format.MessageToString(
beam_fn_api_pb2.ApiServiceDescriptor(url=self._control_address))
p = subprocess.Popen(
self._worker_command_line,
shell=True,
env=dict(os.environ,
CONTROL_API_SERVICE_DESCRIPTOR=control_descriptor))
try:
p.wait()
if p.returncode:
raise RuntimeError(
"Worker subprocess exited with return code %s" %
p.returncode)
finally:
if p.poll() is None:
p.kill()