Fix ModelManager tag mismatch on release after model update (#40469)

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()

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()

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)

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):