sdks/python: make milvus as extra dependency

Author: mohamedawnallah

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

@@ -467,7 +467,9 @@ def create_user_yaml(service_port: int, max_vector_field_num=5):
         os.remove(path)
 
 
-@pytest.mark.uses_testcontainer
+# Mark Milvus itests with 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.non_self_hosted
 @unittest.skipUnless(
     platform.system() == "Linux",
     "Test runs only on Linux due to lack of support, as yet, for nested "
@@ -483,22 +485,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.
+    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.
@@ -588,7 +587,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": "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 = [

sdks/python/tox.ini

@@ -85,11 +85,13 @@ commands_pre =
   pip --version
   # pip check
   bash {toxinidir}/scripts/run_tox_cleanup.sh
+extras = test,milvus
 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,milvus
 commands =
   python apache_beam/examples/complete/autocomplete_test.py
   bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"
@@ -109,7 +111,7 @@ deps =
   pip==25.0.1
   accelerate>=1.6.0
 setenv =
-extras = test,gcp,dataframe,ml_test
+extras = test,gcp,dataframe,milvus,ml_test
 commands =
   # Log tensorflow version for debugging
   /bin/sh -c "pip freeze | grep -E tensorflow"
@@ -121,7 +123,7 @@ commands =
 deps =
   accelerate>=1.6.0
 setenv =
-extras = test,gcp,dataframe,p312_ml_test
+extras = test,gcp,dataframe,milvus,p312_ml_test
 commands =
   # Log tensorflow version for debugging
   /bin/sh -c "pip freeze | grep -E tensorflow"
@@ -546,6 +548,14 @@ commands =
   /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{39,310,311,312,313}-milvus]
+extras = test,gcp,milvus
+commands =
+  # Log versions for debugging.
+  /bin/sh -c "pip freeze | grep -E 'pymilvus|testcontainers'"
+  bash {toxinidir}/scripts/run_pytest.sh {envname} "{posargs}"
+
+
 [testenv:py{39,310}-TFHubEmbeddings-{014,015}]
 deps =
   014: tensorflow-hub>=0.14.0,<0.15.0