From 25d6d38d092d150f640c2bfd179f9ca4d8330bd0 Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Tue, 22 Sep 2026 19:44:21 +0000 Subject: [PATCH 1/3] fix flakey python tests --- .../internal/cloudpickle_pickler.py | 69 +++++-------------- .../internal/cloudpickle_pickler_test.py | 8 +++ .../apache_beam/yaml/yaml_transform_test.py | 34 ++++----- 3 files changed, 44 insertions(+), 67 deletions(-) diff --git a/sdks/python/apache_beam/internal/cloudpickle_pickler.py b/sdks/python/apache_beam/internal/cloudpickle_pickler.py index cea4f01f803c..e6f008c5c768 100644 --- a/sdks/python/apache_beam/internal/cloudpickle_pickler.py +++ b/sdks/python/apache_beam/internal/cloudpickle_pickler.py @@ -55,38 +55,13 @@ except (ImportError, ModuleNotFoundError): pass - -def _get_proto_enum_descriptor_class(): - try: - from google.protobuf.internal import api_implementation - except ImportError: - return None - - implementation_type = api_implementation.Type() - - if implementation_type == 'upb': - try: - from google._upb._message import EnumDescriptor - return EnumDescriptor - except ImportError: - pass - elif implementation_type == 'cpp': - try: - from google.protobuf.pyext._message import EnumDescriptor - return EnumDescriptor - except ImportError: - pass - elif implementation_type == 'python': - try: - from google.protobuf.internal.python_message import EnumDescriptor - return EnumDescriptor - except ImportError: - pass - - return None - - -EnumDescriptor = _get_proto_enum_descriptor_class() +try: + from google.protobuf import descriptor_pb2 + MessageDescriptor = type(descriptor_pb2.DescriptorProto.DESCRIPTOR) + EnumDescriptor = type(descriptor_pb2.FieldDescriptorProto.Type.DESCRIPTOR) +except (ImportError, ModuleNotFoundError): + MessageDescriptor = None + EnumDescriptor = None # Pickling, especially unpickling, causes broken module imports on Python 3 # if executed concurrently, see: BEAM-8651, http://bugs.python.org/issue38884. @@ -117,30 +92,22 @@ def _typealias_reduce(obj): return _return_obj, (underlying, ) -def _reconstruct_enum_descriptor(full_name): - for _, module in list(sys.modules.items()): - if not hasattr(module, 'DESCRIPTOR'): - continue +def _reconstruct_message_descriptor(full_name): + from google.protobuf import descriptor_pool + return descriptor_pool.Default().FindMessageTypeByName(full_name) + - if hasattr(module.DESCRIPTOR, 'enum_types_by_name'): - for (_, enum_desc) in module.DESCRIPTOR.enum_types_by_name.items(): - if enum_desc.full_name == full_name: - return enum_desc +def _pickle_message_descriptor(obj): + return _reconstruct_message_descriptor, (obj.full_name, ) - for _, attr_value in vars(module).items(): - if not hasattr(attr_value, 'DESCRIPTOR'): - continue - if hasattr(attr_value.DESCRIPTOR, 'enum_types_by_name'): - for (_, enum_desc) in attr_value.DESCRIPTOR.enum_types_by_name.items(): - if enum_desc.full_name == full_name: - return enum_desc - raise ImportError(f'Could not find enum descriptor: {full_name}') +def _reconstruct_enum_descriptor(full_name): + from google.protobuf import descriptor_pool + return descriptor_pool.Default().FindEnumTypeByName(full_name) def _pickle_enum_descriptor(obj): - full_name = obj.full_name - return _reconstruct_enum_descriptor, (full_name, ) + return _reconstruct_enum_descriptor, (obj.full_name, ) def dumps( @@ -204,6 +171,8 @@ def _dumps( pickler.dispatch_table[LOCK_TYPE] = _lock_reducer except NameError: pass + if MessageDescriptor is not None: + pickler.dispatch_table[MessageDescriptor] = _pickle_message_descriptor if EnumDescriptor is not None: pickler.dispatch_table[EnumDescriptor] = _pickle_enum_descriptor pickler.dump(o) diff --git a/sdks/python/apache_beam/internal/cloudpickle_pickler_test.py b/sdks/python/apache_beam/internal/cloudpickle_pickler_test.py index 99fbb03ac2e4..e270d5a512b1 100644 --- a/sdks/python/apache_beam/internal/cloudpickle_pickler_test.py +++ b/sdks/python/apache_beam/internal/cloudpickle_pickler_test.py @@ -132,6 +132,14 @@ def fn(): self.assertEqual(fn(), loads(dumps(fn))()) + def test_pickle_message_descriptor(self): + desc = proto2_coder_test_messages_pb2.MessageA.DESCRIPTOR + + def fn(): + return desc + + self.assertEqual(fn(), loads(dumps(fn))()) + def test_basics(self): self.assertEqual([1, 'a', ('z', )], loads(dumps([1, 'a', ('z', )]))) fun = lambda x: 'xyz-%s' % x diff --git a/sdks/python/apache_beam/yaml/yaml_transform_test.py b/sdks/python/apache_beam/yaml/yaml_transform_test.py index 4cfe625fc9cd..a86eaf3b4d9c 100644 --- a/sdks/python/apache_beam/yaml/yaml_transform_test.py +++ b/sdks/python/apache_beam/yaml/yaml_transform_test.py @@ -295,10 +295,10 @@ def test_csv_to_json(self): pd.testing.assert_frame_equal(data, result) def test_circular_reference_validation(self): - with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( - pickle_library='cloudpickle')) as p: - # pylint: disable=expression-not-assigned - with self.assertRaisesRegex(ValueError, r'Circular reference detected.*'): + # pylint: disable=expression-not-assigned + with self.assertRaisesRegex(ValueError, r'Circular reference detected.*'): + with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle')) as p: p | YamlTransform( ''' type: composite @@ -318,10 +318,10 @@ def test_circular_reference_validation(self): providers=TEST_PROVIDERS) def test_circular_reference_multi_inputs_validation(self): - with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( - pickle_library='cloudpickle')) as p: - # pylint: disable=expression-not-assigned - with self.assertRaisesRegex(ValueError, r'Circular reference detected.*'): + # pylint: disable=expression-not-assigned + with self.assertRaisesRegex(ValueError, r'Circular reference detected.*'): + with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle')) as p: p | YamlTransform( ''' type: composite @@ -362,10 +362,10 @@ def test_name_is_not_ambiguous(self): assert_that(result, equal_to([0, 1, 9, 16])) def test_name_is_ambiguous(self): - with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( - pickle_library='cloudpickle')) as p: - # pylint: disable=expression-not-assigned - with self.assertRaisesRegex(ValueError, r'Circular reference detected.*'): + # pylint: disable=expression-not-assigned + with self.assertRaisesRegex(ValueError, r'Circular reference detected.*'): + with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle')) as p: p | YamlTransform( ''' type: composite @@ -389,11 +389,11 @@ def test_name_is_ambiguous(self): providers=TEST_PROVIDERS) def test_empty_inputs_throws_error(self): - with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( - pickle_library='cloudpickle')) as p: - with self.assertRaisesRegex(ValueError, - 'Missing inputs for transform at ' - '"EmptyInputOkButYamlDoesntKnow" at line .*'): + with self.assertRaisesRegex(ValueError, + 'Missing inputs for transform at ' + '"EmptyInputOkButYamlDoesntKnow" at line .*'): + with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( + pickle_library='cloudpickle')) as p: _ = p | YamlTransform( ''' type: composite From ee27e338c568cce887c5c841f19c3f142b650386 Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Tue, 22 Sep 2026 19:48:41 +0000 Subject: [PATCH 2/3] make consistent changes for the tests --- sdks/python/apache_beam/yaml/yaml_transform_test.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/sdks/python/apache_beam/yaml/yaml_transform_test.py b/sdks/python/apache_beam/yaml/yaml_transform_test.py index a86eaf3b4d9c..5a1e7b9e2205 100644 --- a/sdks/python/apache_beam/yaml/yaml_transform_test.py +++ b/sdks/python/apache_beam/yaml/yaml_transform_test.py @@ -295,11 +295,10 @@ def test_csv_to_json(self): pd.testing.assert_frame_equal(data, result) def test_circular_reference_validation(self): - # pylint: disable=expression-not-assigned with self.assertRaisesRegex(ValueError, r'Circular reference detected.*'): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: - p | YamlTransform( + _ = p | YamlTransform( ''' type: composite transforms: @@ -318,11 +317,10 @@ def test_circular_reference_validation(self): providers=TEST_PROVIDERS) def test_circular_reference_multi_inputs_validation(self): - # pylint: disable=expression-not-assigned with self.assertRaisesRegex(ValueError, r'Circular reference detected.*'): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: - p | YamlTransform( + _ = p | YamlTransform( ''' type: composite transforms: @@ -362,11 +360,10 @@ def test_name_is_not_ambiguous(self): assert_that(result, equal_to([0, 1, 9, 16])) def test_name_is_ambiguous(self): - # pylint: disable=expression-not-assigned with self.assertRaisesRegex(ValueError, r'Circular reference detected.*'): with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions( pickle_library='cloudpickle')) as p: - p | YamlTransform( + _ = p | YamlTransform( ''' type: composite transforms: From 581b7085f717f95240f7be345967fd06e7cf08ec Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Tue, 22 Sep 2026 20:02:00 +0000 Subject: [PATCH 3/3] remove cloudpickle changes for later --- .../internal/cloudpickle_pickler.py | 69 ++++++++++++++----- .../internal/cloudpickle_pickler_test.py | 8 --- 2 files changed, 50 insertions(+), 27 deletions(-) diff --git a/sdks/python/apache_beam/internal/cloudpickle_pickler.py b/sdks/python/apache_beam/internal/cloudpickle_pickler.py index e6f008c5c768..cea4f01f803c 100644 --- a/sdks/python/apache_beam/internal/cloudpickle_pickler.py +++ b/sdks/python/apache_beam/internal/cloudpickle_pickler.py @@ -55,13 +55,38 @@ except (ImportError, ModuleNotFoundError): pass -try: - from google.protobuf import descriptor_pb2 - MessageDescriptor = type(descriptor_pb2.DescriptorProto.DESCRIPTOR) - EnumDescriptor = type(descriptor_pb2.FieldDescriptorProto.Type.DESCRIPTOR) -except (ImportError, ModuleNotFoundError): - MessageDescriptor = None - EnumDescriptor = None + +def _get_proto_enum_descriptor_class(): + try: + from google.protobuf.internal import api_implementation + except ImportError: + return None + + implementation_type = api_implementation.Type() + + if implementation_type == 'upb': + try: + from google._upb._message import EnumDescriptor + return EnumDescriptor + except ImportError: + pass + elif implementation_type == 'cpp': + try: + from google.protobuf.pyext._message import EnumDescriptor + return EnumDescriptor + except ImportError: + pass + elif implementation_type == 'python': + try: + from google.protobuf.internal.python_message import EnumDescriptor + return EnumDescriptor + except ImportError: + pass + + return None + + +EnumDescriptor = _get_proto_enum_descriptor_class() # Pickling, especially unpickling, causes broken module imports on Python 3 # if executed concurrently, see: BEAM-8651, http://bugs.python.org/issue38884. @@ -92,22 +117,30 @@ def _typealias_reduce(obj): return _return_obj, (underlying, ) -def _reconstruct_message_descriptor(full_name): - from google.protobuf import descriptor_pool - return descriptor_pool.Default().FindMessageTypeByName(full_name) - +def _reconstruct_enum_descriptor(full_name): + for _, module in list(sys.modules.items()): + if not hasattr(module, 'DESCRIPTOR'): + continue -def _pickle_message_descriptor(obj): - return _reconstruct_message_descriptor, (obj.full_name, ) + if hasattr(module.DESCRIPTOR, 'enum_types_by_name'): + for (_, enum_desc) in module.DESCRIPTOR.enum_types_by_name.items(): + if enum_desc.full_name == full_name: + return enum_desc + for _, attr_value in vars(module).items(): + if not hasattr(attr_value, 'DESCRIPTOR'): + continue -def _reconstruct_enum_descriptor(full_name): - from google.protobuf import descriptor_pool - return descriptor_pool.Default().FindEnumTypeByName(full_name) + if hasattr(attr_value.DESCRIPTOR, 'enum_types_by_name'): + for (_, enum_desc) in attr_value.DESCRIPTOR.enum_types_by_name.items(): + if enum_desc.full_name == full_name: + return enum_desc + raise ImportError(f'Could not find enum descriptor: {full_name}') def _pickle_enum_descriptor(obj): - return _reconstruct_enum_descriptor, (obj.full_name, ) + full_name = obj.full_name + return _reconstruct_enum_descriptor, (full_name, ) def dumps( @@ -171,8 +204,6 @@ def _dumps( pickler.dispatch_table[LOCK_TYPE] = _lock_reducer except NameError: pass - if MessageDescriptor is not None: - pickler.dispatch_table[MessageDescriptor] = _pickle_message_descriptor if EnumDescriptor is not None: pickler.dispatch_table[EnumDescriptor] = _pickle_enum_descriptor pickler.dump(o) diff --git a/sdks/python/apache_beam/internal/cloudpickle_pickler_test.py b/sdks/python/apache_beam/internal/cloudpickle_pickler_test.py index e270d5a512b1..99fbb03ac2e4 100644 --- a/sdks/python/apache_beam/internal/cloudpickle_pickler_test.py +++ b/sdks/python/apache_beam/internal/cloudpickle_pickler_test.py @@ -132,14 +132,6 @@ def fn(): self.assertEqual(fn(), loads(dumps(fn))()) - def test_pickle_message_descriptor(self): - desc = proto2_coder_test_messages_pb2.MessageA.DESCRIPTOR - - def fn(): - return desc - - self.assertEqual(fn(), loads(dumps(fn))()) - def test_basics(self): self.assertEqual([1, 'a', ('z', )], loads(dumps([1, 'a', ('z', )]))) fun = lambda x: 'xyz-%s' % x