AMOOOMA opened a new issue, #40715: URL: https://github.com/apache/beam/issues/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`: https://github.com/apache/beam/blob/5b6aa703f0ab6a48d92cc63ea89226485cc265f9/sdks/python/apache_beam/transforms/core.py#L2414 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 https://github.com/apache/beam/pull/40671#discussion_r4342559127. ### Issue Priority Priority: 3 (nice-to-have improvement) ### Issue Components - [X] 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 -- 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]
