Cherrypick #37725

Author: Amar3tto

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

@@ -16,6 +16,7 @@
 
 """Tests for apache_beam.ml.rag.embeddings.vertex_ai."""
 
+import pytest
 import shutil
 import tempfile
 import unittest
@@ -52,6 +53,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):

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.')

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"

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