[YAML] Fix flaky ErrorHandlingTest: don't pickle Scope into user DoFn closures (#40467)

* [YAML] Add test reproducing unpicklable provider state leaking into DoFn closures

A user transform's DoFn that captures `self` also captures the closures
YAML patches onto it, which hold the Scope and every provider. Any
unpicklable provider state (e.g. a protobuf Descriptor) then makes
pickling fail. This test injects such state to reproduce the flaky
ErrorHandlingTest failures deterministically.

* [YAML] Don't pickle Scope (and all providers) into user DoFn closures

Author: shunping

sdks/python/apache_beam/yaml/yaml_transform.py

@@ -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)

sdks/python/apache_beam/yaml/yaml_transform_test.py

@@ -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: