From dc9d26a943d26d70434f6c5af8f7be9989ff3d92 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Fri, 12 Jun 2026 04:30:35 +0000 Subject: [PATCH 1/4] Make ModelManager import more robust --- sdks/python/apache_beam/ml/inference/base.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index b2441281dd18..1f3f1b7321f0 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -68,10 +68,12 @@ try: # pylint: disable=wrong-import-order, wrong-import-position import resource +except ImportError: + resource = None # type: ignore[assignment] +try: from apache_beam.ml.inference.model_manager import ModelManager except ImportError: - resource = None # type: ignore[assignment] ModelManager = None # type: ignore[assignment] _NANOSECOND_TO_MILLISECOND = 1_000_000 @@ -1997,6 +1999,10 @@ def load(): # Ensure the tag we're loading is valid, if not replace it with a valid tag self._cur_tag = self._model_metadata.get_valid_tag(model_tag) if self.use_model_manager: + if ModelManager is None: + raise ImportError( + "Model Manager is not available. Please ensure that the " + "all required packages for inference is installed and up to date.") logging.info("Using Model Manager to manage models automatically.") model_manager = multi_process_shared.MultiProcessShared( lambda: ModelManager(**self._model_manager_args), From e4882e83697b8b640e2f56e717862ea2e75d6c27 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Fri, 12 Jun 2026 04:38:12 +0000 Subject: [PATCH 2/4] Make ModelManager import more robust 2 --- sdks/python/apache_beam/ml/inference/base.py | 28 ++++++++++++++------ 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 1f3f1b7321f0..393e77ec8a73 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -71,10 +71,23 @@ except ImportError: resource = None # type: ignore[assignment] -try: - from apache_beam.ml.inference.model_manager import ModelManager -except ImportError: - ModelManager = None # type: ignore[assignment] + +def _try_import_model_manager(throw_error: bool = True): + try: + from apache_beam.ml.inference.model_manager import ModelManager + return ModelManager + except ImportError as e: + if throw_error: + raise ImportError( + "Model Manager is not available. Please ensure that the " + "all required packages for inference is installed and up to date." + ) from e + else: + return None + + +# ModelManager is an optional dependency so we don't throw an error here. +ModelManager = _try_import_model_manager(throw_error=False) _NANOSECOND_TO_MILLISECOND = 1_000_000 _NANOSECOND_TO_MICROSECOND = 1_000 @@ -1999,10 +2012,9 @@ def load(): # Ensure the tag we're loading is valid, if not replace it with a valid tag self._cur_tag = self._model_metadata.get_valid_tag(model_tag) if self.use_model_manager: - if ModelManager is None: - raise ImportError( - "Model Manager is not available. Please ensure that the " - "all required packages for inference is installed and up to date.") + # Force an import here to avoid missing ModelManager when needed. + # Throw an error if ModelManager is not available since it's required for this code path. + ModelManager = _try_import_model_manager(throw_error=True) logging.info("Using Model Manager to manage models automatically.") model_manager = multi_process_shared.MultiProcessShared( lambda: ModelManager(**self._model_manager_args), From 9f606d598c875c17b16d4e7091ccd03227b6d8ae Mon Sep 17 00:00:00 2001 From: "RuiLong J." Date: Thu, 11 Jun 2026 22:20:45 -0700 Subject: [PATCH 3/4] Apply suggestion from @gemini-code-assist[bot] Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- sdks/python/apache_beam/ml/inference/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 393e77ec8a73..18cbf04d3202 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -79,8 +79,8 @@ def _try_import_model_manager(throw_error: bool = True): except ImportError as e: if throw_error: raise ImportError( - "Model Manager is not available. Please ensure that the " - "all required packages for inference is installed and up to date." + "Model Manager is not available. Please ensure that " + "all required packages for inference are installed and up to date." ) from e else: return None From b2578df32f709567222601ddb19e4cb5cd01f541 Mon Sep 17 00:00:00 2001 From: AMOOOMA Date: Sat, 13 Jun 2026 00:37:30 +0000 Subject: [PATCH 4/4] Add model tag to annotations as well --- sdks/python/apache_beam/ml/inference/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 18cbf04d3202..97dd22103969 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -1458,6 +1458,7 @@ def annotations(self): 'model_handler_type': ( f'{self._model_handler.__class__.__module__}' f'.{self._model_handler.__class__.__qualname__}'), + 'model_identifier': self._model_tag, **super().annotations() }