Scope RunInference model_identifier annotation and preserve ModelManager metric_tag across model updates

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

@@ -1539,7 +1539,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()
     }
@@ -1647,47 +1646,48 @@ 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)
+    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')
@@ -2010,10 +2010,12 @@ def __init__(
       self,
       models: Union[list[Any], ModelManager],
       model_tag: str,
-      loader_func: Optional[Callable[[], Any]] = None):
+      loader_func: Optional[Callable[[], Any]] = None,
+      metric_tag: Optional[str] = None):
     self.models = models
     self.use_model_manager = not isinstance(models, list)
     self.model_tag = model_tag
+    self.metric_tag = metric_tag or model_tag
     self.loader_func = loader_func
     if not self.use_model_manager and len(models) > 1:
       self.model_router = multi_process_shared.MultiProcessShared(
@@ -2024,7 +2026,8 @@ def __init__(
   def next_model(self):
     if self.use_model_manager:
       loader_wrapper = _ProxyLoader(self.loader_func, self.model_tag)
-      return self.models.acquire_model(self.model_tag, loader_wrapper)
+      return self.models.acquire_model(
+          self.model_tag, loader_wrapper, metric_tag=self.metric_tag)
 
     if len(self.models) == 1:
       # Short circuit if there's no routing strategy needed in order to
@@ -2131,7 +2134,10 @@ def load():
           tag='model_manager',
           always_proxy=True).acquire()
       model_wrapper = _SharedModelWrapper(
-          model_manager, self._cur_tag, self._model_handler.load_model)
+          model_manager,
+          self._cur_tag,
+          self._model_handler.load_model,
+          metric_tag=self._model_tag)
     elif self._model_handler.share_model_across_processes():
       models = []
       for copy_tag in _get_tags_for_copies(self._cur_tag,

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

@@ -2074,6 +2074,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)
 
@@ -2818,6 +2837,7 @@ def test_run_inference_dofn_model_update_acquires_and_releases_matching_tag(
       # Second batch on 'model_v1' reuses the 'model_v1' instance (3, 4)
       self.assertEqual(dofn.process([0, 0], None, meta_v1), [3, 4])
       self.assertEqual(dofn._cur_tag, 'model_v1')
+      self.assertEqual(dofn._model.metric_tag, 'initial_tag')
       self.assertEqual(
           len(dofn._model.all_models()), 1, 'Expected 1 model for model_v1')
 
@@ -2827,6 +2847,7 @@ def test_run_inference_dofn_model_update_acquires_and_releases_matching_tag(
       # Second batch on 'model_v2' reuses the 'model_v2' instance (3, 4)
       self.assertEqual(dofn.process([0, 0], None, meta_v2), [3, 4])
       self.assertEqual(dofn._cur_tag, 'model_v2')
+      self.assertEqual(dofn._model.metric_tag, 'initial_tag')
       self.assertEqual(
           len(dofn._model.all_models()), 1, 'Expected 1 model for model_v2')
     finally:

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

@@ -185,6 +185,8 @@ def __init__(
     self.min_data_points = min_data_points
     self.verbose_logging = verbose_logging
     self.estimates: dict[str, float] = {}
+    # Maps a model tag to the metric tags its metrics are reported under.
+    self._metric_tags: dict[str, set[str]] = defaultdict(set)
     self.history = defaultdict(lambda: deque(maxlen=20))
     self.known_models = set()
     self._lock = threading.Lock()
@@ -193,6 +195,18 @@ def logging_info(self, message: str, *args):
     if self.verbose_logging:
       logger.info(message, *args)
 
+  def set_metric_tag(self, model_tag: str, metric_tag: str):
+    with self._lock:
+      self._metric_tags[model_tag].add(metric_tag)
+
+  # Should hold _lock when calling
+  def _report_estimate(self, model_tag: str, cost: float):
+    for metric_tag in self._metric_tags.get(model_tag) or {model_tag}:
+      Metrics.distribution(
+          _MODEL_MANAGER_METRICS_NAMESPACE,
+          f"memory_estimate_mb_{metric_tag}",
+          process_wide=True).update(int(cost))
+
   def is_unknown(self, model_tag: str) -> bool:
     with self._lock:
       return model_tag not in self.estimates
@@ -205,10 +219,7 @@ def set_initial_estimate(self, model_tag: str, cost: float):
     with self._lock:
       self.estimates[model_tag] = cost
       self.known_models.add(model_tag)
-      Metrics.distribution(
-          _MODEL_MANAGER_METRICS_NAMESPACE,
-          f"memory_estimate_mb_{model_tag}",
-          process_wide=True).update(int(cost))
+      self._report_estimate(model_tag, cost)
       self.logging_info("Initial Profile for %s: %s MB", model_tag, cost)
 
   def add_observation(
@@ -296,10 +307,7 @@ def _solve(self):
         self.logging_info(
             "Updated Estimate for %s: %.1f MB", model, self.estimates[model])
 
-        Metrics.distribution(
-            _MODEL_MANAGER_METRICS_NAMESPACE,
-            f"memory_estimate_mb_{model}",
-            process_wide=True).update(int(self.estimates[model]))
+        self._report_estimate(model, self.estimates[model])
       self.logging_info("System Bias: %s MB", bias)
 
     except Exception as e:
@@ -358,6 +366,10 @@ def __init__(
 
     # Resource State
     self._models = defaultdict(list)
+    # Maps a stable metric_tag (e.g. model_identifier) to the currently active
+    # model tag so metrics continue to be emitted under metric_tag across
+    # runtime side-input model updates.
+    self._metric_tag_to_active_tag: dict[str, str] = {}
     # Idle LRU used to track released models that
     # can be freed or reused upon request.
     self._idle_lru = OrderedDict()
@@ -384,17 +396,17 @@ def __init__(
     self._monitor.start()
 
   def _update_model_count_metric(self):
-    for tag, instances in self._models.items():
+    for metric_tag, active_tag in self._metric_tag_to_active_tag.items():
       Metrics.distribution(
           _MODEL_MANAGER_METRICS_NAMESPACE,
-          f"num_loaded_models_{tag}",
-          process_wide=True).update(len(instances))
+          f"num_loaded_models_{metric_tag}",
+          process_wide=True).update(len(self._models.get(active_tag, ())))
 
   def _clear_all_model_metrics(self):
-    for tag in self._models:
+    for metric_tag in self._metric_tag_to_active_tag:
       Metrics.distribution(
           _MODEL_MANAGER_METRICS_NAMESPACE,
-          f"num_loaded_models_{tag}",
+          f"num_loaded_models_{metric_tag}",
           process_wide=True).update(0)
 
   def logging_info(self, message: str, *args):
@@ -488,12 +500,23 @@ def _wait_in_queue(self, ticket: QueueTicket):
     finally:
       self._cv.acquire()
 
-  def acquire_model(self, tag: str, loader_func: Callable[[], Any]) -> Any:
+  def acquire_model(
+      self,
+      tag: str,
+      loader_func: Callable[[], Any],
+      metric_tag: Optional[str] = None) -> Any:
+    effective_metric_tag = metric_tag or tag
+    self._estimator.set_metric_tag(tag, effective_metric_tag)
     current_priority = 0 if self._estimator.is_unknown(tag) else 1
     ticket_num = next(self._ticket_counter)
     my_ticket = QueueTicket(current_priority, ticket_num, tag)
 
     with self._cv:
+      if self._metric_tag_to_active_tag.get(effective_metric_tag) != tag:
+        self._metric_tag_to_active_tag[effective_metric_tag] = tag
+        if self._models.get(tag):
+          self._update_model_count_metric()
+
       # FAST PATH: Grab from idle LRU if available
       if not self._isolation_mode:
         cached_instance = self._try_grab_from_lru(tag)
@@ -803,6 +826,7 @@ def force_reset(self):
     logger.warning("Force Reset Triggered")
     self._delete_all_models()
     self._models = defaultdict(list)
+    self._metric_tag_to_active_tag = {}
     self._idle_lru = OrderedDict()
     self._active_counts = Counter()
     self._wait_queue = []

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

@@ -401,6 +401,48 @@ def _get_est_dist_mean(tag):
     self.manager.force_reset()
     self.assertEqual(_get_count_dist_max(tag1), 2)
 
+  def test_model_manager_metrics_with_metric_tag(self):
+    """Test that metrics only report under metric_tag across version updates."""
+    stable_id = "stable_model_identifier"
+    v1_tag = "gs://bucket/model_v1"
+    v2_tag = "gs://bucket/model_v2"
+
+    def _get_count_dist(tag):
+      dist = MetricsEnvironment.process_wide_container().get_distribution(
+          MetricName('BeamML_ModelManager', f'num_loaded_models_{tag}'))
+      return dist.get_cumulative()
+
+    def _get_est_dist(tag):
+      dist = MetricsEnvironment.process_wide_container().get_distribution(
+          MetricName('BeamML_ModelManager', f'memory_estimate_mb_{tag}'))
+      return dist.get_cumulative()
+
+    # Acquire v1 with metric_tag=stable_id
+    m1 = self.manager.acquire_model(
+        v1_tag, lambda: MockModel(v1_tag, 1000.0, self.mock_monitor),
+        metric_tag=stable_id)
+    self.manager.release_model(v1_tag, m1)
+    self.assertEqual(_get_count_dist(stable_id).max, 1)
+    self.assertEqual(_get_est_dist(stable_id).max, 1000)
+    self.assertEqual(_get_count_dist(v1_tag).count, 0)
+    self.assertEqual(_get_est_dist(v1_tag).count, 0)
+
+    # Acquire 2 instances of v2 with metric_tag=stable_id
+    self.manager._estimator.set_metric_tag(v2_tag, stable_id)
+    self.manager._estimator.set_initial_estimate(v2_tag, 2000.0)
+    m2_a = self.manager.acquire_model(
+        v2_tag, lambda: MockModel(v2_tag, 2000.0, self.mock_monitor),
+        metric_tag=stable_id)
+    m2_b = self.manager.acquire_model(
+        v2_tag, lambda: MockModel(v2_tag, 2000.0, self.mock_monitor),
+        metric_tag=stable_id)
+    self.manager.release_model(v2_tag, m2_a)
+    self.manager.release_model(v2_tag, m2_b)
+    self.assertEqual(_get_count_dist(stable_id).max, 2)
+    self.assertEqual(_get_est_dist(stable_id).max, 2000)
+    self.assertEqual(_get_count_dist(v2_tag).count, 0)
+    self.assertEqual(_get_est_dist(v2_tag).count, 0)
+
   def test_single_model_convergence_with_fluctuations(self):
     """
     Tests that the estimator converges to the true usage with fluctuations.

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

@@ -572,7 +572,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")