tvalentyn opened a new issue, #40469: URL: https://github.com/apache/beam/issues/40469
### What happened? (Description is mostly AI-generated, but seems reasonable) With use_model_manager=True, _RunInferenceDoFn acquires models from ModelManager under the tag of the loaded model, self._cur_tag ([wrapper](https://github.com/apache/beam/blob/69bdf5f29174adebe19006e1397acd6681ffce9d/sdks/python/apache_beam/ml/inference/base.py#L2038-L2039), [acquire](https://github.com/apache/beam/blob/69bdf5f29174adebe19006e1397acd6681ffce9d/sdks/python/apache_beam/ml/inference/base.py#L1932-L1935)). But it [releases](https://github.com/apache/beam/blob/69bdf5f29174adebe19006e1397acd6681ffce9d/sdks/python/apache_beam/ml/inference/base.py#L2104-L2106) them under the transform's original tag, self._model_tag: The two tags differ after a model update from model_metadata_pcoll or watch_model_pattern, because [_cur_tag becomes the new model path](https://github.com/apache/beam/blob/69bdf5f29174adebe19006e1397acd6681ffce9d/sdks/python/apache_beam/ml/inference/base.py#L2022-L2026). [ModelManager.release_model()](https://github.com/apache/beam/blob/69bdf5f29174adebe19006e1397acd6681ffce9d/sdks/python/apache_beam/ml/inference/model_manager.py#L590-L620) then: - files the idle instance under the wrong tag, so acquire_model() never reuses it; - stores the model's measured memory under the wrong tag, so the real tag stays "unknown"; - doesn't decrement _active_counts for the real tag. From then on, every acquire_model() sees an unknown tag and enters [isolation mode](https://github.com/apache/beam/blob/69bdf5f29174adebe19006e1397acd6681ffce9d/sdks/python/apache_beam/ml/inference/model_manager.py#L408-L424). Isolation mode deletes all models the ModelManager holds, including other transforms' models, and then loads the model again in a new process. Result: the model is loaded once per batch, and every other ModelManager model is flushed each time. Other issue that is worth fixing in the same PR: [_SharedModelWrapper.all_models()](https://github.com/apache/beam/blob/69bdf5f29174adebe19006e1397acd6681ffce9d/sdks/python/apache_beam/ml/inference/base.py#L1948-L1951) calls self.models.all_models()[self.model_tag]. But [ModelManager.all_models()](https://github.com/apache/beam/blob/69bdf5f29174adebe19006e1397acd6681ffce9d/sdks/python/apache_beam/ml/inference/model_manager.py#L404-L405) takes the tag as an argument. ### Issue Priority Priority: 2 (default / most bugs should be filed as P2) ### Issue Components - [ ] 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]
