Brijesh-Thakkar opened a new pull request, #40922:
URL: https://github.com/apache/beam/pull/40922
Fixes #40715.
`RunInference.expand()` scoped the `model_identifier` annotation with
`pipeline.transform_annotations(...)` (added in #40671). That context applies
to everything applied inside it, so with `with_exception_handling(threshold<1)`
the annotation landed on all 25 annotated transforms: `CountTotal`, `CountBad`,
the `CombineGlobally` sub-steps, `Map(check_threshold)` and the wrapper
composites, none of which run `_RunInferenceDoFn`. Even with `threshold=1` the
outer composites were annotated.
Now exactly one transform, the inner inference ParDo
(`RunInference/BeamML_RunInference/ParDo(_RunInferenceDoFn)/ParDo(_RunInferenceDoFn)`),
carries `model_identifier`.
### Changes
- `ml/inference/base.py`: drop the `transform_annotations` context manager
and set `model_identifier` on `run_inference_pardo.annotations` instead (same
instance-override pattern as `annotate_yaml`). Most of the line count is the
dedent from removing the `with` block; `git diff -w` is ~16 lines.
- `transforms/core.py`: `ParDo.with_exception_handling` now forwards
`self.annotations()` to a new
`_ExceptionHandlingWrapper(pardo_annotations=...)` argument, and `_build_pardo`
applies them to the inner ParDo only (not to the threshold-checking transforms).
### Please weigh in
The `core.py` change is **not RunInference-specific**: any annotations set
on a `ParDo` are now carried onto the inner ParDo built by
`with_exception_handling`. Previously they were silently dropped, since the
wrapper built a fresh ParDo. I think that is the right semantics and avoids a
RunInference-only hook, but it touches a core path (`transforms/`), so I'd
appreciate reviewers' opinions. The alternative is an explicit
RunInference-only hook on the wrapper.
### Testing
-
`base_test.py::test_run_inference_with_exception_handling_only_annotates_inference_step`
(new): `RunInference(...).with_exception_handling(threshold=0.5)` annotates
exactly the inference ParDo, and the `CountBad` steps exist but are unannotated.
-
`core_test.py::ExceptionHandlingTest::test_pardo_annotations_only_on_inner_pardo_with_threshold`
(new): annotations set on a `ParDo` land on exactly one transform under
`with_exception_handling(threshold=0.5)`.
- Both new tests fail without the fix and pass with it.
- The existing
`test_run_inference_only_attaches_model_identifier_to_inference_step` (no
exception handling) is unchanged and still passes.
- `pytest apache_beam/ml/inference/base_test.py
apache_beam/transforms/core_test.py
apache_beam/transforms/error_handling_test.py
apache_beam/transforms/ptransform_test.py`: 365 passed, 11 skipped.
- `yapf --diff` clean on the changed files; `ruff` reports the same 5
pre-existing diagnostics as master, none in changed lines.
------------------------
Thank you for your contribution! Follow this checklist to help us
incorporate your contribution quickly and easily:
- [ ] If this contribution is large, please file an Apache [Individual
Contributor License Agreement](https://www.apache.org/licenses/icla.pdf).
- [x] Replace these instructions with a description of your change and why
you have made it.
- [x] Mention the appropriate issue in your description (for example:
`addresses #123` or `fixes #456`).
- [ ] Update `CHANGES.md` with noteworthy changes.
See
[CONTRIBUTING.md](https://github.com/apache/beam/blob/master/CONTRIBUTING.md)
for more details.
--
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]
