workflows+gradle: capture and run ML tests that require non self-hosted environment on ubuntu-latest (#35585)

mohamedawnallah · damccorm · web-flow · commit fbd330f24afe · 2025-07-15T09:41:59.000-04:00
* .github+sdks: run ml tests that require non self-hosted env

* .github: trigger post commit python

* Update sdks/python/pytest.ini

Co-authored-by: Danny McCormick &lt;dannymccormick@google.com&gt;

* .github+sdks: add descriptive comments about runner changes

* workflows+gradle: use `require_docker_in_docker` marker

---------

Co-authored-by: Danny McCormick &lt;dannymccormick@google.com&gt;
diff --git a/.github/workflows/beam_PostCommit_Python.yml b/.github/workflows/beam_PostCommit_Python.yml
@@ -54,14 +54,20 @@ env:
 jobs:
   beam_PostCommit_Python:
     name: ${{ matrix.job_name }} (${{ matrix.job_phrase }} ${{ matrix.python_version }})
-    runs-on: [self-hosted, ubuntu-20.04, highmem22]
+    runs-on: ${{ matrix.os }}
     timeout-minutes: 240
     strategy:
       fail-fast: false
       matrix:
         job_name: [beam_PostCommit_Python]
         job_phrase: [Run Python PostCommit]
         python_version: ['3.9', '3.10', '3.11', '3.12']
+        # Run on both self-hosted and GitHub-hosted runners.
+        # Some tests (marked require_docker_in_docker) can't run on Beam's
+        # self-hosted runners due to Docker-in-Docker environment constraint.
+        # These tests will only execute on ubuntu-latest (GitHub-hosted).
+        # Context: https://github.com/apache/beam/pull/35585
+        os: [[self-hosted, ubuntu-20.04, highmem22], [ubuntu-latest]]
     if: |
       github.event_name == 'workflow_dispatch' ||
       github.event_name == 'pull_request_target' ||
@@ -99,6 +105,11 @@ jobs:
           arguments: |
             -Pjava21Home=$JAVA_HOME_21_X64 \
             -PuseWheelDistribution \
+            -Pposargs="${{
+              contains(matrix.os, 'self-hosted') &&
+                '-m ''not require_non_self_hosted''' ||
+                '-m ''require_non_self_hosted'''
+            }}" \
             -PpythonVersion=${{ matrix.python_version }} \
         env:
           CLOUDSDK_CONFIG: ${{ env.KUBELET_GCLOUD_CONFIG_PATH}}
@@ -115,4 +126,4 @@ jobs:
           commit: '${{ env.prsha || env.GITHUB_SHA }}'
           comment_mode: ${{ github.event_name == 'issue_comment'  && 'always' || 'off' }}
           files: '**/pytest*.xml'
-          large_files: true
+          large_files: true
diff --git a/.github/workflows/beam_PreCommit_Python_ML.yml b/.github/workflows/beam_PreCommit_Python_ML.yml
@@ -58,14 +58,20 @@ env:
 jobs:
   beam_PreCommit_Python_ML:
     name: ${{ matrix.job_name }} (${{ matrix.job_phrase }} ${{ matrix.python_version }})
-    runs-on: [self-hosted, ubuntu-20.04, main]
+    runs-on: ${{ matrix.os }}
     timeout-minutes: 180
     strategy:
       fail-fast: false
       matrix:
         job_name: ['beam_PreCommit_Python_ML']
         job_phrase: ['Run Python_ML PreCommit']
         python_version: ['3.9','3.10','3.11','3.12']
+        # Run on both self-hosted and GitHub-hosted runners.
+        # Some tests (marked require_docker_in_docker) can't run on Beam's
+        # self-hosted runners due to Docker-in-Docker environment constraint.
+        # These tests will only execute on ubuntu-latest (GitHub-hosted).
+        # Context: https://github.com/apache/beam/pull/35585
+        os: [[self-hosted, ubuntu-20.04, main], [ubuntu-latest]]
     if: |
       github.event_name == 'push' ||
       github.event_name == 'pull_request_target' ||
@@ -96,7 +102,11 @@ jobs:
         with:
           gradle-command: :sdks:python:test-suites:tox:py${{steps.set_py_ver_clean.outputs.py_ver_clean}}:testPy${{steps.set_py_ver_clean.outputs.py_ver_clean}}ML
           arguments: |
-            -Pposargs=apache_beam/ml/ \
+            -Pposargs="${{
+              contains(matrix.os, 'self-hosted') &&
+                'apache_beam/ml/ -m ''not require_non_self_hosted''' ||
+                'apache_beam/ml/ -m ''require_non_self_hosted'''
+            }}" \
             -PpythonVersion=${{ matrix.python_version }}
       - name: Archive Python Test Results
         uses: actions/upload-artifact@v4
@@ -111,4 +121,4 @@ jobs:
           commit: '${{ env.prsha || env.GITHUB_SHA }}'
           comment_mode: ${{ github.event_name == 'issue_comment'  && 'always' || 'off' }}
           files: '**/pytest*.xml'
-          large_files: true
+          large_files: true
diff --git a/sdks/python/pytest.ini b/sdks/python/pytest.ini
@@ -70,6 +70,7 @@ markers =
     uses_mock_api: tests that uses the mock API cluster.
     uses_feast: tests that uses feast in some way
     gemini_postcommit: gemini postcommits that need additional deps.
+    require_docker_in_docker: tests that require running Docker inside Docker (Docker-in-Docker), which is not supported on Beam’s self-hosted runners. Context: https://github.com/apache/beam/pull/35585
 
 # Default timeout intended for unit tests.
 # If certain tests need a different value, please see the docs on how to
diff --git a/sdks/python/test-suites/direct/common.gradle b/sdks/python/test-suites/direct/common.gradle
@@ -419,6 +419,33 @@ task feastIntegrationTest {
     }
 }
 
+// Integration tests that require Docker-in-Docker capabilities.
+// These tests are marked with the `require_docker_in_docker` pytest marker
+// because they rely on Docker-in-Docker configurations that are not supported
+// on Beam's self-hosted GitHub Actions runners. Docker-in-Docker works on
+// ubuntu-latest GitHub-hosted runners but not on self-hosted environments due
+// to containerization architecture differences.
+// Context: https://github.com/apache/beam/pull/35585
+task dockerInDockerIntegrationTest {
+  dependsOn 'installGcpTest'
+  dependsOn ':sdks:python:sdist'
+
+  doLast {
+      def testOpts = basicTestOpts
+      def argMap = [
+          "test_opts": testOpts,
+          "suite": "postCommitIT-direct-py${pythonVersionSuffix}",
+          "collect": "require_docker_in_docker",
+          "runner": "TestDirectRunner",
+      ]
+      def cmdArgs = mapToArgString(argMap)
+      exec {
+        executable 'sh'
+        args '-c', ". ${envdir}/bin/activate && ${runScriptsDir}/run_integration_test.sh $cmdArgs"
+      }
+    }
+}
+
 // Add all the RunInference framework IT tests to this gradle task that runs on Direct Runner Post commit suite.
 project.tasks.register("inferencePostCommitIT") {
   dependsOn = [
@@ -429,6 +456,7 @@ project.tasks.register("inferencePostCommitIT") {
   'transformersInferenceTest',
   'testcontainersTest',
   'feastIntegrationTest',
+  'dockerInDockerIntegrationTest',
   // (TODO) https://github.com/apache/beam/issues/25799
    // uncomment tfx bsl tests once tfx supports protobuf 4.x
   // 'tfxInferenceTest',
