sdks/python: make milvus as extra dependency

Author: mohamedawnallah

.github/workflows/beam_PostCommit_Python.yml

@@ -54,7 +54,9 @@ 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:
+      - [self-hosted, ubuntu-20.04, highmem22]
+      - ubuntu-latest
     timeout-minutes: 240
     strategy:
       fail-fast: false
@@ -99,7 +101,13 @@ jobs:
           arguments: |
             -Pjava21Home=$JAVA_HOME_21_X64 \
             -PuseWheelDistribution \
-            -PpythonVersion=${{ matrix.python_version }} \
+            # Conditionally add -Pposargs based on the runner type.
+            -Pposargs="${{
+              contains(matrix.runner_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}}
       - name: Archive Python Test Results
@@ -115,4 +123,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

.github/workflows/beam_PreCommit_Python_ML.yml

@@ -58,7 +58,9 @@ 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:
+      - [self-hosted, ubuntu-20.04, main]
+      - ubuntu-latest
     timeout-minutes: 180
     strategy:
       fail-fast: false
@@ -96,7 +98,12 @@ 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/ \
+            # Conditionally add -Pposargs based on the runner type.
+            -Pposargs="${{
+              contains(matrix.runner_os, 'self-hosted') &&
+                '-m ''not require_non_self_hosted''' ||
+                '-m ''require_non_self_hosted'''
+            }}" \
             -PpythonVersion=${{ matrix.python_version }}
       - name: Archive Python Test Results
         uses: actions/upload-artifact@v4
@@ -111,4 +118,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

sdks/python/apache_beam/ml/rag/enrichment/milvus_search_it_test.py

@@ -467,7 +467,10 @@ def create_user_yaml(service_port: int, max_vector_field_num=5):
         os.remove(path)
 
 
-@pytest.mark.uses_testcontainer
+# Mark Milvus itests with require_non_self_hosted due to Docker-in-Docker
+# runtime issues on the self-hosted envs. This typically runs on the
+# ubuntu-latest env.
+@pytest.mark.require_non_self_hosted
 @unittest.skipUnless(
     platform.system() == "Linux",
     "Test runs only on Linux due to lack of support, as yet, for nested "
@@ -483,22 +486,17 @@ class TestMilvusSearchEnrichment(unittest.TestCase):
 
   @classmethod
   def setUpClass(cls):
-    try:
-      cls._db = MilvusEnrichmentTestHelper.start_db_container(
-          cls._version, vector_client_max_retries=1, tc_max_retries=1)
-      cls._connection_params = MilvusConnectionParameters(
-          uri=cls._db.uri,
-          user=cls._db.user,
-          password=cls._db.password,
-          db_id=cls._db.id,
-          token=cls._db.token)
-      cls._collection_load_params = MilvusCollectionLoadParameters()
-      cls._collection_name = MilvusEnrichmentTestHelper.initialize_db_with_data(
-          cls._connection_params)
-    except Exception as e:
-      pytest.skip(
-          f"Skipping all tests in {cls.__name__} due to DB startup failure: {e}"
-      )
+    cls._db = MilvusEnrichmentTestHelper.start_db_container(
+        cls._version, vector_client_max_retries=1)
+    cls._connection_params = MilvusConnectionParameters(
+        uri=cls._db.uri,
+        user=cls._db.user,
+        password=cls._db.password,
+        db_id=cls._db.id,
+        token=cls._db.token)
+    cls._collection_load_params = MilvusCollectionLoadParameters()
+    cls._collection_name = MilvusEnrichmentTestHelper.initialize_db_with_data(
+        cls._connection_params)
 
   @classmethod
   def tearDownClass(cls):

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_non_self_hosted: tests requiring specialized environments or infrastructure not present in self-hosted setups.
 
 # Default timeout intended for unit tests.
 # If certain tests need a different value, please see the docs on how to

sdks/python/setup.py

@@ -401,7 +401,6 @@ def get_portability_package_data():
           'typing-extensions>=3.7.0',
           'zstandard>=0.18.0,<1',
           'pyyaml>=3.12,<7.0.0',
-          'pymilvus>=2.5.10,<3.0.0',
           # Dynamic dependencies must be specified in a separate list, otherwise
           # Dependabot won't be able to parse the main list. Any dynamic
           # dependencies will not receive updates from Dependabot.
@@ -445,7 +444,8 @@ def get_portability_package_data():
               'mysql-connector-python>=9.3.0',
               'python-tds>=1.16.1',
               'sqlalchemy-pytds>=1.0.2',
-              'oracledb>=3.1.1'
+              'oracledb>=3.1.1',
+              'milvus'
           ],
           'gcp': [
               'cachetools>=3.1.0,<6',
@@ -588,7 +588,8 @@ def get_portability_package_data():
               'transformers==4.25.1'
           ],
           'xgboost': ['xgboost>=1.6.0,<2.1.3', 'datatable==1.0.0'],
-          'tensorflow-hub': ['tensorflow-hub>=0.14.0,<0.16.0']
+          'tensorflow-hub': ['tensorflow-hub>=0.14.0,<0.16.0'],
+          'milvus': ['pymilvus>=2.5.10,<3.0.0']
       },
       zip_safe=False,
       # PyPI package information.

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

@@ -419,6 +419,27 @@ task feastIntegrationTest {
     }
 }
 
+// Integration tests that runs on non self-hosted environments.
+task nonSelfHostedIntegrationTest {
+  dependsOn 'installGcpTest'
+  dependsOn ':sdks:python:sdist'
+
+  doLast {
+      def testOpts = basicTestOpts
+      def argMap = [
+          "test_opts": testOpts,
+          "suite": "postCommitIT-direct-py${pythonVersionSuffix}",
+          "collect": "require_non_self_hosted",
+          "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 +450,7 @@ project.tasks.register("inferencePostCommitIT") {
   'transformersInferenceTest',
   'testcontainersTest',
   'feastIntegrationTest',
+  'nonSelfHostedIntegrationTest',
   // (TODO) https://github.com/apache/beam/issues/25799
    // uncomment tfx bsl tests once tfx supports protobuf 4.x
   // 'tfxInferenceTest',

sdks/python/tox.ini

@@ -85,11 +85,13 @@ commands_pre =
   pip --version
   # pip check
   bash {toxinidir}/scripts/run_tox_cleanup.sh
+extras = test
 commands =
   python apache_beam/examples/complete/autocomplete_test.py
   bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"
 
 [testenv:py{39,310,311,312,313}-win]
+extras = test
 commands =
   python apache_beam/examples/complete/autocomplete_test.py
   bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"