feat: Adding packages for Milvus Online Store (#4854)

franciscojavierarceo · web-flow · commit 49171bd53fb8 · 2024-12-17T14:46:01.000-05:00
Signed-off-by: Francisco Javier Arceo &lt;farceo@redhat.com&gt;
diff --git a/sdk/python/feast/infra/online_stores/milvus_online_store/__init__.py b/sdk/python/feast/infra/online_stores/milvus_online_store/__init__.py
diff --git a/sdk/python/feast/infra/online_stores/milvus_online_store/milvus_repo_configuration.py b/sdk/python/feast/infra/online_stores/milvus_online_store/milvus_repo_configuration.py
@@ -0,0 +1,12 @@
+from tests.integration.feature_repos.integration_test_repo_config import (
+    IntegrationTestRepoConfig,
+)
+from tests.integration.feature_repos.universal.online_store.milvus import (
+    MilvusOnlineStoreCreator,
+)
+
+FULL_REPO_CONFIGS = [
+    IntegrationTestRepoConfig(
+        online_store="milvus", online_store_creator=MilvusOnlineStoreCreator
+    ),
+]
diff --git a/sdk/python/requirements/py3.10-ci-requirements.txt b/sdk/python/requirements/py3.10-ci-requirements.txt
@@ -180,6 +180,8 @@ elasticsearch==8.16.0
     # via feast (setup.py)
 entrypoints==0.4
     # via altair
+environs==9.5.0
+    # via pymilvus
 exceptiongroup==1.2.2
     # via
     #   anyio
@@ -275,6 +277,7 @@ grpcio==1.68.0
     #   grpcio-status
     #   grpcio-testing
     #   grpcio-tools
+    #   pymilvus
     #   qdrant-client
 grpcio-health-checking==1.62.3
     # via feast (setup.py)
@@ -443,13 +446,17 @@ markupsafe==3.0.2
     #   nbconvert
     #   werkzeug
 marshmallow==3.23.1
-    # via great-expectations
+    # via
+    #   environs
+    #   great-expectations
 matplotlib-inline==0.1.7
     # via
     #   ipykernel
     #   ipython
 mdurl==0.1.2
     # via markdown-it-py
+milvus-lite==2.4.10
+    # via pymilvus
 minio==7.1.0
     # via feast (setup.py)
 mistune==3.0.2
@@ -548,6 +555,7 @@ pandas==2.2.3
     #   google-cloud-bigquery
     #   great-expectations
     #   ibis-framework
+    #   pymilvus
     #   snowflake-connector-python
 pandocfilters==1.5.1
     # via nbconvert
@@ -614,6 +622,7 @@ protobuf==4.25.5
     #   grpcio-tools
     #   mypy-protobuf
     #   proto-plus
+    #   pymilvus
     #   substrait
 psutil==5.9.0
     # via
@@ -658,7 +667,7 @@ pybindgen==0.22.1
     # via feast (setup.py)
 pycparser==2.22
     # via cffi
-pydantic==2.10.1
+pydantic==2.10.2
     # via
     #   feast (setup.py)
     #   fastapi
@@ -679,6 +688,8 @@ pyjwt[crypto]==2.10.0
     #   msal
     #   singlestoredb
     #   snowflake-connector-python
+pymilvus==2.4.9
+    # via feast (setup.py)
 pymssql==2.3.2
     # via feast (setup.py)
 pymysql==1.1.1
@@ -738,7 +749,9 @@ python-dateutil==2.9.0.post0
     #   pandas
     #   trino
 python-dotenv==1.0.1
-    # via uvicorn
+    # via
+    #   environs
+    #   uvicorn
 python-json-logger==2.0.7
     # via jupyter-events
 python-keycloak==4.2.2
@@ -839,6 +852,7 @@ setuptools==75.6.0
     #   jupyterlab
     #   kubernetes
     #   pip-tools
+    #   pymilvus
     #   singlestoredb
 singlestoredb==1.7.2
     # via feast (setup.py)
@@ -939,6 +953,7 @@ tqdm==4.67.1
     # via
     #   feast (setup.py)
     #   great-expectations
+    #   milvus-lite
 traitlets==5.14.3
     # via
     #   comm
@@ -1018,6 +1033,8 @@ tzlocal==5.2
     # via
     #   great-expectations
     #   trino
+ujson==5.10.0
+    # via pymilvus
 uri-template==1.3.0
     # via jsonschema
 urllib3==2.2.3
diff --git a/sdk/python/requirements/py3.10-requirements.txt b/sdk/python/requirements/py3.10-requirements.txt
@@ -95,7 +95,7 @@ pyarrow==18.0.0
     # via
     #   feast (setup.py)
     #   dask-expr
-pydantic==2.10.1
+pydantic==2.10.2
     # via
     #   feast (setup.py)
     #   fastapi
diff --git a/sdk/python/requirements/py3.11-ci-requirements.txt b/sdk/python/requirements/py3.11-ci-requirements.txt
@@ -178,6 +178,8 @@ elasticsearch==8.16.0
     # via feast (setup.py)
 entrypoints==0.4
     # via altair
+environs==9.5.0
+    # via pymilvus
 execnet==2.1.1
     # via pytest-xdist
 executing==2.1.0
@@ -268,6 +270,7 @@ grpcio==1.68.0
     #   grpcio-status
     #   grpcio-testing
     #   grpcio-tools
+    #   pymilvus
     #   qdrant-client
 grpcio-health-checking==1.62.3
     # via feast (setup.py)
@@ -434,13 +437,17 @@ markupsafe==3.0.2
     #   nbconvert
     #   werkzeug
 marshmallow==3.23.1
-    # via great-expectations
+    # via
+    #   environs
+    #   great-expectations
 matplotlib-inline==0.1.7
     # via
     #   ipykernel
     #   ipython
 mdurl==0.1.2
     # via markdown-it-py
+milvus-lite==2.4.10
+    # via pymilvus
 minio==7.1.0
     # via feast (setup.py)
 mistune==3.0.2
@@ -539,6 +546,7 @@ pandas==2.2.3
     #   google-cloud-bigquery
     #   great-expectations
     #   ibis-framework
+    #   pymilvus
     #   snowflake-connector-python
 pandocfilters==1.5.1
     # via nbconvert
@@ -605,6 +613,7 @@ protobuf==4.25.5
     #   grpcio-tools
     #   mypy-protobuf
     #   proto-plus
+    #   pymilvus
     #   substrait
 psutil==5.9.0
     # via
@@ -649,7 +658,7 @@ pybindgen==0.22.1
     # via feast (setup.py)
 pycparser==2.22
     # via cffi
-pydantic==2.10.1
+pydantic==2.10.2
     # via
     #   feast (setup.py)
     #   fastapi
@@ -670,6 +679,8 @@ pyjwt[crypto]==2.10.0
     #   msal
     #   singlestoredb
     #   snowflake-connector-python
+pymilvus==2.4.9
+    # via feast (setup.py)
 pymssql==2.3.2
     # via feast (setup.py)
 pymysql==1.1.1
@@ -729,7 +740,9 @@ python-dateutil==2.9.0.post0
     #   pandas
     #   trino
 python-dotenv==1.0.1
-    # via uvicorn
+    # via
+    #   environs
+    #   uvicorn
 python-json-logger==2.0.7
     # via jupyter-events
 python-keycloak==4.2.2
@@ -830,6 +843,7 @@ setuptools==75.6.0
     #   jupyterlab
     #   kubernetes
     #   pip-tools
+    #   pymilvus
     #   singlestoredb
 singlestoredb==1.7.2
     # via feast (setup.py)
@@ -920,6 +934,7 @@ tqdm==4.67.1
     # via
     #   feast (setup.py)
     #   great-expectations
+    #   milvus-lite
 traitlets==5.14.3
     # via
     #   comm
@@ -994,6 +1009,8 @@ tzlocal==5.2
     # via
     #   great-expectations
     #   trino
+ujson==5.10.0
+    # via pymilvus
 uri-template==1.3.0
     # via jsonschema
 urllib3==2.2.3
diff --git a/sdk/python/requirements/py3.11-requirements.txt b/sdk/python/requirements/py3.11-requirements.txt
@@ -93,7 +93,7 @@ pyarrow==18.0.0
     # via
     #   feast (setup.py)
     #   dask-expr
-pydantic==2.10.1
+pydantic==2.10.2
     # via
     #   feast (setup.py)
     #   fastapi
diff --git a/sdk/python/requirements/py3.9-ci-requirements.txt b/sdk/python/requirements/py3.9-ci-requirements.txt
@@ -182,6 +182,8 @@ elasticsearch==8.16.0
     # via feast (setup.py)
 entrypoints==0.4
     # via altair
+environs==9.5.0
+    # via pymilvus
 exceptiongroup==1.2.2
     # via
     #   anyio
@@ -277,6 +279,7 @@ grpcio==1.68.0
     #   grpcio-status
     #   grpcio-testing
     #   grpcio-tools
+    #   pymilvus
     #   qdrant-client
 grpcio-health-checking==1.62.3
     # via feast (setup.py)
@@ -452,13 +455,17 @@ markupsafe==3.0.2
     #   nbconvert
     #   werkzeug
 marshmallow==3.23.1
-    # via great-expectations
+    # via
+    #   environs
+    #   great-expectations
 matplotlib-inline==0.1.7
     # via
     #   ipykernel
     #   ipython
 mdurl==0.1.2
     # via markdown-it-py
+milvus-lite==2.4.10
+    # via pymilvus
 minio==7.1.0
     # via feast (setup.py)
 mistune==3.0.2
@@ -556,6 +563,7 @@ pandas==2.2.3
     #   google-cloud-bigquery
     #   great-expectations
     #   ibis-framework
+    #   pymilvus
     #   snowflake-connector-python
 pandocfilters==1.5.1
     # via nbconvert
@@ -622,6 +630,7 @@ protobuf==4.25.5
     #   grpcio-tools
     #   mypy-protobuf
     #   proto-plus
+    #   pymilvus
     #   substrait
 psutil==5.9.0
     # via
@@ -666,7 +675,7 @@ pybindgen==0.22.1
     # via feast (setup.py)
 pycparser==2.22
     # via cffi
-pydantic==2.10.1
+pydantic==2.10.2
     # via
     #   feast (setup.py)
     #   fastapi
@@ -687,6 +696,8 @@ pyjwt[crypto]==2.10.0
     #   msal
     #   singlestoredb
     #   snowflake-connector-python
+pymilvus==2.4.9
+    # via feast (setup.py)
 pymssql==2.3.2
     # via feast (setup.py)
 pymysql==1.1.1
@@ -746,7 +757,9 @@ python-dateutil==2.9.0.post0
     #   pandas
     #   trino
 python-dotenv==1.0.1
-    # via uvicorn
+    # via
+    #   environs
+    #   uvicorn
 python-json-logger==2.0.7
     # via jupyter-events
 python-keycloak==4.2.2
@@ -847,6 +860,7 @@ setuptools==75.6.0
     #   jupyterlab
     #   kubernetes
     #   pip-tools
+    #   pymilvus
     #   singlestoredb
 singlestoredb==1.7.2
     # via feast (setup.py)
@@ -947,6 +961,7 @@ tqdm==4.67.1
     # via
     #   feast (setup.py)
     #   great-expectations
+    #   milvus-lite
 traitlets==5.14.3
     # via
     #   comm
@@ -1028,6 +1043,8 @@ tzlocal==5.2
     # via
     #   great-expectations
     #   trino
+ujson==5.10.0
+    # via pymilvus
 uri-template==1.3.0
     # via jsonschema
 urllib3==1.26.20
diff --git a/sdk/python/requirements/py3.9-requirements.txt b/sdk/python/requirements/py3.9-requirements.txt
@@ -97,7 +97,7 @@ pyarrow==18.0.0
     # via
     #   feast (setup.py)
     #   dask-expr
-pydantic==2.10.1
+pydantic==2.10.2
     # via
     #   feast (setup.py)
     #   fastapi
diff --git a/sdk/python/tests/integration/feature_repos/universal/online_store/milvus.py b/sdk/python/tests/integration/feature_repos/universal/online_store/milvus.py
@@ -0,0 +1,35 @@
+from typing import Any, Dict
+
+from testcontainers.milvus import MilvusContainer
+
+from tests.integration.feature_repos.universal.online_store_creator import (
+    OnlineStoreCreator,
+)
+
+
+class MilvusOnlineStoreCreator(OnlineStoreCreator):
+    def __init__(self, project_name: str, **kwargs):
+        super().__init__(project_name)
+        self.fixed_port = 19530
+        self.container = MilvusContainer("milvusdb/milvus:v2.4.4").with_exposed_ports(
+            self.fixed_port
+        )
+
+    def create_online_store(self) -> Dict[str, Any]:
+        self.container.start()
+        # Wait for Milvus server to be ready
+        host = "localhost"
+        port = self.container.get_exposed_port(self.fixed_port)
+        return {
+            "type": "milvus",
+            "host": host,
+            "port": int(port),
+            "index_type": "IVF_FLAT",
+            "metric_type": "L2",
+            "embedding_dim": 2,
+            "vector_enabled": True,
+            "nlist": 1,
+        }
+
+    def teardown(self):
+        self.container.stop()
diff --git a/setup.py b/setup.py
