sdks/python: mark milvus itests with require_docker_in_docker

Author: mohamedawnallah

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

@@ -34,18 +34,6 @@
 
 import pytest
 import yaml
-from pymilvus import CollectionSchema
-from pymilvus import DataType
-from pymilvus import FieldSchema
-from pymilvus import Function
-from pymilvus import FunctionType
-from pymilvus import MilvusClient
-from pymilvus import RRFRanker
-from pymilvus.milvus_client import IndexParams
-from testcontainers.core.config import MAX_TRIES as TC_MAX_TRIES
-from testcontainers.core.config import testcontainers_config
-from testcontainers.core.generic import DbContainer
-from testcontainers.milvus import MilvusContainer
 
 import apache_beam as beam
 from apache_beam.ml.rag.types import Chunk
@@ -54,7 +42,21 @@
 from apache_beam.testing.test_pipeline import TestPipeline
 from apache_beam.testing.util import assert_that
 
+# pylint: disable=ungrouped-imports
 try:
+  from pymilvus import (
+      CollectionSchema,
+      DataType,
+      FieldSchema,
+      Function,
+      FunctionType,
+      MilvusClient,
+      RRFRanker)
+  from pymilvus.milvus_client import IndexParams
+  from testcontainers.core.config import MAX_TRIES as TC_MAX_TRIES
+  from testcontainers.core.config import testcontainers_config
+  from testcontainers.core.generic import DbContainer
+  from testcontainers.milvus import MilvusContainer
   from apache_beam.transforms.enrichment import Enrichment
   from apache_beam.ml.rag.enrichment.milvus_search import (
       MilvusSearchEnrichmentHandler,
@@ -467,7 +469,7 @@ def create_user_yaml(service_port: int, max_vector_field_num=5):
         os.remove(path)
 
 
-@pytest.mark.uses_testcontainer
+@pytest.mark.require_docker_in_docker
 @unittest.skipUnless(
     platform.system() == "Linux",
     "Test runs only on Linux due to lack of support, as yet, for nested "
@@ -483,22 +485,16 @@ 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)
+    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):
@@ -1368,4 +1364,4 @@ def assert_chunks_equivalent(
 
 
 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()

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.