github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/runners/portability/expansion_service_main.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  # pytype: skip-file
    18  
    19  import argparse
    20  import logging
    21  import signal
    22  import sys
    23  
    24  import grpc
    25  
    26  from apache_beam.internal import pickler
    27  from apache_beam.options.pipeline_options import PipelineOptions
    28  from apache_beam.options.pipeline_options import SetupOptions
    29  from apache_beam.portability.api import beam_artifact_api_pb2_grpc
    30  from apache_beam.portability.api import beam_expansion_api_pb2_grpc
    31  from apache_beam.runners.portability import artifact_service
    32  from apache_beam.runners.portability import expansion_service
    33  from apache_beam.transforms import fully_qualified_named_transform
    34  from apache_beam.utils import thread_pool_executor
    35  
    36  _LOGGER = logging.getLogger(__name__)
    37  
    38  
    39  def main(argv):
    40    parser = argparse.ArgumentParser()
    41    parser.add_argument(
    42        '-p', '--port', type=int, help='port on which to serve the job api')
    43    parser.add_argument('--fully_qualified_name_glob', default=None)
    44    known_args, pipeline_args = parser.parse_known_args(argv)
    45    pipeline_options = PipelineOptions(
    46        pipeline_args + ["--experiments=beam_fn_api", "--sdk_location=container"])
    47  
    48    # Set this before any pipeline construction occurs.
    49    # See https://github.com/apache/beam/issues/21615
    50    pickler.set_library(pipeline_options.view_as(SetupOptions).pickle_library)
    51  
    52    with fully_qualified_named_transform.FullyQualifiedNamedTransform.with_filter(
    53        known_args.fully_qualified_name_glob):
    54  
    55      server = grpc.server(thread_pool_executor.shared_unbounded_instance())
    56      beam_expansion_api_pb2_grpc.add_ExpansionServiceServicer_to_server(
    57          expansion_service.ExpansionServiceServicer(pipeline_options), server)
    58      beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server(
    59          artifact_service.ArtifactRetrievalService(
    60              artifact_service.BeamFilesystemHandler(None).file_reader),
    61          server)
    62      server.add_insecure_port('[::]:{}'.format(known_args.port))
    63      server.start()
    64      _LOGGER.info('Listening for expansion requests at %d', known_args.port)
    65  
    66      def cleanup(unused_signum, unused_frame):
    67        _LOGGER.info('Shutting down expansion service.')
    68        server.stop(None)
    69  
    70      signal.signal(signal.SIGTERM, cleanup)
    71      signal.signal(signal.SIGINT, cleanup)
    72      # blocking main thread forever.
    73      signal.pause()
    74  
    75  
    76  if __name__ == '__main__':
    77    logging.getLogger().setLevel(logging.INFO)
    78    main(sys.argv)