sdks/python/apache_beam/runners/pipeline_utils.py - beam - Git at Google


            
            #

            
            # Licensed to the Apache Software Foundation (ASF) under one or more

            
            # contributor license agreements.  See the NOTICE file distributed with

            
            # this work for additional information regarding copyright ownership.

            
            # The ASF licenses this file to You under the Apache License, Version 2.0

            
            # (the "License"); you may not use this file except in compliance with

            
            # the License.  You may obtain a copy of the License at

            
            #

            
            #    http://www.apache.org/licenses/LICENSE-2.0

            
            #

            
            # Unless required by applicable law or agreed to in writing, software

            
            # distributed under the License is distributed on an "AS IS" BASIS,

            
            # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

            
            # See the License for the specific language governing permissions and

            
            # limitations under the License.

            
            #


            """Pipeline manipulation utilities useful for many runners.


            For internal use only; no backwards-compatibility guarantees.

            
            """


            # pytype: skip-fileimport collections


            import collections

            
            import copy


            from apache_beam.portability import common_urns

            
            from apache_beam.portability.api import beam_runner_api_pb2

            
            from apache_beam.transforms import environments

            
            from apache_beam.typehints import typehints


            def group_by_key_input_visitor(deterministic_key_coders=True):

            
              # Importing here to avoid a circular dependency

            
              # pylint: disable=wrong-import-order, wrong-import-position

            
              from apache_beam.pipeline import PipelineVisitor

            
              from apache_beam.transforms.core import GroupByKey


              class GroupByKeyInputVisitor(PipelineVisitor):

            
                """A visitor that replaces `Any` element type for input `PCollection` of

            
                a `GroupByKey` with a `KV` type.


                TODO(BEAM-115): Once Python SDK is compatible with the new Runner API,

            
                we could directly replace the coder instead of mutating the element type.

            
                """

            
                def __init__(self, deterministic_key_coders=True):

            
                  self.deterministic_key_coders = deterministic_key_coders


                def enter_composite_transform(self, transform_node):

            
                  self.visit_transform(transform_node)


                def visit_transform(self, transform_node):

            
                  if isinstance(transform_node.transform, GroupByKey):

            
                    pcoll = transform_node.inputs[0]

            
                    pcoll.element_type = typehints.coerce_to_kv_type(

            
                        pcoll.element_type, transform_node.full_label)

            
                    pcoll.requires_deterministic_key_coder = (

            
                        self.deterministic_key_coders and transform_node.full_label)

            
                    key_type, value_type = pcoll.element_type.tuple_types

            
                    if transform_node.outputs:

            
                      key = next(iter(transform_node.outputs.keys()))

            
                      transform_node.outputs[key].element_type = typehints.KV[

            
                          key_type, typehints.Iterable[value_type]]

            
                      transform_node.outputs[key].requires_deterministic_key_coder = (

            
                          self.deterministic_key_coders and transform_node.full_label)


              return GroupByKeyInputVisitor(deterministic_key_coders)


            def validate_pipeline_graph(pipeline_proto):

            
              """Ensures this is a correctly constructed Beam pipeline.

            
              """

            
              def get_coder(pcoll_id):

            
                return pipeline_proto.components.coders[

            
                    pipeline_proto.components.pcollections[pcoll_id].coder_id]


              def validate_transform(transform_id):

            
                transform_proto = pipeline_proto.components.transforms[transform_id]


                # Currently the only validation we perform is that GBK operations have

            
                # their coders set properly.

            
                if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:

            
                  if len(transform_proto.inputs) != 1:

            
                    raise ValueError("Unexpected number of inputs: %s" % transform_proto)

            
                  if len(transform_proto.outputs) != 1:

            
                    raise ValueError("Unexpected number of outputs: %s" % transform_proto)

            
                  input_coder = get_coder(next(iter(transform_proto.inputs.values())))

            
                  output_coder = get_coder(next(iter(transform_proto.outputs.values())))

            
                  if input_coder.spec.urn != common_urns.coders.KV.urn:

            
                    raise ValueError(

            
                        "Bad coder for input of %s: %s" % (transform_id, input_coder))

            
                  if output_coder.spec.urn != common_urns.coders.KV.urn:

            
                    raise ValueError(

            
                        "Bad coder for output of %s: %s" % (transform_id, output_coder))

            
                  input_key_coder_id = input_coder.component_coder_ids[0]

            
                  output_key_coder_id = output_coder.component_coder_ids[0]

            
                  if input_key_coder_id != output_key_coder_id:

            
                    raise ValueError(

            
                        "Input key coder %s does not match output key coder %s for "

            
                        "transform %s" %

            
                        (input_key_coder_id, output_key_coder_id, transform_id))

            
                  output_values_coder_id = output_coder.component_coder_ids[1]

            
                  output_values_coder = pipeline_proto.components.coders[

            
                      output_values_coder_id]

            
                  if output_values_coder.spec.urn != common_urns.coders.ITERABLE.urn:

            
                    raise ValueError(

            
                        "Output value coder %s for transform %s must be an iterable "

            
                        "coder, but uses URN %s" % (

            
                            output_values_coder_id,

            
                            transform_id,

            
                            output_values_coder.spec.urn))

            
                  input_value_coder_id = input_coder.component_coder_ids[1]

            
                  output_value_coder_id = output_values_coder.component_coder_ids[0]

            
                  if output_value_coder_id != input_value_coder_id:

            
                    raise ValueError(

            
                        "Input value coder %s does not match output value coder %s for "

            
                        "transform %s" %

            
                        (input_value_coder_id, output_value_coder_id, transform_id))

            
                elif transform_proto.spec.urn == common_urns.primitives.ASSIGN_WINDOWS.urn:

            
                  if not transform_proto.inputs:

            
                    raise ValueError("Missing input for transform: %s" % transform_proto)

            
                elif transform_proto.spec.urn == common_urns.primitives.PAR_DO.urn:

            
                  if not transform_proto.inputs:

            
                    raise ValueError("Missing input for transform: %s" % transform_proto)


                for t in transform_proto.subtransforms:

            
                  validate_transform(t)


              for t in pipeline_proto.root_transform_ids:

            
                validate_transform(t)


            def _dep_key(dep):

            
              if dep.type_urn == common_urns.artifact_types.FILE.urn:

            
                payload = beam_runner_api_pb2.ArtifactFilePayload.FromString(

            
                    dep.type_payload)

            
                if payload.sha256:

            
                  type_info = 'sha256', payload.sha256

            
                else:

            
                  type_info = 'path', payload.path

            
              elif dep.type_urn == common_urns.artifact_types.URL.urn:

            
                payload = beam_runner_api_pb2.ArtifactUrlPayload.FromString(

            
                    dep.type_payload)

            
                if payload.sha256:

            
                  type_info = 'sha256', payload.sha256

            
                else:

            
                  type_info = 'url', payload.url

            
              else:

            
                type_info = dep.type_urn, dep.type_payload

            
              return type_info, dep.role_urn, dep.role_payload


            def _expanded_dep_keys(dep):

            
              if (dep.type_urn == common_urns.artifact_types.FILE.urn and

            
                  dep.role_urn == common_urns.artifact_roles.STAGING_TO.urn):

            
                payload = beam_runner_api_pb2.ArtifactFilePayload.FromString(

            
                    dep.type_payload)

            
                role = beam_runner_api_pb2.ArtifactStagingToRolePayload.FromString(

            
                    dep.role_payload)

            
                if role.staged_name == 'submission_environment_dependencies.txt':

            
                  return

            
                elif role.staged_name == 'requirements.txt':

            
                  with open(payload.path) as fin:

            
                    for line in fin:

            
                      yield 'requirements.txt', line.strip()

            
                  return


              yield _dep_key(dep)


            def _base_env_key(env, include_deps=True):

            
              return (

            
                  env.urn,

            
                  env.payload,

            
                  tuple(sorted(env.capabilities)),

            
                  tuple(sorted(env.resource_hints.items())),

            
                  tuple(sorted(_dep_key(dep)

            
                               for dep in env.dependencies)) if include_deps else None)


            def _env_key(env):

            
              return tuple(

            
                  sorted(

            
                      _base_env_key(e)

            
                      for e in environments.expand_anyof_environments(env)))


            def merge_common_environments(pipeline_proto, inplace=False):

            
              canonical_environments = collections.defaultdict(list)

            
              for env_id, env in pipeline_proto.components.environments.items():

            
                canonical_environments[_env_key(env)].append(env_id)


              if len(canonical_environments) == len(pipeline_proto.components.environments):

            
                # All environments are already sufficiently distinct.

            
                return pipeline_proto


              environment_remappings = {

            
                  e: es[0]

            
                  for es in canonical_environments.values()

            
                  for e in es

            
              }


              return update_environments(pipeline_proto, environment_remappings, inplace)


            def merge_superset_dep_environments(pipeline_proto):

            
              """Merges all environemnts A and B where A and B are equivalent except that

            
              A has a superset of the dependencies of B.

            
              """

            
              docker_envs = {}

            
              for env_id, env in pipeline_proto.components.environments.items():

            
                docker_env = environments.resolve_anyof_environment(

            
                    env, common_urns.environments.DOCKER.urn)

            
                if docker_env.urn == common_urns.environments.DOCKER.urn:

            
                  docker_envs[env_id] = docker_env


              has_base_and_dep = collections.defaultdict(set)

            
              env_scores = {

            
                  env_id: (len(env.dependencies), env_id)

            
                  for (env_id, env) in docker_envs.items()

            
              }


              for env_id, env in docker_envs.items():

            
                base_key = _base_env_key(env, include_deps=False)

            
                has_base_and_dep[base_key, None].add(env_id)

            
                for dep in env.dependencies:

            
                  for dep_key in _expanded_dep_keys(dep):

            
                    has_base_and_dep[base_key, dep_key].add(env_id)


              environment_remappings = {}

            
              for env_id, env in docker_envs.items():

            
                base_key = _base_env_key(env, include_deps=False)

            
                # This is the set of all environments that have at least all of env's deps.

            
                candidates = set.intersection(

            
                    has_base_and_dep[base_key, None],

            
                    *[

            
                        has_base_and_dep[base_key, dep_key] for dep in env.dependencies

            
                        for dep_key in _expanded_dep_keys(dep)

            
                    ])

            
                # Choose the maximal one.

            
                best = max(candidates, key=env_scores.get)

            
                if best != env_id:

            
                  environment_remappings[env_id] = best


              return update_environments(pipeline_proto, environment_remappings)


            def update_environments(pipeline_proto, environment_remappings, inplace=False):

            
              if not environment_remappings:

            
                return pipeline_proto


              if not inplace:

            
                pipeline_proto = copy.copy(pipeline_proto)


              for t in pipeline_proto.components.transforms.values():

            
                if t.environment_id not in pipeline_proto.components.environments:

            
                  # TODO(https://github.com/apache/beam/issues/30876): Remove this

            
                  #  workaround.

            
                  continue

            
                if t.environment_id and t.environment_id in environment_remappings:

            
                  t.environment_id = environment_remappings[t.environment_id]

            
              for w in pipeline_proto.components.windowing_strategies.values():

            
                if w.environment_id not in pipeline_proto.components.environments:

            
                  # TODO(https://github.com/apache/beam/issues/30876): Remove this

            
                  #  workaround.

            
                  continue

            
                if w.environment_id and w.environment_id in environment_remappings:

            
                  w.environment_id = environment_remappings[w.environment_id]

            
              for e in set(environment_remappings.keys()) - set(

            
                  environment_remappings.values()):

            
                del pipeline_proto.components.environments[e]

            
              return pipeline_proto
	#
	# Licensed to the Apache Software Foundation (ASF) under one or more
	# contributor license agreements. See the NOTICE file distributed with
	# this work for additional information regarding copyright ownership.
	# The ASF licenses this file to You under the Apache License, Version 2.0
	# (the "License"); you may not use this file except in compliance with
	# the License. You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.
	#

	"""Pipeline manipulation utilities useful for many runners.

	For internal use only; no backwards-compatibility guarantees.
	"""

	# pytype: skip-fileimport collections

	import collections
	import copy

	from apache_beam.portability import common_urns
	from apache_beam.portability.api import beam_runner_api_pb2
	from apache_beam.transforms import environments
	from apache_beam.typehints import typehints


	def group_by_key_input_visitor(deterministic_key_coders=True):
	# Importing here to avoid a circular dependency
	# pylint: disable=wrong-import-order, wrong-import-position
	from apache_beam.pipeline import PipelineVisitor
	from apache_beam.transforms.core import GroupByKey

	class GroupByKeyInputVisitor(PipelineVisitor):
	"""A visitor that replaces `Any` element type for input `PCollection` of
	a `GroupByKey` with a `KV` type.

	TODO(BEAM-115): Once Python SDK is compatible with the new Runner API,
	we could directly replace the coder instead of mutating the element type.
	"""
	def __init__(self, deterministic_key_coders=True):
	self.deterministic_key_coders = deterministic_key_coders

	def enter_composite_transform(self, transform_node):
	self.visit_transform(transform_node)

	def visit_transform(self, transform_node):
	if isinstance(transform_node.transform, GroupByKey):
	pcoll = transform_node.inputs[0]
	pcoll.element_type = typehints.coerce_to_kv_type(
	pcoll.element_type, transform_node.full_label)
	pcoll.requires_deterministic_key_coder = (
	self.deterministic_key_coders and transform_node.full_label)
	key_type, value_type = pcoll.element_type.tuple_types
	if transform_node.outputs:
	key = next(iter(transform_node.outputs.keys()))
	transform_node.outputs[key].element_type = typehints.KV[
	key_type, typehints.Iterable[value_type]]
	transform_node.outputs[key].requires_deterministic_key_coder = (
	self.deterministic_key_coders and transform_node.full_label)

	return GroupByKeyInputVisitor(deterministic_key_coders)


	def validate_pipeline_graph(pipeline_proto):
	"""Ensures this is a correctly constructed Beam pipeline.
	"""
	def get_coder(pcoll_id):
	return pipeline_proto.components.coders[
	pipeline_proto.components.pcollections[pcoll_id].coder_id]

	def validate_transform(transform_id):
	transform_proto = pipeline_proto.components.transforms[transform_id]

	# Currently the only validation we perform is that GBK operations have
	# their coders set properly.
	if transform_proto.spec.urn == common_urns.primitives.GROUP_BY_KEY.urn:
	if len(transform_proto.inputs) != 1:
	raise ValueError("Unexpected number of inputs: %s" % transform_proto)
	if len(transform_proto.outputs) != 1:
	raise ValueError("Unexpected number of outputs: %s" % transform_proto)
	input_coder = get_coder(next(iter(transform_proto.inputs.values())))
	output_coder = get_coder(next(iter(transform_proto.outputs.values())))
	if input_coder.spec.urn != common_urns.coders.KV.urn:
	raise ValueError(
	"Bad coder for input of %s: %s" % (transform_id, input_coder))
	if output_coder.spec.urn != common_urns.coders.KV.urn:
	raise ValueError(
	"Bad coder for output of %s: %s" % (transform_id, output_coder))
	input_key_coder_id = input_coder.component_coder_ids[0]
	output_key_coder_id = output_coder.component_coder_ids[0]
	if input_key_coder_id != output_key_coder_id:
	raise ValueError(
	"Input key coder %s does not match output key coder %s for "
	"transform %s" %
	(input_key_coder_id, output_key_coder_id, transform_id))
	output_values_coder_id = output_coder.component_coder_ids[1]
	output_values_coder = pipeline_proto.components.coders[
	output_values_coder_id]
	if output_values_coder.spec.urn != common_urns.coders.ITERABLE.urn:
	raise ValueError(
	"Output value coder %s for transform %s must be an iterable "
	"coder, but uses URN %s" % (
	output_values_coder_id,
	transform_id,
	output_values_coder.spec.urn))
	input_value_coder_id = input_coder.component_coder_ids[1]
	output_value_coder_id = output_values_coder.component_coder_ids[0]
	if output_value_coder_id != input_value_coder_id:
	raise ValueError(
	"Input value coder %s does not match output value coder %s for "
	"transform %s" %
	(input_value_coder_id, output_value_coder_id, transform_id))
	elif transform_proto.spec.urn == common_urns.primitives.ASSIGN_WINDOWS.urn:
	if not transform_proto.inputs:
	raise ValueError("Missing input for transform: %s" % transform_proto)
	elif transform_proto.spec.urn == common_urns.primitives.PAR_DO.urn:
	if not transform_proto.inputs:
	raise ValueError("Missing input for transform: %s" % transform_proto)

	for t in transform_proto.subtransforms:
	validate_transform(t)

	for t in pipeline_proto.root_transform_ids:
	validate_transform(t)


	def _dep_key(dep):
	if dep.type_urn == common_urns.artifact_types.FILE.urn:
	payload = beam_runner_api_pb2.ArtifactFilePayload.FromString(
	dep.type_payload)
	if payload.sha256:
	type_info = 'sha256', payload.sha256
	else:
	type_info = 'path', payload.path
	elif dep.type_urn == common_urns.artifact_types.URL.urn:
	payload = beam_runner_api_pb2.ArtifactUrlPayload.FromString(
	dep.type_payload)
	if payload.sha256:
	type_info = 'sha256', payload.sha256
	else:
	type_info = 'url', payload.url
	else:
	type_info = dep.type_urn, dep.type_payload
	return type_info, dep.role_urn, dep.role_payload


	def _expanded_dep_keys(dep):
	if (dep.type_urn == common_urns.artifact_types.FILE.urn and
	dep.role_urn == common_urns.artifact_roles.STAGING_TO.urn):
	payload = beam_runner_api_pb2.ArtifactFilePayload.FromString(
	dep.type_payload)
	role = beam_runner_api_pb2.ArtifactStagingToRolePayload.FromString(
	dep.role_payload)
	if role.staged_name == 'submission_environment_dependencies.txt':
	return
	elif role.staged_name == 'requirements.txt':
	with open(payload.path) as fin:
	for line in fin:
	yield 'requirements.txt', line.strip()
	return

	yield _dep_key(dep)


	def _base_env_key(env, include_deps=True):
	return (
	env.urn,
	env.payload,
	tuple(sorted(env.capabilities)),
	tuple(sorted(env.resource_hints.items())),
	tuple(sorted(_dep_key(dep)
	for dep in env.dependencies)) if include_deps else None)


	def _env_key(env):
	return tuple(
	sorted(
	_base_env_key(e)
	for e in environments.expand_anyof_environments(env)))


	def merge_common_environments(pipeline_proto, inplace=False):
	canonical_environments = collections.defaultdict(list)
	for env_id, env in pipeline_proto.components.environments.items():
	canonical_environments[_env_key(env)].append(env_id)

	if len(canonical_environments) == len(pipeline_proto.components.environments):
	# All environments are already sufficiently distinct.
	return pipeline_proto

	environment_remappings = {
	e: es[0]
	for es in canonical_environments.values()
	for e in es
	}

	return update_environments(pipeline_proto, environment_remappings, inplace)


	def merge_superset_dep_environments(pipeline_proto):
	"""Merges all environemnts A and B where A and B are equivalent except that
	A has a superset of the dependencies of B.
	"""
	docker_envs = {}
	for env_id, env in pipeline_proto.components.environments.items():
	docker_env = environments.resolve_anyof_environment(
	env, common_urns.environments.DOCKER.urn)
	if docker_env.urn == common_urns.environments.DOCKER.urn:
	docker_envs[env_id] = docker_env

	has_base_and_dep = collections.defaultdict(set)
	env_scores = {
	env_id: (len(env.dependencies), env_id)
	for (env_id, env) in docker_envs.items()
	}

	for env_id, env in docker_envs.items():
	base_key = _base_env_key(env, include_deps=False)
	has_base_and_dep[base_key, None].add(env_id)
	for dep in env.dependencies:
	for dep_key in _expanded_dep_keys(dep):
	has_base_and_dep[base_key, dep_key].add(env_id)

	environment_remappings = {}
	for env_id, env in docker_envs.items():
	base_key = _base_env_key(env, include_deps=False)
	# This is the set of all environments that have at least all of env's deps.
	candidates = set.intersection(
	has_base_and_dep[base_key, None],
	*[
	has_base_and_dep[base_key, dep_key] for dep in env.dependencies
	for dep_key in _expanded_dep_keys(dep)
	])
	# Choose the maximal one.
	best = max(candidates, key=env_scores.get)
	if best != env_id:
	environment_remappings[env_id] = best

	return update_environments(pipeline_proto, environment_remappings)


	def update_environments(pipeline_proto, environment_remappings, inplace=False):
	if not environment_remappings:
	return pipeline_proto

	if not inplace:
	pipeline_proto = copy.copy(pipeline_proto)

	for t in pipeline_proto.components.transforms.values():
	if t.environment_id not in pipeline_proto.components.environments:
	# TODO(https://github.com/apache/beam/issues/30876): Remove this
	# workaround.
	continue
	if t.environment_id and t.environment_id in environment_remappings:
	t.environment_id = environment_remappings[t.environment_id]
	for w in pipeline_proto.components.windowing_strategies.values():
	if w.environment_id not in pipeline_proto.components.environments:
	# TODO(https://github.com/apache/beam/issues/30876): Remove this
	# workaround.
	continue
	if w.environment_id and w.environment_id in environment_remappings:
	w.environment_id = environment_remappings[w.environment_id]
	for e in set(environment_remappings.keys()) - set(
	environment_remappings.values()):
	del pipeline_proto.components.environments[e]
	return pipeline_proto