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 65ea034a6f1 Fix ModelManager tag mismatch on release after model
update (#40469) (#40473)
65ea034a6f1 is described below
commit 65ea034a6f1cf485c9dae2e5a0774cf3f2d2aa0e
Author: RuiLong J. <[email protected]>
AuthorDate: Wed Oct 14 14:54:32 2026 -0700
Fix ModelManager tag mismatch on release after model update (#40469)
(#40473)
* Fix ModelManager tag mismatch on release after model update (#40469)
* Fix _SharedModelWrapper.force_reset to call _force_reset()
* Add TODO(#40539) in _SharedModelWrapper.all_models()
* Isolate ModelManagerTest from xdist and expose force_reset on ModelManager
* Rename ModelManager._force_reset to force_reset
---
sdks/python/apache_beam/ml/inference/base.py | 7 +-
sdks/python/apache_beam/ml/inference/base_test.py | 98 ++++++++++++++++++++++
.../apache_beam/ml/inference/model_manager.py | 2 +-
.../apache_beam/ml/inference/model_manager_test.py | 6 +-
4 files changed, 107 insertions(+), 6 deletions(-)
diff --git a/sdks/python/apache_beam/ml/inference/base.py
b/sdks/python/apache_beam/ml/inference/base.py
index 8c4b10f8573..7a19eba3511 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -1950,7 +1950,10 @@ class _SharedModelWrapper():
def all_models(self):
if self.use_model_manager:
- return self.models.all_models()[self.model_tag]
+ # 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
def force_reset(self):
@@ -2106,7 +2109,7 @@ class _RunInferenceDoFn(beam.DoFn, Generic[ExampleT,
PredictionT]):
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 8cf408db76f..a8ffc81e5ac 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 @@ class SimpleFakeModelHandler(base.ModelHandler[int, int,
FakeModel]):
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
@@ -2417,9 +2435,15 @@ def try_import_model_manager():
return False
[email protected]_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()
@@ -2479,6 +2503,80 @@ class ModelManagerTest(unittest.TestCase):
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'])
+ wrapper.force_reset()
+ self.assertEqual(wrapper.all_models(), [])
+ 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()
+ 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__':
unittest.main()
diff --git a/sdks/python/apache_beam/ml/inference/model_manager.py
b/sdks/python/apache_beam/ml/inference/model_manager.py
index 9c0d90d7a02..69db694cc97 100644
--- a/sdks/python/apache_beam/ml/inference/model_manager.py
+++ b/sdks/python/apache_beam/ml/inference/model_manager.py
@@ -795,7 +795,7 @@ class ModelManager:
self._monitor.refresh()
self._monitor.reset_peak()
- def _force_reset(self):
+ 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 370961eb61d..fc84f4ba859 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 @@ class TestModelManager(unittest.TestCase):
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 @@ class TestModelManager(unittest.TestCase):
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 @@ class TestModelManager(unittest.TestCase):
# 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):
