From 5209774cc1fda288e20d4e76a2b5943cb1424d09 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Wed, 21 Oct 2026 21:18:53 +0000 Subject: [PATCH 1/4] Preserve ModelManager metric_tag across RunInference side-input model updates --- sdks/python/apache_beam/ml/inference/base.py | 12 +++-- .../apache_beam/ml/inference/base_test.py | 2 + .../apache_beam/ml/inference/model_manager.py | 52 ++++++++++++++----- .../ml/inference/model_manager_test.py | 42 +++++++++++++++ 4 files changed, 91 insertions(+), 17 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 4c2d9cb4c80e..2780f616c8be 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -2078,10 +2078,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( @@ -2092,7 +2094,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 +2204,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, diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index 5d6c5bd517c7..5554a4120652 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/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: diff --git a/sdks/python/apache_beam/ml/inference/model_manager.py b/sdks/python/apache_beam/ml/inference/model_manager.py index bc16feb6f60f..c7d2bee8956f 100644 --- a/sdks/python/apache_beam/ml/inference/model_manager.py +++ b/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 = [] diff --git a/sdks/python/apache_beam/ml/inference/model_manager_test.py b/sdks/python/apache_beam/ml/inference/model_manager_test.py index d1f2b35bc1d9..9c4fd46856e0 100644 --- a/sdks/python/apache_beam/ml/inference/model_manager_test.py +++ b/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. From edad87e4155c3c29f7dab4acaaae25e9aa04580b Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Wed, 21 Oct 2026 23:09:43 +0000 Subject: [PATCH 2/4] Trigger CI rerun From 092e30e5ef367beb836be90bc8cda3c8dc2ddf18 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Wed, 21 Oct 2026 23:56:27 +0000 Subject: [PATCH 3/4] Move estimator metric_tag mapping to the active version on model update; document model_tag vs metric_tag --- sdks/python/apache_beam/ml/inference/base.py | 11 +++++ .../apache_beam/ml/inference/model_manager.py | 35 ++++++++++++--- .../ml/inference/model_manager_test.py | 45 +++++++++++++++++++ 3 files changed, 86 insertions(+), 5 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 2780f616c8be..b3f81632796e 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -2073,6 +2073,17 @@ 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. """ def __init__( self, diff --git a/sdks/python/apache_beam/ml/inference/model_manager.py b/sdks/python/apache_beam/ml/inference/model_manager.py index c7d2bee8956f..aad7d39da8fe 100644 --- a/sdks/python/apache_beam/ml/inference/model_manager.py +++ b/sdks/python/apache_beam/ml/inference/model_manager.py @@ -199,6 +199,15 @@ def set_metric_tag(self, model_tag: str, metric_tag: str): with self._lock: self._metric_tags[model_tag].add(metric_tag) + def unset_metric_tag(self, model_tag: str, metric_tag: str): + with self._lock: + self._metric_tags[model_tag].discard(metric_tag) + + 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.get(model_tag) or {model_tag}: @@ -413,6 +422,26 @@ 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 + self._estimator.set_metric_tag(tag, metric_tag) + if prev_tag is not None: + self._estimator.unset_metric_tag(prev_tag, metric_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] @@ -506,16 +535,12 @@ def acquire_model( 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() + self._maybe_switch_active_tag(effective_metric_tag, tag) # FAST PATH: Grab from idle LRU if available if not self._isolation_mode: diff --git a/sdks/python/apache_beam/ml/inference/model_manager_test.py b/sdks/python/apache_beam/ml/inference/model_manager_test.py index 9c4fd46856e0..abcb819c1fee 100644 --- a/sdks/python/apache_beam/ml/inference/model_manager_test.py +++ b/sdks/python/apache_beam/ml/inference/model_manager_test.py @@ -443,6 +443,51 @@ def _get_est_dist(tag): self.assertEqual(_get_count_dist(v2_tag).count, 0) self.assertEqual(_get_est_dist(v2_tag).count, 0) + def test_model_manager_metric_tag_switch_keeps_old_version_estimates(self): + """Test that a version switch moves metric reporting without forgetting + old versions that are still loaded.""" + stable_id = "switching_model_identifier" + v1_tag = "gs://bucket/switch_model_v1" + v2_tag = "gs://bucket/switch_model_v2" + + def _get_est_dist(tag): + dist = MetricsEnvironment.process_wide_container().get_distribution( + MetricName('BeamML_ModelManager', f'memory_estimate_mb_{tag}')) + return dist.get_cumulative() + + self.manager._estimator.set_initial_estimate(v1_tag, 1000.0) + self.manager._estimator.set_initial_estimate(v2_tag, 3000.0) + + # v1 is active for stable_id; keep one copy active across the switch. + m1 = self.manager.acquire_model( + v1_tag, lambda: MockModel(v1_tag, 1000.0, self.mock_monitor), + metric_tag=stable_id) + self.assertEqual(_get_est_dist(stable_id).max, 1000) + + # Switch to v2: stable_id now reports v2 and the switch re-emits v2's + # estimate immediately. + m2 = self.manager.acquire_model( + v2_tag, lambda: MockModel(v2_tag, 3000.0, self.mock_monitor), + metric_tag=stable_id) + self.assertEqual(self.manager._metric_tag_to_active_tag[stable_id], v2_tag) + self.assertEqual(_get_est_dist(stable_id).max, 3000) + self.assertEqual(self.manager._estimator._metric_tags[v1_tag], set()) + self.assertEqual(self.manager._estimator._metric_tags[v2_tag], {stable_id}) + + # v1 is still loaded, so the estimator must keep knowing about it for + # admission control, and its estimate no longer reports under stable_id. + MetricsEnvironment.process_wide_container().reset() + self.assertFalse(self.manager._estimator.is_unknown(v1_tag)) + self.assertIn(v1_tag, self.manager._estimator.known_models) + self.manager._estimator.set_initial_estimate(v1_tag, 1500.0) + self.assertEqual(_get_est_dist(stable_id).count, 0) + self.assertEqual(_get_est_dist(v1_tag).max, 1500) + + self.manager.release_model(v1_tag, m1) + self.manager.release_model(v2_tag, m2) + self.assertFalse(m1.deleted) + self.assertEqual(len(self.manager.all_models(v1_tag)), 1) + def test_single_model_convergence_with_fluctuations(self): """ Tests that the estimator converges to the true usage with fluctuations. From 5766f0c28db0e1efdec75c506480d4ce7ca22ab7 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Thu, 22 Oct 2026 00:07:29 +0000 Subject: [PATCH 4/4] Share metric_tag mapping between ModelManager and ResourceEstimator; note that model updates should only change weights --- sdks/python/apache_beam/ml/inference/base.py | 5 +- .../apache_beam/ml/inference/model_manager.py | 72 +++++++++++-------- .../ml/inference/model_manager_test.py | 59 ++++----------- 3 files changed, 59 insertions(+), 77 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index b3f81632796e..248186ff4fb9 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -2083,7 +2083,10 @@ class _SharedModelWrapper(): 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. + 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, diff --git a/sdks/python/apache_beam/ml/inference/model_manager.py b/sdks/python/apache_beam/ml/inference/model_manager.py index aad7d39da8fe..64a355cf372e 100644 --- a/sdks/python/apache_beam/ml/inference/model_manager.py +++ b/sdks/python/apache_beam/ml/inference/model_manager.py @@ -180,29 +180,27 @@ 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] = {} - # 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() + 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 set_metric_tag(self, model_tag: str, metric_tag: str): - with self._lock: - self._metric_tags[model_tag].add(metric_tag) - - def unset_metric_tag(self, model_tag: str, metric_tag: str): - with self._lock: - self._metric_tags[model_tag].discard(metric_tag) - def report_estimate(self, model_tag: str): with self._lock: if model_tag in self.estimates: @@ -210,7 +208,7 @@ def report_estimate(self, model_tag: str): # 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}: + for metric_tag in self._metric_tags_fn(model_tag): Metrics.distribution( _MODEL_MANAGER_METRICS_NAMESPACE, f"memory_estimate_mb_{metric_tag}", @@ -362,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 @@ -375,9 +374,16 @@ 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. + # 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. @@ -404,19 +410,30 @@ 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 metric_tag, active_tag in self._metric_tag_to_active_tag.items(): - Metrics.distribution( - _MODEL_MANAGER_METRICS_NAMESPACE, - f"num_loaded_models_{metric_tag}", - process_wide=True).update(len(self._models.get(active_tag, ()))) + for tag, instances in self._models.items(): + 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 metric_tag in self._metric_tag_to_active_tag: - Metrics.distribution( - _MODEL_MANAGER_METRICS_NAMESPACE, - f"num_loaded_models_{metric_tag}", - process_wide=True).update(0) + for tag in self._models: + 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: @@ -434,9 +451,6 @@ def _maybe_switch_active_tag(self, metric_tag: str, tag: str): if prev_tag == tag: return self._metric_tag_to_active_tag[metric_tag] = tag - self._estimator.set_metric_tag(tag, metric_tag) - if prev_tag is not None: - self._estimator.unset_metric_tag(prev_tag, metric_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): diff --git a/sdks/python/apache_beam/ml/inference/model_manager_test.py b/sdks/python/apache_beam/ml/inference/model_manager_test.py index abcb819c1fee..edd0c3e5af3b 100644 --- a/sdks/python/apache_beam/ml/inference/model_manager_test.py +++ b/sdks/python/apache_beam/ml/inference/model_manager_test.py @@ -402,7 +402,7 @@ def _get_est_dist_mean(tag): 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.""" + """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" @@ -417,7 +417,9 @@ def _get_est_dist(tag): MetricName('BeamML_ModelManager', f'memory_estimate_mb_{tag}')) return dist.get_cumulative() - # Acquire v1 with metric_tag=stable_id + 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) @@ -427,9 +429,11 @@ def _get_est_dist(tag): 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) + # 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) @@ -443,50 +447,11 @@ def _get_est_dist(tag): self.assertEqual(_get_count_dist(v2_tag).count, 0) self.assertEqual(_get_est_dist(v2_tag).count, 0) - def test_model_manager_metric_tag_switch_keeps_old_version_estimates(self): - """Test that a version switch moves metric reporting without forgetting - old versions that are still loaded.""" - stable_id = "switching_model_identifier" - v1_tag = "gs://bucket/switch_model_v1" - v2_tag = "gs://bucket/switch_model_v2" - - def _get_est_dist(tag): - dist = MetricsEnvironment.process_wide_container().get_distribution( - MetricName('BeamML_ModelManager', f'memory_estimate_mb_{tag}')) - return dist.get_cumulative() - - self.manager._estimator.set_initial_estimate(v1_tag, 1000.0) - self.manager._estimator.set_initial_estimate(v2_tag, 3000.0) - - # v1 is active for stable_id; keep one copy active across the switch. - m1 = self.manager.acquire_model( - v1_tag, lambda: MockModel(v1_tag, 1000.0, self.mock_monitor), - metric_tag=stable_id) - self.assertEqual(_get_est_dist(stable_id).max, 1000) - - # Switch to v2: stable_id now reports v2 and the switch re-emits v2's - # estimate immediately. - m2 = self.manager.acquire_model( - v2_tag, lambda: MockModel(v2_tag, 3000.0, self.mock_monitor), - metric_tag=stable_id) - self.assertEqual(self.manager._metric_tag_to_active_tag[stable_id], v2_tag) - self.assertEqual(_get_est_dist(stable_id).max, 3000) - self.assertEqual(self.manager._estimator._metric_tags[v1_tag], set()) - self.assertEqual(self.manager._estimator._metric_tags[v2_tag], {stable_id}) - - # v1 is still loaded, so the estimator must keep knowing about it for - # admission control, and its estimate no longer reports under stable_id. - MetricsEnvironment.process_wide_container().reset() - self.assertFalse(self.manager._estimator.is_unknown(v1_tag)) - self.assertIn(v1_tag, self.manager._estimator.known_models) - self.manager._estimator.set_initial_estimate(v1_tag, 1500.0) - self.assertEqual(_get_est_dist(stable_id).count, 0) - self.assertEqual(_get_est_dist(v1_tag).max, 1500) - - self.manager.release_model(v1_tag, m1) - self.manager.release_model(v2_tag, m2) + # 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.assertEqual(len(self.manager.all_models(v1_tag)), 1) + 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): """