Example #1
0
  def run_pipeline(self, pipeline):
    portable_options = pipeline.options.view_as(PortableOptions)
    job_endpoint = portable_options.job_endpoint

    # TODO: https://issues.apache.org/jira/browse/BEAM-5525
    # portable runner specific default
    if pipeline.options.view_as(SetupOptions).sdk_location == 'default':
      pipeline.options.view_as(SetupOptions).sdk_location = 'container'

    if not job_endpoint:
      docker = DockerizedJobServer()
      job_endpoint = docker.start()

    proto_context = pipeline_context.PipelineContext(
        default_environment=PortableRunner._create_environment(
            portable_options))
    proto_pipeline = pipeline.to_runner_api(context=proto_context)

    # Some runners won't detect the GroupByKey transform unless it has no
    # subtransforms.  Remove all sub-transforms until BEAM-4605 is resolved.
    for _, transform_proto in list(
        proto_pipeline.components.transforms.items()):
      if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
        for sub_transform in transform_proto.subtransforms:
          del proto_pipeline.components.transforms[sub_transform]
        del transform_proto.subtransforms[:]

    # TODO: Define URNs for options.
    # convert int values: https://issues.apache.org/jira/browse/BEAM-5509
    options = {'beam:option:' + k + ':v1': (str(v) if type(v) == int else v)
               for k, v in pipeline._options.get_all_options().items()
               if v is not None}

    channel = grpc.insecure_channel(job_endpoint)
    grpc.channel_ready_future(channel).result()
    job_service = beam_job_api_pb2_grpc.JobServiceStub(channel)

    # Sends the PrepareRequest but retries in case the channel is not ready
    def send_prepare_request(max_retries=5):
      num_retries = 0
      while True:
        try:
          # This reports channel is READY but connections may fail
          # Seems to be only an issue on Mac with port forwardings
          grpc.channel_ready_future(channel).result()
          return job_service.Prepare(
              beam_job_api_pb2.PrepareJobRequest(
                  job_name='job', pipeline=proto_pipeline,
                  pipeline_options=job_utils.dict_to_struct(options)))
        except grpc._channel._Rendezvous as e:
          num_retries += 1
          if num_retries > max_retries:
            raise e

    prepare_response = send_prepare_request()
    if prepare_response.artifact_staging_endpoint.url:
      stager = portable_stager.PortableStager(
          grpc.insecure_channel(prepare_response.artifact_staging_endpoint.url),
          prepare_response.staging_session_token)
      retrieval_token, _ = stager.stage_job_resources(
          pipeline._options,
          staging_location='')
    else:
      retrieval_token = None
    run_response = job_service.Run(
        beam_job_api_pb2.RunJobRequest(
            preparation_id=prepare_response.preparation_id,
            retrieval_token=retrieval_token))
    return PipelineResult(job_service, run_response.job_id)
  def run_pipeline(self, pipeline):
    portable_options = pipeline.options.view_as(PortableOptions)
    job_endpoint = portable_options.job_endpoint

    # TODO: https://issues.apache.org/jira/browse/BEAM-5525
    # portable runner specific default
    if pipeline.options.view_as(SetupOptions).sdk_location == 'default':
      pipeline.options.view_as(SetupOptions).sdk_location = 'container'

    if not job_endpoint:
      docker = DockerizedJobServer()
      job_endpoint = docker.start()

    proto_context = pipeline_context.PipelineContext(
        default_environment=PortableRunner._create_environment(
            portable_options))
    proto_pipeline = pipeline.to_runner_api(context=proto_context)

    if not self.is_embedded_fnapi_runner:
      # Java has different expectations about coders
      # (windowed in Fn API, but *un*windowed in runner API), whereas the
      # embedded FnApiRunner treats them consistently, so we must guard this
      # for now, until FnApiRunner is fixed.
      # See also BEAM-2717.
      for pcoll in proto_pipeline.components.pcollections.values():
        if pcoll.coder_id not in proto_context.coders:
          # This is not really a coder id, but a pickled coder.
          coder = coders.registry.get_coder(pickler.loads(pcoll.coder_id))
          pcoll.coder_id = proto_context.coders.get_id(coder)
      proto_context.coders.populate_map(proto_pipeline.components.coders)

    # Some runners won't detect the GroupByKey transform unless it has no
    # subtransforms.  Remove all sub-transforms until BEAM-4605 is resolved.
    for _, transform_proto in list(
        proto_pipeline.components.transforms.items()):
      if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
        for sub_transform in transform_proto.subtransforms:
          del proto_pipeline.components.transforms[sub_transform]
        del transform_proto.subtransforms[:]

    # TODO: Define URNs for options.
    # convert int values: https://issues.apache.org/jira/browse/BEAM-5509
    options = {'beam:option:' + k + ':v1': (str(v) if type(v) == int else v)
               for k, v in pipeline._options.get_all_options().items()
               if v is not None}

    channel = grpc.insecure_channel(job_endpoint)
    grpc.channel_ready_future(channel).result()
    job_service = beam_job_api_pb2_grpc.JobServiceStub(channel)

    # Sends the PrepareRequest but retries in case the channel is not ready
    def send_prepare_request(max_retries=5):
      num_retries = 0
      while True:
        try:
          # This reports channel is READY but connections may fail
          # Seems to be only an issue on Mac with port forwardings
          grpc.channel_ready_future(channel).result()
          return job_service.Prepare(
              beam_job_api_pb2.PrepareJobRequest(
                  job_name='job', pipeline=proto_pipeline,
                  pipeline_options=job_utils.dict_to_struct(options)))
        except grpc._channel._Rendezvous as e:
          num_retries += 1
          if num_retries > max_retries:
            raise e

    prepare_response = send_prepare_request()
    if prepare_response.artifact_staging_endpoint.url:
      stager = portable_stager.PortableStager(
          grpc.insecure_channel(prepare_response.artifact_staging_endpoint.url),
          prepare_response.staging_session_token)
      retrieval_token, _ = stager.stage_job_resources(
          pipeline._options,
          staging_location='')
    else:
      retrieval_token = None
    run_response = job_service.Run(
        beam_job_api_pb2.RunJobRequest(
            preparation_id=prepare_response.preparation_id,
            retrieval_token=retrieval_token))
    return PipelineResult(job_service, run_response.job_id)
Example #3
0
    def run_pipeline(self, pipeline, options):
        portable_options = options.view_as(PortableOptions)
        job_endpoint = portable_options.job_endpoint

        # TODO: https://issues.apache.org/jira/browse/BEAM-5525
        # portable runner specific default
        if options.view_as(SetupOptions).sdk_location == 'default':
            options.view_as(SetupOptions).sdk_location = 'container'

        if not job_endpoint:
            # TODO Provide a way to specify a container Docker URL
            # https://issues.apache.org/jira/browse/BEAM-6328
            docker = DockerizedJobServer()
            job_endpoint = docker.start()

        # This is needed as we start a worker server if one is requested
        # but none is provided.
        if portable_options.environment_type == 'LOOPBACK':
            portable_options.environment_config, server = (
                BeamFnExternalWorkerPoolServicer.start(
                    sdk_worker_main._get_worker_count(options)))
            cleanup_callbacks = [functools.partial(server.stop, 1)]
        else:
            cleanup_callbacks = []

        proto_pipeline = pipeline.to_runner_api(
            default_environment=PortableRunner._create_environment(
                portable_options))

        # Some runners won't detect the GroupByKey transform unless it has no
        # subtransforms.  Remove all sub-transforms until BEAM-4605 is resolved.
        for _, transform_proto in list(
                proto_pipeline.components.transforms.items()):
            if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
                for sub_transform in transform_proto.subtransforms:
                    del proto_pipeline.components.transforms[sub_transform]
                del transform_proto.subtransforms[:]

        # Preemptively apply combiner lifting, until all runners support it.
        # This optimization is idempotent.
        if not options.view_as(StandardOptions).streaming:
            stages = list(
                fn_api_runner_transforms.leaf_transform_stages(
                    proto_pipeline.root_transform_ids,
                    proto_pipeline.components))
            stages = fn_api_runner_transforms.lift_combiners(
                stages,
                fn_api_runner_transforms.TransformContext(
                    proto_pipeline.components))
            proto_pipeline = fn_api_runner_transforms.with_stages(
                proto_pipeline, stages)

        # TODO: Define URNs for options.
        # convert int values: https://issues.apache.org/jira/browse/BEAM-5509
        p_options = {
            'beam:option:' + k + ':v1': (str(v) if type(v) == int else v)
            for k, v in options.get_all_options().items() if v is not None
        }

        channel = grpc.insecure_channel(job_endpoint)
        grpc.channel_ready_future(channel).result()
        job_service = beam_job_api_pb2_grpc.JobServiceStub(channel)

        # Sends the PrepareRequest but retries in case the channel is not ready
        def send_prepare_request(max_retries=5):
            num_retries = 0
            while True:
                try:
                    # This reports channel is READY but connections may fail
                    # Seems to be only an issue on Mac with port forwardings
                    grpc.channel_ready_future(channel).result()
                    return job_service.Prepare(
                        beam_job_api_pb2.PrepareJobRequest(
                            job_name='job',
                            pipeline=proto_pipeline,
                            pipeline_options=job_utils.dict_to_struct(
                                p_options)))
                except grpc._channel._Rendezvous as e:
                    num_retries += 1
                    if num_retries > max_retries:
                        raise e

        prepare_response = send_prepare_request()
        if prepare_response.artifact_staging_endpoint.url:
            stager = portable_stager.PortableStager(
                grpc.insecure_channel(
                    prepare_response.artifact_staging_endpoint.url),
                prepare_response.staging_session_token)
            retrieval_token, _ = stager.stage_job_resources(
                options, staging_location='')
        else:
            retrieval_token = None

        try:
            state_stream = job_service.GetStateStream(
                beam_job_api_pb2.GetJobStateRequest(
                    job_id=prepare_response.preparation_id))
            # If there's an error, we don't always get it until we try to read.
            # Fortunately, there's always an immediate current state published.
            state_stream = itertools.chain([next(state_stream)], state_stream)
            message_stream = job_service.GetMessageStream(
                beam_job_api_pb2.JobMessagesRequest(
                    job_id=prepare_response.preparation_id))
        except Exception:
            # TODO(BEAM-6442): Unify preparation_id and job_id for all runners.
            state_stream = message_stream = None

        # Run the job and wait for a result.
        run_response = job_service.Run(
            beam_job_api_pb2.RunJobRequest(
                preparation_id=prepare_response.preparation_id,
                retrieval_token=retrieval_token))

        if state_stream is None:
            state_stream = job_service.GetStateStream(
                beam_job_api_pb2.GetJobStateRequest(
                    job_id=run_response.job_id))
            message_stream = job_service.GetMessageStream(
                beam_job_api_pb2.JobMessagesRequest(
                    job_id=run_response.job_id))

        return PipelineResult(job_service, run_response.job_id, message_stream,
                              state_stream, cleanup_callbacks)
Example #4
0
  def run_pipeline(self, pipeline, options):
    portable_options = options.view_as(PortableOptions)
    job_endpoint = portable_options.job_endpoint

    # TODO: https://issues.apache.org/jira/browse/BEAM-5525
    # portable runner specific default
    if options.view_as(SetupOptions).sdk_location == 'default':
      options.view_as(SetupOptions).sdk_location = 'container'

    if not job_endpoint:
      # TODO Provide a way to specify a container Docker URL
      # https://issues.apache.org/jira/browse/BEAM-6328
      docker = DockerizedJobServer()
      job_endpoint = docker.start()
      job_service = None
    elif job_endpoint == 'embed':
      job_service = local_job_service.LocalJobServicer()
    else:
      job_service = None

    # This is needed as we start a worker server if one is requested
    # but none is provided.
    if portable_options.environment_type == 'LOOPBACK':
      portable_options.environment_config, server = (
          BeamFnExternalWorkerPoolServicer.start(
              sdk_worker_main._get_worker_count(options)))
      globals()['x'] = server
      cleanup_callbacks = [functools.partial(server.stop, 1)]
    else:
      cleanup_callbacks = []

    proto_pipeline = pipeline.to_runner_api(
        default_environment=PortableRunner._create_environment(
            portable_options))

    # Some runners won't detect the GroupByKey transform unless it has no
    # subtransforms.  Remove all sub-transforms until BEAM-4605 is resolved.
    for _, transform_proto in list(
        proto_pipeline.components.transforms.items()):
      if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
        for sub_transform in transform_proto.subtransforms:
          del proto_pipeline.components.transforms[sub_transform]
        del transform_proto.subtransforms[:]

    # Preemptively apply combiner lifting, until all runners support it.
    # This optimization is idempotent.
    pre_optimize = options.view_as(DebugOptions).lookup_experiment(
        'pre_optimize', 'combine').lower()
    if not options.view_as(StandardOptions).streaming:
      flink_known_urns = frozenset([
          common_urns.composites.RESHUFFLE.urn,
          common_urns.primitives.IMPULSE.urn,
          common_urns.primitives.FLATTEN.urn,
          common_urns.primitives.GROUP_BY_KEY.urn])
      if pre_optimize == 'combine':
        proto_pipeline = fn_api_runner_transforms.optimize_pipeline(
            proto_pipeline,
            phases=[fn_api_runner_transforms.lift_combiners],
            known_runner_urns=flink_known_urns,
            partial=True)
      elif pre_optimize == 'all':
        proto_pipeline = fn_api_runner_transforms.optimize_pipeline(
            proto_pipeline,
            phases=[fn_api_runner_transforms.annotate_downstream_side_inputs,
                    fn_api_runner_transforms.annotate_stateful_dofns_as_roots,
                    fn_api_runner_transforms.fix_side_input_pcoll_coders,
                    fn_api_runner_transforms.lift_combiners,
                    fn_api_runner_transforms.fix_flatten_coders,
                    # fn_api_runner_transforms.sink_flattens,
                    fn_api_runner_transforms.greedily_fuse,
                    fn_api_runner_transforms.read_to_impulse,
                    fn_api_runner_transforms.extract_impulse_stages,
                    fn_api_runner_transforms.remove_data_plane_ops,
                    fn_api_runner_transforms.sort_stages],
            known_runner_urns=flink_known_urns)
      elif pre_optimize == 'none':
        pass
      else:
        raise ValueError('Unknown value for pre_optimize: %s' % pre_optimize)

    if not job_service:
      channel = grpc.insecure_channel(job_endpoint)
      grpc.channel_ready_future(channel).result()
      job_service = beam_job_api_pb2_grpc.JobServiceStub(channel)
    else:
      channel = None

    # fetch runner options from job service
    # retries in case the channel is not ready
    def send_options_request(max_retries=5):
      num_retries = 0
      while True:
        try:
          # This reports channel is READY but connections may fail
          # Seems to be only an issue on Mac with port forwardings
          if channel:
            grpc.channel_ready_future(channel).result()
          return job_service.DescribePipelineOptions(
              beam_job_api_pb2.DescribePipelineOptionsRequest())
        except grpc._channel._Rendezvous as e:
          num_retries += 1
          if num_retries > max_retries:
            raise e
          time.sleep(1)

    options_response = send_options_request()

    def add_runner_options(parser):
      for option in options_response.options:
        try:
          # no default values - we don't want runner options
          # added unless they were specified by the user
          add_arg_args = {'action' : 'store', 'help' : option.description}
          if option.type == beam_job_api_pb2.PipelineOptionType.BOOLEAN:
            add_arg_args['action'] = 'store_true'\
              if option.default_value != 'true' else 'store_false'
          elif option.type == beam_job_api_pb2.PipelineOptionType.INTEGER:
            add_arg_args['type'] = int
          elif option.type == beam_job_api_pb2.PipelineOptionType.ARRAY:
            add_arg_args['action'] = 'append'
          parser.add_argument("--%s" % option.name, **add_arg_args)
        except Exception as e:
          # ignore runner options that are already present
          # only in this case is duplicate not treated as error
          if 'conflicting option string' not in str(e):
            raise
          logging.debug("Runner option '%s' was already added" % option.name)

    all_options = options.get_all_options(add_extra_args_fn=add_runner_options)
    # TODO: Define URNs for options.
    # convert int values: https://issues.apache.org/jira/browse/BEAM-5509
    p_options = {'beam:option:' + k + ':v1': (str(v) if type(v) == int else v)
                 for k, v in all_options.items()
                 if v is not None}

    prepare_response = job_service.Prepare(
        beam_job_api_pb2.PrepareJobRequest(
            job_name='job', pipeline=proto_pipeline,
            pipeline_options=job_utils.dict_to_struct(p_options)))
    if prepare_response.artifact_staging_endpoint.url:
      stager = portable_stager.PortableStager(
          grpc.insecure_channel(prepare_response.artifact_staging_endpoint.url),
          prepare_response.staging_session_token)
      retrieval_token, _ = stager.stage_job_resources(
          options,
          staging_location='')
    else:
      retrieval_token = None

    try:
      state_stream = job_service.GetStateStream(
          beam_job_api_pb2.GetJobStateRequest(
              job_id=prepare_response.preparation_id))
      # If there's an error, we don't always get it until we try to read.
      # Fortunately, there's always an immediate current state published.
      state_stream = itertools.chain(
          [next(state_stream)],
          state_stream)
      message_stream = job_service.GetMessageStream(
          beam_job_api_pb2.JobMessagesRequest(
              job_id=prepare_response.preparation_id))
    except Exception:
      # TODO(BEAM-6442): Unify preparation_id and job_id for all runners.
      state_stream = message_stream = None

    # Run the job and wait for a result.
    run_response = job_service.Run(
        beam_job_api_pb2.RunJobRequest(
            preparation_id=prepare_response.preparation_id,
            retrieval_token=retrieval_token))

    if state_stream is None:
      state_stream = job_service.GetStateStream(
          beam_job_api_pb2.GetJobStateRequest(
              job_id=run_response.job_id))
      message_stream = job_service.GetMessageStream(
          beam_job_api_pb2.JobMessagesRequest(
              job_id=run_response.job_id))

    return PipelineResult(job_service, run_response.job_id, message_stream,
                          state_stream, cleanup_callbacks)
Example #5
0
    def run_pipeline(self, pipeline):
        portable_options = pipeline.options.view_as(PortableOptions)
        job_endpoint = portable_options.job_endpoint
        if not job_endpoint:
            docker = DockerizedJobServer()
            job_endpoint = docker.start()

        proto_context = pipeline_context.PipelineContext(
            default_environment=PortableRunner._create_environment(
                portable_options))
        proto_pipeline = pipeline.to_runner_api(context=proto_context)

        if not self.is_embedded_fnapi_runner:
            # Java has different expectations about coders
            # (windowed in Fn API, but *un*windowed in runner API), whereas the
            # embedded FnApiRunner treats them consistently, so we must guard this
            # for now, until FnApiRunner is fixed.
            # See also BEAM-2717.
            for pcoll in proto_pipeline.components.pcollections.values():
                if pcoll.coder_id not in proto_context.coders:
                    # This is not really a coder id, but a pickled coder.
                    coder = coders.registry.get_coder(
                        pickler.loads(pcoll.coder_id))
                    pcoll.coder_id = proto_context.coders.get_id(coder)
            proto_context.coders.populate_map(proto_pipeline.components.coders)

        # Some runners won't detect the GroupByKey transform unless it has no
        # subtransforms.  Remove all sub-transforms until BEAM-4605 is resolved.
        for _, transform_proto in list(
                proto_pipeline.components.transforms.items()):
            if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
                for sub_transform in transform_proto.subtransforms:
                    del proto_pipeline.components.transforms[sub_transform]
                del transform_proto.subtransforms[:]

        # TODO: Define URNs for options.
        # convert int values: https://issues.apache.org/jira/browse/BEAM-5509
        options = {
            'beam:option:' + k + ':v1': (str(v) if type(v) == int else v)
            for k, v in pipeline._options.get_all_options().items()
            if v is not None
        }

        channel = grpc.insecure_channel(job_endpoint)
        grpc.channel_ready_future(channel).result()
        job_service = beam_job_api_pb2_grpc.JobServiceStub(channel)

        # Sends the PrepareRequest but retries in case the channel is not ready
        def send_prepare_request(max_retries=5):
            num_retries = 0
            while True:
                try:
                    # This reports channel is READY but connections may fail
                    # Seems to be only an issue on Mac with port forwardings
                    grpc.channel_ready_future(channel).result()
                    return job_service.Prepare(
                        beam_job_api_pb2.PrepareJobRequest(
                            job_name='job',
                            pipeline=proto_pipeline,
                            pipeline_options=job_utils.dict_to_struct(
                                options)))
                except grpc._channel._Rendezvous as e:
                    num_retries += 1
                    if num_retries > max_retries:
                        raise e

        prepare_response = send_prepare_request()
        if prepare_response.artifact_staging_endpoint.url:
            stager = portable_stager.PortableStager(
                grpc.insecure_channel(
                    prepare_response.artifact_staging_endpoint.url),
                prepare_response.staging_session_token)
            retrieval_token, _ = stager.stage_job_resources(
                pipeline._options, staging_location='')
        else:
            retrieval_token = None
        run_response = job_service.Run(
            beam_job_api_pb2.RunJobRequest(
                preparation_id=prepare_response.preparation_id,
                retrieval_token=retrieval_token))
        return PipelineResult(job_service, run_response.job_id)
Example #6
0
  def run_pipeline(self, pipeline, options):
    portable_options = options.view_as(PortableOptions)
    job_endpoint = portable_options.job_endpoint

    # TODO: https://issues.apache.org/jira/browse/BEAM-5525
    # portable runner specific default
    if options.view_as(SetupOptions).sdk_location == 'default':
      options.view_as(SetupOptions).sdk_location = 'container'

    if not job_endpoint:
      # TODO Provide a way to specify a container Docker URL
      # https://issues.apache.org/jira/browse/BEAM-6328
      docker = DockerizedJobServer()
      job_endpoint = docker.start()
      job_service = None
    elif job_endpoint == 'embed':
      job_service = local_job_service.LocalJobServicer()
    else:
      job_service = None

    # This is needed as we start a worker server if one is requested
    # but none is provided.
    if portable_options.environment_type == 'LOOPBACK':
      portable_options.environment_config, server = (
          BeamFnExternalWorkerPoolServicer.start(
              sdk_worker_main._get_worker_count(options)))
      globals()['x'] = server
      cleanup_callbacks = [functools.partial(server.stop, 1)]
    else:
      cleanup_callbacks = []

    proto_pipeline = pipeline.to_runner_api(
        default_environment=PortableRunner._create_environment(
            portable_options))

    # Some runners won't detect the GroupByKey transform unless it has no
    # subtransforms.  Remove all sub-transforms until BEAM-4605 is resolved.
    for _, transform_proto in list(
        proto_pipeline.components.transforms.items()):
      if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
        for sub_transform in transform_proto.subtransforms:
          del proto_pipeline.components.transforms[sub_transform]
        del transform_proto.subtransforms[:]

    # Preemptively apply combiner lifting, until all runners support it.
    # This optimization is idempotent.
    pre_optimize = options.view_as(DebugOptions).lookup_experiment(
        'pre_optimize', 'combine').lower()
    if not options.view_as(StandardOptions).streaming:
      flink_known_urns = frozenset([
          common_urns.composites.RESHUFFLE.urn,
          common_urns.primitives.IMPULSE.urn,
          common_urns.primitives.FLATTEN.urn,
          common_urns.primitives.GROUP_BY_KEY.urn])
      if pre_optimize == 'combine':
        proto_pipeline = fn_api_runner_transforms.optimize_pipeline(
            proto_pipeline,
            phases=[fn_api_runner_transforms.lift_combiners],
            known_runner_urns=flink_known_urns,
            partial=True)
      elif pre_optimize == 'all':
        proto_pipeline = fn_api_runner_transforms.optimize_pipeline(
            proto_pipeline,
            phases=[fn_api_runner_transforms.annotate_downstream_side_inputs,
                    fn_api_runner_transforms.annotate_stateful_dofns_as_roots,
                    fn_api_runner_transforms.fix_side_input_pcoll_coders,
                    fn_api_runner_transforms.lift_combiners,
                    fn_api_runner_transforms.fix_flatten_coders,
                    # fn_api_runner_transforms.sink_flattens,
                    fn_api_runner_transforms.greedily_fuse,
                    fn_api_runner_transforms.read_to_impulse,
                    fn_api_runner_transforms.extract_impulse_stages,
                    fn_api_runner_transforms.remove_data_plane_ops,
                    fn_api_runner_transforms.sort_stages],
            known_runner_urns=flink_known_urns)
      elif pre_optimize == 'none':
        pass
      else:
        raise ValueError('Unknown value for pre_optimize: %s' % pre_optimize)

    if not job_service:
      channel = grpc.insecure_channel(job_endpoint)
      grpc.channel_ready_future(channel).result()
      job_service = beam_job_api_pb2_grpc.JobServiceStub(channel)
    else:
      channel = None

    # fetch runner options from job service
    # retries in case the channel is not ready
    def send_options_request(max_retries=5):
      num_retries = 0
      while True:
        try:
          # This reports channel is READY but connections may fail
          # Seems to be only an issue on Mac with port forwardings
          if channel:
            grpc.channel_ready_future(channel).result()
          return job_service.DescribePipelineOptions(
              beam_job_api_pb2.DescribePipelineOptionsRequest())
        except grpc._channel._Rendezvous as e:
          num_retries += 1
          if num_retries > max_retries:
            raise e

    options_response = send_options_request()

    def add_runner_options(parser):
      for option in options_response.options:
        try:
          # no default values - we don't want runner options
          # added unless they were specified by the user
          add_arg_args = {'action' : 'store', 'help' : option.description}
          if option.type == beam_job_api_pb2.PipelineOptionType.BOOLEAN:
            add_arg_args['action'] = 'store_true'\
              if option.default_value != 'true' else 'store_false'
          elif option.type == beam_job_api_pb2.PipelineOptionType.INTEGER:
            add_arg_args['type'] = int
          elif option.type == beam_job_api_pb2.PipelineOptionType.ARRAY:
            add_arg_args['action'] = 'append'
          parser.add_argument("--%s" % option.name, **add_arg_args)
        except Exception as e:
          # ignore runner options that are already present
          # only in this case is duplicate not treated as error
          if 'conflicting option string' not in str(e):
            raise
          logging.debug("Runner option '%s' was already added" % option.name)

    all_options = options.get_all_options(add_extra_args_fn=add_runner_options)
    # TODO: Define URNs for options.
    # convert int values: https://issues.apache.org/jira/browse/BEAM-5509
    p_options = {'beam:option:' + k + ':v1': (str(v) if type(v) == int else v)
                 for k, v in all_options.items()
                 if v is not None}

    prepare_response = job_service.Prepare(
        beam_job_api_pb2.PrepareJobRequest(
            job_name='job', pipeline=proto_pipeline,
            pipeline_options=job_utils.dict_to_struct(p_options)))
    if prepare_response.artifact_staging_endpoint.url:
      stager = portable_stager.PortableStager(
          grpc.insecure_channel(prepare_response.artifact_staging_endpoint.url),
          prepare_response.staging_session_token)
      retrieval_token, _ = stager.stage_job_resources(
          options,
          staging_location='')
    else:
      retrieval_token = None

    try:
      state_stream = job_service.GetStateStream(
          beam_job_api_pb2.GetJobStateRequest(
              job_id=prepare_response.preparation_id))
      # If there's an error, we don't always get it until we try to read.
      # Fortunately, there's always an immediate current state published.
      state_stream = itertools.chain(
          [next(state_stream)],
          state_stream)
      message_stream = job_service.GetMessageStream(
          beam_job_api_pb2.JobMessagesRequest(
              job_id=prepare_response.preparation_id))
    except Exception:
      # TODO(BEAM-6442): Unify preparation_id and job_id for all runners.
      state_stream = message_stream = None

    # Run the job and wait for a result.
    run_response = job_service.Run(
        beam_job_api_pb2.RunJobRequest(
            preparation_id=prepare_response.preparation_id,
            retrieval_token=retrieval_token))

    if state_stream is None:
      state_stream = job_service.GetStateStream(
          beam_job_api_pb2.GetJobStateRequest(
              job_id=run_response.job_id))
      message_stream = job_service.GetMessageStream(
          beam_job_api_pb2.JobMessagesRequest(
              job_id=run_response.job_id))

    return PipelineResult(job_service, run_response.job_id, message_stream,
                          state_stream, cleanup_callbacks)
Example #7
0
  def run_pipeline(self, pipeline):
    portable_options = pipeline.options.view_as(PortableOptions)
    job_endpoint = portable_options.job_endpoint

    # TODO: https://issues.apache.org/jira/browse/BEAM-5525
    # portable runner specific default
    if pipeline.options.view_as(SetupOptions).sdk_location == 'default':
      pipeline.options.view_as(SetupOptions).sdk_location = 'container'

    if not job_endpoint:
      docker = DockerizedJobServer()
      job_endpoint = docker.start()

    proto_pipeline = pipeline.to_runner_api()

    # Some runners won't detect the GroupByKey transform unless it has no
    # subtransforms.  Remove all sub-transforms until BEAM-4605 is resolved.
    for _, transform_proto in list(
        proto_pipeline.components.transforms.items()):
      if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
        for sub_transform in transform_proto.subtransforms:
          del proto_pipeline.components.transforms[sub_transform]
        del transform_proto.subtransforms[:]

    # TODO: Define URNs for options.
    # convert int values: https://issues.apache.org/jira/browse/BEAM-5509
    options = {'beam:option:' + k + ':v1': (str(v) if type(v) == int else v)
               for k, v in pipeline._options.get_all_options().items()
               if v is not None}

    channel = grpc.insecure_channel(job_endpoint)
    grpc.channel_ready_future(channel).result()
    job_service = beam_job_api_pb2_grpc.JobServiceStub(channel)

    # Sends the PrepareRequest but retries in case the channel is not ready
    def send_prepare_request(max_retries=5):
      num_retries = 0
      while True:
        try:
          # This reports channel is READY but connections may fail
          # Seems to be only an issue on Mac with port forwardings
          grpc.channel_ready_future(channel).result()
          return job_service.Prepare(
              beam_job_api_pb2.PrepareJobRequest(
                  job_name='job', pipeline=proto_pipeline,
                  pipeline_options=job_utils.dict_to_struct(options)))
        except grpc._channel._Rendezvous as e:
          num_retries += 1
          if num_retries > max_retries:
            raise e

    prepare_response = send_prepare_request()
    if prepare_response.artifact_staging_endpoint.url:
      stager = portable_stager.PortableStager(
          grpc.insecure_channel(prepare_response.artifact_staging_endpoint.url),
          prepare_response.staging_session_token)
      retrieval_token, _ = stager.stage_job_resources(
          pipeline._options,
          staging_location='')
    else:
      retrieval_token = None
    run_response = job_service.Run(
        beam_job_api_pb2.RunJobRequest(
            preparation_id=prepare_response.preparation_id,
            retrieval_token=retrieval_token))
    return PipelineResult(job_service, run_response.job_id)
Example #8
0
 def init_dockerized_job_server(self):
   # TODO Provide a way to specify a container Docker URL
   # https://issues.apache.org/jira/browse/BEAM-6328
   docker = DockerizedJobServer()
   self._job_endpoint = docker.start()
 def init_dockerized_job_server(self):
   # TODO Provide a way to specify a container Docker URL
   # https://issues.apache.org/jira/browse/BEAM-6328
   docker = DockerizedJobServer()
   self._job_endpoint = docker.start()