Build cagra index (iter1)

mayya-sharipova · mayya-sharipova · commit 8cee75e07148 · 2025-06-27T15:57:49.000-04:00
diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/GPUVectorsFormat.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/GPUVectorsFormat.java
@@ -77,18 +77,19 @@ static GPUVectorsReader getGPUReader(KnnVectorsReader vectorsReader, String fiel
     }
 
     /** Tells whether the platform supports cuvs. */
-    public static boolean supported() {
-        try (var resources = CuVSResources.create()) {
-            return true;
+    public static CuVSResources cuVSResourcesOrNull() {
+        try {
+            var resources = CuVSResources.create();
+            return resources;
         } catch (UnsupportedOperationException uoe) {
             var msg = uoe.getMessage() == null ? "" : ": " + uoe.getMessage();
-            LOG.warn("cuvs is not supported on this platform or java version" + msg);
+            LOG.warn("GPU based vector search is not supported on this platform or java version" + msg);
         } catch (Throwable t) {
             if (t instanceof ExceptionInInitializerError ex) {
                 t = ex.getCause();
             }
             LOG.warn("Exception occurred during creation of cuvs resources. " + t);
         }
-        return false;
+        return null;
     }
 }
diff --git a/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/GPUVectorsWriter.java b/x-pack/plugin/gpu/src/main/java/org/elasticsearch/xpack/gpu/codec/GPUVectorsWriter.java
@@ -7,6 +7,10 @@
 
 package org.elasticsearch.xpack.gpu.codec;
 
+import com.nvidia.cuvs.CagraIndex;
+import com.nvidia.cuvs.CagraIndexParams;
+import com.nvidia.cuvs.CuVSResources;
+
 import org.apache.lucene.codecs.CodecUtil;
 import org.apache.lucene.codecs.KnnFieldVectorsWriter;
 import org.apache.lucene.codecs.KnnVectorsWriter;
@@ -15,33 +19,46 @@
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.FloatVectorValues;
 import org.apache.lucene.index.IndexFileNames;
+import org.apache.lucene.index.KnnVectorValues;
 import org.apache.lucene.index.MergeState;
 import org.apache.lucene.index.SegmentWriteState;
 import org.apache.lucene.index.Sorter;
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
 import org.apache.lucene.store.IndexOutput;
+import org.elasticsearch.common.lucene.store.IndexOutputOutputStream;
 import org.elasticsearch.core.IOUtils;
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 
 import static org.apache.lucene.codecs.lucene99.Lucene99HnswVectorsReader.SIMILARITY_FUNCTIONS;
+import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
 
 /**
  * Writer for GPU-accelerated vectors.
  */
 public class GPUVectorsWriter extends KnnVectorsWriter {
+    private static final Logger logger = LogManager.getLogger(GPUVectorsWriter.class);
+    // 2 for now based on https://github.com/rapidsai/cuvs/issues/666, but can be increased later
+    private static final int MIN_NUM_VECTORS_FOR_GPU_BUILD = 2;
 
     private final List<FieldWriter> fieldWriters = new ArrayList<>();
     private final IndexOutput gpuIdx;
     private final IndexOutput gpuMeta;
     private final FlatVectorsWriter rawVectorDelegate;
     private final SegmentWriteState segmentWriteState;
+    private final CuVSResources cuVSResources;
 
     @SuppressWarnings("this-escape")
     public GPUVectorsWriter(SegmentWriteState state, FlatVectorsWriter rawVectorDelegate) throws IOException {
+        this.cuVSResources = GPUVectorsFormat.cuVSResourcesOrNull();
+        if (cuVSResources == null) {
+            throw new IllegalArgumentException("GPU based vector search is not supported on this platform or java version");
+        }
         this.segmentWriteState = state;
         this.rawVectorDelegate = rawVectorDelegate;
         final String metaFileName = IndexFileNames.segmentFileName(
@@ -95,30 +112,89 @@ public final KnnFieldVectorsWriter<?> addField(FieldInfo fieldInfo) throws IOExc
     @Override
     public final void flush(int maxDoc, Sorter.DocMap sortMap) throws IOException {
         rawVectorDelegate.flush(maxDoc, sortMap);
+        // TODO: implement the case when sortMap != null
+
         for (FieldWriter fieldWriter : fieldWriters) {
-            // TODO: Implement GPU-specific vector merging instead of bogus implementation
+            // TODO: can we use MemorySegment instead of passing array of vectors
+            float[][] vectors = fieldWriter.delegate.getVectors().toArray(float[][]::new);
             long dataOffset = gpuIdx.alignFilePointer(Float.BYTES);
-            var vectors = fieldWriter.delegate.getVectors();
-            for (int i = 0; i < vectors.size(); i++) {
-                gpuIdx.writeVInt(0);
+            try {
+                buildAndwriteGPUIndex(fieldWriter.fieldInfo.getVectorSimilarityFunction(), vectors);
+                long dataLength = gpuIdx.getFilePointer() - dataOffset;
+                writeMeta(fieldWriter.fieldInfo, dataOffset, dataLength);
+            } catch (IOException e) {
+                throw e;
+            } catch (Throwable t) {
+                throw new IOException("Failed to write GPU index: ", t);
+            }
+        }
+    }
+
+    private void buildAndwriteGPUIndex(VectorSimilarityFunction similarityFunction, float[][] vectors) throws Throwable {
+        // TODO: should we Lucene HNSW index write here
+        if (vectors.length < MIN_NUM_VECTORS_FOR_GPU_BUILD) {
+            if (logger.isDebugEnabled()) {
+                logger.debug("Skip building carga index; vectors length {} < {}", vectors.length, MIN_NUM_VECTORS_FOR_GPU_BUILD);
+            }
+            return;
+        }
+
+        CagraIndexParams.CuvsDistanceType distanceType = switch (similarityFunction) {
+            case EUCLIDEAN -> CagraIndexParams.CuvsDistanceType.L2Expanded;
+            case DOT_PRODUCT, MAXIMUM_INNER_PRODUCT -> CagraIndexParams.CuvsDistanceType.InnerProduct;
+            case COSINE -> CagraIndexParams.CuvsDistanceType.CosineExpanded;
+        };
+
+        // TODO: expose cagra index params of intermediate graph degree, graph degre, algorithm, NNDescentNumIterations
+        CagraIndexParams params = new CagraIndexParams.Builder().withNumWriterThreads(1) // TODO: how many CPU threads we can use?
+            .withCagraGraphBuildAlgo(CagraIndexParams.CagraGraphBuildAlgo.NN_DESCENT)
+            .withMetric(distanceType)
+            .build();
+
+        // build index on GPU
+        long startTime = System.nanoTime();
+        var index = CagraIndex.newBuilder(cuVSResources).withDataset(vectors).withIndexParams(params).build();
+        if (logger.isDebugEnabled()) {
+            logger.debug("Carga index created in: {} ms; #num vectors: {}", (System.nanoTime() - startTime) / 1_000_000.0, vectors.length);
+        }
+
+        // TODO: do serialization through MemorySegment instead of a temp file
+        // serialize index for CPU consumption
+        startTime = System.nanoTime();
+        var gpuIndexOutputStream = new IndexOutputOutputStream(gpuIdx);
+        try {
+            index.serialize(gpuIndexOutputStream);
+            if (logger.isDebugEnabled()) {
+                logger.debug("Carga index serialized in: {} ms", (System.nanoTime() - startTime) / 1_000_000.0);
             }
-            long dataLength = gpuIdx.getFilePointer() - dataOffset;
-            writeMeta(fieldWriter.fieldInfo, dataOffset, dataLength);
+        } finally {
+            index.destroyIndex();
         }
     }
 
     @Override
     public final void mergeOneField(FieldInfo fieldInfo, MergeState mergeState) throws IOException {
         if (fieldInfo.getVectorEncoding().equals(VectorEncoding.FLOAT32)) {
             rawVectorDelegate.mergeOneField(fieldInfo, mergeState);
-            // TODO: Implement GPU-specific vector merging instead of bogus implementation
-            FloatVectorValues floatVectorValues = KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState);
+            FloatVectorValues vectorValues = KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState);
+            // TODO: more efficient way to pass merged vector values to gpuIndex construction
+            KnnVectorValues.DocIndexIterator iter = vectorValues.iterator();
+            List<float[]> vectorList = new ArrayList<>();
+            for (int docV = iter.nextDoc(); docV != NO_MORE_DOCS; docV = iter.nextDoc()) {
+                vectorList.add(vectorValues.vectorValue(iter.index()));
+            }
+            float[][] vectors = vectorList.toArray(new float[0][]);
+
             long dataOffset = gpuIdx.alignFilePointer(Float.BYTES);
-            for (int i = 0; i < floatVectorValues.size(); i++) {
-                gpuIdx.writeVInt(0);
+            try {
+                buildAndwriteGPUIndex(fieldInfo.getVectorSimilarityFunction(), vectors);
+                long dataLength = gpuIdx.getFilePointer() - dataOffset;
+                writeMeta(fieldInfo, dataOffset, dataLength);
+            } catch (IOException e) {
+                throw e;
+            } catch (Throwable t) {
+                throw new IOException("Failed to write GPU index: ", t);
             }
-            long dataLength = gpuIdx.getFilePointer() - dataOffset;
-            writeMeta(fieldInfo, dataOffset, dataLength);
         } else {
             rawVectorDelegate.mergeOneField(fieldInfo, mergeState);
         }
@@ -157,6 +233,7 @@ public final void finish() throws IOException {
     @Override
     public final void close() throws IOException {
         IOUtils.close(rawVectorDelegate, gpuMeta, gpuIdx);
+        cuVSResources.close();
     }
 
     @Override
diff --git a/x-pack/plugin/gpu/src/main/plugin-metadata/entitlement-policy.yaml b/x-pack/plugin/gpu/src/main/plugin-metadata/entitlement-policy.yaml
@@ -0,0 +1,2 @@
+com.nvidia.cuvs:
+  - load_native_libraries
diff --git a/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/GPUVectorsFormatTests.java b/x-pack/plugin/gpu/src/test/java/org/elasticsearch/xpack/gpu/codec/GPUVectorsFormatTests.java
@@ -25,7 +25,7 @@ public class GPUVectorsFormatTests extends BaseKnnVectorsFormatTestCase {
 
     @BeforeClass
     public static void beforeClass() {
-        assumeTrue("cuvs not supported", GPUVectorsFormat.supported());
+        assumeTrue("cuvs not supported", GPUVectorsFormat.cuVSResourcesOrNull() != null);
     }
 
     static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new GPUVectorsFormat());
diff --git a/x-pack/plugin/gpu/src/yamlRestTest/java/org/elasticsearch/xpack/gpu/GPUClientYamlTestSuiteIT.java b/x-pack/plugin/gpu/src/yamlRestTest/java/org/elasticsearch/xpack/gpu/GPUClientYamlTestSuiteIT.java
@@ -19,7 +19,7 @@ public class GPUClientYamlTestSuiteIT extends ESClientYamlSuiteTestCase {
     public static ElasticsearchCluster cluster = ElasticsearchCluster.local()
         .nodes(1)
         .module("gpu")
-        .setting("xpack.license.self_generated.type", "basic")
+        .setting("xpack.license.self_generated.type", "trial")
         .setting("xpack.security.enabled", "false")
         .build();
 

Original file line number	Diff line number	Diff line change
`@@ -77,18 +77,19 @@ static GPUVectorsReader getGPUReader(KnnVectorsReader vectorsReader, String fiel`
`77`	`77`	`}`
`78`	`78`
`79`	`79`	`/** Tells whether the platform supports cuvs. */`
`80`		`- public static boolean supported() {`
`81`		`- try (var resources = CuVSResources.create()) {`
`82`		`- return true;`
	`80`	`+ public static CuVSResources cuVSResourcesOrNull() {`
	`81`	`+ try {`
	`82`	`+ var resources = CuVSResources.create();`
	`83`	`+ return resources;`
`83`	`84`	`} catch (UnsupportedOperationException uoe) {`
`84`	`85`	`var msg = uoe.getMessage() == null ? "" : ": " + uoe.getMessage();`
`85`		`- LOG.warn("cuvs is not supported on this platform or java version" + msg);`
	`86`	`+ LOG.warn("GPU based vector search is not supported on this platform or java version" + msg);`
`86`	`87`	`} catch (Throwable t) {`
`87`	`88`	`if (t instanceof ExceptionInInitializerError ex) {`
`88`	`89`	`t = ex.getCause();`
`89`	`90`	`}`
`90`	`91`	`LOG.warn("Exception occurred during creation of cuvs resources. " + t);`
`91`	`92`	`}`
`92`		`- return false;`
	`93`	`+ return null;`
`93`	`94`	`}`
`94`	`95`	`}`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+com.nvidia.cuvs:`
	`2`	`+ - load_native_libraries`
Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,7 @@ public class GPUVectorsFormatTests extends BaseKnnVectorsFormatTestCase {`
`25`	`25`
`26`	`26`	`@BeforeClass`
`27`	`27`	`public static void beforeClass() {`
`28`		`- assumeTrue("cuvs not supported", GPUVectorsFormat.supported());`
	`28`	`+ assumeTrue("cuvs not supported", GPUVectorsFormat.cuVSResourcesOrNull() != null);`
`29`	`29`	`}`
`30`	`30`
`31`	`31`	`static final Codec codec = TestUtil.alwaysKnnVectorsFormat(new GPUVectorsFormat());`