Skip to content

Commit b89fadd

Browse files
authored
fix: Updating milvus connect function to work with remote instance (feast-dev#5401)
* Updating milvus connect function to work with remote instance Signed-off-by: Fiona Waters <[email protected]> * Update test configuration to use path for db Signed-off-by: Fiona Waters <[email protected]> --------- Signed-off-by: Fiona Waters <[email protected]>
1 parent 2540846 commit b89fadd

File tree

3 files changed

+8
-28
lines changed

3 files changed

+8
-28
lines changed

sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -88,8 +88,8 @@ class MilvusOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
8888
"""
8989

9090
type: Literal["milvus"] = "milvus"
91-
path: Optional[StrictStr] = "online_store.db"
92-
host: Optional[StrictStr] = "localhost"
91+
path: Optional[StrictStr] = ""
92+
host: Optional[StrictStr] = "http://localhost"
9393
port: Optional[int] = 19530
9494
index_type: Optional[str] = "FLAT"
9595
metric_type: Optional[str] = "COSINE"
@@ -126,13 +126,16 @@ def _get_db_path(self, config: RepoConfig) -> str:
126126

127127
def _connect(self, config: RepoConfig) -> MilvusClient:
128128
if not self.client:
129-
if config.provider == "local":
129+
if config.provider == "local" and config.online_store.path:
130130
db_path = self._get_db_path(config)
131131
print(f"Connecting to Milvus in local mode using {db_path}")
132132
self.client = MilvusClient(db_path)
133133
else:
134+
print(
135+
f"Connecting to Milvus remotely at {config.online_store.host}:{config.online_store.port}"
136+
)
134137
self.client = MilvusClient(
135-
url=f"{config.online_store.host}:{config.online_store.port}",
138+
uri=f"{config.online_store.host}:{config.online_store.port}",
136139
token=f"{config.online_store.username}:{config.online_store.password}"
137140
if config.online_store.username and config.online_store.password
138141
else "",

sdk/python/tests/integration/feature_repos/repo_configuration.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@
8989
)
9090

9191
DYNAMO_CONFIG = {"type": "dynamodb", "region": "us-west-2"}
92-
MILVUS_CONFIG = {"type": "milvus", "embedding_dim": "2"}
92+
MILVUS_CONFIG = {"type": "milvus", "embedding_dim": 2, "path": "online_store.db"}
9393
REDIS_CONFIG = {"type": "redis", "connection_string": "localhost:6379,db=0"}
9494
REDIS_CLUSTER_CONFIG = {
9595
"type": "redis",
Lines changed: 0 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,5 @@
11
from typing import Any, Dict
22

3-
import docker
4-
from testcontainers.core.container import DockerContainer
5-
from testcontainers.core.waiting_utils import wait_for_logs
6-
73
from tests.integration.feature_repos.universal.online_store_creator import (
84
OnlineStoreCreator,
95
)
@@ -12,33 +8,14 @@
128
class MilvusOnlineStoreCreator(OnlineStoreCreator):
139
def __init__(self, project_name: str, **kwargs):
1410
super().__init__(project_name)
15-
self.fixed_port = 19530
16-
self.container = DockerContainer("milvusdb/milvus:v2.4.9").with_exposed_ports(
17-
self.fixed_port
18-
)
19-
self.client = docker.from_env()
2011

2112
def create_online_store(self) -> Dict[str, Any]:
22-
self.container.start()
23-
# Wait for Milvus server to be ready
24-
# log_string_to_wait_for = "Ready to accept connections"
25-
log_string_to_wait_for = ""
26-
wait_for_logs(
27-
container=self.container, predicate=log_string_to_wait_for, timeout=30
28-
)
29-
host = "localhost"
30-
port = self.container.get_exposed_port(self.fixed_port)
3113
return {
3214
"type": "milvus",
33-
"host": host,
34-
"port": int(port),
3515
"path": "online_store.db",
3616
"index_type": "IVF_FLAT",
3717
"metric_type": "L2",
3818
"embedding_dim": 2,
3919
"vector_enabled": True,
4020
"nlist": 1,
4121
}
42-
43-
def teardown(self):
44-
self.container.stop()

0 commit comments

Comments
 (0)