Add flag for disabling dill check in coders. (#36453)

* init.

* Trigger test.

* Remove pipeline option.

* BQ file loads test.

* Fixes.

* Fix test.

Author: claudevdm

.github/trigger_files/beam_PreCommit_Python_Dill.json

@@ -0,0 +1,4 @@
+{
+    "comment": "Modify this file in a trivial way to cause this test suite to run",
+    "revision": 1
+}

.github/workflows/beam_PreCommit_Python_Dill.yml

@@ -106,9 +106,9 @@ jobs:
           arguments: |
             -Pposargs="${{
               contains(matrix.os, 'self-hosted') &&
-              'apache_beam/internal/ apache_beam/ml/ apache_beam/transforms/ apache_beam/typehints/ apache_beam/runners/portability/ -m (uses_dill and not require_docker_in_docker)' ||
-                'apache_beam/internal/ apache_beam/ml/ apache_beam/transforms/ apache_beam/typehints/ apache_beam/runners/portability/ -m (uses_dill and require_docker_in_docker)'
-            }}"  \
+              'apache_beam/internal/ apache_beam/io/gcp/ apache_beam/options/ apache_beam/transforms/ apache_beam/typehints/ apache_beam/runners/portability/ -m (uses_dill and not require_docker_in_docker)' ||
+                'apache_beam/internal/ apache_beam/io/gcp/ apache_beam/options/ apache_beam/transforms/ apache_beam/typehints/ apache_beam/runners/portability/ -m (uses_dill and require_docker_in_docker)'
+            }}" \
             -PpythonVersion=${{ matrix.python_version }}
       - name: Archive Python Test Results
         uses: actions/upload-artifact@v4

sdks/python/apache_beam/coders/coders.py

@@ -1004,8 +1004,16 @@ def to_type_hint(self):
     return Any
 
 
-def _should_force_use_dill(update_compat_version):
+def _should_force_use_dill(registry):
+  # force_dill_deterministic_coders is for testing purposes. If there is a
+  # DeterministicFastPrimitivesCoder in the pipeline graph but the dill
+  # encoding path is not really triggered dill does not have to be installed.
+  # and this check can be skipped.
+  if getattr(registry, 'force_dill_deterministic_coders', False):
+    return True
+
   from apache_beam.transforms.util import is_v1_prior_to_v2
+  update_compat_version = registry.update_compatibility_version
   if not update_compat_version:
     return False
 
@@ -1035,11 +1043,11 @@ def _update_compatible_deterministic_fast_primitives_coder(coder, step_label):
    relative filepaths in code objects and dynamic functions.
   """
   from apache_beam.coders import typecoders
-  update_compat_version = typecoders.registry.update_compatibility_version
-  if _should_force_use_dill(update_compat_version):
+
+  if _should_force_use_dill(typecoders.registry):
     return DeterministicFastPrimitivesCoder(coder, step_label)
   return DeterministicFastPrimitivesCoderV2(
-      coder, step_label, update_compat_version)
+      coder, step_label, typecoders.registry.update_compatibility_version)
 
 
 class FastPrimitivesCoder(FastCoder):

sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py

@@ -63,11 +63,6 @@
 except ImportError:
   raise unittest.SkipTest('GCP dependencies are not installed')
 
-try:
-  import dill
-except ImportError:
-  dill = None
-
 _LOGGER = logging.getLogger(__name__)
 
 _DESTINATION_ELEMENT_PAIRS = [
@@ -411,13 +406,6 @@ def test_partition_files_dofn_size_split(self):
           label='CheckSinglePartition')
 
 
-def maybe_skip(compat_version):
-  if compat_version and not dill:
-    raise unittest.SkipTest(
-        'Dill dependency not installed which is required for compat_version'
-        ' <= 2.67.0')
-
-
 class TestBigQueryFileLoads(_TestCaseWithTempDirCleanUp):
   def test_trigger_load_jobs_with_empty_files(self):
     destination = "project:dataset.table"
@@ -497,9 +485,9 @@ def test_records_traverse_transform_with_mocks(self):
       param(compat_version=None),
       param(compat_version="2.64.0"),
   ])
-  @pytest.mark.uses_dill
   def test_reshuffle_before_load(self, compat_version):
-    maybe_skip(compat_version)
+    from apache_beam.coders import typecoders
+    typecoders.registry.force_dill_deterministic_coders = True
     destination = 'project1:dataset1.table1'
 
     job_reference = bigquery_api.JobReference()
@@ -525,13 +513,17 @@ def test_reshuffle_before_load(self, compat_version):
         validate=False,
         temp_file_format=bigquery_tools.FileFormat.JSON)
 
-    options = PipelineOptions(update_compatibility_version=compat_version)
+    options = PipelineOptions(
+        update_compatibility_version=compat_version,
+        # Disable unrelated compatibility change.
+        force_cloudpickle_deterministic_coders=True)
     # Need to test this with the DirectRunner to avoid serializing mocks
     with TestPipeline('DirectRunner', options=options) as p:
       _ = p | beam.Create(_ELEMENTS) | transform
 
     reshuffle_before_load = compat_version is None
     assert transform.reshuffle_before_load == reshuffle_before_load
+    typecoders.registry.force_dill_deterministic_coders = False
 
   def test_load_job_id_used(self):
     job_reference = bigquery_api.JobReference()
@@ -1008,7 +1000,9 @@ def dynamic_destination_resolver(element, *side_inputs):
   ])
   def test_triggering_frequency(
       self, is_streaming, with_auto_sharding, compat_version):
-    maybe_skip(compat_version)
+    from apache_beam.coders import typecoders
+    typecoders.registry.force_dill_deterministic_coders = True
+
     destination = 'project1:dataset1.table1'
 
     job_reference = bigquery_api.JobReference()
@@ -1114,6 +1108,8 @@ def __call__(self):
           label='CheckDestinations')
       assert_that(jobs, equal_to(expected_jobs), label='CheckJobs')
 
+    typecoders.registry.force_dill_deterministic_coders = False
+
 
 class BigQueryFileLoadsIT(unittest.TestCase):
 

sdks/python/apache_beam/options/pipeline_options.py

@@ -874,6 +874,18 @@ def _add_argparse_args(cls, parser):
         'their condition met. Some operations, such as GroupByKey, disallow '
         'this. This exists for cases where such loss is acceptable and for '
         'backwards compatibility. See BEAM-9487.')
+    parser.add_argument(
+        '--force_cloudpickle_deterministic_coders',
+        default=False,
+        action='store_true',
+        help=(
+            'Force the use of cloudpickle-based deterministic coders '
+            'instead of dill-based coders, even when '
+            'update_compatibility_version  would normally trigger dill usage '
+            'for backward compatibility. This flag overrides automatic coder '
+            'selection to always use the modern cloudpickle serialization '
+            ' path. Warning: May break pipeline update compatibility with '
+            ' SDK versions prior to 2.68.0.'))
 
   def validate(self, unused_validator):
     errors = []

sdks/python/apache_beam/transforms/util_test.py

@@ -92,11 +92,6 @@
 from apache_beam.utils.windowed_value import PaneInfoTiming
 from apache_beam.utils.windowed_value import WindowedValue
 
-try:
-  import dill
-except ImportError:
-  dill = None
-
 try:
   from google.cloud import secretmanager
 except ImportError:
@@ -131,13 +126,6 @@ def is_deterministic(self):
     return True
 
 
-def maybe_skip(compat_version):
-  if compat_version and not dill:
-    raise unittest.SkipTest(
-        'Dill dependency not installed which is required for compat_version'
-        ' <= 2.67.0')
-
-
 class CoGroupByKeyTest(unittest.TestCase):
   def test_co_group_by_key_on_tuple(self):
     with TestPipeline() as pipeline:
@@ -1219,10 +1207,10 @@ def test_reshuffle_streaming_global_window_with_buckets(self):
       param(compat_version=None),
       param(compat_version="2.64.0"),
   ])
-  @pytest.mark.uses_dill
   def test_reshuffle_custom_window_preserves_metadata(self, compat_version):
     """Tests that Reshuffle preserves pane info."""
-    maybe_skip(compat_version)
+    from apache_beam.coders import typecoders
+    typecoders.registry.force_dill_deterministic_coders = True
     element_count = 12
     timestamp_value = timestamp.Timestamp(0)
     l = [
@@ -1286,7 +1274,6 @@ def test_reshuffle_custom_window_preserves_metadata(self, compat_version):
                           expected_timestamp, [GlobalWindow()],
                           PANE_INFO_UNKNOWN)
     ])
-
     options = PipelineOptions(update_compatibility_version=compat_version)
     options.view_as(StandardOptions).streaming = True
 
@@ -1317,16 +1304,17 @@ def test_reshuffle_custom_window_preserves_metadata(self, compat_version):
           equal_to(expected),
           label='CheckMetadataPreserved',
           reify_windows=True)
+    typecoders.registry.force_dill_deterministic_coders = False
 
   @parameterized.expand([
       param(compat_version=None),
       param(compat_version="2.64.0"),
   ])
-  @pytest.mark.uses_dill
   def test_reshuffle_default_window_preserves_metadata(self, compat_version):
     """Tests that Reshuffle preserves timestamp, window, and pane info
     metadata."""
-    maybe_skip(compat_version)
+    from apache_beam.coders import typecoders
+    typecoders.registry.force_dill_deterministic_coders = True
     no_firing = PaneInfo(
         is_first=True,
         is_last=True,
@@ -1400,6 +1388,7 @@ def test_reshuffle_default_window_preserves_metadata(self, compat_version):
           equal_to(expected),
           label='CheckMetadataPreserved',
           reify_windows=True)
+    typecoders.registry.force_dill_deterministic_coders = False
 
   @pytest.mark.it_validatesrunner
   def test_reshuffle_preserves_timestamps(self):

sdks/python/tox.ini

@@ -590,4 +590,4 @@ commands =
   # Log dill version for debugging
   /bin/sh -c "pip freeze | grep -E dill"
   # Run all dill-specific tests
-  /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 1 -m uses_dill {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'
+  /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 1 -m uses_dill -vv -ra {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'