Allow runner to override default pickler.

Author: Claude

sdks/python/apache_beam/pipeline.py

@@ -188,9 +188,6 @@ def __init__(
 
     FileSystems.set_options(self._options)
 
-    pickle_library = self._options.view_as(SetupOptions).pickle_library
-    pickler.set_library(pickle_library)
-
     if runner is None:
       runner = self._options.view_as(StandardOptions).runner
       if runner is None:
@@ -225,6 +222,16 @@ def __init__(
 
     # Default runner to be used.
     self.runner = runner
+
+    if (self._options.view_as(SetupOptions).pickle_library == 'default' and
+        self.runner.default_pickle_library_override() is not None):
+      logging.info(
+          "Default pickling library set to : %s.",
+          runner.default_pickle_library_override())
+      self._options.view_as(
+          SetupOptions).pickle_library = runner.default_pickle_library_override(
+          )
+    pickler.set_library(self._options.view_as(SetupOptions).pickle_library)
     # Stack of transforms generated by nested apply() calls. The stack will
     # contain a root node as an enclosing (parent) node for top transforms.
     self.transforms_stack = [

sdks/python/apache_beam/pipeline_test.py

@@ -62,6 +62,7 @@
 from apache_beam.transforms.window import TimestampedValue
 from apache_beam.utils import windowed_value
 from apache_beam.utils.timestamp import MIN_TIMESTAMP
+from apache_beam.runners.direct import direct_runner
 
 
 class FakeUnboundedSource(SourceBase):
@@ -156,6 +157,20 @@ def test_create(self):
       pcoll3 = pcoll2 | 'do' >> FlatMap(lambda x: [x + 10])
       assert_that(pcoll3, equal_to([14, 15, 16]), label='pcoll3')
 
+  @mock.patch('logging.info')
+  def test_runner_overrides_default_pickler(self, mock_info):
+    with mock.patch.object(direct_runner.SwitchingDirectRunner,
+                           'default_pickle_library_override') as mock_fn:
+      mock_fn.return_value = 'dill'
+      with TestPipeline() as pipeline:
+        pcoll = pipeline | 'label1' >> Create([1, 2, 3])
+        assert_that(pcoll, equal_to([1, 2, 3]))
+
+        from apache_beam.internal import pickler
+        from apache_beam.internal import dill_pickler
+        self.assertIs(pickler.desired_pickle_lib, dill_pickler)
+    mock_info.assert_any_call('Default pickling library set to : %s.', 'dill')
+
   def test_flatmap_builtin(self):
     with TestPipeline() as pipeline:
       pcoll = pipeline | 'label1' >> Create([1, 2, 3])
@@ -279,7 +294,7 @@ def test_no_wait_until_finish(self, mock_info):
     with Pipeline(runner='DirectRunner',
                   options=PipelineOptions(["--no_wait_until_finish"])) as p:
       _ = p | beam.Create(['test'])
-    mock_info.assert_called_once_with(
+    mock_info.assert_any_call(
         'Job execution continues without waiting for completion. '
         'Use "wait_until_finish" in PipelineResult to block until finished.')
     p.result.wait_until_finish()

sdks/python/apache_beam/runners/direct/direct_runner.py

@@ -66,6 +66,10 @@ class SwitchingDirectRunner(PipelineRunner):
   which supports streaming execution and certain primitives not yet
   implemented in the FnApiRunner.
   """
+  def default_pickle_library_override(self):
+    """Default pickle library, can be overridden by runner implementation."""
+    return 'cloudpickle'
+
   def is_fnapi_compatible(self):
     return BundleBasedDirectRunner.is_fnapi_compatible()
 

sdks/python/apache_beam/runners/runner.py

@@ -225,6 +225,10 @@ def check_requirements(
               beam_runner_api_pb2.TimeDomain.PROCESSING_TIME):
             raise NotImplementedError(timer.time_domain)
 
+  def default_pickle_library_override(self):
+    """Default pickle library, can be overridden by runner implementation."""
+    return None
+
 
 # FIXME: replace with PipelineState(str, enum.Enum)
 class PipelineState(object):