From f4f1aee14b0b13c749da3208f281f1777ffa0391 Mon Sep 17 00:00:00 2001 From: Valentyn Tymofieiev Date: Mon, 19 Oct 2026 12:17:30 -0700 Subject: [PATCH 1/5] Honor model updates in first batch. Don't reload more than necessary. --- CHANGES.md | 1 + sdks/python/apache_beam/ml/inference/base.py | 38 ++++--- .../apache_beam/ml/inference/base_test.py | 99 +++++++++++++++++++ 3 files changed, 122 insertions(+), 16 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index d9ca6f0b83dd..ad0c0d483165 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -105,6 +105,7 @@ * (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)). * (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)). diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index ce014e4b04f9..9f6c0b926e41 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -2078,6 +2078,9 @@ def __init__( self._model = None self._metrics_namespace = metrics_namespace self._load_model_at_runtime = load_model_at_runtime + # The side input value that the currently loaded model reflects: a model + # id from ModelMetadata, a list[KeyModelPathMapping], or None if no side + # input value has been applied. Used to skip redundant model updates. self._side_input_path = None # _model_tag is the original tag passed in. # _cur_tag is the tag of the actually loaded model @@ -2149,11 +2152,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 = ''): @@ -2237,26 +2240,29 @@ 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 si_model_metadata or + isinstance(si_model_metadata, beam.pvalue.EmptySideInput)): + # No side input value: keep the loaded model, if any. if (not self._model_metadata.is_valid_tag(self._cur_tag) or self._model is 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) + needs_update = ( + self._model is None or # No model loaded yet. + self._side_input_path != si_model_metadata or # Mapping changed. + not self._model_metadata.is_valid_tag(self._cur_tag)) # Invalidated. + if needs_update: + self.update_model(si_model_metadata) + self._side_input_path = list(si_model_metadata) elif self._side_input_path != si_model_metadata.model_id: - self._side_input_path = 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) diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index 64f5c3b1525d..99126ca8eacd 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/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): @@ -1440,6 +1475,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 @@ -1722,6 +1792,35 @@ 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_run_inference_side_input_in_batch_multi_process_shared(self): first_ts = math.floor(time.time()) - 30 interval = 7 From aa4ee5d457506730eca42340cb6102b949ae4c90 Mon Sep 17 00:00:00 2001 From: Valentyn Tymofieiev Date: Mon, 19 Oct 2026 12:30:59 -0700 Subject: [PATCH 2/5] Drive-by changes.md updates. --- CHANGES.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index ad0c0d483165..3bad51f77e23 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -94,7 +94,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 @@ -106,6 +105,8 @@ * (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)). @@ -117,7 +118,7 @@ * (Python) `ReadFromText` with `escapechar` no longer skips a delimiter that directly follows an escaped delimiter, which merged two records into one ([#40459](https://github.com/apache/beam/issues/40459)). * (Python) DataFrame API: `agg` with a list that includes `'count'` or `'size'` now labels those rows `count`/`size` instead of `sum` ([#40504](https://github.com/apache/beam/issues/40504)). * (Python) DataFrame API: `Series.value_counts(bins=...)` no longer always fails with a `TypeError` ([#40503](https://github.com/apache/beam/issues/40503)). -* Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). + ## Security Fixes From 3a6a20f825899b9896dad6d4aead7f309a608529 Mon Sep 17 00:00:00 2001 From: Valentyn Tymofieiev Date: Mon, 19 Oct 2026 12:49:31 -0700 Subject: [PATCH 3/5] Also reload the model previously supplied via side input if it was invalidated. --- CHANGES.md | 1 - sdks/python/apache_beam/ml/inference/base.py | 3 ++- .../apache_beam/ml/inference/base_test.py | 23 +++++++++++++++++++ 3 files changed, 25 insertions(+), 2 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 3bad51f77e23..e140b82e604f 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -119,7 +119,6 @@ * (Python) DataFrame API: `agg` with a list that includes `'count'` or `'size'` now labels those rows `count`/`size` instead of `sum` ([#40504](https://github.com/apache/beam/issues/40504)). * (Python) DataFrame API: `Series.value_counts(bins=...)` no longer always fails with a `TypeError` ([#40503](https://github.com/apache/beam/issues/40503)). - ## Security Fixes * Fixed [CVE-YYYY-NNNN](https://www.cve.org/CVERecord?id=CVE-YYYY-NNNN) (Java/Python/Go) ([#X](https://github.com/apache/beam/issues/X)). diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 9f6c0b926e41..4e5ba938e7e2 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -2258,7 +2258,8 @@ def process( if needs_update: self.update_model(si_model_metadata) self._side_input_path = list(si_model_metadata) - elif self._side_input_path != si_model_metadata.model_id: + elif (self._side_input_path != si_model_metadata.model_id or # New model. + not self._model_metadata.is_valid_tag(self._cur_tag)): # Invalidated. self._metrics_collector = self.get_metrics_collector( prefix=si_model_metadata.model_name) self.update_model(si_model_metadata.model_id) diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index 99126ca8eacd..4fb06d84b22c 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -1821,6 +1821,29 @@ def test_run_inference_side_input_per_key_models_used_from_first_batch(self): 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 From 3cc35f6f589d7cd219f4775e4090f24939f645ca Mon Sep 17 00:00:00 2001 From: Valentyn Tymofieiev Date: Mon, 19 Oct 2026 13:11:00 -0700 Subject: [PATCH 4/5] review feedback: Also reload models if they were invalidated in the remaining branch --- sdks/python/apache_beam/ml/inference/base.py | 24 ++++++++++++-------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 4e5ba938e7e2..0b542889bff0 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -2243,8 +2243,10 @@ def process( if (not si_model_metadata or isinstance(si_model_metadata, beam.pvalue.EmptySideInput)): # No side input value: keep the loaded model, if any. - if (not self._model_metadata.is_valid_tag(self._cur_tag) or - self._model is None): + needs_update = ( + self._model is None or # No model loaded yet. + not self._model_metadata.is_valid_tag(self._cur_tag)) # Invalidated. + if needs_update: self.update_model(side_input_model_path=None) return self._run_inference(batch, inference_args) @@ -2253,17 +2255,21 @@ def process( # TODO(https://github.com/apache/beam/issues/27628): Update metrics here needs_update = ( self._model is None or # No model loaded yet. - self._side_input_path != si_model_metadata or # Mapping changed. + self._side_input_path != si_model_metadata or # New mapping. not self._model_metadata.is_valid_tag(self._cur_tag)) # Invalidated. if needs_update: self.update_model(si_model_metadata) self._side_input_path = list(si_model_metadata) - elif (self._side_input_path != si_model_metadata.model_id or # New model. - not self._model_metadata.is_valid_tag(self._cur_tag)): # Invalidated. - self._metrics_collector = self.get_metrics_collector( - prefix=si_model_metadata.model_name) - self.update_model(si_model_metadata.model_id) - self._side_input_path = si_model_metadata.model_id + else: + needs_update = ( + self._model is None or # No model loaded yet. + self._side_input_path != si_model_metadata.model_id or # New model. + not self._model_metadata.is_valid_tag(self._cur_tag)) # Invalidated. + if needs_update: + self._metrics_collector = self.get_metrics_collector( + prefix=si_model_metadata.model_name) + self.update_model(si_model_metadata.model_id) + self._side_input_path = si_model_metadata.model_id return self._run_inference(batch, inference_args) From f7f88f645410a04c4637b8f7e0a29a8efc021bd2 Mon Sep 17 00:00:00 2001 From: Valentyn Tymofieiev Date: Tue, 20 Oct 2026 18:38:14 -0700 Subject: [PATCH 5/5] Address comments, add a helper. --- sdks/python/apache_beam/ml/inference/base.py | 49 +++++++++++--------- 1 file changed, 26 insertions(+), 23 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 0b542889bff0..33c802ab7840 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -2078,10 +2078,9 @@ def __init__( self._model = None self._metrics_namespace = metrics_namespace self._load_model_at_runtime = load_model_at_runtime - # The side input value that the currently loaded model reflects: a model - # id from ModelMetadata, a list[KeyModelPathMapping], or None if no side - # input value has been applied. Used to skip redundant model updates. - 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 @@ -2186,6 +2185,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: @@ -2243,33 +2258,21 @@ def process( if (not si_model_metadata or isinstance(si_model_metadata, beam.pvalue.EmptySideInput)): # No side input value: keep the loaded model, if any. - needs_update = ( - self._model is None or # No model loaded yet. - not self._model_metadata.is_valid_tag(self._cur_tag)) # Invalidated. - if needs_update: + 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, list) and hasattr(si_model_metadata[0], 'keys'): # TODO(https://github.com/apache/beam/issues/27628): Update metrics here - needs_update = ( - self._model is None or # No model loaded yet. - self._side_input_path != si_model_metadata or # New mapping. - not self._model_metadata.is_valid_tag(self._cur_tag)) # Invalidated. - if needs_update: + if self._should_update_model(si_model_metadata): self.update_model(si_model_metadata) self._side_input_path = list(si_model_metadata) - else: - needs_update = ( - self._model is None or # No model loaded yet. - self._side_input_path != si_model_metadata.model_id or # New model. - not self._model_metadata.is_valid_tag(self._cur_tag)) # Invalidated. - if needs_update: - self._metrics_collector = self.get_metrics_collector( - prefix=si_model_metadata.model_name) - self.update_model(si_model_metadata.model_id) - self._side_input_path = si_model_metadata.model_id + elif self._should_update_model(si_model_metadata.model_id): + self._metrics_collector = self.get_metrics_collector( + prefix=si_model_metadata.model_name) + self.update_model(si_model_metadata.model_id) + self._side_input_path = si_model_metadata.model_id return self._run_inference(batch, inference_args)