add in gpu milvus (#166)

NVIDIA · Oct 16, 2024 · 733f40f · 733f40f
1 parent ac81779
commit 733f40f
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 40 deletions.
diff --git a/docker-compose.yaml b/docker-compose.yaml
@@ -19,13 +19,13 @@ services:
       - NIM_HTTP_API_PORT=8000
       - NIM_TRITON_LOG_VERBOSE=1
       - NGC_API_KEY=${NIM_NGC_API_KEY:-${NGC_API_KEY:-ngcapikey}}
-      - CUDA_VISIBLE_DEVICES=1
+      - CUDA_VISIBLE_DEVICES=0
     deploy:
       resources:
         reservations:
           devices:
             - driver: nvidia
-              count: all
+              device_ids: ["1"]
               capabilities: [gpu]
     runtime: nvidia
 
@@ -46,7 +46,7 @@ services:
         reservations:
           devices:
             - driver: nvidia
-              count: all
+              device_ids: ["0"]
               capabilities: [gpu]
     runtime: nvidia
 
@@ -62,13 +62,13 @@ services:
       - NIM_HTTP_API_PORT=8000
       - NIM_TRITON_LOG_VERBOSE=1
       - NGC_API_KEY=${NIM_NGC_API_KEY:-${NGC_API_KEY:-ngcapikey}}
-      - CUDA_VISIBLE_DEVICES=1
+      - CUDA_VISIBLE_DEVICES=0
     deploy:
       resources:
         reservations:
           devices:
             - driver: nvidia
-              count: all
+              device_ids: ["1"]
               capabilities: [gpu]
     runtime: nvidia
 
@@ -84,13 +84,13 @@ services:
       - NIM_HTTP_API_PORT=8000
       - NIM_TRITON_LOG_VERBOSE=1
       - NGC_API_KEY=${NIM_NGC_API_KEY:-${NGC_API_KEY:-ngcapikey}}
-      - CUDA_VISIBLE_DEVICES=1
+      - CUDA_VISIBLE_DEVICES=0
     deploy:
       resources:
         reservations:
           devices:
             - driver: nvidia
-              count: all
+              device_ids: ["1"]
               capabilities: [gpu]
     runtime: nvidia
 
@@ -106,13 +106,13 @@ services:
       - NIM_HTTP_API_PORT=8000
       - NIM_TRITON_LOG_VERBOSE=1
       - NGC_API_KEY=${NIM_NGC_API_KEY:-${NGC_API_KEY:-ngcapikey}}
-      - CUDA_VISIBLE_DEVICES=1
+      - CUDA_VISIBLE_DEVICES=0
     deploy:
       resources:
         reservations:
           devices:
             - driver: nvidia
-              count: all
+              device_ids: ["1"]
               capabilities: [gpu]
     runtime: nvidia
 
@@ -167,9 +167,10 @@ services:
         reservations:
           devices:
             - driver: nvidia
-              count: all
+              device_ids: ["1"]
               capabilities: [gpu]
-
+    depends_on:
+      - milvus
 
   otel-collector:
     image: otel/opentelemetry-collector-contrib:0.91.0
@@ -258,32 +259,39 @@ services:
 #      timeout: 20s
 #      retries: 3
 
-#  milvus:
-#    # Turn on to leverage the `vdb_upload` task
-#    restart: always
-#    container_name: milvus-standalone
-#    image: milvusdb/milvus:v2.3.5
-#    command: ["milvus", "run", "standalone"]
-#    hostname: milvus
-#    security_opt:
-#      - seccomp:unconfined
-#    environment:
-#      ETCD_ENDPOINTS: etcd:2379
-#      MINIO_ADDRESS: minio:9000
-#    volumes:
-#      - ./.volumes/milvus:/var/lib/milvus
-#    healthcheck:
-#      test: ["CMD", "curl", "-f", "http://localhost:9091/healthz"]
-#      interval: 30s
-#      start_period: 90s
-#      timeout: 20s
-#      retries: 3
-#    ports:
-#      - "19530:19530"
-#      - "9091:9091"
-#    depends_on:
-#      - "etcd"
-#      - "minio"
+  # milvus:
+  #   # Turn on to leverage the `vdb_upload` task
+  #   restart: always
+  #   container_name: milvus-standalone
+  #   image: milvusdb/milvus:v2.4.9-gpu
+  #   command: ["milvus", "run", "standalone"]
+  #   hostname: milvus
+  #   security_opt:
+  #     - seccomp:unconfined
+  #   environment:
+  #     ETCD_ENDPOINTS: etcd:2379
+  #     MINIO_ADDRESS: minio:9000
+  #   volumes:
+  #     - ./.volumes/milvus:/var/lib/milvus
+  #   healthcheck:
+  #     test: ["CMD", "curl", "-f", "http://localhost:9091/healthz"]
+  #     interval: 30s
+  #     start_period: 90s
+  #     timeout: 20s
+  #     retries: 3
+  #   ports:
+  #     - "19530:19530"
+  #     - "9091:9091"
+  #   deploy:
+  #     resources:
+  #       reservations:
+  #         devices:
+  #           - driver: nvidia
+  #             device_ids: ["1"]
+  #             capabilities: [gpu]
+  #   depends_on:
+  #     - "etcd"
+  #     - "minio"
 
 #  attu:
 #    # Turn on to leverage the `vdb_upload` task

diff --git a/src/nv_ingest/schemas/vdb_task_sink_schema.py b/src/nv_ingest/schemas/vdb_task_sink_schema.py
@@ -37,10 +37,11 @@ def build_default_milvus_config(embedding_size: int = 1024) -> typing.Dict[str,
         "index_conf": {
             "field_name": "vector",
             "metric_type": "L2",
-            "index_type": "HNSW",
+            "index_type": "GPU_CAGRA",
             "params": {
-                "M": 8,
-                "efConstruction": 64,
+                'intermediate_graph_degree':128,
+                'graph_degree': 64,
+                "build_algo": "NN_DESCENT",
             },
         },
         "schema_conf": {