Skip to content

Commit f5797cd

Browse files
feat: Add support for testing remote-milvus provider (#712) (#722)
* feat: Add remote milvus provider for vector_stores testing This PR includes the work done by @Bobbins228 in his PR #652 (I had problems rebasing from the main branch, so I decided to send a new PR) In addition to Mark's work, I've refactored the code to introduce the vector_io_provider parameter. Using the vector_io_provider parameter it is now possible to run vector_stores tests with multiple providers. Their deployment code is now in the tests/fixtures/vector_io pytest plugin * fix: Parametrize MILVUS_IMAGE and MILVUS_TOKEN * feat: Explain in README.md how to test new providers + small fixes * fix: remove usage of deprecated MILVUS_DB_PATH env var It is not used in the Red Hat Distribution since RHOAI 2.23 * fix: Add wait_for_resource=True when creating milvus related services --------- Signed-off-by: Jorge Garcia Oncins <jgarciao@redhat.com> Co-authored-by: Jorge <jgarciao@users.noreply.github.com>
1 parent 3274497 commit f5797cd

File tree

5 files changed

+371
-42
lines changed

5 files changed

+371
-42
lines changed

tests/conftest.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,12 @@
6161

6262
LOGGER = get_logger(name=__name__)
6363

64-
pytest_plugins = ["tests.fixtures.inference", "tests.fixtures.guardrails", "tests.fixtures.trustyai"]
64+
pytest_plugins = [
65+
"tests.fixtures.inference",
66+
"tests.fixtures.guardrails",
67+
"tests.fixtures.trustyai",
68+
"tests.fixtures.vector_io",
69+
]
6570

6671

6772
@pytest.fixture(scope="session")

tests/fixtures/vector_io.py

Lines changed: 239 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,239 @@
1+
from typing import Generator, Any, Callable, Dict
2+
import pytest
3+
import os
4+
import secrets
5+
from _pytest.fixtures import FixtureRequest
6+
from kubernetes.dynamic import DynamicClient
7+
from utilities.constants import Timeout
8+
9+
10+
from ocp_resources.deployment import Deployment
11+
12+
from ocp_resources.namespace import Namespace
13+
from ocp_resources.service import Service
14+
15+
16+
MILVUS_IMAGE = os.getenv(
17+
"LLS_VECTOR_IO_MILVUS_IMAGE",
18+
"docker.io/milvusdb/milvus@sha256:3d772c3eae3a6107b778636cea5715b9353360b92e5dcfdcaf4ca7022f4f497c", # Milvus 2.6.3
19+
)
20+
MILVUS_TOKEN = os.getenv("LLS_VECTOR_IO_MILVUS_TOKEN", secrets.token_urlsafe(32))
21+
ETCD_IMAGE = os.getenv(
22+
"LLS_VECTOR_IO_ETCD_IMAGE",
23+
"quay.io/coreos/etcd@sha256:3397341272b9e0a6f44d7e3fc7c321c6efe6cbe82ce866b9b01d0c704bfc5bf3", # etcd v3.6.5
24+
)
25+
26+
27+
@pytest.fixture(scope="class")
28+
def vector_io_provider_deployment_config_factory(
29+
request: FixtureRequest,
30+
) -> Callable[[str], list[Dict[str, str]]]:
31+
"""
32+
Factory fixture for deploying vector I/O providers and returning their configuration.
33+
34+
This fixture returns a factory function that can deploy different vector I/O providers
35+
(such as Milvus) in the cluster and return the necessary environment variables
36+
for configuring the LlamaStack server to use these providers.
37+
38+
Args:
39+
request: Pytest fixture request object for accessing other fixtures
40+
41+
Returns:
42+
Callable[[str], list[Dict[str, str]]]: Factory function that takes a provider name
43+
and returns a list of environment variable dictionaries
44+
45+
Supported Providers:
46+
- "milvus" (or None): Local Milvus instance with embedded database
47+
- "milvus-remote": Remote Milvus service requiring external deployment
48+
49+
Environment Variables by Provider:
50+
- "milvus": no env vars available
51+
- "milvus-remote":
52+
* MILVUS_ENDPOINT: Remote Milvus service endpoint URL
53+
* MILVUS_TOKEN: Authentication token for remote service
54+
* MILVUS_CONSISTENCY_LEVEL: Consistency level for operations
55+
56+
Example:
57+
def test_with_milvus(vector_io_provider_deployment_config_factory):
58+
env_vars = vector_io_provider_deployment_config_factory("milvus-remote")
59+
# env_vars contains MILVUS_ENDPOINT, MILVUS_TOKEN, etc.
60+
"""
61+
62+
def _factory(provider_name: str) -> list[Dict[str, str]]:
63+
env_vars: list[dict[str, str]] = []
64+
65+
if provider_name is None or provider_name == "milvus":
66+
# Default case - no additional environment variables needed
67+
pass
68+
elif provider_name == "milvus-remote":
69+
request.getfixturevalue(argname="milvus_service")
70+
env_vars.append({"name": "MILVUS_ENDPOINT", "value": "http://vector-io-milvus-service:19530"})
71+
env_vars.append({"name": "MILVUS_TOKEN", "value": MILVUS_TOKEN})
72+
env_vars.append({"name": "MILVUS_CONSISTENCY_LEVEL", "value": "Bounded"})
73+
74+
return env_vars
75+
76+
return _factory
77+
78+
79+
@pytest.fixture(scope="class")
80+
def etcd_deployment(
81+
unprivileged_client: DynamicClient,
82+
unprivileged_model_namespace: Namespace,
83+
) -> Generator[Deployment, Any, Any]:
84+
"""Deploy an etcd instance for vector I/O provider testing."""
85+
with Deployment(
86+
client=unprivileged_client,
87+
namespace=unprivileged_model_namespace.name,
88+
name="vector-io-etcd-deployment",
89+
replicas=1,
90+
selector={"matchLabels": {"app": "etcd"}},
91+
strategy={"type": "Recreate"},
92+
template=get_etcd_deployment_template(),
93+
teardown=True,
94+
) as deployment:
95+
deployment.wait_for_replicas(deployed=True, timeout=Timeout.TIMEOUT_2MIN)
96+
yield deployment
97+
98+
99+
@pytest.fixture(scope="class")
100+
def etcd_service(
101+
unprivileged_client: DynamicClient,
102+
unprivileged_model_namespace: Namespace,
103+
) -> Generator[Service, Any, Any]:
104+
"""Create a service for the etcd deployment."""
105+
with Service(
106+
client=unprivileged_client,
107+
namespace=unprivileged_model_namespace.name,
108+
name="vector-io-etcd-service",
109+
ports=[
110+
{
111+
"port": 2379,
112+
"targetPort": 2379,
113+
}
114+
],
115+
selector={"app": "etcd"},
116+
wait_for_resource=True,
117+
) as service:
118+
yield service
119+
120+
121+
@pytest.fixture(scope="class")
122+
def remote_milvus_deployment(
123+
unprivileged_client: DynamicClient,
124+
unprivileged_model_namespace: Namespace,
125+
etcd_deployment: Deployment,
126+
etcd_service: Service,
127+
) -> Generator[Deployment, Any, Any]:
128+
"""Deploy a remote Milvus instance for vector I/O provider testing."""
129+
with Deployment(
130+
client=unprivileged_client,
131+
namespace=unprivileged_model_namespace.name,
132+
name="vector-io-milvus-deployment",
133+
replicas=1,
134+
selector={"matchLabels": {"app": "milvus-standalone"}},
135+
strategy={"type": "Recreate"},
136+
template=get_milvus_deployment_template(),
137+
teardown=True,
138+
) as deployment:
139+
deployment.wait_for_replicas(deployed=True, timeout=Timeout.TIMEOUT_2MIN)
140+
yield deployment
141+
142+
143+
@pytest.fixture(scope="class")
144+
def milvus_service(
145+
unprivileged_client: DynamicClient,
146+
unprivileged_model_namespace: Namespace,
147+
remote_milvus_deployment: Deployment,
148+
) -> Generator[Service, Any, Any]:
149+
"""Create a service for the remote Milvus deployment."""
150+
with Service(
151+
client=unprivileged_client,
152+
namespace=unprivileged_model_namespace.name,
153+
name="vector-io-milvus-service",
154+
ports=[
155+
{
156+
"name": "grpc",
157+
"port": 19530,
158+
"targetPort": 19530,
159+
},
160+
],
161+
selector={"app": "milvus-standalone"},
162+
wait_for_resource=True,
163+
) as service:
164+
yield service
165+
166+
167+
def get_milvus_deployment_template() -> Dict[str, Any]:
168+
"""Return the Kubernetes deployment template for Milvus standalone."""
169+
return {
170+
"metadata": {"labels": {"app": "milvus-standalone"}},
171+
"spec": {
172+
"containers": [
173+
{
174+
"name": "milvus-standalone",
175+
"image": MILVUS_IMAGE,
176+
"args": ["milvus", "run", "standalone"],
177+
"ports": [{"containerPort": 19530, "protocol": "TCP"}],
178+
"volumeMounts": [
179+
{
180+
"name": "milvus-data",
181+
"mountPath": "/var/lib/milvus",
182+
}
183+
],
184+
"env": [
185+
{"name": "DEPLOY_MODE", "value": "standalone"},
186+
{"name": "ETCD_ENDPOINTS", "value": "vector-io-etcd-service:2379"},
187+
{"name": "MINIO_ADDRESS", "value": ""},
188+
{"name": "COMMON_STORAGETYPE", "value": "local"},
189+
],
190+
}
191+
],
192+
"volumes": [
193+
{
194+
"name": "milvus-data",
195+
"emptyDir": {},
196+
}
197+
],
198+
},
199+
}
200+
201+
202+
def get_etcd_deployment_template() -> Dict[str, Any]:
203+
"""Return the Kubernetes deployment template for etcd."""
204+
return {
205+
"metadata": {"labels": {"app": "etcd"}},
206+
"spec": {
207+
"containers": [
208+
{
209+
"name": "etcd",
210+
"image": ETCD_IMAGE,
211+
"command": [
212+
"etcd",
213+
"--advertise-client-urls=http://vector-io-etcd-service:2379",
214+
"--listen-client-urls=http://0.0.0.0:2379",
215+
"--data-dir=/etcd",
216+
],
217+
"ports": [{"containerPort": 2379}],
218+
"volumeMounts": [
219+
{
220+
"name": "etcd-data",
221+
"mountPath": "/etcd",
222+
}
223+
],
224+
"env": [
225+
{"name": "ETCD_AUTO_COMPACTION_MODE", "value": "revision"},
226+
{"name": "ETCD_AUTO_COMPACTION_RETENTION", "value": "1000"},
227+
{"name": "ETCD_QUOTA_BACKEND_BYTES", "value": "4294967296"},
228+
{"name": "ETCD_SNAPSHOT_COUNT", "value": "50000"},
229+
],
230+
}
231+
],
232+
"volumes": [
233+
{
234+
"name": "etcd-data",
235+
"emptyDir": {},
236+
}
237+
],
238+
},
239+
}

tests/llama_stack/README.md

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,11 @@ def test_vector_stores_functionality():
2727
# Test implementation
2828
```
2929

30+
## Adding Support for New API Providers
31+
32+
To add support for testing new LlamaStack API providers (e.g., a new vector_io provider), create deployment fixtures in the appropriate `/tests/fixtures/` file, update the corresponding provider factory function to return the required environment variables, and add the new provider as a test parameter in the relevant test files. For example, to add a new vector_io provider, add deployment fixtures in `/tests/fixtures/vector_io.py`, update the `vector_io_provider_deployment_config_factory` function, and add a new `pytest.param` entry in `/tests/llama_stack/vector_io/test_vector_stores.py`.
33+
34+
3035
### Available Team Markers (to be expanded)
3136

3237
- `@pytest.mark.llama_stack` - LlamaStack Core team tests
@@ -40,10 +45,13 @@ def test_vector_stores_functionality():
4045

4146
LlamaStack tests require setting the following environment variables (for example in a .env file at the root folder):
4247
```bash
43-
OC_BINARY_PATH=/usr/local/sbin/oc
48+
OC_BINARY_PATH=/usr/local/sbin/oc # Optional
4449
LLS_CORE_VLLM_URL=<LLAMA-3.2-3b-ENDPOINT>/v1 (ends with /v1)
4550
LLS_CORE_INFERENCE_MODEL=<LLAMA-3.2-3b-MODEL_NAME>
4651
LLS_CORE_VLLM_API_TOKEN=<LLAMA-3.2-3b-TOKEN>
52+
LLS_VECTOR_IO_MILVUS_IMAGE=<CUSTOM-MILVUS-IMAGE> # Optional
53+
LLS_VECTOR_IO_MILVUS_TOKEN=<CUSTOM-MILVUS-TOKEN> # Optional
54+
LLS_VECTOR_IO_ETCD_IMAGE=<CUSTOM-ETCD-IMAGE> # Optional
4755
```
4856

4957
### Run All Llama Stack Tests

0 commit comments

Comments
 (0)