diff --git a/CHANGES.md b/CHANGES.md index b9accf9cf88f..66ee680b3ab1 100644 --- a/CHANGES.md +++ b/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)). diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 3d2f2657c0a8..4c2d9cb4c80e 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/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) diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index d1c37811d9d3..5d6c5bd517c7 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): @@ -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