Use pytest.importorskip for Vertex AI tests & remove vertex_ai_skip helper

Author: aIbrahiim

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

@@ -27,7 +27,8 @@
 from apache_beam.io.filesystems import FileSystems
 from apache_beam.ml.inference.base import RunInference
 from apache_beam.testing.test_pipeline import TestPipeline
-from apache_beam.testing.vertex_ai_skip import skip_if_vertex_ai_disabled
+
+pytest.importorskip("vertexai", reason="Vertex AI dependencies not available")
 
 # pylint: disable=ungrouped-imports
 try:
@@ -54,7 +55,6 @@
 _INVOKE_OUTPUT_DIR = "gs://apache-beam-ml/testing/outputs/vertex_invoke"
 
 
-@skip_if_vertex_ai_disabled
 @pytest.mark.vertex_ai_postcommit
 class VertexAIInference(unittest.TestCase):
   def test_vertex_ai_run_flower_image_classification(self):

sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py

@@ -34,7 +34,8 @@
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that
 from apache_beam.testing.util import equal_to
-from apache_beam.testing.vertex_ai_skip import skip_if_vertex_ai_disabled
+
+pytest.importorskip("vertexai", reason="Vertex AI dependencies not available")
 
 # pylint: disable=ungrouped-imports
 try:
@@ -43,9 +44,10 @@
   from apache_beam.ml.rag.embeddings.vertex_ai import VertexAIImageEmbeddings
   from apache_beam.ml.rag.embeddings.vertex_ai import VertexAITextEmbeddings
   from apache_beam.ml.rag.embeddings.vertex_ai import _create_image_adapter
-  VERTEX_AI_AVAILABLE = True
 except ImportError:
-  VERTEX_AI_AVAILABLE = False
+  VertexAIImageEmbeddings = None  # type: ignore
+  VertexAITextEmbeddings = None  # type: ignore
+  _create_image_adapter = None  # type: ignore
 
 
 def chunk_approximately_equals(expected, actual):
@@ -62,8 +64,6 @@ def chunk_approximately_equals(expected, actual):
 
 
 @pytest.mark.vertex_ai_postcommit
-@unittest.skipIf(
-    not VERTEX_AI_AVAILABLE, "Vertex AI dependencies not available")
 class VertexAITextEmbeddingsTest(unittest.TestCase):
   def setUp(self):
     self.artifact_location = tempfile.mkdtemp(prefix='vertex_ai_')
@@ -117,8 +117,6 @@ def test_embedding_pipeline(self):
           embeddings, equal_to(expected, equals_fn=chunk_approximately_equals))
 
 
-@unittest.skipIf(
-    not VERTEX_AI_AVAILABLE, "Vertex AI dependencies not available")
 class VertexAIImageAdapterTest(unittest.TestCase):
   def test_image_adapter_missing_content(self):
     adapter = _create_image_adapter()
@@ -150,10 +148,7 @@ def test_image_adapter_output(self):
     self.assertEqual(result[0].embedding.dense_embedding, [0.1, 0.2, 0.3])
 
 
-@skip_if_vertex_ai_disabled
 @pytest.mark.vertex_ai_postcommit
-@unittest.skipIf(
-    not VERTEX_AI_AVAILABLE, "Vertex AI dependencies not available")
 class VertexAIImageEmbeddingsTest(unittest.TestCase):
   def setUp(self):
     self.artifact_location = tempfile.mkdtemp(prefix='vertex_ai_img_')

…/transforms/embeddings/vertex_ai_test.py …ansforms/embeddings/vertex_ai_it_test.pysdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py renamed to sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_it_test.py sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py renamed to sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_it_test.py

@@ -26,7 +26,8 @@
 from apache_beam.ml.inference.base import RunInference
 from apache_beam.ml.transforms import base
 from apache_beam.ml.transforms.base import MLTransform
-from apache_beam.testing.vertex_ai_skip import skip_if_vertex_ai_disabled
+
+pytest.importorskip("vertexai", reason="Vertex AI dependencies not available")
 
 # pylint: disable=ungrouped-imports
 # isort: off
@@ -61,10 +62,7 @@
 model_name: str = "text-embedding-005"
 
 
-@skip_if_vertex_ai_disabled
 @pytest.mark.vertex_ai_postcommit
-@unittest.skipIf(
-    VertexAITextEmbeddings is None, 'Vertex AI Python SDK is not installed.')
 class VertexAIEmbeddingsTest(unittest.TestCase):
   def setUp(self) -> None:
     self.artifact_location = tempfile.mkdtemp(prefix='_vertex_ai_test')
@@ -266,9 +264,7 @@ def test_mltransform_to_ptransform_with_vertex(self):
           ptransform_list[i]._model_handler._underlying.model_name, model_name)
 
 
-@skip_if_vertex_ai_disabled
-@unittest.skipIf(
-    VertexAIImageEmbeddings is None, 'Vertex AI Python SDK is not installed.')
+@pytest.mark.vertex_ai_postcommit
 class VertexAIImageEmbeddingsTest(unittest.TestCase):
   def setUp(self) -> None:
     self.artifact_location = tempfile.mkdtemp(prefix='_vertex_ai_image_test')
@@ -314,11 +310,7 @@ def _make_text_chunk(input: str) -> Chunk:
   return Chunk(content=Content(text=input))
 
 
-@skip_if_vertex_ai_disabled
 @pytest.mark.vertex_ai_postcommit
-@unittest.skipIf(
-    VertexAIMultiModalEmbeddings is None,
-    'Vertex AI Python SDK is not installed.')
 class VertexAIMultiModalEmbeddingsTest(unittest.TestCase):
   def setUp(self) -> None:
     self.artifact_location = tempfile.mkdtemp(

sdks/python/apache_beam/testing/vertex_ai_skip.py

@@ -504,7 +504,7 @@ task vertexAIInferenceTest {
       def cmdArgs = mapToArgString(argMap)
       exec {
         executable 'sh'
-        args '-c', ". ${envdir}/bin/activate && RUN_VERTEX_AI_TESTS=1 ${runScriptsDir}/run_integration_test.sh $cmdArgs"
+        args '-c', ". ${envdir}/bin/activate && ${runScriptsDir}/run_integration_test.sh $cmdArgs"
       }
     }
 }

sdks/python/test-suites/dataflow/common.gradle

@@ -159,10 +159,6 @@ toxTask "testPy310transformers-447", "py310-transformers-447", "${posargs}"
 test.dependsOn "testPy310transformers-447"
 postCommitPyDep.dependsOn "testPy310transformers-447"
 
-toxTask "testPy310transformers-448", "py310-transformers-448", "${posargs}"
-test.dependsOn "testPy310transformers-448"
-postCommitPyDep.dependsOn "testPy310transformers-448"
-
 toxTask "testPy310transformers-latest", "py310-transformers-latest", "${posargs}"
 test.dependsOn "testPy310transformers-latest"
 postCommitPyDep.dependsOn "testPy310transformers-latest"

sdks/python/test-suites/tox/py310/build.gradle

@@ -512,7 +512,7 @@ commands =
   # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
   /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_xgboost {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'
 
-[testenv:py{310,311}-transformers-{428,447,448,latest}]
+[testenv:py{310,311}-transformers-{428,447,latest}]
 deps =
   # Environment dependencies are defined in the `setenv` section and installed in the `commands` section.
 pip_pre = False
@@ -522,7 +522,6 @@ setenv =
     # sentence-transformers 2.2.2 is the latest version that supports transformers 4.28.x
     428: DEPS = sentence-transformers==2.2.2 'transformers>=4.28.0,<4.29.0' 'torch>=1.9.0,<1.14.0'
     447: DEPS = 'transformers>=4.47.0,<4.48.0' 'torch>=1.9.0,<1.14.0'
-    448: DEPS = 'transformers>=4.48.0,<4.49.0' 'torch>=1.9.0,<1.14.0'
     latest: DEPS = 'transformers>=4.55.0' 'torch>=2.0.0' 'accelerate>=1.6.0'
 commands =
   /bin/sh -c "pip install .[{extras}] {env:DEPS} {env:COMMON_DEPS}"