Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions sdks/python/apache_beam/yaml/yaml_transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -192,6 +192,11 @@ def __init__(
self.input_providers = input_providers
self._all_followers = None

def __reduce__(self):
# Scope is construction-only, but closures attached in create_ptransform
# can pull it (and all providers) into user DoFn pickles. Stub it out.
return str, ('Pickled YAML scope stub.', )

def followers(self, transform_name):
if self._all_followers is None:
self._all_followers = collections.defaultdict(list)
Expand Down
40 changes: 40 additions & 0 deletions sdks/python/apache_beam/yaml/yaml_transform_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -834,7 +834,47 @@ def test_output_schema_with_main_transform_error_handling_success(self):
''')


class _UnpicklableFactory:
"""A transform factory holding an object that cloudpickle cannot serialize.

Used to simulate provider state (e.g. lazily populated caches in the
process-global standard providers) that should never be pulled into the
pickled closure of an unrelated user transform.
"""
def __init__(self):
from apache_beam.portability.api import beam_runner_api_pb2
self._descriptor = beam_runner_api_pb2.Pipeline.DESCRIPTOR

def __call__(self):
return beam.Map(lambda x: x)


class ErrorHandlingTest(unittest.TestCase):
def test_closure_does_not_capture_unrelated_providers(self):
# SizeLimiter's DoFn closes over `self`. Pickling that DoFn must not drag in
# the YAML Scope and every provider it knows about.
with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
pickle_library='cloudpickle')) as p:
result = p | YamlTransform(
'''
type: composite
transforms:
- type: Create
config:
elements: ['a', 'b', 'biiiiig']
- type: SizeLimiter
input: Create
config:
limit: 5
error_handling:
output: errors
output:
good: SizeLimiter
bad: SizeLimiter.errors
''',
providers=dict(TEST_PROVIDERS, Unpicklable=_UnpicklableFactory()))
assert_that(result['good'], equal_to(['a', 'b']), label="CheckGood")

def test_error_handling_outputs(self):
with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
pickle_library='cloudpickle')) as p:
Expand Down
Loading