praneetnadella commented on code in PR #40711:
URL: https://github.com/apache/beam/pull/40711#discussion_r4473080991
##########
sdks/python/apache_beam/internal/code_object_pickler_test.py:
##########
@@ -214,6 +367,163 @@ def test_roundtrip(self, callable, unused_path):
actual = code_object_pickler.get_code_from_identifier(path)
self.assertEqual(actual, callable.__code__)
+ def test_two_lambdas_where_one_has_no_arguments_roundtrips(self):
+ zero_arg_lambda, one_arg_lambda = get_zero_arg_and_one_arg_lambdas()
+ zero_arg_path = code_object_pickler.get_code_object_identifier(
+ zero_arg_lambda)
+ one_arg_path = code_object_pickler.get_code_object_identifier(
+ one_arg_lambda)
+
+ self.assertEqual(
+ zero_arg_path,
+ f"{prefix}.get_zero_arg_and_one_arg_lambdas.__code__.co_consts"
+ "[<lambda>, ()]",
+ )
+ self.assertEqual(
+ one_arg_path,
+ f"{prefix}.get_zero_arg_and_one_arg_lambdas.__code__.co_consts"
+ "[<lambda>, ('x',)]",
+ )
+ self.assertEqual(
+ code_object_pickler.get_code_from_identifier(zero_arg_path),
+ zero_arg_lambda.__code__,
+ )
+ self.assertEqual(
+ code_object_pickler.get_code_from_identifier(one_arg_path),
+ one_arg_lambda.__code__,
+ )
+
+ def test_same_co_code_lambdas_roundtrip_distinct_code_objects(self):
+ lambda_a, lambda_b = get_same_co_code_lambdas()
+ path_a = code_object_pickler.get_code_object_identifier(lambda_a)
+ path_b = code_object_pickler.get_code_object_identifier(lambda_b)
+
+ self.assertNotEqual(path_a, path_b)
+ self.assertEqual(
+ code_object_pickler.get_code_from_identifier(path_a),
+ lambda_a.__code__,
+ )
+ self.assertEqual(
+ code_object_pickler.get_code_from_identifier(path_b),
+ lambda_b.__code__,
+ )
+
+ def test_shared_decorator_wrapper_on_outer_and_inner_returns_none(self):
+ inner = outer_decorated_with_shared_wrapper()
+ self.assertIsNone(code_object_pickler.get_code_object_identifier(inner))
+
+ def test_module_reexporting_same_named_leaf_resolves_longest_module(self):
+ parent_mod = types.ModuleType("_fake_pkg")
+ child_mod = types.ModuleType("_fake_pkg.leaf")
+ child_mod.leaf = top_level_function
+ # Simulate `from _fake_pkg.leaf import leaf` in `_fake_pkg/__init__.py`
+ parent_mod.leaf = top_level_function
+ sys.modules["_fake_pkg"] = parent_mod
+ sys.modules["_fake_pkg.leaf"] = child_mod
+ try:
+ actual = code_object_pickler.get_code_from_identifier(
+ "_fake_pkg.leaf.leaf.__code__")
+ self.assertEqual(actual, top_level_function.__code__)
+ finally:
+ sys.modules.pop("_fake_pkg", None)
+ sys.modules.pop("_fake_pkg.leaf", None)
+
+ def test_same_named_nested_functions_with_different_args_roundtrip(self):
+ fn_true = get_same_named_nested_functions_different_args(True)
+ fn_false = get_same_named_nested_functions_different_args(False)
+ path_true = code_object_pickler.get_code_object_identifier(fn_true)
+ path_false = code_object_pickler.get_code_object_identifier(fn_false)
+
+ self.assertNotEqual(path_true, path_false)
+ self.assertEqual(
+ code_object_pickler.get_code_from_identifier(path_true),
+ fn_true.__code__,
+ )
+ self.assertEqual(
+ code_object_pickler.get_code_from_identifier(path_false),
+ fn_false.__code__,
+ )
+
+ def test_same_named_nested_functions_with_same_args_roundtrip(self):
+ fn_true = get_same_named_nested_functions_same_args(True)
+ fn_false = get_same_named_nested_functions_same_args(False)
+ path_true = code_object_pickler.get_code_object_identifier(fn_true)
+ path_false = code_object_pickler.get_code_object_identifier(fn_false)
+
+ self.assertNotEqual(path_true, path_false)
+ self.assertEqual(
+ code_object_pickler.get_code_from_identifier(path_true),
+ fn_true.__code__,
+ )
+ self.assertEqual(
+ code_object_pickler.get_code_from_identifier(path_false),
+ fn_false.__code__,
+ )
+
+ def test_same_nested_functions_with_complex_consts_roundtrip(self):
+ fn_true = get_same_named_nested_functions_with_complex_consts(True)
+ fn_false = get_same_named_nested_functions_with_complex_consts(False)
+ path_true = code_object_pickler.get_code_object_identifier(fn_true)
+ path_false = code_object_pickler.get_code_object_identifier(fn_false)
+
+ self.assertNotEqual(path_true, path_false)
+ self.assertEqual(
+ code_object_pickler.get_code_from_identifier(path_true),
+ fn_true.__code__,
+ )
+ self.assertEqual(
+ code_object_pickler.get_code_from_identifier(path_false),
+ fn_false.__code__,
+ )
+
+ def test_get_code_from_identifier_imports_unloaded_module(self):
+ mod_name = "apache_beam.internal.test_data.module_1"
+ saved = sys.modules.pop(mod_name, None)
+ try:
+ code = code_object_pickler.get_code_from_identifier(
+ f"{mod_name}.my_function.__code__")
+ self.assertEqual(code, module_1.my_function.__code__)
+ finally:
+ if saved is not None:
+ sys.modules[mod_name] = saved
+ else:
+ sys.modules.pop(mod_name, None)
+
+ def test_closures_with_different_param_kinds_have_distinct_identifiers(self):
+ lambda_pos, lambda_vararg = get_same_named_lambdas_different_param_kinds()
+ path_pos = code_object_pickler.get_code_object_identifier(lambda_pos)
+ path_vararg = code_object_pickler.get_code_object_identifier(lambda_vararg)
+
+ self.asserNotEqual(path_pos, path_vararg)
Review Comment:
Typo: `asserNotEqual` should be `assertNotEqual`, otherwise this test raises
`AttributeError`.
##########
sdks/python/apache_beam/internal/code_object_pickler.py:
##########
@@ -357,13 +352,13 @@ def _get_code_object_from_single_name_pattern(
raise AttributeError(f'Could not find code object with path: {path}')
-def _get_code_object_from_lambda_with_args_pattern(
- obj: types.ModuleType, lambda_with_args_result: re.Match[str], path: str):
- """Returns the code object from a lambda with args pattern.
+def _get_code_object_from_object_with_args_pattern(
+ obj: types.ModuleType, object_with_args_result: re.Match[str], path: str):
+ """Returns the code object from an object with args pattern.
Args:
obj: The object to search within.
- lambda_with_args_result: The result of the lambda with args pattern search.
+ object_with_args_result: The result of the lambda with args pattern search.
Review Comment:
nit: The docstrings for `_get_code_object_from_object_with_args_pattern` /
`_get_code_object_from_object_with_hash_pattern` still say "lambda with
args/hash pattern"; maybe update them to "object" to match the rename.
(note this is same below)
##########
sdks/python/apache_beam/internal/code_object_pickler.py:
##########
@@ -435,19 +430,37 @@ 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:
+ max_module_parts = len(parts)
+ for idx, part in enumerate(parts):
+ if part == '__code__' or part.startswith(('co_consts[', '__defaults__[')):
+ max_module_parts = idx
+ break
+ obj = None
+ start_index = 0
+ for i in range(len(parts), 0, -1):
Review Comment:
Following up on my earlier comment, I should have been more precise:
`max_module_parts` is computed above but this loop still uses `len(parts)`.
Should it be `range(max_module_parts, 0, -1)` with `max_module_parts = max(1,
idx - 1)`? `parts[idx - 1]` is always the function or lambda name, so otherwise
we still call `import_module("pkg.foo")`. If `pkg/__init__.py` defines `foo`
and an unimported `pkg/foo.py` exists, that import replaces `pkg.foo` with the
module and then fails with `AttributeError`.
##########
sdks/python/apache_beam/internal/code_object_pickler.py:
##########
@@ -279,51 +268,57 @@ def _search_lambda(
Returns:
The code object identifier, or None if not found.
"""
- # There are multiple lambdas in the code object, so we need to calculate
- # the signature and the hash to identify the correct lambda.
- lambda_code_objects_by_name = collections.defaultdict(list)
+ if not qual_name_parts:
+ return None
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)}'
- lambda_code_objects_by_name[lambda_name].append(code_object)
- # Check if there are any lambdas with the same signature.
- # If there are, we need to calculate the hash to identify the correct
- # lambda.
- for lambda_name, lambda_objects in lambda_code_objects_by_name.items():
- if len(lambda_objects) > 1:
- for lambda_object in lambda_objects:
- path = _search(callable, lambda_object, qual_name_parts)
- if path is not None:
- return _extend_path(
- f'co_consts[{lambda_name},'
- f' {_create_bytecode_hash(lambda_object)}]',
- path,
- )
- else:
- # If there is only one lambda with this signature, we can
- # use the signature to identify the correct lambda.
- path = _search(callable, code_objects[0], qual_name_parts)
- if path is not None:
- return _extend_path(f'co_consts[{lambda_name}]', path)
- else:
- # For non lambda objects, we can use the name to identify the object.
+ # If there is only one code object with this name, we can use the default
path
+ # 'co_consts[<name>]'. This is the most common case and it is faster than
+ # calculating the signature and the hash.
+ if len(code_objects) == 1:
path = _search(callable, code_objects[0], qual_name_parts)
if path is not None:
return _extend_path(f'co_consts[{name}]', path)
+ return None
+
+ # There are multiple code objects with the same name, so we need to calculate
+ # the signature and the hash to identify the correct code object.
+ code_objects_by_sig = collections.defaultdict(list)
+ for code_object in code_objects:
+ sig_name = f'{name}, {_signature(code_object)}'
+ code_objects_by_sig[sig_name].append(code_object)
+ # Check if there are any code objects with the same signature.
+ # If there are, we need to calculate the hash to identify the correct
+ # code object.
+ for sig_name, sig_objects in code_objects_by_sig.items():
+ if len(sig_objects) > 1:
+ for code_object in sig_objects:
+ path = _search(callable, code_object, qual_name_parts)
+ if path is not None:
+ return _extend_path(
+ f'co_consts[{sig_name},'
+ f' {_create_bytecode_hash(code_object)}]',
+ path,
+ )
+ else:
+ # If there is only one code object with this signature, we can
+ # use the signature to identify the correct code object.
+ path = _search(callable, sig_objects[0], qual_name_parts)
+ if path is not None:
+ return _extend_path(f'co_consts[{sig_name}]', path)
# Matches a path like: co_consts[my_function]
_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*'[^']+')*,?\))\]")
-# Matches a path like: co_consts[<lambda>, ('x',), 1234567890]
-_LAMBDA_WITH_HASH_PATTERN = re.compile(
- r"co_consts\[(<[^>]+>),\s*(\([^\)]*\)),?\s*(.*)\]")
+# Matches a path like: co_consts[<lambda>, ()] or co_consts[my_function,
('x',)]
+_OBJECT_WITH_ARGS_PATTERN = re.compile(
Review Comment:
nit: IIUC now that these patterns match named functions too, should the name
group be `[^\s,\[\]]+` instead of `[a-zA-Z0-9\<\>_-]+` (here and in
`_SINGLE_NAME_PATTERN` / `_OBJECT_WITH_HASH_PATTERN`)? A nested function with a
valid non-ASCII name like `def α_fn` gets an identifier that can't be parsed
back, so `cloudpickle.loads` fails instead of rebuilding it (`\w` alone still
misses combining marks like `x̄`). Please verify the regex
##########
sdks/python/apache_beam/internal/code_object_pickler_test.py:
##########
@@ -76,6 +83,92 @@ def function_decorator(f):
return lambda x: f(f(x))
+def get_zero_arg_and_one_arg_lambdas():
+ return (lambda: 1, lambda x: 2)
+
+
+def get_same_co_code_lambdas():
+ tag_a = "train"
+ tag_b = "test"
+ return (lambda x: (tag_a, x), lambda x: (tag_b, x))
+
+
+def get_single_lambda_and_nested_function_with_lambda():
+ unused_outer_lambda = lambda x: x + 1
+
+ def inner_fn():
+ return lambda y: y + 2
+
+ return inner_fn()
+
+
+def shared_wrapper_decorator(fn):
+ @functools.wraps(fn)
+ def wrapper(*args, **kwargs):
+ return fn(*args, **kwargs)
+
+ return wrapper
+
+
+@shared_wrapper_decorator
+def outer_decorated_with_shared_wrapper():
+ @shared_wrapper_decorator
+ def inner_decorated():
+ return 1
+
+ return inner_decorated
+
+
+def get_same_named_nested_functions_different_args(branch):
+ if branch:
+
+ def nested_fn(x):
+ return x + 1
+ else:
+
+ def nested_fn(x, y):
+ return x + y
+
+ return nested_fn
+
+
+def get_same_named_nested_functions_same_args(branch):
+ if branch:
+
+ def nested_fn(x):
+ return x + 1
+ else:
+
+ def nested_fn(x):
+ return x + 2
+
+ return nested_fn
+
+
+def get_same_named_nested_functions_with_complex_consts(branch):
+ if branch:
+
+ def nested_fn(x):
+ def inner(x):
Review Comment:
nit: `inner` takes `x` here but is called as `inner()`, and the other branch
defines `def inner():`. Should this be `def inner():` so the two branches
differ only in their constants?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
