[BEAM-8123] Add cloudpickle as optional library

CHANGES.md

@@ -33,6 +33,7 @@
 
 * X feature added (Java/Python) ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
 * Remote packages can now be downloaded from locations supported by apache_beam.io.filesystems. The files will be downloaded on Stager and uploaded to staging location. For more information, see [BEAM-11275](https://issues.apache.org/jira/browse/BEAM-11275)
+* Added support for cloudpickle as a pickling library for Python SDK. To use cloudpickle set pipeline options, --pickler_lib=cloudpickle.
 
 ## Breaking Changes
 

sdks/python/apache_beam/coders/coders.py

@@ -80,7 +80,7 @@
 try:
   # Import dill from the pickler module to make sure our monkey-patching of dill
   # occurs.
-  from apache_beam.internal.pickler import dill
+  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.

sdks/python/apache_beam/coders/coders_test_common.py

@@ -679,17 +679,16 @@ def iterable_state_read(token, element_coder_impl):
         read_state=iterable_state_read,
         write_state=iterable_state_write,
         write_state_threshold=1)
-    context = pipeline_context.PipelineContext(
-        iterable_state_read=iterable_state_read,
-        iterable_state_write=iterable_state_write)
-    self.check_coder(
-        coder, [1, 2, 3], context=context, test_size_estimation=False)
+    # Note: do not use check_coder
+    # see https://github.com/cloudpipe/cloudpickle/issues/452
+    self._observe(coder)
+    self.assertEqual([1, 2, 3], coder.decode(coder.encode([1, 2, 3])))
     # Ensure that state was actually used.
     self.assertNotEqual(state, {})
-    self.check_coder(
-        coders.TupleCoder((coder, coder)), ([1], [2, 3]),
-        context=context,
-        test_size_estimation=False)
+    tupleCoder = coders.TupleCoder((coder, coder))
+    self._observe(tupleCoder)
+    self.assertEqual(([1], [2, 3]),
+                     tupleCoder.decode(tupleCoder.encode(([1], [2, 3]))))
 
   def test_nullable_coder(self):
     self.check_coder(coders.NullableCoder(coders.VarIntCoder()), None, 2 * 64)

sdks/python/apache_beam/internal/cloudpickle_pickler.py

@@ -0,0 +1,111 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Pickler for values, functions, and classes.
+
+For internal use only. No backwards compatibility guarantees.
+
+Uses the cloudpickle library to pickle data, functions, lambdas
+and classes.
+
+dump_session and load_session are no-ops.
+"""
+
+# pytype: skip-file
+
+import base64
+import bz2
+import io
+import threading
+import zlib
+
+import cloudpickle
+
+try:
+  from absl import flags
+except (ImportError, ModuleNotFoundError):
+  pass
+
+# Pickling, especially unpickling, causes broken module imports on Python 3
+# if executed concurrently, see: BEAM-8651, http://bugs.python.org/issue38884.
+_pickle_lock = threading.RLock()
+
+
+def dumps(o, enable_trace=True, use_zlib=False):
+  # type: (...) -> bytes
+
+  """For internal use only; no backwards-compatibility guarantees."""
+  with _pickle_lock:
+    with io.BytesIO() as file:
+      pickler = cloudpickle.CloudPickler(file)
+      try:
+        pickler.dispatch_table[type(flags.FLAGS)] = _pickle_absl_flags
+      except NameError:
+        pass
+      pickler.dump(o)
+      s = file.getvalue()
+
+  # Compress as compactly as possible (compresslevel=9) to decrease peak memory
+  # usage (of multiple in-memory copies) and to avoid hitting protocol buffer
+  # limits.
+  # WARNING: Be cautious about compressor change since it can lead to pipeline
+  # representation change, and can break streaming job update compatibility on
+  # runners such as Dataflow.
+  if use_zlib:
+    c = zlib.compress(s, 9)
+  else:
+    c = bz2.compress(s, compresslevel=9)
+  del s  # Free up some possibly large and no-longer-needed memory.
+
+  return base64.b64encode(c)
+
+
+def loads(encoded, enable_trace=True, use_zlib=False):
+  """For internal use only; no backwards-compatibility guarantees."""
+
+  c = base64.b64decode(encoded)
+
+  if use_zlib:
+    s = zlib.decompress(c)
+  else:
+    s = bz2.decompress(c)
+
+  del c  # Free up some possibly large and no-longer-needed memory.
+
+  with _pickle_lock:
+    unpickled = cloudpickle.loads(s)
+    return unpickled
+
+
+def _pickle_absl_flags(obj):
+  return _create_absl_flags, tuple([])
+
+
+def _create_absl_flags():
+  return flags.FLAGS
+
+
+def dump_session(file_path):
+  # It is possible to dump session with cloudpickle. However, since references
+  # are saved it should not be necessary. See https://s.apache.org/beam-picklers
+  pass
+
+
+def load_session(file_path):
+  # It is possible to load_session with cloudpickle. However, since references
+  # are saved it should not be necessary. See https://s.apache.org/beam-picklers
+  pass

sdks/python/apache_beam/internal/cloudpickle_pickler_test.py

@@ -0,0 +1,112 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Unit tests for the cloudpickle_pickler module."""
+
+# pytype: skip-file
+
+import sys
+import types
+import unittest
+
+from apache_beam.internal import module_test
+from apache_beam.internal.cloudpickle_pickler import dumps
+from apache_beam.internal.cloudpickle_pickler import loads
+
+
+class PicklerTest(unittest.TestCase):
+
+  NO_MAPPINGPROXYTYPE = not hasattr(types, "MappingProxyType")
+
+  def test_basics(self):
+    self.assertEqual([1, 'a', (u'z', )], loads(dumps([1, 'a', (u'z', )])))
+    fun = lambda x: 'xyz-%s' % x
+    self.assertEqual('xyz-abc', loads(dumps(fun))('abc'))
+
+  def test_lambda_with_globals(self):
+    """Tests that the globals of a function are preserved."""
+
+    # The point of the test is that the lambda being called after unpickling
+    # relies on having the re module being loaded.
+    self.assertEqual(['abc', 'def'],
+                     loads(dumps(
+                         module_test.get_lambda_with_globals()))('abc def'))
+
+  def test_lambda_with_main_globals(self):
+    self.assertEqual(unittest, loads(dumps(lambda: unittest))())
+
+  def test_lambda_with_closure(self):
+    """Tests that the closure of a function is preserved."""
+    self.assertEqual(
+        'closure: abc',
+        loads(dumps(module_test.get_lambda_with_closure('abc')))())
+
+  def test_class_object_pickled(self):
+    self.assertEqual(['abc', 'def'],
+                     loads(dumps(module_test.Xyz))().foo('abc def'))
+
+  def test_class_instance_pickled(self):
+    self.assertEqual(['abc', 'def'],
+                     loads(dumps(module_test.XYZ_OBJECT)).foo('abc def'))
+
+  def test_pickling_preserves_closure_of_a_function(self):
+    self.assertEqual(
+        'X:abc', loads(dumps(module_test.TopClass.NestedClass('abc'))).datum)
+    self.assertEqual(
+        'Y:abc',
+        loads(dumps(module_test.TopClass.MiddleClass.NestedClass('abc'))).datum)
+
+  def test_pickle_dynamic_class(self):
+    self.assertEqual(
+        'Z:abc', loads(dumps(module_test.create_class('abc'))).get())
+
+  def test_generators(self):
+    with self.assertRaises(TypeError):
+      dumps((_ for _ in range(10)))
+
+  def test_recursive_class(self):
+    self.assertEqual(
+        'RecursiveClass:abc',
+        loads(dumps(module_test.RecursiveClass('abc').datum)))
+
+  def test_function_with_external_reference(self):
+    out_of_scope_var = 'expected_value'
+
+    def foo():
+      return out_of_scope_var
+
+    self.assertEqual('expected_value', loads(dumps(foo))())
+
+  @unittest.skipIf(NO_MAPPINGPROXYTYPE, 'test if MappingProxyType introduced')
+  def test_dump_and_load_mapping_proxy(self):
+    self.assertEqual(
+        'def', loads(dumps(types.MappingProxyType({'abc': 'def'})))['abc'])
+    self.assertEqual(
+        types.MappingProxyType, type(loads(dumps(types.MappingProxyType({})))))
+
+  # pylint: disable=exec-used
+  @unittest.skipIf(sys.version_info < (3, 7), 'Python 3.7 or above only')
+  def test_dataclass(self):
+    exec(
+        '''
+from apache_beam.internal.module_test import DataClass
+self.assertEqual(DataClass(datum='abc'), loads(dumps(DataClass(datum='abc'))))
+    ''')
+
+
+if __name__ == '__main__':
+  unittest.main()