Use milvus gpu with GPU_IVF_FLAT flag (NVIDIA#243)

* Use GPU flag for document ingestion as default

* Port milvus changes to release branch

Author: nv-pranjald

RAG/examples/local_deploy/docker-compose-vectordb.yaml

@@ -54,7 +54,7 @@ services:
 
   milvus:
     container_name: milvus-standalone
-    image: milvusdb/milvus:v2.4.12
+    image: milvusdb/milvus:v2.4.15-gpu
     command: ["milvus", "run", "standalone"]
     environment:
       ETCD_ENDPOINTS: etcd:2379
@@ -74,6 +74,13 @@ services:
     depends_on:
       - "etcd"
       - "minio"
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              capabilities: ["gpu"]
+              device_ids: ['${VECTORSTORE_GPU_DEVICE_ID:-0}']
     profiles: ["nemo-retriever", "milvus", ""]
 
   elasticsearch:

RAG/src/chain_server/configuration.py

@@ -40,7 +40,7 @@ class VectorStoreConfig(ConfigWizard):
         "nprobe", default=16, help_txt="Number of units to query",  # IVF Flat milvus
     )
     index_type: str = configfield(
-        "index_type", default="IVF_FLAT", help_txt="Index of the vector db",  # IVF Flat for milvus
+        "index_type", default="GPU_IVF_FLAT", help_txt="Index of the vector db",  # IVF Flat for milvus
     )
 
 

RAG/src/chain_server/utils.py

@@ -314,13 +314,17 @@ def create_vectorstore_langchain(document_embedder: "Embeddings", collection_nam
         )
     elif config.vector_store.name == "milvus":
         logger.info(f"Using milvus collection: {collection_name}")
-        # vectorstore url can be updated using environment variable APP_VECTORSTORE_URL, it should be in http://ip:port format
+        if not collection_name:
+            collection_name = os.getenv('COLLECTION_NAME', "vector_db")
+        logger.info(f"Using milvus collection: {collection_name}")
         url = urlparse(config.vector_store.url)
         vectorstore = Milvus(
             document_embedder,
             connection_args={"host": url.hostname, "port": url.port},
             collection_name=collection_name,
-            auto_id=True,
+            index_params={"index_type": config.vector_store.index_type, "metric_type": "L2", "nlist": config.vector_store.nlist},
+            search_params={"nprobe": config.vector_store.nprobe},
+            auto_id = True
         )
     else:
         raise ValueError(f"{config.vector_store.name} vector database is not supported")