praneetnadella commented on code in PR #40711:
URL: https://github.com/apache/beam/pull/40711#discussion_r4462220244
##########
sdks/python/apache_beam/internal/code_object_pickler.py:
##########
@@ -505,4 +526,12 @@ def _create_bytecode_hash(code_object: types.CodeType):
Returns:
The hash of the code object.
"""
- return hashlib.md5(code_object.co_code).hexdigest()
+ hash_input = repr((
+ code_object.co_code,
+ tuple(_normalize_const(c) for c in code_object.co_consts),
+ code_object.co_names,
+ code_object.co_varnames,
+ code_object.co_freevars,
+ code_object.co_cellvars,
+ )).encode('utf-8')
Review Comment:
nit: Should we also include `co.co_argcount`, `co.co_posonlyargcount`,
`co.co_kwonlyargcount`, and `co.co_flags` in `hash_input`? Two closures with
the same parameter names and body that differ only in signature kind (like
`lambda x, y: (x, y)` vs `lambda x, *y: (x, y)`) currently produce the same
`_signature` and bytecode hash.
##########
sdks/python/apache_beam/internal/code_object_pickler.py:
##########
@@ -496,6 +504,19 @@ def _get_arg_count(obj: types.CodeType):
)
+def _normalize_const(const):
+ """Normalizes a constant from co_consts for deterministic hashing."""
+ if inspect.iscode(const):
+ return ('__code__', _create_bytecode_hash(const))
+ if isinstance(const, tuple):
+ return tuple(_normalize_const(x) for x in const)
+ if isinstance(const, frozenset):
+ return (
Review Comment:
Should we sort normalized `frozenset` elements into a tuple in
`_normalize_const` (e.g., `('frozenset', tuple(sorted((_normalize_const(x) for
x in const), key=repr)))`)? Because `_normalize_const` converts elements into
tuples of strings and `_create_bytecode_hash` hashes `repr(normalized_consts)`,
`frozenset` iteration order changes across processes with different
`PYTHONHASHSEED` values (even for integer set literals like `x in {1, 2}`).
##########
sdks/python/apache_beam/internal/code_object_pickler_test.py:
##########
@@ -214,6 +362,146 @@ 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:
+ with self.assertRaisesRegex(AttributeError,
+ "Could not find code object with path"):
+ code_object_pickler.get_code_from_identifier(mod_name)
Review Comment:
nit: Should `test_get_code_from_identifier_imports_unloaded_module` pass a
full code identifier like `f"{mod_name}.my_function.__code__"` and assert that
it returns `module_1.my_function.__code__`? Passing bare `mod_name` stops at
the `ModuleType` and raises `AttributeError`, so the test only checks the error
branch (just remember to also restore `test_data.module_1 = saved` in `finally`
if re-imported).
##########
sdks/python/apache_beam/internal/code_object_pickler.py:
##########
@@ -435,19 +430,32 @@ 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:
+ obj = None
+ start_index = 0
+ for i in range(len(parts), 0, -1):
+ module_name = '.'.join(parts[:i])
+ if module_name in sys.modules and sys.modules[module_name] is not None:
+ obj = sys.modules[module_name]
+ start_index = i
+ break
+ try:
+ obj = importlib.import_module(module_name)
Review Comment:
nit: Should we cap the reverse `import_module` search before `__code__` (or
`__defaults__`) segments? Right now, looking up an identifier like
`...module_1.my_function.__code__.co_consts[<lambda>]` calls `import_module` on
`.co_consts[<lambda>]`, `.__code__`, and `.my_function` before reaching
`module_1` in `sys.modules`.
--
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]
