|
| 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 | + } |
0 commit comments