Include MLTransform IT test suite as a part of PostCommits

apache · Feb 13, 2024 · ecb443c · ecb443c
1 parent 514b03b
commit ecb443c
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 7 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -510,25 +510,23 @@ tasks.register("python38PostCommit") {
   // The default container uses Python 3.8. The goal here is to
   // duild Docker images for TensorRT tests during run time for python versions
   // other than 3.8 and add these tests in other python postcommit suites.
-  dependsOn(":sdks:python:test-suites:dataflow:py38:inferencePostCommitIT")
-  dependsOn(":sdks:python:test-suites:direct:py38:inferencePostCommitIT")
+  dependsOn(":sdks:python:test-suites:dataflow:py38:BeamMLPostCommitIT")
+  dependsOn(":sdks:python:test-suites:direct:py38:BeamMLPostCommitIT")
 }
 
 tasks.register("python39PostCommit") {
   dependsOn(":sdks:python:test-suites:dataflow:py39:postCommitIT")
   dependsOn(":sdks:python:test-suites:direct:py39:postCommitIT")
   dependsOn(":sdks:python:test-suites:direct:py39:hdfsIntegrationTest")
   dependsOn(":sdks:python:test-suites:portable:py39:postCommitPy39")
-  // TODO (https://github.com/apache/beam/issues/23966)
-  // Move this to Python 3.10 test suite once tfx-bsl has python 3.10 wheel.
-  dependsOn(":sdks:python:test-suites:direct:py39:inferencePostCommitIT")
 }
 
 tasks.register("python310PostCommit") {
   dependsOn(":sdks:python:test-suites:dataflow:py310:postCommitIT")
   dependsOn(":sdks:python:test-suites:direct:py310:postCommitIT")
   dependsOn(":sdks:python:test-suites:direct:py310:hdfsIntegrationTest")
   dependsOn(":sdks:python:test-suites:portable:py310:postCommitPy310")
+  dependsOn(":sdks:python:test-suites:direct:py310:BeamMLPostCommitIT")
 }
 
 tasks.register("python311PostCommit") {

diff --git a/sdks/python/apache_beam/examples/ml_transform/ml_transform_it_test.py b/sdks/python/apache_beam/examples/ml_transform/ml_transform_it_test.py
@@ -57,6 +57,7 @@ def _publish_metrics(pipeline, metric_value, metrics_table, metric_name):
   )])
 
 
+@pytest.mark.it_postcommit
 @pytest.mark.uses_tft
 class LargeMovieReviewDatasetProcessTest(unittest.TestCase):
   def test_process_large_movie_review_dataset(self):

diff --git a/sdks/python/test-suites/dataflow/common.gradle b/sdks/python/test-suites/dataflow/common.gradle
@@ -475,7 +475,7 @@ task tftTests {
     def argMap = [
         "test_opts": testOpts,
         "suite": "TFTransformTests-df${pythonVersionSuffix}",
-        "collect": "uses_tft",
+        "collect": "uses_tft and not it_postcommit",
         "requirements_file": "apache_beam/testing/benchmarks/cloudml/requirements.txt",
         "pipeline_opts": opts,
     ]
@@ -487,14 +487,32 @@ task tftTests {
   }
 }
 
+task MLTransformITTest {
+  dependsOn "installTFTRequirements"
+  doLast {
+    def opts = project.findProperty('opts')
+    opts += " --sdk_location=${project.ext.sdkLocation}"
+    def testOpts = basicPytestOpts + ["--numprocesses=8", "--dist=loadfile"]
+    def argMap = [
+        "test_opts": testOpts,
+        "suite": "MLTransformITTest-df${pythonVersionSuffix}",
+        // Add MLTransform IT tests pytest markers.
+        "collect": "uses_tft and it_postcommit",
+        "requirements_file": "apache_beam/testing/benchmarks/cloudml/requirements.txt",
+        "pipeline_opts": opts,
+    ]
+  }
+}
+
 // add all RunInference E2E tests that run on DataflowRunner
 // As of now, this test suite is enable in py38 suite as the base NVIDIA image used for Tensor RT
 // contains Python 3.8.
 // TODO: https://github.com/apache/beam/issues/22651
-project.tasks.register("inferencePostCommitIT") {
+project.tasks.register("BeamMLPostCommitIT") {
   dependsOn = [
   'tensorRTtests',
   'vertexAIInferenceTest',
+  'MLTransformITTest'
   ]
 }