This is an automated email from the ASF dual-hosted git repository.
tvalentyn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 1ba4129cb10 Preserve ModelManager metric_tag across RunInference
side-input model updates (#40714)
1ba4129cb10 is described below
commit 1ba4129cb1005b0004c965934c34faff939c00e2
Author: RuiLong J. <[email protected]>
AuthorDate: Fri Oct 23 12:18:31 2026 -0700
Preserve ModelManager metric_tag across RunInference side-input model
updates (#40714)
* Preserve ModelManager metric_tag across RunInference side-input model
updates
* Trigger CI rerun
* Move estimator metric_tag mapping to the active version on model update;
document model_tag vs metric_tag
* Share metric_tag mapping between ModelManager and ResourceEstimator; note
that model updates should only change weights
---
sdks/python/apache_beam/ml/inference/base.py | 26 +++++-
sdks/python/apache_beam/ml/inference/base_test.py | 2 +
.../apache_beam/ml/inference/model_manager.py | 101 +++++++++++++++++----
.../apache_beam/ml/inference/model_manager_test.py | 52 +++++++++++
4 files changed, 159 insertions(+), 22 deletions(-)
diff --git a/sdks/python/apache_beam/ml/inference/base.py
b/sdks/python/apache_beam/ml/inference/base.py
index c2e084d0e03..5ee65afa061 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -2107,15 +2107,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(
@@ -2126,7 +2142,8 @@ class _SharedModelWrapper():
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
@@ -2235,7 +2252,10 @@ class _RunInferenceDoFn(beam.DoFn, Generic[ExampleT,
PredictionT]):
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 f1bb5d1245e..dec1c2906cd 100644
--- a/sdks/python/apache_beam/ml/inference/base_test.py
+++ b/sdks/python/apache_beam/ml/inference/base_test.py
@@ -3136,6 +3136,7 @@ class ModelManagerTest(unittest.TestCase):
# 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')
@@ -3145,6 +3146,7 @@ class ModelManagerTest(unittest.TestCase):
# 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 79a96a360fa..acedca39af8 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 @@ class ResourceEstimator:
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 @@ class ResourceEstimator:
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 @@ class ResourceEstimator:
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 @@ class ResourceEstimator:
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 @@ class ModelManager:
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 @@ class ModelManager:
# 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 @@ class ModelManager:
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 @@ class ModelManager:
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)
@@ -800,6 +862,7 @@ class ModelManager:
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 6b94164baf8..eadca456b40 100644
--- a/sdks/python/apache_beam/ml/inference/model_manager_test.py
+++ b/sdks/python/apache_beam/ml/inference/model_manager_test.py
@@ -424,6 +424,58 @@ class TestModelManager(unittest.TestCase):
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.
