From 41807129a06da806798ad0413bb75bfee272b992 Mon Sep 17 00:00:00 2001 From: Divyansh Yadav Date: Fri, 16 Oct 2026 22:58:18 +0530 Subject: [PATCH] Only feed total memory readings to the ModelManager resource estimator --- .../apache_beam/ml/inference/model_manager.py | 9 +++----- .../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 69db694cc970..c906ceb6d75b 100644 --- a/sdks/python/apache_beam/ml/inference/model_manager.py +++ b/sdks/python/apache_beam/ml/inference/model_manager.py @@ -753,13 +753,10 @@ def _spawn_new_model( 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 fc84f4ba8598..7b6337eaec8f 100644 --- a/sdks/python/apache_beam/ml/inference/model_manager_test.py +++ b/sdks/python/apache_beam/ml/inference/model_manager_test.py @@ -265,6 +265,29 @@ def test_model_manager_concurrent_execution(self): 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