Repository navigation
perf(ops): inject traceback argument info on error only, not on every call - #23183
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request optimizes Keras operations by moving traceback argument injection from an eager wrapper-based approach to a lazy, on-exception mechanism, thereby eliminating per-call overhead on the happy path. The review feedback identifies two key robustness issues on the error-handling path: first, inspect.signature can raise a ValueError for certain callables, which must be caught alongside TypeError to avoid masking the original exception; second, instantiating custom exceptions with the augmented message can raise exceptions other than TypeError, so catching all exceptions is recommended to safely fall back to a RuntimeError.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #23183 +/- ##
==========================================
- Coverage 80.46% 78.11% -2.36%
==========================================
Files 466 466
Lines 70017 70033 +16
Branches 11582 11585 +3
==========================================
- Hits 56339 54705 -1634
- Misses 10824 12471 +1647
- Partials 2854 2857 +3
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
789de50 to
0510fdb
Compare
0510fdb to
9eaf543
Compare
be82b6b to
95c54c3
Compare
In inject_argument_info_in_error, a failed attempt to set the _keras_call_info_injected flag on the augmented exception was silently swallowed, but the augmented (and now unmarked) exception was still returned and raised. For exception types that reject attribute assignment after construction, this meant the flag never actually stuck, so each enclosing Operation.__call__ frame in a nested call chain treated the exception as not-yet-augmented and wrapped it again, duplicating the "Exception encountered when calling ..." text once per nesting level. Change the except block to return None instead of pass when the flag assignment fails, so the caller falls back to raising the plain original exception instead of an unmarked, re-augmentable one. Added a regression test that raises a custom exception whose __setattr__ rejects new attributes after construction through two nested Operation subclasses' call() methods, and asserts the augmented message text does not appear more than once.
95c54c3 to
fb8f4c0
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request optimizes the traceback argument injection mechanism by moving the overhead from the happy path to the error path, replacing the function wrapper with a lazy exception-augmentation helper. The review feedback correctly identifies that passing the local call_fn helper to the new inject_argument_info_in_error function obscures the specific signature of the underlying call method (resulting in generic *args, **kwargs in error messages). To maintain detailed and informative error messages as required by the Keras API design guidelines, the reviewer suggests passing self.call instead of call_fn in both core.py and operation.py.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Cond.__call__ and Operation.__call__ passed a local call_fn closure (generic *args, **kwargs signature) into inject_argument_info_in_error, so error messages showed generic args=/kwargs= instead of the real parameter names. Pass self.call instead so the specific, user-facing signature is always inspected, including on the symbolic-call dispatch path.
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
Style: drop the module docstring, "# ---" section dividers, and numbered inline comment headers from traceback_utils_test.py (no other *_test.py file in keras/src/utils or keras/src/ops uses this decorative structure); restore an Args/Returns docstring on inject_argument_info_in_error, matching the pre-refactor docstring's level of detail and current house style (Args:, not the legacy Arguments:); remove a stale "Wrap self.call..." comment in Operation.__call__ left over from before this PR's refactor, which no longer described the code beneath it; and reword the duplicated "avoids creating a closure" comment in Cond.__call__, where it read as ambiguous next to the call_fn closure defined a few lines above it. Robustness: when args/kwargs can't be bound to the callee's real signature (the arg-count-mismatch edge case noted in review), fall back to a raw args=/kwargs= dump instead of silently dropping all argument context. This only runs on the already-slow exception path, so it adds no overhead to the hot no-exception path this PR exists to speed up. Tests: add coverage for the new fallback (bindable-with-fallback and still-nothing-to-report cases), and add a regression test for the tf.errors.OpError reconstruction branch, which previously had no direct test coverage.
…st.py _LockedError, InnerLockedLayer, and OuterLockedLayer had real logic (overridden __setattr__/call/compute_output_spec) but were defined inline inside a single test method, unlike every other helper class in this file. Move them to module level for consistency with the file's own convention and with the rnn_test.py precedent.
| if any_symbolic_tensors(args, kwargs): | ||
| return self.symbolic_call(*args, **kwargs) |
There was a problem hiding this comment.
This should be moved to def cond below like all the other ops.
Once, that's done, can we just remove __call__ here and just inherit the behavior from Operation for free instead?
hertschuh
added
stat:awaiting response from contributor
and removed
awaiting review
labels
- Remove redundant Cond.__call__; inherit error injection from Operation - Change inject_argument_info_in_error to always return an exception with the traceback already set, simplifying callers to a single raise - Shorten docstrings and remove references to previous implementation - Remove happy-path test covered by existing layer tests - Update unbindable-args test to expect original exception returned
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request refactors the traceback argument injection mechanism in Keras operations to be lazy, moving the overhead of argument formatting and exception wrapping from the hot path of every call to the error path only. It replaces inject_argument_info_in_traceback with inject_argument_info_in_error and updates Operation.__call__ to catch exceptions and inject argument info dynamically. Feedback on the changes highlights a potential TypeError in the fallback formatter when handling nested structures, and suggests a more robust fallback when formatting individual bound arguments fails instead of discarding the entire argument context.
| args_repr = ", ".join( | ||
| tree.map_structure(format_argument_value, arg) for arg in args | ||
| ) |
There was a problem hiding this comment.
In _format_raw_arguments_fallback, tree.map_structure(format_argument_value, arg) returns a structure (such as a list, tuple, or dict of strings) rather than a single string when arg is a nested structure. Passing this directly to ", ".join(...) will raise a TypeError (e.g., TypeError: sequence item 0: expected str instance, list found), causing the entire fallback block to fail and return None.
Wrapping the mapped structure in str() ensures it is always converted to a string representation before joining.
| args_repr = ", ".join( | |
| tree.map_structure(format_argument_value, arg) for arg in args | |
| ) | |
| args_repr = ", ".join( | |
| str(tree.map_structure(format_argument_value, arg)) for arg in args | |
| ) |
| try: | ||
| value = tree.map_structure( | ||
| format_argument_value, | ||
| bound_signature.arguments[arg.name], | ||
| ) | ||
| else: | ||
| value = arg.default | ||
| arguments_context.append(f" • {arg.name}={value}") | ||
| if arguments_context: | ||
| arguments_context = "\n".join(arguments_context) | ||
| # Get original error message and append information to it. | ||
| if tf_errors is not None and isinstance(e, tf_errors.OpError): | ||
| message = e.message | ||
| elif e.args: | ||
| # Canonically, the 1st argument in an exception is the error | ||
| # message. This works for all built-in Python exceptions. | ||
| message = e.args[0] | ||
| else: | ||
| message = "" | ||
| display_name = f"{object_name if object_name else fn.__name__}" | ||
| message = ( | ||
| f"Exception encountered when calling {display_name}.\n\n" | ||
| f"\x1b[1m{message}\x1b[0m\n\n" | ||
| f"Arguments received by {display_name}:\n" | ||
| f"{arguments_context}" | ||
| ) | ||
|
|
||
| # Reraise exception, with added context | ||
| if tf_errors is not None and isinstance(e, tf_errors.OpError): | ||
| new_e = e.__class__(e.node_def, e.op, message, e.error_code) | ||
| else: | ||
| try: | ||
| # For standard exceptions such as ValueError, TypeError, | ||
| # etc. | ||
| new_e = e.__class__(message) | ||
| except TypeError: | ||
| # For any custom error that doesn't have a standard | ||
| # signature. | ||
| new_e = RuntimeError(message) | ||
| new_e._keras_call_info_injected = True | ||
| except Exception: | ||
| # Formatting this specific bound value failed. Retrying | ||
| # via the raw args/kwargs fallback would call the same | ||
| # `format_argument_value` on the same value and fail the | ||
| # same way, so there is nothing to recover here. | ||
| return e |
There was a problem hiding this comment.
If formatting a specific bound argument fails (for example, if a custom or mock object raises an exception during is_tensor or shape/dtype inspection), the current implementation completely discards the traceback argument injection and returns the original exception e unmodified.
Instead of giving up on all arguments, we can fall back to repr(value) or a placeholder for that specific argument, allowing the successfully formatted arguments to still be displayed in the traceback.
| try: | |
| value = tree.map_structure( | |
| format_argument_value, | |
| bound_signature.arguments[arg.name], | |
| ) | |
| else: | |
| value = arg.default | |
| arguments_context.append(f" • {arg.name}={value}") | |
| if arguments_context: | |
| arguments_context = "\n".join(arguments_context) | |
| # Get original error message and append information to it. | |
| if tf_errors is not None and isinstance(e, tf_errors.OpError): | |
| message = e.message | |
| elif e.args: | |
| # Canonically, the 1st argument in an exception is the error | |
| # message. This works for all built-in Python exceptions. | |
| message = e.args[0] | |
| else: | |
| message = "" | |
| display_name = f"{object_name if object_name else fn.__name__}" | |
| message = ( | |
| f"Exception encountered when calling {display_name}.\n\n" | |
| f"\x1b[1m{message}\x1b[0m\n\n" | |
| f"Arguments received by {display_name}:\n" | |
| f"{arguments_context}" | |
| ) | |
| # Reraise exception, with added context | |
| if tf_errors is not None and isinstance(e, tf_errors.OpError): | |
| new_e = e.__class__(e.node_def, e.op, message, e.error_code) | |
| else: | |
| try: | |
| # For standard exceptions such as ValueError, TypeError, | |
| # etc. | |
| new_e = e.__class__(message) | |
| except TypeError: | |
| # For any custom error that doesn't have a standard | |
| # signature. | |
| new_e = RuntimeError(message) | |
| new_e._keras_call_info_injected = True | |
| except Exception: | |
| # Formatting this specific bound value failed. Retrying | |
| # via the raw args/kwargs fallback would call the same | |
| # `format_argument_value` on the same value and fail the | |
| # same way, so there is nothing to recover here. | |
| return e | |
| try: | |
| value = tree.map_structure( | |
| format_argument_value, | |
| bound_signature.arguments[arg.name], | |
| ) | |
| except Exception: | |
| try: | |
| value = repr(bound_signature.arguments[arg.name]) | |
| except Exception: | |
| value = "<unformattable>" |
hertschuh
left a comment
hertschuh
left a comment
There was a problem hiding this comment.
Sorry for the back and forth:
| @traceback_utils.filter_traceback | ||
| def __call__(self, *args, **kwargs): | ||
| def call_fn(*args, **kwargs): | ||
| if any_symbolic_tensors(args, kwargs): | ||
| return self.symbolic_call(*args, **kwargs) | ||
| else: | ||
| return self.call(*args, **kwargs) | ||
|
|
||
| if traceback_utils.is_traceback_filtering_enabled(): | ||
| # Wrap self.call to provide helpful info in case of exception | ||
| call_fn = traceback_utils.inject_argument_info_in_traceback( | ||
| call_fn, | ||
| object_name=(f"{self.__class__.__name__}.call()"), | ||
| ) | ||
| return call_fn(*args, **kwargs) | ||
|
|
||
| # Plain flow. | ||
| return call_fn(*args, **kwargs) |
There was a problem hiding this comment.
Oh, now I understand why __call__ was overriden. This is a special case (for which there is no unit test).
Can you redo your change?
`Cond.call` receives Python callables rather than tensors, which rematerialization cannot trace, so `Cond` needs its own `__call__` instead of inheriting `Operation`'s. Restore it, using the same on-error injection the rest of this change introduces, and cover it with a `RematScope` regression test on both scope modes. `_format_raw_arguments_fallback` joined the result of `tree.map_structure` directly, but that mirrors the input structure, so a nested argument came back as a list or dict and the join raised `TypeError`. The surrounding `except` swallowed it and the whole raw dump disappeared. Render each result before joining, and extend the fallback test with nested arguments, which the previous flat-scalar test could not catch. Also drop the comment above the injection try block in `Operation.__call__`, which described code that is no longer there, and the test docstring clause about third-party `call` overrides, which does not affect signature inspection.
hertschuh
left a comment
hertschuh
left a comment
There was a problem hiding this comment.
Thanks for the optimization!
google-ml-butler
Bot
added
kokoro:force-run
ready to pull
… call (keras-team#23183) * perf(ops): inject traceback argument info on error only, not on every call Part of keras-team#22561. * Fix duplicated error message when exception blocks attribute assignment In inject_argument_info_in_error, a failed attempt to set the _keras_call_info_injected flag on the augmented exception was silently swallowed, but the augmented (and now unmarked) exception was still returned and raised. For exception types that reject attribute assignment after construction, this meant the flag never actually stuck, so each enclosing Operation.__call__ frame in a nested call chain treated the exception as not-yet-augmented and wrapped it again, duplicating the "Exception encountered when calling ..." text once per nesting level. Change the except block to return None instead of pass when the flag assignment fails, so the caller falls back to raising the plain original exception instead of an unmarked, re-augmentable one. Added a regression test that raises a custom exception whose __setattr__ rejects new attributes after construction through two nested Operation subclasses' call() methods, and asserts the augmented message text does not appear more than once. * fix: use self.call's real signature for error-message argument injection Cond.__call__ and Operation.__call__ passed a local call_fn closure (generic *args, **kwargs signature) into inject_argument_info_in_error, so error messages showed generic args=/kwargs= instead of the real parameter names. Pass self.call instead so the specific, user-facing signature is always inspected, including on the symbolic-call dispatch path. * Clean up house-style deviations and restore lost error context Style: drop the module docstring, "# ---" section dividers, and numbered inline comment headers from traceback_utils_test.py (no other *_test.py file in keras/src/utils or keras/src/ops uses this decorative structure); restore an Args/Returns docstring on inject_argument_info_in_error, matching the pre-refactor docstring's level of detail and current house style (Args:, not the legacy Arguments:); remove a stale "Wrap self.call..." comment in Operation.__call__ left over from before this PR's refactor, which no longer described the code beneath it; and reword the duplicated "avoids creating a closure" comment in Cond.__call__, where it read as ambiguous next to the call_fn closure defined a few lines above it. Robustness: when args/kwargs can't be bound to the callee's real signature (the arg-count-mismatch edge case noted in review), fall back to a raw args=/kwargs= dump instead of silently dropping all argument context. This only runs on the already-slow exception path, so it adds no overhead to the hot no-exception path this PR exists to speed up. Tests: add coverage for the new fallback (bindable-with-fallback and still-nothing-to-report cases), and add a regression test for the tf.errors.OpError reconstruction branch, which previously had no direct test coverage. * Hoist test-local helper classes to module level in traceback_utils_test.py _LockedError, InnerLockedLayer, and OuterLockedLayer had real logic (overridden __setattr__/call/compute_output_spec) but were defined inline inside a single test method, unlike every other helper class in this file. Move them to module level for consistency with the file's own convention and with the rnn_test.py precedent. * refactor(ops): simplify traceback argument injection per review - Remove redundant Cond.__call__; inherit error injection from Operation - Change inject_argument_info_in_error to always return an exception with the traceback already set, simplifying callers to a single raise - Shorten docstrings and remove references to previous implementation - Remove happy-path test covered by existing layer tests - Update unbindable-args test to expect original exception returned * Keep Cond out of rematerialization and harden the raw-argument dump `Cond.call` receives Python callables rather than tensors, which rematerialization cannot trace, so `Cond` needs its own `__call__` instead of inheriting `Operation`'s. Restore it, using the same on-error injection the rest of this change introduces, and cover it with a `RematScope` regression test on both scope modes. `_format_raw_arguments_fallback` joined the result of `tree.map_structure` directly, but that mirrors the input structure, so a nested argument came back as a list or dict and the join raised `TypeError`. The surrounding `except` swallowed it and the whole raw dump disappeared. Render each result before joining, and extend the fallback test with nested arguments, which the previous flat-scalar test could not catch. Also drop the comment above the injection try block in `Operation.__call__`, which described code that is no longer there, and the test docstring clause about third-party `call` overrides, which does not affect signature inspection. * Drop the PR-context parenthetical from the test class docstring
Closes #23270
Summary
Operation.__call__andCond.__call__built a freshinject_argument_info_in_tracebackclosure (viafunctools.wraps) on every single call, even though it only ever did work on the exception path. This PR moves that wrapping into theexceptblock itself, so the happy path is just a plain call with no closure allocation.Why this is safe
inject_argument_info_in_errorhelper (intraceback_utils.py) produces the same augmented exception, double-injection guard (_keras_call_info_injected), and TFOpErrorhandling as before.signature.bind()cannot match args againstfn's real signature (e.g. a symbolic-path arg-count mismatch),_format_raw_arguments_fallbackdumps rawargs/kwargsrather than dropping the argument context. It is only reachable from inside theexceptblock, so it cannot affect the hot path._keras_call_info_injectedattribute set against slotted exception types, avoiding a secondaryAttributeError.inject_argument_info_in_tracebackis deleted; no callers remain.Benchmarks (GPU, 3-way: master / this PR alone / this PR + parents)
In isolation this lands in noise at the whole-model level, as expected — removing one per-call closure allocation is a small fraction of a forward pass. The primary evidence is the deterministic call-count reduction on the fast path itself. Stacked on #23182, LLM generate-32 crosses into a real, measured win (7/8 paired runs agree, ~8.5% off master past the noise threshold), showing the compounding effect of the two per-call fixes together. Outputs are bit-identical to master in both configs (max|Δ| = 0.0, generated token ids identical).
torch.compile performance
Compiled mode is slower than eager on every workload, in both states, on this hardware — this PR does not change that. That is not a regression this PR introduces: master itself is 1.58x–2.32x slower under
torch.compilethan eager on the same three workloads on this harness, so a compiled-slower-than-eager result is a pre-existing characteristic of these small synthetic benchmarks (short sequences, small batch, few iterations) undertorch.compile, not something this PR's change to the traceback path could plausibly move. The eager-vs-compiled ratios above are the only numbers that should be read across (within-run pairs); the isolated and stacked runs were captured in separate processes at different times under different ambient load, so raw absolute milliseconds should not be compared state-to-state — e.g. stacked eager LLM generate-32 (129.6 ms) reads higher than isolated eager LLM generate-32 (98.2 ms) here purely from run-to-run noise, not from adding a compile-neutral parent PR; that same magnitude of spread shows up on master's own repeated runs and on sibling PR #23182's compile runs. Equivalence checks against master (bit-exact outputs) passed for both states.Part of the perf series starting at #23182; recommended merge order #23182 → #23183 → #23184 → ... → #23190.
Tests
keras/src/utils/traceback_utils_test.py(11 tests): error message content, double-injection guard, signature-bind-failure fallback including nested arguments, and TFOpErrorpreservation.keras/src/ops/core_test.pyadds twoCondtests: errors useCond.call()'s real parameter names (pred/true_fn/false_fn), andcondkeeps working inside aRematScope, which covers whyCondoverrides__call__at all.Contributor Agreement
torch.compile impact
This PR is compile-neutral: deterministic dynamo decomposition shows a 0 graph-break delta at its position in the stack. That's expected — moving the traceback-injection closure into the
exceptblock only changes eager-mode dispatch and never touches the compiled graph itself. The compile-side win comes from removing the attention graph break in the torch backend (#23327), not from this PR.