def main(unused_argv):
PyPIArtifactRegistry.register_artifact('beautifulsoup4', '>=4.9,<5.0')
parser = argparse.ArgumentParser()
parser.add_argument('-p',
'--port',
type=int,
help='port on which to serve the job api')
parser.add_argument('--fully_qualified_name_glob', default=None)
options = parser.parse_args()
global server
with fully_qualified_named_transform.FullyQualifiedNamedTransform.with_filter(
options.fully_qualified_name_glob):
server = grpc.server(thread_pool_executor.shared_unbounded_instance())
beam_expansion_api_pb2_grpc.add_ExpansionServiceServicer_to_server(
expansion_service.ExpansionServiceServicer(
PipelineOptions([
"--experiments", "beam_fn_api", "--sdk_location",
"container"
])), server)
beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server(
artifact_service.ArtifactRetrievalService(
artifact_service.BeamFilesystemHandler(None).file_reader),
server)
server.add_insecure_port('localhost:{}'.format(options.port))
server.start()
_LOGGER.info('Listening for expansion requests at %d', options.port)
signal.signal(signal.SIGTERM, cleanup)
signal.signal(signal.SIGINT, cleanup)
# blocking main thread forever.
signal.pause()
def main(argv):
parser = argparse.ArgumentParser()
parser.add_argument('-p',
'--port',
type=int,
help='port on which to serve the job api')
parser.add_argument('--fully_qualified_name_glob', default=None)
known_args, pipeline_args = parser.parse_known_args(argv)
pipeline_options = PipelineOptions(
pipeline_args +
["--experiments=beam_fn_api", "--sdk_location=container"])
with fully_qualified_named_transform.FullyQualifiedNamedTransform.with_filter(
known_args.fully_qualified_name_glob):
server = grpc.server(thread_pool_executor.shared_unbounded_instance())
beam_expansion_api_pb2_grpc.add_ExpansionServiceServicer_to_server(
expansion_service.ExpansionServiceServicer(pipeline_options),
server)
beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server(
artifact_service.ArtifactRetrievalService(
artifact_service.BeamFilesystemHandler(None).file_reader),
server)
server.add_insecure_port('localhost:{}'.format(known_args.port))
server.start()
_LOGGER.info('Listening for expansion requests at %d', known_args.port)
def cleanup(unused_signum, unused_frame):
_LOGGER.info('Shutting down expansion service.')
server.stop(None)
signal.signal(signal.SIGTERM, cleanup)
signal.signal(signal.SIGINT, cleanup)
# blocking main thread forever.
signal.pause()
def __init__(self, staging_dir=None):
super(LocalJobServicer, self).__init__()
self._cleanup_staging_dir = staging_dir is None
self._staging_dir = staging_dir or tempfile.mkdtemp()
self._artifact_service = artifact_service.ArtifactStagingService(
artifact_service.BeamFilesystemHandler(self._staging_dir).file_writer)
self._artifact_staging_endpoint = None # type: Optional[endpoints_pb2.ApiServiceDescriptor]
def _stage_via_portable_service(self, artifact_staging_channel,
staging_session_token):
artifact_service.offer_artifacts(
beam_artifact_api_pb2_grpc.ArtifactStagingServiceStub(
channel=artifact_staging_channel),
artifact_service.ArtifactRetrievalService(
artifact_service.BeamFilesystemHandler(None).file_reader),
staging_session_token)
def stage(self, pipeline, artifact_staging_endpoint,
staging_session_token):
# type: (...) -> Optional[Any]
"""Stage artifacts"""
if artifact_staging_endpoint:
artifact_service.offer_artifacts(
beam_artifact_api_pb2_grpc.ArtifactStagingServiceStub(
channel=grpc.insecure_channel(artifact_staging_endpoint)),
artifact_service.ArtifactRetrievalService(
artifact_service.BeamFilesystemHandler(None).file_reader),
staging_session_token)
