(beam) branch master updated: Make ModelManager import more robust (#38936)

[damccorm]

This is an automated email from the ASF dual-hosted git repository.

damccorm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 1d703fe904f Make ModelManager import more robust (#38936)
1d703fe904f is described below

commit 1d703fe904ff685bcd48c1507d4e5779b0f8bbce
Author: RuiLong J. <ruilong.ji...@outlook.com>
AuthorDate: Wed Jun 17 11:58:08 2026 -0700

    Make ModelManager import more robust (#38936)
    
    * Make ModelManager import more robust
    
    * Make ModelManager import more robust 2
    
    * Apply suggestion from @gemini-code-assist[bot]
    
    Co-authored-by: gemini-code-assist[bot] 
<176961590+gemini-code-assist[bot]@users.noreply.github.com>
    
    * Add model tag to annotations as well
    
    * Change to not use function import
    
    ---------
    
    Co-authored-by: gemini-code-assist[bot] 
<176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 sdks/python/apache_beam/ml/inference/base.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/sdks/python/apache_beam/ml/inference/base.py 
b/sdks/python/apache_beam/ml/inference/base.py
index b2441281dd1..f81382bbeec 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -1443,6 +1443,7 @@ class 
RunInference(beam.PTransform[beam.PCollection[Union[ExampleT,
         'model_handler_type': (
             f'{self._model_handler.__class__.__module__}'
             f'.{self._model_handler.__class__.__qualname__}'),
+        'model_identifier': self._model_tag,
         **super().annotations()
     }
 
@@ -1997,6 +1998,11 @@ class _RunInferenceDoFn(beam.DoFn, Generic[ExampleT, 
PredictionT]):
     # 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:
+      # 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.
+      global ModelManager
+      if ModelManager is None:
+        from apache_beam.ml.inference.model_manager import ModelManager
       logging.info("Using Model Manager to manage models automatically.")
       model_manager = multi_process_shared.MultiProcessShared(
           lambda: ModelManager(**self._model_manager_args),