Catch a broad Exception when setting _BeamPipeline.

PiperOrigin-RevId: 617013620
tensorflow · Mar 19, 2024 · 281ddf5 · 281ddf5
1 parent 2283b52
commit 281ddf5
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/tfx/dsl/component/experimental/annotations.py b/tfx/dsl/component/experimental/annotations.py
@@ -27,7 +27,7 @@
   import apache_beam as beam  # pytype: disable=import-error  # pylint: disable=g-import-not-at-top
 
   _BeamPipeline = beam.Pipeline
-except ModuleNotFoundError:
+except Exception:  # pylint: disable=broad-exception-caught
   beam = None
   _BeamPipeline = Any
 

diff --git a/tfx/dsl/component/experimental/decorators.py b/tfx/dsl/component/experimental/decorators.py
@@ -37,7 +37,7 @@
 try:
   import apache_beam as beam  # pytype: disable=import-error  # pylint: disable=g-import-not-at-top
   _BeamPipeline = beam.Pipeline
-except ModuleNotFoundError:
+except Exception:  # pylint: disable=broad-exception-caught
   beam = None
   _BeamPipeline = Any
 

diff --git a/tfx/dsl/component/experimental/function_parser.py b/tfx/dsl/component/experimental/function_parser.py
@@ -35,7 +35,7 @@
 try:
   import apache_beam as beam  # pytype: disable=import-error  # pylint: disable=g-import-not-at-top
   _BeamPipeline = beam.Pipeline
-except ModuleNotFoundError:
+except Exception:  # pylint: disable=broad-exception-caught
   beam = None
   _BeamPipeline = Any
 

diff --git a/tfx/dsl/components/base/base_beam_executor.py b/tfx/dsl/components/base/base_beam_executor.py
@@ -28,7 +28,7 @@
 try:
   import apache_beam as beam  # pytype: disable=import-error  # pylint: disable=g-import-not-at-top
   _BeamPipeline = beam.Pipeline
-except ModuleNotFoundError:
+except Exception:  # pylint: disable=broad-exception-caught
   beam = None
   _BeamPipeline = Any