[RunInference] Honor model updates in also in the first batch.

CHANGES.md

@@ -99,7 +99,6 @@
 ## Deprecations
 
 * Support for Python 3.10 was removed following its end-of-life ([#40250](https://github.com/apache/beam/issues/40250)).
-* X behavior is deprecated and will be removed in X versions ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Bugfixes
 
@@ -111,6 +110,9 @@
 * (Go) Wrap last dataWriter/timerWriter send errors on Close ([#40405](https://github.com/apache/beam/issues/40405)).
 * (Go) Fixed a deadlock recreating a data channel while holding the channel lock ([#40414](https://github.com/apache/beam/issues/40414)).
 * (Python) Fixed `KeyedModelHandler` not enforcing `max_models_per_worker_hint` when several copies of the model handler run in the same SDK process ([#40468](https://github.com/apache/beam/issues/40468)).
+* (Python) Fixed a bug where `RunInference` was serving the first batch on each worker with the original models instead of the paths from a `KeyModelPathMapping` side input ([#40665](https://github.com/apache/beam/issues/40665)).
+* (Python) Fixed a bug in `RunInference` where adaptive throttling in remote inference codepath was accidentally disabled ([#40598](https://github.com/apache/beam/pull/40598)).
+* (Python) Fixed an inefficiency in `RunInference` that caused unnecessary model reloading in some cases when `large_model=True` was not set ([#40441](https://github.com/apache/beam/pull/40441)).
 * (Java) Fixed the declared schema of the error output of the Kafka write SchemaTransform, which wrapped the error schema a second time and did not match the rows it emits ([#39760](https://github.com/apache/beam/issues/39760)).
 * (Go) Fixed pubsubio importing a `google.golang.org/genproto` package removed in recent releases, which broke builds of Go modules depending on a current `genproto` version ([#40018](https://github.com/apache/beam/issues/40018)).
 * (Python) `VertexAIFeatureStoreEnrichmentHandler` now logs a warning instead of failing when the caller lacks permission to read the feature store metadata ([#40442](https://github.com/apache/beam/issues/40442)).

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

@@ -2146,7 +2146,9 @@ def __init__(
     self._model = None
     self._metrics_namespace = metrics_namespace
     self._load_model_at_runtime = load_model_at_runtime
-    self._side_input_path = None
+    # Most recently applied side input model path(s), if any.
+    self._side_input_path: Optional[Union[str,
+                                          list[KeyModelPathMapping]]] = None
     # _model_tag is the original tag passed in.
     # _cur_tag is the tag of the actually loaded model
     self._model_tag = model_tag
@@ -2217,11 +2219,11 @@ def load():
     # of calling load(). For sanity check, call update_model_path again.
     if isinstance(side_input_model_path, str):
       self._model_handler.update_model_path(side_input_model_path)
-    else:
-      if self._model is not None:
-        models = self._model.all_models()
-        for m in models:
-          self._model_handler.update_model_paths(m, side_input_model_path)
+    elif side_input_model_path:
+      # For per-key models, side_input_model_path is a list of mappings.
+      for m in model_wrapper.all_models():
+        self._model_handler.update_model_paths(
+            m, model_paths=side_input_model_path)
     return model_wrapper
 
   def get_metrics_collector(self, prefix: str = ''):
@@ -2251,6 +2253,22 @@ def update_model(
                                             list[KeyModelPathMapping]]] = None):
     self._model = self._load_model(side_input_model_path=side_input_model_path)
 
+  def _should_update_model(
+      self,
+      new_side_input_path: Optional[Union[str,
+                                          list[KeyModelPathMapping]]]) -> bool:
+    """Returns whether the model needs to be (re)loaded before inference.
+
+    Args:
+      new_side_input_path: Model path(s) from the side input, or None if there
+        is no side input value.
+    """
+    return (
+        self._model is None or (
+            new_side_input_path is not None and
+            self._side_input_path != new_side_input_path) or
+        not self._model_metadata.is_valid_tag(self._cur_tag))
+
   def _run_inference(self, batch, inference_args):
     start_time = _to_microseconds(self._clock.time_ns())
     try:
@@ -2305,26 +2323,24 @@ def process(
       side input is empty or the model has not been updated, the method
       simply runs inference on the batch of data.
     """
-    if not si_model_metadata:
-      if (not self._model_metadata.is_valid_tag(self._cur_tag) or
-          self._model is None):
+    if (not si_model_metadata or
+        isinstance(si_model_metadata, beam.pvalue.EmptySideInput)):
+      # No side input value: keep the loaded model, if any.
+      if self._should_update_model(new_side_input_path=None):
         self.update_model(side_input_model_path=None)
       return self._run_inference(batch, inference_args)
 
-    if isinstance(si_model_metadata, beam.pvalue.EmptySideInput):
-      self.update_model(side_input_model_path=None)
-    elif isinstance(si_model_metadata, list) and hasattr(si_model_metadata[0],
-                                                         'keys'):
+    if isinstance(si_model_metadata, list) and hasattr(si_model_metadata[0],
+                                                       'keys'):
       # TODO(https://github.com/apache/beam/issues/27628): Update metrics here
-      self.update_model(si_model_metadata)
-    elif self._side_input_path != si_model_metadata.model_id:
-      self._side_input_path = si_model_metadata.model_id
+      if self._should_update_model(si_model_metadata):
+        self.update_model(si_model_metadata)
+        self._side_input_path = list(si_model_metadata)
+    elif self._should_update_model(si_model_metadata.model_id):
       self._metrics_collector = self.get_metrics_collector(
           prefix=si_model_metadata.model_name)
-      lock = threading.Lock()
-      with lock:
-        self.update_model(si_model_metadata.model_id)
-        return self._run_inference(batch, inference_args)
+      self.update_model(si_model_metadata.model_id)
+      self._side_input_path = si_model_metadata.model_id
 
     return self._run_inference(batch, inference_args)
 

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

@@ -274,6 +274,41 @@ def share_model_across_processes(self):
     return self._multi_process_shared
 
 
+class FakeModelWithPath:
+  """A model that reports the path it was loaded from."""
+  def __init__(self, path: str):
+    self._path = path
+
+  def predict(self, example: int) -> str:
+    return self._path
+
+
+class FakeModelHandlerWithPath(base.ModelHandler[int, str, FakeModelWithPath]):
+  """Predicts the path of the model that actually served each example."""
+  def __init__(self, path: str):
+    self._path = path
+
+  def load_model(self):
+    return FakeModelWithPath(self._path)
+
+  def run_inference(
+      self,
+      batch: Sequence[int],
+      model: FakeModelWithPath,
+      inference_args=None) -> Iterable[str]:
+    return [model.predict(example) for example in batch]
+
+  def update_model_path(self, model_path: Optional[str] = None):
+    if model_path:
+      self._path = model_path
+
+
+class SingleElementBatchesKeyedModelHandler(base.KeyedModelHandler):
+  """Runs inference on one element at a time, so pipelines have many batches."""
+  def batch_elements_kwargs(self):
+    return {'min_batch_size': 1, 'max_batch_size': 1}
+
+
 class FakeModelHandlerNoEnvVars(base.ModelHandler[int, int, FakeModel]):
   def __init__(
       self, clock=None, min_batch_size=1, max_batch_size=9999, **kwargs):
@@ -1463,6 +1498,41 @@ def test_run_inference_empty_side_input(self):
 
       assert_that(result_pcoll, equal_to(expected))
 
+  def _count_update_model_calls(self):
+    """Returns per-DoFn-instance update_model() call counts and a patcher.
+
+    The counts are only collected when the pipeline runs in this process,
+    e.g. on FnApiRunner.
+    """
+    calls: dict[str, int] = {}
+    original_update_model = base._RunInferenceDoFn.update_model
+
+    def counting_update_model(dofn, *args, **kwargs):
+      instance_id = dofn.__dict__.setdefault(
+          '_test_instance_id', uuid.uuid4().hex)
+      calls[instance_id] = calls.get(instance_id, 0) + 1
+      return original_update_model(dofn, *args, **kwargs)
+
+    return calls, unittest.mock.patch.object(
+        base._RunInferenceDoFn, 'update_model', counting_update_model)
+
+  def test_run_inference_empty_side_input_loads_model_once(self):
+    examples = list(range(10))
+    update_model_calls, count_calls = self._count_update_model_calls()
+    with count_calls, TestPipeline('FnApiRunner') as pipeline:
+      side_pcoll = pipeline | 'side' >> beam.Create([])
+      result_pcoll = (
+          pipeline
+          | beam.Create(examples)
+          | base.RunInference(
+              FakeModelHandler(min_batch_size=1, max_batch_size=1),
+              model_metadata_pcoll=side_pcoll))
+      assert_that(result_pcoll, equal_to([e + 1 for e in examples]))
+
+    # The model is loaded once per DoFn instance, not once per batch.
+    self.assertTrue(update_model_calls)
+    self.assertEqual({1}, set(update_model_calls.values()))
+
   def test_run_inference_side_input_in_batch(self):
     first_ts = math.floor(time.time()) - 30
     interval = 7
@@ -1745,6 +1815,58 @@ def process(self, element):
 
       assert_that(result_pcoll, equal_to(expected_result))
 
+  def test_run_inference_side_input_per_key_models_used_from_first_batch(self):
+    model_handler = SingleElementBatchesKeyedModelHandler([
+        base.KeyModelMapping(['key1'],
+                             FakeModelHandlerWithPath('original/key1')),
+        base.KeyModelMapping(['key2'],
+                             FakeModelHandlerWithPath('original/key2')),
+    ])
+    key_model_paths = [
+        base.KeyModelPathMapping(keys=['key1'], update_path='updated/key1'),
+        base.KeyModelPathMapping(keys=['key2'], update_path='updated/key2'),
+    ]
+    examples = [('key1', 1), ('key2', 2), ('key1', 3), ('key2', 4)]
+    update_model_calls, count_calls = self._count_update_model_calls()
+    with count_calls, TestPipeline('FnApiRunner') as pipeline:
+      side_input = pipeline | 'SideInput' >> beam.Create([key_model_paths])
+      result_pcoll = (
+          pipeline
+          | 'Examples' >> beam.Create(examples)
+          | base.RunInference(model_handler, model_metadata_pcoll=side_input))
+      # Every example, including the ones in the first batch, is served by
+      # the model loaded from the path given in the side input.
+      assert_that(
+          result_pcoll,
+          equal_to([(key, 'updated/' + key) for key, _ in examples]))
+
+    # The unchanged mapping is applied once per DoFn instance, not per batch.
+    self.assertTrue(update_model_calls)
+    self.assertEqual({1}, set(update_model_calls.values()))
+
+  def test_side_input_model_reloaded_after_invalidation(self):
+    dofn = base._RunInferenceDoFn(
+        FakeModelHandler(multi_process_shared=True),
+        FakeClock(),
+        metrics_namespace=None,
+        load_model_at_runtime=True,
+        model_tag='test_model_' + uuid.uuid4().hex)
+    dofn.setup()
+    model_id = 'side_input_model_' + uuid.uuid4().hex
+    side_input = base.ModelMetadata(model_id=model_id, model_name='model')
+    self.assertEqual([2], dofn.process([1], None, side_input))
+    self.assertEqual(model_id, dofn._cur_tag)
+
+    # Simulate an inference timeout: the model is marked invalid and, since
+    # its grace period is 0s here, deleted in finish_bundle().
+    dofn._model_metadata.try_mark_current_model_invalid(0)
+    dofn.finish_bundle()
+
+    # The same side input value reloads the model under a new tag instead of
+    # serving the deleted one.
+    self.assertEqual([3], dofn.process([2], None, side_input))
+    self.assertEqual(f'{model_id}_reload_1', dofn._cur_tag)
+
   def test_run_inference_side_input_in_batch_multi_process_shared(self):
     first_ts = math.floor(time.time()) - 30
     interval = 7