[Feature Request]: Attach RunInference model_identifier annotation only to the inference ParDo when with_exception_handling(threshold<1) is used #40715

What would you like to happen?

RunInference.expand() scopes the model_identifier transform annotation to the BeamML_RunInference step via pipeline.transform_annotations(model_identifier=...) (see #40671). Because transform_annotations applies to every sub-transform applied inside the context, when with_exception_handling(threshold=<1.0) is used, the threshold-checking transforms that _ExceptionHandlingWrapper._post_process_result adds (CountTotal, CountBad, CombineGlobally, CheckThreshold, ...) also inherit the model_identifier annotation, even though they never run _RunInferenceDoFn:

This is harmless for Dataflow today (only ParDo steps consume the annotation and the threshold steps are tiny), but it is imprecise. We should refactor so the annotation is attached only to the inner inference ParDo, e.g.:

• expose a hook on _ExceptionHandlingWrapper / ParDo.with_exception_handling to annotate only the wrapped ParDo built in _build_pardo, or
• apply the annotation to run_inference_pardo directly (e.g. via PTransform.annotations() on a thin ParDo subclass) instead of using the pipeline-level transform_annotations context manager.

Follow-up from the discussion in #40671 (comment).

Issue Priority

Priority: 3 (nice-to-have improvement)

Issue Components

• Component: Python SDK
• Component: Java SDK
• Component: Go SDK
• Component: Typescript SDK
• Component: IO connector
• Component: Beam YAML
• Component: Beam examples
• Component: Beam playground
• Component: Beam katas
• Component: Website
• Component: Infrastructure
• Component: Spark Runner
• Component: Flink Runner
• Component: Samza Runner
• Component: Twister2 Runner
• Component: Hazelcast Jet Runner
• Component: Google Cloud Dataflow Runner