Make dill optional

Author: claudevdm

plugins/beam-code-completion-plugin/.run/Run IDE with Plugin.run.xml

@@ -40,4 +40,4 @@
         <DebugAllEnabled>false</DebugAllEnabled>
         <method v="2" />
     </configuration>
-</component>
+</component>

sdks/python/apache_beam/coders/coders.py

@@ -85,9 +85,7 @@
   # occurs.
   from apache_beam.internal.dill_pickler import dill
 except ImportError:
-  # We fall back to using the stock dill library in tests that don't use the
-  # full Python SDK.
-  import dill
+  dill = None  # type: ignore
 
 __all__ = [
     'Coder',
@@ -900,6 +898,12 @@ def to_type_hint(self):
 
 class DillCoder(_PickleCoderBase):
   """Coder using dill's pickle functionality."""
+  def __init__(self):
+    if not dill:
+      raise RuntimeError("This pipeline contains a DillCoder which requires " \
+      "the dill package. Install the dill package with the dill extra " \
+        "e.g. apache-beam[dill]")
+
   def _create_impl(self):
     return coder_impl.CallbackCoderImpl(maybe_dill_dumps, maybe_dill_loads)
 

sdks/python/apache_beam/coders/coders_test_common.py

@@ -59,6 +59,11 @@
 except ImportError:
   dataclasses = None  # type: ignore
 
+try:
+  import dill
+except ImportError:
+  dill = None
+
 MyNamedTuple = collections.namedtuple('A', ['x', 'y'])  # type: ignore[name-match]
 AnotherNamedTuple = collections.namedtuple('AnotherNamedTuple', ['x', 'y'])
 MyTypedNamedTuple = NamedTuple('MyTypedNamedTuple', [('f1', int), ('f2', str)])
@@ -116,6 +121,7 @@ class UnFrozenDataClass:
 # These tests need to all be run in the same process due to the asserts
 # in tearDownClass.
 @pytest.mark.no_xdist
+@pytest.mark.uses_dill
 class CodersTest(unittest.TestCase):
 
   # These class methods ensure that we test each defined coder in both
@@ -173,6 +179,9 @@ def tearDownClass(cls):
         coders.BigIntegerCoder,  # tested in DecimalCoder
         coders.TimestampPrefixingOpaqueWindowCoder,
     ])
+    if not dill:
+      standard -= set(
+          [coders.DillCoder, coders.DeterministicFastPrimitivesCoder])
     cls.seen_nested -= set(
         [coders.ProtoCoder, coders.ProtoPlusCoder, CustomCoder])
     assert not standard - cls.seen, str(standard - cls.seen)
@@ -241,8 +250,13 @@ def test_memoizing_pickle_coder(self):
       param(compat_version="2.67.0"),
   ])
   def test_deterministic_coder(self, compat_version):
+
     typecoders.registry.update_compatibility_version = compat_version
     coder = coders.FastPrimitivesCoder()
+    if not dill and compat_version:
+      with self.assertRaises(RuntimeError):
+        coder.as_deterministic_coder(step_label="step")
+      self.skipTest('Dill not installed')
     deterministic_coder = coder.as_deterministic_coder(step_label="step")
 
     self.check_coder(deterministic_coder, *self.test_values_deterministic)
@@ -321,6 +335,11 @@ def test_deterministic_map_coder_is_update_compatible(self, compat_version):
     coder = coders.MapCoder(
         coders.FastPrimitivesCoder(), coders.FastPrimitivesCoder())
 
+    if not dill and compat_version:
+      with self.assertRaises(RuntimeError):
+        coder.as_deterministic_coder(step_label="step")
+      self.skipTest('Dill not installed')
+
     deterministic_coder = coder.as_deterministic_coder(step_label="step")
 
     assert isinstance(
@@ -331,6 +350,11 @@ def test_deterministic_map_coder_is_update_compatible(self, compat_version):
     self.check_coder(deterministic_coder, *values)
 
   def test_dill_coder(self):
+    if not dill:
+      with self.assertRaises(RuntimeError):
+        coders.DillCoder()
+      self.skipTest('Dill not installed')
+
     cell_value = (lambda x: lambda: x)(0).__closure__[0]
     self.check_coder(coders.DillCoder(), 'a', 1, cell_value)
     self.check_coder(
@@ -661,6 +685,8 @@ def test_param_windowed_value_coder(self):
   def test_cross_process_encoding_of_special_types_is_deterministic(
       self, compat_version):
     """Test cross-process determinism for all special deterministic types"""
+    if compat_version:
+      pytest.importorskip("dill")
 
     if sys.executable is None:
       self.skipTest('No Python interpreter found')

sdks/python/apache_beam/internal/pickler.py

@@ -29,7 +29,11 @@
 """
 
 from apache_beam.internal import cloudpickle_pickler
-from apache_beam.internal import dill_pickler
+
+try:
+  from apache_beam.internal import dill_pickler
+except ImportError:
+  dill_pickler = None
 
 USE_CLOUDPICKLE = 'cloudpickle'
 USE_DILL = 'dill'
@@ -74,6 +78,12 @@ def load_session(file_path):
 def set_library(selected_library=DEFAULT_PICKLE_LIB):
   """ Sets pickle library that will be used. """
   global desired_pickle_lib
+
+  if selected_library == USE_DILL and not dill_pickler:
+    if not dill_pickler:
+      raise ImportError(
+          "Dill is not installed. To use dill pickler, install apache-beam with"
+          " the dill extras package e.g. apache-beam[dill].")
   # If switching to or from dill, update the pickler hook overrides.
   if (selected_library == USE_DILL) != (desired_pickle_lib == dill_pickler):
     dill_pickler.override_pickler_hooks(selected_library == USE_DILL)

sdks/python/apache_beam/internal/pickler_test.py

@@ -19,6 +19,7 @@
 
 # pytype: skip-file
 
+import pytest
 import random
 import sys
 import threading
@@ -34,6 +35,12 @@
 from apache_beam.internal.pickler import loads
 
 
+def maybe_skip_if_no_dill(pickle_library):
+  if pickle_library == 'dill':
+    pytest.importorskip("dill")
+
+
+@pytest.mark.uses_dill
 class PicklerTest(unittest.TestCase):
 
   NO_MAPPINGPROXYTYPE = not hasattr(types, "MappingProxyType")
@@ -43,6 +50,7 @@ class PicklerTest(unittest.TestCase):
       param(pickle_lib='cloudpickle'),
   ])
   def test_basics(self, pickle_lib):
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
 
     self.assertEqual([1, 'a', ('z', )], loads(dumps([1, 'a', ('z', )])))
@@ -55,6 +63,7 @@ def test_basics(self, pickle_lib):
   ])
   def test_lambda_with_globals(self, pickle_lib):
     """Tests that the globals of a function are preserved."""
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
 
     # The point of the test is that the lambda being called after unpickling
@@ -68,6 +77,7 @@ def test_lambda_with_globals(self, pickle_lib):
       param(pickle_lib='cloudpickle'),
   ])
   def test_lambda_with_main_globals(self, pickle_lib):
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
     self.assertEqual(unittest, loads(dumps(lambda: unittest))())
 
@@ -77,6 +87,7 @@ def test_lambda_with_main_globals(self, pickle_lib):
   ])
   def test_lambda_with_closure(self, pickle_lib):
     """Tests that the closure of a function is preserved."""
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
     self.assertEqual(
         'closure: abc',
@@ -88,6 +99,7 @@ def test_lambda_with_closure(self, pickle_lib):
   ])
   def test_class(self, pickle_lib):
     """Tests that a class object is pickled correctly."""
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
     self.assertEqual(['abc', 'def'],
                      loads(dumps(module_test.Xyz))().foo('abc def'))
@@ -98,6 +110,7 @@ def test_class(self, pickle_lib):
   ])
   def test_object(self, pickle_lib):
     """Tests that a class instance is pickled correctly."""
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
     self.assertEqual(['abc', 'def'],
                      loads(dumps(module_test.XYZ_OBJECT)).foo('abc def'))
@@ -108,6 +121,7 @@ def test_object(self, pickle_lib):
   ])
   def test_nested_class(self, pickle_lib):
     """Tests that a nested class object is pickled correctly."""
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
     self.assertEqual(
         'X:abc', loads(dumps(module_test.TopClass.NestedClass('abc'))).datum)
@@ -121,6 +135,7 @@ def test_nested_class(self, pickle_lib):
   ])
   def test_dynamic_class(self, pickle_lib):
     """Tests that a nested class object is pickled correctly."""
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
     self.assertEqual(
         'Z:abc', loads(dumps(module_test.create_class('abc'))).get())
@@ -130,6 +145,7 @@ def test_dynamic_class(self, pickle_lib):
       param(pickle_lib='cloudpickle'),
   ])
   def test_generators(self, pickle_lib):
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
     with self.assertRaises(TypeError):
       dumps((_ for _ in range(10)))
@@ -139,6 +155,7 @@ def test_generators(self, pickle_lib):
       param(pickle_lib='cloudpickle'),
   ])
   def test_recursive_class(self, pickle_lib):
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
     self.assertEqual(
         'RecursiveClass:abc',
@@ -149,6 +166,7 @@ def test_recursive_class(self, pickle_lib):
       param(pickle_lib='cloudpickle'),
   ])
   def test_pickle_rlock(self, pickle_lib):
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
     rlock_instance = threading.RLock()
     rlock_type = type(rlock_instance)
@@ -160,6 +178,7 @@ def test_pickle_rlock(self, pickle_lib):
       param(pickle_lib='cloudpickle'),
   ])
   def test_save_paths(self, pickle_lib):
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
     f = loads(dumps(lambda x: x))
     co_filename = f.__code__.co_filename
@@ -171,6 +190,7 @@ def test_save_paths(self, pickle_lib):
       param(pickle_lib='cloudpickle'),
   ])
   def test_dump_and_load_mapping_proxy(self, pickle_lib):
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
     self.assertEqual(
         'def', loads(dumps(types.MappingProxyType({'abc': 'def'})))['abc'])
@@ -184,6 +204,7 @@ def test_dump_and_load_mapping_proxy(self, pickle_lib):
       param(pickle_lib='cloudpickle'),
   ])
   def test_dataclass(self, pickle_lib):
+    maybe_skip_if_no_dill(pickle_lib)
     exec(
         '''
 from apache_beam.internal.module_test import DataClass
@@ -195,6 +216,7 @@ def test_dataclass(self, pickle_lib):
       param(pickle_lib='cloudpickle'),
   ])
   def test_class_states_not_changed_at_subsequent_loading(self, pickle_lib):
+    maybe_skip_if_no_dill(pickle_lib)
     pickler.set_library(pickle_lib)
 
     class Local:
@@ -255,6 +277,7 @@ def maybe_get_sets_with_different_iteration_orders(self):
     return set1, set2
 
   def test_best_effort_determinism(self):
+    maybe_skip_if_no_dill('dill')
     pickler.set_library('dill')
     set1, set2 = self.maybe_get_sets_with_different_iteration_orders()
     self.assertEqual(
@@ -267,6 +290,7 @@ def test_best_effort_determinism(self):
       self.skipTest('Set iteration orders matched. Test results inconclusive.')
 
   def test_disable_best_effort_determinism(self):
+    maybe_skip_if_no_dill('dill')
     pickler.set_library('dill')
     set1, set2 = self.maybe_get_sets_with_different_iteration_orders()
     # The test relies on the sets having different iteration orders for the

sdks/python/apache_beam/ml/anomaly/specifiable_test.py

@@ -19,6 +19,7 @@
 import dataclasses
 import logging
 import os
+import pytest
 import unittest
 from typing import Optional
 
@@ -323,7 +324,10 @@ def __init__(self, arg):
       self.my_arg = arg * 10
       type(self).counter += 1
 
-  def test_on_pickle(self):
+  @pytest.mark.uses_dill
+  def test_on_dill_pickle(self):
+    pytest.importorskip("dill")
+
     FooForPickle = TestInitCallCount.FooForPickle
 
     import dill
@@ -339,6 +343,9 @@ def test_on_pickle(self):
     self.assertEqual(FooForPickle.counter, 1)
     self.assertEqual(new_foo_2.__dict__, foo.__dict__)
 
+  def test_on_pickle(self):
+    FooForPickle = TestInitCallCount.FooForPickle
+
     # Note that pickle does not support classes/functions nested in a function.
     import pickle
     FooForPickle.counter = 0

sdks/python/apache_beam/options/pipeline_options.py

@@ -1619,7 +1619,7 @@ def _add_argparse_args(cls, parser):
         help=(
             'Chooses which pickle library to use. Options are dill, '
             'cloudpickle or default.'),
-        choices=['cloudpickle', 'default', 'dill'])
+        choices=['cloudpickle', 'default', 'dill', 'dill_unsafe'])
     parser.add_argument(
         '--save_main_session',
         default=False,
@@ -1701,6 +1701,7 @@ def _add_argparse_args(cls, parser):
   def validate(self, validator):
     errors = []
     errors.extend(validator.validate_container_prebuilding_options(self))
+    errors.extend(validator.validate_pickle_library(self))
     return errors
 
 

sdks/python/apache_beam/options/pipeline_options_validator.py

@@ -119,6 +119,15 @@ class PipelineOptionsValidator(object):
   ERR_REPEATABLE_OPTIONS_NOT_SET_AS_LIST = (
       '(%s) is a string. Programmatically set PipelineOptions like (%s) '
       'options need to be specified as a list.')
+  ERR_DILL_NOT_INSTALLED = (
+      'Option pickle_library=dill requires dill==0.3.1.1. Install apache-beam '
+      'with the dill extra e.g. apache-beam[gcp, dill]. Dill package was not '
+      'found')
+  ERR_UNSAFE_DILL_VERSION = (
+      'Dill version 0.3.1.1 is required when using pickle_library=dill. Other'
+      'versions of dill are untested with Apache Beam. To install the supported'
+      'dill version instal apache-beam[dill] extra. To use an unsupported dill '
+      'version, use pickle_library=dill_unsafe. %s')
 
   # GCS path specific patterns.
   GCS_URI = '(?P<SCHEME>[^:]+)://(?P<BUCKET>[^/]+)(/(?P<OBJECT>.*))?'
@@ -196,6 +205,25 @@ def validate_gcs_path(self, view, arg_name):
       return self._validate_error(self.ERR_INVALID_GCS_OBJECT, arg, arg_name)
     return []
 
+  def validate_pickle_library(self, view):
+    """Validates a GCS path against gs://bucket/object URI format."""
+    if view.pickle_library == 'default' or view.pickle_library == 'cloudpickle':
+      return []
+
+    if view.pickle_library == 'dill_unsafe':
+      return []
+
+    if view.pickle_library == 'dill':
+      try:
+        import dill
+        if dill.__version__ != "0.3.1.1":
+          return self._validate_error(
+              self.ERR_UNSAFE_DILL_VERSION,
+              f"Dill version found {dill.__version__}")
+      except ImportError:
+        return self._validate_error(self.ERR_DILL_NOT_INSTALLED)
+    return []
+
   def validate_cloud_options(self, view):
     """Validates job_name and project arguments."""
     errors = []