From 5fc758990ef831a19282c197abf704918c6a0c77 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Thu, 8 Oct 2026 20:45:03 +0000 Subject: [PATCH 1/5] Fix ModelManager tag mismatch on release after model update (#40469) --- sdks/python/apache_beam/ml/inference/base.py | 4 +- .../apache_beam/ml/inference/base_test.py | 88 +++++++++++++++++++ 2 files changed, 90 insertions(+), 2 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 8c4b10f8573d..4143ec60e300 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -1950,7 +1950,7 @@ def release_model(self, model_tag: str, model: Any): def all_models(self): if self.use_model_manager: - return self.models.all_models()[self.model_tag] + return self.models.all_models(self.model_tag) return self.models def force_reset(self): @@ -2106,7 +2106,7 @@ def _run_inference(self, batch, inference_args): finally: # Always release the model so that it can be reloaded. if self.use_model_manager: - self._model.release_model(self._model_tag, unique_tag) + self._model.release_model(self._cur_tag, unique_tag) except BaseException as e: if self._metrics_collector: self._metrics_collector.failed_batches_counter.inc() diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index 8cf408db76f4..7a138ec727f6 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -2408,6 +2408,24 @@ def run_inference( yield model.predict(example) +class IncrementingFakeModelHandler(base.ModelHandler[int, + int, + FakeIncrementingModel]): + def load_model(self): + return FakeIncrementingModel() + + def run_inference( + self, + batch: Sequence[int], + model: FakeIncrementingModel, + inference_args=None) -> Iterable[int]: + for _ in batch: + yield model.predict(0) + + def batch_elements_kwargs(self): + return {'min_batch_size': 2, 'max_batch_size': 2} + + def try_import_model_manager(): try: # pylint: disable=unused-import @@ -2479,6 +2497,76 @@ def batch_elements_kwargs(self): OOMFakeModelHandler(), use_model_manager=True) assert_that(actual, equal_to([2, 6, 4, 11]), label='assert:inferences') + @unittest.skipIf( + not try_import_model_manager(), 'Model Manager not available') + def test_shared_model_wrapper_all_models_with_model_manager(self): + from apache_beam.ml.inference.model_manager import ModelManager + mm = ModelManager() + try: + tag = 'test_tag' + instance = mm.acquire_model(tag, lambda: 'fake_model_instance') + mm.release_model(tag, instance) + wrapper = base._SharedModelWrapper(mm, tag, lambda: 'fake_model_instance') + self.assertEqual(wrapper.all_models(), ['fake_model_instance']) + finally: + mm.shutdown() + + @unittest.skipIf( + not try_import_model_manager(), 'Model Manager not available') + def test_run_inference_with_model_manager_and_model_metadata_pcoll_reuses_model( + self): + with TestPipeline('FnApiRunner') as pipeline: + model_metadata = pipeline | 'ModelMetadata' >> beam.Create( + [base.ModelMetadata(model_id='model_v2', model_name='v2')]) + actual = ( + pipeline + | 'start' >> beam.Create([0, 0, 0, 0]) + | base.RunInference( + IncrementingFakeModelHandler(), + use_model_manager=True, + model_metadata_pcoll=model_metadata)) + # Since the model is reused across the 2 batches (4 elements total), + # FakeIncrementingModel increments from 1 to 4 instead of resetting to 1 + # on the second batch. + assert_that(actual, equal_to([1, 2, 3, 4]), label='assert:inferences') + + @unittest.skipIf( + not try_import_model_manager(), 'Model Manager not available') + def test_run_inference_dofn_model_update_acquires_and_releases_matching_tag( + self): + dofn = base._RunInferenceDoFn( + model_handler=IncrementingFakeModelHandler(), + clock=FakeClock(), + metrics_namespace='test', + load_model_at_runtime=False, + model_tag='initial_tag', + use_model_manager=True) + dofn.setup() + + # 1. Before any update, batches use 'initial_tag' and reuse the same model + self.assertEqual(dofn.process([0, 0], None, None), [1, 2]) + self.assertEqual(dofn.process([0, 0], None, None), [3, 4]) + self.assertEqual( + len(dofn._model.all_models()), 1, 'Expected 1 model for initial_tag') + + # 2. Update to 'model_v1' via side input metadata + meta_v1 = base.ModelMetadata(model_id='model_v1', model_name='v1') + self.assertEqual(dofn.process([0, 0], None, meta_v1), [1, 2]) + # 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( + len(dofn._model.all_models()), 1, 'Expected 1 model for model_v1') + + # 3. Update to 'model_v2' via side input metadata + meta_v2 = base.ModelMetadata(model_id='model_v2', model_name='v2') + self.assertEqual(dofn.process([0, 0], None, meta_v2), [1, 2]) + # 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( + len(dofn._model.all_models()), 1, 'Expected 1 model for model_v2') + if __name__ == '__main__': unittest.main() From 3cdcfb8be102b5a233ce6c58431b5247286985e7 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Thu, 8 Oct 2026 23:18:28 +0000 Subject: [PATCH 2/5] Fix _SharedModelWrapper.force_reset to call _force_reset() --- sdks/python/apache_beam/ml/inference/base.py | 2 +- sdks/python/apache_beam/ml/inference/base_test.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 4143ec60e300..f39050b08bba 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -1955,7 +1955,7 @@ def all_models(self): def force_reset(self): if self.use_model_manager: - self.models.force_reset() + self.models._force_reset() class _RunInferenceDoFn(beam.DoFn, Generic[ExampleT, PredictionT]): diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index 7a138ec727f6..500b66772b26 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -2508,6 +2508,8 @@ def test_shared_model_wrapper_all_models_with_model_manager(self): mm.release_model(tag, instance) wrapper = base._SharedModelWrapper(mm, tag, lambda: 'fake_model_instance') self.assertEqual(wrapper.all_models(), ['fake_model_instance']) + wrapper.force_reset() + self.assertEqual(wrapper.all_models(), []) finally: mm.shutdown() From 8a5e611ccc8c11a9583a21e0765c5057addbe81b Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Mon, 12 Oct 2026 18:36:28 +0000 Subject: [PATCH 3/5] Add TODO(#40539) in _SharedModelWrapper.all_models() --- sdks/python/apache_beam/ml/inference/base.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index f39050b08bba..ead7fd534894 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -1950,6 +1950,9 @@ def release_model(self, model_tag: str, model: Any): def all_models(self): if self.use_model_manager: + # TODO(https://github.com/apache/beam/issues/40539): Support multi-model + # KeyedModelHandler with ModelManager (e.g. resolving proxies for + # update_model_paths). return self.models.all_models(self.model_tag) return self.models From d3b2fddd6af5f7ac27810b0d53905c8cd6b46b68 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Mon, 12 Oct 2026 20:52:45 +0000 Subject: [PATCH 4/5] Isolate ModelManagerTest from xdist and expose force_reset on ModelManager --- sdks/python/apache_beam/ml/inference/base.py | 2 +- .../apache_beam/ml/inference/base_test.py | 56 +++++++++++-------- .../apache_beam/ml/inference/model_manager.py | 3 + 3 files changed, 36 insertions(+), 25 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index ead7fd534894..7a19eba35117 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -1958,7 +1958,7 @@ def all_models(self): def force_reset(self): if self.use_model_manager: - self.models._force_reset() + self.models.force_reset() class _RunInferenceDoFn(beam.DoFn, Generic[ExampleT, PredictionT]): diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index 500b66772b26..a8ffc81e5acf 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -2435,9 +2435,15 @@ def try_import_model_manager(): return False +@pytest.mark.no_xdist class ModelManagerTest(unittest.TestCase): """Tests for RunInference with Model Manager integration.""" def tearDown(self): + if multi_process_shared._process_level_singleton_manager.has_singleton( + 'model_manager'): + multi_process_shared.MultiProcessShared( + lambda: None, tag='model_manager', + always_proxy=True).acquire().force_reset() for p in multiprocessing.active_children(): p.terminate() p.join() @@ -2544,30 +2550,32 @@ def test_run_inference_dofn_model_update_acquires_and_releases_matching_tag( model_tag='initial_tag', use_model_manager=True) dofn.setup() - - # 1. Before any update, batches use 'initial_tag' and reuse the same model - self.assertEqual(dofn.process([0, 0], None, None), [1, 2]) - self.assertEqual(dofn.process([0, 0], None, None), [3, 4]) - self.assertEqual( - len(dofn._model.all_models()), 1, 'Expected 1 model for initial_tag') - - # 2. Update to 'model_v1' via side input metadata - meta_v1 = base.ModelMetadata(model_id='model_v1', model_name='v1') - self.assertEqual(dofn.process([0, 0], None, meta_v1), [1, 2]) - # 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( - len(dofn._model.all_models()), 1, 'Expected 1 model for model_v1') - - # 3. Update to 'model_v2' via side input metadata - meta_v2 = base.ModelMetadata(model_id='model_v2', model_name='v2') - self.assertEqual(dofn.process([0, 0], None, meta_v2), [1, 2]) - # 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( - len(dofn._model.all_models()), 1, 'Expected 1 model for model_v2') + try: + # 1. Before any update, batches use 'initial_tag' and reuse the same model + self.assertEqual(dofn.process([0, 0], None, None), [1, 2]) + self.assertEqual(dofn.process([0, 0], None, None), [3, 4]) + self.assertEqual( + len(dofn._model.all_models()), 1, 'Expected 1 model for initial_tag') + + # 2. Update to 'model_v1' via side input metadata + meta_v1 = base.ModelMetadata(model_id='model_v1', model_name='v1') + self.assertEqual(dofn.process([0, 0], None, meta_v1), [1, 2]) + # 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( + len(dofn._model.all_models()), 1, 'Expected 1 model for model_v1') + + # 3. Update to 'model_v2' via side input metadata + meta_v2 = base.ModelMetadata(model_id='model_v2', model_name='v2') + self.assertEqual(dofn.process([0, 0], None, meta_v2), [1, 2]) + # 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( + len(dofn._model.all_models()), 1, 'Expected 1 model for model_v2') + finally: + dofn._model.force_reset() if __name__ == '__main__': diff --git a/sdks/python/apache_beam/ml/inference/model_manager.py b/sdks/python/apache_beam/ml/inference/model_manager.py index 9c0d90d7a02f..ab2d54f3e605 100644 --- a/sdks/python/apache_beam/ml/inference/model_manager.py +++ b/sdks/python/apache_beam/ml/inference/model_manager.py @@ -795,6 +795,9 @@ def _delete_all_models(self): self._monitor.refresh() self._monitor.reset_peak() + def force_reset(self): + self._force_reset() + def _force_reset(self): logger.warning("Force Reset Triggered") self._delete_all_models() From 755c6596a16b3b0d080e28d29c8bcaa6504a8ead Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Mon, 12 Oct 2026 21:49:28 +0000 Subject: [PATCH 5/5] Rename ModelManager._force_reset to force_reset --- sdks/python/apache_beam/ml/inference/model_manager.py | 3 --- sdks/python/apache_beam/ml/inference/model_manager_test.py | 6 +++--- 2 files changed, 3 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 ab2d54f3e605..69db694cc970 100644 --- a/sdks/python/apache_beam/ml/inference/model_manager.py +++ b/sdks/python/apache_beam/ml/inference/model_manager.py @@ -796,9 +796,6 @@ def _delete_all_models(self): self._monitor.reset_peak() def force_reset(self): - self._force_reset() - - def _force_reset(self): logger.warning("Force Reset Triggered") self._delete_all_models() self._models = defaultdict(list) 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 370961eb61dd..fc84f4ba8598 100644 --- a/sdks/python/apache_beam/ml/inference/model_manager_test.py +++ b/sdks/python/apache_beam/ml/inference/model_manager_test.py @@ -148,7 +148,7 @@ def loader(): self.manager.release_model(model_name, instance) # Force delete all models - self.manager._force_reset() + self.manager.force_reset() # Verify that the MultiProcessShared instance is deleted # and the counter is reseted @@ -327,7 +327,7 @@ def dummy_loader(): self.manager.release_model(model_name, instance) instance = self.manager.acquire_model(model_name, dummy_loader) except RuntimeError: - self.manager._force_reset() + self.manager.force_reset() self.assertTrue(len(self.manager._models[model_name]) == 0) self.assertEqual(self.manager._total_active_jobs, 0) self.assertEqual(self.manager._pending_reservations, 0.0) @@ -397,7 +397,7 @@ def _get_est_dist_mean(tag): # Max is still 2 from earlier in the test run self.assertEqual(_get_count_dist_max(tag1), 2) - self.manager._force_reset() + self.manager.force_reset() self.assertEqual(_get_count_dist_max(tag1), 2) def test_single_model_convergence_with_fluctuations(self):