AMOOOMA opened a new issue, #40539: URL: https://github.com/apache/beam/issues/40539
### What would you like to happen? Currently, `RunInference(..., use_model_manager=True)` does not properly integrate with a multi-model `KeyedModelHandler` (i.e. when `KeyedModelHandler` is initialized with a `list[KeyModelMapping]`): 1. **Competing model managers**: When `KeyedModelHandler` wraps multiple `KeyModelMapping`s, `KeyedModelHandler.load_model()` returns a `_ModelHandlerManager` instance rather than loading an ML model directly. Inside `KeyedModelHandler.run_inference()`, `_ModelHandlerManager.load(key)` spawns and caches per-key models via `MultiProcessShared` using its own count-based LRU (`max_models_per_worker_hint`), bypassing `ModelManager`'s GPU memory monitoring, resource estimation, and eviction. 2. **Per-key side-input model updates (`KeyModelPathMapping`)**: In `_RunInferenceDoFn._load_model()`, when `side_input_model_path` is a `list[KeyModelPathMapping]`, it iterates over `self._model.all_models()` and calls `self._model_handler.update_model_paths(m, side_input_model_path)`. With `use_model_manager=True`, `ModelManager.all_models(tag)` returns `unique_tag` strings rather than `_ModelHandlerManager` proxies, and `ModelManager` is not aware of the per-key cohort updates. We should either: - Integrate multi-model `KeyedModelHandler` with `ModelManager` so per-key models (and `KeyModelPathMapping` updates) are managed directly by `ModelManager`, or - Explicitly validate/document the interaction when `use_model_manager=True` is used with a multi-model `KeyedModelHandler`. Related: #40469 -- 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]
