From 590684f81b06373ac53826b51cd9483923fac108 Mon Sep 17 00:00:00 2001 From: aIbrahiim Date: Fri, 27 Feb 2026 22:56:41 +0200 Subject: [PATCH 1/6] Exclude vertex_ai_postcommit tests from embeddings dep suite --- .../ml/rag/embeddings/vertex_ai_test.py | 4 ++ .../transforms/embeddings/vertex_ai_test.py | 5 ++ .../python/test-suites/tox/py310/build.gradle | 46 +++++++++---------- sdks/python/tox.ini | 9 ++-- 4 files changed, 37 insertions(+), 27 deletions(-) diff --git a/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py b/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py index 1a479c8e7dba..06cc8a241763 100644 --- a/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py +++ b/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py @@ -23,6 +23,8 @@ import unittest import zlib +import pytest + import apache_beam as beam from apache_beam.ml.rag.types import Chunk from apache_beam.ml.rag.types import Content @@ -58,6 +60,7 @@ def chunk_approximately_equals(expected, actual): all(isinstance(x, float) for x in actual.embedding.dense_embedding)) +@pytest.mark.vertex_ai_postcommit @unittest.skipIf( not VERTEX_AI_AVAILABLE, "Vertex AI dependencies not available") class VertexAITextEmbeddingsTest(unittest.TestCase): @@ -146,6 +149,7 @@ def test_image_adapter_output(self): self.assertEqual(result[0].embedding.dense_embedding, [0.1, 0.2, 0.3]) +@pytest.mark.vertex_ai_postcommit @unittest.skipIf( not VERTEX_AI_AVAILABLE, "Vertex AI dependencies not available") class VertexAIImageEmbeddingsTest(unittest.TestCase): diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py b/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py index 50507c54e36d..6e21b112f5a4 100644 --- a/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py +++ b/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py @@ -20,6 +20,8 @@ import unittest import uuid +import pytest + import apache_beam as beam from apache_beam.ml.inference.base import RunInference from apache_beam.ml.transforms import base @@ -58,6 +60,7 @@ model_name: str = "text-embedding-005" +@pytest.mark.vertex_ai_postcommit @unittest.skipIf( VertexAITextEmbeddings is None, 'Vertex AI Python SDK is not installed.') class VertexAIEmbeddingsTest(unittest.TestCase): @@ -261,6 +264,7 @@ def test_mltransform_to_ptransform_with_vertex(self): ptransform_list[i]._model_handler._underlying.model_name, model_name) +@pytest.mark.vertex_ai_postcommit @unittest.skipIf( VertexAIImageEmbeddings is None, 'Vertex AI Python SDK is not installed.') class VertexAIImageEmbeddingsTest(unittest.TestCase): @@ -308,6 +312,7 @@ def _make_text_chunk(input: str) -> Chunk: return Chunk(content=Content(text=input)) +@pytest.mark.vertex_ai_postcommit @unittest.skipIf( VertexAIMultiModalEmbeddings is None, 'Vertex AI Python SDK is not installed.') diff --git a/sdks/python/test-suites/tox/py310/build.gradle b/sdks/python/test-suites/tox/py310/build.gradle index 751faa682ae3..217564257673 100644 --- a/sdks/python/test-suites/tox/py310/build.gradle +++ b/sdks/python/test-suites/tox/py310/build.gradle @@ -49,29 +49,9 @@ project.tasks.register("postCommitPyDep") {} // For versions that we would like to prioritize for testing, // for example versions released in a timeframe of last 1-2 years. -toxTask "testPy310pyarrow-9", "py310-pyarrow-9", "${posargs}" -test.dependsOn "testPy310pyarrow-9" -postCommitPyDep.dependsOn "testPy310pyarrow-9" - -toxTask "testPy310pyarrow-10", "py310-pyarrow-10", "${posargs}" -test.dependsOn "testPy310pyarrow-10" -postCommitPyDep.dependsOn "testPy310pyarrow-10" - -toxTask "testPy310pyarrow-11", "py310-pyarrow-11", "${posargs}" -test.dependsOn "testPy310pyarrow-11" -postCommitPyDep.dependsOn "testPy310pyarrow-11" - -toxTask "testPy310pyarrow-12", "py310-pyarrow-12", "${posargs}" -test.dependsOn "testPy310pyarrow-12" -postCommitPyDep.dependsOn "testPy310pyarrow-12" - -toxTask "testPy310pyarrow-13", "py310-pyarrow-13", "${posargs}" -test.dependsOn "testPy310pyarrow-13" -postCommitPyDep.dependsOn "testPy310pyarrow-13" - -toxTask "testPy310pyarrow-14", "py310-pyarrow-14", "${posargs}" -test.dependsOn "testPy310pyarrow-14" -postCommitPyDep.dependsOn "testPy310pyarrow-14" +toxTask "testPy310pyarrow-6", "py310-pyarrow-6", "${posargs}" +test.dependsOn "testPy310pyarrow-6" +postCommitPyDep.dependsOn "testPy310pyarrow-6" toxTask "testPy310pyarrow-15", "py310-pyarrow-15", "${posargs}" test.dependsOn "testPy310pyarrow-15" @@ -89,6 +69,26 @@ toxTask "testPy310pyarrow-18", "py310-pyarrow-18", "${posargs}" test.dependsOn "testPy310pyarrow-18" postCommitPyDep.dependsOn "testPy310pyarrow-18" +toxTask "testPy310pyarrow-19", "py310-pyarrow-19", "${posargs}" +test.dependsOn "testPy310pyarrow-19" +postCommitPyDep.dependsOn "testPy310pyarrow-19" + +toxTask "testPy310pyarrow-20", "py310-pyarrow-20", "${posargs}" +test.dependsOn "testPy310pyarrow-20" +postCommitPyDep.dependsOn "testPy310pyarrow-20" + +toxTask "testPy310pyarrow-21", "py310-pyarrow-21", "${posargs}" +test.dependsOn "testPy310pyarrow-21" +postCommitPyDep.dependsOn "testPy310pyarrow-21" + +toxTask "testPy310pyarrow-22", "py310-pyarrow-22", "${posargs}" +test.dependsOn "testPy310pyarrow-22" +postCommitPyDep.dependsOn "testPy310pyarrow-22" + +toxTask "testPy310pyarrow-23", "py310-pyarrow-23", "${posargs}" +test.dependsOn "testPy310pyarrow-23" +postCommitPyDep.dependsOn "testPy310pyarrow-23" + // Create a test task for each supported minor version of pandas toxTask "testPy310pandas-14", "py310-pandas-14", "${posargs}" test.dependsOn "testPy310pandas-14" diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini index 32be7b52b034..03fa34054c2d 100644 --- a/sdks/python/tox.ini +++ b/sdks/python/tox.ini @@ -515,13 +515,14 @@ commands = [testenv:py{310,311}-transformers-{428,447,448,latest}] deps = # Environment dependencies are defined in the `setenv` section and installed in the `commands` section. -extras = test,gcp,ml_test +pip_pre = False +extras = test setenv = COMMON_DEPS = tensorflow==2.12.0 protobuf==4.25.5 pip==25.0.1 # 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' - 455: DEPS = 'transformers>=4.55.0,<4.56.0' 'torch>=2.0.0,<2.1.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}" @@ -561,7 +562,7 @@ commands = /bin/sh -c "pip freeze | grep -E google-cloud-aiplatform" /bin/sh -c "pip freeze | grep -E transformers" # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories. - /bin/sh -c 'pytest apache_beam/ml/transforms/embeddings -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret' + /bin/sh -c 'pytest apache_beam/ml/transforms/embeddings -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m "not vertex_ai_postcommit" {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret' [testenv:py{310,311}-TFHubEmbeddings-{014,015}] @@ -579,7 +580,7 @@ commands = # Log aiplatform and its dependencies version for debugging /bin/sh -c "pip freeze | grep -E tensorflow" # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories. - bash {toxinidir}/scripts/run_pytest.sh {envname} 'apache_beam/ml/transforms/embeddings' + bash {toxinidir}/scripts/run_pytest.sh {envname} 'apache_beam/ml/transforms/embeddings' "-m not(vertex_ai_postcommit)" [testenv:py{310,312}-dill] extras = test,dill From 417922c2f04b506385b98d21f1e0bbdc5a61a777 Mon Sep 17 00:00:00 2001 From: aIbrahiim Date: Tue, 3 Mar 2026 15:09:17 +0200 Subject: [PATCH 2/6] centralize skipping vertex ai test --- .../inference/vertex_ai_inference_it_test.py | 5 ++- .../ml/rag/embeddings/vertex_ai_test.py | 3 ++ .../transforms/embeddings/vertex_ai_test.py | 5 ++- .../apache_beam/testing/vertex_ai_skip.py | 38 +++++++++++++++++++ .../python/test-suites/dataflow/common.gradle | 2 +- sdks/python/tox.ini | 4 +- 6 files changed, 51 insertions(+), 6 deletions(-) create mode 100644 sdks/python/apache_beam/testing/vertex_ai_skip.py diff --git a/sdks/python/apache_beam/ml/inference/vertex_ai_inference_it_test.py b/sdks/python/apache_beam/ml/inference/vertex_ai_inference_it_test.py index 11643992c392..8908477dae7f 100644 --- a/sdks/python/apache_beam/ml/inference/vertex_ai_inference_it_test.py +++ b/sdks/python/apache_beam/ml/inference/vertex_ai_inference_it_test.py @@ -27,6 +27,7 @@ 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 # pylint: disable=ungrouped-imports try: @@ -53,8 +54,9 @@ _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): - @pytest.mark.vertex_ai_postcommit def test_vertex_ai_run_flower_image_classification(self): output_file = '/'.join([_OUTPUT_DIR, str(uuid.uuid4()), 'output.txt']) @@ -73,7 +75,6 @@ def test_vertex_ai_run_flower_image_classification(self): test_pipeline.get_full_options_as_args(**extra_opts)) self.assertEqual(FileSystems().exists(output_file), True) - @pytest.mark.vertex_ai_postcommit @unittest.skipIf( not _INVOKE_ENDPOINT_ID, "Invoke endpoint not configured. Set _INVOKE_ENDPOINT_ID.") diff --git a/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py b/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py index 06cc8a241763..f66d1f9d9b7b 100644 --- a/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py +++ b/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py @@ -34,6 +34,7 @@ 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 # pylint: disable=ungrouped-imports try: @@ -60,6 +61,7 @@ def chunk_approximately_equals(expected, actual): all(isinstance(x, float) for x in actual.embedding.dense_embedding)) +@skip_if_vertex_ai_disabled @pytest.mark.vertex_ai_postcommit @unittest.skipIf( not VERTEX_AI_AVAILABLE, "Vertex AI dependencies not available") @@ -149,6 +151,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") diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py b/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py index 6e21b112f5a4..bc8ca4f0673e 100644 --- a/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py +++ b/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py @@ -26,6 +26,7 @@ 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 # pylint: disable=ungrouped-imports # isort: off @@ -60,6 +61,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.') @@ -264,7 +266,7 @@ def test_mltransform_to_ptransform_with_vertex(self): ptransform_list[i]._model_handler._underlying.model_name, model_name) -@pytest.mark.vertex_ai_postcommit +@skip_if_vertex_ai_disabled @unittest.skipIf( VertexAIImageEmbeddings is None, 'Vertex AI Python SDK is not installed.') class VertexAIImageEmbeddingsTest(unittest.TestCase): @@ -312,6 +314,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, diff --git a/sdks/python/apache_beam/testing/vertex_ai_skip.py b/sdks/python/apache_beam/testing/vertex_ai_skip.py new file mode 100644 index 000000000000..9207ea188e0a --- /dev/null +++ b/sdks/python/apache_beam/testing/vertex_ai_skip.py @@ -0,0 +1,38 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Centralized control for Vertex AI integration tests. + +Test modules should use skip_if_vertex_ai_disabled on classes or tests that +require Vertex AI Python dependencies to be installed. +""" + +import pytest + + +def _is_vertex_ai_available() -> bool: + """Return True if Vertex AI client dependencies are importable.""" + try: + import vertexai + except ImportError: + return False + return True + + +skip_if_vertex_ai_disabled = pytest.mark.skipif( + not _is_vertex_ai_available(), + reason='Vertex AI dependencies not available.') diff --git a/sdks/python/test-suites/dataflow/common.gradle b/sdks/python/test-suites/dataflow/common.gradle index 7d65da6ee3ad..b21efb2ae828 100644 --- a/sdks/python/test-suites/dataflow/common.gradle +++ b/sdks/python/test-suites/dataflow/common.gradle @@ -504,7 +504,7 @@ task vertexAIInferenceTest { def cmdArgs = mapToArgString(argMap) exec { executable 'sh' - args '-c', ". ${envdir}/bin/activate && ${runScriptsDir}/run_integration_test.sh $cmdArgs" + args '-c', ". ${envdir}/bin/activate && RUN_VERTEX_AI_TESTS=1 ${runScriptsDir}/run_integration_test.sh $cmdArgs" } } } diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini index 03fa34054c2d..039f9dd73a30 100644 --- a/sdks/python/tox.ini +++ b/sdks/python/tox.ini @@ -562,7 +562,7 @@ commands = /bin/sh -c "pip freeze | grep -E google-cloud-aiplatform" /bin/sh -c "pip freeze | grep -E transformers" # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories. - /bin/sh -c 'pytest apache_beam/ml/transforms/embeddings -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m "not vertex_ai_postcommit" {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret' + /bin/sh -c 'pytest apache_beam/ml/transforms/embeddings -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret' [testenv:py{310,311}-TFHubEmbeddings-{014,015}] @@ -580,7 +580,7 @@ commands = # Log aiplatform and its dependencies version for debugging /bin/sh -c "pip freeze | grep -E tensorflow" # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories. - bash {toxinidir}/scripts/run_pytest.sh {envname} 'apache_beam/ml/transforms/embeddings' "-m not(vertex_ai_postcommit)" + bash {toxinidir}/scripts/run_pytest.sh {envname} 'apache_beam/ml/transforms/embeddings' [testenv:py{310,312}-dill] extras = test,dill From 81fe00cff1e35754dcf8b2c6f1be27c593004770 Mon Sep 17 00:00:00 2001 From: aIbrahiim Date: Tue, 3 Mar 2026 21:46:15 +0200 Subject: [PATCH 3/6] Fix pylint vertex ai test --- sdks/python/apache_beam/testing/vertex_ai_skip.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/testing/vertex_ai_skip.py b/sdks/python/apache_beam/testing/vertex_ai_skip.py index 9207ea188e0a..f2fe386f2d0e 100644 --- a/sdks/python/apache_beam/testing/vertex_ai_skip.py +++ b/sdks/python/apache_beam/testing/vertex_ai_skip.py @@ -27,7 +27,7 @@ def _is_vertex_ai_available() -> bool: """Return True if Vertex AI client dependencies are importable.""" try: - import vertexai + import vertexai # pylint: disable=unused-import except ImportError: return False return True From d64c21493d58e72d537d091fbfe6f4193286f46d Mon Sep 17 00:00:00 2001 From: aIbrahiim Date: Wed, 4 Mar 2026 11:55:52 +0200 Subject: [PATCH 4/6] remove RUN_VERTEX_AI_TESTS flag --- sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py b/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py index f66d1f9d9b7b..d6bef97a7495 100644 --- a/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py +++ b/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py @@ -61,7 +61,6 @@ def chunk_approximately_equals(expected, actual): all(isinstance(x, float) for x in actual.embedding.dense_embedding)) -@skip_if_vertex_ai_disabled @pytest.mark.vertex_ai_postcommit @unittest.skipIf( not VERTEX_AI_AVAILABLE, "Vertex AI dependencies not available") From 0f805b56d3b7fd6411747f34ac45caa578256202 Mon Sep 17 00:00:00 2001 From: aIbrahiim Date: Thu, 5 Mar 2026 02:35:38 +0200 Subject: [PATCH 5/6] Use pytest.importorskip for Vertex AI tests & remove vertex_ai_skip helper --- .../inference/vertex_ai_inference_it_test.py | 4 +- .../ml/rag/embeddings/vertex_ai_test.py | 15 +++----- ...vertex_ai_test.py => vertex_ai_it_test.py} | 14 ++----- .../apache_beam/testing/vertex_ai_skip.py | 38 ------------------- .../python/test-suites/dataflow/common.gradle | 2 +- .../python/test-suites/tox/py310/build.gradle | 4 -- sdks/python/tox.ini | 3 +- 7 files changed, 12 insertions(+), 68 deletions(-) rename sdks/python/apache_beam/ml/transforms/embeddings/{vertex_ai_test.py => vertex_ai_it_test.py} (97%) delete mode 100644 sdks/python/apache_beam/testing/vertex_ai_skip.py diff --git a/sdks/python/apache_beam/ml/inference/vertex_ai_inference_it_test.py b/sdks/python/apache_beam/ml/inference/vertex_ai_inference_it_test.py index 8908477dae7f..f2f0c2e8a924 100644 --- a/sdks/python/apache_beam/ml/inference/vertex_ai_inference_it_test.py +++ b/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): diff --git a/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py b/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py index d6bef97a7495..43f52d9186fc 100644 --- a/sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py +++ b/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_') diff --git a/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py b/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_it_test.py similarity index 97% rename from sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py rename to sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_it_test.py index bc8ca4f0673e..ec6bdcfbae47 100644 --- a/sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py +++ b/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( diff --git a/sdks/python/apache_beam/testing/vertex_ai_skip.py b/sdks/python/apache_beam/testing/vertex_ai_skip.py deleted file mode 100644 index f2fe386f2d0e..000000000000 --- a/sdks/python/apache_beam/testing/vertex_ai_skip.py +++ /dev/null @@ -1,38 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one or more -# contributor license agreements. See the NOTICE file distributed with -# this work for additional information regarding copyright ownership. -# The ASF licenses this file to You under the Apache License, Version 2.0 -# (the "License"); you may not use this file except in compliance with -# the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# - -"""Centralized control for Vertex AI integration tests. - -Test modules should use skip_if_vertex_ai_disabled on classes or tests that -require Vertex AI Python dependencies to be installed. -""" - -import pytest - - -def _is_vertex_ai_available() -> bool: - """Return True if Vertex AI client dependencies are importable.""" - try: - import vertexai # pylint: disable=unused-import - except ImportError: - return False - return True - - -skip_if_vertex_ai_disabled = pytest.mark.skipif( - not _is_vertex_ai_available(), - reason='Vertex AI dependencies not available.') diff --git a/sdks/python/test-suites/dataflow/common.gradle b/sdks/python/test-suites/dataflow/common.gradle index b21efb2ae828..7d65da6ee3ad 100644 --- a/sdks/python/test-suites/dataflow/common.gradle +++ b/sdks/python/test-suites/dataflow/common.gradle @@ -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" } } } diff --git a/sdks/python/test-suites/tox/py310/build.gradle b/sdks/python/test-suites/tox/py310/build.gradle index 217564257673..8902f53a5482 100644 --- a/sdks/python/test-suites/tox/py310/build.gradle +++ b/sdks/python/test-suites/tox/py310/build.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" diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini index 039f9dd73a30..71b9b27cfcf3 100644 --- a/sdks/python/tox.ini +++ b/sdks/python/tox.ini @@ -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}" From cd821e92110cb28508246a85148f075a1e52882b Mon Sep 17 00:00:00 2001 From: aIbrahiim Date: Thu, 5 Mar 2026 16:05:14 +0200 Subject: [PATCH 6/6] Stabilize py310 ml tox env and disabling pre release deps --- sdks/python/tox.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/sdks/python/tox.ini b/sdks/python/tox.ini index 71b9b27cfcf3..49c05d3ff94a 100644 --- a/sdks/python/tox.ini +++ b/sdks/python/tox.ini @@ -109,6 +109,7 @@ deps = pip==25.0.1 accelerate>=1.6.0 onnx<1.19.0 +pip_pre = False setenv = extras = test,gcp,dataframe,ml_test commands =