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))
       output_values_coder = pipeline_proto.components.coders[
           output_coder.component_coder_ids[1]]
       if (input_coder.component_coder_ids[0]
           != output_coder.component_coder_ids[0] or
           output_values_coder.spec.urn != common_urns.coders.ITERABLE.urn or
           output_values_coder.component_coder_ids[0]
           != input_coder.component_coder_ids[1]):
         raise ValueError(
             "Incompatible input coder %s and output coder %s for transform %s" %
             (transform_id, input_coder, output_coder))
     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))
	output_values_coder = pipeline_proto.components.coders[
	output_coder.component_coder_ids[1]]
	if (input_coder.component_coder_ids[0]
	!= output_coder.component_coder_ids[0] or
	output_values_coder.spec.urn != common_urns.coders.ITERABLE.urn or
	output_values_coder.component_coder_ids[0]
	!= input_coder.component_coder_ids[1]):
	raise ValueError(
	"Incompatible input coder %s and output coder %s for transform %s" %
	(transform_id, input_coder, output_coder))
	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