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 start_test_artifact_server(): server = grpc.server(futures.ThreadPoolExecutor(max_workers=1)) add_ArtifactRetrievalServiceServicer_to_server( ArtifactService(), server) port = get_unused_port() server.add_insecure_port('[::]:' + str(port)) server.start() return server, port
def test_with_grpc(self): server = grpc.server(UnboundedThreadPoolExecutor()) try: beam_artifact_api_pb2_grpc.add_ArtifactStagingServiceServicer_to_server( self._service, server) beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server( self._service, server) port = server.add_insecure_port('[::]:0') server.start() channel = grpc.insecure_channel('localhost:%d' % port) self._run_staging( beam_artifact_api_pb2_grpc.ArtifactStagingServiceStub(channel), beam_artifact_api_pb2_grpc.ArtifactRetrievalServiceStub( channel)) channel.close() finally: server.stop(1)
def __init__( self, state, # type: StateServicer provision_info, # type: Optional[ExtendedProvisionInfo] worker_manager, # type: WorkerHandlerManager ): # type: (...) -> None self.state = state self.provision_info = provision_info self.control_server = grpc.server( thread_pool_executor.shared_unbounded_instance()) self.control_port = self.control_server.add_insecure_port('[::]:0') self.control_address = 'localhost:%s' % self.control_port # Options to have no limits (-1) on the size of the messages # received or sent over the data plane. The actual buffer size # is controlled in a layer above. no_max_message_sizes = [("grpc.max_receive_message_length", -1), ("grpc.max_send_message_length", -1)] self.data_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=no_max_message_sizes) self.data_port = self.data_server.add_insecure_port('[::]:0') self.state_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=no_max_message_sizes) self.state_port = self.state_server.add_insecure_port('[::]:0') self.control_handler = BeamFnControlServicer(worker_manager) beam_fn_api_pb2_grpc.add_BeamFnControlServicer_to_server( self.control_handler, self.control_server) # If we have provision info, serve these off the control port as well. if self.provision_info: if self.provision_info.provision_info: beam_provision_api_pb2_grpc.add_ProvisionServiceServicer_to_server( BasicProvisionService(self.provision_info.provision_info, worker_manager), self.control_server) def open_uncompressed(f): # type: (str) -> BinaryIO return filesystems.FileSystems.open( f, compression_type=CompressionTypes.UNCOMPRESSED) beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server( artifact_service.ArtifactRetrievalService( file_reader=open_uncompressed), self.control_server) self.data_plane_handler = data_plane.BeamFnDataServicer( DATA_BUFFER_TIME_LIMIT_MS) beam_fn_api_pb2_grpc.add_BeamFnDataServicer_to_server( self.data_plane_handler, self.data_server) beam_fn_api_pb2_grpc.add_BeamFnStateServicer_to_server( GrpcStateServicer(state), self.state_server) self.logging_server = grpc.server( thread_pool_executor.shared_unbounded_instance(), options=no_max_message_sizes) self.logging_port = self.logging_server.add_insecure_port('[::]:0') beam_fn_api_pb2_grpc.add_BeamFnLoggingServicer_to_server( BasicLoggingService(), self.logging_server) _LOGGER.info('starting control server on port %s', self.control_port) _LOGGER.info('starting data server on port %s', self.data_port) _LOGGER.info('starting state server on port %s', self.state_port) _LOGGER.info('starting logging server on port %s', self.logging_port) self.logging_server.start() self.state_server.start() self.data_server.start() self.control_server.start()
def __init__( self, state, # type: StateServicer provision_info, # type: Optional[ExtendedProvisionInfo] worker_manager, # type: WorkerHandlerManager ): # type: (...) -> None self.state = state self.provision_info = provision_info self.control_server = grpc.server(UnboundedThreadPoolExecutor()) self.control_port = self.control_server.add_insecure_port('[::]:0') self.control_address = 'localhost:%s' % self.control_port # Options to have no limits (-1) on the size of the messages # received or sent over the data plane. The actual buffer size # is controlled in a layer above. no_max_message_sizes = [("grpc.max_receive_message_length", -1), ("grpc.max_send_message_length", -1)] self.data_server = grpc.server(UnboundedThreadPoolExecutor(), options=no_max_message_sizes) self.data_port = self.data_server.add_insecure_port('[::]:0') self.state_server = grpc.server(UnboundedThreadPoolExecutor(), options=no_max_message_sizes) self.state_port = self.state_server.add_insecure_port('[::]:0') self.control_handler = BeamFnControlServicer() beam_fn_api_pb2_grpc.add_BeamFnControlServicer_to_server( self.control_handler, self.control_server) # If we have provision info, serve these off the control port as well. if self.provision_info: if self.provision_info.provision_info: beam_provision_api_pb2_grpc.add_ProvisionServiceServicer_to_server( BasicProvisionService(self.provision_info.provision_info, worker_manager), self.control_server) if self.provision_info.artifact_staging_dir: service = artifact_service.BeamFilesystemArtifactService( self.provision_info.artifact_staging_dir ) # type: beam_artifact_api_pb2_grpc.ArtifactRetrievalServiceServicer else: service = EmptyArtifactRetrievalService() beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server( service, self.control_server) self.data_plane_handler = data_plane.BeamFnDataServicer( DATA_BUFFER_TIME_LIMIT_MS) beam_fn_api_pb2_grpc.add_BeamFnDataServicer_to_server( self.data_plane_handler, self.data_server) beam_fn_api_pb2_grpc.add_BeamFnStateServicer_to_server( GrpcStateServicer(state), self.state_server) self.logging_server = grpc.server(UnboundedThreadPoolExecutor(), options=no_max_message_sizes) self.logging_port = self.logging_server.add_insecure_port('[::]:0') beam_fn_api_pb2_grpc.add_BeamFnLoggingServicer_to_server( BasicLoggingService(), self.logging_server) _LOGGER.info('starting control server on port %s', self.control_port) _LOGGER.info('starting data server on port %s', self.data_port) _LOGGER.info('starting state server on port %s', self.state_port) _LOGGER.info('starting logging server on port %s', self.logging_port) self.logging_server.start() self.state_server.start() self.data_server.start() self.control_server.start()