Scope RunInference model_identifier annotation to BeamML_RunInference

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

@@ -1607,7 +1607,6 @@ def annotations(self):
         'model_handler_type': (
             f'{self._model_handler.__class__.__module__}'
             f'.{self._model_handler.__class__.__qualname__}'),
-        'model_identifier': self._model_tag,
         **extra,
         **super().annotations()
     }
@@ -1715,47 +1714,54 @@ def expand(
         ) if self._model_metadata_pcoll else None).with_resource_hints(
             **resource_hints)
 
-    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 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)
 
     if self._monitoring_transform is not None:
-      with results.pipeline.transform_annotations(model_identifier=''):
-        _ = (
-            results
-            | 'BeamML_RunInference_MonitoringOutlet' >>
-            self._monitoring_transform)
+      _ = (
+          results
+          |
+          'BeamML_RunInference_MonitoringOutlet' >> self._monitoring_transform)
 
     results, bad_postprocessed = self._apply_fns(
       results, postprocess_fns, 'BeamML_RunInference_Postprocess')

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

@@ -2219,6 +2219,25 @@ def test_run_inference_watch_file_pattern_keyword_arg_side_input_label(self):
         FakeModelHandler(), watch_model_pattern='fake/path/*')
     assert side_input_str in str(result_pcoll.producer.side_inputs[0])
 
+  def test_run_inference_only_attaches_model_identifier_to_inference_step(self):
+    pipeline = beam.Pipeline()
+    from apache_beam.ml.inference.utils import WatchFilePattern
+    file_pattern_side_input = (
+        pipeline
+        | 'WatchFilePattern' >> WatchFilePattern(file_pattern='fake/path/*'))
+    pcoll = pipeline | 'start' >> beam.Create([1, 2, 3])
+    _ = pcoll | base.RunInference(
+        FakeModelHandler(),
+        model_metadata_pcoll=file_pattern_side_input,
+        model_identifier='my_custom_model_id')
+    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'])
+
   def test_model_status_provides_valid_tags(self):
     ms = base._ModelStatus(False)
 

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

@@ -347,7 +347,6 @@ def __init__(
 
   def annotations(self) -> dict[str, Any]:
     return {
-        'model_identifier': '',
         **super().annotations(),
     }
 

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

@@ -542,8 +542,7 @@ def test_annotations_shadow_model_identifier(self):
         bigquery_table=self.bq_table,
     )
     annotations = transform.annotations()
-    self.assertIn("model_identifier", annotations)
-    self.assertEqual(annotations["model_identifier"], "")
+    self.assertNotIn("model_identifier", annotations)
 
   def test_run_inference_monitoring_outlet_shadows_model_identifier(self):
     class DummyMonitoring(beam.PTransform):
@@ -572,7 +571,7 @@ def load_model(self):
     ]
     self.assertTrue(len(outlet_transforms) > 0)
     for t in outlet_transforms:
-      self.assertEqual(t.annotations.get("model_identifier"), b"")
+      self.assertFalse(t.annotations.get("model_identifier"))
 
   @mock.patch(
       "apache_beam.ml.inference.vertex_ai_model_monitoring_v2.WriteToBigQuery")