From 140afdc12aeee12d6418668db7f6f748d8acc396 Mon Sep 17 00:00:00 2001 From: Stephan Hoyer Date: Wed, 8 Oct 2025 10:19:01 -0700 Subject: [PATCH 1/3] 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. --- .../internal/cloudpickle_pickler.py | 46 +++++++++++++------ .../apache_beam/internal/dill_pickler.py | 45 ++++++++++++------ sdks/python/apache_beam/internal/pickler.py | 5 ++ sdks/python/apache_beam/io/filebasedsource.py | 4 +- .../runners/direct/transform_evaluator.py | 2 +- sdks/python/apache_beam/transforms/core.py | 4 +- .../apache_beam/transforms/ptransform.py | 6 +-- 7 files changed, 75 insertions(+), 37 deletions(-) diff --git a/sdks/python/apache_beam/internal/cloudpickle_pickler.py b/sdks/python/apache_beam/internal/cloudpickle_pickler.py index e55818bfb226..e01ec9c774ca 100644 --- a/sdks/python/apache_beam/internal/cloudpickle_pickler.py +++ b/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 = _dump(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([]) diff --git a/sdks/python/apache_beam/internal/dill_pickler.py b/sdks/python/apache_beam/internal/dill_pickler.py index 9a3d43826610..41820043bff3 100644 --- a/sdks/python/apache_beam/internal/dill_pickler.py +++ b/sdks/python/apache_beam/internal/dill_pickler.py @@ -381,6 +381,28 @@ 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 +422,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 +436,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 +453,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. diff --git a/sdks/python/apache_beam/internal/pickler.py b/sdks/python/apache_beam/internal/pickler.py index 6f8dba463bc3..c1a54e6e961e 100644 --- a/sdks/python/apache_beam/internal/pickler.py +++ b/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. diff --git a/sdks/python/apache_beam/io/filebasedsource.py b/sdks/python/apache_beam/io/filebasedsource.py index 49b1b1d125f1..b80e4fb8a841 100644 --- a/sdks/python/apache_beam/io/filebasedsource.py +++ b/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, diff --git a/sdks/python/apache_beam/runners/direct/transform_evaluator.py b/sdks/python/apache_beam/runners/direct/transform_evaluator.py index ee97b729ac28..3443a519e54c 100644 --- a/sdks/python/apache_beam/runners/direct/transform_evaluator.py +++ b/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 [] diff --git a/sdks/python/apache_beam/transforms/core.py b/sdks/python/apache_beam/transforms/core.py index 2126169a57fb..7ba8aa128c24 100644 --- a/sdks/python/apache_beam/transforms/core.py +++ b/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 diff --git a/sdks/python/apache_beam/transforms/ptransform.py b/sdks/python/apache_beam/transforms/ptransform.py index cac8a8fbd957..55453a3e92eb 100644 --- a/sdks/python/apache_beam/transforms/ptransform.py +++ b/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 From 3cda7894af56c99ffec0db95b3416f2645f75f37 Mon Sep 17 00:00:00 2001 From: Stephan Hoyer Date: Wed, 8 Oct 2025 12:51:19 -0700 Subject: [PATCH 2/3] Fix typo --- sdks/python/apache_beam/internal/cloudpickle_pickler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/internal/cloudpickle_pickler.py b/sdks/python/apache_beam/internal/cloudpickle_pickler.py index e01ec9c774ca..d2fa4d72395a 100644 --- a/sdks/python/apache_beam/internal/cloudpickle_pickler.py +++ b/sdks/python/apache_beam/internal/cloudpickle_pickler.py @@ -121,7 +121,7 @@ def dumps( enable_best_effort_determinism=False, config: cloudpickle.CloudPickleConfig = DEFAULT_CONFIG) -> bytes: """For internal use only; no backwards-compatibility guarantees.""" - s = _dump(o, enable_best_effort_determinism, config) + 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 From ca431c4152cd56db4f9375effd158dc257a2cfab Mon Sep 17 00:00:00 2001 From: Stephan Hoyer Date: Wed, 8 Oct 2025 13:12:23 -0700 Subject: [PATCH 3/3] formatting --- sdks/python/apache_beam/internal/dill_pickler.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/sdks/python/apache_beam/internal/dill_pickler.py b/sdks/python/apache_beam/internal/dill_pickler.py index 41820043bff3..e88cb3c1e138 100644 --- a/sdks/python/apache_beam/internal/dill_pickler.py +++ b/sdks/python/apache_beam/internal/dill_pickler.py @@ -398,10 +398,7 @@ def dumps( return base64.b64encode(c) -def _dumps( - o, - enable_trace=True, - enable_best_effort_determinism=False) -> bytes: +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: