diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 4c2d9cb4c80e..248186ff4fb9 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/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, 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..64a355cf372e 100644 --- a/sdks/python/apache_beam/ml/inference/model_manager.py +++ b/sdks/python/apache_beam/ml/inference/model_manager.py @@ -180,7 +180,14 @@ 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 @@ -188,11 +195,25 @@ def __init__( 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 = [] 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..edd0c3e5af3b 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,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.