From 02ab9dafd00e7c9d0ad62272de2694a5797e9a04 Mon Sep 17 00:00:00 2001 From: kushmiD Date: Fri, 1 Nov 2024 09:17:36 -0700 Subject: [PATCH] Support caching in Apache Beam by using relative co_filename paths. (#32979) * The motivation for this change is to support caching in Apache Beam. Apache Beam does the following: - Pickle Python code - Send the pickled source code to "worker" VMs - The workers unpickle and execute the code In the environment that these Beam pipelines execute, the source code is in a temporary directory whose name is random and changes. The source code paths relative to the temporary directory are constant. Using absolute paths prevents pickled code from being cached because the absolute path keeps changing. Using relative paths enables this caching and promises significant resource savings and speed-ups. Additionally the absolute paths leak information about the directory structure of the machine pickling the source code. When the pickled code is passed across the network to another machine, the absolute paths may no longer be valid when the other machine has a different directory structure. The reason for using relative paths rather than omitting the path entirely is because Python uses the co_filename attribute to create stack traces. * The motivation for this change is to support caching in Apache Beam for Google. Apache Beam does the following: - Pickle Python code - Send the pickled source code to "worker" VMs - The workers unpickle and execute the code In the environment that these Beam pipelines execute, the source code is in a temporary directory whose name is random and changes. The source code paths relative to the temporary directory are constant. Using absolute paths prevents pickled code from being cached because the absolute path keeps changing. Using relative paths enables this caching and promises significant resource savings and speed-ups. Additionally the absolute paths leak information about the directory structure of the machine pickling the source code. When the pickled code is passed across the network to another machine, the absolute paths may no longer be valid when the other machine has a different directory structure. The reason for using relative paths rather than omitting the path entirely is because Python uses the co_filename attribute to create stack traces. * Simplify. --------- Co-authored-by: Robert Bradshaw --- .../apache_beam/internal/dill_pickler.py | 19 +++++++++++++------ .../apache_beam/internal/pickler_test.py | 5 +++++ 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/sdks/python/apache_beam/internal/dill_pickler.py b/sdks/python/apache_beam/internal/dill_pickler.py index 7f7ac5b214fa..e1d6b7e74e49 100644 --- a/sdks/python/apache_beam/internal/dill_pickler.py +++ b/sdks/python/apache_beam/internal/dill_pickler.py @@ -46,9 +46,15 @@ settings = {'dill_byref': None} -if sys.version_info >= (3, 10) and dill.__version__ == "0.3.1.1": - # Let's make dill 0.3.1.1 support Python 3.11. +patch_save_code = sys.version_info >= (3, 10) and dill.__version__ == "0.3.1.1" + +def get_normalized_path(path): + """Returns a normalized path. This function is intended to be overridden.""" + return path + + +if patch_save_code: # The following function is based on 'save_code' from 'dill' # Author: Mike McKerns (mmckerns @caltech and @uqfoundation) # Copyright (c) 2008-2015 California Institute of Technology. @@ -66,6 +72,7 @@ @dill.register(CodeType) def save_code(pickler, obj): + co_filename = get_normalized_path(obj.co_filename) if hasattr(obj, "co_endlinetable"): # python 3.11a (20 args) args = ( obj.co_argcount, @@ -78,7 +85,7 @@ def save_code(pickler, obj): obj.co_consts, obj.co_names, obj.co_varnames, - obj.co_filename, + co_filename, obj.co_name, obj.co_qualname, obj.co_firstlineno, @@ -100,7 +107,7 @@ def save_code(pickler, obj): obj.co_consts, obj.co_names, obj.co_varnames, - obj.co_filename, + co_filename, obj.co_name, obj.co_qualname, obj.co_firstlineno, @@ -120,7 +127,7 @@ def save_code(pickler, obj): obj.co_consts, obj.co_names, obj.co_varnames, - obj.co_filename, + co_filename, obj.co_name, obj.co_firstlineno, obj.co_linetable, @@ -138,7 +145,7 @@ def save_code(pickler, obj): obj.co_consts, obj.co_names, obj.co_varnames, - obj.co_filename, + co_filename, obj.co_name, obj.co_firstlineno, obj.co_lnotab, diff --git a/sdks/python/apache_beam/internal/pickler_test.py b/sdks/python/apache_beam/internal/pickler_test.py index 824c4c59c0ce..c26a8ee3e653 100644 --- a/sdks/python/apache_beam/internal/pickler_test.py +++ b/sdks/python/apache_beam/internal/pickler_test.py @@ -94,6 +94,11 @@ def test_pickle_rlock(self): self.assertIsInstance(loads(dumps(rlock_instance)), rlock_type) + def test_save_paths(self): + f = loads(dumps(lambda x: x)) + co_filename = f.__code__.co_filename + self.assertTrue(co_filename.endswith('pickler_test.py')) + @unittest.skipIf(NO_MAPPINGPROXYTYPE, 'test if MappingProxyType introduced') def test_dump_and_load_mapping_proxy(self): self.assertEqual(