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 e333810ea3a Only feed total memory readings to the ModelManager
resource estimator (#40646)
e333810ea3a is described below
commit e333810ea3a1723b74b0adcd2006bfc5d3bb9306
Author: Divyansh Yadav <[email protected]>
AuthorDate: Fri Oct 23 03:04:42 2026 +0530
Only feed total memory readings to the ModelManager resource estimator
(#40646)
---
.../apache_beam/ml/inference/model_manager.py | 9 +++------
.../apache_beam/ml/inference/model_manager_test.py | 23 ++++++++++++++++++++++
2 files changed, 26 insertions(+), 6 deletions(-)
diff --git a/sdks/python/apache_beam/ml/inference/model_manager.py
b/sdks/python/apache_beam/ml/inference/model_manager.py
index bc16feb6f60..79a96a360fa 100644
--- a/sdks/python/apache_beam/ml/inference/model_manager.py
+++ b/sdks/python/apache_beam/ml/inference/model_manager.py
@@ -757,13 +757,10 @@ class ModelManager:
try:
with self._cv:
self.logging_info("Loading Model: %s (Unknown: %s)", tag, is_unknown)
- baseline_snap, _, _ = self._monitor.get_stats()
+ # No estimator observation here: the estimator is fit on total memory
+ # readings, which release_model records along with the full snapshot
+ # of loaded models. A load delta would contradict those readings.
instance = loader_func()
- _, peak_during_load, _ = self._monitor.get_stats()
-
- snapshot = {tag: 1}
- self._estimator.add_observation(
- snapshot, peak_during_load - baseline_snap)
if not is_unknown:
self._pending_reservations = max(
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 d1f2b35bc1d..6b94164baf8 100644
--- a/sdks/python/apache_beam/ml/inference/model_manager_test.py
+++ b/sdks/python/apache_beam/ml/inference/model_manager_test.py
@@ -266,6 +266,29 @@ class TestModelManager(unittest.TestCase):
self.manager.release_model(model_b, inst_b)
self.assertEqual(self.manager._total_active_jobs, 0)
+ def test_model_manager_observations_use_total_memory(self):
+ """
+ Test that the estimator only receives total memory readings for the
+ full snapshot of loaded models, not the memory delta of a single load.
+ """
+ model_a = "model_a"
+ model_b = "model_b"
+
+ self.manager._estimator.set_initial_estimate(model_a, 1000.0)
+ self.manager._estimator.set_initial_estimate(model_b, 2000.0)
+
+ inst_a = self.manager.acquire_model(
+ model_a, lambda: MockModel(model_a, 1000.0, self.mock_monitor))
+ inst_b = self.manager.acquire_model(
+ model_b, lambda: MockModel(model_b, 2000.0, self.mock_monitor))
+ self.manager.release_model(model_b, inst_b)
+ self.manager.release_model(model_a, inst_a)
+
+ history = self.manager._estimator.history
+ snapshot_key = ((model_a, 1), (model_b, 1))
+ self.assertEqual(list(history.keys()), [snapshot_key])
+ self.assertEqual(list(history[snapshot_key]), [3000.0, 3000.0])
+
def test_model_manager_concurrent_mixed_workload_convergence(self):
"""
Simulates a production environment with multiple model types running
