Add pickler.roundtrip() shortcut for testing pickle (#36441)

* Add pickler.roundtrip() shortcut for testing pickle

`pickler.roundtrip(...)`` is equivalent to `pickler.loads(pickler.dumps(...))`,
but avoids the overhead of compression. For pipelines that serialize large
objects, this can significantly reduce the overhead of pipeline creation.

* Fix typo

* formatting

Author: shoyer

sdks/python/apache_beam/internal/cloudpickle_pickler.py

@@ -121,6 +121,28 @@ def dumps(
     enable_best_effort_determinism=False,
     config: cloudpickle.CloudPickleConfig = DEFAULT_CONFIG) -> bytes:
   """For internal use only; no backwards-compatibility guarantees."""
+  s = _dumps(o, enable_best_effort_determinism, config)
+
+  # 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 _dumps(
+    o,
+    enable_best_effort_determinism=False,
+    config: cloudpickle.CloudPickleConfig = DEFAULT_CONFIG) -> bytes:
+
   if enable_best_effort_determinism:
     # TODO: Add support once https://github.com/cloudpipe/cloudpickle/pull/563
     # is merged in.
@@ -145,21 +167,7 @@ def dumps(
       if EnumDescriptor is not None:
         pickler.dispatch_table[EnumDescriptor] = _pickle_enum_descriptor
       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)
+      return file.getvalue()
 
 
 def loads(encoded, enable_trace=True, use_zlib=False):
@@ -173,12 +181,20 @@ def loads(encoded, enable_trace=True, use_zlib=False):
     s = bz2.decompress(c)
 
   del c  # Free up some possibly large and no-longer-needed memory.
+  return _loads(s)
+
 
+def _loads(s):
   with _pickle_lock:
     unpickled = cloudpickle.loads(s)
     return unpickled
 
 
+def roundtrip(o):
+  """Internal utility for testing round-trip pickle serialization."""
+  return _loads(_dumps(o))
+
+
 def _pickle_absl_flags(obj):
   return _create_absl_flags, tuple([])
 

sdks/python/apache_beam/internal/dill_pickler.py

@@ -381,6 +381,25 @@ def dumps(
     use_zlib=False,
     enable_best_effort_determinism=False) -> bytes:
   """For internal use only; no backwards-compatibility guarantees."""
+  s = _dumps(o, enable_trace, enable_best_effort_determinism)
+
+  # 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 _dumps(o, enable_trace=True, enable_best_effort_determinism=False) -> bytes:
+  """For internal use only; no backwards-compatibility guarantees."""
   with _pickle_lock:
     if enable_best_effort_determinism:
       old_save_set = dill.dill.Pickler.dispatch[set]
@@ -400,20 +419,7 @@ def dumps(
       if enable_best_effort_determinism:
         dill.dill.pickle(set, old_save_set)
         dill.dill.pickle(frozenset, old_save_frozenset)
-
-  # 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)
+  return s
 
 
 def loads(encoded, enable_trace=True, use_zlib=False):
@@ -427,7 +433,10 @@ def loads(encoded, enable_trace=True, use_zlib=False):
     s = bz2.decompress(c)
 
   del c  # Free up some possibly large and no-longer-needed memory.
+  return _loads(s, enable_trace)
+
 
+def _loads(s, enable_trace=True):
   with _pickle_lock:
     try:
       return dill.loads(s)
@@ -441,6 +450,11 @@ def loads(encoded, enable_trace=True, use_zlib=False):
       dill.dill._trace(False)  # pylint: disable=protected-access
 
 
+def roundtrip(o):
+  """Internal utility for testing round-trip pickle serialization."""
+  return _loads(_dumps(o))
+
+
 def dump_session(file_path):
   """For internal use only; no backwards-compatibility guarantees.
 

sdks/python/apache_beam/internal/pickler.py

@@ -63,6 +63,11 @@ def loads(encoded, enable_trace=True, use_zlib=False):
       encoded, enable_trace=enable_trace, use_zlib=use_zlib)
 
 
+def roundtrip(o):
+  """Internal utility for testing round-trip pickle serialization."""
+  return desired_pickle_lib.roundtrip(o)
+
+
 def dump_session(file_path):
   """For internal use only; no backwards-compatibility guarantees.
 

sdks/python/apache_beam/io/filebasedsource.py

@@ -147,7 +147,7 @@ def _get_concat_source(self) -> concat_source.ConcatSource:
       # with each _SingleFileSource. To prevent this FileBasedSource from having
       # a reference to ConcatSource (resulting in quadratic space complexity)
       # we clone it here.
-      file_based_source_ref = pickler.loads(pickler.dumps(self))
+      file_based_source_ref = pickler.roundtrip(self)
 
       for file_metadata in files_metadata:
         file_name = file_metadata.path
@@ -284,7 +284,7 @@ def split(self, desired_bundle_size, start_offset=None, stop_offset=None):
             split.stop - split.start,
             _SingleFileSource(
                 # Copying this so that each sub-source gets a fresh instance.
-                pickler.loads(pickler.dumps(self._file_based_source)),
+                pickler.roundtrip(self._file_based_source),
                 self._file_name,
                 split.start,
                 split.stop,

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

@@ -822,7 +822,7 @@ def start_bundle(self):
 
     # TODO(aaltay): Consider storing the serialized form as an optimization.
     dofn = (
-        pickler.loads(pickler.dumps(transform.dofn))
+        pickler.roundtrip(transform.dofn)
         if self._perform_dofn_pickle_test else transform.dofn)
 
     args = transform.args if hasattr(transform, 'args') else []

sdks/python/apache_beam/transforms/core.py

@@ -3267,7 +3267,7 @@ def __init__(self):
         try:
           self._combine_fn_copy = copy.deepcopy(combine_fn)
         except Exception:
-          self._combine_fn_copy = pickler.loads(pickler.dumps(combine_fn))
+          self._combine_fn_copy = pickler.roundtrip(combine_fn)
 
         self.setup = self._combine_fn_copy.setup
         self.create_accumulator = self._combine_fn_copy.create_accumulator
@@ -3288,7 +3288,7 @@ def __init__(self):
         try:
           self._combine_fn_copy = copy.deepcopy(combine_fn)
         except Exception:
-          self._combine_fn_copy = pickler.loads(pickler.dumps(combine_fn))
+          self._combine_fn_copy = pickler.roundtrip(combine_fn)
 
         self.setup = self._combine_fn_copy.setup
         self.create_accumulator = self._combine_fn_copy.create_accumulator

sdks/python/apache_beam/transforms/ptransform.py

@@ -875,12 +875,12 @@ def __init__(self, fn, *args, **kwargs):
 
     # Ensure fn and side inputs are picklable for remote execution.
     try:
-      self.fn = pickler.loads(pickler.dumps(self.fn))
+      self.fn = pickler.roundtrip(self.fn)
     except RuntimeError as e:
       raise RuntimeError('Unable to pickle fn %s: %s' % (self.fn, e))
 
-    self.args = pickler.loads(pickler.dumps(self.args))
-    self.kwargs = pickler.loads(pickler.dumps(self.kwargs))
+    self.args = pickler.roundtrip(self.args)
+    self.kwargs = pickler.roundtrip(self.kwargs)
 
     # For type hints, because loads(dumps(class)) != class.
     self.fn = self._cached_fn