diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 8c4b10f8573d..7a19eba35117 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -1950,7 +1950,10 @@ 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] + # 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 @@ 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..a8ffc81e5acf 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 @@ -2417,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() @@ -2479,6 +2503,80 @@ 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']) + 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 9c0d90d7a02f..69db694cc970 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 @@ def _delete_all_models(self): 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 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):