Preserve ModelManager metric_tag across RunInference side-input model updates

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

@@ -2073,15 +2073,31 @@ class _SharedModelWrapper():
 
     This allows us to round robin calls to models sitting in different
     processes so that we can more efficiently use resources (e.g. GPUs).
+
+    Args:
+      models: Either a list of loaded models or a ModelManager.
+      model_tag: Tag identifying the model version currently in use. When a
+        model is updated via side input this is the new model path, so the
+        ModelManager caches and evicts each version independently.
+      loader_func: Function used by ModelManager to load a new instance.
+      metric_tag: Stable tag (the RunInference model_identifier) that
+        ModelManager reports num_loaded_models_* and memory_estimate_mb_*
+        metrics under, so the metric name does not change when the model
+        version changes. Defaults to model_tag. Since these metrics span model
+        versions, runtime model updates should ideally only change weights
+        rather than swap to a model with a different memory footprint, which
+        would skew the reported estimates.
   """
   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(
@@ -2092,7 +2108,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
@@ -2201,7 +2218,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

@@ -3101,6 +3101,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')
 
@@ -3110,6 +3111,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

@@ -180,19 +180,40 @@ def __init__(
       self,
       smoothing_factor: float = 0.2,
       min_data_points: int = 5,
-      verbose_logging: bool = False):
+      verbose_logging: bool = False,
+      metric_tags_fn: Optional[Callable[[str], list[str]]] = None):
+    """
+    Args:
+      metric_tags_fn: Resolves a model tag to the metric tags its
+        memory_estimate_mb_* metric is reported under. Defaults to the model
+        tag itself.
+    """
     self.smoothing_factor = smoothing_factor
     self.min_data_points = min_data_points
     self.verbose_logging = verbose_logging
     self.estimates: dict[str, float] = {}
     self.history = defaultdict(lambda: deque(maxlen=20))
     self.known_models = set()
     self._lock = threading.Lock()
+    self._metric_tags_fn = metric_tags_fn or (lambda model_tag: [model_tag])
 
   def logging_info(self, message: str, *args):
     if self.verbose_logging:
       logger.info(message, *args)
 
+  def report_estimate(self, model_tag: str):
+    with self._lock:
+      if model_tag in self.estimates:
+        self._report_estimate(model_tag, self.estimates[model_tag])
+
+  # Should hold _lock when calling
+  def _report_estimate(self, model_tag: str, cost: float):
+    for metric_tag in self._metric_tags_fn(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 +226,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 +314,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:
@@ -345,7 +360,8 @@ def __init__(
     self._estimator = ResourceEstimator(
         min_data_points=min_data_points,
         smoothing_factor=smoothing_factor,
-        verbose_logging=verbose_logging)
+        verbose_logging=verbose_logging,
+        metric_tags_fn=self._metric_tags_for)
     self._monitor = monitor if monitor else GPUMonitor(
         poll_interval=poll_interval, peak_window_seconds=peak_window_seconds)
     self._slack_percentage = slack_percentage
@@ -358,6 +374,17 @@ def __init__(
 
     # Resource State
     self._models = defaultdict(list)
+    # Maps a stable metric_tag (e.g. model_identifier) to the model tag that is
+    # currently active for it. Metrics for the active tag are reported under
+    # metric_tag so metric names stay stable across runtime side-input model
+    # updates; superseded tags that are still loaded report under their own
+    # tag while they linger in the LRU.
+    # NOTE: Because metrics under metric_tag are cumulative distributions that
+    # span model versions, a runtime model update should ideally only change
+    # weights (same architecture / memory footprint). Swapping to a model of a
+    # different size blends the two versions into one memory_estimate_mb_* and
+    # num_loaded_models_* series and skews the reported estimates.
+    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()
@@ -383,24 +410,52 @@ def __init__(
 
     self._monitor.start()
 
+  def _metric_tags_for(self, tag: str) -> list[str]:
+    """Returns the metric tags that metrics for model tag are reported under."""
+    metric_tags = [
+        metric_tag
+        for metric_tag, active_tag in self._metric_tag_to_active_tag.items()
+        if active_tag == tag
+    ]
+    return metric_tags or [tag]
+
   def _update_model_count_metric(self):
     for tag, instances in self._models.items():
-      Metrics.distribution(
-          _MODEL_MANAGER_METRICS_NAMESPACE,
-          f"num_loaded_models_{tag}",
-          process_wide=True).update(len(instances))
+      for metric_tag in self._metric_tags_for(tag):
+        Metrics.distribution(
+            _MODEL_MANAGER_METRICS_NAMESPACE,
+            f"num_loaded_models_{metric_tag}",
+            process_wide=True).update(len(instances))
 
   def _clear_all_model_metrics(self):
     for tag in self._models:
-      Metrics.distribution(
-          _MODEL_MANAGER_METRICS_NAMESPACE,
-          f"num_loaded_models_{tag}",
-          process_wide=True).update(0)
+      for metric_tag in self._metric_tags_for(tag):
+        Metrics.distribution(
+            _MODEL_MANAGER_METRICS_NAMESPACE,
+            f"num_loaded_models_{metric_tag}",
+            process_wide=True).update(0)
 
   def logging_info(self, message: str, *args):
     if self._verbose_logging:
       logger.info(message, *args)
 
+  # Should hold _cv lock when calling
+  def _maybe_switch_active_tag(self, metric_tag: str, tag: str):
+    """Tracks which model tag is active for metric_tag.
+
+    When a different tag is requested for a metric_tag (e.g. a side-input
+    model update), metrics for metric_tag start reflecting the new tag. The
+    previous tag's instances stay in the LRU and are evicted on demand.
+    """
+    prev_tag = self._metric_tag_to_active_tag.get(metric_tag)
+    if prev_tag == tag:
+      return
+    self._metric_tag_to_active_tag[metric_tag] = tag
+    # Re-emit under metric_tag right away if this tag is already profiled.
+    self._estimator.report_estimate(tag)
+    if prev_tag is not None or self._models.get(tag):
+      self._update_model_count_metric()
+
   def all_models(self, tag) -> list[Any]:
     return self._models[tag]
 
@@ -488,12 +543,19 @@ 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
     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:
+      self._maybe_switch_active_tag(effective_metric_tag, tag)
+
       # FAST PATH: Grab from idle LRU if available
       if not self._isolation_mode:
         cached_instance = self._try_grab_from_lru(tag)
@@ -803,6 +865,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,58 @@ 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 for the active version report under metric_tag."""
+    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()
+
+    MetricsEnvironment.process_wide_container().reset()
+
+    # Acquire v1 with metric_tag=stable_id: reported under stable_id only.
+    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)
+
+    # Switch stable_id to v2 and load 2 copies: stable_id now tracks v2. v2 is
+    # pre-profiled so the switch does not enter isolation mode (which flushes
+    # all loaded models) and v1 keeps lingering in the LRU.
+    self.manager._estimator.set_initial_estimate(v2_tag, 2000.0)
+    MetricsEnvironment.process_wide_container().reset()
+    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)
+
+    # v1 is still loaded (idle in the LRU) and still known to the estimator
+    # for admission control; while it lingers it reports under its own tag.
+    self.assertFalse(m1.deleted)
+    self.assertFalse(self.manager._estimator.is_unknown(v1_tag))
+    self.assertEqual(_get_count_dist(v1_tag).max, 1)
+
   def test_single_model_convergence_with_fluctuations(self):
     """
     Tests that the estimator converges to the true usage with fluctuations.