Skip model manager import if error

Author: AMOOOMA

sdks/python/apache_beam/ml/inference/base.py

@@ -59,13 +59,14 @@
 from apache_beam.utils import multi_process_shared
 from apache_beam.utils import retry
 from apache_beam.utils import shared
-from apache_beam.ml.inference.model_manager import ModelManager
 
 try:
   # pylint: disable=wrong-import-order, wrong-import-position
   import resource
+  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
 _NANOSECOND_TO_MICROSECOND = 1_000

sdks/python/apache_beam/ml/inference/base_test.py

@@ -317,6 +317,15 @@ def validate_inference_args(self, inference_args):
     pass
 
 
+def try_import_model_manager():
+  try:
+    # pylint: disable=unused-import
+    from apache_beam.ml.inference.model_manager import ModelManager
+    return True
+  except ImportError:
+    return False
+
+
 class RunInferenceBaseTest(unittest.TestCase):
   def test_run_inference_impl_simple_examples(self):
     with TestPipeline() as pipeline:
@@ -1894,6 +1903,7 @@ def test_model_status_provides_valid_garbage_collection(self):
 
     self.assertEqual(0, len(tags))
 
+  @skipIf(try_import_model_manager() is False, 'Model Manager not available')
   def test_run_inference_impl_with_model_manager(self):
     with TestPipeline() as pipeline:
       examples = [1, 5, 3, 10]
@@ -1903,6 +1913,7 @@ def test_run_inference_impl_with_model_manager(self):
           SimpleFakeModelHanlder(), use_model_manager=True)
       assert_that(actual, equal_to(expected), label='assert:inferences')
 
+  @skipIf(try_import_model_manager() is False, 'Model Manager not available')
   def test_run_inference_impl_with_model_manager_args(self):
     with TestPipeline() as pipeline:
       examples = [1, 5, 3, 10]