tvalentyn opened a new issue, #40468: URL: https://github.com/apache/beam/issues/40468
### What happened? The intent of `max_models_per_worker` hint is to limit the [number of models that will be loaded per SDK process](https://github.com/apache/beam/blob/dabcf50ffbf532bb30a1233927cd1edb6ae067bb/website/www/site/content/en/documentation/ml/multi-model-pipelines.md?plain=1#L126). The logic to increment allowed max number of workers https://github.com/apache/beam/blob/dabcf50ffbf532bb30a1233927cd1edb6ae067bb/sdks/python/apache_beam/ml/inference/base.py#L852-L857 appears to be flawed since: 1) The lock acquisition always succeeds (it's a new lock instance) 2) If there is more than 1 process bundle descriptor over the life time of the SDK process, we might have more than 1 copy of the deserialized RunInference DoFn with unpickled ModelHandler instance, which won't persist `self._max_models_per_worker_hint = None` from a prior initialization: AI repro: ``` from apache_beam.internal import pickler from apache_beam.ml.inference import base class Model: def predict(self, x): return x class Handler(base.ModelHandler): def load_model(self): return Model() def run_inference(self, batch, model, inference_args=None): return [model.predict(x) for x in batch] mhs = [base.KeyModelMapping([k], Handler()) for k in ('a', 'b', 'c')] keyed_handler = base.KeyedModelHandler(mhs, max_models_per_worker_hint=1) # RunInference shares one _ModelHandlerManager per transform across all # DoFn instances (and processes) via MultiProcessShared. manager = keyed_handler.load_model() # 5 DoFn instances in one process (harness threads, re-created bundle # processors); each deserializes its own copy of the model handler. for _ in range(5): handler_copy = pickler.roundtrip(keyed_handler) handler_copy.override_metrics('ns') handler_copy.run_inference([('a', 1), ('b', 2), ('c', 3)], manager) print('model limit:', manager._max_models) # 5, expected 1 print('models in memory:', len(manager._tag_map)) # 3 ``` ### Issue Priority Priority: 2 (default / most bugs should be filed as P2) ### Issue Components - [x] Component: Python SDK - [ ] Component: Java SDK - [ ] Component: Go SDK - [ ] Component: Typescript SDK - [ ] Component: IO connector - [ ] Component: Beam YAML - [ ] Component: Beam examples - [ ] Component: Beam playground - [ ] Component: Beam katas - [ ] Component: Website - [ ] Component: Infrastructure - [ ] Component: Spark Runner - [ ] Component: Flink Runner - [ ] Component: Prism Runner - [ ] Component: Twister2 Runner - [ ] Component: Hazelcast Jet Runner - [ ] Component: Google Cloud Dataflow Runner -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
