[Brijesh-Thakkar]

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.
• Replace these instructions with a description of your change and why you have made it.
• Mention the appropriate issue in your description (for example: addresses #123 or fixes #456).
• Update CHANGES.md with noteworthy changes.

See CONTRIBUTING.md for more details.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[github-actions]

📋 Review Summary

This PR resolves #40715 by attaching the model_identifier annotation directly to the run_inference_pardo instance and forwarding ParDo.annotations() to the inner ParDo inside _ExceptionHandlingWrapper._build_pardo. The changes are clean, well-scoped, and properly covered by unit tests.

🔍 General Feedback

• Forwarding ParDo.annotations() to the inner ParDo in ParDo.with_exception_handling / _ExceptionHandlingWrapper._build_pardo has clean semantics consistent with how resource_hints and type_hints are already forwarded, while avoiding leaking ParDo-specific annotations onto the auxiliary threshold-checking transforms or outer composite wrappers.
• Replacing the pipeline.transform_annotations(...) context manager in RunInference.expand() with an instance-level annotations override follows the established pattern used elsewhere in the SDK (e.g., annotate_yaml).

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

[github-actions]

Assigning reviewers:

R: @claudevdm for label python.
Backup reviewer: shunping (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

This pull request likely touches a core component ("core" label). Please review with scrutiny.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).