feat: Update pymilvus to 2.4.x (#151)

* feat: update pymilvus to 2.4.x * fix: Update expediagroup implementation of milvus online store as eg-milvus to differentiate from the opensource implementation * update requirements * fix tests * fix tests * fix: Just update testcontainers and pymilvus deps
ExpediaGroup · Nov 1, 2024 · 9b5ae73 · 9b5ae73
1 parent 22bcbc4
commit 9b5ae73
Show file tree

Hide file tree

Showing 13 changed files with 95 additions and 95 deletions.
diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml
@@ -46,7 +46,7 @@ jobs:
           CGO_LDFLAGS_ALLOW=".*" COMPILE_GO=True python setup.py build_ext --inplace
       - name: Test Milvus tests
         if: matrix.os == 'ubuntu-latest'
-        run: python -m pytest -n 1 --color=yes sdk/python/tests/expediagroup/test_milvus_online_store.py
+        run: python -m pytest -n 1 --color=yes sdk/python/tests/expediagroup/test_eg_milvus_online_store.py
       - name: Test Python
         if: matrix.os == 'ubuntu-latest'
         run: make test-python-unit

diff --git a/Makefile b/Makefile
@@ -119,7 +119,7 @@ test-milvus-integration-local:
 		FEAST_IS_LOCAL_TEST=True \
 		FEAST_LOCAL_ONLINE_CONTAINER=True \
 		python -m pytest -n 8 --integration \
-			-k "milvus" \
+			-k "eg-milvus" \
 		sdk/python/tests \
 	) || echo "This script uses Docker, and it isn't running - please start the Docker Daemon and try again!";
 

diff --git a/docs/reference/online-stores/milvus.md b/docs/reference/online-stores/milvus.md
@@ -13,7 +13,7 @@ Features:
 
 ## Getting started
 In order to use this online store, you'll need to install the milvus extra:
-- `pip install 'feast[milvus]'`
+- `pip install 'feast[eg-milvus]'`
 
 ## Example
 

diff --git a/milvus/Dockerfile b/milvus/Dockerfile
@@ -1,12 +1,12 @@
 FROM python:3.11-slim
 
-RUN python3 -m pip install milvus==2.2.12
+RUN python3 -m pip install milvus==2.4.13
 
 # this is needed to divert logs to stdout
-RUN mkdir -p /root/.milvus.io/milvus-server/2.2.12/logs/
-RUN touch /root/.milvus.io/milvus-server/2.2.12/logs/milvus-stdout.log
-RUN touch /root/.milvus.io/milvus-server/2.2.12/logs/milvus-stderr.log
-RUN ln -sf /dev/stdout /root/.milvus.io/milvus-server/2.2.12/logs/milvus-stdout.log \
-    && ln -sf /dev/stderr /root/.milvus.io/milvus-server/2.2.12/logs/milvus-stderr.log
+RUN mkdir -p /root/.milvus.io/milvus-server/2.4.13/logs/
+RUN touch /root/.milvus.io/milvus-server/2.4.13/logs/milvus-stdout.log
+RUN touch /root/.milvus.io/milvus-server/2.4.13/logs/milvus-stderr.log
+RUN ln -sf /dev/stdout /root/.milvus.io/milvus-server/2.4.13/logs/milvus-stdout.log \
+    && ln -sf /dev/stderr /root/.milvus.io/milvus-server/2.4.13/logs/milvus-stderr.log
 
 CMD ["milvus-server"]
diff --git a/...ediagroup/vectordb/milvus_online_store.py → ...agroup/vectordb/eg_milvus_online_store.py b/...ediagroup/vectordb/milvus_online_store.py → ...agroup/vectordb/eg_milvus_online_store.py
@@ -53,10 +53,10 @@
 BATCH_SIZE = 10_000
 
 
-class MilvusOnlineStoreConfig(FeastConfigBaseModel):
+class EGMilvusOnlineStoreConfig(FeastConfigBaseModel):
     """Online store config for the Milvus online store"""
 
-    type: Literal["milvus"] = "milvus"
+    type: Literal["eg-milvus"] = "eg-milvus"
     """Online store type selector"""
 
     alias: str = "default"
@@ -75,7 +75,7 @@ class MilvusOnlineStoreConfig(FeastConfigBaseModel):
     """ the port to connect to a Milvus instance. Should be the one used for GRPC (default: 19530) """
 
 
-class MilvusConnectionManager:
+class EGMilvusConnectionManager:
     def __init__(self, online_config: RepoConfig):
         self.online_config = online_config
 
@@ -107,7 +107,7 @@ def __exit__(self, exc_type, exc_value, traceback):
             print(f"An exception of type {exc_type} occurred: {exc_value}")
 
 
-class MilvusOnlineStore(OnlineStore):
+class EGMilvusOnlineStore(OnlineStore):
     def online_write_batch(
         self,
         config: RepoConfig,
@@ -117,7 +117,7 @@ def online_write_batch(
         ],
         progress: Optional[Callable[[int], Any]],
     ) -> None:
-        with MilvusConnectionManager(config.online_store):
+        with EGMilvusConnectionManager(config.online_store):
             self._create_collection_if_not_exists(table)
             print("Starting the process to batch write data into Milvus.")
             collection_to_load_data = Collection(table.name)
@@ -139,7 +139,7 @@ def online_read(
         entity_keys: List[EntityKeyProto],
         requested_features: Optional[List[str]] = None,
     ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
-        with MilvusConnectionManager(config.online_store):
+        with EGMilvusConnectionManager(config.online_store):
             quer_expr = self._construct_milvus_query(entity_keys)
             use_iter_search = len(entity_keys) > MAX_SEARCH_SIZE
             collection = Collection(table.name)
@@ -179,7 +179,7 @@ def update(
         entities_to_keep: Sequence[Entity],
         partial: bool,
     ):
-        with MilvusConnectionManager(config.online_store):
+        with EGMilvusConnectionManager(config.online_store):
             for table_to_keep in tables_to_keep:
                 self._create_collection_if_not_exists(table_to_keep)
 
@@ -204,7 +204,7 @@ def teardown(
         tables: Sequence[FeatureView],
         entities: Sequence[Entity],
     ):
-        with MilvusConnectionManager(config.online_store):
+        with EGMilvusConnectionManager(config.online_store):
             for table in tables:
                 collection_name = table.name
                 if utility.has_collection(collection_name):

diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py
@@ -65,7 +65,7 @@
     "scylladb": "feast.infra.online_stores.contrib.cassandra_online_store.cassandra_online_store.CassandraOnlineStore",
     "mysql": "feast.infra.online_stores.contrib.mysql_online_store.mysql.MySQLOnlineStore",
     "hazelcast": "feast.infra.online_stores.contrib.hazelcast_online_store.hazelcast_online_store.HazelcastOnlineStore",
-    "milvus": "feast.expediagroup.vectordb.milvus_online_store.MilvusOnlineStore",
+    "eg-milvus": "feast.expediagroup.vectordb.eg_milvus_online_store.EGMilvusOnlineStore",
     "elasticsearch": "feast.expediagroup.vectordb.elasticsearch_online_store.ElasticsearchOnlineStore",
     "ikv": "feast.infra.online_stores.contrib.ikv_online_store.ikv.IKVOnlineStore",
     # "elasticsearch": "feast.infra.online_stores.contrib.elasticsearch.ElasticSearchOnlineStore",

diff --git a/sdk/python/requirements/py3.10-ci-requirements.txt b/sdk/python/requirements/py3.10-ci-requirements.txt
@@ -664,7 +664,7 @@ pyjwt[crypto]==2.9.0
     #   msal
     #   singlestoredb
     #   snowflake-connector-python
-pymilvus==2.4.6
+pymilvus==2.4.9
     # via eg-feast (setup.py)
 pymssql==2.3.1
     # via eg-feast (setup.py)
@@ -888,7 +888,7 @@ terminado==0.18.1
     # via
     #   jupyter-server
     #   jupyter-server-terminals
-testcontainers==4.4.0
+testcontainers==4.8.2
     # via eg-feast (setup.py)
 thriftpy2==0.5.2
     # via happybase

diff --git a/sdk/python/requirements/py3.11-ci-requirements.txt b/sdk/python/requirements/py3.11-ci-requirements.txt
@@ -655,7 +655,7 @@ pyjwt[crypto]==2.9.0
     #   msal
     #   singlestoredb
     #   snowflake-connector-python
-pymilvus==2.4.6
+pymilvus==2.4.9
     # via eg-feast (setup.py)
 pymssql==2.3.1
     # via eg-feast (setup.py)
@@ -879,7 +879,7 @@ terminado==0.18.1
     # via
     #   jupyter-server
     #   jupyter-server-terminals
-testcontainers==4.4.0
+testcontainers==4.8.2
     # via eg-feast (setup.py)
 thriftpy2==0.5.2
     # via happybase

diff --git a/sdk/python/requirements/py3.9-ci-requirements.txt b/sdk/python/requirements/py3.9-ci-requirements.txt
@@ -672,7 +672,7 @@ pyjwt[crypto]==2.9.0
     #   msal
     #   singlestoredb
     #   snowflake-connector-python
-pymilvus==2.4.6
+pymilvus==2.4.9
     # via eg-feast (setup.py)
 pymssql==2.3.1
     # via eg-feast (setup.py)
@@ -896,7 +896,7 @@ terminado==0.18.1
     # via
     #   jupyter-server
     #   jupyter-server-terminals
-testcontainers==4.4.0
+testcontainers==4.8.2
     # via eg-feast (setup.py)
 thriftpy2==0.5.2
     # via happybase

diff --git a/...pediagroup/milvus_online_store_creator.py → ...iagroup/eg_milvus_online_store_creator.py b/...pediagroup/milvus_online_store_creator.py → ...iagroup/eg_milvus_online_store_creator.py
@@ -1,19 +1,18 @@
 from typing import Dict
 
-from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
+from testcontainers.milvus import MilvusContainer
 
 from tests.integration.feature_repos.universal.online_store_creator import (
     OnlineStoreCreator,
 )
 
 
-class MilvusOnlineStoreCreator(OnlineStoreCreator):
+class EGMilvusOnlineStoreCreator(OnlineStoreCreator):
     def __init__(self, project_name: str, **kwargs):
         super().__init__(project_name)
-        self.container = DockerContainer("mbackes/milvus:2.3.0").with_exposed_ports(
-            "19530"
-        )
+        with MilvusContainer("milvusdb/milvus:v2.4.11") as milvus_container:
+            self.container = milvus_container
 
     def create_online_store(self) -> Dict[str, str]:
         self.container.start()
@@ -23,13 +22,14 @@ def create_online_store(self) -> Dict[str, str]:
         wait_for_logs(
             container=self.container, predicate=log_string_to_wait_for, timeout=60
         )
-        exposed_port = self.container.get_exposed_port("19530")
+        host = self.container.get_container_host_ip()
+        exposed_port = int(self.container.get_exposed_port(self.container.port))
 
         return {
             "alias": "default",
-            "type": "milvus",
-            "host": "localhost",
-            "port": str(exposed_port),
+            "type": "eg-milvus",
+            "host": host,
+            "port": exposed_port,
             "username": "user",
             "password": "password",
         }