Move options override to before validation.

Author: Claude

sdks/python/apache_beam/pipeline.py

@@ -207,6 +207,17 @@ def __init__(
           'Runner %s is not a PipelineRunner object or the '
           'name of a registered runner.' % runner)
 
+    # Runner can oerride the default runner to be used.
+    if (self._options.view_as(SetupOptions).pickle_library == 'default' and
+        runner.default_pickle_library_override() is not None):
+      logging.info(
+          "Runner defaulting to pickling library: %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)
+
     # Validate pipeline options
     errors = PipelineOptionsValidator(self._options, runner).validate()
     if errors:
@@ -226,15 +237,6 @@ 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

@@ -41,6 +41,7 @@
 from apache_beam.portability.api import beam_runner_api_pb2
 from apache_beam.pvalue import AsSingleton
 from apache_beam.pvalue import TaggedOutput
+from apache_beam.runners.runner import PipelineRunner
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to
@@ -62,7 +63,6 @@
 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):
@@ -159,7 +159,7 @@ def test_create(self):
 
   @mock.patch('logging.info')
   def test_runner_overrides_default_pickler(self, mock_info):
-    with mock.patch.object(direct_runner.SwitchingDirectRunner,
+    with mock.patch.object(PipelineRunner,
                            'default_pickle_library_override') as mock_fn:
       mock_fn.return_value = 'dill'
       with TestPipeline() as pipeline:
@@ -169,7 +169,8 @@ def test_runner_overrides_default_pickler(self, mock_info):
         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')
+    mock_info.assert_any_call(
+        'Runner defaulting to pickling library: %s.', 'dill')
 
   def test_flatmap_builtin(self):
     with TestPipeline() as pipeline:

sdks/python/apache_beam/runners/dataflow/dataflow_runner.py

@@ -98,6 +98,9 @@ class DataflowRunner(PipelineRunner):
   def __init__(self, cache=None):
     self._default_environment = None
 
+  def default_pickle_library_override(self):
+    return 'cloudpickle'
+
   def is_fnapi_compatible(self):
     return False
 

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

@@ -66,10 +66,6 @@ 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()