Attach RunInference model_identifier only to the inference ParDo with exception handling

sdks/python/apache_beam/ml/inference/base.py

@@ -1745,48 +1745,48 @@ def expand(
         ) if self._model_metadata_pcoll else None).with_resource_hints(
             **resource_hints)
 
-    # Attach model_identifier only to the inference step so that runners can
-    # associate it with the ParDo that actually loads and runs the model.
-    # TODO(https://github.com/apache/beam/issues/40715): When
-    # with_exception_handling(threshold<1) is used, the threshold-checking
-    # transforms added by _ExceptionHandlingWrapper also inherit this
-    # annotation. Refactor to annotate only the inner inference ParDo.
-    with batched_elements_pcoll.pipeline.transform_annotations(
-        model_identifier=self._model_tag):
-      if self._with_exception_handling:
-        # On timeouts, report back to the central model metadata
-        # that the model is invalid
-        model_tag = self._model_tag
-        share_across_processes = (
-            self._model_handler.share_model_across_processes())
-        timeout = self._timeout
-
-        def failure_callback(exception: Exception, element: Any):
-          if type(exception) is not TimeoutError:
-            return
-          model_metadata = load_model_status(model_tag, share_across_processes)
-          model_metadata.try_mark_current_model_invalid(timeout)
-          logging.warning(
-              "Inference failed with a timeout, marking the current " +
-              "model for garbage collection")
-
-        callback = None
-        if (self._timeout is not None and
-            self._model_handler.should_garbage_collect_on_timeout()):
-          callback = failure_callback
-        results, bad_inference = (
-            batched_elements_pcoll
-            | 'BeamML_RunInference' >>
-            run_inference_pardo.with_exception_handling(
-            exc_class=self._exc_class,
-            use_subprocess=self._use_subprocess,
-            threshold=self._threshold,
-            timeout = self._timeout,
-            on_failure_callback=callback))
-      else:
-        results = (
-            batched_elements_pcoll
-            | 'BeamML_RunInference' >> run_inference_pardo)
+    # Attach model_identifier only to the inference ParDo (also when it is
+    # wrapped by with_exception_handling) so that runners can associate it with
+    # the step that actually loads and runs the model.
+    default_annotations = run_inference_pardo.annotations
+    run_inference_pardo.annotations = lambda: {
+        **default_annotations(), 'model_identifier': self._model_tag
+    }
+
+    if self._with_exception_handling:
+      # On timeouts, report back to the central model metadata
+      # that the model is invalid
+      model_tag = self._model_tag
+      share_across_processes = (
+          self._model_handler.share_model_across_processes())
+      timeout = self._timeout
+
+      def failure_callback(exception: Exception, element: Any):
+        if type(exception) is not TimeoutError:
+          return
+        model_metadata = load_model_status(model_tag, share_across_processes)
+        model_metadata.try_mark_current_model_invalid(timeout)
+        logging.warning(
+            "Inference failed with a timeout, marking the current " +
+            "model for garbage collection")
+
+      callback = None
+      if (self._timeout is not None and
+          self._model_handler.should_garbage_collect_on_timeout()):
+        callback = failure_callback
+      results, bad_inference = (
+          batched_elements_pcoll
+          | 'BeamML_RunInference' >>
+          run_inference_pardo.with_exception_handling(
+          exc_class=self._exc_class,
+          use_subprocess=self._use_subprocess,
+          threshold=self._threshold,
+          timeout = self._timeout,
+          on_failure_callback=callback))
+    else:
+      results = (
+          batched_elements_pcoll
+          | 'BeamML_RunInference' >> run_inference_pardo)
 
     if self._monitoring_transform is not None:
       _ = (

sdks/python/apache_beam/ml/inference/base_test.py

@@ -2256,6 +2256,31 @@ def test_run_inference_only_attaches_model_identifier_to_inference_step(self):
     self.assertEqual(
         transforms_with_model_id, ['RunInference/BeamML_RunInference'])
 
+  def test_run_inference_with_exception_handling_only_annotates_inference_step(
+      self):
+    pipeline = beam.Pipeline()
+    pcoll = pipeline | 'start' >> beam.Create([1, 2, 3])
+    _ = pcoll | base.RunInference(
+        FakeModelHandler(),
+        model_identifier='my_custom_model_id').with_exception_handling(
+            threshold=0.5)
+    proto = pipeline.to_runner_api()
+    transforms_with_model_id = [
+        t.unique_name for t in proto.components.transforms.values()
+        if t.annotations.get('model_identifier') == b'my_custom_model_id'
+    ]
+    self.assertEqual(
+        transforms_with_model_id,
+        [
+            'RunInference/BeamML_RunInference/'
+            'ParDo(_RunInferenceDoFn)/ParDo(_RunInferenceDoFn)'
+        ])
+    # The threshold-checking transforms exist but are not annotated.
+    self.assertTrue(
+        any(
+            'CountBad' in t.unique_name
+            for t in proto.components.transforms.values()))
+
   def test_model_status_provides_valid_tags(self):
     ms = base._ModelStatus(False)
 

sdks/python/apache_beam/transforms/core.py

@@ -1688,7 +1688,8 @@ def with_exception_handling(
         on_failure_callback,
         allow_unsafe_userstate_in_process,
         self.get_resource_hints(),
-        self.get_type_hints())
+        self.get_type_hints(),
+        self.annotations())
 
   def with_error_handler(self, error_handler, **exception_handling_kwargs):
     """An alias for `with_exception_handling(error_handler=error_handler, ...)`
@@ -2351,7 +2352,8 @@ def __init__(
       on_failure_callback,
       allow_unsafe_userstate_in_process,
       resource_hints,
-      pardo_type_hints=None):
+      pardo_type_hints=None,
+      pardo_annotations=None):
     if partial and use_subprocess:
       raise ValueError('partial and use_subprocess are mutually incompatible.')
     self._fn = fn
@@ -2370,6 +2372,7 @@ def __init__(
     self._allow_unsafe_userstate_in_process = allow_unsafe_userstate_in_process
     self._resource_hints = resource_hints
     self._pardo_type_hints = pardo_type_hints
+    self._pardo_annotations = pardo_annotations
     self._extra_tags = None
 
   def with_outputs(self, *tags, main=None):
@@ -2407,6 +2410,13 @@ def _build_pardo(self, pcoll):
         **self._kwargs,
     )
     pardo.get_resource_hints().update(self._resource_hints)
+    if self._pardo_annotations:
+      # Attach to the inner ParDo only, not to the threshold-checking
+      # transforms added in _post_process_result.
+      default_annotations = pardo.annotations
+      pardo.annotations = lambda: {
+          **default_annotations(), **self._pardo_annotations
+      }
     return pardo
 
   def _post_process_result(self, pcoll, result):

sdks/python/apache_beam/transforms/core_test.py

@@ -430,6 +430,24 @@ def test_timer_exception_handling(self):
       assert_that(good, equal_to([0, 1, 2]), 'good')
       assert_that(bad_elements, equal_to([(1, 5), (1, 10)]), 'bad')
 
+  def test_pardo_annotations_only_on_inner_pardo_with_threshold(self):
+    pardo = beam.ParDo(TestDoFn9())
+    default_annotations = pardo.annotations
+    pardo.annotations = lambda: {**default_annotations(), 'my_key': 'my_value'}
+    pipeline = beam.Pipeline()
+    _ = (
+        pipeline
+        | beam.Create([1, 2, 3])
+        | 'Handled' >> pardo.with_exception_handling(threshold=0.5))
+    transforms = pipeline.to_runner_api().components.transforms.values()
+    annotated = [
+        t.unique_name for t in transforms
+        if t.annotations.get('my_key') == b'my_value'
+    ]
+    self.assertEqual(len(annotated), 1, annotated)
+    self.assertTrue(annotated[0].startswith('Handled/'), annotated)
+    self.assertTrue(any('CountBad' in t.unique_name for t in transforms))
+
   def test_tags_with_exception_handling_then_resource_hint(self):
     class TagHint(ResourceHint):
       urn = 'beam:resources:tags:v1'