Integrate lambda name pickling with Cloudpickle

sdks/python/apache_beam/internal/cloudpickle/cloudpickle.py

@@ -107,6 +107,12 @@ def uuid_generator(_):
   return uuid.uuid4().hex
 
 
+@dataclasses.dataclass
+class GetCodeObjectParams:
+  get_code_object_identifier: typing.Optional[callable]
+  get_code_from_identifier: typing.Optional[callable]
+
+
 @dataclasses.dataclass
 class CloudPickleConfig:
   """Configuration for cloudpickle behavior.
@@ -127,10 +133,18 @@ class CloudPickleConfig:
 
         filepath_interceptor: Used to modify filepaths in `co_filename` and
             function.__globals__['__file__'].
+
+        get_code_object_params: Use identifiers derived from code
+            location when pickling dynamic functions (e.g. lambdas). Enabling
+            this setting results in pickled payloads becoming more stable to
+            code changes: when a particular lambda function is slightly
+            modified  but the location of the function in the codebase has not
+            changed, the pickled representation might stay the same.
     """
   id_generator: typing.Optional[callable] = uuid_generator
   skip_reset_dynamic_type_state: bool = False
   filepath_interceptor: typing.Optional[callable] = None
+  get_code_object_params: typing.Optional[GetCodeObjectParams] = None
 
 
 DEFAULT_CONFIG = CloudPickleConfig()
@@ -567,6 +581,12 @@ def _make_function(code, globals, name, argdefs, closure):
   return types.FunctionType(code, globals, name, argdefs, closure)
 
 
+def _make_function_from_identifier(
+    get_code_from_identifier, code_path, globals, name, argdefs, closure):
+  fcode = get_code_from_identifier(code_path)
+  return _make_function(fcode, globals, name, argdefs, closure)
+
+
 def _make_empty_cell():
   if False:
     # trick the compiler into creating an empty cell in our lambda
@@ -1305,6 +1325,30 @@ class Pickler(pickle.Pickler):
 
   dispatch_table = ChainMap(_dispatch_table, copyreg.dispatch_table)
 
+  def _stable_identifier_function_reduce(self, func):
+    code_object_params = self.config.get_code_object_params
+    if code_object_params is None:
+      return self._dynamic_function_reduce(func)
+    code_path = code_object_params.get_code_object_identifier(func)
+    if not code_path:
+      return self._dynamic_function_reduce(func)
+    newargs = (
+        code_path,
+        func.__globals__,
+        func.__name__,
+        func.__defaults__,
+        func.__closure__)
+    state = _function_getstate(func)
+    return (
+        functools.partial(
+            _make_function_from_identifier,
+            code_object_params.get_code_from_identifier),
+        newargs,
+        state,
+        None,
+        None,
+        _function_setstate)
+
   # function reducers are defined as instance methods of cloudpickle.Pickler
   # objects, as they rely on a cloudpickle.Pickler attribute (globals_ref)
   def _dynamic_function_reduce(self, func):
@@ -1324,6 +1368,8 @@ def _function_reduce(self, obj):
         """
     if _should_pickle_by_reference(obj):
       return NotImplemented
+    elif self.config.get_code_object_params is not None:
+      return self._stable_identifier_function_reduce(obj)
     else:
       return self._dynamic_function_reduce(obj)
 

sdks/python/apache_beam/internal/cloudpickle_pickler.py

@@ -36,11 +36,18 @@
 import zlib
 
 from apache_beam.internal.cloudpickle import cloudpickle
+from apache_beam.internal import code_object_pickler
 
 DEFAULT_CONFIG = cloudpickle.CloudPickleConfig(
     skip_reset_dynamic_type_state=True)
 NO_DYNAMIC_CLASS_TRACKING_CONFIG = cloudpickle.CloudPickleConfig(
     id_generator=None, skip_reset_dynamic_type_state=True)
+STABLE_CODE_IDENTIFIER_CONFIG = cloudpickle.CloudPickleConfig(
+    skip_reset_dynamic_type_state=True,
+    get_code_object_params=cloudpickle.GetCodeObjectParams(
+        get_code_object_identifier=code_object_pickler.
+        get_code_object_identifier,
+        get_code_from_identifier=code_object_pickler.get_code_from_identifier))
 
 try:
   from absl import flags
@@ -119,6 +126,7 @@ def dumps(
     enable_trace=True,
     use_zlib=False,
     enable_best_effort_determinism=False,
+    enable_stable_code_identifier_pickling=False,
     config: cloudpickle.CloudPickleConfig = DEFAULT_CONFIG) -> bytes:
   """For internal use only; no backwards-compatibility guarantees."""
   if enable_best_effort_determinism:
@@ -129,6 +137,8 @@ def dumps(
         'This has only been implemented for dill.')
   with _pickle_lock:
     with io.BytesIO() as file:
+      if enable_stable_code_identifier_pickling:
+        config = STABLE_CODE_IDENTIFIER_CONFIG
       pickler = cloudpickle.CloudPickler(file, config=config)
       try:
         pickler.dispatch_table[type(flags.FLAGS)] = _pickle_absl_flags

sdks/python/apache_beam/internal/code_object_pickler.py

@@ -19,7 +19,7 @@
 
 This module provides helper functions to improve pickling code objects,
 especially lambdas, in a consistent way by using code object identifiers. These
-helper functions will be used to patch pickler implementations used by Beam
+helper functions are used to patch pickler implementations used by Beam
 (e.g. Cloudpickle).
 
 A code object identifier is a unique identifier for a code object that provides
@@ -81,8 +81,9 @@ def get_code_object_identifier(callable: types.FunctionType):
       - __main__.ClassWithNestedLambda.process.__code__.co_consts[
         <lambda>, ('x',), 1234567890]
   """
-  if not hasattr(callable, '__module__') or not hasattr(callable,
-                                                        '__qualname__'):
+  if (not hasattr(callable, '__module__') or
+      not hasattr(callable, '__qualname__') or not callable.__module__ or
+      callable.__module__ not in sys.modules):
     return None
   code_path: str = _extend_path(
       callable.__module__,
@@ -100,7 +101,7 @@ def _extend_path(prefix: str, current_path: Optional[str]):
 
   Args:
     prefix: The prefix of the path.
-    suffix: The rest of the path.
+    current_path: The rest of the path.
 
   Returns:
     The extended path.
@@ -189,6 +190,8 @@ def _search_module_or_class(
           if path is not None:
             return _extend_path(name, _extend_path(f'__defaults__[{i}]', path))
   else:
+    if not hasattr(node, first_part):
+      return None
     return _extend_path(
         first_part, _search(callable, getattr(node, first_part), rest))
 
@@ -281,6 +284,8 @@ def _search_lambda(
   lambda_code_objects_by_name = collections.defaultdict(list)
   name = qual_name_parts[0]
   code_objects = code_objects_by_name[name]
+  if not code_objects:
+    return None
   if name == '<lambda>':
     for code_object in code_objects:
       lambda_name = f'<lambda>, {_signature(code_object)}'
@@ -315,10 +320,10 @@ def _search_lambda(
 _SINGLE_NAME_PATTERN = re.compile(r'co_consts\[([a-zA-Z0-9\<\>_-]+)]')
 # Matches a path like: co_consts[<lambda>, ('x',)]
 _LAMBDA_WITH_ARGS_PATTERN = re.compile(
-    r"co_consts\[(<[^>]+>),\s*(\('[^']*'\s*,\s*\))\]")
+    r"co_consts\[(<.*?>),\s(\('[^']+'(?:,\s*'[^']+')*,?\))\]")
 # Matches a path like: co_consts[<lambda>, ('x',), 1234567890]
 _LAMBDA_WITH_HASH_PATTERN = re.compile(
-    r"co_consts\[(<[^>]+>),\s*(\('[^']*'\s*,\s*\)),\s*(.+)\]")
+    r"co_consts\[(<[^>]+>),\s*(\([^\)]*\)),?\s*(.*)\]")
 # Matches a path like: __defaults__[0]
 _DEFAULT_PATTERN = re.compile(r'(__defaults__)\[(\d+)\]')
 # Matches an argument like: 'x'
@@ -345,9 +350,10 @@ def _get_code_object_from_single_name_pattern(
     raise ValueError(f'Invalid pattern for single name: {name_result.group(0)}')
   # Groups are indexed starting at 1, group(0) is the entire match.
   name = name_result.group(1)
-  for co_const in obj.co_consts:
-    if inspect.iscode(co_const) and co_const.co_name == name:
-      return co_const
+  if hasattr(obj, 'co_consts'):
+    for co_const in obj.co_consts:
+      if inspect.iscode(co_const) and co_const.co_name == name:
+        return co_const
   raise AttributeError(f'Could not find code object with path: {path}')
 
 
@@ -368,15 +374,16 @@ def _get_code_object_from_lambda_with_args_pattern(
   """
   name = lambda_with_args_result.group(1)
   code_objects = collections.defaultdict(list)
-  for co_const in obj.co_consts:
-    if inspect.iscode(co_const) and co_const.co_name == name:
-      code_objects[co_const.co_name].append(co_const)
-  for name, objects in code_objects.items():
-    for obj_ in objects:
-      args = tuple(
-          re.findall(_ARGUMENT_PATTERN, lambda_with_args_result.group(2)))
-      if obj_.co_varnames == args:
-        return obj_
+  if hasattr(obj, 'co_consts'):
+    for co_const in obj.co_consts:
+      if inspect.iscode(co_const) and co_const.co_name == name:
+        code_objects[co_const.co_name].append(co_const)
+    for name, objects in code_objects.items():
+      for obj_ in objects:
+        args = tuple(
+            re.findall(_ARGUMENT_PATTERN, lambda_with_args_result.group(2)))
+        if obj_.co_varnames[:_get_arg_count(obj_)] == args:
+          return obj_
   raise AttributeError(f'Could not find code object with path: {path}')
 
 
@@ -397,17 +404,18 @@ def _get_code_object_from_lambda_with_hash_pattern(
   """
   name = lambda_with_hash_result.group(1)
   code_objects = collections.defaultdict(list)
-  for co_const in obj.co_consts:
-    if inspect.iscode(co_const) and co_const.co_name == name:
-      code_objects[co_const.co_name].append(co_const)
-  for name, objects in code_objects.items():
-    for obj_ in objects:
-      args = tuple(
-          re.findall(_ARGUMENT_PATTERN, lambda_with_hash_result.group(2)))
-      if obj_.co_varnames == args:
-        hash_value = lambda_with_hash_result.group(3)
-        if hash_value == str(_create_bytecode_hash(obj_)):
-          return obj_
+  if hasattr(obj, 'co_consts'):
+    for co_const in obj.co_consts:
+      if inspect.iscode(co_const) and co_const.co_name == name:
+        code_objects[co_const.co_name].append(co_const)
+    for name, objects in code_objects.items():
+      for obj_ in objects:
+        args = tuple(
+            re.findall(_ARGUMENT_PATTERN, lambda_with_hash_result.group(2)))
+        if obj_.co_varnames[:_get_arg_count(obj_)] == args:
+          hash_value = lambda_with_hash_result.group(3)
+          if hash_value == str(_create_bytecode_hash(obj_)):
+            return obj_
   raise AttributeError(f'Could not find code object with path: {path}')
 
 
@@ -427,6 +435,8 @@ def get_code_from_identifier(code_object_identifier: str):
   if not code_object_identifier:
     raise ValueError('Path must not be empty.')
   parts = code_object_identifier.split('.')
+  if parts[0] not in sys.modules:
+    raise AttributeError(f'Module {parts[0]} not found in sys.modules')
   obj = sys.modules[parts[0]]
   for part in parts[1:]:
     if name_result := _SINGLE_NAME_PATTERN.fullmatch(part):
@@ -447,7 +457,11 @@ def get_code_from_identifier(code_object_identifier: str):
       obj = getattr(obj, '__defaults__')[index]
     else:
       obj = getattr(obj, part)
-  return obj
+  if isinstance(obj, types.CodeType):
+    return obj
+  else:
+    raise AttributeError(
+        f'Could not find code object with path: {code_object_identifier}')
 
 
 def _signature(obj: types.CodeType):
@@ -462,12 +476,24 @@ def _signature(obj: types.CodeType):
   Returns:
     A tuple of the names of the arguments of the code object.
   """
-  arg_count = (
+  return obj.co_varnames[:_get_arg_count(obj)]
+
+
+def _get_arg_count(obj: types.CodeType):
+  """Returns the number of arguments of a code object.
+
+  Args:
+    obj: A code object, function, method, or cell.
+
+  Returns:
+    The number of arguments of the code object, or None if the object is not a
+    code object.
+  """
+  return (
       obj.co_argcount + obj.co_kwonlyargcount +
       (obj.co_flags & 4 == 4)  # PyCF_VARARGS
       + (obj.co_flags & 8 == 8)  # PyCF_VARKEYWORDS
   )
-  return obj.co_varnames[:arg_count]
 
 
 def _create_bytecode_hash(code_object: types.CodeType):

sdks/python/apache_beam/internal/code_object_pickler_test.py

@@ -274,12 +274,14 @@ def test_adding_lambda_variable_in_class_preserves_object(self):
         module_2_modified.AddLambdaVariable.my_method(self).__code__,
     )
 
-  def test_removing_lambda_variable_in_class_changes_object(self):
-    with self.assertRaisesRegex(AttributeError, "object has no attribute"):
-      code_object_pickler.get_code_from_identifier(
-          code_object_pickler.get_code_object_identifier(
-              module_2.RemoveLambdaVariable.my_method(self)).replace(
-                  "module_2", "module_2_modified"))
+  def test_removing_lambda_variable_in_class_preserves_object(self):
+    self.assertEqual(
+        code_object_pickler.get_code_from_identifier(
+            code_object_pickler.get_code_object_identifier(
+                module_2.RemoveLambdaVariable.my_method(self)).replace(
+                    "module_2", "module_2_modified")),
+        module_2_modified.RemoveLambdaVariable.my_method(self).__code__,
+    )
 
   def test_adding_nested_function_in_class_preserves_object(self):
     self.assertEqual(
@@ -391,11 +393,14 @@ def test_adding_lambda_variable_in_function_preserves_object(self):
         module_1_lambda_variable_added.my_function().__code__,
     )
 
-  def test_removing_lambda_variable_in_function_raises_exception(self):
-    with self.assertRaisesRegex(AttributeError, "object has no attribute"):
-      code_object_pickler.get_code_from_identifier(
-          code_object_pickler.get_code_object_identifier(
-              module_3.my_function()).replace("module_3", "module_3_modified"))
+  def test_removing_lambda_variable_in_function_preserves_object(self):
+    self.assertEqual(
+        code_object_pickler.get_code_from_identifier(
+            code_object_pickler.get_code_object_identifier(
+                module_3.my_function()).replace(
+                    "module_3", "module_3_modified")),
+        module_3_modified.my_function().__code__,
+    )
 
 
 class CodePathStabilityTest(unittest.TestCase):

sdks/python/apache_beam/internal/pickler.py

@@ -47,8 +47,18 @@ def dumps(
     o,
     enable_trace=True,
     use_zlib=False,
-    enable_best_effort_determinism=False) -> bytes:
-
+    enable_best_effort_determinism=False,
+    enable_stable_code_identifier_pickling=False) -> bytes:
+
+  if (desired_pickle_lib == cloudpickle_pickler):
+    return cloudpickle_pickler.dumps(
+        o,
+        enable_trace=enable_trace,
+        use_zlib=use_zlib,
+        enable_best_effort_determinism=enable_best_effort_determinism,
+        enable_stable_code_identifier_pickling=
+        enable_stable_code_identifier_pickling,
+    )
   return desired_pickle_lib.dumps(
       o,
       enable_trace=enable_trace,

sdks/python/apache_beam/internal/pickler_test.py

@@ -302,6 +302,21 @@ def test_disable_best_effort_determinism(self):
         dumps(set1, enable_best_effort_determinism=False),
         dumps(set2, enable_best_effort_determinism=False))
 
+  def test_enable_stable_code_identifier_pickling(self):
+    pickler.set_library('cloudpickle')
+    pickled = pickler.loads(
+        pickler.dumps(lambda x: x, enable_stable_code_identifier_pickling=True))
+    pickled_type = type(pickled)
+    self.assertIsInstance(pickled, pickled_type)
+
+  def test_disable_stable_code_identifier_pickling(self):
+    pickler.set_library('cloudpickle')
+    pickled = pickler.loads(
+        pickler.dumps(
+            lambda x: x, enable_stable_code_identifier_pickling=False))
+    pickled_type = type(pickled)
+    self.assertIsInstance(pickled, pickled_type)
+
 
 if __name__ == '__main__':
   unittest.main()

sdks/python/apache_beam/runners/pipeline_context.py

@@ -227,6 +227,7 @@ def __init__(
     self.iterable_state_write = iterable_state_write
     self._requirements = set(requirements)
     self.enable_best_effort_deterministic_pickling = False
+    self.enable_stable_code_identifier_pickling = False
 
   def add_requirement(self, requirement: str) -> None:
     self._requirements.add(requirement)