| # |
| # 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. |
| # |
| |
| # pytype: skip-file |
| |
| import hashlib |
| import os |
| import random |
| import tempfile |
| import unittest |
| |
| import apache_beam as beam |
| from apache_beam.portability import common_urns |
| from apache_beam.portability.api import beam_runner_api_pb2 |
| from apache_beam.runners.pipeline_utils import merge_common_environments |
| from apache_beam.runners.pipeline_utils import merge_superset_dep_environments |
| from apache_beam.runners.pipeline_utils import validate_pipeline_graph |
| from apache_beam.runners.portability.expansion_service_test import FibTransform |
| |
| |
| class PipelineUtilitiesTest(unittest.TestCase): |
| @staticmethod |
| def _pipeline_with_gbk_coders( |
| input_key_coder='key_coder', |
| output_key_coder='key_coder', |
| output_values_coder_urn=common_urns.coders.ITERABLE.urn, |
| input_value_coder='input_value_coder', |
| output_value_coder='input_value_coder'): |
| leaf_coders = { |
| coder_id: beam_runner_api_pb2.Coder( |
| spec=beam_runner_api_pb2.FunctionSpec( |
| urn=common_urns.coders.BYTES.urn)) |
| for coder_id in { |
| input_key_coder, |
| output_key_coder, |
| input_value_coder, |
| output_value_coder, } |
| } |
| output_values_coder = beam_runner_api_pb2.Coder( |
| spec=beam_runner_api_pb2.FunctionSpec(urn=output_values_coder_urn), |
| component_coder_ids=([output_value_coder] if output_values_coder_urn |
| == common_urns.coders.ITERABLE.urn else [])) |
| return beam_runner_api_pb2.Pipeline( |
| components=beam_runner_api_pb2.Components( |
| coders={ |
| **leaf_coders, |
| 'input_coder': beam_runner_api_pb2.Coder( |
| spec=beam_runner_api_pb2.FunctionSpec( |
| urn=common_urns.coders.KV.urn), |
| component_coder_ids=[input_key_coder, input_value_coder]), |
| 'output_coder': beam_runner_api_pb2.Coder( |
| spec=beam_runner_api_pb2.FunctionSpec( |
| urn=common_urns.coders.KV.urn), |
| component_coder_ids=[ |
| output_key_coder, 'output_values_coder' |
| ]), |
| 'output_values_coder': output_values_coder, |
| }, |
| pcollections={ |
| 'input': beam_runner_api_pb2.PCollection( |
| coder_id='input_coder'), |
| 'output': beam_runner_api_pb2.PCollection( |
| coder_id='output_coder'), |
| }, |
| transforms={ |
| 'gbk': beam_runner_api_pb2.PTransform( |
| spec=beam_runner_api_pb2.FunctionSpec( |
| urn=common_urns.primitives.GROUP_BY_KEY.urn), |
| inputs={'input': 'input'}, |
| outputs={'output': 'output'}), |
| }), |
| root_transform_ids=['gbk']) |
| |
| def test_validate_pipeline_graph_accepts_valid_gbk_coders(self): |
| validate_pipeline_graph(self._pipeline_with_gbk_coders()) |
| |
| def test_validate_pipeline_graph_reports_gbk_key_coder_mismatch(self): |
| pipeline = self._pipeline_with_gbk_coders( |
| input_key_coder='input_key_coder', output_key_coder='output_key_coder') |
| |
| with self.assertRaises(ValueError) as error: |
| validate_pipeline_graph(pipeline) |
| |
| self.assertEqual( |
| str(error.exception), |
| 'Input key coder input_key_coder does not match output key coder ' |
| 'output_key_coder for transform gbk') |
| |
| def test_validate_pipeline_graph_reports_non_iterable_gbk_output(self): |
| pipeline = self._pipeline_with_gbk_coders( |
| output_values_coder_urn=common_urns.coders.BYTES.urn) |
| |
| with self.assertRaises(ValueError) as error: |
| validate_pipeline_graph(pipeline) |
| |
| self.assertEqual( |
| str(error.exception), |
| 'Output value coder output_values_coder for transform gbk must be an ' |
| 'iterable coder, but uses URN beam:coder:bytes:v1') |
| |
| def test_validate_pipeline_graph_reports_gbk_value_coder_mismatch(self): |
| pipeline = self._pipeline_with_gbk_coders( |
| input_value_coder='input_value_coder', |
| output_value_coder='output_value_coder') |
| |
| with self.assertRaises(ValueError) as error: |
| validate_pipeline_graph(pipeline) |
| |
| self.assertEqual( |
| str(error.exception), |
| 'Input value coder input_value_coder does not match output value coder ' |
| 'output_value_coder for transform gbk') |
| |
| def test_equal_environments_merged(self): |
| pipeline_proto = merge_common_environments( |
| beam_runner_api_pb2.Pipeline( |
| components=beam_runner_api_pb2.Components( |
| environments={ |
| 'a1': beam_runner_api_pb2.Environment(urn='A'), |
| 'a2': beam_runner_api_pb2.Environment(urn='A'), |
| 'b1': beam_runner_api_pb2.Environment( |
| urn='B', payload=b'x'), |
| 'b2': beam_runner_api_pb2.Environment( |
| urn='B', payload=b'x'), |
| 'b3': beam_runner_api_pb2.Environment( |
| urn='B', payload=b'y'), |
| }, |
| transforms={ |
| 't1': beam_runner_api_pb2.PTransform( |
| unique_name='t1', environment_id='a1'), |
| 't2': beam_runner_api_pb2.PTransform( |
| unique_name='t2', environment_id='a2'), |
| }, |
| windowing_strategies={ |
| 'w1': beam_runner_api_pb2.WindowingStrategy( |
| environment_id='b1'), |
| 'w2': beam_runner_api_pb2.WindowingStrategy( |
| environment_id='b2'), |
| }))) |
| self.assertEqual(len(pipeline_proto.components.environments), 3) |
| self.assertTrue(('a1' in pipeline_proto.components.environments) |
| ^ ('a2' in pipeline_proto.components.environments)) |
| self.assertTrue(('b1' in pipeline_proto.components.environments) |
| ^ ('b2' in pipeline_proto.components.environments)) |
| self.assertEqual( |
| len( |
| set( |
| t.environment_id |
| for t in pipeline_proto.components.transforms.values())), |
| 1) |
| self.assertEqual( |
| len( |
| set( |
| w.environment_id for w in |
| pipeline_proto.components.windowing_strategies.values())), |
| 1) |
| |
| def _make_dep(self, path): |
| hasher = hashlib.sha256() |
| if os.path.exists(path): |
| with open(path, 'rb') as fin: |
| hasher.update(fin.read()) |
| else: |
| # A fake file, identified only by its path. |
| hasher.update(path.encode('utf-8')) |
| return beam_runner_api_pb2.ArtifactInformation( |
| type_urn=common_urns.artifact_types.FILE.urn, |
| type_payload=beam_runner_api_pb2.ArtifactFilePayload( |
| path=path, sha256=hasher.hexdigest()).SerializeToString(), |
| role_urn=common_urns.artifact_roles.STAGING_TO.urn, |
| role_payload=beam_runner_api_pb2.ArtifactStagingToRolePayload( |
| staged_name=os.path.basename(path)).SerializeToString()) |
| |
| def _docker_env(self, id, deps=()): |
| return beam_runner_api_pb2.Environment( |
| urn=common_urns.environments.DOCKER.urn, |
| payload=id.encode('utf8'), |
| dependencies=[self._make_dep(path) for path in deps], |
| ) |
| |
| def test_subset_deps_environments_merged(self): |
| environments = { |
| 'A': self._docker_env('A'), |
| 'Ax': self._docker_env('A', ['x']), |
| 'Ay': self._docker_env('A', ['y']), |
| 'Axy': self._docker_env('A', ['x', 'y']), |
| 'Bx': self._docker_env('B', ['x']), |
| 'Bxy': self._docker_env('B', ['x', 'y']), |
| 'Byz': self._docker_env('B', ['y', 'z']), |
| } |
| transforms = { |
| env_id: beam_runner_api_pb2.PTransform( |
| unique_name=env_id, environment_id=env_id) |
| for env_id in environments.keys() |
| } |
| pipeline_proto = merge_superset_dep_environments( |
| beam_runner_api_pb2.Pipeline( |
| components=beam_runner_api_pb2.Components( |
| environments=environments, transforms=transforms))) |
| |
| # These can all be merged into the same environment. |
| self.assertEqual( |
| pipeline_proto.components.transforms['A'].environment_id, 'Axy') |
| self.assertEqual( |
| pipeline_proto.components.transforms['Ax'].environment_id, 'Axy') |
| self.assertEqual( |
| pipeline_proto.components.transforms['Ay'].environment_id, 'Axy') |
| self.assertEqual( |
| pipeline_proto.components.transforms['Axy'].environment_id, 'Axy') |
| # Despite having the same dependencies, these must be merged into their own. |
| self.assertEqual( |
| pipeline_proto.components.transforms['Bx'].environment_id, 'Bxy') |
| self.assertEqual( |
| pipeline_proto.components.transforms['Bxy'].environment_id, 'Bxy') |
| # This is not a subset of any, must be left alone. |
| self.assertEqual( |
| pipeline_proto.components.transforms['Byz'].environment_id, 'Byz') |
| |
| def test_subset_deps_environments_merged_with_requirements_txt(self): |
| with tempfile.TemporaryDirectory() as tmpdir: |
| |
| def make_file(basename, content): |
| subdir = tempfile.mkdtemp(dir=tmpdir) |
| path = os.path.join(subdir, basename) |
| with open(path, 'w') as fout: |
| fout.write(content) |
| return path |
| |
| def make_py_deps(*pkgs): |
| return [ |
| make_file('requirements.txt', '\n'.join(pkgs)), |
| make_file( |
| 'submission_environment_dependencies.txt', str( |
| random.random())), |
| ] + [make_file(pkg, pkg) for pkg in pkgs] |
| |
| environments = { |
| 'A': self._docker_env('A'), |
| 'Ax': self._docker_env('A', make_py_deps('x')), |
| 'Ay': self._docker_env('A', make_py_deps('y')), |
| 'Axy': self._docker_env('A', make_py_deps('x', 'y')), |
| } |
| transforms = { |
| env_id: beam_runner_api_pb2.PTransform( |
| unique_name=env_id, environment_id=env_id) |
| for env_id in environments.keys() |
| } |
| pipeline_proto = merge_superset_dep_environments( |
| beam_runner_api_pb2.Pipeline( |
| components=beam_runner_api_pb2.Components( |
| environments=environments, transforms=transforms))) |
| |
| # These can all be merged into the same environment. |
| self.assertEqual( |
| pipeline_proto.components.transforms['A'].environment_id, 'Axy') |
| self.assertEqual( |
| pipeline_proto.components.transforms['Ax'].environment_id, 'Axy') |
| self.assertEqual( |
| pipeline_proto.components.transforms['Ay'].environment_id, 'Axy') |
| self.assertEqual( |
| pipeline_proto.components.transforms['Axy'].environment_id, 'Axy') |
| |
| def test_external_merged(self): |
| p = beam.Pipeline() |
| # This transform recursively creates several external environments. |
| _ = p | FibTransform(4) |
| pipeline_proto = p.to_runner_api() |
| # All our external environments are equal and consolidated. |
| # We also have a placeholder "default" environment that has not been |
| # resolved do anything concrete yet. |
| envs = pipeline_proto.components.environments |
| self.assertEqual( |
| len(envs), 2, f'should be 2 environments, instead got: {envs}') |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
