[RunInference] Reject use_model_manager with a multi-model KeyedModelHandler

sdks/python/apache_beam/ml/inference/base.py

@@ -742,6 +742,10 @@ def __init__(
     at the same time. For more information about memory management, see
     `Use a keyed `ModelHandler <https://beam.apache.org/documentation/ml/about-ml/#use-a-keyed-modelhandler-object>_`.  # pylint: disable=line-too-long
 
+    A KeyedModelHandler with many models defined per key cannot be used with
+    `RunInference(..., use_model_manager=True)`, since it manages its per-key
+    models itself.
+
 
     Args:
       unkeyed: Either (a) an implementation of ModelHandler that does not
@@ -1370,6 +1374,19 @@ def __call__(self, *args, **kwargs):
       raise e
 
 
+def _is_multi_model_keyed_handler(model_handler: ModelHandler) -> bool:
+  """Returns whether model_handler is a KeyedModelHandler with many models,
+  looking through any batching or pre/postprocessing wrappers."""
+  while isinstance(model_handler,
+                   (_PrebatchedModelHandler,
+                    _PreProcessingModelHandler,
+                    _PostProcessingModelHandler)):
+    model_handler = model_handler._base
+  return (
+      isinstance(model_handler, KeyedModelHandler) and
+      not model_handler._single_model)
+
+
 class RunInference(beam.PTransform[beam.PCollection[Union[ExampleT,
                                                           Iterable[ExampleT]]],
                                    beam.PCollection[PredictionT]]):
@@ -1416,10 +1433,26 @@ def __init__(
           the same tag for different models will lead to non-deterministic
           results, so exercise caution when using this parameter. This only
           impacts models which are already being shared across processes.
+        use_model_manager: Whether to load and manage models with a
+          ModelManager, which tracks memory usage and evicts models as needed.
+          This is not supported with a KeyedModelHandler that maps keys to
+          multiple model handlers, since that handler manages its own per-key
+          models (see `max_models_per_worker_hint`); a ValueError is raised
+          in that case.
+        model_manager_args: Keyword arguments used to construct the
+          ModelManager when use_model_manager is True.
         monitoring_transform: A PTransform that receives a copy of the
           un-postprocessed PCollection of PredictionResult objects produced
           directly by inference.
     """
+    if use_model_manager and _is_multi_model_keyed_handler(model_handler):
+      raise ValueError(
+          'use_model_manager=True is not supported with a KeyedModelHandler '
+          'that maps keys to multiple model handlers. That handler loads and '
+          'caches its per-key models itself, so they would not be tracked by '
+          'the ModelManager. Use max_models_per_worker_hint on the '
+          'KeyedModelHandler to limit how many models are loaded at once, or '
+          'set use_model_manager=False.')
     self._model_handler = model_handler
     self._inference_args = inference_args
     self._clock = clock

sdks/python/apache_beam/ml/inference/base_test.py

@@ -571,6 +571,29 @@ def mult_two(example: str) -> int:
     with self.assertRaises(ValueError):
       base.KeyedModelHandler(mhs)
 
+  def test_keyed_many_model_handlers_with_model_manager_raises(self):
+    mhs = [
+        base.KeyModelMapping([0],
+                             FakeModelHandler(
+                                 state=200, multi_process_shared=True)),
+        base.KeyModelMapping([1, 2, 3],
+                             FakeModelHandler(multi_process_shared=True))
+    ]
+    keyed_mh = base.KeyedModelHandler(mhs)
+    with self.assertRaisesRegex(ValueError, r'use_model_manager=True'):
+      base.RunInference(keyed_mh, use_model_manager=True)
+
+    # Wrapped handlers should be rejected as well.
+    with self.assertRaisesRegex(ValueError, r'use_model_manager=True'):
+      base.RunInference(
+          keyed_mh.with_preprocess_fn(lambda x: x).with_no_batching(),
+          use_model_manager=True)
+
+    # Without the model manager, or with a single keyed model, it is allowed.
+    base.RunInference(keyed_mh)
+    base.RunInference(
+        base.KeyedModelHandler(FakeModelHandler()), use_model_manager=True)
+
   def test_keyed_model_handler_get_num_bytes(self):
     mh = base.KeyedModelHandler(FakeModelHandler(num_bytes_per_element=10))
     batch = [('key1', 1), ('key2', 2), ('key1', 3)]